rustytorch/crates/core/rtx-cubecl/src/client.rs

//! `CubeCL` compute client wrapper.
//!
//! Provides a unified client interface for executing `CubeCL` kernels across different backends.

#[cfg(feature = "wgpu")]
use cubecl::CubeCount;
#[cfg(feature = "wgpu")]
use cubecl_runtime::client::ComputeClient;
use cubecl_runtime::server::{Binding, Handle};

use crate::device::CubeclDevice;
use crate::error::{CubeclError, Result};
use crate::runtime::RuntimeBackend;

/// Type alias for the WGPU runtime compute client.
#[cfg(feature = "wgpu")]
pub type WgpuClient = ComputeClient<cubecl_wgpu::WgpuRuntime>;

/// Handle to a tensor buffer on the compute device.
///
/// This wraps the `CubeCL` `Handle` type which represents an allocated buffer.
#[derive(Clone, Debug)]
pub struct TensorHandle {
    /// The underlying `CubeCL` handle.
    pub(crate) handle: Handle,
    /// Number of elements in the buffer.
    pub(crate) numel: usize,
    /// Size in bytes.
    pub(crate) size_bytes: usize,
}

impl TensorHandle {
    /// Create a new tensor handle.
    #[inline]
    pub fn new(handle: Handle, numel: usize, size_bytes: usize) -> Self {
        Self {
            handle,
            numel,
            size_bytes,
        }
    }

    /// Get the number of elements.
    #[inline]
    pub fn numel(&self) -> usize {
        self.numel
    }

    /// Get the size in bytes.
    #[inline]
    pub fn size_bytes(&self) -> usize {
        self.size_bytes
    }

    /// Get the binding for kernel execution.
    #[inline]
    pub fn binding(&self) -> Binding {
        self.handle.clone().binding()
    }
}

/// Unified `CubeCL` client that can work with any runtime.
///
/// This provides a runtime-agnostic interface for:
/// - Allocating tensor buffers
/// - Copying data to/from device
/// - Executing kernels
#[derive(Clone)]
pub struct CubeclClient {
    /// The backend type.
    backend: RuntimeBackend,
    /// Type-erased client for runtime dispatch.
    inner: ClientInner,
}

/// Type-erased client wrapper.
#[derive(Clone)]
enum ClientInner {
    #[cfg(feature = "wgpu")]
    Wgpu(WgpuClient),
    /// CPU fallback (no actual compute, just for testing)
    Cpu,
}

impl CubeclClient {
    /// Create a new client for the given device.
    pub fn new(device: &CubeclDevice) -> Result<Self> {
        match device.backend {
            #[cfg(feature = "wgpu")]
            RuntimeBackend::Wgpu => {
                use cubecl_wgpu::{WgpuDevice, WgpuRuntime};

                let wgpu_device = match device.index {
                    0 => WgpuDevice::BestAvailable,
                    n => WgpuDevice::DiscreteGpu(n),
                };
                let client = WgpuRuntime::client(&wgpu_device);
                Ok(Self {
                    backend: RuntimeBackend::Wgpu,
                    inner: ClientInner::Wgpu(client),
                })
            }
            #[cfg(not(feature = "wgpu"))]
            RuntimeBackend::Wgpu => Err(CubeclError::DeviceError(
                "WGPU backend not enabled".to_string(),
            )),
            RuntimeBackend::Cpu => Ok(Self {
                backend: RuntimeBackend::Cpu,
                inner: ClientInner::Cpu,
            }),
            _ => Err(CubeclError::DeviceError(format!(
                "{} backend not yet implemented",
                device.backend
            ))),
        }
    }

    /// Get the backend type.
    #[inline]
    pub fn backend(&self) -> RuntimeBackend {
        self.backend
    }

    /// Allocate an empty buffer of the given size in bytes.
    pub fn empty(&self, size_bytes: usize) -> Result<TensorHandle> {
        match &self.inner {
            #[cfg(feature = "wgpu")]
            ClientInner::Wgpu(client) => {
                let handle = client.empty(size_bytes);
                Ok(TensorHandle::new(
                    handle,
                    size_bytes / 4, // Assuming f32
                    size_bytes,
                ))
            }
            ClientInner::Cpu => {
                // CPU fallback - no actual allocation
                Err(CubeclError::UnsupportedError(
                    "CPU backend cannot allocate GPU buffers".to_string(),
                ))
            }
        }
    }

    /// Create a buffer with the given data.
    pub fn create(&self, data: &[u8]) -> Result<TensorHandle> {
        match &self.inner {
            #[cfg(feature = "wgpu")]
            ClientInner::Wgpu(client) => {
                let handle = client.create(data);
                Ok(TensorHandle::new(
                    handle,
                    data.len() / 4, // Assuming f32
                    data.len(),
                ))
            }
            ClientInner::Cpu => Err(CubeclError::UnsupportedError(
                "CPU backend cannot create GPU buffers".to_string(),
            )),
        }
    }

    /// Create a buffer from f32 data.
    pub fn create_f32(&self, data: &[f32]) -> Result<TensorHandle> {
        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        let mut handle = self.create(&bytes)?;
        handle.numel = data.len();
        Ok(handle)
    }

    /// Read data from a buffer as bytes.
    pub fn read(&self, handle: &TensorHandle) -> Result<Vec<u8>> {
        match &self.inner {
            #[cfg(feature = "wgpu")]
            ClientInner::Wgpu(client) => {
                let data = client.read_one(handle.binding());
                Ok(data)
            }
            ClientInner::Cpu => Err(CubeclError::UnsupportedError(
                "CPU backend cannot read GPU buffers".to_string(),
            )),
        }
    }

    /// Read data from a buffer as f32.
    pub fn read_f32(&self, handle: &TensorHandle) -> Result<Vec<f32>> {
        let bytes = self.read(handle)?;
        let floats: Vec<f32> = bytes
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
            .collect();
        Ok(floats)
    }

    /// Synchronize the device (wait for all operations to complete).
    pub fn sync(&self) {
        match &self.inner {
            #[cfg(feature = "wgpu")]
            ClientInner::Wgpu(client) => {
                cubecl_common::future::block_on(client.sync());
            }
            ClientInner::Cpu => {}
        }
    }

    /// Execute a kernel.
    #[cfg(feature = "wgpu")]
    pub fn execute_wgpu<K>(&self, kernel: K, count: CubeCount, bindings: Vec<Binding>)
    where
        K: cubecl_runtime::server::ComputeKernel,
    {
        match &self.inner {
            ClientInner::Wgpu(client) => {
                client.execute(Box::new(kernel), count, bindings);
            }
            _ => {}
        }
    }

    /// Get the WGPU client if this is a WGPU backend.
    #[cfg(feature = "wgpu")]
    pub fn as_wgpu(&self) -> Option<&WgpuClient> {
        match &self.inner {
            ClientInner::Wgpu(client) => Some(client),
            _ => None,
        }
    }
}

impl std::fmt::Debug for CubeclClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CubeclClient")
            .field("backend", &self.backend)
            .finish()
    }
}

/// Global client registry for device-to-client mapping.
///
/// This allows reusing clients across tensor operations.
mod registry {
    use super::{CubeclClient, CubeclDevice, CubeclError, Result};
    use std::collections::HashMap;
    use std::sync::RwLock;

    static CLIENTS: std::sync::LazyLock<RwLock<HashMap<CubeclDevice, CubeclClient>>> =
        std::sync::LazyLock::new(|| RwLock::new(HashMap::new()));

    /// Get or create a client for the given device.
    pub fn get_or_create(device: &CubeclDevice) -> Result<CubeclClient> {
        // Check if we already have a client
        {
            let clients = CLIENTS
                .read()
                .map_err(|_| CubeclError::RuntimeError("Client registry lock poisoned".into()))?;
            if let Some(client) = clients.get(device) {
                return Ok(client.clone());
            }
        }

        // Create a new client
        let client = CubeclClient::new(device)?;

        // Store it
        {
            let mut clients = CLIENTS
                .write()
                .map_err(|_| CubeclError::RuntimeError("Client registry lock poisoned".into()))?;
            clients.insert(device.clone(), client.clone());
        }

        Ok(client)
    }
}

pub use registry::get_or_create;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cpu_client() {
        let device = CubeclDevice::cpu();
        let client = CubeclClient::new(&device).unwrap();
        assert_eq!(client.backend(), RuntimeBackend::Cpu);
    }
}