From 37e792fd510202fd06b120adbd956388aaa6efe7 Mon Sep 17 00:00:00 2001 From: NCIA-AlphaSparse Date: Thu, 15 Oct 2026 16:33:01 +0800 Subject: [PATCH 1/4] musa debug docs --- docs/debug/MUSA.md | 39 +++++++++++++++++++++++++++++++++++---- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/docs/debug/MUSA.md b/docs/debug/MUSA.md index f4d4248..7c9b5e0 100644 --- a/docs/debug/MUSA.md +++ b/docs/debug/MUSA.md @@ -8,15 +8,46 @@ 环境按 `../MUSA.md` 第 2 节配好（`FLAGSPARSE_BACKEND=mthreads`、`MUSA_HOME`）后，跑一轮 split： ```bash +# 1. 更新代码：至少要包含 a330b62（见下面"2026-10-14 那轮的两个 NoBaseline"） +git pull --ff-only +git log --oneline -1 +git merge-base --is-ancestor a330b62 HEAD && echo "包含 a330b62 修复" + +# 2. 环境自检 +export PYTHONPATH=$PWD/src FLAGSPARSE_BACKEND=mthreads MUSA_HOME=/usr/local/musa # MUSA_HOME 以本机实际安装位置为准 +python3 -c "import flagsparse; print(flagsparse.__file__)" # 必须是 <仓库>/src/flagsparse/__init__.py +python3 -c "from flagsparse.sparse_operations import _common as c; print(c._backend_name(), c._accel_device_type(), c._accel_fallback_reason())" +# 期望：mthreads musa None +ls tests/data/*.mtx | wc -l # 期望 10 +mthreads-gmi # 确认没有别的任务在用 GPU + +# 3. 全量 split（每次用新的结果目录） +D=results_musa_split_65_<日期> setsid timeout -s KILL 43200 python3 -u run_flagsparse_split_delivery.py \ --mode normal --benchmark-input tests/data --timeout 3600 \ --h800-reference \ - --results-dir results_musa_split_65_<日期> \ - > results_musa_split_65_<日期>.log 2>&1 < /dev/null & -# 跑完后 -python3 tools/delivery_table.py results_musa_split_65_<日期> + --results-dir $D \ + > $D.log 2>&1 < /dev/null & + +# 4. 跑完后 +tail -5 $D.log # H800 折算三类各补了多少、还剩哪些 +python3 tools/delivery_table.py $D # 期望 0 missing，性能里没有 NoBaseline +python3 tools/delivery_table.py $D --show-baseline | grep -E "scatter_i8_int|spmm_csc_f16" +# 期望这两行有加速比，来源 h800_scaled ``` +- `flagsparse.__file__` 那一步不能省：系统 site-packages 里残留的旧安装会顶替 `src/`，基线悄悄变成 `N/A`。 +- `MUSA_HOME` 设错不会报错，muSPARSE 基线同样会悄悄退成 `none`（`../MUSA.md` 第 2 节）。 + +**2026-10-14 那轮的两个 NoBaseline**（`results_musa_split_65_20261014`，63 Passed + 2 NoBaseline）： +`scatter_i8_int` 和 `spmm_csc_f16_int_non_non_row` muSPARSE 都不支持（C API 报 `Scatter dtype unsupported`、 +`unsupported SpMM type or operation`），只能靠 Python 侧的折算补，而那轮 Python 侧这两处都崩了：scatter int8 的 +torch 参考被 muDNN 拒绝（`Now muDNN only support float64, float32, half, bfloat16, int32, and int64`，4 行全是 +`ERROR`）；spmm_csc 的 `torch.sparse.mm` 参考在 MUSA 上没有实现（`Could not run 'aten::addmm' with arguments from +the 'Sparsemusa' backend`，CSV 为空）。`a330b62`（参考改在 CPU / SciPy 上算）修了这两处，但它比那轮结果晚提交， +所以那份报告里的这两行是旧代码的结果，按上面的命令重跑即可。CUDA 上以 `FLAGSPARSE_BACKEND=mthreads` 模拟已确认 +两行都能测出性能；折算这一步只在真机上走得到。 + - 精度那一半调用 `run_flagsparse_pytest.py --phase accuracy --delivery-only`，覆盖全部 65 个变体。 - 性能那一半取 C API 的 `ctest`（对 muSPARSE），**现在也覆盖全部 65 个变体**：`origin/q4` 的 C API 工作 2026-10-08 已移植进 main（见下面"C API 现在覆盖全部 65 个变体"一节），每个变体按自己的 id 单独测一次。 From 9e92e76b16d19ca9bf46b2a169ca82c67f12351c Mon Sep 17 00:00:00 2001 From: NCIC-AlphaSparse Date: Sat, 17 Oct 2026 22:34:03 +0800 Subject: [PATCH 2/4] musa opt 1 --- capi/flagsparse_codegen/sparse_gather.py | 7 +- capi/flagsparse_codegen/spgemm_csr.py | 44 + capi/flagsparse_codegen/spmm_csr.py | 22 +- capi/flagsparse_codegen/spmv_csr.py | 15 +- capi/flagsparse_codegen/spmv_sell.py | 36 + capi/src/core/sparse_gather.hpp | 6 +- capi/src/ops/gather_scatter.cpp | 13 +- capi/src/ops/spgemm.cpp | 192 ++++- capi/src/ops/spmm.cpp | 15 +- capi/src/ops/spmv.cpp | 27 +- capi/src/ops/spvv.cpp | 12 +- docs/debug/MUSA.md | 376 ++++++++- .../sparse_operations/_musa_spmv_launch.py | 45 + .../sparse_operations/_spmm_row_gather.py | 17 +- .../sparse_operations/_spmv_csr_nnz.py | 43 +- .../sparse_operations/_spmv_csr_transpose.py | 29 +- .../sparse_operations/gather_scatter.py | 25 + .../sparse_operations/mixed_spmx.py | 155 +++- src/flagsparse/sparse_operations/spmm_coo.py | 11 +- src/flagsparse/sparse_operations/spmm_csc.py | 67 +- src/flagsparse/sparse_operations/spmm_csr.py | 7 +- src/flagsparse/sparse_operations/spmv_coo.py | 13 +- src/flagsparse/sparse_operations/spmv_csc.py | 2 +- src/flagsparse/sparse_operations/spmv_csr.py | 1 + src/flagsparse/sparse_operations/spmv_sell.py | 169 +++- .../sparse_operations/vector_ops.py | 52 +- tests/pytest/test_musa_sparse_routes.py | 797 ++++++++++++++++++ tests/pytest/test_spmv_csr_accuracy.py | 6 +- 28 files changed, 2081 insertions(+), 123 deletions(-) create mode 100644 src/flagsparse/sparse_operations/_musa_spmv_launch.py create mode 100644 tests/pytest/test_musa_sparse_routes.py diff --git a/capi/flagsparse_codegen/sparse_gather.py b/capi/flagsparse_codegen/sparse_gather.py index ea7dde2..05c64fa 100644 --- a/capi/flagsparse_codegen/sparse_gather.py +++ b/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.) diff --git a/capi/flagsparse_codegen/spgemm_csr.py b/capi/flagsparse_codegen/spgemm_csr.py index 2de77a6..446459f 100644 --- a/capi/flagsparse_codegen/spgemm_csr.py +++ b/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", ] diff --git a/capi/flagsparse_codegen/spmm_csr.py b/capi/flagsparse_codegen/spmm_csr.py index 4176c07..4a84ece 100644 --- a/capi/flagsparse_codegen/spmm_csr.py +++ b/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, diff --git a/capi/flagsparse_codegen/spmv_csr.py b/capi/flagsparse_codegen/spmv_csr.py index f975a65..76756c8 100644 --- a/capi/flagsparse_codegen/spmv_csr.py +++ b/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: diff --git a/capi/flagsparse_codegen/spmv_sell.py b/capi/flagsparse_codegen/spmv_sell.py index 6f2924a..7840e9d 100644 --- a/capi/flagsparse_codegen/spmv_sell.py +++ b/capi/flagsparse_codegen/spmv_sell.py @@ -14,6 +14,9 @@ from flagsparse.sparse_operations.spmv_sell import _spmv_sell_real_kernel # noqa: E402 from flagsparse.sparse_operations.spmv_sell import _spmv_sell_complex_kernel # noqa: E402 +from flagsparse.sparse_operations.spmv_sell import ( # noqa: E402 + _musa_spmv_sell_kernel, +) @triton.jit @@ -36,3 +39,36 @@ def spmv_sell_i32acc( y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, SLICE=SLICE, BLOCK_R=BLOCK_R, ACC=tl.int32, ) + + +@triton.jit +def spmv_sell_musa_f32acc( + y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, + SLICE: tl.constexpr, ROWS: tl.constexpr, +): + _musa_spmv_sell_kernel( + y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, + SLICE=SLICE, ROWS=ROWS, SLOTS=4, COMPLEX=False, ACC=tl.float32, + ) + + +@triton.jit +def spmv_sell_musa_i32acc( + y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, + SLICE: tl.constexpr, ROWS: tl.constexpr, +): + _musa_spmv_sell_kernel( + y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, + SLICE=SLICE, ROWS=ROWS, SLOTS=4, COMPLEX=False, ACC=tl.int32, + ) + + +@triton.jit +def spmv_sell_musa_c32( + y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, + SLICE: tl.constexpr, ROWS: tl.constexpr, SLOTS: tl.constexpr, +): + _musa_spmv_sell_kernel( + y_ptr, values_ptr, cols_ptr, offsets_ptr, x_ptr, n_rows, + SLICE=SLICE, ROWS=ROWS, SLOTS=SLOTS, COMPLEX=True, ACC=tl.float32, + ) diff --git a/capi/src/core/sparse_gather.hpp b/capi/src/core/sparse_gather.hpp index 21dbfeb..8a7fcb7 100644 --- a/capi/src/core/sparse_gather.hpp +++ b/capi/src/core/sparse_gather.hpp @@ -119,9 +119,11 @@ inline flagsparseStatus_t launch_sparse_gather(flagsparseHandle_t handle, SpMatD if (auto s = prepare_sparse_gather(A, buffer)) return s; } const bool complex = A->value_type == FLAGSPARSE_C_32F; - if (n == 1 && indirect && !complex) - return launch_sparse_scatter(handle, A, B, C, m, ar, br, buffer); const double average = static_cast(A->nnz) / std::max(m, 1); + // Longer rows amortize indirect gather loads and avoid order-dependent + // scatter accumulation in cancellation-sensitive fp32 outputs. + if (n == 1 && indirect && !complex && average < 12.) + return launch_sparse_scatter(handle, A, B, C, m, ar, br, buffer); const bool stream = n == 8; const int r = n == 1 ? (average <= 4. ? 128 : 16) : (stream ? (average <= 8. ? 8 : 16) : 4); const int k = n == 1 ? (average <= 4. ? 4 : 8) : (stream ? 1 : 16); diff --git a/capi/src/ops/gather_scatter.cpp b/capi/src/ops/gather_scatter.cpp index beeb621..2dd6ddf 100644 --- a/capi/src/ops/gather_scatter.cpp +++ b/capi/src/ops/gather_scatter.cpp @@ -23,6 +23,7 @@ // Triton has and the one the Python package already uses. #include +#include #include #include @@ -45,6 +46,11 @@ constexpr int kNumStages = 2; // Deriving it from the device keeps it right across backends, whose SM counts // and warp sizes differ. int64_t gather_max_programs(const Context* c) { + const bool musa = std::strcmp(adaptor::backend_name(), "musa") == 0; + static thread_local int cached_device = -1; + static thread_local int64_t cached_programs = 0; + if (musa && cached_device == c->device_index && cached_programs > 0) + return cached_programs; const int sm = adaptor::multiprocessor_count(c->device_index); const int threads_per_sm = adaptor::max_threads_per_multiprocessor(c->device_index); const int warp = adaptor::warp_size(c->device_index); @@ -52,7 +58,12 @@ int64_t gather_max_programs(const Context* c) { const int block_threads = std::max(1, kNumWarps * (warp > 0 ? warp : 32)); const int blocks_per_sm = threads_per_sm > 0 ? std::max(1, threads_per_sm / block_threads) : 1; - return std::max(1, static_cast(sm) * blocks_per_sm); + const int64_t programs = std::max(1, static_cast(sm) * blocks_per_sm); + if (musa) { + cached_device = c->device_index; + cached_programs = programs; + } + return programs; } struct Plan { diff --git a/capi/src/ops/spgemm.cpp b/capi/src/ops/spgemm.cpp index ab6cbc0..c093dd4 100644 --- a/capi/src/ops/spgemm.cpp +++ b/capi/src/ops/spgemm.cpp @@ -31,6 +31,7 @@ // reports NOT_SUPPORTED instead of quietly returning a short row. #include +#include #include #include #include @@ -48,6 +49,12 @@ using namespace flagsparse; namespace { constexpr int kNumStages = 1; +constexpr int kMusaScanBlock = 1024; + +bool musa_backend() { + static const bool musa = std::string(adaptor::backend_name()) == "musa"; + return musa; +} // (max products per row, CAP, BLOCK, num_warps). The threshold is 0.75 * CAP, // the hash table's load limit; a row past the last bucket cannot be served. @@ -79,6 +86,9 @@ struct SpGEMMDescr { int64_t max_row_work = 0; int64_t max_row_nnz = 0; std::vector host_indptr; // the CSR indptr compute derived + std::vector group_counts; + std::vector group_offsets; + std::vector host_counts; }; SpGEMMDescr* spgemm(flagsparseSpGEMMDescr_t d) { @@ -89,11 +99,15 @@ size_t align_up(size_t v) { return (v + 255u) & ~static_cast(255u); } size_t buffer1_bytes(int64_t rows, int64_t nnz_a) { return align_up(static_cast(rows) * sizeof(std::int32_t)) + - align_up(static_cast(nnz_a) * sizeof(std::int32_t)); + align_up(static_cast(nnz_a) * sizeof(std::int32_t)) + + (std::string(adaptor::backend_name()) == "musa" + ? align_up((static_cast(rows) + 20) * sizeof(std::int32_t)) : 0); } size_t buffer2_bytes(int64_t rows) { - return align_up(static_cast(rows) * sizeof(std::int32_t)) * 2 + - align_up(static_cast(rows + 1) * sizeof(std::int32_t)); + const size_t base = align_up(static_cast(rows) * sizeof(std::int32_t)) * 2 + + align_up(static_cast(rows + 1) * sizeof(std::int32_t)); + const size_t blocks = (static_cast(rows) + kMusaScanBlock - 1) / kMusaScanBlock; + return base + (musa_backend() ? align_up((4 * blocks + 3) * sizeof(int64_t)) : 0); } // Read an index array back to the host as int64, whatever width it is stored in. @@ -245,7 +259,8 @@ flagsparseStatus_t launch_hash(flagsparseHandle_t handle, const char* kernel, void* c_indptr, void* c_ind, void* c_data, const char* c_offsets_type, const Bucket& bucket, bool report, bool fill, - flagsparseDataType_t computeType) { + flagsparseDataType_t computeType, + void* row_ids = nullptr, int64_t row_count = 0) { const char* vt = triton_dtype(computeType); const char* ait = triton_index_dtype(A->indices_type); const char* aot = triton_index_dtype(A->offsets_type); @@ -273,7 +288,7 @@ flagsparseStatus_t launch_hash(flagsparseHandle_t handle, const char* kernel, sig += "*i32:16,"; // ovf sig += std::to_string(bucket.cap) + ","; sig += std::to_string(bucket.block) + ","; - sig += "False"; // USE_ROWS + sig += row_ids ? "True" : "False"; // USE_ROWS if (fill) { sig += report ? ",True" : ",False"; } // REPORT std::vector args; @@ -288,8 +303,9 @@ flagsparseStatus_t launch_hash(flagsparseHandle_t handle, const char* kernel, if (fill) args.push_back(ptr(B->values)); args.push_back(ptr(B->indices)); args.push_back(ptr(B->offsets)); - args.push_back(ptr(rw)); // rows_ptr placeholder - args.push_back(jit::Arg::i(static_cast(A->rows))); + args.push_back(ptr(row_ids ? row_ids : rw)); + const int64_t launch_rows = row_ids ? row_count : A->rows; + args.push_back(jit::Arg::i(static_cast(launch_rows))); args.push_back(ptr(rw)); if (fill) { args.push_back(ptr(c_indptr)); args.push_back(ptr(c_ind)); args.push_back(ptr(c_data)); } @@ -299,11 +315,55 @@ flagsparseStatus_t launch_hash(flagsparseHandle_t handle, const char* kernel, std::string err; const flagsparseStatus_t st = jit::launch( jit::codegen_module("spgemm_csr.py"), kernel, sig, ctx(handle)->stream, - A->rows, 1, 1, bucket.warps, kNumStages, args, &err); + launch_rows, 1, 1, bucket.warps, kNumStages, args, &err); if (st != FLAGSPARSE_STATUS_SUCCESS) ctx(handle)->last_error = err; return st; } +flagsparseStatus_t musa_build_indptr(flagsparseHandle_t handle, void* counts, void* overflow, + void* indptr, void* scratch, int64_t rows, + int64_t* total, int64_t* maximum) { + const int blocks = static_cast((rows + kMusaScanBlock - 1) / kMusaScanBlock); + auto* summary = static_cast(scratch); + auto* offsets = summary + 3 * blocks; + auto* metadata = summary + 4 * blocks; + const auto ptr = [](void* p) { + return jit::Arg::ptr(reinterpret_cast(p)); + }; + const auto module = jit::codegen_module("spgemm_csr.py"); + std::string err; + auto st = jit::launch( + module, "musa_spgemm_scan_rows", "*i32:16,*i32:16,*i32:16,*i64:16,i32,i32,1024", + ctx(handle)->stream, blocks, 1, 1, 4, 1, + {ptr(counts), ptr(overflow), ptr(indptr), ptr(summary), + jit::Arg::i(static_cast(rows)), jit::Arg::i(blocks)}, &err); + if (st != FLAGSPARSE_STATUS_SUCCESS) { ctx(handle)->last_error = err; return st; } + int scan_width = 1; + while (scan_width < blocks) scan_width *= 2; + st = jit::launch( + module, "musa_spgemm_scan_blocks", "*i64:16,*i64,*i64,i32," + std::to_string(scan_width), + ctx(handle)->stream, 1, 1, 1, 4, 1, + {ptr(summary), ptr(offsets), ptr(metadata), jit::Arg::i(blocks)}, &err); + if (st != FLAGSPARSE_STATUS_SUCCESS) { ctx(handle)->last_error = err; return st; } + adaptor::synchronize(); + std::array host{}; + if (auto s = adaptor::memcpy_d2h(host.data(), + reinterpret_cast(metadata), sizeof(host))) return s; + if (host[2] != 0) { + ctx(handle)->last_error = "SpGEMM row overflowed the shared-memory hash table"; + return FLAGSPARSE_STATUS_NOT_SUPPORTED; + } + if (host[0] > INT32_MAX) return FLAGSPARSE_STATUS_NOT_SUPPORTED; + st = jit::launch( + module, "musa_spgemm_add_offsets", "*i32:16,*i64,i32,1024", + ctx(handle)->stream, blocks, 1, 1, 4, 1, + {ptr(indptr), ptr(offsets), jit::Arg::i(static_cast(rows))}, &err); + if (st != FLAGSPARSE_STATUS_SUCCESS) { ctx(handle)->last_error = err; return st; } + *total = host[0]; + *maximum = host[1]; + return FLAGSPARSE_STATUS_SUCCESS; +} + } // namespace extern "C" { @@ -391,6 +451,31 @@ flagsparseStatus_t flagsparseSpGEMM_workEstimation( a_pref.size() * sizeof(std::int32_t))) return s; } + d->group_counts.clear(); + d->group_offsets.clear(); + if (std::string(adaptor::backend_name()) == "musa" && max_work <= 6144) { + // Match the existing hash capacity to each row's product work, + // rather than allocating the worst row's table for every row. + std::vector row_ids; + for (const auto& group : kBuckets) { + d->group_offsets.push_back(row_ids.size()); + const int64_t lower = &group == &kBuckets[0] + ? -1 : (&group - 1)->max_work; + for (int64_t r = 0; r < A->rows; ++r) { + if (rw[r] > lower && rw[r] <= group.max_work) + row_ids.push_back(static_cast(r)); + } + d->group_counts.push_back(static_cast( + row_ids.size() - d->group_offsets.back())); + while (row_ids.size() % 4) row_ids.push_back(0); + } + void* rows_dev = base + align_up(static_cast(A->rows) * sizeof(std::int32_t)) + + align_up(static_cast(A->nnz) * sizeof(std::int32_t)); + if (flagsparseStatus_t s = adaptor::memcpy_h2d( + reinterpret_cast(rows_dev), row_ids.data(), + row_ids.size() * sizeof(std::int32_t))) return s; + } + d->buffer1 = externalBuffer1; d->matrix_a = static_cast(A); d->matrix_b = static_cast(B); @@ -495,42 +580,60 @@ flagsparseStatus_t flagsparseSpGEMM_compute( void* rw = b1; void* a_pref = b1 + align_up(static_cast(A->rows) * sizeof(std::int32_t)); - if (flagsparseStatus_t s = launch_hash( - handle, "_spgemm_hash_count_kernel", A, B, rw, a_pref, row_nnz, ovf, - nullptr, nullptr, nullptr, "i32", *bucket, false, false, computeType)) { - return s; - } - adaptor::synchronize(); - - std::vector host_nnz(static_cast(A->rows)); - std::vector host_ovf(static_cast(A->rows)); - if (flagsparseStatus_t s = adaptor::memcpy_d2h( - host_nnz.data(), reinterpret_cast(row_nnz), - host_nnz.size() * sizeof(std::int32_t))) return s; - if (flagsparseStatus_t s = adaptor::memcpy_d2h( - host_ovf.data(), reinterpret_cast(ovf), - host_ovf.size() * sizeof(std::int32_t))) return s; - for (std::int32_t f : host_ovf) { - if (f != 0) { - ctx(handle)->last_error = - "SpGEMM row overflowed the shared-memory hash table after " - "product-count sizing; expand-sort-compress fallback is not " - "implemented in the C API"; - return FLAGSPARSE_STATUS_NOT_SUPPORTED; + if (std::string(adaptor::backend_name()) == "musa" && !d->group_counts.empty()) { + auto* row_ids = b1 + align_up(static_cast(A->rows) * sizeof(std::int32_t)) + + align_up(static_cast(A->nnz) * sizeof(std::int32_t)); + for (size_t i = 0; i < d->group_counts.size(); ++i) { + if (!d->group_counts[i]) continue; + if (flagsparseStatus_t s = launch_hash( + handle, "_spgemm_hash_count_kernel", A, B, rw, a_pref, row_nnz, ovf, + nullptr, nullptr, nullptr, "i32", kBuckets[i], false, false, computeType, + row_ids + d->group_offsets[i] * sizeof(std::int32_t), d->group_counts[i])) + return s; } + } else { + if (flagsparseStatus_t s = launch_hash( + handle, "_spgemm_hash_count_kernel", A, B, rw, a_pref, row_nnz, ovf, + nullptr, nullptr, nullptr, "i32", *bucket, false, false, computeType)) return s; } - - d->host_indptr.assign(static_cast(A->rows) + 1, 0); int64_t total = 0, max_row = 0; - for (int64_t r = 0; r < A->rows; ++r) { - total += host_nnz[static_cast(r)]; - max_row = std::max(max_row, host_nnz[static_cast(r)]); - if (total > static_cast(INT32_MAX)) return FLAGSPARSE_STATUS_NOT_SUPPORTED; - d->host_indptr[static_cast(r) + 1] = static_cast(total); + if (musa_backend()) { + void* scratch = base + 2 * rows_bytes + + align_up(static_cast(A->rows + 1) * sizeof(int32_t)); + if (auto s = musa_build_indptr(handle, row_nnz, ovf, c_indptr, scratch, + A->rows, &total, &max_row)) return s; + } else { + adaptor::synchronize(); + std::vector host_nnz(static_cast(A->rows)); + std::vector host_ovf(static_cast(A->rows)); + if (flagsparseStatus_t s = adaptor::memcpy_d2h( + host_nnz.data(), reinterpret_cast(row_nnz), + host_nnz.size() * sizeof(std::int32_t))) return s; + if (flagsparseStatus_t s = adaptor::memcpy_d2h( + host_ovf.data(), reinterpret_cast(ovf), + host_ovf.size() * sizeof(std::int32_t))) return s; + for (std::int32_t f : host_ovf) { + if (f != 0) { + ctx(handle)->last_error = + "SpGEMM row overflowed the shared-memory hash table after " + "product-count sizing; expand-sort-compress fallback is not " + "implemented in the C API"; + return FLAGSPARSE_STATUS_NOT_SUPPORTED; + } + } + d->host_indptr.assign(static_cast(A->rows) + 1, 0); + for (int64_t r = 0; r < A->rows; ++r) { + total += host_nnz[static_cast(r)]; + max_row = std::max(max_row, host_nnz[static_cast(r)]); + if (total > static_cast(INT32_MAX)) return FLAGSPARSE_STATUS_NOT_SUPPORTED; + d->host_indptr[static_cast(r) + 1] = static_cast(total); + } + } + if (!musa_backend()) { + if (flagsparseStatus_t s = adaptor::memcpy_h2d( + reinterpret_cast(c_indptr), d->host_indptr.data(), + d->host_indptr.size() * sizeof(std::int32_t))) return s; } - if (flagsparseStatus_t s = adaptor::memcpy_h2d( - reinterpret_cast(c_indptr), d->host_indptr.data(), - d->host_indptr.size() * sizeof(std::int32_t))) return s; // This is the whole point of the two-call dance: the caller now learns // C's nnz through flagsparseSpMatGetSize and allocates for it. @@ -568,6 +671,15 @@ flagsparseStatus_t flagsparseSpGEMM_copy( "allocate, then attach them with flagsparseCsrSetPointers."; return FLAGSPARSE_STATUS_INVALID_VALUE; } + if (musa_backend() && !d->host_fallback && A->rows > 0) { + const size_t rows_bytes = align_up(static_cast(A->rows) * sizeof(int32_t)); + auto* indptr = static_cast(d->buffer2) + 2 * rows_bytes; + d->host_indptr.resize(static_cast(A->rows) + 1); + adaptor::synchronize(); + if (auto s = adaptor::memcpy_d2h(d->host_indptr.data(), + reinterpret_cast(indptr), + d->host_indptr.size() * sizeof(int32_t))) return s; + } // C's own offsets array gets the indptr compute derived. if (C->offsets_type == FLAGSPARSE_INDEX_32I) { if (flagsparseStatus_t s = adaptor::memcpy_h2d( diff --git a/capi/src/ops/spmm.cpp b/capi/src/ops/spmm.cpp index 8e3107b..99a51e2 100644 --- a/capi/src/ops/spmm.cpp +++ b/capi/src/ops/spmm.cpp @@ -665,7 +665,7 @@ flagsparseStatus_t run_csr(flagsparseHandle_t handle, SpMatDescr* A, const DnMat // and its 1-warp launch were measured there and nowhere else. const bool batched8 = musa_backend() && !ops.complex_op && ops.n == 8 && max_row_nnz <= 256; - const int batch_rows = 4; + const int batch_rows = ops.acc_is_fp64 ? 16 : 32; const char* it = triton_index_dtype(A->indices_type); const char* ot = triton_index_dtype(A->offsets_type); @@ -681,7 +681,16 @@ flagsparseStatus_t run_csr(flagsparseHandle_t handle, SpMatDescr* A, const DnMat if (batched8) { sig += std::to_string(batch_rows) + ","; sig += ops.acc_is_fp64 ? "True," : "False,"; - sig += ops.has_beta ? "True" : "False"; + sig += ops.has_beta ? "True," : "False,"; + // Each operand's row/column contribution is bounded separately, so + // their sum fits int32 as well. Large strided views keep int64 math. + const bool narrow = A->nnz <= INT32_MAX - 128 && ops.m <= INT32_MAX - 128 && + ops.stride_bk <= INT32_MAX / (2 * std::max(1, A->cols)) && + ops.stride_bn <= INT32_MAX / (2 * std::max(1, ops.n)) && + ops.stride_cm <= INT32_MAX / (2 * std::max(1, ops.m)) && + ops.stride_cn <= INT32_MAX / (2 * std::max(1, ops.n)); + sig += narrow ? "False" : "True"; + sig += ops.acc_is_fp64 ? ",1" : ",4"; } else { sig += std::to_string(cfg.block_n) + ","; sig += std::to_string(cfg.block_nnz) + ","; @@ -710,7 +719,7 @@ flagsparseStatus_t run_csr(flagsparseHandle_t handle, SpMatDescr* A, const DnMat ops.complex_op ? "spmm_csr_complex" : (batched8 ? "spmm_csr_batched8_real" : "spmm_csr_real"), sig, ctx(handle)->stream, batched8 ? (ops.m + batch_rows - 1) / batch_rows : ops.m, - grid_n, 1, batched8 ? 1 : cfg.num_warps, batched8 ? 1 : cfg.num_stages, + grid_n, 1, batched8 ? 4 : cfg.num_warps, batched8 ? 1 : cfg.num_stages, args, &err); if (st != FLAGSPARSE_STATUS_SUCCESS) ctx(handle)->last_error = err; return st; diff --git a/capi/src/ops/spmv.cpp b/capi/src/ops/spmv.cpp index 1218702..197af29 100644 --- a/capi/src/ops/spmv.cpp +++ b/capi/src/ops/spmv.cpp @@ -328,7 +328,10 @@ flagsparseStatus_t launch_csr_row_tile_short(flagsparseHandle_t handle, SpMatDes sig += std::to_string(rows_per_program) + ","; sig += std::to_string(lanes_per_row) + ","; sig += ops.acc_is_fp64 ? "True," : "False,"; - sig += ops.has_beta ? "True" : "False"; + sig += ops.has_beta ? "True," : "False,"; + const bool wide = A->rows > INT32_MAX - 128 || A->nnz > INT32_MAX - 128 || + A->cols > INT32_MAX - 128; + sig += wide ? "True" : "False"; std::vector args; args.reserve(8); args.push_back(jit::Arg::ptr(reinterpret_cast(A->values))); @@ -596,7 +599,8 @@ flagsparseStatus_t run_mixed_sell(flagsparseHandle_t handle, const SpMatDescr* A if (A->rows == 0) return FLAGSPARSE_STATUS_SUCCESS; const char* it = triton_index_dtype(A->indices_type); const char* ot = triton_index_dtype(A->offsets_type); - const int block_r = next_power_of_two(A->slice_size); + const bool parallel = std::string(adaptor::backend_name()) == "musa"; + const int block_r = parallel ? 32 : next_power_of_two(A->slice_size); const int64_t n_slices = (A->rows + A->slice_size - 1) / A->slice_size; std::string sig; @@ -617,9 +621,10 @@ flagsparseStatus_t run_mixed_sell(flagsparseHandle_t handle, const SpMatDescr* A std::string err; const flagsparseStatus_t st = jit::launch( - jit::codegen_module("spmv_sell.py"), "spmv_sell_i32acc", sig, - ctx(handle)->stream, n_slices, 1, 1, /*num_warps=*/1, - /*num_stages=*/2, args, &err); + jit::codegen_module("spmv_sell.py"), + parallel ? "spmv_sell_musa_i32acc" : "spmv_sell_i32acc", sig, + ctx(handle)->stream, parallel ? (A->rows + 31) / 32 : n_slices, 1, 1, parallel ? 4 : 1, + parallel ? 1 : 2, args, &err); if (st != FLAGSPARSE_STATUS_SUCCESS) ctx(handle)->last_error = err; return st; } @@ -630,7 +635,8 @@ flagsparseStatus_t run_sell(flagsparseHandle_t handle, const SpMatDescr* A, if (A->rows == 0) return FLAGSPARSE_STATUS_SUCCESS; const char* it = triton_index_dtype(A->indices_type); const char* ot = triton_index_dtype(A->offsets_type); - const int block_r = next_power_of_two(A->slice_size); + const bool parallel = std::string(adaptor::backend_name()) == "musa"; + const int block_r = parallel ? 32 : next_power_of_two(A->slice_size); const int64_t n_slices = (A->rows + A->slice_size - 1) / A->slice_size; const bool complex = ops.complex_op; @@ -642,7 +648,7 @@ flagsparseStatus_t run_sell(flagsparseHandle_t handle, const SpMatDescr* A, sig += "*"; sig += ot; sig += ":16,"; sig += "*"; sig += ops.vt; sig += ":16,i32,"; sig += std::to_string(A->slice_size) + "," + std::to_string(block_r); - if (complex) sig += ",1"; // GROUP in the shared complex SELL kernel + if (complex) sig += parallel ? ",4" : ",1"; // SLOTS or GROUP std::vector args; args.reserve(6); @@ -661,9 +667,10 @@ flagsparseStatus_t run_sell(flagsparseHandle_t handle, const SpMatDescr* A, } const flagsparseStatus_t st = jit::launch( jit::codegen_module("spmv_sell.py"), - complex ? "_spmv_sell_complex_kernel" : "spmv_sell_f32acc", sig, - ctx(handle)->stream, n_slices, 1, 1, /*num_warps=*/1, - /*num_stages=*/2, args, &err); + complex ? (parallel ? "spmv_sell_musa_c32" : "_spmv_sell_complex_kernel") + : (parallel ? "spmv_sell_musa_f32acc" : "spmv_sell_f32acc"), sig, + ctx(handle)->stream, parallel ? (A->rows + 31) / 32 : n_slices, 1, 1, parallel ? 4 : 1, + parallel ? 1 : 2, args, &err); if (st != FLAGSPARSE_STATUS_SUCCESS) ctx(handle)->last_error = err; return st; } diff --git a/capi/src/ops/spvv.cpp b/capi/src/ops/spvv.cpp index 6a80251..7dbd6a0 100644 --- a/capi/src/ops/spvv.cpp +++ b/capi/src/ops/spvv.cpp @@ -91,7 +91,8 @@ flagsparseStatus_t run(flagsparseHandle_t handle, flagsparseOperation_t opX, if (c32) { sig += opX == FLAGSPARSE_OPERATION_CONJUGATE_TRANSPOSE ? "True," : "False,"; } - sig += std::to_string(kBlock); + const bool small_musa = std::string(adaptor::backend_name()) == "musa" && X->nnz <= 32; + sig += std::to_string(small_musa ? 32 : kBlock); std::vector args; args.reserve(5); @@ -105,16 +106,17 @@ flagsparseStatus_t run(flagsparseHandle_t handle, flagsparseOperation_t opX, flagsparseStatus_t st = jit::launch( jit::codegen_module("spvv.py"), c32 ? "spvv_c32" : (computeType == FLAGSPARSE_R_32I ? "spvv_mixed_i32acc" : "spvv_mixed_f32acc"), - sig, ctx(handle)->stream, 1, 1, 1, /*num_warps=*/4, /*num_stages=*/2, + sig, ctx(handle)->stream, 1, 1, 1, small_musa ? 1 : 4, /*num_stages=*/2, args, &err); if (st != FLAGSPARSE_STATUS_SUCCESS) { ctx(handle)->last_error = err; return st; } if (!result_on_device) { - // The runtime copies are synchronous and do not carry the handle's - // stream, so finish the launch before copying the scalar to host. - adaptor::synchronize(); + // MUSA's synchronous D2H copy is ordered after a default-stream launch. + // A different handle stream still needs the explicit synchronization. + if (std::string(adaptor::backend_name()) != "musa" || ctx(handle)->stream != nullptr) + adaptor::synchronize(); st = adaptor::memcpy_d2h(result, reinterpret_cast(device_result), dtype_size(computeType)); diff --git a/docs/debug/MUSA.md b/docs/debug/MUSA.md index 7c9b5e0..77c02fc 100644 --- a/docs/debug/MUSA.md +++ b/docs/debug/MUSA.md @@ -262,4 +262,378 @@ C API 精度 10/10 测试族通过。唯一的精度瑕疵是 `spmm_csc_f32_int_ - 其他后端是否也需要：是 / 否 / 不确定 ``` -（暂无） +### 2026-10-15 优化与复测口径 + +本轮按用户要求优化代码，验收目标是**每个交付变体各自的平均加速比 > 0.8x**，不是 65 项的整体平均。 +矩阵仍为 `tests/data` 的原始 10 个文件，split runner、测试规模、容差和基线优先级不变。 +本机 CPU cgroup 配额为 15 核，后续配对测量统一设置 +`OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8`，原始源码也用相同设置复测， +避免默认大线程池触发 CPU 限流。H800 带宽折算仍是估算，不能当作本机厂商实测。 +下列定向性能结果只用于调优，未替代最终完整 split；未达标和缺基线的条目必须继续如实报告。 + +### 2026-10-15 src/flagsparse/sparse_operations/mixed_spmx.py::MUSA SpMV/SpMM 路由 +- 改了什么：COO/CSC 转 CSR 使用有界、张量版本感知的缓存；短行 SpMV 调整行块，严重不均衡短行复用 NNZ 分块内核；SpMM 复用行批量 gather；避免输出自拷贝。 +- 为什么改（附报错原文）：无报错，原 COO 半精度/整数原子路径及逐行 SpMM 的实测加速比低于 0.8x。只缓存结构和必要的值重排，不缓存计算结果；输入变更失效。 +- 怎么验证的：10 矩阵定向复测 `results_musa_opt_20261015_nnz_threads8/`；Python 定向回归 80 passed、3 skipped，含乱序/重复 COO、值/索引/向量修改、无版本推理张量、非连续输出。 +- 其他后端是否也需要：否；新路由仅 MUSA。 + +### 2026-10-15 src/flagsparse/sparse_operations/_spmm_row_gather.py::_kernel/compute +- 改了什么：增加可选累加类型以支持 int8 到 int32；MUSA 使用 8 行、4 个稀疏槽、4 warps。 +- 为什么改（附报错原文）：无报错；MUSA 多行批量配置比逐行混合精度内核更快，整数输出必须保持整数累加。 +- 怎么验证的：上述 80 项回归与 10 矩阵 SpMM 定向性能复测。 +- 其他后端是否也需要：否；默认累加类型和其他后端的配置保持不变。 + +### 2026-10-15 src/flagsparse/sparse_operations/_spmv_csr_nnz.py::_musa_segmented_add +- 改了什么：MUSA 专用分段扫描的零分支使用整数文字 `0`，避免 int32 被浮点文字隐式提升；通过默认 False 的编译期 MUSA 开关选择，原共享归约函数保留原样。 +- 为什么改（附报错原文）：无报错；复用到 MUSA int8/int32 时必须保持精确整数归约。 +- 怎么验证的：整数 NNZ 路径包含重复列、空行、输入变更和非连续输出的精确回归。 +- 其他后端是否也需要：否；其他后端仍使用原 `_segmented_add` 及浮点零分支。 + +### 2026-10-15 src/flagsparse/sparse_operations/_spmv_csr_transpose.py::make_plan/gather +- 改了什么：MUSA 调整行块；半精度实数 gather 复用混合精度 CSR 执行器。 +- 为什么改（附报错原文）：无报错；原半精度 CSR 标量路径很慢。实验中的全行排序未带来收益，未用于最终混合 SpMV 路径。 +- 怎么验证的：80 项 Python 定向回归；含复数转置用例，未修改精度容差。 +- 其他后端是否也需要：否；配置与路由均以 MUSA 为条件。 + +### 2026-10-15 src/flagsparse/sparse_operations/spmv_csr.py::prepare_spmv_csr +- 改了什么：MUSA 半精度默认配置建立 non-gather plan，运行时使用优化的行批量路径。 +- 为什么改（附报错原文）：无报错；原 FP16 SpMV 性能明显落后。 +- 怎么验证的：半精度交付精度用例及 10 矩阵定向性能测试。 +- 其他后端是否也需要：否；显式配置和其他后端保持原行为。 + +### 2026-10-15 src/flagsparse/sparse_operations/spmv_csc.py::_spmv_csc_via_mixed +- 改了什么：MUSA 半精度非转置调用缓存化的 `spmv_csc_half`。 +- 为什么改（附报错原文）：无报错；原路径每次展开列号，抵消了缓存收益。 +- 怎么验证的：CSC 值/索引变更回归，10 矩阵定向复测约 4.61x（H800 折算）。 +- 其他后端是否也需要：否；原 CUDA/MACA 路由不变。 + +### 2026-10-15 src/flagsparse/sparse_operations/spmm_csr.py::默认半精度执行 +- 改了什么：MUSA 非 accuracy 模式 FP16 复用 row-gather。 +- 为什么改（附报错原文）：无报错；原路径逐行/逐槽执行缓慢。 +- 怎么验证的：80 项回归及 10 矩阵定向性能复测，FP16 变体约 2.08x（H800 折算）。 +- 其他后端是否也需要：否；CUDA 复数路由与其他后端不变。 + +### 2026-10-15 src/flagsparse/sparse_operations/spmm_coo.py::PreparedCooSpmmRoute +- 改了什么：MUSA FP16 准备 CSR gather 指针，复用现有行批量执行器。 +- 为什么改（附报错原文）：无报错；避免原子累加/逐非零元程序的高开销。 +- 怎么验证的：COO 精度用例和 10 矩阵定向性能复测，FP16 约 2.19x（H800 折算）。 +- 其他后端是否也需要：否。 + +### 2026-10-15 src/flagsparse/sparse_operations/spmv_sell.py::csr_to_sell/并行执行 +- 改了什么：MUSA 复数转换按实虚分量 index_copy；实数使用并行稀疏槽归约，复数使用已有并行内核。 +- 为什么改（附报错原文）：原转换报 `"IndexPut" not implemented for 'ComplexFloat'`；串行稀疏槽性能较慢。 +- 怎么验证的：10 矩阵、5 种数据/输出组合均 PASS；包含部分 slice、空行和 padding 的回归。 +- 其他后端是否也需要：否；复数规避复用 Ascend 的已有方式，其他后端执行路由不变。 + +### 2026-10-15 src/flagsparse/sparse_operations/gather_scatter.py::_triton_gather_impl +- 改了什么：MUSA FP16 使用原生 index_select；保留输出形状/类型检查，非连续输入先连续化，非连续输出经临时缓冲区复制。 +- 为什么改（附报错原文）：无异常；原 Triton 启动约 25us，原生索引约 7--16us。实测原生非连续输出有错误，不能直接传入。 +- 怎么验证的：4 个原始 Gather 性能规模均 PASS；重复索引、非连续输入/输出、原地修改、错误形状/类型回归。 +- 其他后端是否也需要：否。 + +### 2026-10-15 capi/flagsparse_codegen/spmv_csr.py::spmv_csr_row_tile_real +- 改了什么：增加 WIDE 开关，小矩阵使用 int32 地址计算，大范围仍保留 int64。 +- 为什么改（附报错原文）：无报错；无条件 int64 地址计算增加 MUSA 指令成本。 +- 怎么验证的：C API SpMV 精度测试族通过，10 矩阵 FP32 定向实测约 0.824x。 +- 其他后端是否也需要：否；默认 WIDE 为 True。 + +### 2026-10-15 capi/src/ops/spmv.cpp::短行 CSR 与 SELL 启动 +- 改了什么：按边界选择 CSR WIDE；MUSA SELL 选择并行稀疏槽内核和对应的 4-warps 配置。 +- 为什么改（附报错原文）：无报错；减少地址计算与串行 slot 循环的开销。 +- 怎么验证的：C API SpMV 精度测试族通过；10 矩阵定向性能复测。 +- 其他后端是否也需要：否；SELL 其他后端沿用原内核和配置。 + +### 2026-10-15 capi/flagsparse_codegen/spmv_sell.py::并行 SELL 包装 +- 改了什么：导出共享复数并行内核，增加 FP32/int32 累加的薄包装。 +- 为什么改（附报错原文）：无报错；C API 复用 Python 中同一个实现，避免复制计算逻辑。 +- 怎么验证的：C API SpMV 精度测试族通过，覆盖 SELL 实数、半精度、复数、整数。 +- 其他后端是否也需要：否；原导出仍保留。 + +### 2026-10-15 capi/flagsparse_codegen/spmm_csr.py::spmm_csr_batched8_real +- 改了什么：按 WIDE 选择地址宽度，增加可调稀疏槽宽度 K，默认值保持 4。 +- 为什么改（附报错原文）：无报错；int64 寻址与固定槽宽限制了 MUSA FP32/FP64 性能。 +- 怎么验证的：C API SpMM 精度测试族通过；10 矩阵定向性能测试。 +- 其他后端是否也需要：否；该内核只由 MUSA 路由选择。 + +### 2026-10-15 capi/src/ops/spmm.cpp::run_csr +- 改了什么：MUSA N=8 按累加类型选择行批量和槽宽，使用 4 warps；大地址保留 64 位。 +- 为什么改（附报错原文）：无报错；原 batch=4/1 warp 实测吞吐低。 +- 怎么验证的：10 矩阵定向实测 FP32 约 0.863x、FP64 约 0.902x，精度测试族通过。 +- 其他后端是否也需要：否。 + +### 2026-10-15 capi/src/ops/spvv.cpp::run +- 改了什么：MUSA nnz<=32 使用 block=32、1 warp，其他情况保持原配置。 +- 为什么改（附报错原文）：无报错；小向量不需要 256 个归约槽。 +- 怎么验证的：C API SpVV 精度测试族通过；定向性能有波动，尚不能据此认定稳定达到 0.8x。 +- 其他后端是否也需要：否。 + +### 2026-10-15 capi/src/ops/spgemm.cpp::workEstimation/compute/launch_hash +- 改了什么：MUSA 按行产品工作量分桶；复用主机计数存储，合并相邻的计数/溢出标志传输与前缀扫描。 +- 为什么改（附报错原文）：无报错；原来全矩阵使用最坏行的哈希容量，短行也承担单行哈希和大表开销。 +- 怎么验证的：C API SpGEMM 精度测试族通过；10 矩阵定向复测约 0.825x，存在波动。排序去重实验精度通过但性能退化，已撤回。每次重新计算，不缓存输出计数。 +- 其他后端是否也需要：否；分桶与主机缓冲复用均限定 MUSA，其他后端保留原两次传输、溢出检查和前缀扫描流程。 + +### 2026-10-15 tests/pytest/test_musa_sparse_routes.py::新增回归 +- 改了什么：新增 MUSA 格式转换缓存失效、整数精确累加、SELL padding/部分 slice、非连续输出、推理张量和 Gather 边界测试。 +- 为什么改（附报错原文）：优化后必须验证输入实时可变，不能只测同一输入重复执行。 +- 怎么验证的：连同 q4 SpMV/SpMM 精度用例共 80 passed、3 skipped；跳过的是不在该缓存路由内的 CSC 组合，不是交付变体跳过。 +- 其他后端是否也需要：否；该文件显式要求 MUSA，原跨后端测试未删除或放宽。 + +### 2026-10-15 MUSA-only 范围收紧 +- 改了什么：按用户补充要求，`mixed_spmx._finish` 仅在 MUSA 避免输出自拷贝，非 MUSA 恢复原行块选择；`_spmm_row_gather.compute` 的 int32 累加仅在 MUSA 启用；`_spmv_csr_nnz` 的 MUSA 归约通过编译期开关隔离；`capi/src/ops/spgemm.cpp` 的缓冲复用和分桶显式限定 MUSA。 +- 为什么改（附报错原文）：无报错；用户要求“修改仅限于musa后端”，不能把共享行为变化带给其他后端。 +- 怎么验证的：新增不执行 GPU 内核的路由隔离回归，分别模拟 MUSA/非 MUSA 检查自拷贝、整数累加和行块参数，6 passed；`git diff --check` 通过。此前 80 项 GPU 回归是隔离前结果；隔离后的 Python 回归、C API 重新编译及 SpMV/SpMM/SpGEMM/SpVV 精度复验排在当前 split 完成后运行，日志 `results_musa_scope_20261015.log`，未完成前不视为通过。 +- 其他后端是否也需要：否；没有 CUDA/ROCm 等其他后端的真机验证，不声称已验证。 +- 测试版本说明：`results_musa_split_65_20261015_opt/` 在隔离修改前启动。修改未调整 MUSA 算法、配置或基线，但源码哈希发生变化；该轮作为隔离前优化结果保留，不标为最终隔离后完整验收。 + +### 2026-10-16 本轮状态与定向数据 + +用户确认停止旧整轮后，旧 runner 及排队验证已停止，原目录与日志全部保留。本轮继续限定 MUSA，未修改测试矩阵、规模、容差或基线优先级。核对 `tests/data` 的 10 个 `.mtx` 大小与上文一致；导入路径为当前仓库 `src/flagsparse/__init__.py`，运行环境为 `mthreads musa None`、torch 2.7.1、Triton 3.6.0、MTT S5000。 + +定向结果不是完整交付验收，所有 H800 折算均仍为估算： + +- `results_musa_opt_20261016_complex3d/`：原 10 矩阵上每条精度 PASS；`spmv_csr_f32c32_int_non` 平均从本轮改动前约 0.703x 提升至 0.819x。FP16→FP32 0.779x、FP16 0.754x、int8→FP32 0.820x、int8→int32 0.812x。COO 对应三项约 0.586x、0.685x、0.645x，仍未达标。 +- `results_musa_opt_20261016_sell_flat/sell.csv`：5 种值/输出组合、10 矩阵，共 50 条 PASS。按已有同规模 Python 折算基线，SELL FP16 0.609x、FP32 0.695x、int8→int32 0.590x，仍未达标。 +- `results_musa_opt_20261016_capi_sell/`：4 个交付 SELL 变体各 10 矩阵，40 条 `status=ok`、精度通过；复数 SELL 按已有 H800 同矩阵折算基线为 0.610x（旧轮约 0.275x），仍未达标。 +- 固定最大行宽展开、SELL 行连续布局、成对复数加载、自动 Graph 重放均仅作临时探针，未显示稳定整体收益，没有接入正式源码。SELL→CSR 在 ASIC 上明显改善、其他矩阵有退化，也未接入。 +- Python 扩大回归（MUSA 专用回归、q4 SpMV/SpMM、原 CSR 精度文件）已得到 559 passed、1659 skipped、23 deselected；跳过为原有后端/算法适用性规则，未新增或放宽规则。最终补充非连续输出后的结果及完整 split 见后续记录。 + +**尚未达到“每个变体各自平均 > 0.8x”目标，不能用上述单项或整体平均代替验收。** + +### 2026-10-16 src/flagsparse/sparse_operations/mixed_spmx.py::_musa_csr_real_by_complex_kernel/_musa_coo_csr +- 改了什么：实数×复数使用三维实虚分量布局；严重不均衡短行可走复数 NNZ 路径。MUSA COO 从已排序行号用 `searchsorted` 建 CSR 指针，稀疏输入检查随版本感知缓存失效重新执行，RHS/输出仍每次检查。 +- 为什么改（附报错原文）：无异常；原实数×复数逐行 kernel 慢。另复现原生 MUSA `bincount` 偶发少计数，4930 个输入得到计数和 4917，导致后续行指针及结果错误。 +- 怎么验证的：10 矩阵定向结果见本轮状态；FP32/FP64 实数×复数、跨块、空行、输入修改回归通过，不缓存输出。 +- 其他后端是否也需要：否；新 kernel 和预处理均仅 MUSA 路由选择。 + +### 2026-10-16 src/flagsparse/sparse_operations/_spmv_csr_nnz.py::_nnz_kernel +- 改了什么：增加默认 False 的 `COMPLEX_RHS`，MUSA 分别扫描实虚分量并处理跨块前缀；MUSA 的前缀补偿零分支保持整数类型。 +- 为什么改（附报错原文）：无异常；实数×复数需要 NNZ 均衡，int8→int32 不允许跨块补偿隐式转 FP32 而丢失大整数低位。 +- 怎么验证的：新增 1101 个 127×127 的精确整数累加，以及 601 个非零元跨块复数、空行和实时修改回归；未开启新参数的其他后端保持原扫描/补偿路径。 +- 其他后端是否也需要：否。 + +### 2026-10-16 src/flagsparse/sparse_operations/_spmv_csr_transpose.py::prepare/gather +- 改了什么：MUSA 转置预处理使用排序后 `searchsorted` 替代 `bincount`；FP16 已按 CSR 排列时提前进入执行器；非连续输出使用临时缓冲区后复制回调用方视图。 +- 为什么改（附报错原文）：复现 MUSA `bincount` 少计数；通用处理存在额外调用开销，且临时输出不能丢失调用方指定的非连续输出。 +- 怎么验证的：扩大 CSR 精度、q4 SpMV/SpMM 回归通过，补充直接/间接 FP16 gather 的非连续输出回归。 +- 其他后端是否也需要：否；其他后端仍使用原预处理和执行路径。 + +### 2026-10-16 src/flagsparse/sparse_operations/spmv_coo.py::PreparedCoo.__init__ +- 改了什么：仅 MUSA 对排序后的转置行号使用 `searchsorted` 构造 CSR 行指针。 +- 为什么改（附报错原文）：q4 COO FP32 转置出现 `AssertionError: max error 24 (scale 21.4)`；诊断时行指针末值 4862 而实际 nnz=4930。独立重复计数探针亦复现少计数，不是归约容差问题。 +- 怎么验证的：新增 30 次重复 COO 转置回归，原 q4 转置及 CSR 精度回归通过。 +- 其他后端是否也需要：否。 + +### 2026-10-16 src/flagsparse/sparse_operations/spmm_coo.py::PreparedCooSpmmRoute.__init__ +- 改了什么：MUSA FP16 的已排序 COO gather 指针也由 `searchsorted` 生成，CUDA 复数仍保持原 `bincount` 路径。 +- 为什么改（附报错原文）：同一 MUSA 计数原语用于该路由；避免把少计数风险留在 SpMM 中。 +- 怎么验证的：q4 SpMM 与 MUSA 混合输入修改回归通过。 +- 其他后端是否也需要：否。 + +### 2026-10-16 src/flagsparse/sparse_operations/spmv_sell.py::_musa_spmv_sell_kernel/flagsparse_spmv_sell +- 改了什么：MUSA 使用按行并行的二维 ROWS×SLOTS 布局；实数、整数和复数共享索引逻辑，复数累加显式映射 FP32/FP64。 +- 为什么改（附报错原文）：旧 SELL 性能低；新路由最初遗漏复数累加类型，回归报 `KeyError: torch.complex64` / `torch.complex128`，现已修正。 +- 怎么验证的：部分 slice、空行、padding、FP64/complex128 回归通过；上述 50 条 Python 和 40 条 C API 定向记录通过精度。 +- 其他后端是否也需要：否；其他后端内核和配置保持不变。 + +### 2026-10-16 capi/flagsparse_codegen/spmv_sell.py / capi/src/ops/spmv.cpp::MUSA SELL +- 改了什么：MUSA 专用薄包装导出新的共享二维 SELL kernel，按 32 行/program、4 slots、4 warps 启动；非 MUSA 保留原包装和网格。 +- 为什么改（附报错原文）：无报错；复用经过 Python 回归的实现降低旧串行 SELL 开销。 +- 怎么验证的：重新构建确认无待编译项，C API SpMV 精度族通过，40 条 SELL 定向记录全部精度通过。 +- 其他后端是否也需要：否。 + +### 2026-10-16 capi/src/ops/spvv.cpp::run / capi/ctest/accuracy/test_spvv.cpp +- 改了什么：MUSA 默认流主机结果依赖有序同步 D2H，省去重复全设备同步；非默认流和其他后端仍显式同步。增加 MUSA 默认流/非默认流、复数输入反复变更的主机结果检查。 +- 为什么改（附报错原文）：无报错；重复同步影响小向量耗时。 +- 怎么验证的：40 次双流复数结果检查及 C API SpVV 精度族通过；`results_musa_opt_20261015_spvv_sync/` 原 10 矩阵目录口径下的向量 fixture 定向约 1.202x，完整验收仍以新 split 为准。 +- 其他后端是否也需要：否。 + +### 2026-10-16 tests/pytest/test_musa_sparse_routes.py::补充回归 +- 改了什么：增加重复 COO 转置、大整数跨块、实数×复数的跨块/输入变更/FP64/部分 tile，以及直接和间接 FP16 gather 非连续输出用例。 +- 为什么改（附报错原文）：覆盖本轮预处理及启动配置变更的真实失败和边界条件；不放宽原容差。 +- 怎么验证的：扩大 GPU 回归见本轮状态与后续最终验证记录。 +- 其他后端是否也需要：否；文件显式限定 MUSA。 + +### 2026-10-16 tests/pytest/test_spmv_csr_accuracy.py::test_spmv_csr_empty_dtype_ops +- 改了什么：仅 MUSA 复数结果改为对 `view_as_real(actual)` 执行原零值断言。 +- 为什么改（附报错原文）：原检查调用 MUSA 不支持的复数 `count_nonzero`，报 `RuntimeError: SetMUTensorDType Unsupported tensor dtype: ComplexFloat` / `ComplexDouble`。 +- 怎么验证的：原测试覆盖的所有 MUSA 空矩阵/数据类型/转置组合通过；仍断言非零元素数严格为 0，未跳过测试或改变精度标准。 +- 其他后端是否也需要：否；其他后端检查保持原样。 + +### 2026-10-16 最终启动前验证与阶段性完整复测 + +- 最终 Python 回归：`561 passed, 1659 skipped, 23 deselected`，15.82 秒；日志 `results_musa_opt_20261016_validation.log`。包含后来补充的直接/间接 FP16 gather 非连续输出检查。 +- C API 交付相关精度族：`accuracy.axpby/spmv/spmm/spvv/spgemm/sddmm/gather/scatter`，8/8 passed，44.97 秒。未把非交付 SpSV/SpSM 纳入本轮验收。 +- `cmake --build capi/build --parallel 8` 确认 `ninja: no work to do.`；`git diff --check` 通过。 +- 新整轮目录：`results_musa_split_65_20261016_opt`，日志 `results_musa_split_65_20261016_opt.log`。使用 `setsid`、12 小时外层限时、原 split runner，`--mode normal --benchmark-input tests/data --timeout 3600 --h800-reference`，并固定三种 CPU 线程池为 8。 +- 此整轮用于检查当前阶段代码的 65 项结果，不预先视为所有加速比达标；定向测量中 FP16、COO、SELL 仍低于 0.8x，完整表生成后仍需按每个变体逐项检查。 + +### 2026-10-16 上述整轮完成后的原始数据复核 + +`results_musa_split_65_20261016_opt/` 已完成，65 项、0 missing，合并表的精度/性能标签均为 Passed， +但逐行复核仍有 16 个变体平均加速比低于 0.8x，另有 2 条 C API 原始精度失败，**不能按标签认定完全通过**。 +Python 普通 CSC SpMM 的 `amazon0601 float32 trans n=32` 也有误差比 2.53 的失败，使 runner 返回 1。 +C API 原始文件已保留在该目录的 `capi_raw_snapshot/`。Python CSC 性能族耗时 1212.88 秒；C API 8 个性能族 +合计 1097.73 秒。补基线日志原文： + +```text +--h800-reference (musa-s5000): NoBaseline variants filled from the Python half=26, from the C API time=0, still without a baseline=0 [] (estimates, marked baseline=h800_scaled) +``` + +后续整轮显式使用独立的 `--capi-bench-out <新结果目录>/capi`，不再复用共享输出目录。 + +### 2026-10-16 src/flagsparse/sparse_operations/gather_scatter.py::_triton_scatter_impl +- 改了什么：MUSA FP16/int8 默认配置使用原生 `index_copy_`，显式 block_size 仍走原 Triton 路径；原生索引提升到 int64 并如实记录执行索引类型。保留原输出分配、清零、reset_output 和空输入行为。 +- 为什么改（附报错原文）：小向量 Triton 启动开销较大；原生 int32 索引实际报 `RuntimeError: index_copy_(): Expected a long tensor for index, but got Int`，不能假定支持。 +- 怎么验证的：`results_musa_opt_20261016_routes_v2/scatter.csv` 8/8 PASS，FP16/int8 对本机 PyTorch 约 0.96x/0.94x；增加非连续输入输出、实时修改、显式调优和非 MUSA 路由隔离回归。后续完整验证见最后记录。 +- 其他后端是否也需要：否；原生路由仅 MUSA 默认配置。 + +### 2026-10-16 src/flagsparse/sparse_operations/spmm_csc.py::FP16 路由和转置 gather +- 改了什么：MUSA 默认非转置 FP16 缓存 CSC 到 CSR 的结构/值重排，使用已有混合 SpMM；原始索引、指针、值、存储或版本变化均失效，不缓存输出。MUSA 未指定调优块参数的 FP32/complex64 转置使用行拥有输出的 gather；支持 auto 和其解析后的 spmm_csc_base，保留其他类型和显式块参数路径。索引回退保留路由标志，元数据反映实际块配置。 +- 为什么改（附报错原文）：FP16 每次转 FP32 并原子散射，旧折算均值仅 0.012703x。普通 FP32 转置出现上述误差比 2.53；旧 complex64 转置在 amazon0601 约 1.2 秒/次。 +- 怎么验证的：`results_musa_opt_20261016_csc_integrated_v2/` FP16 原 10 矩阵全 PASS，折算均值 3.3125x；`results_musa_opt_20261016_routes_v4/csc_trans.csv` 普通 FP32/complex64、trans/conj、原 10 矩阵共 40 条 PASS。amazon0601 FP32 约 65.82 ms 到 0.917 ms，complex64 约 1233 ms 到 13.32 ms，均为同脚本本身耗时，不是新的厂商基线。 +- 其他后端是否也需要：否；新增选择条件显式限定 MUSA，其他后端保持原实现。 + +### 2026-10-16 src/flagsparse/sparse_operations/mixed_spmx.py::_musa_csr_spmv/_musa_row_tile +- 改了什么：MUSA 网格使用整数上取整，去掉复数路径未用网格；超短实数/整数行采用 256 行、2 slots，大规模中短行采用 128 行、8 slots，小规模中短行仍为 64 行。 +- 为什么改（附报错原文）：无异常；减少主机调用开销和超短行 program 数量。 +- 怎么验证的：`results_musa_opt_20261016_wide_rows_v1/measurements.json` 配置探针保留首尾原配置作对照；roadNet-TX FP16/int8 约快 10%～15%。原有完整精度标准不变，新增实时存储替换回归。整轮验收仍以新 split 为准。 +- 其他后端是否也需要：否；仅修改 MUSA 专用 helper。 + +### 2026-10-16 src/flagsparse/sparse_operations/spmv_sell.py::MUSA 分发与计划 +- 改了什么：MUSA 提前分发，省去通用路径未用的网格/环境变量处理；对 padding 超过有效非零元 2 倍的实数/整数 SELL 压缩为 CSR。版本感知计划复用稀疏校验和连续化结果，RHS 与输出仍每次检查；validate=False 计划升级到 True 时重新校验，输入修改后失效。 +- 为什么改（附报错原文）：无异常；ASIC 和 filter3D 的 SELL padding 分别约 3.98 倍、2.48 倍，读取无效槽开销大；重复包装校验影响短 kernel。 +- 怎么验证的：原 10 矩阵 50 条精度 PASS，包含非交付 FP16 到 FP32；`results_musa_opt_20261016_routes_v4/sell.csv` 保留全部原始时间（包括 wave FP16 的较慢单次记录），不删慢行。新增高 padding、部分 slice、指针/列/值修改、推理张量、非连续 RHS/out、校验升级和非法索引回归。 +- 其他后端是否也需要：否；其他后端的校验和计算分支保持原样。 + +### 2026-10-16 capi/flagsparse_codegen/sparse_gather.py::sparse_gather_spmv_f32 +- 改了什么：复数各 lane 的累加仍为 FP32，仅最终跨 lane 归约提升到 FP64，再转回 FP32 执行 alpha/beta。 +- 为什么改（附报错原文）：ASIC 第 405878 行的 3 个乘积包含小项和约正负 74 的抵消；原 FP32 归约丢小项，严格误差比 1.363375。CPU 与 GPU 探针均复现，未修改容差。 +- 怎么验证的：`results_musa_opt_20261016_precision_v1/measurements.json`；加宽最后归约后 ASIC 该变体整体严格误差比约 0.0495。C API 新增双分量抵消回归，精度族通过。 +- 其他后端是否也需要：否；该模块由 MUSA 专用 sparse_gather 分支使用。 + +### 2026-10-16 capi/src/core/sparse_gather.hpp::launch_sparse_gather +- 改了什么：间接实数 SpMV 平均行长不足 12 时保留 scatter，较长行采用确定性 gather。 +- 为什么改（附报错原文）：cfd2 CSC FP32 的原子加法顺序导致原始严格精度失败；较长行可摊薄间接读取开销。 +- 怎么验证的：`results_musa_opt_20261016_capi_precision_v2/` 原 10 矩阵，CSC FP32 10 条可比且精度通过，muSPARSE 平均约 0.974x；复数 CSR 共轭转置 10 条精度通过，muSPARSE 平均约 0.860x。 +- 其他后端是否也需要：否；调用方只在 MUSA 选择此执行器。 + +### 2026-10-16 capi/src/ops/gather_scatter.cpp::gather_max_programs +- 改了什么：仅 MUSA 使用按设备区分的 thread_local program 上限缓存，避免每次查询硬件 SM/warp/线程属性。 +- 为什么改（附报错原文）：无异常；小 Gather 的主机属性查询占比高。 +- 怎么验证的：上述 C API 定向目录中 FP32/FP64/complex64/complex128 各 10 条精度通过，对 muSPARSE 平均约 0.970x/0.973x/0.966x/0.971x。 +- 其他后端是否也需要：否；其他后端仍保留原查询行为。 + +### 2026-10-16 capi/ctest/accuracy/test_spmv.cpp::MusaComplexGatherCancellation +- 改了什么：MUSA 专用宏下增加来自 ASIC 的三项抵消回归，使用原有 complex64 严格容差，检查两个分量。 +- 为什么改（附报错原文）：覆盖上文精度缺陷，防止只通过总体状态标签而漏掉原始失败。 +- 怎么验证的：重新构建并运行 accuracy.spmv/gather/scatter，3/3 passed；后续扩大验证见最后记录。 +- 其他后端是否也需要：否；新测试显式限定 MUSA。 + +### 2026-10-16 定向调优中未接入的实验 + +- `results_musa_opt_20261016_sorted_ell_v1/`：按行长重排的连续布局有局部收益，但多矩阵退化，未接入。 +- `results_musa_opt_20261016_coo_cache_pair_v1/`：同进程交替配对表明，拆分 COO 拓扑/值缓存没有稳定收益，部分短行慢约 1%～10%；仅撤销本次实验，保留此前版本感知缓存。 +- `results_musa_opt_20261016_loop_v1/`、`results_musa_opt_20261016_load_v1/`：循环展开、软件流水和读取缓存提示没有可靠整体收益，未接入。 +- `results_musa_opt_20261016_sell_complex_compact_v1/`：复数 SELL 压缩在 filter3D 有收益，但 ASIC 退化，10 矩阵平均约 0.744x；未接入。稀疏校验计划也仅在 FP16/BF16/FP32/int8 使用，复数保持原校验与 kernel。 +- 当前验证过的扩大 Python 回归为 845 passed、1658 skipped，另增 6 项原生 Scatter 调优/后端隔离回归通过。仍未达到每个变体 >0.8x；任何定向均值不能替代新的完整 split 结果。 + +### 2026-10-16 本次源码固定后的验证与下一轮完整复测 + +- Python 扩大回归：851 passed、1658 skipped，18.69 秒；`results_musa_opt_20261016_routes_v4/accuracy.xml`。包含 CSR/CSC/SELL、q4、Gather/Scatter 和 MUSA 路由隔离回归；没有降低容差或新增跳过规则。 +- C API 交付相关精度族：8/8 passed，39.42 秒；`results_musa_opt_20261016_routes_v4/capi_accuracy.log`。构建确认 `ninja: no work to do.`，`git diff --check` 通过。 +- `results_musa_opt_20261016_spmv_final_v1/`：CSR 5 项、COO 3 项各原 10 矩阵，共 80 条 PASS。CSR FP16/FP16到FP32/FP32乘complex64/int8到FP32/int8到int32 的折算均值分别为 0.810x/0.825x/0.842x/0.887x/0.876x；COO 对应三项仍为 0.679x/0.634x/0.691x，未达标。这些 H800 折算仍是估算。 +- `results_musa_opt_20261016_routes_v4/scatter_final.csv`：默认 20 次 warmup、200 次迭代，FP16/int8 × 4 个原向量规模，8/8 PASS；对本机 PyTorch 各行约 0.92x～0.97x，不能与厂商或 H800 折算基线混称。 +- SELL 最新定向数据仍有低于 0.8x 的类型，所有慢行原样保留。因此下一轮是阶段性完整复测，**不是全部 65 项达标声明**。 +- `tests/data` 的原 10 个文件大小再次核对一致，torch 2.7.1、Triton 3.6.0、MTT S5000、驱动 3.3.5-server，HEAD `4d5ab61`；导入当前 `src/`，`mthreads musa None`。 + +新目录 `results_musa_split_65_20261016_opt2/`，保留本地源码补丁、未跟踪回归测试补丁和代码 SHA256。 +命令如下；是否完成及是否达标必须检查该目录的新结果，不能沿用旧表： + +```bash +setsid env PYTHONPATH=$PWD/src FLAGSPARSE_BACKEND=mthreads MUSA_HOME=/usr/local/musa \ + OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 \ + timeout -s KILL 43200 python3 -u run_flagsparse_split_delivery.py \ + --mode normal --benchmark-input tests/data --timeout 3600 --h800-reference \ + --results-dir results_musa_split_65_20261016_opt2 \ + --capi-bench-out results_musa_split_65_20261016_opt2/capi \ + > results_musa_split_65_20261016_opt2.log 2>&1 < /dev/null & +``` + +### 2026-10-17 上一轮结果与本轮优化范围 + +`results_musa_split_65_20261016_opt2/` 已完成，65 项、0 missing，合并精度/性能标签均为 Passed， +但实际仍有 8 个变体均值小于 0.8x：COO FP16到FP32/FP16/int8到int32、SELL FP16/int8到int32/complex64、 +AXPBY FP16、SpGEMM FP32。继续按用户要求只优化 MUSA 后端；原始 10 矩阵、容差、迭代次数和基线优先级不变。 +下述定向数据不是完整验收结论，尤其不能把 H800 折算值当成本机厂商实测。 + +### 2026-10-17 src/flagsparse/sparse_operations/vector_ops.py::_musa_axpby +- 改了什么：MUSA FP16/BF16 的中小向量用单个融合 kernel 完成缩放与稀疏更新，缓存有界、版本感知的索引反向映射；值和输出每次重新读取。大向量和其他后端保留原路由。 +- 为什么改（附报错原文）：无报错；旧轮 AXPBY FP16 均值为 0.723508x，分开的缩放/更新启动开销较高。 +- 怎么验证的：原 4 个向量规模全部 PASS，对本机 PyTorch 平均 1.221323x；原始数据 `results_musa_opt_20261017_layout_v1/axpby_public.csv`。新增 FP16/BF16、索引宽度、输入变更、推理张量、空输入和后端隔离回归，均纳入本轮扩大验证。 +- 其他后端是否也需要：否。 + +### 2026-10-17 capi/flagsparse_codegen/spgemm_csr.py::musa_spgemm_scan_* +- 改了什么：增加 MUSA 分块行计数前缀和、块级前缀和及偏移回填 kernel，同时归约总 nnz、最大行 nnz 与溢出标志。 +- 为什么改（附报错原文）：无报错；SpGEMM compute 阶段原先回传每行 counts/overflow 并在主机执行前缀和，旧轮均值为 0.715145x。 +- 怎么验证的：MUSA C API 重新构建，accuracy.spgemm 通过，含下述跨块边界回归；原 10 矩阵定向记录 `results_musa_opt_20261017_layout_v1/capi_spgemm/`，6 个规则内可比矩阵对本机 muSPARSE 平均 1.076070x。非方阵、输出预算与厂商精度状态按原规则保留。 +- 其他后端是否也需要：否；新 kernel 只由 MUSA 分支调用。 + +### 2026-10-17 capi/src/ops/spgemm.cpp::musa_build_indptr/compute/copy +- 改了什么：MUSA compute 在设备上构建 indptr，只回传 3 个 int64 元数据；copy 阶段仍按原排序流程回传所需行指针。不缓存输出计数，每次重新执行 count 和 scan。其他后端的 counts/overflow 回传与主机前缀和保持原样。 +- 为什么改（附报错原文）：无报错；减少 compute 阶段逐行回传及主机扫描成本。 +- 怎么验证的：同上 SpGEMM 原 10 矩阵定向性能和 C API 精度测试；完整 split 仍需重新验收。 +- 其他后端是否也需要：否。 + +### 2026-10-17 capi/ctest/accuracy/test_spgemm.cpp::MusaMultiBlockRowPrefix +- 改了什么：增加 MUSA 专用 1024、1025、3079 行的 FP32/FP64 用例，检查 scan 块边界、部分末块和完整结果。 +- 为什么改（附报错原文）：无报错；原有小规模精度用例不能覆盖新增前缀和的跨块路径。 +- 怎么验证的：accuracy.spgemm 通过，未更改原有精度判定。 +- 其他后端是否也需要：否；其他后端不运行新增 MUSA 专用用例。 + +### 2026-10-17 src/flagsparse/sparse_operations/_musa_spmv_launch.py::launch +- 改了什么：新增最多 64 项的编译启动器 LRU，只缓存 CompiledKernel 的公开 launcher，不保留输入张量或输出。key 包含 kernel、网格、数据类型、设备、地址对齐、非零指针别名、运行时标量、constexpr 和 kernel.debug；每次执行仍解析当前流并运行 Triton profiling hooks。存在 pre-run hook、运行时 debug、instrumentation 或自定义编译阶段 hook 时回退正常 JIT。 +- 为什么改（附报错原文）：无报错；MUSA 短行 SpMV 受 JIT 参数绑定及特化处理的主机开销限制，公开 API 配对测量确认启动器复用有收益。 +- 怎么验证的：新增 key 隔离、LRU 淘汰、调试/插桩回退、未对齐连续输入/输出、输出独立性、实时 RHS、非默认流及非 MUSA 路由隔离回归；本轮扩大回归全部通过。 +- 其他后端是否也需要：否；模块导入不初始化 MUSA 或导入 Triton knobs，只有 MUSA 路由调用 launch。 + +### 2026-10-17 src/flagsparse/sparse_operations/mixed_spmx.py::_musa_csr_spmv/_musa_row_tile +- 改了什么：MUSA 实数/整数短行 CSR 和 NNZ 路由复用上述启动器；按平均行长、行数与整数累加选择行块和 slots。普通实数乘复数行 kernel 保持原启动方式。 +- 为什么改（附报错原文）：无报错；旧轮 COO 三个半精度/整数变体低于 0.8x，格式转换缓存后的重复 JIT 启动仍占比较高。 +- 怎么验证的：最终公开 API 定向复测 `results_musa_opt_20261017_dispatch_v2/coo_public_final.csv`，原 10 矩阵共 30 条精度 PASS；FP16到FP32/FP16/int8到int32 的 H800 折算均值分别为 0.821573x/0.897036x/0.902609x。首项裕量较小，不能以定向数据代替完整 split。 +- 其他后端是否也需要：否。 + +### 2026-10-17 src/flagsparse/sparse_operations/spmv_sell.py::MUSA kernel 启动 +- 改了什么：MUSA SELL 复用上述编译启动器；保留已有版本感知计划、高 padding 实数压缩及 dtype/规模分块策略。 +- 为什么改（附报错原文）：无报错；旧轮 SELL FP16/int8到int32/complex64 均值为 0.735946x/0.718123x/0.787854x，短 kernel 的主机启动成本占比较高。 +- 怎么验证的：`results_musa_opt_20261017_dispatch_v2/sell_public.csv` 原 10 矩阵共 30 条精度 PASS；三项 H800 折算均值为 0.969267x/0.961707x/0.893333x。最终重复测量 `sell_public_final.csv` 仍为 30 条 PASS，均值为 0.981482x/0.984545x/0.911817x。 +- 其他后端是否也需要：否。 + +### 2026-10-17 tests/pytest/test_musa_sparse_routes.py::AXPBY 与启动缓存回归 +- 改了什么：新增上述 AXPBY、启动缓存特化/回退、公开 COO/SELL 缓冲区及流切换、后端路由隔离用例；保留之前所有回归。 +- 为什么改（附报错原文）：无报错；优化不能依赖输入地址不变、复用旧输出或忽略运行时调试设置。 +- 怎么验证的：最终扩大 Python 回归 `879 passed, 1658 skipped`，18.79 秒，`results_musa_opt_20261017_dispatch_v2/accuracy_final.xml`；跳过均沿用已有适用性规则，未放宽容差。C API 交付相关精度族 8/8 passed，39.89 秒，`capi_accuracy.log`；`git diff --check` 通过。 +- 其他后端是否也需要：否；此回归文件显式要求 MUSA。 + +### 2026-10-17 未接入的实验 + +- MUSA mixed COO 块内 segmented atomic 实验没有公开 API 的整体收益，新增实验 kernel 和 wrapper 已撤回；不改变当前 CSR 缓存路线。 +- 裸 kernel 启动探针低估了公开 API 的主机开销，最终采用同一公开 API 的原规模配对测量与仓库原 benchmark 校验收益。 +- SELL 的额外 tile 实验没有稳定整体收益，保留原配置；性能与精度原始记录见 `results_musa_opt_20261017_layout_v1/` 和 `results_musa_opt_20261017_dispatch_v2/`。 + +### 2026-10-17 冻结源码与完整复测 + +- 新结果目录 `results_musa_split_65_20261017_opt3/`，不复用旧轮；包含源码补丁、未跟踪 helper/回归测试补丁和 SHA256。 +- 源码导入路径 `/root/gcx/FlagSparse/src/flagsparse/__init__.py`；HEAD `4d5ab61`；torch 2.7.1、torch_musa 2.7.1+5ee0a64、Triton 3.6.0、MTT S5000，驱动 3.3.5-server，mcc 5.1.0，torch.version.musa=40305。 +- 再次核对 `tests/data` 原 10 个矩阵名称与大小一致；固定三个 CPU 线程池为 8。不跳过 C API 构建，不改变原 runner 的规模、容差、迭代次数或基线优先级。 +- 这是本轮优化后的完整验收启动记录，不是已完成或全部达标声明。必须等新整轮结束后逐变体检查实际平均加速比及原始精度记录。 + +```bash +setsid env PYTHONPATH=$PWD/src FLAGSPARSE_BACKEND=mthreads MUSA_HOME=/usr/local/musa \ + OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 \ + timeout -s KILL 43200 python3 -u run_flagsparse_split_delivery.py \ + --mode normal --benchmark-input tests/data --timeout 3600 --h800-reference \ + --results-dir results_musa_split_65_20261017_opt3 \ + --capi-bench-out results_musa_split_65_20261017_opt3/capi \ + > results_musa_split_65_20261017_opt3.log 2>&1 < /dev/null & +``` diff --git a/src/flagsparse/sparse_operations/_musa_spmv_launch.py b/src/flagsparse/sparse_operations/_musa_spmv_launch.py new file mode 100644 index 0000000..5251b6f --- /dev/null +++ b/src/flagsparse/sparse_operations/_musa_spmv_launch.py @@ -0,0 +1,45 @@ +# Copyright 2026 FlagOS Contributors +# SPDX-License-Identifier: Apache-2.0 +"""Reuse compiled MUSA SpMV launchers, never operands or computed outputs.""" + +from collections import OrderedDict + + +_LAUNCHERS = OrderedDict() +_MAX_LAUNCHERS = 64 + + +def launch(kernel, grid, tensors, scalars, constants): + """Keep the pointer types, alignment and aliases of Triton's specialization. + + Sparse plan caches still validate tensor versions/storage on every call. + This cache only bypasses repeated JIT argument binding for the same kernel + signature. CompiledKernel's public launcher resolves the current stream + and runs Triton's profiling hooks on every invocation. + """ + from triton import knobs + + if ( + kernel.pre_run_hooks or knobs.runtime.debug + or knobs.compilation.instrumentation_mode + or knobs.runtime.add_stages_inspection_hook is not None + ): + return kernel[grid](*tensors, *scalars, *constants, num_warps=4, num_stages=1) + addresses = tuple(t.data_ptr() for t in tensors) + signature = tuple((t.dtype, t.device, p % 16) for t, p in zip(tensors, addresses)) + aliases = () + if len(set(addresses)) != len(addresses): + aliases = tuple((i, j) for j, p in enumerate(addresses) if p + for i in range(j) if addresses[i] == p) + grid = tuple(grid) + (1,) * (3 - len(grid)) + key = (kernel, grid, signature, aliases, scalars, constants, kernel.debug) + runner = _LAUNCHERS.get(key) + if runner is None: + compiled = kernel[grid](*tensors, *scalars, *constants, num_warps=4, num_stages=1) + if compiled is not None: + _LAUNCHERS[key] = compiled[grid] + while len(_LAUNCHERS) > _MAX_LAUNCHERS: + _LAUNCHERS.popitem(last=False) + return compiled + _LAUNCHERS.move_to_end(key) + runner(*tensors, *scalars, *constants) diff --git a/src/flagsparse/sparse_operations/_spmm_row_gather.py b/src/flagsparse/sparse_operations/_spmm_row_gather.py index 9584b7c..383c119 100644 --- a/src/flagsparse/sparse_operations/_spmm_row_gather.py +++ b/src/flagsparse/sparse_operations/_spmm_row_gather.py @@ -4,20 +4,21 @@ import torch import triton import triton.language as tl +from ._common import _is_mthreads_runtime @triton.jit def _kernel(V, I, P, B, C, M, N, BK, BN, CM, CN, COMPLEX: tl.constexpr, R: tl.constexpr, K: tl.constexpr, - TILE_N: tl.constexpr): + TILE_N: tl.constexpr, ACC: tl.constexpr = tl.float32): rows = tl.program_id(0) * R + tl.arange(0, R) valid = rows < M starts = tl.load(P + rows, valid, other=0) ends = tl.load(P + rows + 1, valid, other=0) ns = tl.program_id(1) * TILE_N + tl.arange(0, TILE_N) ks = tl.arange(0, K) - real = tl.zeros((R, TILE_N), tl.float32) - imag = tl.zeros((R, TILE_N), tl.float32) + real = tl.zeros((R, TILE_N), ACC) + imag = tl.zeros((R, TILE_N), ACC) for base in range(0, tl.max(ends - starts, 0), K): pos = starts[:, None] + base + ks[None, :] mask = valid[:, None] & (pos < ends[:, None]) @@ -32,8 +33,8 @@ def _kernel(V, I, P, B, C, M, N, BK, BN, CM, CN, real += tl.sum(vr[:, :, None] * br - vi[:, :, None] * bi, 1) imag += tl.sum(vr[:, :, None] * bi + vi[:, :, None] * br, 1) else: - v = tl.load(V + pos, mask, other=0.).to(tl.float32) - b = tl.load(B + dense_pos, dense_mask, other=0.).to(tl.float32) + v = tl.load(V + pos, mask, other=0.).to(ACC) + b = tl.load(B + dense_pos, dense_mask, other=0.).to(ACC) real += tl.sum(v[:, :, None] * b, 1) out = rows[:, None] * CM + ns[None, :] * CN store_mask = valid[:, None] & (ns[None, :] < N) @@ -53,10 +54,14 @@ def compute(values, cols, ptr, b, out): v, dense, target = (torch.view_as_real(t.resolve_conj()) for t in (values, b, out)) if is_complex else (values, b, out) mean = values.numel() / max(m, 1) batch, lanes = (8, 4) if mean <= 4 else (4, 8) + musa = _is_mthreads_runtime() + if musa: + batch, lanes = 8, 4 + acc = tl.int32 if musa and values.dtype == torch.int8 and out.dtype == torch.int32 else tl.float32 tile_n = min(32, triton.next_power_of_2(n)) _kernel[(triton.cdiv(m, batch), triton.cdiv(n, tile_n))]( v, cols, ptr, dense, target, m, n, b.stride(0), b.stride(1), out.stride(0), out.stride(1), is_complex, batch, lanes, tile_n, - num_warps=4, num_stages=1, + ACC=acc, num_warps=4, num_stages=1, ) return out diff --git a/src/flagsparse/sparse_operations/_spmv_csr_nnz.py b/src/flagsparse/sparse_operations/_spmv_csr_nnz.py index 742c8ae..1718c8f 100644 --- a/src/flagsparse/sparse_operations/_spmv_csr_nnz.py +++ b/src/flagsparse/sparse_operations/_spmv_csr_nnz.py @@ -12,20 +12,36 @@ def _segmented_add(left_row, left_value, right_row, right_value): return right_row, right_value + tl.where(left_row == right_row, left_value, 0.) +@triton.jit +def _musa_segmented_add(left_row, left_value, right_row, right_value): + return right_row, right_value + tl.where(left_row == right_row, left_value, 0) + + @triton.jit def _nnz_kernel(A, CI, RP, ROW, X, Y, EMPTY, NE: tl.constexpr, - NNZ: tl.constexpr, B: tl.constexpr, ACC: tl.constexpr): + NNZ: tl.constexpr, B: tl.constexpr, ACC: tl.constexpr, + MUSA: tl.constexpr = False, COMPLEX_RHS: tl.constexpr = False): pid = tl.program_id(0) lane = tl.arange(0, B) pos = pid * B + lane empty_row = tl.load(EMPTY + pos, pos < NE, 0) - tl.store(Y + empty_row, 0., pos < NE) + if COMPLEX_RHS: + tl.store(Y + 2 * empty_row, 0., pos < NE) + tl.store(Y + 2 * empty_row + 1, 0., pos < NE) + else: + tl.store(Y + empty_row, 0., pos < NE) valid = pos < NNZ row = tl.load(ROW + pos, valid, -1) col = tl.load(CI + pos, valid, 0) value = tl.load(A + pos, valid, 0).to(ACC) - vector = tl.load(X + col, valid, 0).to(ACC) - _, total = tl.associative_scan((row, value * vector), 0, _segmented_add) + vector = tl.load(X + (2 * col if COMPLEX_RHS else col), valid, 0).to(ACC) + if MUSA: + _, total = tl.associative_scan((row, value * vector), 0, _musa_segmented_add) + else: + _, total = tl.associative_scan((row, value * vector), 0, _segmented_add) + if COMPLEX_RHS: + vector_im = tl.load(X + 2 * col + 1, valid, 0).to(ACC) + _, total_im = tl.associative_scan((row, value * vector_im), 0, _musa_segmented_add) next_row = tl.load(ROW + pos + 1, pos + 1 < NNZ, -1) # Only the block containing a row's last nonzero owns its output. boundary = valid & (row != next_row) @@ -33,6 +49,8 @@ def _nnz_kernel(A, CI, RP, ROW, X, Y, EMPTY, NE: tl.constexpr, start = tl.load(RP + first_row) end = tl.load(RP + first_row + 1) prefix = tl.full((), 0., ACC) + if COMPLEX_RHS: + prefix_im = tl.full((), 0., ACC) if (start < pid * B) & (end <= (pid + 1) * B) & (pid * B < NNZ): # Recompute the short prefix crossing a block boundary, avoiding an # atomic update or a second launch to merge a partial-sum buffer. @@ -41,10 +59,21 @@ def _nnz_kernel(A, CI, RP, ROW, X, Y, EMPTY, NE: tl.constexpr, active = index < pid * B column = tl.load(CI + index, active, 0) a = tl.load(A + index, active, 0).to(ACC) - x = tl.load(X + column, active, 0).to(ACC) + x = tl.load(X + (2 * column if COMPLEX_RHS else column), active, 0).to(ACC) prefix += tl.sum(a * x, 0) - total += tl.where(row == first_row, prefix, 0.) - tl.store(Y + row, total, boundary) + if COMPLEX_RHS: + xi = tl.load(X + 2 * column + 1, active, 0).to(ACC) + prefix_im += tl.sum(a * xi, 0) + if MUSA: + total += tl.where(row == first_row, prefix, 0) + else: + total += tl.where(row == first_row, prefix, 0.) + if COMPLEX_RHS: + total_im += tl.where(row == first_row, prefix_im, 0) + tl.store(Y + 2 * row, total, boundary) + tl.store(Y + 2 * row + 1, total_im, boundary) + else: + tl.store(Y + row, total, boundary) def prepare(prepared): diff --git a/src/flagsparse/sparse_operations/_spmv_csr_transpose.py b/src/flagsparse/sparse_operations/_spmv_csr_transpose.py index a42be1e..e6d6a2e 100644 --- a/src/flagsparse/sparse_operations/_spmv_csr_transpose.py +++ b/src/flagsparse/sparse_operations/_spmv_csr_transpose.py @@ -6,7 +6,7 @@ import triton import triton.language as tl -from ._common import _gather_values, _is_maca_runtime +from ._common import _gather_values, _is_maca_runtime, _is_mthreads_runtime def prepare(indices, indptr, shape): @@ -19,8 +19,12 @@ def prepare(indices, indptr, shape): ) order = torch.argsort(indices, stable=True).to(indices.dtype) cols = source_rows[order] - counts = torch.bincount(indices.long(), minlength=n) - ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))).to(indptr.dtype) + if _is_mthreads_runtime(): + boundaries = torch.arange(n + 1, device=indices.device, dtype=indices.dtype) + ptr = torch.searchsorted(indices[order], boundaries).to(indptr.dtype) + else: + counts = torch.bincount(indices.long(), minlength=n) + ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))).to(indptr.dtype) return make_plan(cols, ptr, order, n) @@ -35,6 +39,9 @@ def make_plan(cols, ptr, order, n): # C550 favors a small one-wave tile over the CUDA 512-element tile. lanes = 32 if avg >= 32 else 16 return cols, ptr, order, row_order, lanes, 128 // lanes, {} + if _is_mthreads_runtime(): + lanes, batch = (4, 128) if avg <= 4 else (8, 64) if avg <= 24 else (8, 16) + return cols, ptr, order, row_order, lanes, batch, {} return cols, ptr, order, row_order, lanes, 512 // lanes, {} @@ -76,6 +83,15 @@ def _gather(V, I, P, ORDER, ROWS, X, Y, M: tl.constexpr, def gather(values, x, plan, n_rows, conjugate=False, out=None): + if _is_mthreads_runtime() and values.dtype == torch.float16 and plan[2] is None: + from .mixed_spmx import _musa_csr_spmv + result = _musa_csr_spmv(values.contiguous(), plan[0], plan[1], x.contiguous(), + n_rows, values.dtype if out is None else out.dtype, + torch.float32, out) + if out is not None and result is not out: + out.copy_(result) + return out + return result y = out if out is not None else torch.empty(n_rows, dtype=values.dtype, device=values.device) if not n_rows: return y @@ -108,6 +124,13 @@ def gather(values, x, plan, n_rows, conjugate=False, out=None): v, vector, target = cache["component_view"], torch.view_as_real(x), torch.view_as_real(y) else: v, vector, target = values, x, y + if _is_mthreads_runtime() and not complex_values and values.dtype == torch.float16: + from .mixed_spmx import _musa_csr_spmv + result = _musa_csr_spmv(values, cols, ptr, x, n_rows, + y.dtype, torch.float32, y) + if result is not y: + y.copy_(result) + return y if maca and values.dtype == torch.complex64: from ._spmv_csr_maca import complex_gather complex_gather[(triton.cdiv(n_rows, batch),)]( diff --git a/src/flagsparse/sparse_operations/gather_scatter.py b/src/flagsparse/sparse_operations/gather_scatter.py index b7b9946..2381033 100644 --- a/src/flagsparse/sparse_operations/gather_scatter.py +++ b/src/flagsparse/sparse_operations/gather_scatter.py @@ -272,6 +272,19 @@ def _triton_gather_impl( raise TypeError("out dtype must match gather output dtype") return out + if _is_mthreads_runtime() and dense_vector.dtype == torch.float16: + # Native indexing avoids the MUSA Triton launch overhead for short gathers. + if out is not None: + if out.shape != (nnz,): + raise ValueError("out shape must match gather output shape") + if out.dtype != dense_vector.dtype: + raise TypeError("out dtype must match gather output dtype") + dense_vector = dense_vector.contiguous() + if out is not None and not out.is_contiguous(): + out.copy_(torch.index_select(dense_vector, 0, kernel_indices)) + return out + return torch.index_select(dense_vector, 0, kernel_indices, out=out) + max_programs = _gather_max_programs(dense_vector.device, DEFAULT_GATHER_NUM_WARPS) grid = lambda meta: (min(max_programs, triton.cdiv(nnz, meta["BLOCK_SIZE"])),) @@ -333,6 +346,11 @@ def _triton_scatter_impl( index_fallback_policy = str(index_fallback_policy).lower() if index_fallback_policy not in ("auto", "strict"): raise ValueError("index_fallback_policy must be 'auto' or 'strict'") + native_scatter = ( + _is_mthreads_runtime() + and block_size is None + and sparse_values.dtype in (torch.float16, torch.int8) + ) block_size = _resolve_scatter_block_size( block_size, sparse_values.dtype, kernel_indices.numel() ) @@ -357,6 +375,13 @@ def _triton_scatter_impl( return dense_values, scatter_meta return dense_values + if native_scatter: + dense_values.index_copy_(0, kernel_indices.to(torch.int64), sparse_values) + scatter_meta["kernel_index_dtype"] = "int64" + if return_metadata: + return dense_values, scatter_meta + return dense_values + try: _launch_triton_scatter_kernel( dense_values, diff --git a/src/flagsparse/sparse_operations/mixed_spmx.py b/src/flagsparse/sparse_operations/mixed_spmx.py index adcd724..9726b42 100644 --- a/src/flagsparse/sparse_operations/mixed_spmx.py +++ b/src/flagsparse/sparse_operations/mixed_spmx.py @@ -46,11 +46,15 @@ import triton import triton.language as tl from ._rocm_sparse_cache import TensorPlanCache +from ._musa_spmv_launch import launch as _launch_musa_spmv _HALF = (torch.float16, torch.bfloat16) _TL_ACC = {torch.float32: tl.float32, torch.float64: tl.float64, torch.int32: tl.int32} _ROCM_CSR_MIXED_CACHE = {} _ROCM_COO_CSR_CACHE = TensorPlanCache() +_MUSA_COO_CSR_CACHE = TensorPlanCache() +_MUSA_CSC_CSR_CACHE = TensorPlanCache() +_MUSA_CSR_NNZ_CACHE = TensorPlanCache(maxsize=16) # value dtype -> widened output dtypes (the first is the default for int8). _WIDE_OUT = { @@ -144,7 +148,8 @@ def _csr_row_ids(indptr, n_rows, nnz): def _finish(y, out, t0, return_time): if out is not None: - out.copy_(y) + if not _is_mthreads_runtime() or y is not out: + out.copy_(y) y = out if return_time: _ACCEL.synchronize() @@ -211,6 +216,27 @@ def _csr_spmv_real_by_complex_kernel( tl.store(y_ri_ptr + 2 * rows + 1, tl.sum(acc_im, axis=1), mask=rmask) +@triton.jit +def _musa_csr_real_by_complex_kernel( + Y, V, C, P, X, M, R: tl.constexpr, K: tl.constexpr, ACC: tl.constexpr, +): + rows = tl.program_id(0) * R + tl.arange(0, R) + start = tl.load(P + rows, rows < M, other=0) + end = tl.load(P + rows + 1, rows < M, other=0) + lanes = tl.arange(0, K) + components = tl.arange(0, 2) + acc = tl.zeros((R, K, 2), ACC) + for offset in range(0, tl.max(end - start, 0), K): + pos = start[:, None] + offset + lanes[None, :] + valid = (rows[:, None] < M) & (pos < end[:, None]) + col = tl.load(C + pos, valid, other=0) + value = tl.load(V + pos, valid, other=0) + vector = tl.load(X + 2 * col[:, :, None] + components[None, None, :], + valid[:, :, None], other=0) + acc += value[:, :, None] * vector + tl.store(Y + 2 * rows[:, None] + components[None, :], tl.sum(acc, 1), rows[:, None] < M) + + def _csr_row_tile(nnz, n_rows): """(ROWS, BLOCK) from the mean row length: one BLOCK covers a typical row. @@ -229,6 +255,89 @@ def _csr_row_tile(nnz, n_rows): return 16, 32 +def _musa_row_tile(nnz, n_rows, integer=False): + mean = nnz / max(1, n_rows) + if mean <= 3: + return 256, 2 + if 4 < mean <= 12: + return (32, 4) if integer or n_rows >= 262144 else (64, 8) + if 12 < mean <= 24: + return (128, 8) if n_rows >= 262144 else (32, 4) + return 16, 8 + + +def _musa_coo_csr(data, row, col, shape): + """Reuse row ownership, invalidating both topology and values on mutation.""" + sources = (data, row, col) + options = tuple(shape) + state = _MUSA_COO_CSR_CACHE.get(sources, options) + if state is None: + _check_1d((data, "data"), (row, "row"), (col, "col")) + if not (data.numel() == row.numel() == col.numel()): + raise ValueError("data, row and col must have the same length") + rows, cols = _index32(row.contiguous()), _index32(col.contiguous()) + if rows.numel() > 1 and not bool(torch.all(rows[1:] >= rows[:-1]).item()): + order = torch.argsort(rows, stable=True) + rows = rows.index_select(0, order) + cols = cols.index_select(0, order) + data = data.index_select(0, order) + # MUSA bincount can lose updates for repeated sorted row IDs. + boundaries = torch.arange(int(shape[0]) + 1, device=rows.device, dtype=rows.dtype) + ptr = torch.searchsorted(rows, boundaries).to(torch.int32) + state = data.contiguous(), cols.contiguous(), ptr.contiguous() + _MUSA_COO_CSR_CACHE.put(sources, options, state) + return state + + +def _musa_csr_spmv(data, cols, ptr, x, n_rows, result_dtype, acc_dtype, out=None): + y = out if out is not None and out.is_contiguous() else torch.empty( + n_rows, dtype=result_dtype, device=x.device + ) + if not n_rows: + return y + if data.numel() / n_rows <= 4: + plan = _MUSA_CSR_NNZ_CACHE.get((ptr,), (n_rows,)) + if plan is None: + lengths = ptr[1:] - ptr[:-1] + if int(lengths.max().item()) > 128: + rows = torch.repeat_interleave( + torch.arange(n_rows, device=ptr.device, dtype=torch.int32), + lengths.long(), output_size=data.numel(), + ) + empty = torch.nonzero(lengths == 0).flatten().to(torch.int32) + plan = rows, empty + else: + plan = () + _MUSA_CSR_NNZ_CACHE.put((ptr,), (n_rows,), plan) + if plan: + from ._spmv_csr_nnz import _nnz_kernel + rows, empty = plan + complex_rhs = _is_complex_dtype(result_dtype) + _launch_musa_spmv( + _nnz_kernel, ((max(data.numel(), empty.numel()) + 511) // 512,), + (data, cols, ptr, rows, + torch.view_as_real(x) if complex_rhs else x, + torch.view_as_real(y) if complex_rhs else y, empty), + (), (empty.numel(), data.numel(), 512, _TL_ACC[acc_dtype], True, complex_rhs), + ) + return y + rows_per, block = _musa_row_tile(data.numel(), n_rows, data.dtype == torch.int8) + if _is_complex_dtype(result_dtype): + mean = data.numel() / n_rows + rows_per, block = (128, 2) if mean <= 3 else (32, 4) if mean <= 12 else (16, 8) + _musa_csr_real_by_complex_kernel[((n_rows + rows_per - 1) // rows_per,)]( + torch.view_as_real(y), data, cols, ptr, + torch.view_as_real(x), n_rows, rows_per, block, _TL_ACC[acc_dtype], + num_warps=4, num_stages=1, + ) + else: + _launch_musa_spmv( + _csr_spmv_mixed_kernel, ((n_rows + rows_per - 1) // rows_per,), + (y, data, cols, ptr, x), (n_rows,), (_TL_ACC[acc_dtype], rows_per, block), + ) + return y + + def _csr_mixed_num_warps(): """Return the mixed CSR SpMV launch width for ROCm tuning.""" if _is_rocm_runtime(): @@ -267,6 +376,9 @@ def spmv_csr_mixed(data, indices, indptr, x, shape, *, op=None, transpose=None, data, x = data.contiguous(), x.contiguous() cols, ptr = _index32(indices.contiguous()), _index32(indptr.contiguous()) t0 = _start(return_time) + if _is_mthreads_runtime(): + y = _musa_csr_spmv(data, cols, ptr, x, n_rows, result_dtype, acc_dtype, out) + return _finish(y, out, t0, return_time) if ( _is_rocm_runtime() and data.dtype == torch.float16 @@ -341,15 +453,26 @@ def _coo_spmv_atomic_kernel( def spmv_coo_mixed(data, row, col, x, shape, *, op=None, transpose=None, out=None, out_dtype=None, return_time=False): _normalize_non_op(op, transpose) - _check_1d((data, "data"), (row, "row"), (col, "col"), (x, "x")) + musa = _is_mthreads_runtime() + if musa: + # Cached sparse operands retain their metadata and mutation signature. + _check_1d((x, "x")) + else: + _check_1d((data, "data"), (row, "row"), (col, "col"), (x, "x")) n_rows, n_cols = int(shape[0]), int(shape[1]) - if not (data.numel() == row.numel() == col.numel()): + if not musa and not (data.numel() == row.numel() == col.numel()): raise ValueError("data, row and col must have the same length") if x.numel() != n_cols: raise ValueError("x shape does not match COO operation") result_dtype, acc_dtype = _resolve_types(data.dtype, x.dtype, out, out_dtype, False) if out is not None and (out.shape != (n_rows,) or out.device != x.device): raise ValueError("out shape/device must match the COO SpMV result") + if musa: + t0 = _start(return_time) + values, cols, ptr = _musa_coo_csr(data, row, col, shape) + y = _musa_csr_spmv(values, cols, ptr, x.contiguous(), n_rows, + result_dtype, acc_dtype, out) + return _finish(y, out, t0, return_time) sources = (data, row, col) data, x = data.contiguous(), x.contiguous() t0 = _start(return_time) @@ -432,6 +555,19 @@ def spmv_csc_half(data, indices, indptr, x, shape, *, out=None, out_dtype=None, raise TypeError("CSC half route requires float32 accumulation") if out is not None and (out.shape != (m,) or out.device != x.device): raise ValueError("out shape/device must match CSC SpMV result") + if _is_mthreads_runtime(): + sources = (data, indices, indptr) + state = _MUSA_CSC_CSR_CACHE.get(sources, (m, n)) + if state is None: + from . import _spmv_csr_transpose + cols, ptr, order, *_ = _spmv_csr_transpose.prepare(indices, indptr, (n, m)) + state = data.index_select(0, order.long()).contiguous(), cols, ptr + _MUSA_CSC_CSR_CACHE.put(sources, (m, n), state) + values, cols, ptr = state + return spmv_csr_mixed( + values, cols, ptr, x, shape, out=out, out_dtype=result_dtype, + return_time=return_time, + ) t0 = _start(return_time) acc = torch.zeros(m, dtype=torch.float32, device=x.device) if n and data.numel(): @@ -496,6 +632,13 @@ def spmm_csr_mixed(data, indices, indptr, B, shape, *, op=None, transpose=None, return _finish(C.to(result_dtype), out, t0, return_time) C = out if out is not None else torch.empty(n_rows, n_dense, dtype=result_dtype, device=B.device) if n_rows and n_dense: + if _is_mthreads_runtime(): + from . import _spmm_row_gather + _spmm_row_gather.compute(data, cols, ptr, B, C) + if return_time: + _ACCEL.synchronize() + return C, (time.perf_counter() - t0) * 1000.0 + return C block_n = min(128, max(16, triton.next_power_of_2(n_dense))) _csr_spmm_mixed_kernel[(n_rows, triton.cdiv(n_dense, block_n))]( C, data, cols, ptr, B, n_dense, @@ -571,6 +714,12 @@ def spmm_coo_mixed(data, row, col, B, shape, *, op=None, transpose=None, n_dense = int(B.shape[1]) if out is not None and (tuple(out.shape) != (n_rows, n_dense) or out.device != B.device): raise ValueError("out shape/device must match the COO SpMM result") + if _is_mthreads_runtime(): + values, cols, ptr = _musa_coo_csr(data, row, col, shape) + return spmm_csr_mixed( + values, cols, ptr, B, shape, out=out, out_dtype=result_dtype, + return_time=return_time, + ) data = data.contiguous() rows, cols = _index32(row.contiguous()), _index32(col.contiguous()) t0 = _start(return_time) diff --git a/src/flagsparse/sparse_operations/spmm_coo.py b/src/flagsparse/sparse_operations/spmm_coo.py index 9ca8e2b..7010e19 100644 --- a/src/flagsparse/sparse_operations/spmm_coo.py +++ b/src/flagsparse/sparse_operations/spmm_coo.py @@ -800,9 +800,14 @@ def __init__( self.op = str(op) self.alg = str(alg) self.gather_ptr = None - if _backend_name() == "cuda" and self.output_dtype == torch.complex64: - counts = torch.bincount(row.long(), minlength=self.n_rows) - self.gather_ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))).to(row.dtype) + if ((_backend_name() == "cuda" and self.output_dtype == torch.complex64) + or (_is_mthreads_runtime() and self.output_dtype == torch.float16)): + if _is_mthreads_runtime(): + boundaries = torch.arange(self.n_rows + 1, device=row.device, dtype=row.dtype) + self.gather_ptr = torch.searchsorted(row, boundaries).to(row.dtype) + else: + counts = torch.bincount(row.long(), minlength=self.n_rows) + self.gather_ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))).to(row.dtype) @triton.jit diff --git a/src/flagsparse/sparse_operations/spmm_csc.py b/src/flagsparse/sparse_operations/spmm_csc.py index 14fb002..d438e1e 100644 --- a/src/flagsparse/sparse_operations/spmm_csc.py +++ b/src/flagsparse/sparse_operations/spmm_csc.py @@ -45,6 +45,7 @@ SPMM_CSC_ALG_BASE = "spmm_csc_base" _SPMM_CSC_RESERVED_OP_MESSAGE = "spmm_csc_base supports op='non', 'trans', and 'conj'" _ROCM_CSC_PREPARED_CACHE = TensorPlanCache() +_MUSA_CSC_HALF_CACHE = TensorPlanCache() class SpmmCscAlgorithmUnavailable(RuntimeError): @@ -158,6 +159,7 @@ class PreparedCscSpmm: "index_fallback_applied", "index_fallback_reason", "rocm_non_gather_plan", + "musa_trans_gather", ) def __init__( @@ -180,6 +182,7 @@ def __init__( index_fallback_policy="auto", index_fallback_applied=False, index_fallback_reason=None, + musa_trans_gather=False, ): self.data = data self.kernel_indices = kernel_indices @@ -203,6 +206,7 @@ def __init__( self.index_fallback_applied = bool(index_fallback_applied) self.index_fallback_reason = index_fallback_reason self.rocm_non_gather_plan = None + self.musa_trans_gather = bool(musa_trans_gather) if (_is_rocm_runtime() and not _is_ascend_runtime() and data.dtype == torch.complex64 and self.op == "non"): from . import _spmv_csr_transpose @@ -608,6 +612,12 @@ def prepare_spmm_csc_route( op=_spmm_csc_op_to_name(op_code), alg=resolved_alg, index_fallback_policy=index_fallback_policy, + musa_trans_gather=( + _is_mthreads_runtime() + and _spmm_csc_op_transposes(op_code) + and data.dtype in (torch.float32, torch.complex64) + and block_n is None and block_nnz is None and max_segments is None + ), ) @@ -1175,6 +1185,17 @@ def _triton_spmm_csc_base_kernel(prepared, B, op_code=None): dtype = prepared.data.dtype n_dense_cols = int(B.shape[1]) out_rows = prepared.n_cols if transposes else prepared.n_rows + if prepared.musa_trans_gather and transposes: + from . import _spmm_row_gather + + C = torch.empty((out_rows, n_dense_cols), dtype=dtype, device=prepared.data.device) + values = prepared.data + if op_code == SPMM_CSC_OP_CONJ_TRANS: + values = values.conj() + return _spmm_row_gather.compute( + values.resolve_conj(), prepared.kernel_indices, prepared.kernel_indptr, + B.resolve_conj(), C, + ) plan = prepared.rocm_non_gather_plan if plan is not None and not transposes: C = torch.empty((out_rows, n_dense_cols), dtype=dtype, device=prepared.data.device) @@ -1517,10 +1538,14 @@ def _run_spmm_csc_base_route(prepared, B, *, timing=False, diagnostics=False): compute_ms = None gather_route = (prepared.rocm_non_gather_plan is not None and not _spmm_csc_op_transposes(prepared.op)) - block_n = (prepared.block_n if prepared.block_n_override else 32) if gather_route else ( - prepared.block_n if prepared.block_n_override - else _select_block_n(int(B.shape[1]), prepared.data.dtype, prepared.data.device) - ) + musa_gather = prepared.musa_trans_gather and _spmm_csc_op_transposes(prepared.op) + if musa_gather: + block_n = min(32, triton.next_power_of_2(int(B.shape[1]))) + elif gather_route: + block_n = prepared.block_n if prepared.block_n_override else 32 + else: + block_n = (prepared.block_n if prepared.block_n_override else + _select_block_n(int(B.shape[1]), prepared.data.dtype, prepared.data.device)) backend_info = _get_device_backend_info(prepared.data.device) if timing: _ACCEL.synchronize() @@ -1537,7 +1562,7 @@ def _run_spmm_csc_base_route(prepared, B, *, timing=False, diagnostics=False): "process_gpu_ms": 0.0 if timing else None, "compute_ms": compute_ms, "block_n": int(block_n), - "block_nnz": 8 if gather_route else int(prepared.block_nnz), + "block_nnz": 4 if musa_gather else 8 if gather_route else int(prepared.block_nnz), "launch_backend": backend_info["backend"], "device_warp_size": int(backend_info["device_warp_size"]), } @@ -1624,6 +1649,7 @@ def _spmm_csc_prepared_with_int32_indices(prepared, reason): index_fallback_policy=prepared.index_fallback_policy, index_fallback_applied=True, index_fallback_reason=str(reason), + musa_trans_gather=prepared.musa_trans_gather, ) @@ -1775,6 +1801,37 @@ def flagsparse_spmm_csc( # thing in float32 (the same "atomic_add on a float32 buffer, cast at the # end" shape COO already uses) and cast back once at this public boundary. if prepared is None and data is not None and data.dtype == torch.float16: + if ( + _is_mthreads_runtime() + and indices is not None and indptr is not None and shape is not None + and B is not None and B.dtype == torch.float16 + and not return_meta and not timing + and block_n is None and block_nnz is None and max_segments is None + and _normalize_spmm_csc_alg(alg) == "auto" + and _normalize_spmm_csc_op(op, bool(transpose)) == SPMM_CSC_OP_NON + and (transpose is None or not bool(transpose)) + and (out is None or out.dtype == torch.float16) + ): + from . import _spmv_csr_transpose, mixed_spmx + + _normalize_spmm_csc_index_fallback_policy(index_fallback_policy) + sources = (data, indices, indptr) + options = tuple(shape) + t0 = mixed_spmx._start(return_time) + state = _MUSA_CSC_HALF_CACHE.get(sources, options) + if state is None: + # Retain the raw CSC validation, but only redo topology work + # when a sparse operand's storage or mutation signature changes. + _, cols, ptr, m, n, *_ = _prepare_spmm_csc_matrix( + data.to(torch.float32), indices, indptr, shape + ) + csr_cols, csr_ptr, order, *_ = _spmv_csr_transpose.prepare(cols, ptr, (n, m)) + state = data.index_select(0, order.long()).contiguous(), csr_cols, csr_ptr + _MUSA_CSC_HALF_CACHE.put(sources, options, state) + if B.device != data.device: + raise ValueError("B must be on the same CUDA device as sparse matrix data") + C = mixed_spmx.spmm_csr_mixed(*state, B, shape, out=out) + return mixed_spmx._finish(C, out, t0, return_time) result = flagsparse_spmm_csc( data.to(torch.float32), indices, diff --git a/src/flagsparse/sparse_operations/spmm_csr.py b/src/flagsparse/sparse_operations/spmm_csr.py index 2cb3485..5d5f617 100644 --- a/src/flagsparse/sparse_operations/spmm_csr.py +++ b/src/flagsparse/sparse_operations/spmm_csr.py @@ -1057,8 +1057,11 @@ def _run_spmm_csr_base_route_impl( and prepared.avg_nnz_per_row <= 32.0 and prepared.max_row_nnz <= 256 ) - use_gather = (_backend_name() == "cuda" and not accuracy - and prepared.data.dtype == torch.complex64 and int(B.shape[1]) == 32) + use_gather = (not accuracy and ( + (_backend_name() == "cuda" and prepared.data.dtype == torch.complex64 + and int(B.shape[1]) == 32) + or (_is_mthreads_runtime() and prepared.data.dtype == torch.float16) + )) if use_gather: from . import _spmm_row_gather C = _spmm_row_gather.compute(prepared.data, prepared.kernel_indices, diff --git a/src/flagsparse/sparse_operations/spmv_coo.py b/src/flagsparse/sparse_operations/spmv_coo.py index 08c4580..bce2636 100644 --- a/src/flagsparse/sparse_operations/spmv_coo.py +++ b/src/flagsparse/sparse_operations/spmv_coo.py @@ -156,10 +156,15 @@ def __init__( ): from . import _spmv_csr_transpose - counts = torch.bincount(self.row_trans.long(), minlength=self.n_cols) - ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))).to( - self.row_trans.dtype - ) + if _is_mthreads_runtime(): + boundaries = torch.arange(self.n_cols + 1, device=self.row_trans.device, + dtype=self.row_trans.dtype) + ptr = torch.searchsorted(self.row_trans, boundaries).to(self.row_trans.dtype) + else: + counts = torch.bincount(self.row_trans.long(), minlength=self.n_cols) + ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))).to( + self.row_trans.dtype + ) self.transpose_csr_plan = _spmv_csr_transpose.make_plan( self.col_trans, ptr, None, self.n_cols ) diff --git a/src/flagsparse/sparse_operations/spmv_csc.py b/src/flagsparse/sparse_operations/spmv_csc.py index 264eecd..f4d5911 100644 --- a/src/flagsparse/sparse_operations/spmv_csc.py +++ b/src/flagsparse/sparse_operations/spmv_csc.py @@ -849,7 +849,7 @@ def _spmv_csc_via_mixed(data, indices, indptr, x, shape, op_code, out, out_dtype data, indices, indptr, x, (n_cols, n_rows), out=out, out_dtype=out_dtype, return_time=timed, ) - elif (_SPMV_CSC_CUDA or _is_maca_runtime()) and data.dtype == torch.float16: + elif (_SPMV_CSC_CUDA or _is_maca_runtime() or _is_mthreads_runtime()) and data.dtype == torch.float16: y = _mixed_spmx.spmv_csc_half( data, indices, indptr, x, shape, out=out, out_dtype=out_dtype, return_time=timed) else: diff --git a/src/flagsparse/sparse_operations/spmv_csr.py b/src/flagsparse/sparse_operations/spmv_csr.py index 880b658..ef5962b 100644 --- a/src/flagsparse/sparse_operations/spmv_csr.py +++ b/src/flagsparse/sparse_operations/spmv_csr.py @@ -1058,6 +1058,7 @@ def prepare_spmv_csr( elif (not transpose and requested_alg == "auto" and config is None and data.dtype in (torch.float16, torch.complex64) and (_backend_name() == "cuda" or + (_is_mthreads_runtime() and data.dtype == torch.float16) or (_is_maca_runtime() and (data.dtype == torch.float16 or data.numel() > 4 * n_rows)))): from . import _spmv_csr_transpose diff --git a/src/flagsparse/sparse_operations/spmv_sell.py b/src/flagsparse/sparse_operations/spmv_sell.py index 888abcc..03b1c2b 100644 --- a/src/flagsparse/sparse_operations/spmv_sell.py +++ b/src/flagsparse/sparse_operations/spmv_sell.py @@ -38,6 +38,7 @@ import triton.language as tl from ._rocm_sparse_cache import TensorPlanCache +from ._musa_spmv_launch import launch as _launch_musa_spmv SPMV_SELL_VALUE_DTYPES = ( torch.float16, @@ -74,6 +75,27 @@ def _sell_acc_dtype(value_dtype, out_dtype): torch.int32: tl.int32, } _ROCM_SELL_CSR_CACHE = TensorPlanCache(maxsize=16) +_MUSA_SELL_PLAN_CACHE = TensorPlanCache(maxsize=16) + + +def _musa_sell_csr_plan(values, cols, offsets, n_rows, slice_size): + """Compact heavily padded real SELL without retaining computed outputs.""" + rows = _sell_entry_rows(cols, offsets, slice_size) + positions = torch.nonzero((cols >= 0) & (rows < n_rows)).flatten() + plan = () + if values.numel() > max(positions.numel(), 1) * 2: + rows = rows.index_select(0, positions) + order = torch.argsort(rows, stable=True) + rows = rows.index_select(0, order) + positions = positions.index_select(0, order) + boundaries = torch.arange(n_rows + 1, device=rows.device, dtype=rows.dtype) + ptr = torch.searchsorted(rows, boundaries).to(cols.dtype) + plan = ( + values.index_select(0, positions), + cols.index_select(0, positions), + ptr, + ) + return plan def _rocm_sell_csr_plan(values, cols, offsets, n_rows, slice_size): @@ -183,6 +205,45 @@ def _spmv_sell_complex_kernel( tl.store(y_ri_ptr + 2 * row + 1, acc_im, mask=mask) +@triton.jit +def _musa_spmv_sell_kernel( + Y, V, C, P, X, M, + SLICE: tl.constexpr, ROWS: tl.constexpr, SLOTS: tl.constexpr, + COMPLEX: tl.constexpr, ACC: tl.constexpr, +): + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + slices = rows // SLICE + start = tl.load(P + slices, rows < M, other=0) + stop = tl.load(P + slices + 1, rows < M, other=0) + width = (stop - start) // SLICE + base = start + rows % SLICE + slots = tl.arange(0, SLOTS) + real = tl.zeros((ROWS, SLOTS), ACC) + imag = tl.zeros((ROWS, SLOTS), ACC) + for offset in range(0, tl.max(width, 0), SLOTS): + j = offset + slots + pos = base[:, None] + j[None, :] * SLICE + active = (rows[:, None] < M) & (j[None, :] < width[:, None]) + col = tl.load(C + pos, active, other=-1) + valid = active & (col >= 0) + if COMPLEX: + vr = tl.load(V + 2 * pos, valid, other=0).to(ACC) + vi = tl.load(V + 2 * pos + 1, valid, other=0).to(ACC) + xr = tl.load(X + 2 * col, valid, other=0).to(ACC) + xi = tl.load(X + 2 * col + 1, valid, other=0).to(ACC) + real += vr * xr - vi * xi + imag += vr * xi + vi * xr + else: + value = tl.load(V + pos, valid, other=0).to(ACC) + vector = tl.load(X + col, valid, other=0).to(ACC) + real += value * vector + if COMPLEX: + tl.store(Y + 2 * rows, tl.sum(real, 1), rows < M) + tl.store(Y + 2 * rows + 1, tl.sum(imag, 1), rows < M) + else: + tl.store(Y + rows, tl.sum(real, 1), rows < M) + + @triton.jit def _spmv_sell_complex_parallel_kernel( Y, V, C, P, X, M, @@ -288,6 +349,45 @@ def _validate_spmv_sell(values, cols, offsets, x, shape, slice_size, validate): return n_rows, n_cols, slice_size, n_slices +def _musa_sell_plan(values, cols, offsets, x, shape, slice_size, validate): + sources = (values, cols, offsets) + options = (int(shape[0]), int(shape[1]), int(slice_size)) + state = _MUSA_SELL_PLAN_CACHE.get(sources, options) + if state is None: + dims = _validate_spmv_sell(values, cols, offsets, x, shape, slice_size, validate) + state = {"dims": dims, "validated": bool(validate), "csr": None, + "arrays": (values.contiguous(), cols.contiguous(), offsets.contiguous())} + _MUSA_SELL_PLAN_CACHE.put(sources, options, state) + elif validate and not state["validated"]: + _validate_spmv_sell(values, cols, offsets, x, shape, slice_size, True) + state["validated"] = True + else: + if not torch.is_tensor(x): + raise TypeError("x must be a torch.Tensor") + if x.ndim != 1: + raise ValueError("x must be a 1D tensor") + if not _is_accel_tensor(x): + raise ValueError("x must be an accelerator tensor") + if x.device != values.device: + raise ValueError("SELL SpMV inputs must be on one device") + if x.dtype != values.dtype: + raise TypeError("x dtype must match values dtype") + if x.numel() != state["dims"][1]: + raise ValueError(f"x length {x.numel()} does not match n_cols {state['dims'][1]}") + return state + + +def _musa_sell_tile(n_values, n_rows, n_slices, slice_size, complex_values): + mean_width = n_values / max(1, n_slices * slice_size) + if mean_width <= 4.5: + return 128, 1 + if 9 < mean_width <= 16 and n_rows <= 524288: + return 64, 2 + if not complex_values and 16 < mean_width <= 24 and n_rows >= 262144: + return 128, 1 + return 32, 4 + + def flagsparse_spmv_sell( values, col_indices, @@ -310,9 +410,25 @@ def flagsparse_spmv_sell( """ if str(op).strip().lower() not in ("non", "n", "non_trans"): raise NotImplementedError("SELL SpMV supports op='non' only") - n_rows, _n_cols, slice_size, n_slices = _validate_spmv_sell( - values, col_indices, slice_offsets, x, shape, slice_size, validate + musa = _is_mthreads_runtime() + cached_musa = ( + musa and torch.is_tensor(values) + and values.dtype in ( + torch.float16, torch.bfloat16, torch.float32, torch.int8, torch.complex64 + ) ) + t0 = None + if cached_musa: + if return_time: + _ACCEL.synchronize() + t0 = time.perf_counter() + state = _musa_sell_plan(values, col_indices, slice_offsets, x, shape, slice_size, validate) + n_rows, _n_cols, slice_size, n_slices = state["dims"] + values, cols, offsets = state["arrays"] + else: + n_rows, _n_cols, slice_size, n_slices = _validate_spmv_sell( + values, col_indices, slice_offsets, x, shape, slice_size, validate + ) allowed = _SPMV_SELL_OUT_DTYPES[values.dtype] if out is not None: if out_dtype is not None and out.dtype != out_dtype: @@ -325,16 +441,51 @@ def flagsparse_spmv_sell( names = ", ".join(str(d).replace("torch.", "") for d in allowed) raise TypeError(f"{values.dtype} SELL SpMV can write: {names}") acc_dtype = _sell_acc_dtype(values.dtype, out_dtype) - values = values.contiguous() - cols = col_indices.contiguous() - offsets = slice_offsets.contiguous() + if not cached_musa: + values = values.contiguous() + cols = col_indices.contiguous() + offsets = slice_offsets.contiguous() x = x.contiguous() - t0 = None - if return_time: + if return_time and not cached_musa: _ACCEL.synchronize() t0 = time.perf_counter() - if _is_ascend_runtime(): + if musa: + y = out if out is not None and out.is_contiguous() else torch.empty( + n_rows, dtype=out_dtype, device=x.device + ) + plan = () + if cached_musa: + if state["csr"] is None: + state["csr"] = () + if ( + n_slices and not _is_complex_dtype(values.dtype) + and values.numel() > n_rows * 4.5 + ): + state["csr"] = _musa_sell_csr_plan(values, cols, offsets, n_rows, slice_size) + plan = state["csr"] + if plan: + from .mixed_spmx import _musa_csr_spmv + + _musa_csr_spmv(*plan, x, n_rows, out_dtype, acc_dtype, y) + elif n_slices: + complex_values = _is_complex_dtype(values.dtype) + rows, slots = _musa_sell_tile( + values.numel(), n_rows, n_slices, slice_size, complex_values + ) + target, matrix, vector = ( + (torch.view_as_real(t) for t in (y, values, x)) + if complex_values else (y, values, x) + ) + _launch_musa_spmv( + _musa_spmv_sell_kernel, ((n_rows + rows - 1) // rows,), + (target, matrix, cols, offsets, vector), (n_rows,), + (slice_size, rows, slots, complex_values, + tl.float64 if acc_dtype == torch.complex128 else ( + tl.float32 if complex_values else _TL[acc_dtype] + )), + ) + elif _is_ascend_runtime(): y = _torch_spmv_sell(values, cols, offsets, x, n_rows, slice_size, acc_dtype) y = y if y.dtype == out_dtype else y.to(out_dtype) else: @@ -444,7 +595,7 @@ def csr_to_sell(data, indices, indptr, n_rows, slice_size, index_dtype=None): ) slot = torch.arange(nnz, device=device, dtype=torch.int64) - ptr[rows] pos = offsets[rows // slice_size] + slot * slice_size + rows % slice_size - if _is_ascend_runtime() and data.is_complex(): + if (_is_ascend_runtime() or _is_mthreads_runtime()) and data.is_complex(): # torch_npu's IndexPut implementation rejects complex tensors. Store # the two real planes instead; IndexCopy supports the corresponding # floating dtype and view_as_complex restores the SELL value layout. diff --git a/src/flagsparse/sparse_operations/vector_ops.py b/src/flagsparse/sparse_operations/vector_ops.py index 7a2937f..c09ada1 100644 --- a/src/flagsparse/sparse_operations/vector_ops.py +++ b/src/flagsparse/sparse_operations/vector_ops.py @@ -37,6 +37,8 @@ import triton import triton.language as tl +from ._rocm_sparse_cache import TensorPlanCache + AXPBY_VALUE_DTYPES = ( torch.float16, torch.bfloat16, @@ -47,6 +49,7 @@ ) SPVV_VALUE_DTYPES = AXPBY_VALUE_DTYPES + (torch.int8,) _VECTOR_BLOCK = 1024 +_MUSA_AXPBY_INDEX_CACHE = TensorPlanCache() def _vector_compute_dtype(value_dtype): @@ -154,6 +157,41 @@ def _axpy_complex_kernel( tl.store(y_ri_ptr + 2 * idx + 1, yi + alpha_re * xi + alpha_im * xr, mask=mask) +@triton.jit +def _musa_axpby_inverse_kernel(inverse, indices, nnz, n, BLOCK: tl.constexpr): + pos = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + row = tl.load(indices + pos, pos < nnz, -1) + tl.store(inverse + row, pos, (pos < nnz) & (row >= 0) & (row < n)) + + +@triton.jit +def _musa_axpby_kernel(y, values, inverse, n, alpha, beta, BLOCK: tl.constexpr): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + pos = tl.load(inverse + row, row < n, -1) + value = tl.load(values + pos, (row < n) & (pos >= 0), 0).to(tl.float32) + previous = tl.load(y + row, row < n, 0).to(tl.float32) + tl.store(y + row, beta * previous + alpha * value, row < n) + + +def _musa_axpby(values, indices, y, alpha, beta): + n = y.numel() + if not n: + return + inverse = _MUSA_AXPBY_INDEX_CACHE.get((indices,), (n,)) + if inverse is None: + inverse = torch.full((n,), -1, dtype=torch.int32, device=y.device) + nnz = indices.numel() + if nnz: + _musa_axpby_inverse_kernel[((nnz + 1023) // 1024,)]( + inverse, indices.contiguous(), nnz, n, BLOCK=1024, + num_warps=4, + ) + _MUSA_AXPBY_INDEX_CACHE.put((indices,), (n,), inverse) + _musa_axpby_kernel[((n + 2047) // 2048,)]( + y, values, inverse, n, alpha, beta, BLOCK=2048, num_warps=4, + ) + + def _torch_axpby(values, indices, y, alpha, beta): """Ascend path, also the reference shape of the Triton path.""" compute = _vector_compute_dtype(y.dtype) @@ -192,11 +230,15 @@ def flagsparse_axpby( if return_time: _ACCEL.synchronize() t0 = time.perf_counter() - # On gfx936, the native half-precision indexed add is substantially faster - # than launching separate fp32 scale and scatter kernels. AXPBY's indices - # are required to be unique, so doing the update in the requested dtype is - # deterministic and matches the documented half output contract. - if _is_rocm_runtime() and y.dtype in (torch.float16, torch.bfloat16): + if ( + _is_mthreads_runtime() and y.dtype in (torch.float16, torch.bfloat16) + and y.numel() <= 4 * 1024 * 1024 + ): + # Cache index ownership only; every call reads the current values/y. + _musa_axpby(values, indices, y, a_re, b_re) + elif _is_rocm_runtime() and y.dtype in (torch.float16, torch.bfloat16): + # Native half indexed add avoids a second launch on gfx936. Indices + # are unique, as required by the sparse-vector API. y.mul_(b_re) if values.numel(): y.index_add_(0, kernel_idx.to(torch.int64), values * a_re) diff --git a/tests/pytest/test_musa_sparse_routes.py b/tests/pytest/test_musa_sparse_routes.py new file mode 100644 index 0000000..7e7ffab --- /dev/null +++ b/tests/pytest/test_musa_sparse_routes.py @@ -0,0 +1,797 @@ +# Copyright 2026 FlagOS Contributors +# SPDX-License-Identifier: Apache-2.0 +"""MUSA row-owned routes must read live values and invalidate changed topology.""" + +import pytest +import torch + +import flagsparse as fs +from flagsparse.sparse_operations._common import _is_mthreads_runtime +from tests.pytest.accuracy_utils import accelerator_available, accelerator_device + +pytestmark = pytest.mark.skipif( + not accelerator_available() or not _is_mthreads_runtime(), reason="requires MUSA" +) + + +def test_musa_coo_transpose_repeated_row_ownership(): + gen = torch.Generator().manual_seed(123) + rows = torch.randint(0, 257, (4930,), generator=gen).int() + cols = torch.randint(0, 129, (4930,), generator=gen).int() + values = torch.randn(4930, generator=gen) + rhs = torch.randn(257, generator=gen) + expected = torch.zeros(129, dtype=torch.float64) + expected.index_add_(0, cols.long(), values.double() * rhs.double()[rows.long()]) + dev = accelerator_device() + for _ in range(30): + actual = fs.flagsparse_spmv_coo( + values.to(dev), rows.to(dev), cols.to(dev), rhs.to(dev), + (257, 129), op="trans", + ) + torch.testing.assert_close(actual.cpu().double(), expected, rtol=1e-5, atol=1e-5) + + +def test_musa_nnz_integer_cross_block_large_sum(): + dev = accelerator_device() + length, n_rows = 1101, 512 + ptr = torch.full((n_rows + 1,), length, dtype=torch.int32, device=dev) + ptr[0] = 0 + cols = torch.zeros(length, dtype=torch.int32, device=dev) + values = torch.full((length,), 127, dtype=torch.int8, device=dev) + rhs = torch.tensor([127], dtype=torch.int8, device=dev) + actual = fs.flagsparse_spmv_csr(values, cols, ptr, rhs, (n_rows, 1)) + expected = torch.zeros(n_rows, dtype=torch.int32) + expected[0] = length * 127 * 127 + assert torch.equal(actual.cpu(), expected) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_musa_nnz_real_by_complex_live_cross_block(dtype): + dev = accelerator_device() + ptr = torch.tensor([0, 601, 700] + [700] * 254, dtype=torch.int32, device=dev) + cols = (torch.arange(700, device=dev) % 7).int() + values = (torch.arange(700, device=dev) % 5 - 2).to(dtype) + rhs = torch.complex(torch.arange(7, device=dev).to(dtype), + torch.arange(7, device=dev).to(dtype) - 2) + for _ in range(3): + host = torch.sparse_csr_tensor(ptr.cpu().long(), cols.cpu().long(), + values.cpu().to(torch.complex128), (256, 7)) + expected = host @ rhs.cpu().to(torch.complex128) + actual = fs.flagsparse_spmv_csr(values, cols, ptr, rhs, (256, 7)) + torch.testing.assert_close(actual.cpu().to(torch.complex128), expected, + atol=1e-5, rtol=1e-5) + values.add_(1) + rhs.mul_(1j) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("shape", [(7, 11, 2), (71, 31, 9), (5, 33, 31)]) +def test_musa_real_by_complex_row_tiles(dtype, shape): + m, n, width = shape + dev = accelerator_device() + values = torch.arange(m * width, device=dev).to(dtype).div_(13) + cols = (torch.arange(m * width, device=dev) % n).int() + ptr = torch.arange(m + 1, dtype=torch.int32, device=dev) * width + rhs = torch.complex(torch.arange(n, device=dev).to(dtype).div_(11), + torch.arange(n, device=dev).to(dtype).div_(7)) + host = torch.sparse_csr_tensor(ptr.cpu().long(), cols.cpu().long(), + values.cpu().to(torch.complex128), (m, n)) + expected = host @ rhs.cpu().to(torch.complex128) + actual = fs.flagsparse_spmv_csr(values, cols, ptr, rhs, (m, n)) + tolerance = 1e-5 if dtype == torch.float32 else 1e-10 + torch.testing.assert_close(actual.cpu().to(torch.complex128), expected, + rtol=tolerance, atol=tolerance) + + +@pytest.mark.parametrize("indirect", [False, True]) +def test_musa_half_gather_strided_output(indirect): + from flagsparse.sparse_operations import _spmv_csr_transpose as gather + dev = accelerator_device() + values = torch.tensor([1, 2, 3], dtype=torch.float16, device=dev) + cols = torch.tensor([0, 2, 1], dtype=torch.int32, device=dev) + ptr = torch.tensor([0, 2, 3, 3], dtype=torch.int32, device=dev) + order = torch.tensor([2, 1, 0], dtype=torch.int32, device=dev) if indirect else None + if indirect: + values = values.flip(0) + rhs = torch.tensor([2, 3, 4], dtype=torch.float16, device=dev) + backing = torch.full((6,), -99, dtype=torch.float16, device=dev) + out = backing[::2] + plan = gather.make_plan(cols, ptr, order, 3) + actual = gather.gather(values, rhs, plan, 3, out=out) + assert actual is out + torch.testing.assert_close(actual.cpu(), torch.tensor([10, 9, 0], dtype=torch.float16)) + assert torch.equal(backing[1::2].cpu(), torch.full((3,), -99, dtype=torch.float16)) + + +@pytest.mark.parametrize("fmt", ["coo", "csc"]) +@pytest.mark.parametrize("matrix_rhs", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) +def test_musa_mixed_live_values_and_indices(fmt, matrix_rhs, dtype): + if fmt == "csc" and dtype == torch.int8: + pytest.skip("CSC int8 is not supported") + dev = accelerator_device() + rows = torch.tensor([2, 0, 2, 0], dtype=torch.int32, device=dev) + cols = torch.tensor([1, 0, 2, 0], dtype=torch.int32, device=dev) + values = torch.tensor([2, 1, 3, -1], dtype=dtype, device=dev) + rhs = torch.ones((3, 7) if matrix_rhs else (3,), dtype=dtype, device=dev) + out_dtype = torch.int32 if dtype == torch.int8 else torch.float32 + if fmt == "csc": + rows = torch.tensor([0, 0, 2, 2], dtype=torch.int32, device=dev) + cols = torch.tensor([0, 2, 3, 4], dtype=torch.int32, device=dev) + + def run(): + kwargs = {"out_dtype": out_dtype} + if fmt == "csc": + if matrix_rhs: + return fs.flagsparse_spmm_csc(values, rows, cols, rhs, (3, 3)) + return fs.flagsparse_spmv_csc(values, rows, cols, rhs, (3, 3), **kwargs) + op = fs.flagsparse_spmm_coo if matrix_rhs else fs.flagsparse_spmv_coo + return op(values, rows, cols, rhs, (3, 3), **kwargs) + + def check(): + a = torch.zeros(3, 3, dtype=torch.float64) + rr, cc, vv = rows.cpu().long(), cols.cpu().long(), values.cpu().double() + if fmt == "csc": + cc = torch.repeat_interleave(torch.arange(3), cc[1:] - cc[:-1]) + a.index_put_((rr, cc), vv, accumulate=True) + expected = a @ rhs.cpu().double() + torch.testing.assert_close(run().cpu().double(), expected, rtol=0, atol=1e-5) + + check() + check() + values.add_(1) + check() + rows[0] = 1 + check() + rhs.mul_(2) + check() + if fmt == "csc": + cols[1] = 1 + check() + with torch.inference_mode(): + rows, cols, values = rows.clone(), cols.clone(), values.clone() + check() + values.add_(1) + cols[1] = 2 + check() + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_musa_csc_half_spmm_strides_and_out(index_dtype): + dev = accelerator_device() + values = torch.tensor([1, 3, 2], dtype=torch.float16, device=dev) + indices = torch.tensor([0, 2, 0], dtype=index_dtype, device=dev) + indptr = torch.tensor([0, 2, 2, 3], dtype=index_dtype, device=dev) + rhs = torch.arange(42, dtype=torch.float16, device=dev).reshape(7, 6).T[::2] + backing = torch.full((3, 14), -99, dtype=torch.float16, device=dev) + out = backing[:, ::2] + expected_a = torch.tensor([[1, 0, 2], [0, 0, 0], [3, 0, 0]], dtype=torch.float32) + for _ in range(2): + actual = fs.flagsparse_spmm_csc(values, indices, indptr, rhs, (3, 3), out=out) + assert actual is out + expected = (expected_a @ rhs.cpu().float()).half() + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + assert torch.equal(backing[:, 1::2].cpu(), torch.full((3, 7), -99, dtype=torch.float16)) + rhs.add_(1) + actual, elapsed = fs.flagsparse_spmm_csc(values, indices, indptr, rhs, (3, 3), return_time=True) + torch.testing.assert_close(actual.cpu(), (expected_a @ rhs.cpu().float()).half()) + assert elapsed >= 0 + with pytest.raises(ValueError, match="transpose conflicts"): + fs.flagsparse_spmm_csc(values, indices, indptr, rhs, (3, 3), op="non", transpose=True) + with pytest.raises(ValueError, match="policy"): + fs.flagsparse_spmm_csc(values, indices, indptr, rhs, (3, 3), index_fallback_policy="invalid") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) +@pytest.mark.parametrize("reset", [False, True]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_musa_scatter_native_strides_and_live_inputs(dtype, reset, index_dtype): + from flagsparse.sparse_operations.gather_scatter import _triton_scatter_impl + + dev = accelerator_device() + values = torch.tensor([1, 99, 2, 99, 3, 99], dtype=dtype, device=dev)[::2] + indices = torch.tensor([4, 99, 1, 99, 3, 99], dtype=index_dtype, device=dev)[::2] + backing = torch.full((12,), -9, dtype=dtype, device=dev) + out = backing[::2] + for _ in range(3): + expected = out.cpu() + if reset: + expected.zero_() + expected.index_copy_(0, indices.cpu().long(), values.cpu()) + actual, meta = _triton_scatter_impl(values, indices, 6, out=out, + reset_output=reset, return_metadata=True) + assert actual is out + assert not meta["index_fallback_applied"] + assert torch.equal(actual.cpu(), expected) + assert torch.equal(backing[1::2].cpu(), torch.full((6,), -9, dtype=dtype)) + values.add_(1) + indices[0] = (int(indices[0]) + 1) % 3 + + +@pytest.mark.parametrize("musa,block", [(False, None), (True, None), (True, 256)]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) +def test_musa_scatter_native_scope_and_explicit_tuning(monkeypatch, musa, block, dtype): + from flagsparse.sparse_operations import gather_scatter as scatter + + values = torch.tensor([1, 2], dtype=dtype) + indices = torch.tensor([2, 0], dtype=torch.int32) + out = torch.zeros(3, dtype=dtype) + launches = [] + + def launch(dense, sparse, index, nnz, block_size=None): + launches.append(block_size) + dense.index_copy_(0, index.long(), sparse) + + monkeypatch.setattr(scatter, "_is_mthreads_runtime", lambda: musa) + monkeypatch.setattr(scatter, "_launch_triton_scatter_kernel", launch) + result, meta = scatter._triton_scatter_impl(values, indices, 3, out=out, + block_size=block, return_metadata=True) + native = musa and block is None + assert result is out + assert bool(launches) != native + assert meta["kernel_index_dtype"] == ("int64" if native else "int32") + assert not meta["index_fallback_applied"] + assert torch.equal(result, torch.tensor([2, 0, 1], dtype=dtype)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64, + torch.complex64, torch.complex128, torch.int8]) +def test_musa_sell_parallel_partial_slice_and_padding(dtype): + dense = torch.tensor([[1, 0, 2], [0, 0, 0], [-2, 3, 4]], dtype=dtype) + csr = dense.to_sparse_csr() + dev = accelerator_device() + values, cols, ptr = fs.csr_to_sell( + csr.values().to(dev), csr.col_indices().int().to(dev), + csr.crow_indices().int().to(dev), 3, 2, + ) + x = torch.tensor([1, 2, -1], dtype=dtype, device=dev) + got = fs.flagsparse_spmv_sell(values, cols, ptr, x, (3, 3), slice_size=2) + torch.testing.assert_close(got.cpu().to(torch.complex128), + dense.to(torch.complex128) @ x.cpu().to(torch.complex128)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.int8]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_musa_sell_compact_live_inputs_and_strides(dtype, index_dtype): + from flagsparse.sparse_operations import spmv_sell as sell + + dev = accelerator_device() + m, n, slice_size = 19, 11, 8 + dense = torch.zeros(m, n, dtype=dtype) + dense[0] = torch.arange(n).to(dtype) + dense[9, 3] = 2 + dense[18, 7] = -1 + csr = dense.to_sparse_csr() + values, cols, offsets = fs.csr_to_sell( + csr.values().to(dev), csr.col_indices().to(device=dev, dtype=index_dtype), + csr.crow_indices().to(device=dev, dtype=index_dtype), m, slice_size, + ) + rhs = (torch.arange(2 * n, device=dev) % 7).to(dtype)[::2] + out_dtype = torch.int32 if dtype == torch.int8 else dtype + backing = torch.full((2 * m,), -99, device=dev, dtype=out_dtype) + out = backing[::2] + + def check(): + host = torch.zeros(m, n, dtype=torch.float64) + rr = sell._sell_entry_rows(cols, offsets, slice_size).cpu() + cc = cols.cpu().long() + valid = (cc >= 0) & (rr < m) + host.index_put_((rr[valid], cc[valid]), values.cpu().double()[valid], accumulate=True) + result = fs.flagsparse_spmv_sell(values, cols, offsets, rhs, (m, n), + slice_size=slice_size, out=out) + assert result is out + torch.testing.assert_close(result.cpu().double(), host @ rhs.cpu().double(), + atol=0, rtol=0) + assert torch.equal(backing[1::2].cpu(), torch.full((m,), -99, dtype=out_dtype)) + + check() + assert sell._MUSA_SELL_PLAN_CACHE.get((values, cols, offsets), (m, n, slice_size))["csr"] + check() + values.add_(1) + check() + cols[0] = 2 + check() + offsets[1] -= slice_size + check() + rhs.add_(1) + check() + timed, elapsed = fs.flagsparse_spmv_sell(values, cols, offsets, rhs, (m, n), + slice_size=slice_size, return_time=True) + torch.testing.assert_close(timed, out) + assert elapsed >= 0 + with torch.inference_mode(): + values, cols, offsets = values.clone(), cols.clone(), offsets.clone() + check() + values.add_(1) + cols[0] = 1 + check() + cols = cols.clone() + cols[0] = n + with pytest.raises(IndexError, match="out of range"): + fs.flagsparse_spmv_sell(values, cols, offsets, rhs, (m, n), slice_size=slice_size) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.complex64]) +@pytest.mark.parametrize("op", ["trans", "conj"]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_musa_csc_transposed_row_gather(dtype, op, index_dtype): + from flagsparse.sparse_operations import spmm_csc as csc + + dev = accelerator_device() + values = torch.tensor([1, 3, 2], dtype=dtype, device=dev) + if dtype.is_complex: + values.add_(1j) + indices = torch.tensor([0, 2, 0], dtype=index_dtype, device=dev) + indptr = torch.tensor([0, 2, 2, 3], dtype=index_dtype, device=dev) + rhs = torch.arange(42, device=dev).to(dtype).reshape(7, 6).T[::2] + if dtype.is_complex: + rhs = (rhs + 2j).conj() + backing = torch.full((3, 14), -99, dtype=dtype, device=dev) + out = backing[:, ::2] + prepared = csc.prepare_spmm_csc_route(values, indices, indptr, (3, 3), op=op) + assert prepared.musa_trans_gather + for _ in range(2): + dense = torch.zeros(3, 3, dtype=torch.complex128) + dense[0, 0], dense[2, 0], dense[0, 2] = values.cpu().to(torch.complex128) + matrix = dense.mH if op == "conj" else dense.T + expected = matrix @ rhs.cpu().to(torch.complex128) + actual, meta = fs.flagsparse_spmm_csc(prepared=prepared, B=rhs, out=out, + return_meta=True) + assert actual is out + torch.testing.assert_close(actual.cpu().to(torch.complex128), expected) + assert meta["block_n"] == 8 and meta["block_nnz"] == 4 + torch.testing.assert_close(backing[:, 1::2].cpu(), torch.full((3, 7), -99, dtype=dtype)) + values.add_(1) + rhs = rhs + 1 + explicit = csc.prepare_spmm_csc_route(values, indices, indptr, (3, 3), op=op, + block_nnz=4) + assert not explicit.musa_trans_gather + named = csc.prepare_spmm_csc_route(values, indices, indptr, (3, 3), op=op, + alg="spmm_csc_base") + assert named.musa_trans_gather + + +def test_musa_sell_cached_validation(): + dev = accelerator_device() + values = torch.tensor([1, 2, 3, 4], dtype=torch.float32, device=dev) + cols = torch.tensor([0, 1, 0, 1], dtype=torch.int32, device=dev) + offsets = torch.tensor([0, 4], dtype=torch.int32, device=dev) + rhs = torch.ones(2, dtype=torch.float32, device=dev) + kwargs = dict(shape=(2, 2), slice_size=2) + fs.flagsparse_spmv_sell(values, cols, offsets, rhs, validate=False, **kwargs) + fs.flagsparse_spmv_sell(values, cols, offsets, rhs, **kwargs) + with pytest.raises(ValueError, match="x length"): + fs.flagsparse_spmv_sell(values, cols, offsets, rhs[:1], **kwargs) + with pytest.raises(TypeError, match="dtype"): + fs.flagsparse_spmv_sell(values, cols, offsets, rhs.half(), **kwargs) + cols[0] = 2 + with pytest.raises(IndexError, match="out of range"): + fs.flagsparse_spmv_sell(values, cols, offsets, rhs, **kwargs) + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_musa_sell_complex_cached_live_inputs(index_dtype): + from flagsparse.sparse_operations import spmv_sell as sell + + dev = accelerator_device() + m, n, slice_size = 5, 4, 2 + dense = torch.tensor([[1 + 2j, 0, 3 - 1j, 0], [0, 0, 0, 0], + [0, 2j, 0, -1], [3, 0, 0, 0], [0, 0, 1j, 0]]) + csr = dense.to_sparse_csr() + values, cols, offsets = fs.csr_to_sell( + csr.values().to(dev), csr.col_indices().to(device=dev, dtype=index_dtype), + csr.crow_indices().to(device=dev, dtype=index_dtype), m, slice_size, + ) + backing_values = torch.empty(values.numel() * 2, dtype=values.dtype, device=dev) + backing_values[::2] = values + values = backing_values[::2] + rhs = torch.tensor([1 + 1j, 2 - 1j, 3, -2j], device=dev) + backing_out = torch.full((2 * m,), -99, dtype=torch.complex64, device=dev) + out = backing_out[::2] + + def check(validate=True): + rr = sell._sell_entry_rows(cols, offsets, slice_size).cpu() + cc = cols.cpu().long() + valid = (rr < m) & (cc >= 0) + host = torch.zeros(m, n, dtype=torch.complex128) + host.index_put_((rr[valid], cc[valid]), + values.cpu().to(torch.complex128)[valid], accumulate=True) + actual = fs.flagsparse_spmv_sell(values, cols, offsets, rhs, (m, n), + slice_size=slice_size, out=out, validate=validate) + assert actual is out + torch.testing.assert_close(actual.cpu().to(torch.complex128), + host @ rhs.cpu().to(torch.complex128)) + torch.testing.assert_close(backing_out[1::2].cpu(), + torch.full((m,), -99, dtype=torch.complex64)) + + check(False) + check() + state = sell._MUSA_SELL_PLAN_CACHE.get((values, cols, offsets), (m, n, slice_size)) + assert state is not None and state["validated"] and not state["csr"] + values.add_(1j) + check() + values.data = (values + 2).data + check() + cols[0] = 1 + check() + offsets[1] -= slice_size + check() + rhs.mul_(1j) + check() + with torch.inference_mode(): + values, cols, offsets = values.clone(), cols.clone(), offsets.clone() + check() + values.add_(1) + check() + cols = cols.clone() + cols[0] = n + with pytest.raises(IndexError, match="out of range"): + fs.flagsparse_spmv_sell(values, cols, offsets, rhs, (m, n), slice_size=slice_size) + + +@pytest.mark.parametrize("sorted_rows", [False, True]) +def test_musa_coo_cache_live_storage(sorted_rows): + dev = accelerator_device() + rows = torch.tensor([0, 0, 2, 2] if sorted_rows else [2, 0, 2, 0], + dtype=torch.int32, device=dev) + cols = torch.tensor([0, 1, 2, 1], dtype=torch.int32, device=dev) + values = torch.tensor([1, 2, 3, 4], dtype=torch.float16, device=dev) + rhs = torch.tensor([2, 3, 4], dtype=torch.float16, device=dev) + + def check(): + expected = torch.zeros(3, dtype=torch.float64) + expected.index_add_(0, rows.cpu().long(), + values.cpu().double() * rhs.cpu().double()[cols.cpu().long()]) + result = fs.flagsparse_spmv_coo(values, rows, cols, rhs, (3, 3), + out_dtype=torch.float32) + torch.testing.assert_close(result.cpu().double(), expected, rtol=0, atol=0) + + check() + values.add_(1) + check() + values.data = (values + 2).data + check() + cols.data = ((cols + 1) % 3).data + check() + rows[0] = 1 + check() + with torch.inference_mode(): + values, rows, cols = values.clone(), rows.clone(), cols.clone() + check() + rows[0] = 0 + values.add_(1) + check() + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) +def test_musa_nnz_route_mutation_and_strided_output(dtype): + dev = accelerator_device() + ptr = torch.tensor([0, 161, 200] + [200] * 62, dtype=torch.int32, device=dev) + cols = (torch.arange(200, device=dev) % 7).int() + values = (torch.arange(200, device=dev) % 5 - 2).to(dtype) + rhs = torch.arange(7, device=dev).to(dtype) + out_dtype = torch.int32 if dtype == torch.int8 else torch.float32 + out = torch.full((128,), -99, dtype=out_dtype, device=dev)[::2] + + def check(): + host = torch.sparse_csr_tensor(ptr.cpu().long(), cols.cpu().long(), + values.cpu().double(), (64, 7)) + actual = fs.flagsparse_spmv_csr(values, cols, ptr, rhs, (64, 7), + out=out, out_dtype=out_dtype) + assert actual is out + torch.testing.assert_close(actual.cpu().double(), host @ rhs.cpu().double(), + atol=1e-5, rtol=0) + + check() + check() + values.add_(1) + rhs.mul_(2) + check() + ptr[1] = 140 + cols[0] = 6 + check() + with torch.inference_mode(): + ptr = ptr.clone() + check() + ptr[1] = 150 + check() + + +def test_musa_half_gather_strides_and_live_input(): + from flagsparse.sparse_operations.gather_scatter import _triton_gather_impl + dev = accelerator_device() + dense = torch.arange(16, device=dev, dtype=torch.float16)[::2] + indices = torch.tensor([7, 1, 1, 0], dtype=torch.int32, device=dev) + out = torch.empty(8, device=dev, dtype=torch.float16)[::2] + assert _triton_gather_impl(dense, indices, out) is out + torch.testing.assert_close(out.cpu(), torch.tensor([14, 2, 2, 0], dtype=torch.float16)) + dense.add_(1) + _triton_gather_impl(dense, indices, out) + torch.testing.assert_close(out.cpu(), torch.tensor([15, 3, 3, 1], dtype=torch.float16)) + with pytest.raises(ValueError, match="shape"): + _triton_gather_impl(dense, indices, out[:1]) + with pytest.raises(TypeError, match="dtype"): + _triton_gather_impl(dense, indices, out.float()) + + +@pytest.mark.parametrize("musa", [False, True]) +def test_musa_scope_output_self_copy(monkeypatch, musa): + from flagsparse.sparse_operations import mixed_spmx + + class Output: + copies = 0 + + def copy_(self, other): + assert other is self + self.copies += 1 + + monkeypatch.setattr(mixed_spmx, "_is_mthreads_runtime", lambda: musa) + out = Output() + assert mixed_spmx._finish(out, out, None, False) is out + assert out.copies == (0 if musa else 1) + + +@pytest.mark.parametrize("musa", [False, True]) +@pytest.mark.parametrize("integer", [False, True]) +def test_musa_scope_row_gather_configuration(monkeypatch, musa, integer): + import triton.language as tl + from flagsparse.sparse_operations import _spmm_row_gather + + launches = [] + + class Kernel: + def __getitem__(self, grid): + def launch(*args, **kwargs): + launches.append((args, kwargs)) + return launch + + monkeypatch.setattr(_spmm_row_gather, "_is_mthreads_runtime", lambda: musa) + monkeypatch.setattr(_spmm_row_gather, "_kernel", Kernel()) + dtype = torch.int8 if integer else torch.float16 + values = torch.ones(64, dtype=dtype) + cols = torch.zeros(64, dtype=torch.int32) + ptr = torch.tensor([0, 16, 32, 48, 64], dtype=torch.int32) + rhs = torch.ones((1, 7), dtype=dtype) + out = torch.empty((4, 7), dtype=torch.int32 if integer else torch.float32) + assert _spmm_row_gather.compute(values, cols, ptr, rhs, out) is out + args, kwargs = launches[0] + assert args[-3:] == ((8, 4, 8) if musa else (4, 8, 8)) + assert kwargs["ACC"] == (tl.int32 if musa and integer else tl.float32) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("inference", [False, True]) +def test_musa_axpby_fused_live_inputs(dtype, index_dtype, inference): + from flagsparse.sparse_operations import vector_ops + + dev = accelerator_device() + with torch.inference_mode(inference): + indices = torch.tensor([3, 99, 0, 99, 6, 99], dtype=index_dtype, device=dev)[::2] + values = torch.tensor([1, 99, -2, 99, 3, 99], dtype=dtype, device=dev)[::2] + y = torch.arange(9, dtype=dtype, device=dev) + for iteration in range(4): + expected = y.cpu().float().mul_(-0.25) + expected.index_add_(0, indices.cpu().long(), values.cpu().float() * 1.5) + actual, elapsed = fs.flagsparse_axpby(values, indices, y, 1.5, -0.25, return_time=True) + assert actual is y + assert elapsed >= 0 + torch.testing.assert_close(y.cpu(), expected.to(dtype), rtol=0, atol=0) + values.add_(0.5) + indices[0] = (iteration + 1) % 3 + 3 + if not inference: + indices.data = indices.clone() + expected = y.cpu().float().mul_(0.5) + expected.index_add_(0, indices.cpu().long(), values.cpu().float() * -2) + fs.flagsparse_axpby(values, indices, y, -2, 0.5, validate=False) + torch.testing.assert_close(y.cpu(), expected.to(dtype), rtol=0, atol=0) + assert len(vector_ops._MUSA_AXPBY_INDEX_CACHE._entries) <= 8 + + +@pytest.mark.parametrize("size", [0, 19]) +def test_musa_axpby_fused_empty_sparse_vector(size): + dev = accelerator_device() + y = torch.arange(size, dtype=torch.float16, device=dev) + expected = y.cpu() * 0.5 + got = fs.flagsparse_axpby(torch.empty(0, dtype=y.dtype, device=dev), + torch.empty(0, dtype=torch.int32, device=dev), y, 2, 0.5) + assert got is y + torch.testing.assert_close(got.cpu(), expected) + + +@pytest.mark.parametrize("musa", [False, True]) +def test_musa_axpby_fused_backend_scope(monkeypatch, musa): + from flagsparse.sparse_operations import vector_ops + + values = torch.tensor([2], dtype=torch.float16) + indices = torch.tensor([1], dtype=torch.int32) + y = torch.ones(3, dtype=torch.float16) + calls = [] + + class Kernel: + def __getitem__(self, grid): + return lambda *args, **kwargs: calls.append("generic") + + monkeypatch.setattr(vector_ops, "_is_mthreads_runtime", lambda: musa) + monkeypatch.setattr(vector_ops, "_is_rocm_runtime", lambda: False) + monkeypatch.setattr(vector_ops, "_is_ascend_runtime", lambda: False) + monkeypatch.setattr(vector_ops, "_is_accel_tensor", lambda _: True) + monkeypatch.setattr(vector_ops, "_musa_axpby", lambda *args: calls.append("musa")) + monkeypatch.setattr(vector_ops, "_scale_real_kernel", Kernel()) + monkeypatch.setattr(vector_ops, "_axpy_real_kernel", Kernel()) + fs.flagsparse_axpby(values, indices, y, 1.5, 0.5, validate=False) + assert calls == (["musa"] if musa else ["generic", "generic"]) + + +@pytest.fixture +def musa_launch_cache(monkeypatch): + from collections import OrderedDict + from flagsparse.sparse_operations import _musa_spmv_launch as cache + + monkeypatch.setattr(cache, "_LAUNCHERS", OrderedDict()) + + class Compiled: + def __getitem__(self, grid): + def run(*args): + kernel.calls.append(("compiled", grid, args)) + return run + + class Kernel: + debug = False + pre_run_hooks = () + + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def run(*args, **kwargs): + assert kwargs == {"num_warps": 4, "num_stages": 1} + self.calls.append(("jit", grid, args)) + return Compiled() + return run + + kernel = Kernel() + return cache, kernel + + +def test_musa_launch_cache_specializations(musa_launch_cache): + from types import SimpleNamespace + + cache, kernel = musa_launch_cache + + def tensor(address, dtype=torch.float16, device="musa:0"): + return SimpleNamespace(dtype=dtype, device=torch.device(device), + data_ptr=lambda: address) + + a, b = tensor(256), tensor(512) + cache.launch(kernel, (2,), (a, b), (17,), (32,)) + cache.launch(kernel, (2, 1, 1), (tensor(768), tensor(1024)), (17,), (32,)) + assert [call[0] for call in kernel.calls] == ["jit", "compiled"] + + for tensors, scalars, constants, grid in [ + ((tensor(258), b), (17,), (32,), (2,)), + ((tensor(256, torch.float32), b), (17,), (32,), (2,)), + ((a, tensor(512, device="musa:1")), (17,), (32,), (2,)), + ((a, a), (17,), (32,), (2,)), + ((a, b), (18,), (32,), (2,)), + ((a, b), (17,), (64,), (2,)), + ((a, b), (17,), (32,), (3,)), + ]: + cache.launch(kernel, grid, tensors, scalars, constants) + assert kernel.calls[-1][0] == "jit" + kernel.debug = True + cache.launch(kernel, (2,), (a, b), (17,), (32,)) + assert kernel.calls[-1][0] == "jit" + assert len(cache._LAUNCHERS) == 9 + + +def test_musa_launch_cache_lru_eviction(musa_launch_cache, monkeypatch): + cache, kernel = musa_launch_cache + monkeypatch.setattr(cache, "_MAX_LAUNCHERS", 2) + tensors = (torch.empty(4),) + for size in [1, 2, 1, 3, 2]: + cache.launch(kernel, (size,), tensors, (), ()) + assert len(cache._LAUNCHERS) <= 2 + assert [call[0] for call in kernel.calls] == [ + "jit", "jit", "compiled", "jit", "jit", + ] + + +@pytest.mark.parametrize("mode", ["pre_run", "debug", "instrumentation", "stages"]) +def test_musa_launch_cache_debug_fallback(musa_launch_cache, monkeypatch, mode): + from triton import knobs + + cache, kernel = musa_launch_cache + tensors = (torch.empty(4),) + cache.launch(kernel, (1,), tensors, (), ()) + if mode == "pre_run": + monkeypatch.setattr(kernel, "pre_run_hooks", (lambda: None,)) + elif mode == "debug": + monkeypatch.setattr(knobs.runtime, "debug", True) + elif mode == "instrumentation": + monkeypatch.setattr(knobs.compilation, "instrumentation_mode", "test") + else: + monkeypatch.setattr(knobs.runtime, "add_stages_inspection_hook", lambda: None) + cache.launch(kernel, (1,), tensors, (), ()) + assert [call[0] for call in kernel.calls] == ["jit", "jit"] + assert len(cache._LAUNCHERS) == 1 + + +@pytest.mark.parametrize("fmt", ["coo", "sell"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) +def test_musa_cached_launch_live_buffers_and_stream(fmt, dtype): + from flagsparse.sparse_operations._common import _ACCEL + + dev = accelerator_device() + dense = torch.tensor([[1, 0, 2], [0, 0, 0], [0, -2, 0]], dtype=dtype) + rows = torch.tensor([0, 0, 2], dtype=torch.int32, device=dev) + cols = torch.tensor([0, 2, 1], dtype=torch.int32, device=dev) + values = torch.tensor([1, 2, -2], dtype=dtype, device=dev) + sell_values, sell_cols, offsets = (t.to(dev) for t in fs.dense_to_sell(dense, 2)) + outputs = (torch.float16, torch.float32) if dtype == torch.float16 else ( + torch.int32, torch.float32 + ) + original = None + + def run(rhs, out_dtype, out=None): + if fmt == "coo": + return fs.flagsparse_spmv_coo(values, rows, cols, rhs, (3, 3), + out=out, out_dtype=out_dtype) + return fs.flagsparse_spmv_sell(sell_values, sell_cols, offsets, rhs, (3, 3), + slice_size=2, out=out, out_dtype=out_dtype) + + stream = _ACCEL.Stream() + stream.wait_stream(_ACCEL.current_stream()) + for out_dtype in outputs: + for offset in (0, 1): + with _ACCEL.stream(stream): + rhs = torch.empty(3 + offset, dtype=dtype, device=dev)[offset:] + rhs.copy_(torch.tensor([2, 3, 4], dtype=dtype, device=dev)) + first = run(rhs, out_dtype) + assert original is None or first.data_ptr() != original.data_ptr() + original = first + rhs.add_(1) + out = torch.empty(3 + offset, dtype=out_dtype, device=dev)[offset:] + actual = run(rhs, out_dtype, out) + assert actual is out + # Exercise a cached launch on the default stream as well. + stream.synchronize() + second = run(rhs, out_dtype) + _ACCEL.synchronize() + torch.testing.assert_close(first.cpu().double(), + dense.double() @ torch.tensor([2., 3., 4.]).double(), + rtol=0, atol=0) + expected = dense.double() @ torch.tensor([3., 4., 5.]).double() + torch.testing.assert_close(actual.cpu().double(), expected, rtol=0, atol=0) + torch.testing.assert_close(second.cpu().double(), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("fmt", ["csr", "sell"]) +@pytest.mark.parametrize("musa", [False, True]) +def test_musa_cached_launch_backend_scope(monkeypatch, fmt, musa): + from flagsparse.sparse_operations import mixed_spmx, spmv_sell + + module = mixed_spmx if fmt == "csr" else spmv_sell + calls = [] + + class Kernel: + def __getitem__(self, grid): + return lambda *args, **kwargs: calls.append("generic") + + for name in ("_is_rocm_runtime", "_is_ascend_runtime", "_is_maca_runtime"): + monkeypatch.setattr(module, name, lambda: False) + monkeypatch.setattr(module, "_is_mthreads_runtime", lambda: musa) + monkeypatch.setattr(module, "_is_accel_tensor", lambda _: True) + monkeypatch.setattr(module, "_launch_musa_spmv", lambda *args: calls.append("musa")) + values = torch.ones(2, dtype=torch.float16) + indices = torch.tensor([0, 1], dtype=torch.int32) + offsets = torch.tensor([0, 2], dtype=torch.int32) + rhs = torch.ones(2, dtype=torch.float16) + if fmt == "csr": + monkeypatch.setattr(module, "_csr_spmv_mixed_kernel", Kernel()) + module.spmv_csr_mixed(values, indices, offsets, rhs, (1, 2), + out_dtype=torch.float32) + else: + monkeypatch.setattr(module, "_spmv_sell_real_kernel", Kernel()) + module.flagsparse_spmv_sell(values, indices, offsets, rhs, (1, 2), slice_size=1) + assert calls == (["musa"] if musa else ["generic"]) diff --git a/tests/pytest/test_spmv_csr_accuracy.py b/tests/pytest/test_spmv_csr_accuracy.py index fbccb69..97c6850 100644 --- a/tests/pytest/test_spmv_csr_accuracy.py +++ b/tests/pytest/test_spmv_csr_accuracy.py @@ -678,7 +678,11 @@ def test_spmv_csr_empty_dtype_ops(alg, dtype, op, shape): prepared = _new_prepared(data, col, ptr, shape, alg, op=op) actual = spmv_mod.flagsparse_spmv_csr_run(prepared, x) assert actual.shape == (shape[0] if op == "non" else shape[1],) - assert torch.count_nonzero(actual) == 0 + checked = ( + torch.view_as_real(actual) + if common._is_mthreads_runtime() and actual.is_complex() else actual + ) + assert torch.count_nonzero(checked) == 0 @pytest.mark.spmv_csr From 3cea4bc3615655dabcf676850de7d770d7da795a Mon Sep 17 00:00:00 2001 From: NCIC-AlphaSparse Date: Sat, 17 Oct 2026 22:50:41 +0800 Subject: [PATCH 3/4] musa opt 1missed --- capi/ctest/accuracy/test_spgemm.cpp | 11 ++++++++ capi/ctest/accuracy/test_spmv.cpp | 42 +++++++++++++++++++++++++++++ capi/ctest/accuracy/test_spvv.cpp | 34 ++++++++++++++++++++++- 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/capi/ctest/accuracy/test_spgemm.cpp b/capi/ctest/accuracy/test_spgemm.cpp index c552b04..f97f785 100644 --- a/capi/ctest/accuracy/test_spgemm.cpp +++ b/capi/ctest/accuracy/test_spgemm.cpp @@ -26,6 +26,7 @@ #include #include +#include "adaptor/adaptor.hpp" #include "common.hpp" using namespace fstest; @@ -213,6 +214,16 @@ TEST_F(SpGEMMAccuracy, Float32) { run_and_check(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(A, B, FLAGSPARSE_R_32F, "musa_multiblock_fp32"); + run_and_check(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) { diff --git a/capi/ctest/accuracy/test_spmv.cpp b/capi/ctest/accuracy/test_spmv.cpp index e98f35e..58eb528 100644 --- a/capi/ctest/accuracy/test_spmv.cpp +++ b/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; + const std::vector values{C(0.07468984f, -0.07468984f), + C(-37.027058f, 37.027058f), + C(37.027058f, -37.027058f)}; + const std::vector ptr{0, 1, 2, 3}, indices{0, 0, 0}; + const std::vector 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(1)[0]; + const double expected = static_cast(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}) { diff --git a/capi/ctest/accuracy/test_spvv.cpp b/capi/ctest/accuracy/test_spvv.cpp index 265364a..4159452 100644 --- a/capi/ctest/accuracy/test_spvv.cpp +++ b/capi/ctest/accuracy/test_spvv.cpp @@ -9,6 +9,10 @@ #include "common.hpp" +#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS) +#include +#endif + using namespace fstest; namespace { @@ -23,10 +27,12 @@ template Out run_spvv(flagsparseHandle_t handle, flagsparseDataType_t input_type, flagsparseDataType_t output_type, const std::vector& sparse, const std::vector& indices, const std::vector& 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; + const std::vector indices{0, 2}; + const std::vector dense{C(3, -2), C(0, 0), C(4, 1)}; + for (flagsparseStream_t selected : { + static_cast(nullptr), + reinterpret_cast(stream.value)}) { + ASSERT_EQ(flagsparseSetStream(handle.h, selected), FLAGSPARSE_STATUS_SUCCESS); + for (int repeat = 0; repeat < 20; ++repeat) { + const std::vector sparse{C(2 + repeat, 1), C(-1, 3)}; + EXPECT_EQ((run_spvv(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 From 20b3f1929881dda74e00d5b6579dbebdadb5cb54 Mon Sep 17 00:00:00 2001 From: NCIC-AlphaSparse Date: Mon, 19 Oct 2026 10:20:20 +0800 Subject: [PATCH 4/4] musa opt2 --- capi/ctest/accuracy/test_scatter.cpp | 73 +++++++ docs/debug/MUSA.md | 132 ++++++++++++ .../sparse_operations/_musa_spmv_launch.py | 106 ++++++++-- .../sparse_operations/gather_scatter.py | 193 +++++++++++++++--- .../sparse_operations/mixed_spmx.py | 111 +++++++++- tests/pytest/test_musa_sparse_routes.py | 158 +++++++++++++- 6 files changed, 711 insertions(+), 62 deletions(-) diff --git a/capi/ctest/accuracy/test_scatter.cpp b/capi/ctest/accuracy/test_scatter.cpp index 8b5ce2c..933f7c2 100644 --- a/capi/ctest/accuracy/test_scatter.cpp +++ b/capi/ctest/accuracy/test_scatter.cpp @@ -25,6 +25,10 @@ #include "common.hpp" +#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS) +#include +#endif + using namespace fstest; namespace { @@ -222,4 +226,73 @@ TEST_F(ScatterAccuracy, SizeMismatchIsRejected) { flagsparseDestroySpVec(vecX); } +#if defined(FLAGSPARSE_MUSA_BASELINE_EXTENSIONS) +template +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 dense(n * bytes), poison(nnz * bytes, 0xa5); + std::vector initial(n * bytes, 0x5a); + std::vector 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(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 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(gathered.size()), gathered); + EXPECT_EQ(d_out.download(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(nullptr), + reinterpret_cast(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(handle.h, dtype, FLAGSPARSE_INDEX_32I, bytes); + musa_live_copies(handle.h, dtype, FLAGSPARSE_INDEX_64I, bytes); + } + } + EXPECT_EQ(flagsparseSetStream(handle.h, nullptr), FLAGSPARSE_STATUS_SUCCESS); +} +#endif + } // namespace diff --git a/docs/debug/MUSA.md b/docs/debug/MUSA.md index 77c02fc..e354951 100644 --- a/docs/debug/MUSA.md +++ b/docs/debug/MUSA.md @@ -3,6 +3,8 @@ `FLAGSPARSE_BACKEND=mthreads`。环境搭建、交付复现见 [`../MUSA.md`](../MUSA.md)。背景见本目录 [`README.md`](README.md)。 +**MUSA 实现约束（用户于 2026-10-18 明确）**：算子计算内核必须使用 Triton，允许使用当前后端支持且经实测验证的 TLE 特性。禁止以手写 MUSA C++/`.mu` 内核、mcc 编译的计算库或 PyTorch 原生索引替代被测算子的计算内核。C API 主机封装、设备内存/流管理和厂商对照基线不属于计算内核；其他后端保持原实现。本文后半部保留的原生实验为撤回历史，不能用于全 Triton 验收。 + ## 要做的事 环境按 `../MUSA.md` 第 2 节配好（`FLAGSPARSE_BACKEND=mthreads`、`MUSA_HOME`）后，跑一轮 split： @@ -637,3 +639,133 @@ setsid env PYTHONPATH=$PWD/src FLAGSPARSE_BACKEND=mthreads MUSA_HOME=/usr/local/ --capi-bench-out results_musa_split_65_20261017_opt3/capi \ > results_musa_split_65_20261017_opt3.log 2>&1 < /dev/null & ``` + +### 2026-10-18 上游同步后的复测与优化口径 + +> 此节及下方原生实现记录为已撤回的实验历史，不代表当前全 Triton 实现。 +> 用户随后明确禁止原生 MUSA 内核；原生定向结果与 opt4 均不用于全 Triton 验收。 + +- 基准提交为 `3cea4bc3615655dabcf676850de7d770d7da795a`。此前优化已被上游吸收；旧 stash、完整轮结果和打包文件均保留。本轮开始先重建 C API，恢复源码与二进制一致。 +- 最近完成的完整轮仍是 `results_musa_split_65_20261017_opt3/`，65/65 精度 Passed、0 missing，但下表 7 个变体未过 0.8x。这里继续按用户要求优化**每个变体各自的矩阵平均值**，不是整个交付集的平均值，也不要求每个矩阵单独大于 0.8x。 +- 性能输入仍为 `tests/data` 的原 10 个矩阵，名称与字节数再次核对一致。固定三个 CPU 线程池为 8，GPU 测量串行；未改 runner、测试规模、迭代次数、容差、跳过规则或基线优先级。 +- 最终源码重复测量如下。COO 三项取原 q4 benchmark（5 warmup / 20 iters）的公开 API 耗时，对照旧完整轮同变体同矩阵的 `h800_scaled_ms`；**属于 H800 带宽折算估算，不是本机厂商实测**。Gather/Scatter 取原 C API benchmark（10 warmup / 100 iters）的本机 muSPARSE 基线。 + +| 变体 | 旧完整轮 | 最终定向复测 | 矩阵精度 | 基线 | +|---|---:|---:|---|---| +| `spmv_coo_f16f32_int_non` | 0.726342x | 0.884297x | 10/10 PASS | H800 带宽折算 | +| `spmv_coo_f16_int_non` | 0.790823x | 0.936304x | 10/10 PASS | H800 带宽折算 | +| `spmv_coo_i8i32_int_non` | 0.772409x | 0.944477x | 10/10 PASS | H800 带宽折算 | +| `gather_f32_int` | 0.794171x | 0.987362x | 10/10 pass | 本机 muSPARSE | +| `scatter_f64_int` | 0.794720x | 0.996308x | 10/10 pass | 本机 muSPARSE | +| `scatter_c32_int` | 0.791277x | 0.995981x | 10/10 pass | 本机 muSPARSE | +| `scatter_c64_int` | 0.795244x | 0.997719x | 10/10 pass | 本机 muSPARSE | + +原始数据：`results_musa_opt_20261018_followup_v1/coo_native_final.csv`、`capi_native_final/*_benchmark.json`；COO 全部 6 个 q4 变体共 60 条均 PASS。早期探针和集成中间轮完整保留，不能选取最好一次替代最终数据。 + +### 2026-10-18 src/flagsparse/sparse_operations/_musa_spmv_native.py::try_spmv/_load +- 改了什么：增加可选 MUSA 原生行子组 SpMV，FP16 在 FP32 中累加，int8 按请求在 int32/FP32 中累加；只在实测过的 mp_31 架构使用。按平均行长选择 2/8 lanes，平均行长 (3,4] 的不规则短行保留原 NNZ 路线，不按矩阵文件名选路。每次分配或使用调用方输出，实时读取 RHS；只缓存编译代码，不缓存结果。首次通过 mcc 编译到 Triton 原子代码缓存；编译器缺失或构建失败回退 Triton。每次解析当前流，原生包装保留/恢复当前设备；调试、插桩或启动钩子存在时回退正常 Triton 路径。 +- 为什么改（附报错原文）：无算子错误；COO FP16 到 FP32 在旧完整轮和拉取后的原规模复测中持续约 0.73x，单纯调整 tile/启动器不足以达标。原生行子组减少动态行块归约及启动开销，配对测量确认有收益。集成时曾把 Triton 空 HookChain 误判为存在钩子，现按其 calls 判断；中间轮原始记录保留。 +- 怎么验证的：原 10 矩阵公开 API 配对探针、原 q4 benchmark 两次正式入口复测；最终三项均值见上表。新增回归明确拒绝意外回退，并检查部分块、空行、非连续输入输出、值/索引/存储修改、推理张量及调试/钩子/编译器不可用回退；原非默认流回归同样通过。 +- 其他后端是否也需要：否；模块仅在 MUSA COO 分支惰性导入。 + +### 2026-10-18 src/flagsparse/sparse_operations/mixed_spmx.py::spmv_coo_mixed +- 改了什么：仅 MUSA COO 路由先尝试上述原生内核，不适用时继续调用原 `_musa_csr_spmv`；保留原有版本/存储感知 COO 转 CSR 缓存、输出和计时语义。CSR、CSC、SpMM 和非 MUSA 路由不变。 +- 为什么改（附报错原文）：无报错；接入经实测确认有收益的 MUSA COO 实现。 +- 怎么验证的：最终 COO 60 条 q4 精度全部 PASS；扩大 Python 回归 1199 passed、1678 skipped（沿用原适用性规则，20.12 秒），日志和 JUnit 为 `native_accuracy_final.log/.xml`。 +- 其他后端是否也需要：否。 + +### 2026-10-18 capi/src/adaptor/backend/musa/gather_scatter.mu 与 gather_scatter.hpp +- 改了什么：新增 MUSA 原生按索引逐元素复制，支持 1/2/4/8/16 字节值和 int32/int64 索引；按原始位模式复制实数/复数，无预处理或结果缓存。每次使用 handle 指定的流，网格步进覆盖全部输入。 +- 为什么改（附报错原文）：无报错；原 Scatter 三项约 0.79x，短向量重复 JIT 调度的主机开销较高。 +- 怎么验证的：原 C API Gather/Scatter 精度族通过；新增逐字节回归包含特殊浮点位模式、未写位置保留、输入/索引更新、部分块、两种索引宽度、默认/非默认流。原 10 矩阵两轮定向测量 40 条各自精度通过，最终均值见上表。 +- 其他后端是否也需要：否。 + +### 2026-10-18 capi/src/adaptor/CMakeLists.txt::MUSA 原生复制构建 +- 改了什么：只在 MUSA 分支且找到 mcc 时生成 mp_31 原生复制目标并链接到库；未找到 mcc 或其他架构继续使用原 Triton 路径，其他后端构建不变。 +- 为什么改（附报错原文）：接入上述原生内核。初版构建规则曾报 `missing and no known rule to make it`，已增加明确 custom target 依赖；原始构建日志保留。 +- 怎么验证的：C API 重新构建通过，基线仍为 `/usr/local/musa/lib/libmusparse.so`；8 个交付相关精度测试族全部通过。 +- 其他后端是否也需要：否。 + +### 2026-10-18 capi/src/ops/gather_scatter.cpp::flagsparseGather/flagsparseScatter +- 改了什么：在 MUSA 原生目标存在且架构为 31 时调用原生复制；保留原 descriptor 校验、零输入语义和其他后端路径。 +- 为什么改（附报错原文）：无报错；减少短向量重复 JIT 调度开销，保证已有相邻 0.8x 的变体有实测余量。 +- 怎么验证的：原 benchmark 原 10 矩阵，Gather/Scatter 四项最终 0.987362x/0.996308x/0.995981x/0.997719x，40 条精度全部 pass；C API 精度族和新增流回归通过。 +- 其他后端是否也需要：否。 + +### 2026-10-18 tests/pytest/test_musa_sparse_routes.py::原生 COO 回归 +- 改了什么：增加原生路径实际执行、不同平均行长、部分块/空行、四种输出组合、输入/存储修改、推理张量和调试/钩子/编译器缺失回退用例；保留所有原用例和精度容差。 +- 为什么改（附报错原文）：无报错；防止原生集成静默回退、复用旧结果或忽略运行时钩子。 +- 怎么验证的：纳入上述扩大 Python 回归，1199 passed、1678 skipped。另跑后端/基线契约检查为 28 passed、1 failed；失败原文 `AssertionError: mthreads: no fallback reason on a CUDA box`，对应测试明确假定 CUDA 机器且 MUSA 插件缺失，当前是真实 MUSA 环境。未修改测试或后端探测以让它通过，原日志 `native_backend_contract.log` 保留。 +- 其他后端是否也需要：否；本文件显式要求 MUSA。 + +### 2026-10-18 capi/ctest/accuracy/test_scatter.cpp::MusaLiveBitCopiesOnDefaultAndNondefaultStreams +- 改了什么：仅 MUSA 增加 Gather 到 Scatter 的原始位复制回归，覆盖 6 种值类型、两种索引宽度、默认/非默认流和多次修改；保留原测试。 +- 为什么改（附报错原文）：无算子错误；原有测试不能覆盖原生复制的非默认流、输入实时性及特殊浮点位模式。新增测试初版误用辅助函数名，已修正，构建日志保留。 +- 怎么验证的：C API 8 个交付相关精度族全部通过（29.30 秒），加上 Gather/Scatter benchmark 共 10/10 passed（42.15 秒），见 `capi_native_final.log`。 +- 其他后端是否也需要：否。 + +### 2026-10-18 冻结与新完整轮 + +- 上述定向结果不覆盖旧完整轮，也不代表 65 项已经全部达标。新完整轮使用 `results_musa_split_65_20261018_opt4/`；该目录保存冻结补丁、源码 SHA256 与运行信息。 +- 不跳过 C API 构建，沿用原 runner 和原 10 矩阵。命令如下；实际进程和启动时间以结果目录内记录及日志为准。 + +```bash +setsid env PYTHONPATH=$PWD/src FLAGSPARSE_BACKEND=mthreads MUSA_HOME=/usr/local/musa \ + OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 \ + timeout -s KILL 43200 python3 -u run_flagsparse_split_delivery.py \ + --mode normal --benchmark-input tests/data --timeout 3600 --h800-reference \ + --results-dir results_musa_split_65_20261018_opt4 \ + --capi-bench-out results_musa_split_65_20261018_opt4/capi \ + > results_musa_split_65_20261018_opt4.log 2>&1 < /dev/null & +``` + +### 2026-10-18 撤回原生 MUSA 实现，恢复全 Triton 约束 + +- 用户要求：计算内核必须保持全 Triton，允许使用 TLE，不允许原生 MUSA 实现。 +- opt4 在 C API SDDMM 性能阶段停止；进程组 1776109 已退出，原日志、结果和冻结补丁全部保留。该轮包含原生实现，不能计作全 Triton 结果。 +- 移除 `_musa_spmv_native.py`、`gather_scatter.mu/.hpp` 及 Python/CMake/C API 原生接入；可恢复副本见 `results_musa_opt_20261018_triton_v1/withdrawn_native/`。COO 混合精度 SpMV 恢复原 Triton 路径。 +- 同时移除此前 MUSA Python FP16 Gather 和 FP16/int8 Scatter 的 PyTorch 原生索引分支，改为支持输入/索引/输出 stride 的 `@triton.jit` 索引复制和清零 kernel，保留其他后端的原路由。 +- 将原生专属回归调整为断言 Triton 路由，并禁止 `index_select`/`gather`/`index_copy_` 调用；保留输入修改、非连续输出和流的精度覆盖。 +- C API 已重新构建；构建规则和库符号核对不再包含原生复制目标。8 个交付精度族全部通过（38.07 秒），日志 `results_musa_opt_20261018_triton_v1/capi_accuracy.log`。 +- Python 路由及 Gather/Scatter 精度回归 188 passed、2 skipped；扩大 SpMV/SpMM/Gather/Scatter 回归 1088 passed、1658 skipped（20.12 秒），日志/JUnit 为 `accuracy_final.log/.xml`。初轮 4 个 CPU mock 因未隔离 Triton 清零函数而失败，修正测试隔离后通过；初轮日志保留。最终源码精度复测记录另行追加。 +- 原 10 矩阵纯 Triton C API 定向性能 `capi_before_v2/*_benchmark.json`：Gather FP32 0.966756x，Scatter FP64/complex64/complex128 分别 0.985301x/0.979894x/0.982658x，40 条精度全部 pass；本机 muSPARSE 对照，未修改 benchmark。首次定向运行未提前创建输出子目录，CTest 通过但未产生 JSON，不用作性能证据；重跑创建目录后的数据完整保留。 +- COO 的归约布局/warp 配置探针未确认稳定收益；未将其集成入生产路径。临时的可调 launch 参数也撤回，保持此前上游 Triton launcher 不变。TLE 安装及 MUSA 特性白名单已检查，此轮未启用未经验证的 TLE kernel。 +- 尚未重新完成全 Triton 的完整 65 项性能验收；COO 等未达标项仍需优化，不能沿用含原生实现的数字声明通过。 +- 最终源码复测为 **1087 passed、1658 skipped**（18.85 秒），见 `accuracy_restored.log/.xml`；相比上面的临时探针版本少 1 个仅测试可调 launcher 编译参数的用例，因为对应生产改动已撤回，并非减少算子精度覆盖。原有 Triton launcher 与上游一致。当前未启动新的完整后台轮。 + +### 2026-10-18 src/flagsparse/sparse_operations/gather_scatter.py::MUSA Triton 索引复制 +- 改了什么：MUSA 的 Gather/Scatter 统一调用 `@triton.jit` 索引复制，清零也使用 Triton；显式传入值/索引/输出 stride，并通过实部视图覆盖复数。移除原有 MUSA FP16 `index_select` 和 FP16/int8 `index_copy_` 分支。 +- 为什么改（附报错原文）：无算子报错；落实用户的全 Triton 约束，并保留原生分支此前提供的非连续输入/输出语义。 +- 怎么验证的：上述扩大 Python 回归；新增测试禁止 PyTorch 原生索引调用，并验证公开 API、实时输入、非连续输出、int32/int64 索引。 +- 其他后端是否也需要：否；新路径严格限于 MUSA，厂商性能基线不改。 + +### 2026-10-18 全 Triton 新完整轮后台启动 + +- 按用户要求于 2026-10-18 00:10:30 UTC 用 setsid 启动完整 65 项精度与性能复测。新目录为 `results_musa_split_65_20261018_triton1/`，日志为 `results_musa_split_65_20261018_triton1.log`；不复用含原生实现的 opt4。 +- 启动后核验 timeout PID/PGID/SID 均为 1914848，PPID=1；split runner PID=1914850，Python 阶段 PID=1914851，日志已进入 delivery-only 测试。此记录更新此前“当前未启动新的完整后台轮”的状态，不代表整轮已完成。 +- HEAD `3cea4bc3615655dabcf676850de7d770d7da795a`；源码导入和 `mthreads musa None` 指纹核对通过；torch 2.7.1、torch_musa 2.7.1+5ee0a64、Triton 3.6.0、MTT S5000、驱动 3.3.5-server。 +- 性能输入仍为 `tests/data` 的原 10 矩阵，名称和字节数一致；固定 OMP/MKL/OPENBLAS 线程池为 8，CTest 串行，启动前设备空闲。没有修改实现、runner、测试规模、容差、迭代次数或基线优先级。 +- 沿用 `--mode normal --benchmark-input tests/data --timeout 3600 --h800-reference`，不跳过 C API 构建，显式 `--capi-bench-out results_musa_split_65_20261018_triton1/capi`；12 小时仅为外层限时。完整命令、PID 与源码/矩阵指纹见本轮 `RUN_INFO.md`、`runner.pid`、`source.patch` 和 `input_source.sha256`。 +- 完成后检查 `summary_split.json` 和 `tools/delivery_table.py --show-baseline`；H800 折算仍属估算。尚不能声明 65 项全部精度通过或每项平均加速比 >0.8x。 + +### 2026-10-19 全 Triton 完整轮结果及继续优化 + +- 最近完成轮为 `results_musa_split_65_20261019_triton2/`：65/65 精度 Passed、65/65 性能测试 Passed、0 missing；按每个变体的矩阵平均加速比检查为 64/65 超过 0.8x，剩余 `spmv_coo_f16f32_int_non` 为 0.760595x。该变体采用 H800 带宽折算估算，不能当作本机 muSPARSE 实测。 +- 本轮继续使用 `tests/data` 原 10 个矩阵、原 q4 benchmark 的 5 warmup / 20 iters 和原 split runner；未改规模、容差、基线或迭代次数。旧轮、归档和全部探索数据保留。 + +### 2026-10-19 src/flagsparse/sparse_operations/mixed_spmx.py::MUSA FP16 到 FP32 COO SpMV +- 改了什么：在 MUSA FP16 到 FP32 COO 路由新增 Triton 固定最大行长循环，避免每个 program 动态归约循环上界；仅在最大行长不超过 128 且平均行长大于 4 时启用。其他形状沿用原 NNZ/动态行块 Triton kernel。版本/存储感知 COO 缓存附带绑定的 Triton 启动特化，校验 RHS 类型、设备、对齐和指针别名，每次读取当前 RHS/输出地址及当前流，不缓存输出。COO SpMM 保留原三项缓存接口。 +- 为什么改（附报错原文）：无算子报错；上一完整轮该变体为 0.760595x，未达到用户要求。探针定位到动态行长归约和重复启动特化开销；行长排序、额外布局/warp 探针没有稳定整体收益，未接入生产路径。 +- 怎么验证的：原 q4 公开 API 六个 COO 变体共 60 条精度全部 PASS；目标变体两次定向十矩阵平均为 0.810783x、0.806586x（同一旧轮 H800 折算分母），记录 `results_musa_opt_20261019_triton_v3/coo_bound_public.csv` 和 `coo_bound_repeat.csv`。最终回归 194 passed、2 skipped，见 `accuracy_final.log/.xml`。启动器缓存键补齐 RHS dtype 后由完整轮再验收，不能以定向结果声明 65/65 性能达标。 +- 其他后端是否也需要：否；新增调用严格位于 MUSA 分支。`audit_scope.py` 确认共享 kernel 定义、非 MUSA COO 路由及 C API 计算/构建代码与上游一致；非 MUSA 索引复制 CPU mock 对照为 96 passed、1 deselected，不代表其他加速器实机性能验证。 + +### 2026-10-19 tests/pytest/test_musa_sparse_routes.py::绑定启动器实时性回归 +- 改了什么：增加固定行长、NNZ、短行三条路线与连续/非连续输出的六个用例，覆盖默认/非默认流、RHS 修改与对齐变化、输出独立性、值/索引修改、替换存储及预热后类型校验。 +- 为什么改（附报错原文）：无算子报错；防止绑定启动器复用错误指针、旧结果、旧流或绕过参数校验。 +- 怎么验证的：与 Gather/Scatter 精度回归一起为 194 passed、2 skipped；既有乱序/非连续 COO 和推理张量回归保留并通过。 +- 其他后端是否也需要：否；测试要求实机 MUSA。 + +### 2026-10-19 全 Triton 第三轮完整复测 + +- 新目录 `results_musa_split_65_20261019_triton3/`，显式将 C API 原始数据存入该目录下 `capi/`，避免覆盖上一轮原始数据。仍按 setsid 在后台启动，不跳过精度、性能或 C API 构建。源码补丁、SHA256、矩阵指纹和命令保存在结果目录。 +- 此处是待验收记录：完整轮约需 50 分钟，完成后再检查 65 项精度、missing、基线来源和各变体平均加速比；不能提前声明全部达标。 diff --git a/src/flagsparse/sparse_operations/_musa_spmv_launch.py b/src/flagsparse/sparse_operations/_musa_spmv_launch.py index 5251b6f..5f9892a 100644 --- a/src/flagsparse/sparse_operations/_musa_spmv_launch.py +++ b/src/flagsparse/sparse_operations/_musa_spmv_launch.py @@ -1,30 +1,85 @@ # Copyright 2026 FlagOS Contributors # SPDX-License-Identifier: Apache-2.0 -"""Reuse compiled MUSA SpMV launchers, never operands or computed outputs.""" +"""Reuse compiled MUSA Triton launchers, never computed outputs.""" from collections import OrderedDict +import os _LAUNCHERS = OrderedDict() _MAX_LAUNCHERS = 64 +_RAW_STREAM = None +try: + _NUM_WARPS = max(1, min(8, int(os.environ.get("FLAGSPARSE_MUSA_TRITON_WARPS", "4")))) +except (TypeError, ValueError): + _NUM_WARPS = 4 -def launch(kernel, grid, tensors, scalars, constants): - """Keep the pointer types, alignment and aliases of Triton's specialization. - Sparse plan caches still validate tensor versions/storage on every call. - This cache only bypasses repeated JIT argument binding for the same kernel - signature. CompiledKernel's public launcher resolves the current stream - and runs Triton's profiling hooks on every invocation. - """ +def _has_hook(hook): + # FlagTree installs empty HookChains instead of None even without a profiler. + return hook is not None and (not hasattr(hook, "calls") or bool(hook.calls)) + + +def fast_path_allowed(kernel): + # Imported here, not at module level: every backend imports this module via + # gather_scatter/mixed_spmx/spmv_sell, and older Triton builds (CoreX 3.2.3 + # ships Triton 2.3.1) have no triton.knobs. Only MUSA ever reaches this call. from triton import knobs - if ( + return not ( kernel.pre_run_hooks or knobs.runtime.debug or knobs.compilation.instrumentation_mode or knobs.runtime.add_stages_inspection_hook is not None - ): - return kernel[grid](*tensors, *scalars, *constants, num_warps=4, num_stages=1) + or _has_hook(knobs.runtime.launch_enter_hook) + or _has_hook(knobs.runtime.launch_exit_hook) + ) + + +class _Launch: + """Own the compiled module but resolve pointers/streams on every launch.""" + + def __init__(self, kernel, compiled, grid, device): + self.kernel = kernel + self.compiled = compiled + self.grid = grid + self.device = device + self.runner = compiled[grid] + self.direct = hasattr(compiled, "run") and device.type == "musa" + + def __call__(self, tensors, scalars, constants, addresses=None): + global _RAW_STREAM + if not fast_path_allowed(self.kernel): + return self.kernel[self.grid]( + *tensors, *scalars, *constants, num_warps=_NUM_WARPS, num_stages=1, + ) + if not self.direct: + return self.runner(*tensors, *scalars, *constants) + if _RAW_STREAM is None: + import torch_musa + + _RAW_STREAM = torch_musa._MUSAC._musa_getCurrentRawStream + if addresses is None: + addresses = tuple(t.data_ptr() for t in tensors) + # This is Triton's normal compiled launcher, including scratch handling. + # Only pointer validation and empty profiling metadata are redundant. + self.compiled.run( + *self.grid, _RAW_STREAM(self.device.index), self.compiled.function, + self.compiled.packed_metadata, None, None, None, + *addresses, *scalars, *constants, + ) + + +def launch(kernel, grid, tensors, scalars, constants): + """Cache dtype/device/alignment/alias specializations, not tensor addresses. + + Returns the bound specialization after executing it. Callers may reuse that + object only while validating the same specialization themselves. Debugging, + instrumentation and profiling always use the full JIT path. + """ + if not fast_path_allowed(kernel): + kernel[grid](*tensors, *scalars, *constants, num_warps=_NUM_WARPS, num_stages=1) + return None addresses = tuple(t.data_ptr() for t in tensors) signature = tuple((t.dtype, t.device, p % 16) for t, p in zip(tensors, addresses)) aliases = () @@ -32,14 +87,21 @@ def launch(kernel, grid, tensors, scalars, constants): aliases = tuple((i, j) for j, p in enumerate(addresses) if p for i in range(j) if addresses[i] == p) grid = tuple(grid) + (1,) * (3 - len(grid)) - key = (kernel, grid, signature, aliases, scalars, constants, kernel.debug) - runner = _LAUNCHERS.get(key) - if runner is None: - compiled = kernel[grid](*tensors, *scalars, *constants, num_warps=4, num_stages=1) - if compiled is not None: - _LAUNCHERS[key] = compiled[grid] - while len(_LAUNCHERS) > _MAX_LAUNCHERS: - _LAUNCHERS.popitem(last=False) - return compiled - _LAUNCHERS.move_to_end(key) - runner(*tensors, *scalars, *constants) + key = (id(kernel), getattr(kernel, "hash", None), grid, signature, aliases, + scalars, constants, kernel.debug, _NUM_WARPS) + bound = _LAUNCHERS.get(key) + if bound is None: + compiled = kernel[grid]( + *tensors, *scalars, *constants, num_warps=_NUM_WARPS, num_stages=1, + ) + if compiled is None: + return None + key = (id(kernel), getattr(kernel, "hash", None), *key[2:]) + bound = _Launch(kernel, compiled, grid, tensors[0].device) + _LAUNCHERS[key] = bound + while len(_LAUNCHERS) > _MAX_LAUNCHERS: + _LAUNCHERS.popitem(last=False) + else: + _LAUNCHERS.move_to_end(key) + bound(tensors, scalars, constants, addresses) + return bound diff --git a/src/flagsparse/sparse_operations/gather_scatter.py b/src/flagsparse/sparse_operations/gather_scatter.py index 2381033..ab76c6f 100644 --- a/src/flagsparse/sparse_operations/gather_scatter.py +++ b/src/flagsparse/sparse_operations/gather_scatter.py @@ -19,12 +19,20 @@ import ctypes import ctypes.util +from collections import OrderedDict HipPointer = _common_mod.HipPointer import triton import triton.language as tl +from ._rocm_sparse_cache import TensorPlanCache +from ._musa_spmv_launch import launch as _musa_launch + +_MUSA_SCATTER_INDEX_CACHE = TensorPlanCache(maxsize=16) +_MUSA_COPY_RUNNERS = OrderedDict() +_MUSA_COPY_RUNNER_LIMIT = 32 + SUPPORTED_SCATTER_VALUE_DTYPES = GATHER_SCATTER_VALUE_DTYPES DEFAULT_GATHER_BLOCK_SIZE = 256 DEFAULT_GATHER_NUM_WARPS = 8 @@ -256,6 +264,131 @@ def _scatter_complex_kernel( tl.store(dense_values_ri_ptr + dense_offsets + 1, values_imag, mask=mask) +@triton.jit +def _musa_indexed_copy_kernel( + DST, SRC, IDX, N, + DST_STRIDE: tl.constexpr, SRC_STRIDE: tl.constexpr, INDEX_STRIDE: tl.constexpr, + COMPONENTS: tl.constexpr, GATHER: tl.constexpr, BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + elements = offsets // COMPONENTS + component = offsets % COMPONENTS + valid = elements < N + indices = tl.load(IDX + elements * INDEX_STRIDE, valid, other=0) + if GATHER: + source = indices * SRC_STRIDE + component + destination = elements * DST_STRIDE + component + else: + source = elements * SRC_STRIDE + component + destination = indices * DST_STRIDE + component + values = tl.load(SRC + source, valid, other=0) + tl.store(DST + destination, values, valid) + + +@triton.jit +def _musa_zero_vector_kernel( + DST, N, STRIDE: tl.constexpr, COMPONENTS: tl.constexpr, BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + elements = offsets // COMPONENTS + tl.store(DST + elements * STRIDE + offsets % COMPONENTS, 0, elements < N) + + +@triton.jit +def _musa_scatter_inverse_fill_kernel(INVERSE, N, BLOCK: tl.constexpr): + pos = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(INVERSE + pos, -1, pos < N) + + +@triton.jit +def _musa_scatter_inverse_kernel( + INVERSE, INDICES, NNZ, N, INDEX_STRIDE: tl.constexpr, BLOCK: tl.constexpr, +): + pos = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + row = tl.load(INDICES + pos * INDEX_STRIDE, pos < NNZ, other=-1) + # Last source position wins even for duplicate destinations. Separate + # initialization and construction launches avoid a cross-program race. + tl.atomic_max(INVERSE + row, pos, (pos < NNZ) & (row >= 0) & (row < N), + sem="relaxed") + + +@triton.jit +def _musa_scatter_dense_kernel( + DST, SRC, INVERSE, N, DST_STRIDE: tl.constexpr, SRC_STRIDE: tl.constexpr, + COMPONENTS: tl.constexpr, BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + row = offsets // COMPONENTS + pos = tl.load(INVERSE + row, row < N, other=-1) + value = tl.load(SRC + pos * SRC_STRIDE + offsets % COMPONENTS, + (row < N) & (pos >= 0), other=0) + tl.store(DST + row * DST_STRIDE + offsets % COMPONENTS, value, row < N) + + +def _musa_scatter_dense(dst, src, indices, block_size): + """Reuse index topology only; read source values and write every output anew.""" + from ._musa_spmv_launch import launch + + n, nnz = dst.numel(), indices.numel() + inverse = _MUSA_SCATTER_INDEX_CACHE.get((indices,), (n,)) + if inverse is None: + inverse = torch.empty(n, dtype=torch.int32, device=dst.device) + launch(_musa_scatter_inverse_fill_kernel, (triton.cdiv(n, 1024),), + (inverse,), (n,), (1024,)) + launch(_musa_scatter_inverse_kernel, (triton.cdiv(nnz, 1024),), + (inverse, indices), (nnz, n), (indices.stride(0), 1024)) + _MUSA_SCATTER_INDEX_CACHE.put((indices,), (n,), inverse) + components = 2 if _is_complex_dtype(src.dtype) else 1 + dst_view = torch.view_as_real(dst) if components == 2 else dst + src_view = torch.view_as_real(src) if components == 2 else src + launch(_musa_scatter_dense_kernel, (triton.cdiv(n * components, block_size),), + (dst_view, src_view, inverse), (n,), + (dst_view.stride(0), src_view.stride(0), components, block_size)) + + +def _musa_indexed_copy(dst, src, indices, gather, block_size): + """Triton-only indexed copies, including noncontiguous real/complex views.""" + components = 2 if _is_complex_dtype(src.dtype) else 1 + dst_view = torch.view_as_real(dst) if components == 2 else dst + src_view = torch.view_as_real(src) if components == 2 else src + nnz = indices.numel() + grid = ((nnz * components + block_size - 1) // block_size,) + scalars = (nnz,) + constants = (dst_view.stride(0), src_view.stride(0), indices.stride(0), + components, gather, block_size) + addresses = (dst_view.data_ptr(), src_view.data_ptr(), indices.data_ptr()) + signature = tuple((t.dtype, t.device, p % 16) for t, p in + zip((dst_view, src_view, indices), addresses)) + aliases = (addresses[0] == addresses[1], addresses[0] == addresses[2], + addresses[1] == addresses[2]) + key = (signature, aliases, nnz, constants) + entry = _MUSA_COPY_RUNNERS.get(key) + if entry is not None: + entry((dst_view, src_view, indices), scalars, constants, addresses) + _MUSA_COPY_RUNNERS.move_to_end(key) + return + runner = _musa_launch( + _musa_indexed_copy_kernel, grid, (dst_view, src_view, indices), + scalars, constants, + ) + if runner is not None: + _MUSA_COPY_RUNNERS[key] = runner + _MUSA_COPY_RUNNERS.move_to_end(key) + while len(_MUSA_COPY_RUNNERS) > _MUSA_COPY_RUNNER_LIMIT: + _MUSA_COPY_RUNNERS.popitem(last=False) + + +def _musa_zero_vector(out): + from ._musa_spmv_launch import launch + + if not out.numel(): + return + components = 2 if _is_complex_dtype(out.dtype) else 1 + view = torch.view_as_real(out) if components == 2 else out + launch(_musa_zero_vector_kernel, (triton.cdiv(out.numel() * components, 256),), + (view,), (out.numel(),), (view.stride(0), components, 256)) + + def _triton_gather_impl( dense_vector, kernel_indices, @@ -272,18 +405,16 @@ def _triton_gather_impl( raise TypeError("out dtype must match gather output dtype") return out - if _is_mthreads_runtime() and dense_vector.dtype == torch.float16: - # Native indexing avoids the MUSA Triton launch overhead for short gathers. + if _is_mthreads_runtime(): if out is not None: if out.shape != (nnz,): raise ValueError("out shape must match gather output shape") if out.dtype != dense_vector.dtype: raise TypeError("out dtype must match gather output dtype") - dense_vector = dense_vector.contiguous() - if out is not None and not out.is_contiguous(): - out.copy_(torch.index_select(dense_vector, 0, kernel_indices)) - return out - return torch.index_select(dense_vector, 0, kernel_indices, out=out) + else: + out = torch.empty(nnz, dtype=dense_vector.dtype, device=dense_vector.device) + _musa_indexed_copy(out, dense_vector, kernel_indices, True, block_size) + return out max_programs = _gather_max_programs(dense_vector.device, DEFAULT_GATHER_NUM_WARPS) grid = lambda meta: (min(max_programs, triton.cdiv(nnz, meta["BLOCK_SIZE"])),) @@ -346,23 +477,28 @@ def _triton_scatter_impl( index_fallback_policy = str(index_fallback_policy).lower() if index_fallback_policy not in ("auto", "strict"): raise ValueError("index_fallback_policy must be 'auto' or 'strict'") - native_scatter = ( - _is_mthreads_runtime() - and block_size is None - and sparse_values.dtype in (torch.float16, torch.int8) - ) + musa = _is_mthreads_runtime() block_size = _resolve_scatter_block_size( block_size, sparse_values.dtype, kernel_indices.numel() ) if out is None: - dense_values = torch.zeros( + allocate = torch.empty if musa else torch.zeros + dense_values = allocate( dense_size, dtype=sparse_values.dtype, device=sparse_values.device ) else: dense_values = out - if reset_output: + if reset_output and not musa: dense_values.zero_() + fused_musa = ( + musa and sparse_values.device.type == "musa" + and (out is None or reset_output) and kernel_indices.numel() > 0 + and kernel_indices.numel() <= _INDEX_LIMIT_INT32 + and sparse_values.dtype in (torch.float16, torch.int8) + ) + if musa and (out is None or reset_output) and not fused_musa: + _musa_zero_vector(dense_values) nnz = kernel_indices.numel() scatter_meta = { @@ -375,21 +511,17 @@ def _triton_scatter_impl( return dense_values, scatter_meta return dense_values - if native_scatter: - dense_values.index_copy_(0, kernel_indices.to(torch.int64), sparse_values) - scatter_meta["kernel_index_dtype"] = "int64" - if return_metadata: - return dense_values, scatter_meta - return dense_values - try: - _launch_triton_scatter_kernel( - dense_values, - sparse_values, - kernel_indices, - nnz, - block_size=block_size, - ) + if fused_musa: + _musa_scatter_dense(dense_values, sparse_values, kernel_indices, block_size) + else: + _launch_triton_scatter_kernel( + dense_values, + sparse_values, + kernel_indices, + nnz, + block_size=block_size, + ) except Exception as exc: if kernel_indices.dtype != torch.int64 or index_fallback_policy != "auto": raise RuntimeError( @@ -406,6 +538,8 @@ def _triton_scatter_impl( fallback_indices = kernel_indices.to(torch.int32) try: + if fused_musa: + _musa_zero_vector(dense_values) _launch_triton_scatter_kernel( dense_values, sparse_values, @@ -434,6 +568,9 @@ def _launch_triton_scatter_kernel( dense_values, sparse_values, kernel_indices, nnz, block_size=None ): block_size = _resolve_scatter_block_size(block_size, sparse_values.dtype, nnz) + if _is_mthreads_runtime(): + _musa_indexed_copy(dense_values, sparse_values, kernel_indices, False, block_size) + return grid = lambda meta: (triton.cdiv(nnz, meta["BLOCK_SIZE"]),) launch_kwargs = {"BLOCK_SIZE": block_size} if _is_rocm_runtime(): diff --git a/src/flagsparse/sparse_operations/mixed_spmx.py b/src/flagsparse/sparse_operations/mixed_spmx.py index 9726b42..85465a2 100644 --- a/src/flagsparse/sparse_operations/mixed_spmx.py +++ b/src/flagsparse/sparse_operations/mixed_spmx.py @@ -256,6 +256,14 @@ def _csr_row_tile(nnz, n_rows): def _musa_row_tile(nnz, n_rows, integer=False): + override = os.environ.get("FLAGSPARSE_MUSA_CSR_TILE") + if override: + try: + rows, block = (int(v.strip()) for v in override.split(",", 1)) + if rows > 0 and block > 0: + return rows, block + except (TypeError, ValueError): + pass mean = nnz / max(1, n_rows) if mean <= 3: return 256, 2 @@ -266,7 +274,95 @@ def _musa_row_tile(nnz, n_rows, integer=False): return 16, 8 -def _musa_coo_csr(data, row, col, shape): +@triton.jit +def _musa_csr_half_fixed_kernel( + Y, V, C, P, X, M, ROWS: tl.constexpr, BLOCK: tl.constexpr, + MAX_LENGTH: tl.constexpr, +): + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + begin = tl.load(P + rows, rows < M, other=0) + end = tl.load(P + rows + 1, rows < M, other=0) + acc = tl.zeros((ROWS, BLOCK), tl.float32) + # The version-checked row pointers bound every row, so no per-program + # reduction is needed to discover the loop limit on each invocation. + for offset in range(0, MAX_LENGTH, BLOCK): + pos = begin[:, None] + offset + tl.arange(0, BLOCK)[None, :] + valid = (rows[:, None] < M) & (pos < end[:, None]) + col = tl.load(C + pos, valid, other=0) + value = tl.load(V + pos, valid, other=0).to(tl.float32) + rhs = tl.load(X + col, valid, other=0).to(tl.float32) + acc += value * rhs + tl.store(Y + rows, tl.sum(acc, 1), rows < M) + + +def _musa_coo_half_spmv(data, cols, ptr, x, n_rows, plans, out=None): + """Reuse topology and launch specializations; always read live operands.""" + plan = plans.get("half_f32") + if plan is None: + lengths = ptr[1:] - ptr[:-1] + maximum = int(lengths.max().item()) + mean = data.numel() / n_rows + nnz_mode = maximum > 128 and mean <= 4 + if nnz_mode: + # Keep row ownership and the existing cross-block prefix merge. + from ._spmv_csr_nnz import _nnz_kernel + + row_ids = torch.repeat_interleave( + torch.arange(n_rows, device=ptr.device, dtype=torch.int32), + lengths.long(), output_size=data.numel(), + ) + empty = torch.nonzero(lengths == 0).flatten().to(torch.int32) + sparse = (data, cols, ptr, row_ids, empty) + kernel = _nnz_kernel + grid = ((max(data.numel(), empty.numel()) + 511) // 512,) + scalars = () + constants = (empty.numel(), data.numel(), 512, tl.float32, True, False) + else: + sparse = (data, cols, ptr) + if maximum <= 128 and mean > 4: + rows, block = (32, 4) if mean <= 12 and n_rows >= 262144 else (16, 8) + kernel = _musa_csr_half_fixed_kernel + constants = (rows, block, maximum) + else: + rows, block = _musa_row_tile(data.numel(), n_rows) + kernel = _csr_spmv_mixed_kernel + constants = (tl.float32, rows, block) + grid = ((n_rows + rows - 1) // rows,) + scalars = (n_rows,) + plan = (kernel, grid, sparse, tuple(t.data_ptr() for t in sparse), + scalars, constants, nnz_mode, {}) + plans["half_f32"] = plan + kernel, grid, sparse, sparse_addresses, scalars, constants, nnz_mode, runners = plan + y = out if out is not None and out.is_contiguous() else torch.empty( + n_rows, dtype=torch.float32, device=x.device + ) + yp, xp = y.data_ptr(), x.data_ptr() + # The sparse cache checks all source versions/storage/metadata. RHS and + # output may change on every call: preserve alignment/device/alias JIT + # specializations while resolving their current pointers and stream. + aliases = tuple((p == yp, p == xp) for p in sparse_addresses) + (yp == xp,) + key = (x.dtype, x.device, xp % 16, yp % 16, aliases, kernel.debug) + runner = runners.get(key) + if nnz_mode: + tensors = (*sparse[:4], x, y, sparse[4]) + addresses = (*sparse_addresses[:4], xp, yp, sparse_addresses[4]) + else: + tensors = (y, *sparse, x) + addresses = (yp, *sparse_addresses, xp) + if runner is None: + runner = _launch_musa_spmv( + kernel, grid, tensors, scalars, constants, + ) + if runner is not None: + if len(runners) >= 16: + runners.pop(next(iter(runners))) + runners[key] = runner + else: + runner(tensors, scalars, constants, addresses) + return y + + +def _musa_coo_csr(data, row, col, shape, *, spmv=False): """Reuse row ownership, invalidating both topology and values on mutation.""" sources = (data, row, col) options = tuple(shape) @@ -284,9 +380,9 @@ def _musa_coo_csr(data, row, col, shape): # MUSA bincount can lose updates for repeated sorted row IDs. boundaries = torch.arange(int(shape[0]) + 1, device=rows.device, dtype=rows.dtype) ptr = torch.searchsorted(rows, boundaries).to(torch.int32) - state = data.contiguous(), cols.contiguous(), ptr.contiguous() + state = data.contiguous(), cols.contiguous(), ptr.contiguous(), {} _MUSA_COO_CSR_CACHE.put(sources, options, state) - return state + return state if spmv else state[:3] def _musa_csr_spmv(data, cols, ptr, x, n_rows, result_dtype, acc_dtype, out=None): @@ -469,9 +565,12 @@ def spmv_coo_mixed(data, row, col, x, shape, *, op=None, transpose=None, raise ValueError("out shape/device must match the COO SpMV result") if musa: t0 = _start(return_time) - values, cols, ptr = _musa_coo_csr(data, row, col, shape) - y = _musa_csr_spmv(values, cols, ptr, x.contiguous(), n_rows, - result_dtype, acc_dtype, out) + values, cols, ptr, plans = _musa_coo_csr(data, row, col, shape, spmv=True) + if data.dtype == torch.float16 and result_dtype == torch.float32 and n_rows: + y = _musa_coo_half_spmv(values, cols, ptr, x.contiguous(), n_rows, plans, out) + else: + y = _musa_csr_spmv(values, cols, ptr, x.contiguous(), n_rows, + result_dtype, acc_dtype, out) return _finish(y, out, t0, return_time) sources = (data, row, col) data, x = data.contiguous(), x.contiguous() diff --git a/tests/pytest/test_musa_sparse_routes.py b/tests/pytest/test_musa_sparse_routes.py index 7e7ffab..7c5aafc 100644 --- a/tests/pytest/test_musa_sparse_routes.py +++ b/tests/pytest/test_musa_sparse_routes.py @@ -185,7 +185,7 @@ def test_musa_csc_half_spmm_strides_and_out(index_dtype): @pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) @pytest.mark.parametrize("reset", [False, True]) @pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) -def test_musa_scatter_native_strides_and_live_inputs(dtype, reset, index_dtype): +def test_musa_scatter_triton_strides_and_live_inputs(dtype, reset, index_dtype): from flagsparse.sparse_operations.gather_scatter import _triton_scatter_impl dev = accelerator_device() @@ -210,7 +210,7 @@ def test_musa_scatter_native_strides_and_live_inputs(dtype, reset, index_dtype): @pytest.mark.parametrize("musa,block", [(False, None), (True, None), (True, 256)]) @pytest.mark.parametrize("dtype", [torch.float16, torch.int8]) -def test_musa_scatter_native_scope_and_explicit_tuning(monkeypatch, musa, block, dtype): +def test_musa_scatter_triton_scope_and_explicit_tuning(monkeypatch, musa, block, dtype): from flagsparse.sparse_operations import gather_scatter as scatter values = torch.tensor([1, 2], dtype=dtype) @@ -223,13 +223,13 @@ def launch(dense, sparse, index, nnz, block_size=None): dense.index_copy_(0, index.long(), sparse) monkeypatch.setattr(scatter, "_is_mthreads_runtime", lambda: musa) + monkeypatch.setattr(scatter, "_musa_zero_vector", lambda dense: dense.zero_()) monkeypatch.setattr(scatter, "_launch_triton_scatter_kernel", launch) result, meta = scatter._triton_scatter_impl(values, indices, 3, out=out, block_size=block, return_metadata=True) - native = musa and block is None assert result is out - assert bool(launches) != native - assert meta["kernel_index_dtype"] == ("int64" if native else "int32") + assert launches == [scatter._resolve_scatter_block_size(block, dtype, 2)] + assert meta["kernel_index_dtype"] == "int32" assert not meta["index_fallback_applied"] assert torch.equal(result, torch.tensor([2, 0, 1], dtype=dtype)) @@ -629,10 +629,11 @@ def musa_launch_cache(monkeypatch): from flagsparse.sparse_operations import _musa_spmv_launch as cache monkeypatch.setattr(cache, "_LAUNCHERS", OrderedDict()) + monkeypatch.setattr(cache, "_RAW_STREAM", lambda device: 123) class Compiled: def __getitem__(self, grid): - def run(*args): + def run(*args, stream=None): kernel.calls.append(("compiled", grid, args)) return run @@ -795,3 +796,148 @@ def __getitem__(self, grid): monkeypatch.setattr(module, "_spmv_sell_real_kernel", Kernel()) module.flagsparse_spmv_sell(values, indices, offsets, rhs, (1, 2), slice_size=1) assert calls == (["musa"] if musa else ["generic"]) + + +@pytest.mark.parametrize("dtype,out_dtype", [ + (torch.float16, torch.float16), (torch.float16, torch.float32), + (torch.int8, torch.int32), (torch.int8, torch.float32), +]) +@pytest.mark.parametrize("width", [2, 9, 29]) +def test_musa_triton_coo_live_partial_rows(monkeypatch, dtype, out_dtype, width): + from flagsparse.sparse_operations import mixed_spmx as mixed + + dev = accelerator_device() + launches = [] + original = mixed._launch_musa_spmv + + def triton_launch(kernel, *args): + from triton.runtime.jit import JITFunction + assert isinstance(kernel, JITFunction) + launches.append(kernel) + return original(kernel, *args) + + monkeypatch.setattr(mixed, "_launch_musa_spmv", triton_launch) + m, n = 67, 41 + lengths = torch.tensor([width if r % 7 else 0 for r in range(m)]) + rr = torch.repeat_interleave(torch.arange(m), lengths).int() + cc = (torch.arange(rr.numel()) * 7 % n).int() + vv = (torch.arange(rr.numel()) % 5 - 2).to(dtype) + # Start unsorted, with noncontiguous sparse operands and vector. + rows = torch.stack((rr.flip(0), rr), 1).to(dev)[:, 0] + cols = torch.stack((cc.flip(0), cc), 1).to(dev)[:, 0] + values = torch.stack((vv.flip(0), vv), 1).to(dev)[:, 0] + rhs = torch.arange(n * 2, dtype=torch.int32).to(dtype).to(dev)[::2] + backing = torch.full((m * 2,), -99, dtype=out_dtype, device=dev) + out = backing[::2] + + def check(): + reference = torch.zeros(m, dtype=torch.float64) + reference.index_add_(0, rows.cpu().long(), values.cpu().double() * + rhs.cpu().double()[cols.cpu().long()]) + actual = fs.flagsparse_spmv_coo(values, rows, cols, rhs, (m, n), + out=out, out_dtype=out_dtype) + assert launches + assert actual is out + torch.testing.assert_close(actual.cpu(), reference.to(out_dtype), rtol=0, atol=0) + assert torch.equal(backing[1::2].cpu(), torch.full((m,), -99, dtype=out_dtype)) + + check() + values.add_(1) + rhs.add_(1) + check() + values.data = (values + 1).data + cols.data = ((cols + 1) % n).data + check() + rows[0] = 0 + check() + with torch.inference_mode(): + values, rows, cols = values.clone(), rows.clone(), cols.clone() + check() + values.sub_(1) + cols[0] = 3 + check() + + +@pytest.mark.parametrize("kind", ["fixed", "nnz", "short"]) +@pytest.mark.parametrize("strided_out", [False, True]) +def test_musa_half_coo_bound_launch_live_storage_and_stream(kind, strided_out): + from flagsparse.sparse_operations._common import _ACCEL + from flagsparse.sparse_operations import mixed_spmx as mixed + + dev = accelerator_device() + m, n, nnz = (67, 19, 493) if kind == "fixed" else (1025, 19, 1101) + rows = (torch.arange(nnz) % m).int().to(dev) + if kind == "nnz": + rows.zero_() + cols = (torch.arange(nnz) % n).int().to(dev) + values = (torch.arange(nnz) % 5 - 2).half().to(dev) + backing = torch.full((m * 2 + 2,), -99, dtype=torch.float32, device=dev) + out = backing[1:2 * m:2] if strided_out else backing[1:m + 1] + stream = _ACCEL.Stream() + stream.wait_stream(_ACCEL.current_stream()) + + def check(rhs, target=None): + reference = torch.zeros(m, dtype=torch.float64) + reference.index_add_(0, rows.cpu().long(), values.cpu().double() * + rhs.cpu().double()[cols.cpu().long()]) + result = fs.flagsparse_spmv_coo(values, rows, cols, rhs, (m, n), + out=target, out_dtype=torch.float32) + _ACCEL.synchronize() + if target is not None: + assert result is target + torch.testing.assert_close(result.cpu().double(), reference, rtol=0, atol=0) + return result + + for offset in (0, 1): + with _ACCEL.stream(stream): + rhs = torch.arange(n + offset, device=dev).half()[offset:] + first = check(rhs) + preserved = first.clone() + rhs.add_(1) + check(rhs, out) + torch.testing.assert_close(first, preserved, rtol=0, atol=0) + stream.synchronize() + check(rhs) + # Unsupported operand types must still fail after warming the cache. + with pytest.raises(TypeError, match="not a supported combination"): + fs.flagsparse_spmv_coo(values, rows, cols, rhs.float(), (m, n), + out_dtype=torch.float32) + check(rhs) + values.add_(1) + rows[0] = m - 1 + check(rhs) + # Replacing storage does not increment the Tensor's version counter. + values.data = (values - 1).data + cols.data = ((cols + 1) % n).data + check(rhs, out) + state = mixed._MUSA_COO_CSR_CACHE.get((values, rows, cols), (m, n)) + assert state is not None and state[3]["half_f32"][-1] + assert torch.equal(backing[:1].cpu(), torch.tensor([-99.])) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.int8, torch.complex64]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_musa_indexed_copies_never_call_native_indexing(monkeypatch, dtype, index_dtype): + from flagsparse.sparse_operations import gather_scatter as copies + + dev = accelerator_device() + dense = torch.arange(18, dtype=torch.int32).to(dtype).to(dev)[::2] + indices = torch.tensor([7, 2, 0], dtype=index_dtype, device=dev) + sparse = torch.empty(6, dtype=dtype, device=dev)[::2] + target = torch.full((18,), -9, dtype=dtype, device=dev)[::2] + + def forbidden(*args, **kwargs): + raise AssertionError("MUSA indexed copies must execute Triton kernels") + + monkeypatch.setattr(torch, "index_select", forbidden) + monkeypatch.setattr(torch, "gather", forbidden) + monkeypatch.setattr(torch.Tensor, "index_select", forbidden) + monkeypatch.setattr(torch.Tensor, "index_copy_", forbidden) + copies._triton_gather_impl(dense, indices, out=sparse) + copies._triton_scatter_impl(sparse, indices, 9, out=target, reset_output=True) + assert torch.equal(sparse.cpu(), torch.tensor([14, 4, 0], dtype=dtype)) + assert torch.equal(target.cpu(), torch.tensor([0, 0, 4, 0, 0, 0, 0, 14, 0], dtype=dtype)) + assert fs.flagsparse_gather(dense, indices, out=sparse) is sparse + fs.flagsparse_scatter(target, indices, sparse, reset_output=True) + assert torch.equal(sparse.cpu(), torch.tensor([14, 4, 0], dtype=dtype)) + assert torch.equal(target.cpu(), torch.tensor([0, 0, 4, 0, 0, 0, 0, 14, 0], dtype=dtype))