singe_cusparse/
spmv.rs

1use singe_cuda::{data_type::DataTypeLike, types::DevicePtr};
2
3use singe_cusparse_sys as sys;
4
5use crate::{
6    context::Context,
7    error::Result,
8    matrix::SparseMatrixDescriptor,
9    scalar::Scalar,
10    try_ffi,
11    types::{Operation, SpMvAlgorithm},
12    utility::to_usize,
13    vector::DenseVectorDescriptor,
14};
15
16pub fn spmv_buffer_size<Compute: DataTypeLike>(
17    ctx: &Context,
18    operation: Operation,
19    alpha: Scalar<'_, Compute>,
20    matrix: &SparseMatrixDescriptor,
21    x: &DenseVectorDescriptor,
22    beta: Scalar<'_, Compute>,
23    y: &mut DenseVectorDescriptor,
24    algorithm: SpMvAlgorithm,
25) -> Result<usize> {
26    matrix.ensure_context(ctx)?;
27    x.ensure_context(ctx)?;
28    y.ensure_context(ctx)?;
29    ctx.bind()?;
30    ctx.require_scalar_pointer_mode(alpha, beta)?;
31
32    let mut size = 0;
33    unsafe {
34        try_ffi!(sys::cusparseSpMV_bufferSize(
35            ctx.as_raw(),
36            operation.into(),
37            alpha.ptr().cast(),
38            matrix.as_raw_const(),
39            x.as_raw_const(),
40            beta.ptr().cast(),
41            y.as_raw(),
42            Compute::data_type().into(),
43            algorithm.into(),
44            &raw mut size,
45        ))?;
46    }
47    to_usize(size, "spmv buffer size")
48}
49
50pub fn spmv_preprocess<Compute: DataTypeLike>(
51    ctx: &Context,
52    operation: Operation,
53    alpha: Scalar<'_, Compute>,
54    matrix: &SparseMatrixDescriptor,
55    x: &DenseVectorDescriptor,
56    beta: Scalar<'_, Compute>,
57    y: &mut DenseVectorDescriptor,
58    algorithm: SpMvAlgorithm,
59    external_buffer: Option<DevicePtr>,
60) -> Result<()> {
61    matrix.ensure_context(ctx)?;
62    x.ensure_context(ctx)?;
63    y.ensure_context(ctx)?;
64    ctx.bind()?;
65    ctx.require_scalar_pointer_mode(alpha, beta)?;
66
67    unsafe {
68        try_ffi!(sys::cusparseSpMV_preprocess(
69            ctx.as_raw(),
70            operation.into(),
71            alpha.ptr().cast(),
72            matrix.as_raw_const(),
73            x.as_raw_const(),
74            beta.ptr().cast(),
75            y.as_raw(),
76            Compute::data_type().into(),
77            algorithm.into(),
78            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
79        ))?;
80    }
81    Ok(())
82}
83
84/// Multiplies sparse matrix `matrix` by dense vector `x`.
85/// Formally, it computes `y = alpha * op(A) * x + beta * y`, where:
86///
87/// * `op(A)` is a sparse matrix of size $m \times k$.
88/// * `X` is a dense vector of size $k$.
89/// * `Y` is a dense vector of size $m$.
90/// * $\alpha$ and $\beta$ are scalars.
91///
92/// `op(A)` is selected by `operation` and may be `A`, `A^T`, or `A^H`.
93///
94/// [`spmv_buffer_size`] returns the workspace size needed by [`spmv_preprocess`] and [`spmv`].
95///
96/// The sparse matrix formats currently supported are listed below:
97///
98/// * [`Format::Coo`](crate::types::Format::Coo)
99/// * [`Format::Csr`](crate::types::Format::Csr)
100/// * [`Format::Csc`](crate::types::Format::Csc)
101/// * [`Format::Bsr`](crate::types::Format::Bsr)
102/// * [`Format::SlicedEllpack`](crate::types::Format::SlicedEllpack)
103///
104/// [`spmv`] supports the following index type for representing the sparse matrix `matrix`:
105///
106/// * 32-bit indices ([`IndexType::I32`](crate::types::IndexType::I32))
107/// * 64-bit indices ([`IndexType::I64`](crate::types::IndexType::I64))
108///
109/// [`spmv`] supports the following data types:
110///
111/// Uniform-precision computation:
112///
113/// | `A`/`X`/`Y`/`compute_type` |
114/// | --- |
115/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |
116/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |
117/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |
118/// | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |
119///
120/// Mixed-precision computation:
121///
122/// | `A`/`X` | `Y` | `compute_type` | Notes |
123/// | --- | --- | --- | --- |
124/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) | [`DataType::I32`](singe_cuda::data_type::DataType::I32) | [`DataType::I32`](singe_cuda::data_type::DataType::I32) |  |
125/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |  |
126/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |  |  |
127/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |  |  |
128/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |  |
129/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |  |
130/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |  |
131/// | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) |  | Deprecated. |
132/// | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) |  | Deprecated. |
133///
134/// | `A` | `X`/`Y`/`compute_type` |
135/// | --- | --- |
136/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |
137///
138/// Mixed Regular/Complex computation:
139///
140/// | `A` | `X`/`Y`/`compute_type` |
141/// | --- | --- |
142/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |
143/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |
144///
145/// [`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.
146///
147/// [`spmv`] supports the following algorithms:
148///
149/// | Algorithm | Notes |
150/// | --- | --- |
151/// | [`SpMvAlgorithm::Default`] | Default algorithm for any sparse matrix format. |
152/// | [`SpMvAlgorithm::Coo1`] | Default algorithm for COO sparse matrix format. May produce slightly different results during different runs with the same input parameters. |
153/// | [`SpMvAlgorithm::Coo2`] | Provides deterministic (bitwise) results for each run. If `operation` is not [`Operation::NonTranspose`], it is identical to [`SpMvAlgorithm::Coo1`]. |
154/// | [`SpMvAlgorithm::Csr1`] | Default algorithm for CSR/CSC sparse matrix format. May produce slightly different results during different runs with the same input parameters. |
155/// | [`SpMvAlgorithm::Csr2`] | Provides deterministic (bitwise) results for each run. If `operation` is not [`Operation::NonTranspose`], it is identical to [`SpMvAlgorithm::Csr1`]. |
156/// | [`SpMvAlgorithm::Sell1`] | Default algorithm for Sliced Ellpack sparse matrix format. Provides deterministic (bitwise) results for each run. |
157/// | [`SpMvAlgorithm::Bsr1`] | Default algorithm for BSR sparse matrix format. Provides deterministic (bitwise) results for each run. Supports only [`Operation::NonTranspose`]. Supports both row-major and column-major block layouts in `A`. |
158///
159/// Calling [`spmv_preprocess`] is optional.
160/// It may accelerate subsequent calls to [`spmv`].
161/// Useful when [`spmv`] is called multiple times with the same sparsity
162/// pattern (`matrix`).
163///
164/// Calling [`spmv_preprocess`] with `buffer` makes that buffer active for `matrix` SpMV calls.
165/// Subsequent calls to [`spmv`] with `matrix` and the active buffer must use the same values for all parameters as the call to [`spmv_preprocess`].
166/// The exceptions are: `alpha`, `beta`, `x`, `y`, and the values (but not indices) of `matrix` may be different.
167/// Importantly, the buffer contents must be unmodified since the call to [`spmv_preprocess`].
168/// When [`spmv`] is called with `matrix` and its active buffer, it may read acceleration data from the buffer.
169///
170/// Calling [`spmv_preprocess`] again with `matrix` and a new buffer makes the new
171/// buffer active and makes the previously active buffer inactive.
172/// For [`spmv`], there can only be one active buffer per sparse matrix at a time.
173/// 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 [`spmv_preprocess`] once per handle with different workspace buffers.
174///
175/// Calling [`spmv`] with an inactive buffer is always permitted.
176/// However, there may be no acceleration from the preprocessing in that case.
177///
178/// For the purposes of thread safety, [`spmv_preprocess`] is writing to `matrix` internal state.
179///
180/// **Performance notes:**
181///
182/// * [`SpMvAlgorithm::Coo1`] and [`SpMvAlgorithm::Csr1`] provide higher performance than [`SpMvAlgorithm::Coo2`] and [`SpMvAlgorithm::Csr2`].
183/// * In general, [`Operation::NonTranspose`] is 3x faster than transpose or conjugate-transpose operations.
184/// * Using [`spmv_preprocess`] helps improve performance of [`spmv`] in CSR.
185///   It is beneficial when [`spmv`] runs multiple times with the same matrix
186///   ([`spmv_preprocess`] is executed only once).
187///
188/// [`spmv`] has the following properties:
189///
190/// * Requires extra storage for CSR/CSC format (all algorithms) and for COO format with [`SpMvAlgorithm::Coo2`] algorithm.
191/// * Provides deterministic (bitwise) results for each run only for [`SpMvAlgorithm::Coo2`], [`SpMvAlgorithm::Csr2`] and [`SpMvAlgorithm::Bsr1`] algorithms, and only with [`Operation::NonTranspose`].
192/// * Supports asynchronous execution.
193/// * `compute-sanitizer` could report false race conditions for this operation when `beta` is `0`.
194///   These reports result from optimization and do not affect computation correctness.
195/// * Allows the indices of `matrix` to be unsorted.
196///
197/// [`spmv`] supports the following optimizations:
198///
199/// * CUDA graph capture.
200/// * Hardware Memory Compression.
201pub fn spmv<Compute: DataTypeLike>(
202    ctx: &Context,
203    operation: Operation,
204    alpha: Scalar<'_, Compute>,
205    matrix: &SparseMatrixDescriptor,
206    x: &DenseVectorDescriptor,
207    beta: Scalar<'_, Compute>,
208    y: &mut DenseVectorDescriptor,
209    algorithm: SpMvAlgorithm,
210    external_buffer: Option<DevicePtr>,
211) -> Result<()> {
212    matrix.ensure_context(ctx)?;
213    x.ensure_context(ctx)?;
214    y.ensure_context(ctx)?;
215    ctx.bind()?;
216    ctx.require_scalar_pointer_mode(alpha, beta)?;
217
218    unsafe {
219        try_ffi!(sys::cusparseSpMV(
220            ctx.as_raw(),
221            operation.into(),
222            alpha.ptr().cast(),
223            matrix.as_raw_const(),
224            x.as_raw_const(),
225            beta.ptr().cast(),
226            y.as_raw(),
227            Compute::data_type().into(),
228            algorithm.into(),
229            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
230        ))?;
231    }
232    Ok(())
233}
singe_cusparse/spmv.rs

singe_cusparse/
spmv.rs
