singe_cusparse/
spmm.rs

1use std::ptr;
2
3use singe_cuda::{data_type::DataTypeLike, types::DevicePtr};
4
5use singe_cusparse_sys as sys;
6
7use crate::{
8    context::Context,
9    error::{Error, Result},
10    matrix::{DenseMatrixDescriptor, SparseMatrixDescriptor},
11    operation::SpMmOpPlan,
12    scalar::Scalar,
13    try_ffi,
14    types::{Operation, SpMmAlgorithm},
15    utility::to_usize,
16};
17
18/// Performs the multiplication of a sparse matrix `matrix_a` and a dense matrix `matrix_b`.
19/// Formally, it computes `C = alpha * op(A) * op(B) + beta * C`, where:
20///
21/// * `op(A)` is a sparse matrix of size $m \times k$.
22/// * `op(B)` is a dense matrix of size $k \times n$.
23/// * `C` is a dense matrix of size $m \times n$.
24/// * $\alpha$ and $\beta$ are scalars.
25///
26/// This can also multiply a dense matrix and a sparse matrix by switching the
27/// dense matrix layouts. `B_C` and `C_C` indicate column-major layout, while
28/// `B_R` and `C_R` indicate row-major layout.
29///
30/// `op(A)` is selected by `op_a` and may be `A`, `A^T`, or `A^H`.
31/// `op(B)` is selected by `op_b` and may be `B`, `B^T`, or `B^H`.
32///
33/// When using the (conjugate) transpose of the sparse matrix `A`, this operation may produce slightly different results during different runs with the same input parameters.
34///
35/// [`spmm_buffer_size`] returns the size of the workspace needed by [`spmm`].
36///
37/// Calling [`spmm_preprocess`] is optional.
38/// It may accelerate subsequent calls to [`spmm`].
39/// Useful when [`spmm`] is called multiple times with the same sparsity
40/// pattern (`matrix_a`).
41/// It provides performance advantages with [`SpMmAlgorithm::Csr1`] or [`SpMmAlgorithm::Csr3`].
42/// It has no effect for all other formats and algorithms.
43///
44/// Calling [`spmm_preprocess`] with `buffer` makes that buffer active for `matrix_a` SpMM calls.
45/// Subsequent calls to [`spmm`] with `matrix_a` and the active buffer must use the same values for all parameters as the call to [`spmm_preprocess`].
46/// The exceptions are: `alpha`, `beta`, `matrix_b`, `matrix_c`, and the values (but not indices) of `matrix_a` may be different.
47/// Importantly, the buffer contents must be unmodified since the call to [`spmm_preprocess`].
48/// When [`spmm`] is called with `matrix_a` and its active buffer, it may read acceleration data from the buffer.
49///
50/// Calling [`spmm_preprocess`] again with `matrix_a` and a new buffer makes the
51/// new buffer active and makes the previously active buffer inactive.
52/// For [`spmm`], there can only be one active buffer per sparse matrix at a time.
53/// To get the effect of multiple active buffers for a single sparse matrix, create multiple matrix handles that all point to the same index and value buffers, and call [`spmm_preprocess`] once per handle with different workspace buffers.
54///
55/// Calling [`spmm`] with an inactive buffer is always permitted.
56/// However, there may be no acceleration from the preprocessing in that case.
57///
58/// For the purposes of thread safety, [`spmm_preprocess`] is writing to `matrix_a` internal state.
59///
60/// [`spmm`] supports the following sparse matrix formats:
61///
62/// * [`Format::Coo`](crate::types::Format::Coo)
63/// * [`Format::Csr`](crate::types::Format::Csr)
64/// * [`Format::Csc`](crate::types::Format::Csc)
65/// * [`Format::Bsr`](crate::types::Format::Bsr)
66/// * [`Format::BlockedEll`](crate::types::Format::BlockedEll)
67///
68/// **(1) COO/CSR/CSC/BSR formats**
69///
70/// [`spmm`] supports the following index type for representing the sparse matrix `matrix_a`:
71///
72/// * 32-bit indices ([`IndexType::I32`](crate::types::IndexType::I32))
73/// * 64-bit indices ([`IndexType::I64`](crate::types::IndexType::I64))
74///
75/// [`spmm`] supports the following data types:
76///
77/// Uniform-precision computation:
78///
79/// | `A`/`B`/`C`/`compute_type` |
80/// | --- |
81/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |
82/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |
83/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |
84/// | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |
85///
86/// Mixed-precision computation:
87///
88/// | `A`/`B` | `C` | `compute_type` | Notes |
89/// | --- | --- | --- | --- |
90/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) | [`DataType::I32`](singe_cuda::data_type::DataType::I32) | [`DataType::I32`](singe_cuda::data_type::DataType::I32) |  |
91/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |  |
92/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |  |  |
93/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |  |  |
94/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |  |
95/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |  |
96/// | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) | Deprecated. |
97/// | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) |  | Deprecated. |
98///
99/// [`DataType::F16`](singe_cuda::data_type::DataType::F16), [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16), [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16), and [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) data types always imply mixed-precision computation.
100///
101/// [`spmm`] supports the following algorithms:
102///
103/// * [`SpMmAlgorithm::Default`]: default algorithm for any sparse matrix format.
104/// * [`SpMmAlgorithm::Coo1`]: algorithm 1 for COO sparse matrix format. May provide better performance for a small number of nonzero values, provides the best performance with column-major layout, supports batched computation, and may produce slightly different results across runs with the same input parameters.
105/// * [`SpMmAlgorithm::Coo2`]: algorithm 2 for COO sparse matrix format. Provides deterministic results and the best performance with column-major layout. In general, it is slower than algorithm 1, supports batched computation, requires additional memory, and is identical to [`SpMmAlgorithm::Coo1`] if `op_a` is not [`Operation::NonTranspose`].
106/// * [`SpMmAlgorithm::Coo3`]: algorithm 3 for COO sparse matrix format. May provide better performance for a large number of nonzero values and may produce slightly different results across runs with the same input parameters.
107/// * [`SpMmAlgorithm::Coo4`]: algorithm 4 for COO sparse matrix format. Provides better performance with row-major layout, supports batched computation, and may produce slightly different results across runs with the same input parameters.
108/// * [`SpMmAlgorithm::Csr1`]: algorithm 1 for CSR/CSC sparse matrix format. Provides the best performance with column-major layout, supports batched computation, requires additional memory, and may produce slightly different results across runs with the same input parameters.
109/// * [`SpMmAlgorithm::Csr2`]: algorithm 2 for CSR/CSC sparse matrix format. Provides the best performance with row-major layout, supports batched computation, requires additional memory, and may produce slightly different results across runs with the same input parameters.
110/// * [`SpMmAlgorithm::Csr3`]: algorithm 3 for CSR sparse matrix format. Provides deterministic results and requires additional memory. It supports only CSR matrices with `op_a` set to [`Operation::NonTranspose`]; `op_b` cannot be [`Operation::ConjugateTranspose`], and the data type cannot be [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) or [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16).
111/// * [`SpMmAlgorithm::Bsr1`]: algorithm 1 for BSR sparse matrix format. Provides deterministic results and requires no additional memory. It supports only `op_a` set to [`Operation::NonTranspose`]; the data type cannot be [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) or [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16), and blocks in `A` cannot be column-major.
112///
113/// When using [`spmm`] for mixed-precision computation on COO or CSR matrices, it defaults to algorithms [`SpMmAlgorithm::Coo2`] and [`SpMmAlgorithm::Csr3`], respectively.
114/// If the required computation is not supported by those algorithms, the
115/// mixed-precision operation fails.
116///
117/// **Performance notes:**
118///
119/// * Row-major layout provides higher performance than column-major.
120/// * [`SpMmAlgorithm::Coo4`] and [`SpMmAlgorithm::Csr2`] are intended for row-major layout, while [`SpMmAlgorithm::Coo1`], [`SpMmAlgorithm::Coo2`], [`SpMmAlgorithm::Coo3`], and [`SpMmAlgorithm::Csr1`] are intended for column-major layout.
121/// * When `beta` is not `1`, most algorithms scale the output matrix before the main computation.
122/// * When `n` is `1`, this operation may use [`spmv`](crate::spmv::spmv).
123///
124/// [`spmm`] with all algorithms support the following batch modes except for [`SpMmAlgorithm::Csr3`]:
125///
126/// * $C\_{i} = A \cdot B\_{i}$
127/// * $C\_{i} = A\_{i} \cdot B$
128/// * $C\_{i} = A\_{i} \cdot B\_{i}$
129///
130/// The number of batches and their strides can be set by using [`SparseMatrixDescriptor::set_coo_strided_batch`], [`SparseMatrixDescriptor::set_csr_strided_batch`], and [`DenseMatrixDescriptor::set_strided_batch`].
131/// The maximum number of batches for [`spmm`] is 65,535.
132///
133/// [`spmm`] has the following properties:
134///
135/// * Requires no extra storage for [`SpMmAlgorithm::Coo1`], [`SpMmAlgorithm::Coo3`], [`SpMmAlgorithm::Coo4`], [`SpMmAlgorithm::Bsr1`].
136/// * Supports asynchronous execution.
137/// * Provides deterministic (bitwise) results for each run only for [`SpMmAlgorithm::Coo2`], [`SpMmAlgorithm::Csr3`], and [`SpMmAlgorithm::Bsr1`] algorithms.
138/// * `compute-sanitizer` could report false race conditions for this operation.
139///   These reports result from optimization and do not affect computation correctness.
140/// * Allows the indices of `matrix_a` to be unsorted.
141///
142/// [`spmm`] supports the following optimizations:
143///
144/// * CUDA graph capture.
145/// * Hardware Memory Compression.
146///
147/// **(2) Blocked-ELLPACK format**
148///
149/// [`spmm`] supports the following data types for [`Format::BlockedEll`](crate::types::Format::BlockedEll) format and the following GPU architectures for exploiting NVIDIA Tensor Cores:
150///
151/// | `A`/`B` | `C` | `compute_type` | `op_b` | Compute Capability |
152/// | --- | --- | --- | --- | --- |
153/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | `N`, `T` | `≥ 70` |
154/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | `N`, `T` | `≥ 70` |
155/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | `N`, `T` | `≥ 70` |
156/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) | [`DataType::I32`](singe_cuda::data_type::DataType::I32) | [`DataType::I32`](singe_cuda::data_type::DataType::I32) | `N` column-major, `T` row-major | `≥ 75` |
157/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | `N`, `T` | `≥ 80` |
158/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | `N`, `T` | `≥ 80` |
159/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | `N`, `T` | `≥ 80` |
160/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) | [`DataType::F64`](singe_cuda::data_type::DataType::F64) | [`DataType::F64`](singe_cuda::data_type::DataType::F64) | `N`, `T` | `≥ 80` |
161///
162/// [`spmm`] supports the following algorithms with [`Format::BlockedEll`](crate::types::Format::BlockedEll) format:
163///
164/// | Algorithm | Notes |
165/// | --- | --- |
166/// | [`SpMmAlgorithm::Default`] | Default algorithm for any sparse matrix format. |
167/// | [`SpMmAlgorithm::BlockedEll1`] | Default algorithm for Blocked-ELL format. |
168///
169/// **Performance notes:**
170///
171/// * Blocked-ELL SpMM provides the best performance with Power-of-2 Block-Sizes.
172/// * Large block sizes, such as 64 or greater, provide the best performance.
173///
174/// This operation has the following limitations:
175///
176/// * The pointer mode must be [`PointerMode::Host`](crate::types::PointerMode::Host).
177/// * Only `op_a` set to [`Operation::NonTranspose`] is supported.
178/// * `op_b` set to [`Operation::ConjugateTranspose`] is not supported.
179/// * Only [`IndexType::I32`](crate::types::IndexType::I32) is supported.
180///
181/// # Errors
182///
183/// Returns an error if the cuSPARSE context cannot be bound, the scalar pointer
184/// modes do not match the handle configuration, the descriptors or selected
185/// algorithm are incompatible, a required workspace buffer is missing or
186/// invalid, the data type combination is unsupported, or cuSPARSE rejects the
187/// operation.
188pub fn spmm<Compute: DataTypeLike>(
189    ctx: &Context,
190    op_a: Operation,
191    op_b: Operation,
192    alpha: Scalar<'_, Compute>,
193    matrix_a: &SparseMatrixDescriptor,
194    matrix_b: &DenseMatrixDescriptor,
195    beta: Scalar<'_, Compute>,
196    matrix_c: &mut DenseMatrixDescriptor,
197    algorithm: SpMmAlgorithm,
198    external_buffer: Option<DevicePtr>,
199) -> Result<()> {
200    matrix_a.ensure_context(ctx)?;
201    matrix_b.ensure_context(ctx)?;
202    matrix_c.ensure_context(ctx)?;
203    ctx.bind()?;
204    ctx.require_scalar_pointer_mode(alpha, beta)?;
205
206    unsafe {
207        try_ffi!(sys::cusparseSpMM(
208            ctx.as_raw(),
209            op_a.into(),
210            op_b.into(),
211            alpha.ptr().cast(),
212            matrix_a.as_raw_const(),
213            matrix_b.as_raw_const(),
214            beta.ptr().cast(),
215            matrix_c.as_raw(),
216            Compute::data_type().into(),
217            algorithm.into(),
218            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
219        ))?;
220    }
221    Ok(())
222}
223
224pub fn spmm_buffer_size<Compute: DataTypeLike>(
225    ctx: &Context,
226    op_a: Operation,
227    op_b: Operation,
228    alpha: Scalar<'_, Compute>,
229    matrix_a: &SparseMatrixDescriptor,
230    matrix_b: &DenseMatrixDescriptor,
231    beta: Scalar<'_, Compute>,
232    matrix_c: &mut DenseMatrixDescriptor,
233    algorithm: SpMmAlgorithm,
234) -> Result<usize> {
235    matrix_a.ensure_context(ctx)?;
236    matrix_b.ensure_context(ctx)?;
237    matrix_c.ensure_context(ctx)?;
238    ctx.bind()?;
239    ctx.require_scalar_pointer_mode(alpha, beta)?;
240
241    let mut size = 0;
242    unsafe {
243        try_ffi!(sys::cusparseSpMM_bufferSize(
244            ctx.as_raw(),
245            op_a.into(),
246            op_b.into(),
247            alpha.ptr().cast(),
248            matrix_a.as_raw_const(),
249            matrix_b.as_raw_const(),
250            beta.ptr().cast(),
251            matrix_c.as_raw(),
252            Compute::data_type().into(),
253            algorithm.into(),
254            &raw mut size,
255        ))?;
256    }
257    to_usize(size, "spmm buffer size")
258}
259
260pub fn spmm_preprocess<Compute: DataTypeLike>(
261    ctx: &Context,
262    op_a: Operation,
263    op_b: Operation,
264    alpha: Scalar<'_, Compute>,
265    matrix_a: &SparseMatrixDescriptor,
266    matrix_b: &DenseMatrixDescriptor,
267    beta: Scalar<'_, Compute>,
268    matrix_c: &mut DenseMatrixDescriptor,
269    algorithm: SpMmAlgorithm,
270    external_buffer: Option<DevicePtr>,
271) -> Result<()> {
272    matrix_a.ensure_context(ctx)?;
273    matrix_b.ensure_context(ctx)?;
274    matrix_c.ensure_context(ctx)?;
275    ctx.bind()?;
276    ctx.require_scalar_pointer_mode(alpha, beta)?;
277
278    unsafe {
279        try_ffi!(sys::cusparseSpMM_preprocess(
280            ctx.as_raw(),
281            op_a.into(),
282            op_b.into(),
283            alpha.ptr().cast(),
284            matrix_a.as_raw_const(),
285            matrix_b.as_raw_const(),
286            beta.ptr().cast(),
287            matrix_c.as_raw(),
288            Compute::data_type().into(),
289            algorithm.into(),
290            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
291        ))?;
292    }
293    Ok(())
294}
295
296pub fn spmm_op(
297    ctx: &Context,
298    plan: &SpMmOpPlan<'_>,
299    external_buffer: Option<DevicePtr>,
300) -> Result<()> {
301    if !ptr::eq(ctx, plan.context()) {
302        return Err(Error::PlanContextMismatch);
303    }
304    plan.execute(external_buffer)
305}
singe_cusparse/spmm.rs

singe_cusparse/
spmm.rs
