numr 0.5.1

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

//! Native ESC (Exact Symbolic Computation) + Hash Accumulation SpGEMM implementation.
//!
//! This implements the SAME algorithm as the CPU backend for perfect backend parity.
//!
//! Algorithm:
//! Phase 1: Symbolic - Count unique columns per output row using bitmap
//! Phase 2: Numeric - Compute values with pre-sized hash accumulator (no resizing)

use super::super::{CudaClient, CudaRuntime};
use crate::dtype::{DType, Element};
use crate::error::{Error, Result};
use crate::runtime::cuda::kernels::{
    exclusive_scan_i32_gpu, launch_cast, spgemm_numeric_phase, spgemm_symbolic_phase,
};
use crate::sparse::CsrData;
use crate::tensor::Tensor;
use cudarc::driver::DeviceRepr;
use cudarc::types::CudaTypeName;

/// Public function to be called from the combined trait implementation
pub(super) fn esc_spgemm_csr(
    client: &CudaClient,
    a_csr: &CsrData<CudaRuntime>,
    b_csr: &CsrData<CudaRuntime>,
) -> Result<CsrData<CudaRuntime>> {
    use crate::algorithm::sparse::validate_spgemm_shapes;

    // Validate shapes
    let ([_m, _n], _k) = validate_spgemm_shapes(a_csr.shape, b_csr.shape)?;

    let dtype = a_csr.values.dtype();

    // Dispatch to typed implementation
    crate::dispatch_dtype!(dtype, T => {
        client.sparse_matmul_csr_esc::<T>(
            &a_csr.row_ptrs,
            &a_csr.col_indices,
            &a_csr.values,
            &b_csr.row_ptrs,
            &b_csr.col_indices,
            &b_csr.values,
            a_csr.shape,
            b_csr.shape,
        )
    }, "esc_spgemm")
}

impl CudaClient {
    /// Sparse × Sparse matrix multiplication using native ESC+Hash algorithm
    ///
    /// This uses the SAME algorithm as the CPU backend:
    /// - Phase 1: Symbolic - count NNZ per row using bitmap
    /// - Phase 2: Numeric - compute values with pre-sized hash accumulator
    ///
    /// Backend parity is guaranteed by algorithmic consistency.
    pub(super) fn sparse_matmul_csr_esc<T: Element + CudaTypeName + Copy + DeviceRepr>(
        &self,
        a_row_ptrs: &Tensor<CudaRuntime>,
        a_col_indices: &Tensor<CudaRuntime>,
        a_values: &Tensor<CudaRuntime>,
        b_row_ptrs: &Tensor<CudaRuntime>,
        b_col_indices: &Tensor<CudaRuntime>,
        b_values: &Tensor<CudaRuntime>,
        a_shape: [usize; 2],
        b_shape: [usize; 2],
    ) -> Result<crate::sparse::CsrData<CudaRuntime>> {
        use crate::runtime::common::sparse_utils::zero_tolerance;

        let [m, _k] = a_shape;
        let [_, n] = b_shape;
        let device = a_values.device();
        let dtype = a_values.dtype();

        // ====================================================================
        // PHASE 1: Symbolic (Count NNZ per output row)
        // ====================================================================
        let row_nnz = unsafe {
            spgemm_symbolic_phase(
                &self.context,
                &self.stream,
                self.device.index,
                device,
                a_row_ptrs,
                a_col_indices,
                b_row_ptrs,
                b_col_indices,
                m,
                n,
            )?
        };

        // Build row_ptrs via exclusive scan
        let (c_row_ptrs_i32, total_nnz_i32) = unsafe {
            exclusive_scan_i32_gpu(
                &self.context,
                &self.stream,
                self.device.index,
                device,
                &row_nnz,
            )?
        };

        // Cast row_ptrs from I32 to I64 (numr internal format)
        let c_row_ptrs = unsafe {
            let output = Tensor::zeros(&[m + 1], DType::I64, device);
            launch_cast(
                &self.context,
                &self.stream,
                self.device.index,
                DType::I32,
                DType::I64,
                c_row_ptrs_i32.ptr(),
                output.ptr(),
                m + 1,
            )?;
            output
        };

        // Get total NNZ from the last element of row_ptrs
        let total_nnz = total_nnz_i32;

        // Allocate output arrays
        let c_col_indices = Tensor::zeros(&[total_nnz], crate::dtype::DType::I64, device);
        let c_values = Tensor::zeros(&[total_nnz], dtype, device);

        // ====================================================================
        // PHASE 2: Numeric (Compute values with pre-sized hash accumulator)
        // ====================================================================
        let threshold = zero_tolerance::<T>();
        unsafe {
            spgemm_numeric_phase::<T>(
                &self.context,
                &self.stream,
                self.device.index,
                a_row_ptrs,
                a_col_indices,
                a_values,
                b_row_ptrs,
                b_col_indices,
                b_values,
                &c_row_ptrs,
                &c_col_indices,
                &c_values,
                m,
                n,
                T::from_f64(threshold),
            )?;
        }

        // Synchronize to ensure kernels complete
        self.stream
            .synchronize()
            .map_err(|e| Error::Internal(format!("CUDA stream synchronization failed: {:?}", e)))?;

        Ok(crate::sparse::CsrData::new(
            c_row_ptrs,
            c_col_indices,
            c_values,
            [m, n],
        )?)
    }
}