cubecl-cuda 0.11.0

use crate::compute::{
    events::{Fence, driver_error},
    storage::{cpu::PinnedMemoryStorage, gpu::GpuStorage},
};
use cubecl_core::{
    MemoryConfiguration,
    ir::MemoryDeviceProperties,
    server::{BufferBinding, DeviceCaptures, Handle, ServerError},
};
use cubecl_server::driver::DriverError;
use cubecl_server::storage::PINNED_MEMORY_ALIGNMENT;
use cubecl_server::{
    config::{memory::CudaAllocator, streaming::StreamPriority},
    logging::ServerLogger,
    memory_management::{
        ErrorGraph, FailureId, MemoryAllocationMode, MemoryManagement, MemoryManagementOptions,
        drop_queue,
    },
    metadata_cache::{MetadataCachePolicy, MetadataInfoCache},
    stream::{EventStreamBackend, StreamCapture, StreamMemory},
};
use std::{mem::MaybeUninit, sync::Arc};

#[derive(Debug)]
pub struct Stream {
    pub sys: cudarc::driver::sys::CUstream,
    pub memory_management_gpu: MemoryManagement<GpuStorage>,
    pub memory_management_cpu: MemoryManagement<PinnedMemoryStorage>,
    pub drop_queue: drop_queue::PendingDropQueue<Fence>,
    /// This stream's graph capture (see [`StreamCapture`]): its position in
    /// the lifecycle, and the memory its recorded launches were given. Enforces
    /// the ordered `graph_prepare` → `begin_capture` → `end_capture`
    /// transitions and gates the deferral of fenced drop-queue flushes while a
    /// capture is recording.
    pub capturing: StreamCapture,
    /// Reusable per-launch info buffers (kernel shapes/strides/scalars), keyed
    /// by the exact info words they were built from. Admission and
    /// least-recently-used eviction are decided by the cache's
    /// [`MetadataCachePolicy`]; the launch path sets its [`CacheMode`] from
    /// the capture lifecycle, so during graph capture every buffer is cached
    /// and none is evicted mid-capture. See [`StreamCapture::cache_mode`].
    pub info_cache: MetadataInfoCache<Handle>,
}

impl StreamMemory for Stream {
    fn failure(&self, binding: &BufferBinding) -> Option<FailureId> {
        self.memory_management_gpu
            .failure(&binding.memory, binding.range())
    }

    fn taint(&mut self, binding: &BufferBinding, failure: FailureId, failures: &mut ErrorGraph) {
        self.memory_management_gpu
            .taint(&binding.memory, binding.range(), failure, failures)
    }

    fn written(&mut self, binding: &BufferBinding, failures: &mut ErrorGraph) {
        self.memory_management_gpu
            .written(&binding.memory, binding.range(), failures)
    }
}

#[derive(new, Debug)]
pub struct CudaStreamBackend {
    mem_props: MemoryDeviceProperties,
    mem_config: MemoryConfiguration,
    mem_alignment: usize,
    logger: Arc<ServerLogger>,
    priority: StreamPriority,
    allocator: CudaAllocator,
    /// The device's captures, which every stream this creates takes its capture state from.
    captures: DeviceCaptures,
}

/// Create a non-blocking CUDA stream, applying the requested priority hint.
///
/// `StreamPriority::Default` preserves the historical `cuStreamCreate` path so
/// existing users see no change. `Low`/`High` go through
/// `cuStreamCreateWithPriority` using the device's range as queried via
/// `cuCtxGetStreamPriorityRange`. CUDA convention: lower number = higher
/// priority, so the queried `greatest` is numerically smallest (most
/// aggressive) and `least` is numerically largest (least aggressive). On
/// devices without priority support both values are 0 and CUDA silently
/// ignores the priority argument — equivalent to the default path.
///
/// Both calls require a current CUDA context; callers in this crate always
/// set the context before invoking stream creation.
///
/// # Errors
///
/// Returns a [`DriverError`] when the driver refuses to create the stream, flagged as poisoning
/// the device when the context is poisoned.
pub(crate) fn create_cuda_stream(
    priority: StreamPriority,
) -> Result<cudarc::driver::sys::CUstream, DriverError> {
    use cudarc::driver::sys::{self, CUstream_flags};

    let use_greatest = match priority {
        StreamPriority::Default => {
            return cudarc::driver::result::stream::create(
                cudarc::driver::result::stream::StreamKind::NonBlocking,
            )
            .map_err(|err| driver_error("cuStreamCreate", err));
        }
        StreamPriority::High => true,
        StreamPriority::Low => false,
    };

    // SAFETY: `cuCtxGetStreamPriorityRange` writes through both pointers on
    // success; we only read the locals after `?` confirms success.
    let value = unsafe {
        let mut least: i32 = 0;
        let mut greatest: i32 = 0;
        sys::cuCtxGetStreamPriorityRange(&mut least, &mut greatest)
            .result()
            .map_err(|err| driver_error("cuCtxGetStreamPriorityRange", err))?;
        if use_greatest { greatest } else { least }
    };

    // SAFETY: `cuStreamCreateWithPriority` writes the new stream handle through
    // the out pointer on success; `?` ensures we only `assume_init` on
    // success.
    unsafe {
        let mut stream = MaybeUninit::uninit();
        sys::cuStreamCreateWithPriority(
            stream.as_mut_ptr(),
            CUstream_flags::CU_STREAM_NON_BLOCKING as u32,
            value,
        )
        .result()
        .map_err(|err| driver_error("cuStreamCreateWithPriority", err))?;
        Ok(stream.assume_init())
    }
}

impl EventStreamBackend for CudaStreamBackend {
    type Stream = Stream;
    type Event = Fence;

    fn create_stream(&self) -> Result<Self::Stream, ServerError> {
        let stream = create_cuda_stream(self.priority)?;

        let storage = GpuStorage::new(self.mem_alignment, stream, self.allocator);

        let memory_management_gpu = MemoryManagement::from_configuration(
            storage,
            &self.mem_props,
            self.mem_config.clone(),
            self.logger.clone(),
            MemoryManagementOptions::new("Main GPU Memory"),
        );
        // We use the same page size and memory pools configuration for CPU pinned memory, since we
        // expect the CPU to have at least the same amount of RAM as GPU memory.
        // The host was never measured, so this pool states no capacity.
        let memory_management_cpu = MemoryManagement::from_configuration(
            PinnedMemoryStorage::new(),
            &MemoryDeviceProperties::new(
                self.mem_props.max_page_size,
                PINNED_MEMORY_ALIGNMENT as u64,
            ),
            self.mem_config.clone(),
            self.logger.clone(),
            MemoryManagementOptions::new("Pinned CPU Memory").mode(MemoryAllocationMode::Auto),
        );

        Ok(Stream {
            sys: stream,
            memory_management_gpu,
            memory_management_cpu,
            drop_queue: Default::default(),
            capturing: StreamCapture::new(&self.captures),
            info_cache: MetadataInfoCache::new(MetadataCachePolicy::default()),
        })
    }

    fn flush(stream: &mut Self::Stream, _failures: &mut ErrorGraph) -> Self::Event {
        Fence::new(stream.sys)
    }

    fn wait_event(stream: &mut Self::Stream, event: Self::Event) -> Result<(), ServerError> {
        event.wait_async(stream.sys)
    }

    fn wait_event_sync(event: Self::Event) -> Result<(), ServerError> {
        event.wait_sync()
    }

    fn handle_cursor(stream: &Self::Stream, binding: &BufferBinding) -> u64 {
        // The slice cursor the sync logic compares against the origin stream's `last_synced`
        // to decide whether to wait. A freed/reallocated slice falls back to `u64::MAX`,
        // which conservatively forces a wait.
        stream
            .memory_management_gpu
            .get_cursor(binding.memory.clone())
            .unwrap_or(u64::MAX)
    }
}