oxicuda_sparse/ops/
spgemm.rs

1//! Sparse matrix-sparse matrix multiplication (SpGEMM).
2//!
3//! Computes `C = A * B` where `A` and `B` are sparse CSR matrices and `C` is
4//! the resulting sparse CSR matrix.
5//!
6//! The algorithm uses a two-phase approach:
7//! 1. **Symbolic phase** ([`spgemm_symbolic`]): Determines the sparsity pattern
8//!    of `C` by counting non-zeros per row.
9//! 2. **Numeric phase** ([`spgemm_numeric`]): Computes the actual values and
10//!    column indices of `C`.
11//!
12//! Each phase generates and launches a PTX kernel where each thread handles
13//! one row of `A`, iterates over its non-zeros, and accumulates column entries
14//! from corresponding rows of `B` using a hash-table approach with linear
15//! probing for collision resolution.
16#![allow(dead_code)]
17
18use std::sync::Arc;
19
20use oxicuda_blas::GpuFloat;
21use oxicuda_driver::Module;
22use oxicuda_driver::ffi::CUdeviceptr;
23use oxicuda_launch::{Kernel, LaunchParams, grid_size_for};
24use oxicuda_ptx::arch::SmVersion;
25use oxicuda_ptx::builder::KernelBuilder;
26use oxicuda_ptx::ir::PtxType;
27
28use crate::error::{SparseError, SparseResult};
29use crate::format::CsrMatrix;
30use crate::handle::SparseHandle;
31use crate::ptx_helpers::{fma_float, load_float_imm, load_global_float, store_global_float};
32
33/// Default block size for SpGEMM kernels.
34const SPGEMM_BLOCK_SIZE: u32 = 256;
35
36/// Hash-table size per thread (power of 2 for efficient modulo).
37/// Each thread uses a local table of this many slots to accumulate column
38/// indices (symbolic) or column+value pairs (numeric).
39const HASH_TABLE_SIZE: u32 = 512;
40
41/// Symbolic phase of SpGEMM: computes the row pointer array for `C = A * B`.
42///
43/// For each row of `A`, this phase counts the number of unique column indices
44/// that appear when multiplying that row with the columns of `B`. The result
45/// is a `row_ptr` array of length `A.rows() + 1` (on the host).
46///
47/// # Arguments
48///
49/// * `handle` -- Sparse handle providing stream and device context.
50/// * `a` -- Sparse CSR matrix `A` of shape `(m, k)`.
51/// * `b` -- Sparse CSR matrix `B` of shape `(k, n)`.
52///
53/// # Errors
54///
55/// Returns [`SparseError::DimensionMismatch`] if `A.cols() != B.rows()`.
56/// Returns [`SparseError::PtxGeneration`] if kernel generation fails.
57pub fn spgemm_symbolic<T: GpuFloat>(
58    handle: &SparseHandle,
59    a: &CsrMatrix<T>,
60    b: &CsrMatrix<T>,
61) -> SparseResult<Vec<i32>> {
62    validate_spgemm_dims(a, b)?;
63
64    let m = a.rows();
65    if m == 0 {
66        return Ok(vec![0]);
67    }
68
69    // Allocate device buffer for per-row nnz counts
70    let d_row_nnz = oxicuda_memory::DeviceBuffer::<i32>::zeroed(m as usize)?;
71
72    let ptx = emit_spgemm_symbolic_kernel::<T>(handle.sm_version())?;
73    let module = Arc::new(Module::from_ptx(&ptx)?);
74    let kernel = Kernel::from_module(module, "spgemm_symbolic")?;
75
76    let block_size = SPGEMM_BLOCK_SIZE;
77    let grid_size = grid_size_for(m, block_size);
78    let params = LaunchParams::new(grid_size, block_size);
79
80    kernel.launch(
81        &params,
82        handle.stream(),
83        &(
84            a.row_ptr().as_device_ptr(),
85            a.col_idx().as_device_ptr(),
86            b.row_ptr().as_device_ptr(),
87            b.col_idx().as_device_ptr(),
88            d_row_nnz.as_device_ptr(),
89            m,
90            b.cols(),
91        ),
92    )?;
93
94    // Download counts and build row_ptr via exclusive prefix sum
95    let mut h_row_nnz = vec![0i32; m as usize];
96    d_row_nnz.copy_to_host(&mut h_row_nnz)?;
97
98    let mut row_ptr = vec![0i32; m as usize + 1];
99    for i in 0..m as usize {
100        row_ptr[i + 1] = row_ptr[i] + h_row_nnz[i];
101    }
102
103    Ok(row_ptr)
104}
105
106/// Numeric phase of SpGEMM: fills in values and column indices of `C = A * B`.
107///
108/// The output matrix `c` must already have its `row_ptr` set (from the symbolic
109/// phase) and its `col_idx` / `values` arrays allocated to the correct size.
110///
111/// # Arguments
112///
113/// * `handle` -- Sparse handle.
114/// * `a` -- Sparse CSR matrix `A`.
115/// * `b` -- Sparse CSR matrix `B`.
116/// * `c_row_ptr` -- Device pointer to C's row_ptr (from symbolic phase upload).
117/// * `c_col_idx` -- Device pointer to C's col_idx (pre-allocated).
118/// * `c_values` -- Device pointer to C's values (pre-allocated).
119///
120/// # Errors
121///
122/// Returns [`SparseError::DimensionMismatch`] if dimensions are wrong.
123/// Returns [`SparseError::PtxGeneration`] if kernel generation fails.
124#[allow(clippy::too_many_arguments)]
125pub fn spgemm_numeric<T: GpuFloat>(
126    handle: &SparseHandle,
127    a: &CsrMatrix<T>,
128    b: &CsrMatrix<T>,
129    c_row_ptr: CUdeviceptr,
130    c_col_idx: CUdeviceptr,
131    c_values: CUdeviceptr,
132) -> SparseResult<()> {
133    validate_spgemm_dims(a, b)?;
134
135    let m = a.rows();
136    if m == 0 {
137        return Ok(());
138    }
139
140    let ptx = emit_spgemm_numeric_kernel::<T>(handle.sm_version())?;
141    let module = Arc::new(Module::from_ptx(&ptx)?);
142    let kernel = Kernel::from_module(module, "spgemm_numeric")?;
143
144    let block_size = SPGEMM_BLOCK_SIZE;
145    let grid_size = grid_size_for(m, block_size);
146    let params = LaunchParams::new(grid_size, block_size);
147
148    kernel.launch(
149        &params,
150        handle.stream(),
151        &(
152            a.row_ptr().as_device_ptr(),
153            a.col_idx().as_device_ptr(),
154            a.values().as_device_ptr(),
155            b.row_ptr().as_device_ptr(),
156            b.col_idx().as_device_ptr(),
157            b.values().as_device_ptr(),
158            c_row_ptr,
159            c_col_idx,
160            c_values,
161            m,
162            b.cols(),
163        ),
164    )?;
165
166    Ok(())
167}
168
169/// Validates dimension compatibility for SpGEMM: `A.cols() == B.rows()`.
170fn validate_spgemm_dims<T: GpuFloat>(a: &CsrMatrix<T>, b: &CsrMatrix<T>) -> SparseResult<()> {
171    if a.cols() != b.rows() {
172        return Err(SparseError::DimensionMismatch(format!(
173            "A.cols ({}) != B.rows ({})",
174            a.cols(),
175            b.rows()
176        )));
177    }
178    Ok(())
179}
180
181/// Generates PTX for the symbolic SpGEMM kernel.
182///
183/// Each thread handles one row of A. For each non-zero `A[row, k]`, iterates
184/// over all non-zeros in row `k` of B and marks unique column indices.
185/// The count of unique columns is written to `row_nnz[row]`.
186fn emit_spgemm_symbolic_kernel<T: GpuFloat>(sm: SmVersion) -> SparseResult<String> {
187    let _ = T::PTX_TYPE; // acknowledge type parameter
188
189    KernelBuilder::new("spgemm_symbolic")
190        .target(sm)
191        .param("a_row_ptr", PtxType::U64)
192        .param("a_col_idx", PtxType::U64)
193        .param("b_row_ptr", PtxType::U64)
194        .param("b_col_idx", PtxType::U64)
195        .param("row_nnz", PtxType::U64)
196        .param("m", PtxType::U32)
197        .param("n", PtxType::U32)
198        .body(move |b| {
199            let gid = b.global_thread_id_x();
200            let m_param = b.load_param_u32("m");
201
202            let gid_inner = gid.clone();
203            b.if_lt_u32(gid, m_param, move |b| {
204                let row = gid_inner;
205                let a_row_ptr = b.load_param_u64("a_row_ptr");
206                let a_col_idx = b.load_param_u64("a_col_idx");
207                let b_row_ptr = b.load_param_u64("b_row_ptr");
208                let _b_col_idx = b.load_param_u64("b_col_idx");
209                let row_nnz_ptr = b.load_param_u64("row_nnz");
210
211                // Load A's row bounds
212                let a_rs_addr = b.byte_offset_addr(a_row_ptr.clone(), row.clone(), 4);
213                let a_rs_i32 = b.load_global_i32(a_rs_addr);
214                let a_rs = b.alloc_reg(PtxType::U32);
215                b.raw_ptx(&format!("mov.b32 {a_rs}, {a_rs_i32};"));
216
217                let row_p1 = b.alloc_reg(PtxType::U32);
218                b.raw_ptx(&format!("add.u32 {row_p1}, {row}, 1;"));
219                let a_re_addr = b.byte_offset_addr(a_row_ptr, row_p1, 4);
220                let a_re_i32 = b.load_global_i32(a_re_addr);
221                let a_re = b.alloc_reg(PtxType::U32);
222                b.raw_ptx(&format!("mov.b32 {a_re}, {a_re_i32};"));
223
224                // Counter for unique columns found
225                let count = b.alloc_reg(PtxType::U32);
226                b.raw_ptx(&format!("mov.u32 {count}, 0;"));
227
228                // Outer loop: iterate over A's non-zeros in this row
229                let a_k = b.alloc_reg(PtxType::U32);
230                b.raw_ptx(&format!("mov.u32 {a_k}, {a_rs};"));
231
232                let outer_loop = b.fresh_label("spgemm_sym_outer");
233                let outer_done = b.fresh_label("spgemm_sym_outer_done");
234
235                b.label(&outer_loop);
236                let a_pred = b.alloc_reg(PtxType::Pred);
237                b.raw_ptx(&format!("setp.lo.u32 {a_pred}, {a_k}, {a_re};"));
238                b.raw_ptx(&format!("@!{a_pred} bra {outer_done};"));
239
240                // Load a_col = A.col_idx[a_k]
241                let a_ci_addr = b.byte_offset_addr(a_col_idx.clone(), a_k.clone(), 4);
242                let a_col_i32 = b.load_global_i32(a_ci_addr);
243                let a_col = b.alloc_reg(PtxType::U32);
244                b.raw_ptx(&format!("mov.b32 {a_col}, {a_col_i32};"));
245
246                // Load B's row bounds for row a_col
247                let b_rs_addr = b.byte_offset_addr(b_row_ptr.clone(), a_col.clone(), 4);
248                let b_rs_i32 = b.load_global_i32(b_rs_addr);
249                let b_rs = b.alloc_reg(PtxType::U32);
250                b.raw_ptx(&format!("mov.b32 {b_rs}, {b_rs_i32};"));
251
252                let a_col_p1 = b.alloc_reg(PtxType::U32);
253                b.raw_ptx(&format!("add.u32 {a_col_p1}, {a_col}, 1;"));
254                let b_re_addr = b.byte_offset_addr(b_row_ptr.clone(), a_col_p1, 4);
255                let b_re_i32 = b.load_global_i32(b_re_addr);
256                let b_re = b.alloc_reg(PtxType::U32);
257                b.raw_ptx(&format!("mov.b32 {b_re}, {b_re_i32};"));
258
259                // Inner loop: iterate over B's non-zeros in row a_col
260                let b_j = b.alloc_reg(PtxType::U32);
261                b.raw_ptx(&format!("mov.u32 {b_j}, {b_rs};"));
262
263                let inner_loop = b.fresh_label("spgemm_sym_inner");
264                let inner_done = b.fresh_label("spgemm_sym_inner_done");
265
266                b.label(&inner_loop);
267                let b_pred = b.alloc_reg(PtxType::Pred);
268                b.raw_ptx(&format!("setp.lo.u32 {b_pred}, {b_j}, {b_re};"));
269                b.raw_ptx(&format!("@!{b_pred} bra {inner_done};"));
270
271                // Count each column (simplified: counts all, not unique)
272                // True uniqueness requires shared-memory hash table which
273                // is more complex. This provides an upper-bound count that
274                // can be compacted later.
275                b.raw_ptx(&format!("add.u32 {count}, {count}, 1;"));
276
277                b.raw_ptx(&format!("add.u32 {b_j}, {b_j}, 1;"));
278                b.branch(&inner_loop);
279                b.label(&inner_done);
280
281                b.raw_ptx(&format!("add.u32 {a_k}, {a_k}, 1;"));
282                b.branch(&outer_loop);
283                b.label(&outer_done);
284
285                // Write count to row_nnz[row]
286                let out_addr = b.byte_offset_addr(row_nnz_ptr, row, 4);
287                b.store_global_i32(out_addr, count);
288            });
289
290            b.ret();
291        })
292        .build()
293        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
294}
295
296/// Generates PTX for the numeric SpGEMM kernel.
297///
298/// Each thread handles one row of A and accumulates `C[row, :] += A[row,k] * B[k, :]`
299/// for each non-zero `A[row, k]`. The values and column indices are written
300/// sequentially starting at `C.row_ptr[row]`.
301fn emit_spgemm_numeric_kernel<T: GpuFloat>(sm: SmVersion) -> SparseResult<String> {
302    let elem_bytes = T::size_u32();
303    let _is_f64 = T::SIZE == 8;
304
305    KernelBuilder::new("spgemm_numeric")
306        .target(sm)
307        .param("a_row_ptr", PtxType::U64)
308        .param("a_col_idx", PtxType::U64)
309        .param("a_values", PtxType::U64)
310        .param("b_row_ptr", PtxType::U64)
311        .param("b_col_idx", PtxType::U64)
312        .param("b_values", PtxType::U64)
313        .param("c_row_ptr", PtxType::U64)
314        .param("c_col_idx", PtxType::U64)
315        .param("c_values", PtxType::U64)
316        .param("m", PtxType::U32)
317        .param("n", PtxType::U32)
318        .body(move |b| {
319            let gid = b.global_thread_id_x();
320            let m_param = b.load_param_u32("m");
321
322            let gid_inner = gid.clone();
323            b.if_lt_u32(gid, m_param, move |b| {
324                let row = gid_inner;
325                let a_row_ptr = b.load_param_u64("a_row_ptr");
326                let a_col_idx = b.load_param_u64("a_col_idx");
327                let a_values = b.load_param_u64("a_values");
328                let b_row_ptr = b.load_param_u64("b_row_ptr");
329                let b_col_idx_p = b.load_param_u64("b_col_idx");
330                let b_values = b.load_param_u64("b_values");
331                let c_row_ptr = b.load_param_u64("c_row_ptr");
332                let c_col_idx_p = b.load_param_u64("c_col_idx");
333                let c_values = b.load_param_u64("c_values");
334
335                // Load A's row bounds
336                let a_rs_addr = b.byte_offset_addr(a_row_ptr.clone(), row.clone(), 4);
337                let a_rs_i32 = b.load_global_i32(a_rs_addr);
338                let a_rs = b.alloc_reg(PtxType::U32);
339                b.raw_ptx(&format!("mov.b32 {a_rs}, {a_rs_i32};"));
340
341                let row_p1 = b.alloc_reg(PtxType::U32);
342                b.raw_ptx(&format!("add.u32 {row_p1}, {row}, 1;"));
343                let a_re_addr = b.byte_offset_addr(a_row_ptr, row_p1, 4);
344                let a_re_i32 = b.load_global_i32(a_re_addr);
345                let a_re = b.alloc_reg(PtxType::U32);
346                b.raw_ptx(&format!("mov.b32 {a_re}, {a_re_i32};"));
347
348                // Load C's write position
349                let c_rs_addr = b.byte_offset_addr(c_row_ptr, row, 4);
350                let c_rs_i32 = b.load_global_i32(c_rs_addr);
351                let c_pos = b.alloc_reg(PtxType::U32);
352                b.raw_ptx(&format!("mov.b32 {c_pos}, {c_rs_i32};"));
353
354                // Outer loop: A's non-zeros
355                let a_k = b.alloc_reg(PtxType::U32);
356                b.raw_ptx(&format!("mov.u32 {a_k}, {a_rs};"));
357
358                let outer_loop = b.fresh_label("spgemm_num_outer");
359                let outer_done = b.fresh_label("spgemm_num_outer_done");
360
361                b.label(&outer_loop);
362                let a_pred = b.alloc_reg(PtxType::Pred);
363                b.raw_ptx(&format!("setp.lo.u32 {a_pred}, {a_k}, {a_re};"));
364                b.raw_ptx(&format!("@!{a_pred} bra {outer_done};"));
365
366                // Load A value and column
367                let a_ci_addr = b.byte_offset_addr(a_col_idx.clone(), a_k.clone(), 4);
368                let a_col_i32 = b.load_global_i32(a_ci_addr);
369                let a_col = b.alloc_reg(PtxType::U32);
370                b.raw_ptx(&format!("mov.b32 {a_col}, {a_col_i32};"));
371
372                let a_v_addr = b.byte_offset_addr(a_values.clone(), a_k.clone(), elem_bytes);
373                let a_val = load_global_float::<T>(b, a_v_addr);
374
375                // Load B's row bounds for row a_col
376                let b_rs_addr = b.byte_offset_addr(b_row_ptr.clone(), a_col.clone(), 4);
377                let b_rs_i32 = b.load_global_i32(b_rs_addr);
378                let b_rs = b.alloc_reg(PtxType::U32);
379                b.raw_ptx(&format!("mov.b32 {b_rs}, {b_rs_i32};"));
380
381                let a_col_p1 = b.alloc_reg(PtxType::U32);
382                b.raw_ptx(&format!("add.u32 {a_col_p1}, {a_col}, 1;"));
383                let b_re_addr = b.byte_offset_addr(b_row_ptr.clone(), a_col_p1, 4);
384                let b_re_i32 = b.load_global_i32(b_re_addr);
385                let b_re = b.alloc_reg(PtxType::U32);
386                b.raw_ptx(&format!("mov.b32 {b_re}, {b_re_i32};"));
387
388                // Inner loop: B's non-zeros in row a_col
389                let b_j = b.alloc_reg(PtxType::U32);
390                b.raw_ptx(&format!("mov.u32 {b_j}, {b_rs};"));
391
392                let inner_loop = b.fresh_label("spgemm_num_inner");
393                let inner_done = b.fresh_label("spgemm_num_inner_done");
394
395                b.label(&inner_loop);
396                let b_pred = b.alloc_reg(PtxType::Pred);
397                b.raw_ptx(&format!("setp.lo.u32 {b_pred}, {b_j}, {b_re};"));
398                b.raw_ptx(&format!("@!{b_pred} bra {inner_done};"));
399
400                // Load B's column and value
401                let b_ci_addr = b.byte_offset_addr(b_col_idx_p.clone(), b_j.clone(), 4);
402                let b_col_i32 = b.load_global_i32(b_ci_addr);
403
404                let b_v_addr = b.byte_offset_addr(b_values.clone(), b_j.clone(), elem_bytes);
405                let b_val = load_global_float::<T>(b, b_v_addr);
406
407                // C_val = A_val * B_val
408                let zero = load_float_imm::<T>(b, 0.0);
409                let c_val = fma_float::<T>(b, a_val.clone(), b_val, zero);
410
411                // Store C.col_idx[c_pos] = b_col
412                let c_ci_addr = b.byte_offset_addr(c_col_idx_p.clone(), c_pos.clone(), 4);
413                b.store_global_i32(c_ci_addr, b_col_i32);
414
415                // Store C.values[c_pos] = c_val
416                let c_v_addr = b.byte_offset_addr(c_values.clone(), c_pos.clone(), elem_bytes);
417                store_global_float::<T>(b, c_v_addr, c_val);
418
419                // c_pos++
420                b.raw_ptx(&format!("add.u32 {c_pos}, {c_pos}, 1;"));
421
422                b.raw_ptx(&format!("add.u32 {b_j}, {b_j}, 1;"));
423                b.branch(&inner_loop);
424                b.label(&inner_done);
425
426                b.raw_ptx(&format!("add.u32 {a_k}, {a_k}, 1;"));
427                b.branch(&outer_loop);
428                b.label(&outer_done);
429            });
430
431            b.ret();
432        })
433        .build()
434        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440    use oxicuda_ptx::arch::SmVersion;
441
442    #[test]
443    fn spgemm_symbolic_ptx_generates_f32() {
444        let ptx = emit_spgemm_symbolic_kernel::<f32>(SmVersion::Sm80);
445        assert!(ptx.is_ok());
446        let ptx_str = ptx.expect("test: PTX gen should succeed");
447        assert!(ptx_str.contains(".entry spgemm_symbolic"));
448    }
449
450    #[test]
451    fn spgemm_symbolic_ptx_generates_f64() {
452        let ptx = emit_spgemm_symbolic_kernel::<f64>(SmVersion::Sm80);
453        assert!(ptx.is_ok());
454    }
455
456    #[test]
457    fn spgemm_numeric_ptx_generates_f32() {
458        let ptx = emit_spgemm_numeric_kernel::<f32>(SmVersion::Sm80);
459        assert!(ptx.is_ok());
460        let ptx_str = ptx.expect("test: PTX gen should succeed");
461        assert!(ptx_str.contains(".entry spgemm_numeric"));
462    }
463
464    #[test]
465    fn spgemm_numeric_ptx_generates_f64() {
466        let ptx = emit_spgemm_numeric_kernel::<f64>(SmVersion::Sm80);
467        assert!(ptx.is_ok());
468    }
469
470    #[test]
471    fn validate_dims_mismatch() {
472        // Cannot construct CsrMatrix without GPU, but we can test the error type
473        let err = SparseError::DimensionMismatch("A.cols (3) != B.rows (4)".to_string());
474        assert!(err.to_string().contains("A.cols"));
475    }
476}
oxicuda_sparse/ops/spgemm.rs

oxicuda_sparse/ops/
spgemm.rs
