Struct SpMmOpPlan

Source

pub struct SpMmOpPlan<'a> { /* private fields */ }

Implementations§

Source §

impl<'a> SpMmOpPlan<'a>

Source

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)>

Source

pub fn context(&self) -> &Context

Source

pub fn execute(&self, external_buffer: Option<DevicePtr>) -> Result<()>

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

SpMmOpPlan::execute supports the following index type for representing matrix_a:

32-bit indices (IndexType::I32)
64-bit indices (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.

Source

pub fn as_raw(&self) -> cusparseSpMMOpPlan_t

Source

pub unsafe fn from_raw( ctx: &'a Context, handle: cusparseSpMMOpPlan_t, ) -> Result<Self>

Wraps an existing cuSPARSE SpMM operation plan and takes ownership of it.

§Safety

handle must be a valid cusparseSpMMOpPlan_t associated with ctx and with matrix descriptors that remain valid for lifetime 'a. Ownership of handle is transferred to the returned plan, and the handle must not be destroyed elsewhere after calling this function.

Source