singe_cusparse/
spgemm.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    operation::SpGemmDescriptor,
10    scalar::Scalar,
11    try_ffi,
12    types::{Operation, SpGemmAlgorithm},
13    utility::to_usize,
14};
15
16pub fn spgemm_work_estimation<Compute: DataTypeLike>(
17    ctx: &Context,
18    op_a: Operation,
19    op_b: Operation,
20    alpha: Scalar<'_, Compute>,
21    matrix_a: &SparseMatrixDescriptor,
22    matrix_b: &SparseMatrixDescriptor,
23    beta: Scalar<'_, Compute>,
24    matrix_c: &mut SparseMatrixDescriptor,
25    algorithm: SpGemmAlgorithm,
26    descriptor: &SpGemmDescriptor,
27    external_buffer: Option<DevicePtr>,
28) -> Result<usize> {
29    descriptor.ensure_context(ctx)?;
30    matrix_a.ensure_context(ctx)?;
31    matrix_b.ensure_context(ctx)?;
32    matrix_c.ensure_context(ctx)?;
33    ctx.bind()?;
34    ctx.require_scalar_pointer_mode(alpha, beta)?;
35
36    let mut size = 0;
37    unsafe {
38        try_ffi!(sys::cusparseSpGEMM_workEstimation(
39            ctx.as_raw(),
40            op_a.into(),
41            op_b.into(),
42            alpha.ptr().cast(),
43            matrix_a.as_raw_const(),
44            matrix_b.as_raw_const(),
45            beta.ptr().cast(),
46            matrix_c.as_raw(),
47            Compute::data_type().into(),
48            algorithm.into(),
49            descriptor.as_raw(),
50            &raw mut size,
51            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
52        ))?;
53    }
54    to_usize(size, "spgemm work estimation buffer size")
55}
56
57pub fn spgemm_estimate_memory<Compute: DataTypeLike>(
58    ctx: &Context,
59    op_a: Operation,
60    op_b: Operation,
61    alpha: Scalar<'_, Compute>,
62    matrix_a: &SparseMatrixDescriptor,
63    matrix_b: &SparseMatrixDescriptor,
64    beta: Scalar<'_, Compute>,
65    matrix_c: &mut SparseMatrixDescriptor,
66    algorithm: SpGemmAlgorithm,
67    descriptor: &SpGemmDescriptor,
68    chunk_fraction: f32,
69    external_buffer: Option<DevicePtr>,
70) -> Result<(usize, usize)> {
71    descriptor.ensure_context(ctx)?;
72    matrix_a.ensure_context(ctx)?;
73    matrix_b.ensure_context(ctx)?;
74    matrix_c.ensure_context(ctx)?;
75    ctx.bind()?;
76    ctx.require_scalar_pointer_mode(alpha, beta)?;
77
78    let mut buffer_size3 = 0;
79    let mut buffer_size2 = 0;
80    unsafe {
81        try_ffi!(sys::cusparseSpGEMM_estimateMemory(
82            ctx.as_raw(),
83            op_a.into(),
84            op_b.into(),
85            alpha.ptr().cast(),
86            matrix_a.as_raw_const(),
87            matrix_b.as_raw_const(),
88            beta.ptr().cast(),
89            matrix_c.as_raw(),
90            Compute::data_type().into(),
91            algorithm.into(),
92            descriptor.as_raw(),
93            chunk_fraction,
94            &raw mut buffer_size3,
95            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
96            &raw mut buffer_size2,
97        ))?;
98    }
99    Ok((
100        to_usize(buffer_size2, "spgemm compute buffer size")?,
101        to_usize(buffer_size3, "spgemm estimate buffer size")?,
102    ))
103}
104
105pub fn spgemm_compute<Compute: DataTypeLike>(
106    ctx: &Context,
107    op_a: Operation,
108    op_b: Operation,
109    alpha: Scalar<'_, Compute>,
110    matrix_a: &SparseMatrixDescriptor,
111    matrix_b: &SparseMatrixDescriptor,
112    beta: Scalar<'_, Compute>,
113    matrix_c: &mut SparseMatrixDescriptor,
114    algorithm: SpGemmAlgorithm,
115    descriptor: &SpGemmDescriptor,
116    external_buffer: Option<DevicePtr>,
117) -> Result<usize> {
118    descriptor.ensure_context(ctx)?;
119    matrix_a.ensure_context(ctx)?;
120    matrix_b.ensure_context(ctx)?;
121    matrix_c.ensure_context(ctx)?;
122    ctx.bind()?;
123    ctx.require_scalar_pointer_mode(alpha, beta)?;
124
125    let mut buffer_size2 = 0;
126    unsafe {
127        try_ffi!(sys::cusparseSpGEMM_compute(
128            ctx.as_raw(),
129            op_a.into(),
130            op_b.into(),
131            alpha.ptr().cast(),
132            matrix_a.as_raw_const(),
133            matrix_b.as_raw_const(),
134            beta.ptr().cast(),
135            matrix_c.as_raw(),
136            Compute::data_type().into(),
137            algorithm.into(),
138            descriptor.as_raw(),
139            &raw mut buffer_size2,
140            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
141        ))?;
142    }
143    to_usize(buffer_size2, "spgemm compute buffer size")
144}
145
146pub fn spgemm_copy<Compute: DataTypeLike>(
147    ctx: &Context,
148    op_a: Operation,
149    op_b: Operation,
150    alpha: Scalar<'_, Compute>,
151    matrix_a: &SparseMatrixDescriptor,
152    matrix_b: &SparseMatrixDescriptor,
153    beta: Scalar<'_, Compute>,
154    matrix_c: &mut SparseMatrixDescriptor,
155    algorithm: SpGemmAlgorithm,
156    descriptor: &SpGemmDescriptor,
157) -> Result<()> {
158    descriptor.ensure_context(ctx)?;
159    matrix_a.ensure_context(ctx)?;
160    matrix_b.ensure_context(ctx)?;
161    matrix_c.ensure_context(ctx)?;
162    ctx.bind()?;
163    ctx.require_scalar_pointer_mode(alpha, beta)?;
164
165    unsafe {
166        try_ffi!(sys::cusparseSpGEMM_copy(
167            ctx.as_raw(),
168            op_a.into(),
169            op_b.into(),
170            alpha.ptr().cast(),
171            matrix_a.as_raw_const(),
172            matrix_b.as_raw_const(),
173            beta.ptr().cast(),
174            matrix_c.as_raw(),
175            Compute::data_type().into(),
176            algorithm.into(),
177            descriptor.as_raw(),
178        ))?;
179    }
180    Ok(())
181}
182
183pub fn spgemmreuse_work_estimation(
184    ctx: &Context,
185    op_a: Operation,
186    op_b: Operation,
187    matrix_a: &SparseMatrixDescriptor,
188    matrix_b: &SparseMatrixDescriptor,
189    matrix_c: &mut SparseMatrixDescriptor,
190    algorithm: SpGemmAlgorithm,
191    descriptor: &SpGemmDescriptor,
192    external_buffer1: Option<DevicePtr>,
193) -> Result<usize> {
194    descriptor.ensure_context(ctx)?;
195    matrix_a.ensure_context(ctx)?;
196    matrix_b.ensure_context(ctx)?;
197    matrix_c.ensure_context(ctx)?;
198    ctx.bind()?;
199    let mut size = 0;
200    unsafe {
201        try_ffi!(sys::cusparseSpGEMMreuse_workEstimation(
202            ctx.as_raw(),
203            op_a.into(),
204            op_b.into(),
205            matrix_a.as_raw_const(),
206            matrix_b.as_raw_const(),
207            matrix_c.as_raw(),
208            algorithm.into(),
209            descriptor.as_raw(),
210            &raw mut size,
211            external_buffer1.unwrap_or(DevicePtr::null()).as_ptr() as _,
212        ))?;
213    }
214    to_usize(size, "spgemmreuse work estimation buffer size")
215}
216
217/// Computes nonzero counts for the reusable sparse matrix-matrix multiplication pipeline.
218///
219/// * This requires temporary extra storage that is allocated internally.
220/// * Supports asynchronous execution if the Stream Ordered Memory Allocator is available.
221/// * Supports CUDA graph capture if the Stream Ordered Memory Allocator is available.
222pub fn spgemmreuse_nnz(
223    ctx: &Context,
224    op_a: Operation,
225    op_b: Operation,
226    matrix_a: &SparseMatrixDescriptor,
227    matrix_b: &SparseMatrixDescriptor,
228    matrix_c: &mut SparseMatrixDescriptor,
229    algorithm: SpGemmAlgorithm,
230    descriptor: &SpGemmDescriptor,
231    external_buffer2: Option<DevicePtr>,
232    external_buffer3: Option<DevicePtr>,
233    external_buffer4: Option<DevicePtr>,
234) -> Result<(usize, usize, usize)> {
235    descriptor.ensure_context(ctx)?;
236    matrix_a.ensure_context(ctx)?;
237    matrix_b.ensure_context(ctx)?;
238    matrix_c.ensure_context(ctx)?;
239    ctx.bind()?;
240    let mut size2 = 0;
241    let mut size3 = 0;
242    let mut size4 = 0;
243    unsafe {
244        try_ffi!(sys::cusparseSpGEMMreuse_nnz(
245            ctx.as_raw(),
246            op_a.into(),
247            op_b.into(),
248            matrix_a.as_raw_const(),
249            matrix_b.as_raw_const(),
250            matrix_c.as_raw(),
251            algorithm.into(),
252            descriptor.as_raw(),
253            &raw mut size2,
254            external_buffer2.unwrap_or(DevicePtr::null()).as_ptr() as _,
255            &raw mut size3,
256            external_buffer3.unwrap_or(DevicePtr::null()).as_ptr() as _,
257            &raw mut size4,
258            external_buffer4.unwrap_or(DevicePtr::null()).as_ptr() as _,
259        ))?;
260    }
261    Ok((
262        to_usize(size2, "spgemmreuse nnz buffer2 size")?,
263        to_usize(size3, "spgemmreuse nnz buffer3 size")?,
264        to_usize(size4, "spgemmreuse nnz buffer4 size")?,
265    ))
266}
267
268pub fn spgemmreuse_copy(
269    ctx: &Context,
270    op_a: Operation,
271    op_b: Operation,
272    matrix_a: &SparseMatrixDescriptor,
273    matrix_b: &SparseMatrixDescriptor,
274    matrix_c: &mut SparseMatrixDescriptor,
275    algorithm: SpGemmAlgorithm,
276    descriptor: &SpGemmDescriptor,
277    external_buffer5: Option<DevicePtr>,
278) -> Result<usize> {
279    descriptor.ensure_context(ctx)?;
280    matrix_a.ensure_context(ctx)?;
281    matrix_b.ensure_context(ctx)?;
282    matrix_c.ensure_context(ctx)?;
283    ctx.bind()?;
284    let mut size = 0;
285    unsafe {
286        try_ffi!(sys::cusparseSpGEMMreuse_copy(
287            ctx.as_raw(),
288            op_a.into(),
289            op_b.into(),
290            matrix_a.as_raw_const(),
291            matrix_b.as_raw_const(),
292            matrix_c.as_raw(),
293            algorithm.into(),
294            descriptor.as_raw(),
295            &raw mut size,
296            external_buffer5.unwrap_or(DevicePtr::null()).as_ptr() as _,
297        ))?;
298    }
299    to_usize(size, "spgemmreuse copy buffer size")
300}
301
302pub fn spgemmreuse_compute<Compute: DataTypeLike>(
303    ctx: &Context,
304    op_a: Operation,
305    op_b: Operation,
306    alpha: Scalar<'_, Compute>,
307    matrix_a: &SparseMatrixDescriptor,
308    matrix_b: &SparseMatrixDescriptor,
309    beta: Scalar<'_, Compute>,
310    matrix_c: &mut SparseMatrixDescriptor,
311    algorithm: SpGemmAlgorithm,
312    descriptor: &SpGemmDescriptor,
313) -> Result<()> {
314    descriptor.ensure_context(ctx)?;
315    matrix_a.ensure_context(ctx)?;
316    matrix_b.ensure_context(ctx)?;
317    matrix_c.ensure_context(ctx)?;
318    ctx.bind()?;
319    ctx.require_scalar_pointer_mode(alpha, beta)?;
320    unsafe {
321        try_ffi!(sys::cusparseSpGEMMreuse_compute(
322            ctx.as_raw(),
323            op_a.into(),
324            op_b.into(),
325            alpha.ptr().cast(),
326            matrix_a.as_raw_const(),
327            matrix_b.as_raw_const(),
328            beta.ptr().cast(),
329            matrix_c.as_raw(),
330            Compute::data_type().into(),
331            algorithm.into(),
332            descriptor.as_raw(),
333        ))?;
334    }
335    Ok(())
336}
singe_cusparse/spgemm.rs

singe_cusparse/
spgemm.rs
