#![warn(missing_docs)]
#![warn(clippy::all)]
#![allow(hidden_glob_reexports)]
pub use ringkernel_core::*;
pub use ringkernel_derive::*;
pub use ringkernel_cpu::CpuRuntime;
#[cfg(feature = "cuda")]
pub use ringkernel_cuda::CudaRuntime;
#[cfg(feature = "wgpu")]
pub use ringkernel_wgpu::WgpuRuntime;
#[cfg(feature = "metal")]
pub use ringkernel_metal::MetalRuntime;
pub use ringkernel_codegen as codegen;
use ringkernel_core::error::Result;
use ringkernel_core::error::RingKernelError;
use ringkernel_core::runtime::Backend;
use ringkernel_core::runtime::KernelHandle;
use ringkernel_core::runtime::KernelId;
use ringkernel_core::runtime::LaunchOptions;
use ringkernel_core::runtime::RingKernelRuntime;
use ringkernel_core::runtime::RuntimeBuilder;
use ringkernel_core::runtime::RuntimeMetrics;
pub mod prelude {
pub use crate::RingKernel;
pub use ringkernel_core::prelude::*;
pub use ringkernel_derive::*;
}
#[doc(hidden)]
pub mod __private {
use ringkernel_core::types::KernelMode;
pub struct KernelRegistration {
pub id: &'static str,
pub mode: KernelMode,
pub grid_size: u32,
pub block_size: u32,
}
inventory::collect!(KernelRegistration);
}
pub struct RingKernel {
inner: Box<dyn RingKernelRuntime>,
}
impl RingKernel {
pub fn builder() -> RingKernelBuilder {
RingKernelBuilder::new()
}
pub async fn new() -> Result<Self> {
Self::builder().build().await
}
pub async fn with_backend(backend: Backend) -> Result<Self> {
Self::builder().backend(backend).build().await
}
pub fn backend(&self) -> Backend {
self.inner.backend()
}
pub fn is_backend_available(&self, backend: Backend) -> bool {
self.inner.is_backend_available(backend)
}
pub async fn launch(&self, kernel_id: &str, options: LaunchOptions) -> Result<KernelHandle> {
self.inner.launch(kernel_id, options).await
}
pub fn get_kernel(&self, kernel_id: &KernelId) -> Option<KernelHandle> {
self.inner.get_kernel(kernel_id)
}
pub fn list_kernels(&self) -> Vec<KernelId> {
self.inner.list_kernels()
}
pub fn metrics(&self) -> RuntimeMetrics {
self.inner.metrics()
}
pub async fn shutdown(self) -> Result<()> {
self.inner.shutdown().await
}
}
pub struct RingKernelBuilder {
inner: RuntimeBuilder,
}
impl RingKernelBuilder {
pub fn new() -> Self {
Self {
inner: RuntimeBuilder::new(),
}
}
pub fn backend(mut self, backend: Backend) -> Self {
self.inner = self.inner.backend(backend);
self
}
pub fn device(mut self, index: usize) -> Self {
self.inner = self.inner.device(index);
self
}
pub fn debug(mut self, enable: bool) -> Self {
self.inner = self.inner.debug(enable);
self
}
pub fn profiling(mut self, enable: bool) -> Self {
self.inner = self.inner.profiling(enable);
self
}
pub async fn build(self) -> Result<RingKernel> {
let backend = self.inner.backend;
let inner: Box<dyn RingKernelRuntime> = match backend {
Backend::Auto => self.build_auto().await?,
Backend::Cpu => Box::new(CpuRuntime::new().await?),
#[cfg(feature = "cuda")]
Backend::Cuda => Box::new(ringkernel_cuda::CudaRuntime::new().await?),
#[cfg(not(feature = "cuda"))]
Backend::Cuda => {
return Err(RingKernelError::BackendUnavailable(
"CUDA feature not enabled".to_string(),
))
}
#[cfg(feature = "metal")]
Backend::Metal => Box::new(ringkernel_metal::MetalRuntime::new().await?),
#[cfg(not(feature = "metal"))]
Backend::Metal => {
return Err(RingKernelError::BackendUnavailable(
"Metal feature not enabled".to_string(),
))
}
#[cfg(feature = "wgpu")]
Backend::Wgpu => Box::new(ringkernel_wgpu::WgpuRuntime::new().await?),
#[cfg(not(feature = "wgpu"))]
Backend::Wgpu => {
return Err(RingKernelError::BackendUnavailable(
"WebGPU feature not enabled".to_string(),
))
}
};
Ok(RingKernel { inner })
}
async fn build_auto(self) -> Result<Box<dyn RingKernelRuntime>> {
#[cfg(feature = "cuda")]
if ringkernel_cuda::is_cuda_available() {
tracing::info!("Auto-selected CUDA backend");
return Ok(Box::new(ringkernel_cuda::CudaRuntime::new().await?));
}
#[cfg(feature = "metal")]
if ringkernel_metal::is_metal_available() {
tracing::info!("Auto-selected Metal backend");
return Ok(Box::new(ringkernel_metal::MetalRuntime::new().await?));
}
#[cfg(feature = "wgpu")]
if ringkernel_wgpu::is_wgpu_available() {
tracing::info!("Auto-selected WebGPU backend");
return Ok(Box::new(ringkernel_wgpu::WgpuRuntime::new().await?));
}
tracing::info!("Auto-selected CPU backend (no GPU available)");
Ok(Box::new(CpuRuntime::new().await?))
}
}
impl Default for RingKernelBuilder {
fn default() -> Self {
Self::new()
}
}
pub fn registered_kernels() -> Vec<&'static str> {
inventory::iter::<__private::KernelRegistration>
.into_iter()
.map(|r| r.id)
.collect()
}
pub mod availability {
pub fn cuda() -> bool {
#[cfg(feature = "cuda")]
{
ringkernel_cuda::is_cuda_available()
}
#[cfg(not(feature = "cuda"))]
{
false
}
}
pub fn metal() -> bool {
#[cfg(feature = "metal")]
{
ringkernel_metal::is_metal_available()
}
#[cfg(not(feature = "metal"))]
{
false
}
}
pub fn wgpu() -> bool {
#[cfg(feature = "wgpu")]
{
ringkernel_wgpu::is_wgpu_available()
}
#[cfg(not(feature = "wgpu"))]
{
false
}
}
pub fn available_backends() -> Vec<super::Backend> {
let mut backends = vec![super::Backend::Cpu];
if cuda() {
backends.push(super::Backend::Cuda);
}
if metal() {
backends.push(super::Backend::Metal);
}
if wgpu() {
backends.push(super::Backend::Wgpu);
}
backends
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_runtime_creation() {
let runtime = RingKernel::new().await.unwrap();
let backend = runtime.backend();
assert!(
backend == Backend::Cuda || backend == Backend::Cpu,
"Expected Cuda or Cpu backend, got {:?}",
backend
);
}
#[tokio::test]
async fn test_builder() {
let runtime = RingKernel::builder()
.backend(Backend::Cpu)
.debug(true)
.build()
.await
.unwrap();
assert_eq!(runtime.backend(), Backend::Cpu);
}
#[tokio::test]
async fn test_launch_kernel() {
let runtime = RingKernel::new().await.unwrap();
let kernel = runtime
.launch("test_kernel", LaunchOptions::default())
.await
.unwrap();
assert_eq!(kernel.id().as_str(), "test_kernel");
}
#[test]
fn test_availability() {
let backends = availability::available_backends();
assert!(backends.contains(&Backend::Cpu));
}
}