Merge pull request #26 from flagos-ai/cuda_update

Cuda update

Author: zyq1105331849

README.md

@@ -96,9 +96,9 @@ python tests/test_spmm_opt.py <dir/> --csv spmm_opt.csv  # optional: --dtype flo
 
 ```bash
 python tests/test_spmm_coo.py <dir_or_file.mtx>
-python tests/test_spmm_coo.py --synthetic                # optional: --route rowrun|atomic|compare, --skip-api-checks, --skip-coo-coverage
-python tests/test_spmm_coo.py <dir/> --csv out.csv      # only --route rowrun or atomic (not compare)
-# same tuning flags as CSR SpMM where applicable: --dense-cols, --block-n, --block-nnz, --warmup, --iters, --no-cusparse
+python tests/test_spmm_coo.py --synthetic                # optional: --op non|trans|conj|all, --route rowrun|atomic|compare
+python tests/test_spmm_coo.py <dir/> --csv out.csv      # only --route rowrun or atomic (not compare); optional: --op all
+# same tuning flags as CSR SpMM where applicable: --op, --dense-cols, --block-n, --block-nnz, --warmup, --iters, --no-cusparse
 ```
 
 **test_sddmm.py** - CSR SDDMM (`.mtx` batch or `--csv`):

README_cn.md

@@ -94,9 +94,9 @@ python tests/test_spmm_opt.py <目录/> --csv spmm_opt.csv  # 可选：--dtype f
 
 ```bash
 python tests/test_spmm_coo.py <目录或文件.mtx>
-python tests/test_spmm_coo.py --synthetic                # 可选：--route rowrun|atomic|compare、--skip-api-checks、--skip-coo-coverage
-python tests/test_spmm_coo.py <目录/> --csv out.csv     # 仅支持 --route rowrun 或 atomic（compare 不能配 --csv）
-# 与 CSR SpMM 类似的调参：--dense-cols、--block-n、--block-nnz、--warmup、--iters、--no-cusparse
+python tests/test_spmm_coo.py --synthetic                # 可选：--op non|trans|conj|all、--route rowrun|atomic|compare
+python tests/test_spmm_coo.py <目录/> --csv out.csv     # 仅支持 --route rowrun 或 atomic（compare 不能配 --csv）；可选：--op all
+# 与 CSR SpMM 类似的调参：--op、--dense-cols、--block-n、--block-nnz、--warmup、--iters、--no-cusparse
 ```
 
 **test_sddmm.py** - CSR SDDMM（`.mtx` 批量或 `--csv`）：

ops_support.csv

@@ -34,17 +34,41 @@ spgemm,CSR,int32,float64,non,triton,SUPPORTED
 spgemm,CSR,int64,float32,non,triton,SUPPORTED
 spgemm,CSR,int64,float64,non,triton,SUPPORTED
 spmm,COO,int32,float16,non,triton,SUPPORTED
+spmm,COO,int32,float16,trans,triton,SUPPORTED
+spmm,COO,int32,float16,conj,triton,SUPPORTED
 spmm,COO,int32,bfloat16,non,triton,SUPPORTED
+spmm,COO,int32,bfloat16,trans,triton,SUPPORTED
+spmm,COO,int32,bfloat16,conj,triton,SUPPORTED
 spmm,COO,int32,float32,non,triton,SUPPORTED
+spmm,COO,int32,float32,trans,triton,SUPPORTED
+spmm,COO,int32,float32,conj,triton,SUPPORTED
 spmm,COO,int32,float64,non,triton,SUPPORTED
+spmm,COO,int32,float64,trans,triton,SUPPORTED
+spmm,COO,int32,float64,conj,triton,SUPPORTED
 spmm,COO,int32,complex64,non,triton,SUPPORTED
+spmm,COO,int32,complex64,trans,triton,SUPPORTED
+spmm,COO,int32,complex64,conj,triton,SUPPORTED
 spmm,COO,int32,complex128,non,triton,SUPPORTED
+spmm,COO,int32,complex128,trans,triton,SUPPORTED
+spmm,COO,int32,complex128,conj,triton,SUPPORTED
 spmm,COO,int64,float16,non,triton,SUPPORTED
+spmm,COO,int64,float16,trans,triton,SUPPORTED
+spmm,COO,int64,float16,conj,triton,SUPPORTED
 spmm,COO,int64,bfloat16,non,triton,SUPPORTED
+spmm,COO,int64,bfloat16,trans,triton,SUPPORTED
+spmm,COO,int64,bfloat16,conj,triton,SUPPORTED
 spmm,COO,int64,float32,non,triton,SUPPORTED
+spmm,COO,int64,float32,trans,triton,SUPPORTED
+spmm,COO,int64,float32,conj,triton,SUPPORTED
 spmm,COO,int64,float64,non,triton,SUPPORTED
+spmm,COO,int64,float64,trans,triton,SUPPORTED
+spmm,COO,int64,float64,conj,triton,SUPPORTED
 spmm,COO,int64,complex64,non,triton,SUPPORTED
+spmm,COO,int64,complex64,trans,triton,SUPPORTED
+spmm,COO,int64,complex64,conj,triton,SUPPORTED
 spmm,COO,int64,complex128,non,triton,SUPPORTED
+spmm,COO,int64,complex128,trans,triton,SUPPORTED
+spmm,COO,int64,complex128,conj,triton,SUPPORTED
 spmm,CSR,int32,float16,non,triton,SUPPORTED
 spmm,CSR,int32,float16,trans,triton,SUPPORTED
 spmm,CSR,int32,float16,conj,triton,SUPPORTED

ops_support.py

@@ -238,6 +238,11 @@ def registry(modules: dict[str, SourceModule]) -> tuple[ApiSpec, ...]:
         if "spmm_csr" in modules
         else ("non", "trans", "conj")
     )
+    spmm_coo_ops = (
+        op_names(modules["spmm_coo"], "SPMM_COO_OP_NAMES")
+        if "spmm_coo" in modules
+        else ("non", "trans", "conj")
+    )
     return (
         ApiSpec(
             "gather",
@@ -263,7 +268,7 @@ def registry(modules: dict[str, SourceModule]) -> tuple[ApiSpec, ...]:
         ApiSpec("spmm", "flagsparse_spmm_csr", "spmm_csr", "CSR", "triton", value_const="SUPPORTED_SPMM_VALUE_DTYPES", index_const="SUPPORTED_INDEX_DTYPES", ops=spmm_ops, notes="op supports non/trans/conj; conj on real dtypes is transpose-equivalent"),
         ApiSpec("spmm", "flagsparse_spmm_csr_opt", "spmm_csr", "CSR", "triton_opt", values=("float32", "float64"), index_const="SUPPORTED_INDEX_DTYPES", ops=("non",), notes="bucketed opt path only supports float32/float64"),
         ApiSpec("spmm", "flagsparse_spmm_csr_opt_alg2", "spmm_csr_opt_alg2", "CSR", "triton_opt_alg2", value_const="SUPPORTED_SPMM_OPT_ALG2_DTYPES", index_const="SUPPORTED_INDEX_DTYPES", ops=("non",), notes="hardware-aware alg2 opt path only supports float32/float64"),
-        ApiSpec("spmm", "flagsparse_spmm_coo", "spmm_coo", "COO", "triton", values=spmm_values, index_const="SUPPORTED_INDEX_DTYPES", ops=("non",), notes="COO SpMM reuses CSR SpMM dtype declaration"),
+        ApiSpec("spmm", "flagsparse_spmm_coo", "spmm_coo", "COO", "triton", values=spmm_values, index_const="SUPPORTED_INDEX_DTYPES", ops=spmm_coo_ops, notes="COO SpMM reuses CSR SpMM dtype declaration; op supports non/trans/conj"),
         ApiSpec("spgemm", "flagsparse_spgemm_csr", "spgemm_csr", "CSR", "triton", value_const="SUPPORTED_SPGEMM_VALUE_DTYPES", index_const="SUPPORTED_INDEX_DTYPES", ops=("non",)),
         ApiSpec("sddmm", "flagsparse_sddmm_csr", "sddmm_csr", "CSR", "triton", value_const="SUPPORTED_SDDMM_VALUE_DTYPES", index_const="SUPPORTED_INDEX_DTYPES", ops=("non",)),
         ApiSpec("spsv", "flagsparse_spsv_csr", "spsv", "CSR", "triton", value_const="SUPPORTED_SPSV_VALUE_DTYPES", index_const="SUPPORTED_SPSV_INDEX_DTYPES", ops=("NON_TRANS", "TRANS"), notes="TRANS support is narrower than NON_TRANS; see combo constants"),

src/flagsparse/__init__.py

@@ -9,6 +9,7 @@
     "pytorch_index_scatter",
     "cusparse_spmv_gather",
     "cusparse_spmv_scatter",
+    "cusparse_native_gather",
     "benchmark_gather_case",
     "benchmark_scatter_case",
     "benchmark_performance",
@@ -128,6 +129,7 @@
     "pytorch_index_scatter",
     "cusparse_spmv_gather",
     "cusparse_spmv_scatter",
+    "cusparse_native_gather",
     "benchmark_gather_case",
     "benchmark_scatter_case",
     "benchmark_performance",

src/flagsparse/sparse_operations/__init__.py

@@ -21,6 +21,7 @@
     prepare_alpha_spmm_alg1_tle_opt2,
 )
 from .gather_scatter import (
+    cusparse_native_gather,
     cusparse_spmv_gather,
     cusparse_spmv_scatter,
     flagsparse_gather,
@@ -140,6 +141,7 @@
     "comprehensive_spsm_test",
     "cusparse_spmv_gather",
     "cusparse_spmv_scatter",
+    "cusparse_native_gather",
     "flagsparse_gather",
     "flagsparse_alpha_spmm_alg1",
     "flagsparse_alpha_spmm_alg1_tle",

src/flagsparse/sparse_operations/benchmarks.py

@@ -3,10 +3,11 @@
 from ._common import *
 
 from .gather_scatter import (
+    DEFAULT_GATHER_BLOCK_SIZE,
+    _PreparedCusparseNativeGather,
     SUPPORTED_SCATTER_VALUE_DTYPES,
     _scatter_dtype_error_message,
     _cusparse_spmv,
-    _make_gather_selector_matrix,
     _make_scatter_selector_matrix,
     _pytorch_scatter_impl,
     _triton_gather_impl,
@@ -76,10 +77,10 @@ def benchmark_gather_case(
     index_dtype=torch.int32,
     warmup=20,
     iters=200,
-    block_size=1024,
+    block_size=DEFAULT_GATHER_BLOCK_SIZE,
     run_cusparse=True,
 ):
-    """Benchmark Triton vs PyTorch indexing vs cuSPARSE-backed COO SpMV."""
+    """Benchmark Triton vs PyTorch indexing vs native cuSPARSE gather."""
     device = torch.device("cuda")
     dense_vector = _build_random_dense(dense_size, value_dtype, device)
     indices = _build_indices(nnz, dense_size, index_dtype, device, unique=False)
@@ -88,7 +89,9 @@ def benchmark_gather_case(
     expected = dense_vector[indices]
 
     pytorch_op = lambda: dense_vector[indices]
-    triton_op = lambda: _triton_gather_impl(dense_vector, kernel_indices, block_size=block_size)
+    triton_op = lambda: _triton_gather_impl(
+        dense_vector, kernel_indices, block_size=block_size
+    )
 
     pytorch_values, pytorch_ms = _benchmark_cuda_op(pytorch_op, warmup=warmup, iters=iters)
     triton_values, triton_ms = _benchmark_cuda_op(triton_op, warmup=warmup, iters=iters)
@@ -110,11 +113,13 @@ def benchmark_gather_case(
         if skip_reason:
             cusparse_reason = skip_reason
         else:
+            cusparse_plan = None
             try:
-                selector_matrix = _make_gather_selector_matrix(
-                    indices, dense_vector.numel(), dense_vector.dtype
+                cusparse_out = torch.empty_like(expected)
+                cusparse_plan = _PreparedCusparseNativeGather(
+                    dense_vector, indices, out=cusparse_out
                 )
-                cusparse_op = lambda: _cusparse_spmv(selector_matrix, dense_vector)
+                cusparse_op = lambda: cusparse_plan.run()
                 cusparse_values, cusparse_ms = _benchmark_cuda_op(
                     cusparse_op, warmup=warmup, iters=iters
                 )
@@ -128,6 +133,9 @@ def benchmark_gather_case(
                 )
             except Exception as exc:
                 cusparse_reason = str(exc)
+            finally:
+                if cusparse_plan is not None:
+                    cusparse_plan.close()
 
     triton_speedup_vs_pytorch = (
         pytorch_ms / triton_ms if triton_ms > 0 else float("inf")