musa opt

capi/ctest/accuracy/test_scatter.cpp

@@ -25,6 +25,10 @@
 
 #include "common.hpp"
 
+#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS)
+#include <musa.h>
+#endif
+
 using namespace fstest;
 
 namespace {
@@ -222,4 +226,73 @@ TEST_F(ScatterAccuracy, SizeMismatchIsRejected) {
     flagsparseDestroySpVec(vecX);
 }
 
+#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS)
+template <typename I>
+void musa_live_copies(flagsparseHandle_t handle, flagsparseDataType_t dtype,
+                       flagsparseIndexType_t itype, std::size_t bytes) {
+    constexpr int n = 259, nnz = 71;  // partial blocks, distinct mutable indices
+    std::vector<std::uint8_t> dense(n * bytes), poison(nnz * bytes, 0xa5);
+    std::vector<std::uint8_t> initial(n * bytes, 0x5a);
+    std::vector<I> indices(nnz);
+    DeviceBuffer d_dense = DeviceBuffer::from(dense);
+    DeviceBuffer d_sparse = DeviceBuffer::from(poison);
+    DeviceBuffer d_out = DeviceBuffer::from(initial);
+    DeviceBuffer d_indices = DeviceBuffer::from(indices);
+    flagsparseDnVecDescr_t source = nullptr, destination = nullptr;
+    flagsparseSpVecDescr_t sparse = nullptr;
+    ASSERT_EQ(flagsparseCreateDnVec(&source, n, d_dense.get(), dtype),
+              FLAGSPARSE_STATUS_SUCCESS);
+    ASSERT_EQ(flagsparseCreateDnVec(&destination, n, d_out.get(), dtype),
+              FLAGSPARSE_STATUS_SUCCESS);
+    ASSERT_EQ(flagsparseCreateSpVec(&sparse, n, nnz, d_indices.get(), d_sparse.get(),
+                                     itype, FLAGSPARSE_INDEX_BASE_ZERO, dtype),
+              FLAGSPARSE_STATUS_SUCCESS);
+    for (int repeat = 0; repeat < 3; ++repeat) {
+        // Treat floating-point values as bytes to check NaNs/sign bits exactly.
+        for (std::size_t j = 0; j < dense.size(); ++j)
+            dense[j] = static_cast<std::uint8_t>(j * 37 + repeat * 19);
+        for (int j = 0; j < nnz; ++j) indices[j] = (j * 3 + repeat) % n;
+        ASSERT_EQ(to_device(d_dense.get(), dense.data(), dense.size()), FLAGSPARSE_STATUS_SUCCESS);
+        ASSERT_EQ(to_device(d_indices.get(), indices.data(), indices.size() * sizeof(I)),
+                  FLAGSPARSE_STATUS_SUCCESS);
+        ASSERT_EQ(to_device(d_out.get(), initial.data(), initial.size()), FLAGSPARSE_STATUS_SUCCESS);
+        dev_sync();  // uploads complete before the nonblocking handle stream
+        ASSERT_EQ(flagsparseGather(handle, source, sparse), FLAGSPARSE_STATUS_SUCCESS);
+        ASSERT_EQ(flagsparseScatter(handle, sparse, destination), FLAGSPARSE_STATUS_SUCCESS);
+        dev_sync();
+        auto expected = initial;
+        std::vector<std::uint8_t> gathered(nnz * bytes);
+        for (int j = 0; j < nnz; ++j)
+            for (std::size_t k = 0; k < bytes; ++k) {
+                gathered[j * bytes + k] = dense[indices[j] * bytes + k];
+                expected[indices[j] * bytes + k] = gathered[j * bytes + k];
+            }
+        EXPECT_EQ(d_sparse.download<std::uint8_t>(gathered.size()), gathered);
+        EXPECT_EQ(d_out.download<std::uint8_t>(expected.size()), expected);
+    }
+    flagsparseDestroySpVec(sparse);
+    flagsparseDestroyDnVec(destination);
+    flagsparseDestroyDnVec(source);
+}
+
+TEST_F(ScatterAccuracy, MusaLiveBitCopiesOnDefaultAndNondefaultStreams) {
+    struct Stream {
+        MUstream value = nullptr;
+        ~Stream() { if (value) muStreamDestroy(value); }
+    } stream;
+    ASSERT_EQ(muStreamCreate(&stream.value, MU_STREAM_NON_BLOCKING), MUSA_SUCCESS);
+    for (flagsparseStream_t selected : {static_cast<flagsparseStream_t>(nullptr),
+                                        reinterpret_cast<flagsparseStream_t>(stream.value)}) {
+        ASSERT_EQ(flagsparseSetStream(handle.h, selected), FLAGSPARSE_STATUS_SUCCESS);
+        for (const auto dtype : {FLAGSPARSE_R_8I, FLAGSPARSE_R_16F, FLAGSPARSE_R_32F,
+                                  FLAGSPARSE_R_64F, FLAGSPARSE_C_32F, FLAGSPARSE_C_64F}) {
+            const std::size_t bytes = value_bytes(dtype);
+            musa_live_copies<int32_t>(handle.h, dtype, FLAGSPARSE_INDEX_32I, bytes);
+            musa_live_copies<int64_t>(handle.h, dtype, FLAGSPARSE_INDEX_64I, bytes);
+        }
+    }
+    EXPECT_EQ(flagsparseSetStream(handle.h, nullptr), FLAGSPARSE_STATUS_SUCCESS);
+}
+#endif
+
 }  // namespace

capi/ctest/accuracy/test_spgemm.cpp

@@ -26,6 +26,7 @@
 #include <array>
 #include <vector>
 
+#include "adaptor/adaptor.hpp"
 #include "common.hpp"
 
 using namespace fstest;
@@ -213,6 +214,16 @@ TEST_F(SpGEMMAccuracy, Float32) {
     run_and_check<float>(A, B, FLAGSPARSE_R_32F, "fp32");
 }
 
+TEST_F(SpGEMMAccuracy, MusaMultiBlockRowPrefix) {
+    if (std::string(flagsparse::adaptor::backend_name()) != "musa") GTEST_SKIP();
+    for (int64_t rows : {1024, 1025, 3079}) {
+        const CsrMatrix A = random_csr(rows, 17, 0.14, 71);
+        const CsrMatrix B = random_csr(17, 23, 0.20, 73);
+        run_and_check<float>(A, B, FLAGSPARSE_R_32F, "musa_multiblock_fp32");
+        run_and_check<double>(A, B, FLAGSPARSE_R_64F, "musa_multiblock_fp64");
+    }
+}
+
 // Rows of A with no entries produce empty rows of C; the discovered indptr has
 // to stay monotone through them.
 TEST_F(SpGEMMAccuracy, EmptyRowsAndEmptyResult) {

capi/ctest/accuracy/test_spmv.cpp

@@ -981,6 +981,48 @@ TEST_F(SpMVAccuracy, CsrAndCooFloat32TransposeMatchHostReference) {
     }
 }
 
+#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS)
+TEST_F(SpMVAccuracy, MusaComplexGatherCancellation) {
+    using C = std::complex<float>;
+    const std::vector<C> values{C(0.07468984f, -0.07468984f),
+                                C(-37.027058f, 37.027058f),
+                                C(37.027058f, -37.027058f)};
+    const std::vector<int32_t> ptr{0, 1, 2, 3}, indices{0, 0, 0};
+    const std::vector<C> x{C(1.5f, 0), C(2, 0), C(2, 0)};
+    DeviceBuffer dv = DeviceBuffer::from(values), dp = DeviceBuffer::from(ptr);
+    DeviceBuffer di = DeviceBuffer::from(indices), dx = DeviceBuffer::from(x);
+    DeviceBuffer dy(sizeof(C));
+    flagsparseSpMatDescr_t a = nullptr;
+    flagsparseDnVecDescr_t vx = nullptr, vy = nullptr;
+    ASSERT_EQ(flagsparseCreateCsr(&a, 3, 1, 3, dp.get(), di.get(), dv.get(),
+                                  FLAGSPARSE_INDEX_32I, FLAGSPARSE_INDEX_32I,
+                                  FLAGSPARSE_INDEX_BASE_ZERO, FLAGSPARSE_C_32F),
+              FLAGSPARSE_STATUS_SUCCESS);
+    ASSERT_EQ(flagsparseCreateDnVec(&vx, 3, dx.get(), FLAGSPARSE_C_32F), FLAGSPARSE_STATUS_SUCCESS);
+    ASSERT_EQ(flagsparseCreateDnVec(&vy, 1, dy.get(), FLAGSPARSE_C_32F), FLAGSPARSE_STATUS_SUCCESS);
+    const C one(1, 0), zero(0, 0);
+    size_t bytes = 0;
+    const auto op = FLAGSPARSE_OPERATION_CONJUGATE_TRANSPOSE;
+    ASSERT_EQ(flagsparseSpMV_bufferSize(handle.h, op, &one, a, vx, &zero, vy,
+                                        FLAGSPARSE_C_32F, FLAGSPARSE_SPMV_ALG_DEFAULT, &bytes),
+              FLAGSPARSE_STATUS_SUCCESS);
+    DeviceBuffer scratch(bytes);
+    ASSERT_EQ(flagsparseSpMV_preprocess(handle.h, op, &one, a, vx, &zero, vy,
+                                        FLAGSPARSE_C_32F, FLAGSPARSE_SPMV_ALG_DEFAULT, scratch.get()),
+              FLAGSPARSE_STATUS_SUCCESS);
+    ASSERT_EQ(flagsparseSpMV(handle.h, op, &one, a, vx, &zero, vy, FLAGSPARSE_C_32F,
+                             FLAGSPARSE_SPMV_ALG_DEFAULT, scratch.get()), FLAGSPARSE_STATUS_SUCCESS);
+    dev_sync();
+    const C actual = dy.download<C>(1)[0];
+    const double expected = static_cast<double>(values[0].real()) * 1.5;
+    const Tolerance tol = default_tolerance(FLAGSPARSE_C_32F);
+    EXPECT_LE(max_error_ratio({actual.real(), actual.imag()}, {expected, expected}, tol), 1.0);
+    flagsparseDestroyDnVec(vy);
+    flagsparseDestroyDnVec(vx);
+    flagsparseDestroySpMat(a);
+}
+#endif
+
 TEST_F(SpMVAccuracy, CsrAndCooComplex64ConjugateTransposeMatchHostReference) {
     const CsrMatrix A = random_csr(67, 89, 0.08, 73);
     for (auto format : {FLAGSPARSE_FORMAT_CSR, FLAGSPARSE_FORMAT_COO}) {

capi/ctest/accuracy/test_spvv.cpp

@@ -9,6 +9,10 @@
 
 #include "common.hpp"
 
+#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS)
+#include <musa.h>
+#endif
+
 using namespace fstest;
 
 namespace {
@@ -23,10 +27,12 @@ template <typename In, typename Out>
 Out run_spvv(flagsparseHandle_t handle, flagsparseDataType_t input_type,
              flagsparseDataType_t output_type, const std::vector<In>& sparse,
              const std::vector<int32_t>& indices, const std::vector<In>& dense,
-             flagsparseOperation_t op = FLAGSPARSE_OPERATION_NON_TRANSPOSE) {
+             flagsparseOperation_t op = FLAGSPARSE_OPERATION_NON_TRANSPOSE,
+             bool synchronize_inputs = false) {
     DeviceBuffer d_sparse = DeviceBuffer::from(sparse);
     DeviceBuffer d_indices = DeviceBuffer::from(indices);
     DeviceBuffer d_dense = DeviceBuffer::from(dense);
+    if (synchronize_inputs) dev_sync();
 
     flagsparseSpVecDescr_t vecX = nullptr;
     flagsparseDnVecDescr_t vecY = nullptr;
@@ -95,4 +101,30 @@ TEST_F(SpVVAccuracy, Complex64Conjugate) {
               C(3, -20));
 }
 
+#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS)
+TEST_F(SpVVAccuracy, MusaRepeatedHostResultsOnDefaultAndNondefaultStreams) {
+    struct Stream {
+        MUstream value = nullptr;
+        ~Stream() { if (value) muStreamDestroy(value); }
+    } stream;
+    ASSERT_EQ(muStreamCreate(&stream.value, MU_STREAM_NON_BLOCKING), MUSA_SUCCESS);
+    using C = std::complex<float>;
+    const std::vector<int32_t> indices{0, 2};
+    const std::vector<C> dense{C(3, -2), C(0, 0), C(4, 1)};
+    for (flagsparseStream_t selected : {
+            static_cast<flagsparseStream_t>(nullptr),
+            reinterpret_cast<flagsparseStream_t>(stream.value)}) {
+        ASSERT_EQ(flagsparseSetStream(handle.h, selected), FLAGSPARSE_STATUS_SUCCESS);
+        for (int repeat = 0; repeat < 20; ++repeat) {
+            const std::vector<C> sparse{C(2 + repeat, 1), C(-1, 3)};
+            EXPECT_EQ((run_spvv<C, C>(handle.h, FLAGSPARSE_C_32F, FLAGSPARSE_C_32F,
+                        sparse, indices, dense,
+                        FLAGSPARSE_OPERATION_CONJUGATE_TRANSPOSE, true)),
+                      C(3 + 3 * repeat, -20 - 2 * repeat));
+        }
+    }
+    EXPECT_EQ(flagsparseSetStream(handle.h, nullptr), FLAGSPARSE_STATUS_SUCCESS);
+}
+#endif
+
 }  // namespace

capi/flagsparse_codegen/sparse_gather.py

@@ -117,9 +117,11 @@ def sparse_gather_spmv_f32(V, P, ORDER, COLS, B, C, ar, ai, br, bi, M, N,
         else:
             value = tl.load(V + src, valid, other=0.)
             re += value * tl.load(B + col, valid, other=0.)
-    real = tl.sum(re, 1)
     if COMPLEX:
-        imag = tl.sum(im, 1)
+        # Widen the final tree so small terms survive cancellation between
+        # lanes without widening the sparse loop's per-lane accumulators.
+        real = tl.sum(re.to(tl.float64), 1).to(tl.float32)
+        imag = tl.sum(im.to(tl.float64), 1).to(tl.float32)
         yr = ar * real - ai * imag
         yi = ar * imag + ai * real
         if HAS_BETA:
@@ -130,6 +132,7 @@ def sparse_gather_spmv_f32(V, P, ORDER, COLS, B, C, ar, ai, br, bi, M, N,
         tl.store(C + 2 * rows, yr, active)
         tl.store(C + 2 * rows + 1, yi, active)
     else:
+        real = tl.sum(re, 1)
         result = ar * real
         if HAS_BETA:
             result += br * tl.load(C + rows, active, other=0.)

capi/flagsparse_codegen/spgemm_csr.py

@@ -47,14 +47,58 @@
 _sys.path.insert(0, _os.path.dirname(_os.path.abspath(__file__)))
 import _bootstrap  # noqa: F401,E402
 
+import triton
+import triton.language as tl
+
 from flagsparse.sparse_operations.spgemm_csr import (  # noqa: E402,F401
     _spgemm_hash_count_kernel,
     _spgemm_hash_fill_kernel,
     _spgemm_hash_fill_kernel_f64,
 )
 
+
+@triton.jit
+def musa_spgemm_scan_rows(counts, overflow, indptr, summary, nrows, nblocks,
+                          BLOCK: tl.constexpr):
+    block = tl.program_id(0)
+    rows = block * BLOCK + tl.arange(0, BLOCK)
+    count = tl.load(counts + rows, rows < nrows, 0)
+    failed = tl.load(overflow + rows, rows < nrows, 0)
+    tl.store(indptr + rows + 1, tl.cumsum(count, 0), rows < nrows)
+    if block == 0:
+        tl.store(indptr, 0)
+    tl.store(summary + block, tl.sum(count.to(tl.int64), 0))
+    tl.store(summary + nblocks + block, tl.max(count, 0).to(tl.int64))
+    tl.store(summary + 2 * nblocks + block, tl.max(failed, 0).to(tl.int64))
+
+
+@triton.jit
+def musa_spgemm_scan_blocks(summary, offsets, metadata, nblocks,
+                            BLOCK: tl.constexpr):
+    block = tl.arange(0, BLOCK)
+    count = tl.load(summary + block, block < nblocks, 0)
+    maximum = tl.load(summary + nblocks + block, block < nblocks, 0)
+    failed = tl.load(summary + 2 * nblocks + block, block < nblocks, 0)
+    tl.store(offsets + block, tl.cumsum(count, 0) - count, block < nblocks)
+    tl.store(metadata, tl.sum(count, 0))
+    tl.store(metadata + 1, tl.max(maximum, 0))
+    tl.store(metadata + 2, tl.max(failed, 0))
+
+
+@triton.jit
+def musa_spgemm_add_offsets(indptr, offsets, nrows, BLOCK: tl.constexpr):
+    block = tl.program_id(0)
+    rows = block * BLOCK + tl.arange(0, BLOCK)
+    local = tl.load(indptr + rows + 1, rows < nrows, 0)
+    offset = tl.load(offsets + block)
+    tl.store(indptr + rows + 1, (local.to(tl.int64) + offset).to(tl.int32), rows < nrows)
+
+
 __all__ = [
     "_spgemm_hash_count_kernel",
     "_spgemm_hash_fill_kernel",
     "_spgemm_hash_fill_kernel_f64",
+    "musa_spgemm_scan_rows",
+    "musa_spgemm_scan_blocks",
+    "musa_spgemm_add_offsets",
 ]

capi/flagsparse_codegen/spmm_csr.py

@@ -173,21 +173,31 @@ def spmm_csr_batched8_real(
     data_ptr, indices_ptr, indptr_ptr, b_ptr, c_ptr, alpha, beta,
     n_rows, n_dense_cols, stride_bk, stride_bn, stride_cm, stride_cn,
     BATCH_ROWS: tl.constexpr, ACC_IS_FP64: tl.constexpr, HAS_BETA: tl.constexpr,
+    WIDE: tl.constexpr = True,
+    K: tl.constexpr = 4,
 ):
-    rows = tl.program_id(0) * BATCH_ROWS + tl.arange(0, BATCH_ROWS)
+    pid = tl.program_id(0).to(tl.int64) if WIDE else tl.program_id(0)
+    rows = pid * BATCH_ROWS + tl.arange(0, BATCH_ROWS)
     active = rows < n_rows
-    start = tl.load(indptr_ptr + rows, mask=active, other=0).to(tl.int64)
-    end = tl.load(indptr_ptr + rows + 1, mask=active, other=0).to(tl.int64)
+    start = tl.load(indptr_ptr + rows, mask=active, other=0)
+    end = tl.load(indptr_ptr + rows + 1, mask=active, other=0)
+    if WIDE:
+        start, end = start.to(tl.int64), end.to(tl.int64)
+    else:
+        start, end = start.to(tl.int32), end.to(tl.int32)
+        stride_bk, stride_bn = tl.cast(stride_bk, tl.int32), tl.cast(stride_bn, tl.int32)
+        stride_cm, stride_cn = tl.cast(stride_cm, tl.int32), tl.cast(stride_cn, tl.int32)
     max_row_nnz = tl.max(end - start)
     offs_n = tl.arange(0, 8)
     acc_dtype = tl.float64 if ACC_IS_FP64 else tl.float32
     acc = tl.zeros((BATCH_ROWS, 8), dtype=acc_dtype)
-    for chunk_start in tl.range(0, max_row_nnz, 4):
-        offs_k = chunk_start + tl.arange(0, 4)
+    for chunk_start in tl.range(0, max_row_nnz, K):
+        offs_k = chunk_start + tl.arange(0, K)
         positions = start[:, None] + offs_k[None, :]
         valid = active[:, None] & (positions < end[:, None])
         values = tl.load(data_ptr + positions, mask=valid, other=0.0).to(acc_dtype)
-        cols = tl.load(indices_ptr + positions, mask=valid, other=0).to(tl.int64)
+        cols = tl.load(indices_ptr + positions, mask=valid, other=0)
+        cols = cols.to(tl.int64) if WIDE else cols.to(tl.int32)
         dense = tl.load(
             b_ptr + cols[:, :, None] * stride_bk + offs_n[None, None, :] * stride_bn,
             mask=valid[:, :, None], other=0.0,

capi/flagsparse_codegen/spmv_csr.py

@@ -53,11 +53,17 @@ def spmv_csr_row_tile_real(
     data_ptr, indices_ptr, indptr_ptr, x_ptr, y_ptr, alpha, beta, n_rows,
     ROWS_PER_PROGRAM: tl.constexpr, LANES_PER_ROW: tl.constexpr,
     ACC_IS_FP64: tl.constexpr, HAS_BETA: tl.constexpr,
+    WIDE: tl.constexpr = True,
 ):
-    rows = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)
+    pid = tl.program_id(0).to(tl.int64) if WIDE else tl.program_id(0)
+    rows = pid * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)
     active = rows < n_rows
-    start = tl.load(indptr_ptr + rows, mask=active, other=0).to(tl.int64)
-    end = tl.load(indptr_ptr + rows + 1, mask=active, other=0).to(tl.int64)
+    start = tl.load(indptr_ptr + rows, mask=active, other=0)
+    end = tl.load(indptr_ptr + rows + 1, mask=active, other=0)
+    if WIDE:
+        start, end = start.to(tl.int64), end.to(tl.int64)
+    else:
+        start, end = start.to(tl.int32), end.to(tl.int32)
     steps = tl.max(tl.cdiv(end - start, LANES_PER_ROW), 0)
     lane = tl.arange(0, LANES_PER_ROW)
     acc_dtype = tl.float64 if ACC_IS_FP64 else tl.float32
@@ -66,7 +72,8 @@ def spmv_csr_row_tile_real(
         positions = start[:, None] + step * LANES_PER_ROW + lane[None, :]
         mask = active[:, None] & (positions < end[:, None])
         values = tl.load(data_ptr + positions, mask=mask, other=0.0).to(acc_dtype)
-        cols = tl.load(indices_ptr + positions, mask=mask, other=0).to(tl.int64)
+        cols = tl.load(indices_ptr + positions, mask=mask, other=0)
+        cols = cols.to(tl.int64) if WIDE else cols.to(tl.int32)
         acc += values * tl.load(x_ptr + cols, mask=mask, other=0.0).to(acc_dtype)
     out = alpha * tl.sum(acc, axis=1)
     if HAS_BETA: