singe-cusparse 0.1.0-alpha.5

use std::{marker::PhantomData, ptr};

use singe_cuda::data_type::DataType;
use singe_cuda::types::DevicePtr;

use crate::{
    context::Context,
    error::{Error, Result},
    matrix::{DenseMatrixDescriptor, SparseMatrixDescriptor},
    sys, try_ffi,
    types::{Operation, SpMmOpAlgorithm},
};

#[derive(Debug)]
pub struct SpGemmDescriptor {
    handle: sys::cusparseSpGEMMDescr_t,
}

#[derive(Debug)]
pub struct SpSvDescriptor {
    handle: sys::cusparseSpSVDescr_t,
}

#[derive(Debug)]
pub struct SpSmDescriptor {
    handle: sys::cusparseSpSMDescr_t,
}

#[derive(Debug)]
pub struct SpMmOpPlan<'a> {
    context: &'a Context,
    handle: sys::cusparseSpMMOpPlan_t,
    _matrix_a: PhantomData<&'a SparseMatrixDescriptor<'a>>,
    _matrix_b: PhantomData<&'a DenseMatrixDescriptor<'a>>,
    _matrix_c: PhantomData<&'a DenseMatrixDescriptor<'a>>,
}

// Operation descriptors and plans own cuSPARSE analysis state but do not expose
// pointer rebinding through shared references, so immutable sharing is allowed.
unsafe impl Send for SpGemmDescriptor {}
unsafe impl Sync for SpGemmDescriptor {}
unsafe impl Send for SpSvDescriptor {}
unsafe impl Sync for SpSvDescriptor {}
unsafe impl Send for SpSmDescriptor {}
unsafe impl Sync for SpSmDescriptor {}
unsafe impl Send for SpMmOpPlan<'_> {}
unsafe impl Sync for SpMmOpPlan<'_> {}

impl SpGemmDescriptor {
    pub fn create() -> Result<Self> {
        let mut handle = ptr::null_mut();
        unsafe {
            try_ffi!(sys::cusparseSpGEMM_createDescr(&raw mut handle))?;
        }

        if handle.is_null() {
            return Err(Error::NullHandle);
        }

        Ok(Self { handle })
    }

    pub fn num_products(&self) -> Result<usize> {
        let mut value = 0_i64;
        unsafe {
            try_ffi!(sys::cusparseSpGEMM_getNumProducts(
                self.as_raw(),
                &raw mut value,
            ))?;
        }

        usize::try_from(value).map_err(|_| Error::OutOfRange {
            name: "num_products".into(),
        })
    }

    pub fn as_raw(&self) -> sys::cusparseSpGEMMDescr_t {
        self.handle
    }
}

impl Drop for SpGemmDescriptor {
    fn drop(&mut self) {
        unsafe {
            if let Err(err) = try_ffi!(sys::cusparseSpGEMM_destroyDescr(self.handle)) {
                #[cfg(debug_assertions)]
                eprintln!("failed to destroy cusparse spgemm descriptor: {err}");
            }
        }
    }
}

impl SpSvDescriptor {
    pub fn create() -> Result<Self> {
        let mut handle = ptr::null_mut();
        unsafe {
            try_ffi!(sys::cusparseSpSV_createDescr(&raw mut handle))?;
        }

        if handle.is_null() {
            return Err(Error::NullHandle);
        }

        Ok(Self { handle })
    }

    pub fn as_raw(&self) -> sys::cusparseSpSVDescr_t {
        self.handle
    }
}

impl Drop for SpSvDescriptor {
    fn drop(&mut self) {
        unsafe {
            if let Err(err) = try_ffi!(sys::cusparseSpSV_destroyDescr(self.handle)) {
                #[cfg(debug_assertions)]
                eprintln!("failed to destroy cusparse spsv descriptor: {err}");
            }
        }
    }
}

impl SpSmDescriptor {
    pub fn create() -> Result<Self> {
        let mut handle = ptr::null_mut();
        unsafe {
            try_ffi!(sys::cusparseSpSM_createDescr(&raw mut handle))?;
        }

        if handle.is_null() {
            return Err(Error::NullHandle);
        }

        Ok(Self { handle })
    }

    pub fn as_raw(&self) -> sys::cusparseSpSMDescr_t {
        self.handle
    }
}

impl<'a> SpMmOpPlan<'a> {
    pub fn create(
        ctx: &'a Context,
        op_a: Operation,
        op_b: Operation,
        matrix_a: &'a SparseMatrixDescriptor<'a>,
        matrix_b: &'a DenseMatrixDescriptor<'a>,
        matrix_c: &'a mut DenseMatrixDescriptor<'a>,
        compute_type: DataType,
        algorithm: SpMmOpAlgorithm,
        add_operation_ltoir: Option<&[u8]>,
        mul_operation_ltoir: Option<&[u8]>,
        epilogue_ltoir: Option<&[u8]>,
    ) -> Result<(Self, usize)> {
        ctx.bind()?;

        let mut handle = ptr::null_mut();
        let mut workspace_size = 0;
        let (add_ptr, add_len) = ltoir_parts(add_operation_ltoir);
        let (mul_ptr, mul_len) = ltoir_parts(mul_operation_ltoir);
        let (epilogue_ptr, epilogue_len) = ltoir_parts(epilogue_ltoir);
        unsafe {
            try_ffi!(sys::cusparseSpMMOp_createPlan(
                ctx.as_raw(),
                &raw mut handle,
                op_a.into(),
                op_b.into(),
                matrix_a.as_raw_const(),
                matrix_b.as_raw_const(),
                matrix_c.as_raw(),
                compute_type.into(),
                algorithm.into(),
                add_ptr.cast(),
                add_len as _,
                mul_ptr.cast(),
                mul_len as _,
                epilogue_ptr.cast(),
                epilogue_len as _,
                &raw mut workspace_size,
            ))?;
        }

        if handle.is_null() {
            return Err(Error::NullHandle);
        }

        Ok((
            Self {
                context: ctx,
                handle,
                _matrix_a: PhantomData,
                _matrix_b: PhantomData,
                _matrix_c: PhantomData,
            },
            usize::try_from(workspace_size).map_err(|_| Error::OutOfRange {
                name: "spmm op workspace size".into(),
            })?,
        ))
    }

    pub fn context(&self) -> &Context {
        self.context
    }

    /// NVRTC and nvJitLink are not currently available on Arm64 Android platforms.
    ///
    /// This operation does not support Android and Tegra platforms except Judy (sm87).
    ///
    /// Experimental: multiplies `matrix_a` and `matrix_b` with custom operators.
    ///
    /// where
    ///
    /// * `op(A)` is a sparse matrix of size $m \times k$.
    /// * `op(B)` is a dense matrix of size $k \times n$.
    /// * `C` is a dense matrix of size $m \times n$.
    /// * $\oplus$, $\otimes$, and $\text{epilogue}$ are custom **add**, **mul**, and **epilogue** operators respectively.
    ///
    /// `op(A)` is selected by `op_a` and may be `A`, `A^T`, or `A^H`.
    /// `op(B)` is selected by `op_b` and may be `B`, `B^T`, or `B^H`.
    ///
    /// Only `op_a == Operation::NonTranspose` is currently supported.
    ///
    /// [`SpMmOpPlan::create`] returns the workspace size together with the compiled kernel state needed by [`SpMmOpPlan::execute`].
    ///
    /// The custom add, multiply, and epilogue operators must accept and return the selected compute type.
    /// The compute type may be `float`, `double`, `cuComplex`, `cuDoubleComplex`, or `int`.
    ///
    /// [`SpMmOpPlan::execute`] supports the following sparse matrix formats:
    ///
    /// * [`Format::Csr`](crate::types::Format::Csr)
    ///
    /// [`SpMmOpPlan::execute`] supports the following index type for representing `matrix_a`:
    ///
    /// * 32-bit indices ([`IndexType::I32`](crate::types::IndexType::I32))
    /// * 64-bit indices ([`IndexType::I64`](crate::types::IndexType::I64))
    ///
    /// [`SpMmOpPlan::execute`] supports the following data types:
    ///
    /// Uniform-precision computation:
    ///
    /// | `A`/`B`/`C`/`compute_type` |
    /// | --- |
    /// | [`DataType::F32`] |
    /// | [`DataType::F64`] |
    /// | [`DataType::ComplexF32`] |
    /// | [`DataType::ComplexF64`] |
    ///
    /// Mixed-precision computation:
    ///
    /// | `A`/`B` | `C` | `compute_type` |
    /// | --- | --- | --- |
    /// | [`DataType::I8`] | [`DataType::I32`] | [`DataType::I32`] |
    /// | [`DataType::I8`] | [`DataType::F32`] | [`DataType::F32`] |
    /// | [`DataType::F16`] | | |
    /// | [`DataType::Bf16`] | | |
    /// | [`DataType::F16`] | [`DataType::F16`] | |
    /// | [`DataType::Bf16`] | [`DataType::Bf16`] | |
    ///
    /// [`SpMmOpPlan::execute`] supports the following algorithms:
    ///
    /// | Algorithm | Notes |
    /// | --- | --- |
    /// | [`SpMmOpAlgorithm::Default`] | Default algorithm for any sparse matrix format. |
    ///
    /// **Performance notes:**
    ///
    /// * Row-major layout provides higher performance than column-major.
    ///
    /// [`SpMmOpPlan::execute`] has the following properties:
    ///
    /// * Requires extra storage.
    /// * Supports asynchronous execution.
    /// * Provides deterministic (bitwise) results for each run.
    /// * Allows the indices of `matrix_a` to be unsorted.
    ///
    /// [`SpMmOpPlan::execute`] supports the following optimizations:
    ///
    /// * CUDA graph capture.
    /// * Hardware Memory Compression.
    ///
    /// # Errors
    ///
    /// Returns an error if the CUDA context cannot be bound or if cuSPARSE
    /// rejects the prepared SpMM operation.
    pub fn execute(&self, external_buffer: Option<DevicePtr>) -> Result<()> {
        self.context.bind()?;

        unsafe {
            try_ffi!(sys::cusparseSpMMOp(
                self.as_raw(),
                external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
            ))?;
        }
        Ok(())
    }

    pub fn as_raw(&self) -> sys::cusparseSpMMOpPlan_t {
        self.handle
    }
}

impl Drop for SpMmOpPlan<'_> {
    fn drop(&mut self) {
        if let Err(err) = self.context.bind() {
            #[cfg(debug_assertions)]
            eprintln!("failed to bind cusparse spmm op plan context during drop: {err}");
            return;
        }

        unsafe {
            if let Err(err) = try_ffi!(sys::cusparseSpMMOp_destroyPlan(self.handle)) {
                #[cfg(debug_assertions)]
                eprintln!("failed to destroy cusparse spmm op plan: {err}");
            }
        }
    }
}

fn ltoir_parts(buffer: Option<&[u8]>) -> (*const u8, usize) {
    match buffer {
        Some(buffer) => (buffer.as_ptr(), buffer.len()),
        None => (ptr::null(), 0),
    }
}

impl Drop for SpSmDescriptor {
    fn drop(&mut self) {
        unsafe {
            if let Err(err) = try_ffi!(sys::cusparseSpSM_destroyDescr(self.handle)) {
                #[cfg(debug_assertions)]
                eprintln!("failed to destroy cusparse spsm descriptor: {err}");
            }
        }
    }
}