[rocSPARSE] Fix sctr/rocsparse_scatter integer overflow for large nnz (AISPARSE-652)

projects/rocsparse/CHANGELOG.md

@@ -6,6 +6,7 @@ Documentation for rocSPARSE is available at
 ## (Unreleased) rocSPARSE 5.2.0
 
 ### Resolved issues
+* Fixed an overflow issue in `rocsparse_Xsctr` and the generic `rocsparse_scatter` routine. When using 64-bit indices and `nnz` >= `2^32`, a 32-bit element-index calculation could overflow and prevent some elements from being scattered. The kernel now computes element indices in 64-bit arithmetic and uses a grid-stride loop with the launch grid clamped to the device limit.
 * Fixed an integer overflow in the ELL SpMV kernels used by `rocsparse_spmv` (ELL format) and the legacy `rocsparse_Xellmv` routines when operating on matrices with more than `2^32` rows and 64-bit index types. The overflow caused rows beyond the `2^32` boundary to be mapped to incorrect row indices, leaving some output values uncomputed. Row indices are now always computed using the matrix index type.
 
 ## (Unreleased) rocSPARSE 5.1.0

projects/rocsparse/clients/testings/testing_scatter.cpp

@@ -130,4 +130,69 @@ INSTANTIATE(int64_t, double);
 INSTANTIATE(int64_t, rocsparse_float_complex);
 INSTANTIATE(int64_t, rocsparse_double_complex);
 
-void testing_scatter_extra(const Arguments& arg) {}
+void testing_scatter_extra(const Arguments& arg)
+{
+    // Regression test for AISPARSE-652.
+    //
+    // Before the fix, sctr_device computed the element index as
+    //   hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x
+    // entirely in 32-bit unsigned arithmetic. Once nnz reaches 2^32 the
+    // block-index multiply wraps around, so the tail of the sparse vector past
+    // the wrap point is never scattered into y. The fix casts the block index
+    // to the (64-bit) index type before the multiply and iterates with a
+    // grid-stride loop.
+    //
+    // This drives the 64-bit-index path of rocsparse_scatter (which dispatches
+    // to sctr_template) with nnz just past the 2^32 boundary and checks that an
+    // element beyond that boundary is actually scattered. Everything is
+    // initialized on the device and a single element is probed, so no host
+    // buffer of this size is ever needed.
+    using I = int64_t;
+    using T = float;
+
+    static constexpr int64_t two_pow_32 = static_cast<int64_t>(1) << 32;
+
+    // nnz just beyond 2^32 so at least one block has a block index whose
+    // (blockIdx * BLOCKSIZE) product overflows 32-bit arithmetic.
+    const I nnz = two_pow_32 + 512;
+
+    // rocsparse_create_spvec_descr rejects nnz > size, so the dense vector has
+    // to be at least as long as the sparse one.
+    const I size = nnz;
+
+    const rocsparse_index_base base = rocsparse_index_base_zero;
+
+    rocsparse_local_handle handle(arg);
+
+    device_vector<I> dx_ind(nnz);
+    device_vector<T> dx_val(nnz);
+    device_vector<T> dy(size);
+
+    // Filler elements all reference dense entry 0 and scatter value 0; y is 0.
+    CHECK_HIP_ERROR(hipMemset(dx_ind, 0, sizeof(I) * nnz));
+    CHECK_HIP_ERROR(hipMemset(dx_val, 0, sizeof(T) * nnz));
+    CHECK_HIP_ERROR(hipMemset(dy, 0, sizeof(T) * size));
+
+    // The probe lives past the 2^32 boundary. It scatters the only non-zero
+    // value into dense entry 1, which no filler element writes, so its result
+    // is isolated from the filler scatters to dense entry 0.
+    const I probe_idx = two_pow_32 + 5;
+    const I probe_ind = 1;
+    const T x_in      = static_cast<T>(3);
+    CHECK_HIP_ERROR(hipMemcpy(
+        static_cast<I*>(dx_ind) + probe_idx, &probe_ind, sizeof(I), hipMemcpyHostToDevice));
+    CHECK_HIP_ERROR(
+        hipMemcpy(static_cast<T*>(dx_val) + probe_idx, &x_in, sizeof(T), hipMemcpyHostToDevice));
+
+    rocsparse_local_spvec x(size, nnz, dx_ind, dx_val, get_indextype<I>(), base, get_datatype<T>());
+    rocsparse_local_dnvec y(size, dy, get_datatype<T>());
+
+    CHECK_ROCSPARSE_ERROR(testing::rocsparse_scatter(handle, x, y));
+
+    // y[1] = x_in = 3. Before the fix the wrapped block index leaves the probe
+    // element unscattered, so y[1] stays 0.
+    T y_out = static_cast<T>(0);
+    CHECK_HIP_ERROR(hipMemcpy(&y_out, static_cast<T*>(dy) + 1, sizeof(T), hipMemcpyDeviceToHost));
+
+    unit_check_scalar<T>(x_in, y_out);
+}

projects/rocsparse/clients/tests/test_scatter.yaml

@@ -149,3 +149,9 @@ Tests:
   nnz: [5, 10, 500]
   baseA: [rocsparse_index_base_zero, rocsparse_index_base_one]
   graph_test: true
+
+- name: scatter_extra
+  category: nightly
+  function: scatter_extra
+  host_memory_gb: 1
+  device_memory_gb: 70

projects/rocsparse/library/src/level1/rocsparse_sctr_impl.cpp

@@ -68,7 +68,12 @@ rocsparse_status rocsparse::sctr_template(rocsparse_handle     handle,
     hipStream_t stream = handle->stream;
 
 #define SCTR_DIM 512
-    dim3 sctr_blocks((nnz - 1) / SCTR_DIM + 1);
+    // Clamp the grid dimension against the device limit; the kernel uses a
+    // grid-stride loop so any grid size covers the full range of nnz.
+    const int64_t num_blocks
+        = rocsparse::min((static_cast<int64_t>(nnz) - 1) / SCTR_DIM + 1,
+                         static_cast<int64_t>(handle->properties.maxGridSize[0]));
+    dim3 sctr_blocks(static_cast<uint32_t>(num_blocks));
     dim3 sctr_threads(SCTR_DIM);
 
     RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::sctr_kernel<SCTR_DIM, I, T>),

projects/rocsparse/library/src/level1/sctr_device.h

@@ -32,13 +32,14 @@ namespace rocsparse
     ROCSPARSE_KERNEL(BLOCKSIZE)
     void sctr_kernel(I nnz, const T* x_val, const I* x_ind, T* y, rocsparse_index_base idx_base)
     {
-        I idx = hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x;
+        // Cast to the (possibly 64-bit) index type I before the multiply so the
+        // thread index does not wrap in 32-bit arithmetic for large nnz.
+        const I stride = static_cast<I>(hipGridDim_x) * BLOCKSIZE;
+        const I gid    = static_cast<I>(hipBlockIdx_x) * BLOCKSIZE + hipThreadIdx_x;
 
-        if(idx >= nnz)
+        for(I idx = gid; idx < nnz; idx += stride)
         {
-            return;
+            y[x_ind[idx] - idx_base] = x_val[idx];
         }
-
-        y[x_ind[idx] - idx_base] = x_val[idx];
     }
 }