singe-cusparse 0.1.0-alpha.7

Safe Rust wrappers for the NVIDIA cuSPARSE sparse linear algebra library.
Documentation
use singe_cuda::{data_type::DataTypeLike, types::DevicePtr};

use singe_cusparse_sys as sys;

use crate::{
    context::Context,
    error::Result,
    matrix::SparseMatrixDescriptor,
    operation::SpSvDescriptor,
    scalar::Scalar,
    try_ffi,
    types::{Operation, SpSvAlgorithm, SpSvUpdate},
    utility::to_usize,
    vector::DenseVectorDescriptor,
};

pub fn spsv_buffer_size<Compute: DataTypeLike>(
    ctx: &Context,
    operation: Operation,
    alpha: Scalar<'_, Compute>,
    matrix: &SparseMatrixDescriptor,
    x: &DenseVectorDescriptor,
    y: &mut DenseVectorDescriptor,
    algorithm: SpSvAlgorithm,
    descriptor: &SpSvDescriptor,
) -> Result<usize> {
    descriptor.ensure_context(ctx)?;
    matrix.ensure_context(ctx)?;
    x.ensure_context(ctx)?;
    y.ensure_context(ctx)?;
    ctx.bind()?;
    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
    }

    let mut size = 0;
    unsafe {
        try_ffi!(sys::cusparseSpSV_bufferSize(
            ctx.as_raw(),
            operation.into(),
            alpha.ptr().cast(),
            matrix.as_raw_const(),
            x.as_raw_const(),
            y.as_raw(),
            Compute::data_type().into(),
            algorithm.into(),
            descriptor.as_raw(),
            &raw mut size,
        ))?;
    }
    to_usize(size, "spsv buffer size")
}

pub fn spsv_analysis<Compute: DataTypeLike>(
    ctx: &Context,
    operation: Operation,
    alpha: Scalar<'_, Compute>,
    matrix: &SparseMatrixDescriptor,
    x: &DenseVectorDescriptor,
    y: &mut DenseVectorDescriptor,
    algorithm: SpSvAlgorithm,
    descriptor: &SpSvDescriptor,
    external_buffer: Option<DevicePtr>,
) -> Result<()> {
    descriptor.ensure_context(ctx)?;
    matrix.ensure_context(ctx)?;
    x.ensure_context(ctx)?;
    y.ensure_context(ctx)?;
    ctx.bind()?;
    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
    }

    unsafe {
        try_ffi!(sys::cusparseSpSV_analysis(
            ctx.as_raw(),
            operation.into(),
            alpha.ptr().cast(),
            matrix.as_raw_const(),
            x.as_raw_const(),
            y.as_raw(),
            Compute::data_type().into(),
            algorithm.into(),
            descriptor.as_raw(),
            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
        ))?;
    }
    Ok(())
}

pub fn spsv_solve<Compute: DataTypeLike>(
    ctx: &Context,
    operation: Operation,
    alpha: Scalar<'_, Compute>,
    matrix: &SparseMatrixDescriptor,
    x: &DenseVectorDescriptor,
    y: &mut DenseVectorDescriptor,
    algorithm: SpSvAlgorithm,
    descriptor: &SpSvDescriptor,
) -> Result<()> {
    descriptor.ensure_context(ctx)?;
    matrix.ensure_context(ctx)?;
    x.ensure_context(ctx)?;
    y.ensure_context(ctx)?;
    ctx.bind()?;
    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
    }

    unsafe {
        try_ffi!(sys::cusparseSpSV_solve(
            ctx.as_raw(),
            operation.into(),
            alpha.ptr().cast(),
            matrix.as_raw_const(),
            x.as_raw_const(),
            y.as_raw(),
            Compute::data_type().into(),
            algorithm.into(),
            descriptor.as_raw(),
        ))?;
    }
    Ok(())
}

pub fn spsv_update_matrix(
    ctx: &Context,
    descriptor: &SpSvDescriptor,
    new_values: DevicePtr,
    update: SpSvUpdate,
) -> Result<()> {
    descriptor.ensure_context(ctx)?;
    ctx.bind()?;
    unsafe {
        try_ffi!(sys::cusparseSpSV_updateMatrix(
            ctx.as_raw(),
            descriptor.as_raw(),
            new_values.as_ptr() as _,
            update.into(),
        ))?;
    }
    Ok(())
}