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

singe_cusparse/
spmv.rs
