From ac27881062e9bf59f4253f0ad1104ad0daa3b047 Mon Sep 17 00:00:00 2001 From: zyq1105331849 <1105331849@qq.com> Date: Sun, 28 Jun 2026 17:02:38 +0800 Subject: [PATCH 1/3] uniform_timing --- .../sparse_operations/alpha_spmm_alg1.py | 209 +++++-- src/flagsparse/sparse_operations/spmm_coo.py | 4 +- src/flagsparse/sparse_operations/spmm_csr.py | 480 ++++++++++++--- .../sparse_operations/spmm_csr_opt_alg2.py | 42 +- tests/test_spmm_coo.py | 568 +++++++++++++++++- tests/test_spmm_csr.py | 14 +- tests/test_spmv_coo.py | 322 ++++++---- tests/test_spmv_opt.py | 227 ++++--- 8 files changed, 1509 insertions(+), 357 deletions(-) diff --git a/src/flagsparse/sparse_operations/alpha_spmm_alg1.py b/src/flagsparse/sparse_operations/alpha_spmm_alg1.py index e3d004e..38ac624 100644 --- a/src/flagsparse/sparse_operations/alpha_spmm_alg1.py +++ b/src/flagsparse/sparse_operations/alpha_spmm_alg1.py @@ -31,7 +31,12 @@ def _load_tle_language(): tle, _TLE_IMPORT_ERROR = _load_tle_language() -SUPPORTED_ALPHA_SPMM_ALG1_DTYPES = (torch.float32, torch.float64) +SUPPORTED_ALPHA_SPMM_ALG1_DTYPES = ( + torch.float32, + torch.float64, + torch.complex64, + torch.complex128, +) _ALPHA_SPMM_ALG1_NUM_WARPS = 8 _ALPHA_SPMM_ALG1_NUM_STAGES = 1 @@ -82,7 +87,10 @@ def _prepare_alpha_spmm_alg1_common(data, indices, indptr, shape): max_row_nnz, ) = _prepare_spmm_csr_matrix(data, indices, indptr, shape) if data.dtype not in SUPPORTED_ALPHA_SPMM_ALG1_DTYPES: - raise TypeError("alpha_spmm_alg1 TLE routes only support float32 and float64") + raise TypeError( + "alpha_spmm_alg1 TLE routes only support float32, float64, " + "complex64, and complex128" + ) return PreparedAlphaSpmmAlg1( data=data, kernel_indices=kernel_indices, @@ -116,7 +124,7 @@ def prepare_alpha_spmm_alg1_tle_opt2(data, indices, indptr, shape): def _alpha_spmm_alg1_acc_dtype(dtype): - return tl.float64 if dtype == torch.float64 else tl.float32 + return tl.float64 if dtype in (torch.float64, torch.complex128) else tl.float32 def _normalize_alpha_spmm_alg1_device_props(device): @@ -418,18 +426,17 @@ def _alpha_spmm_alg1_rowmajor_kernel( tl.store(c_ptr + row * stride_cm + offs3 * stride_cn, acc3, mask=mask3) -def _run_alpha_spmm_alg1(prepared, B, meta): +def _run_alpha_spmm_alg1(prepared, B, meta, out=None): if prepared.n_rows == 0 or int(B.shape[1]) == 0: - return torch.zeros( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, - ) - - C_out = torch.empty( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, + if out is not None: + out.zero_() + return out + return torch.zeros((prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device) + if prepared.data.is_complex(): + return _run_alpha_spmm_alg1_complex_kernel(prepared, B, out) + + C_out = out if out is not None else torch.empty( + (prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device ) grid = (meta["grid_m"], meta["grid_n"]) _alpha_spmm_alg1_rowmajor_kernel[grid]( @@ -905,23 +912,117 @@ def build_alpha_spmm_alg1_tle_opt2_meta(prepared, B): ) -def _run_alpha_spmm_alg1_tle(prepared, B, meta): +@triton.jit +def _alpha_spmm_alg1_complex_rowmajor_kernel( + data_ri_ptr, + indices_ptr, + indptr_ptr, + b_ri_ptr, + c_ri_ptr, + n_rows, + n_dense_cols, + stride_bk, + stride_bn, + stride_br, + stride_cm, + stride_cn, + stride_cr, + BLOCK_N: tl.constexpr, + BLOCK_NNZ: tl.constexpr, + ACC_DTYPE: tl.constexpr, +): + row = tl.program_id(0) + pid_n = tl.program_id(1) + if row >= n_rows: + return + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_n < n_dense_cols + start = tl.load(indptr_ptr + row) + end = tl.load(indptr_ptr + row + 1) + row_nnz = end - start + acc_re = tl.zeros([BLOCK_N], dtype=ACC_DTYPE) + acc_im = tl.zeros([BLOCK_N], dtype=ACC_DTYPE) + + for chunk_start in tl.range(0, row_nnz, BLOCK_NNZ): + for kk in tl.static_range(0, BLOCK_NNZ): + idx = start + chunk_start + kk + valid = idx < end + a_re = tl.load(data_ri_ptr + idx * 2, mask=valid, other=0.0) + a_im = tl.load(data_ri_ptr + idx * 2 + 1, mask=valid, other=0.0) + a_col = tl.load(indices_ptr + idx, mask=valid, other=0) + b_re = tl.load( + b_ri_ptr + a_col * stride_bk + offs_n * stride_bn, + mask=mask_n & valid, + other=0.0, + ) + b_im = tl.load( + b_ri_ptr + a_col * stride_bk + offs_n * stride_bn + stride_br, + mask=mask_n & valid, + other=0.0, + ) + acc_re = acc_re + a_re.to(ACC_DTYPE) * b_re.to(ACC_DTYPE) - a_im.to(ACC_DTYPE) * b_im.to(ACC_DTYPE) + acc_im = acc_im + a_re.to(ACC_DTYPE) * b_im.to(ACC_DTYPE) + a_im.to(ACC_DTYPE) * b_re.to(ACC_DTYPE) + + tl.store(c_ri_ptr + row * stride_cm + offs_n * stride_cn, acc_re, mask=mask_n) + tl.store( + c_ri_ptr + row * stride_cm + offs_n * stride_cn + stride_cr, + acc_im, + mask=mask_n, + ) + + +def _run_alpha_spmm_alg1_complex_kernel(prepared, B, out): + if out is None: + C_out = torch.empty((prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device) + else: + C_out = out + data_ri = torch.view_as_real(prepared.data).contiguous().reshape(-1) + B_ri = torch.view_as_real(B) + C_ri = torch.view_as_real(C_out) + block_n = 32 + block_nnz = 256 + grid = (prepared.n_rows, triton.cdiv(int(B.shape[1]), block_n)) + acc_dtype = tl.float64 if B_ri.dtype == torch.float64 else tl.float32 + _alpha_spmm_alg1_complex_rowmajor_kernel[grid]( + data_ri, + prepared.kernel_indices, + prepared.kernel_indptr, + B_ri, + C_ri, + prepared.n_rows, + int(B.shape[1]), + B_ri.stride(0), + B_ri.stride(1), + B_ri.stride(2), + C_ri.stride(0), + C_ri.stride(1), + C_ri.stride(2), + BLOCK_N=block_n, + BLOCK_NNZ=block_nnz, + ACC_DTYPE=acc_dtype, + num_warps=4, + num_stages=3, + ) + return C_out + + +def _run_alpha_spmm_alg1_tle(prepared, B, meta, out=None): if tle is None: raise RuntimeError( "flagsparse_alpha_spmm_alg1_tle requires FlagTree/TLE-Struct runtime " f"support ({alpha_spmm_alg1_tle_unavailable_reason()})" ) if prepared.n_rows == 0 or int(B.shape[1]) == 0: - return torch.zeros( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, - ) - - C_out = torch.empty( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, + if out is not None: + out.zero_() + return out + return torch.zeros((prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device) + if prepared.data.is_complex(): + return _run_alpha_spmm_alg1_complex_kernel(prepared, B, out) + + C_out = out if out is not None else torch.empty( + (prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device ) grid = (meta["grid_m"], meta["grid_n"]) _alpha_spmm_alg1_tle_rowmajor_kernel[grid]( @@ -948,23 +1049,22 @@ def _run_alpha_spmm_alg1_tle(prepared, B, meta): return C_out -def _run_alpha_spmm_alg1_tle_opt(prepared, B, meta): +def _run_alpha_spmm_alg1_tle_opt(prepared, B, meta, out=None): if tle is None: raise RuntimeError( "flagsparse_alpha_spmm_alg1_tle_opt requires FlagTree/TLE-Struct runtime " f"support ({alpha_spmm_alg1_tle_opt_unavailable_reason()})" ) if prepared.n_rows == 0 or int(B.shape[1]) == 0: - return torch.zeros( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, - ) - - C_out = torch.empty( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, + if out is not None: + out.zero_() + return out + return torch.zeros((prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device) + if prepared.data.is_complex(): + return _run_alpha_spmm_alg1_complex_kernel(prepared, B, out) + + C_out = out if out is not None else torch.empty( + (prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device ) grid = (meta["grid_m"], meta["grid_n"]) _alpha_spmm_alg1_tle_opt_rowmajor_kernel[grid]( @@ -991,23 +1091,22 @@ def _run_alpha_spmm_alg1_tle_opt(prepared, B, meta): return C_out -def _run_alpha_spmm_alg1_tle_opt2(prepared, B, meta): +def _run_alpha_spmm_alg1_tle_opt2(prepared, B, meta, out=None): if tle is None: raise RuntimeError( "flagsparse_alpha_spmm_alg1_tle_opt2 requires FlagTree/TLE-Struct runtime " f"support ({alpha_spmm_alg1_tle_opt2_unavailable_reason()})" ) if prepared.n_rows == 0 or int(B.shape[1]) == 0: - return torch.zeros( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, - ) - - C_out = torch.empty( - (prepared.n_rows, int(B.shape[1])), - dtype=prepared.data.dtype, - device=prepared.data.device, + if out is not None: + out.zero_() + return out + return torch.zeros((prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device) + if prepared.data.is_complex(): + return _run_alpha_spmm_alg1_complex_kernel(prepared, B, out) + + C_out = out if out is not None else torch.empty( + (prepared.n_rows, int(B.shape[1])), dtype=prepared.data.dtype, device=prepared.data.device ) grid = (meta["grid_m"], meta["grid_n"]) _alpha_spmm_alg1_tle_opt2_rowmajor_kernel[grid]( @@ -1062,8 +1161,8 @@ def flagsparse_alpha_spmm_alg1( _build_alpha_spmm_alg1_runtime_meta(prepared, B), "alpha_spmm_alg1", ) - C = _run_alpha_spmm_alg1(prepared, B, meta) - if out is not None: + C = _run_alpha_spmm_alg1(prepared, B, meta, out=out) + if out is not None and C is not out: out.copy_(C) C = out if return_meta: @@ -1099,8 +1198,8 @@ def flagsparse_alpha_spmm_alg1_tle( _build_alpha_spmm_alg1_runtime_meta(prepared, B), "alpha_spmm_alg1_tle", ) - C = _run_alpha_spmm_alg1_tle(prepared, B, meta) - if out is not None: + C = _run_alpha_spmm_alg1_tle(prepared, B, meta, out=out) + if out is not None and C is not out: out.copy_(C) C = out if return_meta: @@ -1136,8 +1235,8 @@ def flagsparse_alpha_spmm_alg1_tle_opt( if meta is None: meta = _build_alpha_spmm_alg1_tle_opt_runtime_meta(prepared, B) meta = _with_alpha_spmm_alg1_route(meta, "alpha_spmm_alg1_tle_opt") - C = _run_alpha_spmm_alg1_tle_opt(prepared, B, meta) - if out is not None: + C = _run_alpha_spmm_alg1_tle_opt(prepared, B, meta, out=out) + if out is not None and C is not out: out.copy_(C) C = out if return_meta: @@ -1173,8 +1272,8 @@ def flagsparse_alpha_spmm_alg1_tle_opt2( if meta is None: meta = _build_alpha_spmm_alg1_tle_opt2_runtime_meta(prepared, B) meta = _with_alpha_spmm_alg1_route(meta, "alpha_spmm_alg1_tle_opt2") - C = _run_alpha_spmm_alg1_tle_opt2(prepared, B, meta) - if out is not None: + C = _run_alpha_spmm_alg1_tle_opt2(prepared, B, meta, out=out) + if out is not None and C is not out: out.copy_(C) C = out if return_meta: diff --git a/src/flagsparse/sparse_operations/spmm_coo.py b/src/flagsparse/sparse_operations/spmm_coo.py index b0c4913..1b527fa 100644 --- a/src/flagsparse/sparse_operations/spmm_coo.py +++ b/src/flagsparse/sparse_operations/spmm_coo.py @@ -558,6 +558,7 @@ def _triton_spmm_coo_rowrun_impl( output_dtype, out=None, dense_layout="row", + seg_starts=None, ): device = data.device dtype = data.dtype @@ -573,7 +574,8 @@ def _triton_spmm_coo_rowrun_impl( return out return _zeros_dense_layout((n_rows, n_dense_cols), output_dtype, device, dense_layout) - seg_starts = _seg_starts_from_sorted_rows(row, int(data.numel()), device) + if seg_starts is None: + seg_starts = _seg_starts_from_sorted_rows(row, int(data.numel()), device) n_segs = int(seg_starts.numel()) - 1 if seg_starts is not None else 0 if n_segs == 0: if out is not None: diff --git a/src/flagsparse/sparse_operations/spmm_csr.py b/src/flagsparse/sparse_operations/spmm_csr.py index 42b2e71..e03a5c8 100644 --- a/src/flagsparse/sparse_operations/spmm_csr.py +++ b/src/flagsparse/sparse_operations/spmm_csr.py @@ -245,6 +245,68 @@ def _spmm_csr_complex_kernel( mask=mask_n, ) + +@triton.jit +def _spmm_csr_selected_rows_complex_kernel( + data_ri_ptr, + indices_ptr, + indptr_ptr, + b_ri_ptr, + c_ri_ptr, + rows_ptr, + n_bucket_rows, + n_dense_cols, + stride_bk, + stride_bn, + stride_br, + stride_cm, + stride_cn, + stride_cr, + BLOCK_N: tl.constexpr, + BLOCK_NNZ: tl.constexpr, + ACC_DTYPE: tl.constexpr, +): + pid_row = tl.program_id(0) + pid_n = tl.program_id(1) + if pid_row >= n_bucket_rows: + return + + row = tl.load(rows_ptr + pid_row) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_n < n_dense_cols + start = tl.load(indptr_ptr + row) + end = tl.load(indptr_ptr + row + 1) + row_nnz = end - start + acc_re = tl.zeros([BLOCK_N], dtype=ACC_DTYPE) + acc_im = tl.zeros([BLOCK_N], dtype=ACC_DTYPE) + + for chunk_start in tl.range(0, row_nnz, BLOCK_NNZ): + for kk in tl.static_range(0, BLOCK_NNZ): + idx = start + chunk_start + kk + valid = idx < end + a_re = tl.load(data_ri_ptr + idx * 2, mask=valid, other=0.0) + a_im = tl.load(data_ri_ptr + idx * 2 + 1, mask=valid, other=0.0) + a_col = tl.load(indices_ptr + idx, mask=valid, other=0) + b_re = tl.load( + b_ri_ptr + a_col * stride_bk + offs_n * stride_bn, + mask=mask_n & valid, + other=0.0, + ) + b_im = tl.load( + b_ri_ptr + a_col * stride_bk + offs_n * stride_bn + stride_br, + mask=mask_n & valid, + other=0.0, + ) + acc_re = acc_re + a_re.to(ACC_DTYPE) * b_re.to(ACC_DTYPE) - a_im.to(ACC_DTYPE) * b_im.to(ACC_DTYPE) + acc_im = acc_im + a_re.to(ACC_DTYPE) * b_im.to(ACC_DTYPE) + a_im.to(ACC_DTYPE) * b_re.to(ACC_DTYPE) + + tl.store(c_ri_ptr + row * stride_cm + offs_n * stride_cn, acc_re, mask=mask_n) + tl.store( + c_ri_ptr + row * stride_cm + offs_n * stride_cn + stride_cr, + acc_im, + mask=mask_n, + ) + def _prepare_spmm_csr_matrix(data, indices, indptr, shape): if len(shape) != 2: raise ValueError("shape must be a 2-tuple: (n_rows, n_cols)") @@ -732,14 +794,12 @@ def _run_spmm_csr_base_route(prepared, B, *, timing=False, diagnostics=False, de ) -def _run_spmm_csr_base_accuracy_route(prepared, B, *, timing=False, diagnostics=False): - if prepared.op != "non": - raise SpmmCsrAlgorithmUnavailable("csr_base_accuracy currently supports op=non only") - if prepared.data.dtype != torch.float32: - raise SpmmCsrAlgorithmUnavailable("csr_base_accuracy currently supports float32 only") +def _run_spmm_csr_base_accuracy_route( + prepared, B, *, timing=False, diagnostics=False, dense_layout="row" +): return _run_spmm_csr_base_route_impl( prepared, B, timing=timing, diagnostics=diagnostics, - dense_layout="row", accuracy=True, route_name="csr_base_accuracy", + dense_layout=dense_layout, accuracy=True, route_name="csr_base_accuracy", ) @@ -849,8 +909,17 @@ def _run_alpha_spmm_alg1_tle_route( run_fn_name, timing=False, diagnostics=False, + dense_layout="row", ): B = _validate_spmm_route_runtime_inputs(prepared, B) + dense_layout = _normalize_dense_layout(dense_layout) + B = _materialize_dense_layout(B, dense_layout) + C_out = _empty_dense_layout( + (prepared.n_rows, int(B.shape[1])), + prepared.data.dtype, + prepared.data.device, + dense_layout, + ) from . import alpha_spmm_alg1 as alpha_mod availability_fn = getattr(alpha_mod, availability_fn_name) unavailable_reason_fn = getattr(alpha_mod, unavailable_reason_fn_name) @@ -869,7 +938,7 @@ def _run_alpha_spmm_alg1_tle_route( start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() - C = run_fn(B=B, prepared=alpha_prepared, meta=launch_meta) + C = run_fn(B=B, prepared=alpha_prepared, meta=launch_meta, out=C_out) if timing: end.record() torch.cuda.synchronize() @@ -881,6 +950,10 @@ def _run_alpha_spmm_alg1_tle_route( "process_cpu_ms": process_cpu_ms, "process_gpu_ms": process_gpu_ms, "compute_ms": compute_ms, + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } if diagnostics: meta["diagnostics"] = { @@ -899,11 +972,17 @@ def _run_alpha_spmm_alg1_tle_route( "block_cols": launch_meta.get("block_cols"), "grid_m": launch_meta.get("grid_m"), "grid_n": launch_meta.get("grid_n"), + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } return C, meta -def _run_alpha_spmm_alg1_tle_opt_route(prepared, B, *, timing=False, diagnostics=False): +def _run_alpha_spmm_alg1_tle_opt_route( + prepared, B, *, timing=False, diagnostics=False, dense_layout="row" +): return _run_alpha_spmm_alg1_tle_route( prepared, B, @@ -914,10 +993,13 @@ def _run_alpha_spmm_alg1_tle_opt_route(prepared, B, *, timing=False, diagnostics run_fn_name="flagsparse_alpha_spmm_alg1_tle_opt", timing=timing, diagnostics=diagnostics, + dense_layout=dense_layout, ) -def _run_alpha_spmm_alg1_tle_opt2_route(prepared, B, *, timing=False, diagnostics=False): +def _run_alpha_spmm_alg1_tle_opt2_route( + prepared, B, *, timing=False, diagnostics=False, dense_layout="row" +): return _run_alpha_spmm_alg1_tle_route( prepared, B, @@ -928,6 +1010,7 @@ def _run_alpha_spmm_alg1_tle_opt2_route(prepared, B, *, timing=False, diagnostic run_fn_name="flagsparse_alpha_spmm_alg1_tle_opt2", timing=timing, diagnostics=diagnostics, + dense_layout=dense_layout, ) @@ -1347,7 +1430,7 @@ def _spmm_csr_alg1_run_split_bucket(plan, B, C_out, block_n, device_props): ) -def _spmm_csr_alg1_compute(plan, B): +def _spmm_csr_alg1_compute(plan, B, *, out=None, dense_layout="row"): if B.ndim != 2: raise ValueError("B must be a 2D dense tensor") if not B.is_cuda: @@ -1358,9 +1441,39 @@ def _spmm_csr_alg1_compute(plan, B): raise TypeError("B dtype must match sparse matrix dtype") if B.shape[0] != plan.n_cols: raise ValueError(f"B.shape[0] must be n_cols={plan.n_cols}, got {B.shape[0]}") - B = B.contiguous() + dense_layout = _normalize_dense_layout(dense_layout) + B = _materialize_dense_layout(B, dense_layout) block_n = _select_spmm_opt_block_n(int(B.shape[1])) - C_out = torch.zeros((plan.n_rows, int(B.shape[1])), dtype=B.dtype, device=B.device) + if _is_complex_dtype(B.dtype): + C_out = out if out is not None else _empty_dense_layout( + (plan.n_rows, int(B.shape[1])), + B.dtype, + B.device, + dense_layout, + ) + C_out.zero_() + for bucket in plan.row_buckets: + rows = plan.long_row_ids if bucket["kind"] == "split" else bucket["rows"] + _spmm_csr_run_selected_rows_complex( + plan.data, + plan.kernel_indices, + plan.kernel_indptr, + B, + C_out, + rows, + block_n=block_n, + block_nnz=int(bucket.get("block_nnz", 256)), + num_warps=4, + num_stages=3, + ) + return C_out + C_out = out if out is not None else _empty_dense_layout( + (plan.n_rows, int(B.shape[1])), + B.dtype, + B.device, + dense_layout, + ) + C_out.zero_() device_props = _normalize_spmm_opt_device_props(plan.data.device) for bucket in plan.row_buckets: if bucket["kind"] == "split": @@ -1370,15 +1483,23 @@ def _spmm_csr_alg1_compute(plan, B): return C_out -def _run_spmm_csr_alg1_route(prepared, B, *, timing=False, diagnostics=False): +def _run_spmm_csr_alg1_route(prepared, B, *, timing=False, diagnostics=False, dense_layout="row"): B = _validate_spmm_route_runtime_inputs(prepared, B) + dense_layout = _normalize_dense_layout(dense_layout) + B = _materialize_dense_layout(B, dense_layout) + C_out = _empty_dense_layout( + (prepared.n_rows, int(B.shape[1])), + prepared.data.dtype, + prepared.data.device, + dense_layout, + ) plan = _spmm_csr_alg1_build_process_plan(prepared, timing=bool(timing)) compute_ms = None if timing: start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() - C = _spmm_csr_alg1_compute(plan, B) + C = _spmm_csr_alg1_compute(plan, B, out=C_out, dense_layout=dense_layout) if timing: end.record() torch.cuda.synchronize() @@ -1390,6 +1511,10 @@ def _run_spmm_csr_alg1_route(prepared, B, *, timing=False, diagnostics=False): "process_cpu_ms": plan.process_cpu_ms, "process_gpu_ms": plan.process_gpu_ms if timing else None, "compute_ms": compute_ms, + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } if diagnostics: meta["diagnostics"] = { @@ -1403,6 +1528,10 @@ def _run_spmm_csr_alg1_route(prepared, B, *, timing=False, diagnostics=False): "num_warps": None, "num_stages": None, "long_part_count": plan.long_part_count, + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } return C, meta @@ -1865,6 +1994,7 @@ def _spmm_csr_alg2_run_bucket( batched_kernel, row_kernel, segmented_kernel = kernels acc_dtype = tl.float64 if high_precision else _spmm_csr_alg2_acc_dtype(plan.data.dtype) out_dtype = tl.float64 if plan.data.dtype == torch.float64 else tl.float32 + accuracy_kernel = bool(accuracy and plan.data.dtype == torch.float32) common_kwargs = { "num_warps": launch["num_warps"], "num_stages": launch["num_stages"], @@ -1893,7 +2023,7 @@ def _spmm_csr_alg2_run_bucket( BLOCK_K=launch["block_k"], ACC_DTYPE=acc_dtype, OUT_DTYPE=out_dtype, - ACCURACY=bool(accuracy), + ACCURACY=accuracy_kernel, **common_kwargs, ) elif bucket["kind"] == "row2d": @@ -1915,7 +2045,7 @@ def _spmm_csr_alg2_run_bucket( BLOCK_K=launch["block_k"], ACC_DTYPE=acc_dtype, OUT_DTYPE=out_dtype, - ACCURACY=bool(accuracy), + ACCURACY=accuracy_kernel, **common_kwargs, ) else: @@ -1938,7 +2068,7 @@ def _spmm_csr_alg2_run_bucket( SEGMENTS=launch["segments"], ACC_DTYPE=acc_dtype, OUT_DTYPE=out_dtype, - ACCURACY=bool(accuracy), + ACCURACY=accuracy_kernel, **common_kwargs, ) @@ -1947,7 +2077,9 @@ def _spmm_csr_alg2_run_bucket( return launch -def _spmm_csr_alg2_compute(plan, B, *, accuracy=False, high_precision=False): +def _spmm_csr_alg2_compute( + plan, B, *, accuracy=False, high_precision=False, out=None, dense_layout="row" +): if B.ndim != 2: raise ValueError("B must be a 2D dense tensor") if not B.is_cuda: @@ -1959,8 +2091,50 @@ def _spmm_csr_alg2_compute(plan, B, *, accuracy=False, high_precision=False): if B.shape[0] != plan.n_cols: raise ValueError(f"B.shape[0] must be n_cols={plan.n_cols}, got {B.shape[0]}") - B = B.contiguous() - C_out = torch.zeros((plan.n_rows, int(B.shape[1])), dtype=B.dtype, device=B.device) + dense_layout = _normalize_dense_layout(dense_layout) + B = _materialize_dense_layout(B, dense_layout) + if _is_complex_dtype(B.dtype): + C_out = out if out is not None else _empty_dense_layout( + (plan.n_rows, int(B.shape[1])), + B.dtype, + B.device, + dense_layout, + ) + C_out.zero_() + device_props = _normalize_spmm_csr_alg2_device_props(plan.data.device) + plan.launch_configs.clear() + for bucket in plan.row_buckets: + launch = _resolve_spmm_csr_alg2_launch( + bucket, + int(B.shape[1]), + plan.data.dtype, + device_props, + ) + rows = bucket["rows"] + _spmm_csr_run_selected_rows_complex( + plan.data, + plan.kernel_indices, + plan.kernel_indptr, + B, + C_out, + rows, + block_n=launch["block_n"], + block_nnz=launch["block_k"], + num_warps=launch["num_warps"], + num_stages=launch["num_stages"], + high_precision=bool(high_precision), + ) + launch["grid_m"] = int(rows.numel()) + launch["grid_n"] = int(triton.cdiv(B.shape[1], launch["block_n"])) + plan.launch_configs.append(launch) + return C_out + C_out = out if out is not None else _empty_dense_layout( + (plan.n_rows, int(B.shape[1])), + B.dtype, + B.device, + dense_layout, + ) + C_out.zero_() device_props = _normalize_spmm_csr_alg2_device_props(plan.data.device) kernels = _spmm_csr_alg2_kernel_bundle() plan.launch_configs.clear() @@ -1973,15 +2147,23 @@ def _spmm_csr_alg2_compute(plan, B, *, accuracy=False, high_precision=False): return C_out -def _run_spmm_csr_alg2_route(prepared, B, *, timing=False, diagnostics=False): +def _run_spmm_csr_alg2_route(prepared, B, *, timing=False, diagnostics=False, dense_layout="row"): B = _validate_spmm_route_runtime_inputs(prepared, B) + dense_layout = _normalize_dense_layout(dense_layout) + B = _materialize_dense_layout(B, dense_layout) + C_out = _empty_dense_layout( + (prepared.n_rows, int(B.shape[1])), + prepared.data.dtype, + prepared.data.device, + dense_layout, + ) plan = _spmm_csr_alg2_build_process_plan(prepared, timing=bool(timing)) compute_ms = None if timing: start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() - C = _spmm_csr_alg2_compute(plan, B) + C = _spmm_csr_alg2_compute(plan, B, out=C_out, dense_layout=dense_layout) if timing: end.record() torch.cuda.synchronize() @@ -1993,6 +2175,10 @@ def _run_spmm_csr_alg2_route(prepared, B, *, timing=False, diagnostics=False): "process_cpu_ms": plan.process_cpu_ms, "process_gpu_ms": plan.process_gpu_ms if timing else None, "compute_ms": compute_ms, + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } if diagnostics: first_launch = plan.launch_configs[0] if plan.launch_configs else {} @@ -2009,40 +2195,59 @@ def _run_spmm_csr_alg2_route(prepared, B, *, timing=False, diagnostics=False): "num_stages": first_launch.get("num_stages"), "grid_m": first_launch.get("grid_m"), "grid_n": first_launch.get("grid_n"), + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } return C, meta -def _run_spmm_csr_alg2_accuracy_route(prepared, B, *, timing=False, diagnostics=False): +def _run_spmm_csr_alg2_accuracy_route( + prepared, B, *, timing=False, diagnostics=False, dense_layout="row" +): return _run_spmm_csr_alg2_accuracy_impl( prepared, B, timing=timing, diagnostics=diagnostics, - high_precision=False, route_name="spmm_csr_alg2_accuracy", + dense_layout=dense_layout, high_precision=False, route_name="spmm_csr_alg2_accuracy", ) -def _run_spmm_csr_alg2_accuracy_hp_route(prepared, B, *, timing=False, diagnostics=False): +def _run_spmm_csr_alg2_accuracy_hp_route( + prepared, B, *, timing=False, diagnostics=False, dense_layout="row" +): return _run_spmm_csr_alg2_accuracy_impl( prepared, B, timing=timing, diagnostics=diagnostics, - high_precision=True, route_name="spmm_csr_alg2_accuracy_hp", + dense_layout=dense_layout, high_precision=True, route_name="spmm_csr_alg2_accuracy_hp", ) def _run_spmm_csr_alg2_accuracy_impl( - prepared, B, *, timing=False, diagnostics=False, + prepared, B, *, timing=False, diagnostics=False, dense_layout="row", high_precision=False, route_name="spmm_csr_alg2_accuracy", ): - if prepared.op != "non": - raise SpmmCsrAlgorithmUnavailable(f"{route_name} currently supports op=non only") - if prepared.data.dtype != torch.float32: - raise SpmmCsrAlgorithmUnavailable(f"{route_name} currently supports float32 only") B = _validate_spmm_route_runtime_inputs(prepared, B) + dense_layout = _normalize_dense_layout(dense_layout) + B = _materialize_dense_layout(B, dense_layout) + C_out = _empty_dense_layout( + (prepared.n_rows, int(B.shape[1])), + prepared.data.dtype, + prepared.data.device, + dense_layout, + ) plan = _spmm_csr_alg2_build_process_plan(prepared, timing=bool(timing)) compute_ms = None if timing: start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() - C = _spmm_csr_alg2_compute(plan, B, accuracy=True, high_precision=high_precision) + C = _spmm_csr_alg2_compute( + plan, + B, + accuracy=True, + high_precision=high_precision, + out=C_out, + dense_layout=dense_layout, + ) if timing: end.record() torch.cuda.synchronize() @@ -2054,6 +2259,10 @@ def _run_spmm_csr_alg2_accuracy_impl( "process_cpu_ms": plan.process_cpu_ms, "process_gpu_ms": plan.process_gpu_ms if timing else None, "compute_ms": compute_ms, + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } if diagnostics: first_launch = plan.launch_configs[0] if plan.launch_configs else {} @@ -2073,6 +2282,10 @@ def _run_spmm_csr_alg2_accuracy_impl( "num_stages": first_launch.get("num_stages"), "grid_m": first_launch.get("grid_m"), "grid_n": first_launch.get("grid_n"), + "dense_layout": dense_layout, + "b_stride": tuple(int(v) for v in B.stride()), + "c_stride": tuple(int(v) for v in C.stride()), + "output_layout": _dense_layout_name(C), } return C, meta @@ -2088,50 +2301,50 @@ def _run_spmm_csr_alg2_accuracy_impl( "csr_base_accuracy": SpmmCsrAlgorithm( name="csr_base_accuracy", display_name="BaseAccuracy", - supported_ops=("non",), - supported_dtypes=(torch.float32,), + supported_ops=tuple(SPMM_OP_NAMES.values()), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_spmm_csr_base_accuracy_route, ), "alpha_alg1_tle_opt": SpmmCsrAlgorithm( name="alpha_alg1_tle_opt", display_name="TLEOpt", supported_ops=tuple(SPMM_OP_NAMES.values()), - supported_dtypes=(torch.float32, torch.float64), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_alpha_spmm_alg1_tle_opt_route, ), "alpha_alg1_tle_opt2": SpmmCsrAlgorithm( name="alpha_alg1_tle_opt2", display_name="TLEOpt2", supported_ops=tuple(SPMM_OP_NAMES.values()), - supported_dtypes=(torch.float32, torch.float64), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_alpha_spmm_alg1_tle_opt2_route, ), "spmm_csr_alg1": SpmmCsrAlgorithm( name="spmm_csr_alg1", display_name="Alg1", supported_ops=tuple(SPMM_OP_NAMES.values()), - supported_dtypes=(torch.float32, torch.float64), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_spmm_csr_alg1_route, ), "spmm_csr_alg2": SpmmCsrAlgorithm( name="spmm_csr_alg2", display_name="Alg2", supported_ops=tuple(SPMM_OP_NAMES.values()), - supported_dtypes=(torch.float32, torch.float64), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_spmm_csr_alg2_route, ), "spmm_csr_alg2_accuracy": SpmmCsrAlgorithm( name="spmm_csr_alg2_accuracy", display_name="Alg2Accuracy", - supported_ops=("non",), - supported_dtypes=(torch.float32,), + supported_ops=tuple(SPMM_OP_NAMES.values()), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_spmm_csr_alg2_accuracy_route, ), "spmm_csr_alg2_accuracy_hp": SpmmCsrAlgorithm( name="spmm_csr_alg2_accuracy_hp", display_name="Alg2AccuracyHP", - supported_ops=("non",), - supported_dtypes=(torch.float32,), + supported_ops=tuple(SPMM_OP_NAMES.values()), + supported_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128), run=_run_spmm_csr_alg2_accuracy_hp_route, ), } @@ -2220,9 +2433,19 @@ def flagsparse_spmm_csr_run( alg_name = prepared.alg if alg is None else _normalize_spmm_csr_alg(alg) algorithm = resolve_spmm_csr_algorithm(alg_name, op_name, prepared.data.dtype) dense_layout = _normalize_dense_layout(dense_layout) - if dense_layout == "col" and algorithm.name != "csr_base": + col_major_algorithms = { + "csr_base", + "csr_base_accuracy", + "alpha_alg1_tle_opt", + "alpha_alg1_tle_opt2", + "spmm_csr_alg1", + "spmm_csr_alg2", + "spmm_csr_alg2_accuracy", + "spmm_csr_alg2_accuracy_hp", + } + if dense_layout == "col" and algorithm.name not in col_major_algorithms: raise SpmmCsrAlgorithmUnavailable( - "col-major layout is currently supported only by csr_base" + "col-major layout is currently supported only by CSR SpMM main algorithms" ) if B is None or not torch.is_tensor(B): raise TypeError("B must be a torch.Tensor") @@ -2245,21 +2468,13 @@ def flagsparse_spmm_csr_run( op_name, timing=bool(timing), ) - if algorithm.name == "csr_base": - C, route_meta = algorithm.run( - runtime_prepared, - B, - timing=bool(timing), - diagnostics=bool(diagnostics), - dense_layout=dense_layout, - ) - else: - C, route_meta = algorithm.run( - runtime_prepared, - B, - timing=bool(timing), - diagnostics=bool(diagnostics), - ) + C, route_meta = algorithm.run( + runtime_prepared, + B, + timing=bool(timing), + diagnostics=bool(diagnostics), + dense_layout=dense_layout, + ) if end is not None: end.record() torch.cuda.synchronize() @@ -2809,6 +3024,115 @@ def _build_spmm_csr_opt_runtime_symbolic_triton(prepared): ) +def _triton_spmm_csr_complex_impl( + data, + indices, + indptr, + B, + n_rows, + n_dense_cols, + block_n, + block_nnz, + num_warps, + num_stages, + out=None, + dense_layout="row", + high_precision=False, +): + dtype = data.dtype + device = data.device + dense_layout = _normalize_dense_layout(dense_layout) + if dtype not in (torch.complex64, torch.complex128): + raise TypeError("_triton_spmm_csr_complex_impl requires a complex dtype") + if out is not None: + if out.shape != (int(n_rows), int(n_dense_cols)) or out.dtype != dtype: + raise ValueError("out shape/dtype must match result") + if out.device != device: + raise ValueError("out must be on the same CUDA device as data") + if n_rows == 0 or n_dense_cols == 0 or B.shape[0] == 0: + if out is not None: + out.zero_() + return out + return _zeros_dense_layout((n_rows, n_dense_cols), dtype, device, dense_layout) + + data_ri = torch.view_as_real(data).contiguous().reshape(-1) + B_ri = torch.view_as_real(B) + C_complex = out if out is not None else _empty_dense_layout( + (n_rows, n_dense_cols), + dtype, + device, + dense_layout, + ) + C_ri = torch.view_as_real(C_complex) + grid = (n_rows, triton.cdiv(n_dense_cols, block_n)) + acc_dtype = tl.float64 if (B_ri.dtype == torch.float64 or high_precision) else tl.float32 + _spmm_csr_complex_kernel[grid]( + data_ri, + indices, + indptr, + B_ri, + C_ri, + n_rows, + n_dense_cols, + B_ri.stride(0), + B_ri.stride(1), + B_ri.stride(2), + C_ri.stride(0), + C_ri.stride(1), + C_ri.stride(2), + BLOCK_N=block_n, + BLOCK_NNZ=block_nnz, + ACC_DTYPE=acc_dtype, + num_warps=num_warps, + num_stages=num_stages, + ) + return C_complex + + +def _spmm_csr_run_selected_rows_complex( + data, + indices, + indptr, + B, + C_out, + rows, + *, + block_n, + block_nnz, + num_warps, + num_stages, + high_precision=False, +): + if rows.numel() == 0: + return + data_ri = torch.view_as_real(data).contiguous().reshape(-1) + B_ri = torch.view_as_real(B) + C_ri = torch.view_as_real(C_out) + acc_dtype = tl.float64 if (B_ri.dtype == torch.float64 or high_precision) else tl.float32 + grid = (rows.numel(), triton.cdiv(int(B.shape[1]), block_n)) + _spmm_csr_selected_rows_complex_kernel[grid]( + data_ri, + indices, + indptr, + B_ri, + C_ri, + rows, + rows.numel(), + int(B.shape[1]), + B_ri.stride(0), + B_ri.stride(1), + B_ri.stride(2), + C_ri.stride(0), + C_ri.stride(1), + C_ri.stride(2), + BLOCK_N=block_n, + BLOCK_NNZ=block_nnz, + ACC_DTYPE=acc_dtype, + num_warps=num_warps, + num_stages=num_stages, + ) + + def _triton_spmm_csr_impl( data, indices, @@ -2869,7 +3193,7 @@ def _triton_spmm_csr_impl( BLOCK_N=block_n, BLOCK_NNZ=block_nnz, ACC_DTYPE=acc_dtype, - ACCURACY=bool(accuracy), + ACCURACY=bool(accuracy and compute_dtype == torch.float32), num_warps=num_warps, num_stages=num_stages, ) @@ -2885,38 +3209,20 @@ def _triton_spmm_csr_impl( return C_cast return C_compute - data_ri = torch.view_as_real(data).contiguous().reshape(-1) - B_ri = torch.view_as_real(B) - C_complex = out if out is not None else _empty_dense_layout( - (n_rows, n_dense_cols), - dtype, - device, - dense_layout, - ) - C_ri = torch.view_as_real(C_complex) - grid = (n_rows, triton.cdiv(n_dense_cols, block_n)) - acc_dtype = tl.float64 if B_ri.dtype == torch.float64 else tl.float32 - _spmm_csr_complex_kernel[grid]( - data_ri, + return _triton_spmm_csr_complex_impl( + data, indices, indptr, - B_ri, - C_ri, + B, n_rows, n_dense_cols, - B_ri.stride(0), - B_ri.stride(1), - B_ri.stride(2), - C_ri.stride(0), - C_ri.stride(1), - C_ri.stride(2), - BLOCK_N=block_n, - BLOCK_NNZ=block_nnz, - ACC_DTYPE=acc_dtype, - num_warps=num_warps, - num_stages=num_stages, + block_n, + block_nnz, + num_warps, + num_stages, + out=out, + dense_layout=dense_layout, ) - return C_complex @triton.jit diff --git a/src/flagsparse/sparse_operations/spmm_csr_opt_alg2.py b/src/flagsparse/sparse_operations/spmm_csr_opt_alg2.py index e0eb212..21e927f 100644 --- a/src/flagsparse/sparse_operations/spmm_csr_opt_alg2.py +++ b/src/flagsparse/sparse_operations/spmm_csr_opt_alg2.py @@ -8,13 +8,19 @@ _prepare_spmm_csr_matrix, _spmm_coo_reference_tolerance, _spmm_validation_metrics, + _triton_spmm_csr_complex_impl, flagsparse_spmm_csr, flagsparse_spmm_csr_opt, prepare_spmm_csr_opt, ) -SUPPORTED_SPMM_OPT_ALG2_DTYPES = (torch.float32, torch.float64) +SUPPORTED_SPMM_OPT_ALG2_DTYPES = ( + torch.float32, + torch.float64, + torch.complex64, + torch.complex128, +) class PreparedCsrSpmmOptAlg2: @@ -903,7 +909,10 @@ def prepare_spmm_csr_opt_alg2_preprocess(data, indices, indptr, shape): def _validate_spmm_opt_alg2_runtime_inputs(prepared, B, out): if not prepared.supports_opt: - raise TypeError("flagsparse_spmm_csr_opt_alg2 only supports float32 and float64") + raise TypeError( + "flagsparse_spmm_csr_opt_alg2 only supports float32, float64, " + "complex64, and complex128" + ) if B is None: raise ValueError("B is required") if B.ndim != 2: @@ -926,7 +935,7 @@ def _validate_spmm_opt_alg2_runtime_inputs(prepared, B, out): def _spmm_opt_alg2_acc_dtype(dtype): - return tl.float64 if dtype == torch.float64 else tl.float32 + return tl.float64 if dtype in (torch.float64, torch.complex128) else tl.float32 def _run_spmm_opt_alg2_bucket(prepared, bucket, B, C_out, device_props): @@ -1021,6 +1030,33 @@ def _run_spmm_opt_alg2_bucket(prepared, bucket, B, C_out, device_props): def _triton_spmm_csr_impl_opt_alg2_prepared(prepared, B, opt_buckets=None, return_meta=False): device_props = _normalize_spmm_opt_alg2_device_props(prepared.data.device) + if B.dtype in (torch.complex64, torch.complex128): + C_out = _triton_spmm_csr_complex_impl( + prepared.data, + prepared.kernel_indices, + prepared.kernel_indptr, + B, + prepared.n_rows, + int(B.shape[1]), + block_n=32, + block_nnz=256, + num_warps=4, + num_stages=3, + dense_layout="row", + ) + meta = None + if return_meta: + meta = { + "device_name": device_props["device_name"], + "sm_count": device_props["sm_count"], + "warp_size": device_props["warp_size"], + "bucket_hits": [], + "launch_configs": [], + "max_row_nnz": prepared.max_row_nnz, + "n_dense_cols": int(B.shape[1]), + "complex_kernel": "csr_base_complex", + } + return C_out, meta C_out = torch.zeros((prepared.n_rows, int(B.shape[1])), dtype=B.dtype, device=B.device) launch_configs = [] bucket_hits = [] diff --git a/tests/test_spmm_coo.py b/tests/test_spmm_coo.py index e498a0d..b951440 100644 --- a/tests/test_spmm_coo.py +++ b/tests/test_spmm_coo.py @@ -449,6 +449,189 @@ def _build_pytorch_reference(data, row, col, shape, B, prepared=None, op="non", +def _cuda_event_benchmark(op, warmup, iters): + out = None + for _ in range(max(0, int(warmup))): + out = op() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + count = max(1, int(iters)) + start.record() + for _ in range(count): + out = op() + end.record() + torch.cuda.synchronize() + return out, start.elapsed_time(end) / count + + +def _prepare_spmm_coo_timing_base(data, row, col, B, shape, op="non", layout="row"): + layout = _normalize_layout_name(layout) + op_code = ast_ops._normalize_spmm_coo_op(op) + data, row, col, shape = ast_ops._materialize_spmm_coo_op(data, row, col, shape, op_code) + native_data, native_row, native_col, native_B, n_rows, n_cols, n_dense_cols = ast_ops._prepare_spmm_coo_inputs( + data, + row, + col, + B, + shape, + dense_layout=layout, + ) + output_dtype = native_data.dtype + compute_dtype = ast_ops._spmm_coo_compute_dtype(output_dtype) + data_compute = native_data if compute_dtype == output_dtype else native_data.to(compute_dtype) + B_compute = native_B if compute_dtype == output_dtype else native_B.to(compute_dtype) + launch = ast_ops._resolve_spmm_coo_launch_config( + n_dense_cols, + int(native_data.numel()), + block_n=None, + block_nnz=DEFAULT_BLOCK_NNZ, + ) + return { + "data": data_compute, + "row": native_row, + "col": native_col, + "B": B_compute, + "n_rows": n_rows, + "n_cols": n_cols, + "n_dense_cols": n_dense_cols, + "output_dtype": output_dtype, + "layout": layout, + "launch": launch, + } + + +def _spmm_coo_rowrun_process(base): + canonical_data, canonical_row, canonical_col = ast_ops._coalesce_coo_entries( + base["data"], + base["row"], + base["col"], + (base["n_rows"], base["n_cols"]), + ) + canonical_data, canonical_row, canonical_col = ast_ops._sort_coo_lex_inplace( + canonical_data, + canonical_row, + canonical_col, + base["n_cols"], + ) + seg_starts = ast_ops._seg_starts_from_sorted_rows( + canonical_row, + int(canonical_data.numel()), + canonical_data.device, + ) + return { + "data": canonical_data, + "row": canonical_row, + "col": canonical_col, + "B": base["B"], + "seg_starts": seg_starts, + } + + +def _spmm_coo_rowrun_compute(base, plan): + launch = base["launch"] + return ast_ops._triton_spmm_coo_rowrun_impl( + plan["data"], + plan["row"], + plan["col"], + plan["B"], + base["n_rows"], + base["n_dense_cols"], + block_n=launch["block_n"], + block_nnz=launch["block_nnz"], + output_dtype=base["output_dtype"], + dense_layout=base["layout"], + seg_starts=plan["seg_starts"], + ) + + +def _spmm_coo_atomic_compute(base): + launch = base["launch"] + return ast_ops._triton_spmm_coo_atomic_impl( + base["data"], + base["row"], + base["col"], + base["B"], + base["n_rows"], + base["n_dense_cols"], + block_n=launch["block_n"], + block_nnz=launch["block_nnz"], + output_dtype=base["output_dtype"], + dense_layout=base["layout"], + ) + + +def _benchmark_spmm_coo_route_policy( + data, + row, + col, + B, + shape, + warmup, + iters, + route="rowrun", + block_n=None, + block_nnz=DEFAULT_BLOCK_NNZ, + op="non", + layout="row", + timing=False, +): + selected_route = _selected_route(route) + base = _prepare_spmm_coo_timing_base(data, row, col, B, shape, op=op, layout=layout) + if block_n is not None or block_nnz != DEFAULT_BLOCK_NNZ: + base["launch"] = ast_ops._resolve_spmm_coo_launch_config( + base["n_dense_cols"], + int(base["data"].numel()), + block_n=block_n, + block_nnz=block_nnz, + ) + + if selected_route == "rowrun": + def full_op(): + plan = _spmm_coo_rowrun_process(base) + return _spmm_coo_rowrun_compute(base, plan) + + torch.cuda.synchronize() + t0 = time.perf_counter() + _ = full_op() + torch.cuda.synchronize() + first_call_ms = (time.perf_counter() - t0) * 1000.0 + values, gpu_ms = _cuda_event_benchmark(full_op, warmup, iters) + process_gpu_ms = None + compute_ms = None + total_ms = gpu_ms + if timing: + plan, process_gpu_ms = _cuda_event_benchmark(lambda: _spmm_coo_rowrun_process(base), warmup, iters) + values, compute_ms = _cuda_event_benchmark(lambda: _spmm_coo_rowrun_compute(base, plan), warmup, iters) + total_ms = process_gpu_ms + compute_ms + else: + def full_op(): + return _spmm_coo_atomic_compute(base) + + torch.cuda.synchronize() + t0 = time.perf_counter() + _ = full_op() + torch.cuda.synchronize() + first_call_ms = (time.perf_counter() - t0) * 1000.0 + values, gpu_ms = _cuda_event_benchmark(full_op, warmup, iters) + process_gpu_ms = 0.0 if timing else None + compute_ms = None + total_ms = gpu_ms + if timing: + values, compute_ms = _cuda_event_benchmark(full_op, warmup, iters) + total_ms = compute_ms + + return { + "values": values, + "ms": total_ms, + "gpu_ms": gpu_ms, + "process_cpu_ms": 0.0, + "process_gpu_ms": process_gpu_ms, + "compute_ms": compute_ms, + "first_call_ms": first_call_ms, + } + + def _benchmark_spmm_coo_route( data, row, @@ -463,10 +646,11 @@ def _benchmark_spmm_coo_route( prepared=None, op="non", layout="row", + timing=False, ): selected_route = _selected_route(route) del prepared - return ast_ops._benchmark_spmm_coo_route( + result = _benchmark_spmm_coo_route_policy( data, row, col, @@ -474,12 +658,14 @@ def _benchmark_spmm_coo_route( shape, warmup, iters, - block_n, - block_nnz, - selected_route, + block_n=block_n, + block_nnz=block_nnz, + route=selected_route, op=op, - dense_layout=layout, + layout=layout, + timing=timing, ) + return result["values"], result["ms"], result["first_call_ms"], result def _summarize_route_output(values, reference, value_dtype, ms=None, first_call_ms=None, cusparse_values=None): metrics = ast_ops._spmm_validation_metrics(values, reference) @@ -640,6 +826,7 @@ def run_one_mtx( route="rowrun", op="non", layout="row", + timing=False, ): route = _normalize_route(route) selected_route = _selected_route(route) @@ -671,6 +858,10 @@ def run_one_mtx( "c_stride": "", "error": None, "triton_ms": None, + "triton_gpu_ms": None, + "process_cpu_ms": None, + "process_gpu_ms": None, + "compute_ms": None, "triton_first_call_ms": None, "cusparse_ms": None, "pytorch_ms": None, @@ -709,7 +900,7 @@ def run_one_mtx( triton_C = None try: - triton_C, triton_ms, triton_first_call_ms = _benchmark_spmm_coo_route( + triton_C, triton_ms, triton_first_call_ms, triton_timing = _benchmark_spmm_coo_route( data, row, col, @@ -723,8 +914,13 @@ def run_one_mtx( prepared=prepared, op=op_name, layout=layout, + timing=timing, ) result["triton_ms"] = triton_ms + result["triton_gpu_ms"] = triton_timing.get("gpu_ms") + result["process_cpu_ms"] = triton_timing.get("process_cpu_ms") + result["process_gpu_ms"] = triton_timing.get("process_gpu_ms") + result["compute_ms"] = triton_timing.get("compute_ms") result["triton_first_call_ms"] = triton_first_call_ms result["c_stride"] = _stride_string(triton_C) except Exception as exc: @@ -822,7 +1018,7 @@ def _run_cusparse_timing(rhs): route_summaries = {} if triton_C is not None: route_outputs[selected_route] = triton_C - route_summaries[selected_route] = _summarize_route_output( + route_summary = _summarize_route_output( triton_C, ref_C, value_dtype, @@ -830,11 +1026,18 @@ def _run_cusparse_timing(rhs): first_call_ms=triton_first_call_ms, cusparse_values=cs_C_t, ) + route_summary.update({ + "gpu_ms": result.get("triton_gpu_ms"), + "process_cpu_ms": result.get("process_cpu_ms"), + "process_gpu_ms": result.get("process_gpu_ms"), + "compute_ms": result.get("compute_ms"), + }) + route_summaries[selected_route] = route_summary for extra_route in ("rowrun", "atomic"): if extra_route in route_outputs: continue try: - extra_C, extra_ms, extra_first_call_ms = _benchmark_spmm_coo_route( + extra_C, extra_ms, extra_first_call_ms, extra_timing = _benchmark_spmm_coo_route( data, row, col, @@ -848,9 +1051,10 @@ def _run_cusparse_timing(rhs): prepared=prepared, op=op_name, layout=layout, + timing=timing, ) route_outputs[extra_route] = extra_C - route_summaries[extra_route] = _summarize_route_output( + route_summary = _summarize_route_output( extra_C, ref_C, value_dtype, @@ -858,6 +1062,13 @@ def _run_cusparse_timing(rhs): first_call_ms=extra_first_call_ms, cusparse_values=cs_C_t, ) + route_summary.update({ + "gpu_ms": extra_timing.get("gpu_ms"), + "process_cpu_ms": extra_timing.get("process_cpu_ms"), + "process_gpu_ms": extra_timing.get("process_gpu_ms"), + "compute_ms": extra_timing.get("compute_ms"), + }) + route_summaries[extra_route] = route_summary except Exception as exc: route_summaries[extra_route] = { "ms": None, @@ -912,6 +1123,7 @@ def run_mtx_batch( route="rowrun", op="non", layout="row", + timing=False, on_result=None, ): results = [] @@ -929,6 +1141,7 @@ def run_mtx_batch( route=route, op=op, layout=layout, + timing=timing, ) results.append(entry) if on_result is not None: @@ -936,35 +1149,299 @@ def run_mtx_batch( return results -def _print_spmm_coo_mtx_header(value_dtype, index_dtype, route, layout=None): +def _benchmark_spmm_coo_synthetic_policy( + n_rows=4096, + n_cols=4096, + nnz=65536, + n_dense_cols=32, + value_dtype=torch.float32, + index_dtype=torch.int32, + warmup=20, + iters=200, + block_n=None, + block_nnz=DEFAULT_BLOCK_NNZ, + run_cusparse=True, + route="rowrun", + compare_routes=False, + op="non", + dense_layout="row", + timing=False, +): + device = torch.device("cuda") + route = _selected_route(route) + layout = _normalize_layout_name(dense_layout) + op_name = ast_ops._spmm_coo_op_to_name(op) + data, row, col = ast_ops._build_random_coo( + n_rows, + n_cols, + int(nnz), + value_dtype, + index_dtype, + device, + ) + shape = (int(n_rows), int(n_cols)) + b_rows = n_rows if ast_ops._spmm_coo_op_transposes(op_name) else n_cols + B = _materialize_dense_layout_for_test( + _build_dense_matrix(b_rows, n_dense_cols, value_dtype, device), + layout, + ) + prepared = _prepare_canonical_case(data, row, col, shape, B, op=op_name, layout=layout) + expected, pytorch_op, pytorch_format, pytorch_reason = _build_pytorch_reference( + data, + row, + col, + shape, + B, + prepared=prepared, + op=op_name, + layout=layout, + ) + launch = ast_ops._resolve_spmm_coo_launch_config( + prepared["n_dense_cols"], + int(prepared["canonical_data"].numel()), + block_n=block_n, + block_nnz=block_nnz, + ) + seg_starts = ast_ops._seg_starts_from_sorted_rows( + prepared["canonical_row"], + int(prepared["canonical_data"].numel()), + device, + ) + n_row_runs = int(seg_starts.numel()) - 1 if seg_starts is not None else 0 + + triton_C, triton_ms, triton_first_call_ms, triton_timing = _benchmark_spmm_coo_route( + data, + row, + col, + B, + shape, + warmup, + iters, + route=route, + block_n=launch["block_n"], + block_nnz=launch["block_nnz"], + prepared=prepared, + op=op_name, + layout=layout, + timing=timing, + ) + triton_summary = _summarize_route_output(triton_C, expected, value_dtype) + + pytorch_values = expected + pytorch_ms = None + try: + pytorch_values, pytorch_ms = ast_ops._benchmark_cuda_op(pytorch_op, warmup=warmup, iters=iters) + except Exception as exc: + pytorch_reason = str(exc) if pytorch_reason is None else f"{pytorch_reason}; timing: {exc}" + + cusparse_ms = None + cusparse_match = None + cusparse_reason = None + cusparse_values = None + cusparse_summary = None + if run_cusparse: + if value_dtype not in (torch.float32, torch.float64, torch.complex64, torch.complex128): + cusparse_reason = "float16/bfloat16 not supported by CuPy sparse; skipped" + else: + try: + import cupy as cp + import cupyx.scipy.sparse as cpx + + data_cp = cp.from_dlpack(torch.utils.dlpack.to_dlpack(prepared["cusparse_data"])) + row_cp = cp.from_dlpack(torch.utils.dlpack.to_dlpack(prepared["cusparse_row"].to(torch.int64))) + col_cp = cp.from_dlpack(torch.utils.dlpack.to_dlpack(prepared["cusparse_col"].to(torch.int64))) + B_cp = cp.from_dlpack(torch.utils.dlpack.to_dlpack(prepared["native_B"])) + A_coo = cpx.coo_matrix((data_cp, (row_cp, col_cp)), shape=(prepared["n_rows"], prepared["n_cols"])) + cusparse_values_cp, cusparse_ms = ast_ops._benchmark_cuda_op( + lambda: A_coo @ B_cp, + warmup=warmup, + iters=iters, + ) + cusparse_values = torch.utils.dlpack.from_dlpack(cusparse_values_cp.toDlpack()) + cusparse_summary = ast_ops._spmm_validation_metrics(cusparse_values, expected) + atol, rtol = _tolerance_for_dtype(value_dtype) + cusparse_match = torch.allclose(cusparse_values, expected, atol=atol, rtol=rtol) + except Exception as exc: + cusparse_reason = str(exc) + + route_results = None + route_samples = None + parity = None + if compare_routes: + route_outputs = {route: triton_C} + route_results = { + route: { + "route": route, + "ms": triton_ms, + "gpu_ms": triton_timing.get("gpu_ms"), + "process_cpu_ms": triton_timing.get("process_cpu_ms"), + "process_gpu_ms": triton_timing.get("process_gpu_ms"), + "compute_ms": triton_timing.get("compute_ms"), + "first_call_ms": triton_first_call_ms, + "match_reference": triton_summary["ok_pt"], + "error_ratio": triton_summary["err_pt"], + "max_abs_error": triton_summary["max_abs_error"], + "max_relative_error": triton_summary["max_relative_error"], + "error": None, + } + } + for extra_route in ("rowrun", "atomic"): + if extra_route in route_outputs: + continue + try: + extra_C, extra_ms, extra_first_call_ms, extra_timing = _benchmark_spmm_coo_route( + data, + row, + col, + B, + shape, + warmup, + iters, + route=extra_route, + block_n=launch["block_n"], + block_nnz=launch["block_nnz"], + prepared=prepared, + op=op_name, + layout=layout, + timing=timing, + ) + extra_summary = _summarize_route_output(extra_C, expected, value_dtype) + route_outputs[extra_route] = extra_C + route_results[extra_route] = { + "route": extra_route, + "ms": extra_ms, + "gpu_ms": extra_timing.get("gpu_ms"), + "process_cpu_ms": extra_timing.get("process_cpu_ms"), + "process_gpu_ms": extra_timing.get("process_gpu_ms"), + "compute_ms": extra_timing.get("compute_ms"), + "first_call_ms": extra_first_call_ms, + "match_reference": extra_summary["ok_pt"], + "error_ratio": extra_summary["err_pt"], + "max_abs_error": extra_summary["max_abs_error"], + "max_relative_error": extra_summary["max_relative_error"], + "error": None, + } + except Exception as exc: + route_results[extra_route] = {"route": extra_route, "ms": None, "first_call_ms": None, "match_reference": False, "error": str(exc)} + parity = {"rowrun_vs_atomic": _empty_pairwise_summary()} + if "rowrun" in route_outputs and "atomic" in route_outputs: + parity["rowrun_vs_atomic"] = _pairwise_route_summary(route_outputs["rowrun"], route_outputs["atomic"], value_dtype) + route_samples = route_outputs + + threshold = ast_ops._spmm_relative_threshold(value_dtype) + return { + "parameters": { + "format": "coo", + "internal_format": f"native-{route}", + "route": route, + "op": op_name, + "dense_layout": layout, + "b_stride": tuple(int(v) for v in prepared["native_B"].stride()), + "c_stride": tuple(int(v) for v in triton_C.stride()), + "n_rows": prepared["n_rows"], + "n_cols": prepared["n_cols"], + "nnz": int(nnz), + "n_dense_cols": n_dense_cols, + "value_dtype": str(value_dtype), + "index_dtype": str(index_dtype), + "warmup": warmup, + "iters": iters, + "block_n": launch["block_n"], + "block_nnz": launch["block_nnz"], + "required_nnz_tiles": launch["required_nnz_tiles"], + "heuristic_warp_size": launch["heuristic_warp_size"], + "heuristic_factor": launch["heuristic_factor"], + "n_row_runs": n_row_runs, + "run_cusparse": run_cusparse, + }, + "performance": { + "pytorch_ms": pytorch_ms, + "triton_ms": triton_ms, + "triton_gpu_ms": triton_timing.get("gpu_ms"), + "process_cpu_ms": triton_timing.get("process_cpu_ms"), + "process_gpu_ms": triton_timing.get("process_gpu_ms"), + "compute_ms": triton_timing.get("compute_ms"), + "triton_first_call_ms": triton_first_call_ms, + "cusparse_ms": cusparse_ms, + "triton_speedup_vs_pytorch": _speedup_ratio(pytorch_ms, triton_ms), + "triton_speedup_vs_cusparse": _speedup_ratio(cusparse_ms, triton_ms), + }, + "verification": { + "triton_match_reference": triton_summary["ok_pt"], + "triton_match_pytorch": triton_summary["ok_pt"], + "triton_max_error": triton_summary["max_abs_error"], + "triton_max_abs_error": triton_summary["max_abs_error"], + "triton_max_relative_error": triton_summary["max_relative_error"], + "triton_sum_relative_error": None, + "triton_relative_threshold": threshold, + "triton_strict_allclose_match": triton_summary["ok_pt"], + "pytorch_match_reference": True, + "cusparse_match_reference": cusparse_match, + "cusparse_match_pytorch": cusparse_match, + "cusparse_max_error": (cusparse_summary["max_abs_error"] if cusparse_summary is not None else None), + "cusparse_max_abs_error": (cusparse_summary["max_abs_error"] if cusparse_summary is not None else None), + "cusparse_max_relative_error": (cusparse_summary["max_relative_error"] if cusparse_summary is not None else None), + "cusparse_sum_relative_error": None, + "cusparse_relative_threshold": threshold, + "cusparse_strict_allclose_match": cusparse_match, + }, + "backend_status": { + "pytorch_unavailable_reason": pytorch_reason, + "pytorch_sparse_format": pytorch_format, + "cusparse_unavailable_reason": cusparse_reason, + "flagsparse_internal_route": f"coo-native-{route}", + }, + "samples": { + "pytorch": pytorch_values, + "triton": triton_C, + "reference": expected, + "cusparse": cusparse_values, + }, + "route_results": route_results, + "parity": parity, + "route_samples": route_samples, + } + + +def _print_spmm_coo_mtx_header(value_dtype, index_dtype, route, layout=None, timing=False): route = _normalize_route(route) print(f"Value dtype: {_dtype_name(value_dtype)} | Index dtype: {_dtype_name(index_dtype)}") if layout is not None: print(f"Dense layout: {layout}") print(f"Formats: FlagSparse={_route_label(route)}, cuSPARSE=COO dense-mm, PyTorch=COO.") + print("Timing: FS(ms)=process_cpu_ms+FS_GPU(ms); --timing adds process_gpu_ms/compute_ms split.") + print("Rowrun process includes COO execution-plan rebuild: coalesce/sort + seg_starts.") + print("Atomic has no current execution-plan preprocessing; host launch config and input normalization are excluded.") print("Timing stays in native dtype. For float32, correctness references use float64 compute then cast.") print("PT/CU show per-reference correctness. Err(PT)/Err(CU)=max(|diff| / (atol + rtol*|ref|)).") print("PyTorch uses COO sparse.mm as the only correctness reference path.") if route == "compare": print("Compare mode also benchmarks native atomic (debug-only) after the main table.") - print("-" * 186) + width = 226 if timing else 202 + print("-" * width) + split = f"{'ProcGPU':>9} {'Compute':>9} " if timing else "" print( f"{'Matrix':<28} {'Op':>5} {'Lay':>4} {'N_rows':>7} {'N_cols':>7} {'NNZ':>10} {'DenseN':>8} " - f"{'FlagSparse(ms)':>14} {'cuSPARSE(ms)':>13} {'PyTorch(ms)':>11} " - f"{'FS/CU':>7} {'FS/PT':>7} {'PT':>6} {'CU':>6} {'Err(PT)':>10} {'Err(CU)':>10}" + f"{'FS(ms)':>9} {'FS_GPU':>9} {'CPUProc':>9} {split}" + f"{'cuSPARSE':>9} {'PyTorch':>9} {'FS/CU':>7} {'FS/PT':>7} {'PT':>6} {'CU':>6} {'Err(PT)':>10} {'Err(CU)':>10}" ) - print("-" * 186) + print("-" * width) -def _print_spmm_coo_mtx_row(entry): +def _print_spmm_coo_mtx_row(entry, timing=False): name = os.path.basename(entry["path"])[:27] n_rows, n_cols = entry["shape"] triton_ms = entry.get("triton_ms") cu_ms = entry.get("cusparse_ms") pt_ms = entry.get("pytorch_ms") + split = ( + f"{_fmt_ms(entry.get('process_gpu_ms')):>9} {_fmt_ms(entry.get('compute_ms')):>9} " + if timing else "" + ) print( f"{name:<28} {entry.get('op', 'non'):>5} {entry.get('layout', 'row'):>4} {n_rows:>7} {n_cols:>7} {entry['nnz']:>10} {entry['dense_cols']:>8} " - f"{_fmt_ms(triton_ms):>14} {_fmt_ms(cu_ms):>13} {_fmt_ms(pt_ms):>11} " + f"{_fmt_ms(triton_ms):>9} {_fmt_ms(entry.get('triton_gpu_ms')):>9} {_fmt_ms(entry.get('process_cpu_ms')):>9} {split}" + f"{_fmt_ms(cu_ms):>9} {_fmt_ms(pt_ms):>9} " f"{_fmt_speedup(cu_ms, triton_ms):>7} {_fmt_speedup(pt_ms, triton_ms):>7} " f"{_fmt_check(entry.get('triton_ok_pt')):>6} {_fmt_check(entry.get('triton_ok_cu')):>6} " f"{_fmt_err(entry.get('err_pt')):>10} {_fmt_err(entry.get('err_cu')):>10}" @@ -977,12 +1454,12 @@ def _print_spmm_coo_mtx_row(entry): print(f" NOTE: {msg}") -def print_mtx_results(results, value_dtype, index_dtype, route="rowrun", layout=None): +def print_mtx_results(results, value_dtype, index_dtype, route="rowrun", layout=None, timing=False): route = _normalize_route(route) - _print_spmm_coo_mtx_header(value_dtype, index_dtype, route, layout=layout) + _print_spmm_coo_mtx_header(value_dtype, index_dtype, route, layout=layout, timing=timing) for entry in results: - _print_spmm_coo_mtx_row(entry) - print("-" * 186) + _print_spmm_coo_mtx_row(entry, timing=timing) + print("-" * (226 if timing else 202)) @@ -1051,6 +1528,7 @@ def run_all_dtypes_export_csv( index_dtypes=None, op_names=None, layout_names=None, + timing=False, ): route = _normalize_route(route) if route == "compare": @@ -1067,7 +1545,7 @@ def run_all_dtypes_export_csv( for op_name in op_names: for layout_name in layout_names: print("=" * 150) - _print_spmm_coo_mtx_header(value_dtype, index_dtype, route, layout=layout_name) + _print_spmm_coo_mtx_header(value_dtype, index_dtype, route, layout=layout_name, timing=timing) results = run_mtx_batch( paths, value_dtype=value_dtype, @@ -1081,9 +1559,10 @@ def run_all_dtypes_export_csv( route=selected_route, op=op_name, layout=layout_name, - on_result=_print_spmm_coo_mtx_row, + timing=timing, + on_result=lambda entry: _print_spmm_coo_mtx_row(entry, timing=timing), ) - print("-" * 186) + print("-" * (226 if timing else 202)) for entry in results: n_rows, n_cols = entry["shape"] status = entry.get("status") @@ -1105,6 +1584,10 @@ def run_all_dtypes_export_csv( "b_stride": entry.get("b_stride"), "c_stride": entry.get("c_stride"), "triton_ms": entry.get("triton_ms"), + "triton_gpu_ms": entry.get("triton_gpu_ms"), + "process_cpu_ms": entry.get("process_cpu_ms"), + "process_gpu_ms": entry.get("process_gpu_ms"), + "compute_ms": entry.get("compute_ms"), "cusparse_ms": entry.get("cusparse_ms"), "pytorch_ms": entry.get("pytorch_ms"), "triton_speedup_vs_cusparse": _speedup_ratio( @@ -1124,10 +1607,14 @@ def run_all_dtypes_export_csv( fieldnames = [ "matrix", "op", "layout", "value_dtype", "index_dtype", "n_rows", "n_cols", "nnz", "b_stride", "c_stride", - "triton_ms", "cusparse_ms", "pytorch_ms", + "triton_ms", "triton_gpu_ms", "process_cpu_ms", + "cusparse_ms", "pytorch_ms", "triton_speedup_vs_cusparse", "triton_speedup_vs_pytorch", "pt_status", "cu_status", "status", "err_pt", "err_cu", "error", "cusparse_reason", ] + if timing: + insert_at = fieldnames.index("cusparse_ms") + fieldnames[insert_at:insert_at] = ["process_gpu_ms", "compute_ms"] with open(csv_path, "w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore") writer.writeheader() @@ -1408,6 +1895,7 @@ def run_comprehensive_synthetic( route="rowrun", op_names=None, layout_names=None, + timing=False, ): if not torch.cuda.is_available(): print("CUDA is not available.") @@ -1426,6 +1914,9 @@ def run_comprehensive_synthetic( f"BLOCK_N: {_fmt_launch_value(block_n)} BLOCK_NNZ: {_fmt_launch_value(block_nnz)} Route: {route} Ops: {','.join(op_names)} Layouts: {','.join(layout_names)}" ) print(f"Formats: FlagSparse={_route_label(route)}, cuSPARSE=COO dense-mm (when supported), PyTorch=COO.") + print("Timing: FS(ms)=process_cpu_ms+FS_GPU(ms); --timing adds process_gpu_ms/compute_ms split.") + print("Rowrun process includes COO execution-plan rebuild: coalesce/sort + seg_starts.") + print("Atomic has no current execution-plan preprocessing; host launch config and input normalization are excluded.") print("For float32, PT checks the float64-based correctness reference while CU reflects native cuSPARSE float32 consistency.") if route == "compare": print("Compare mode also benchmarks native atomic (debug-only) for each synthetic case.") @@ -1436,19 +1927,22 @@ def run_comprehensive_synthetic( for value_dtype in VALUE_DTYPES: for index_dtype in INDEX_DTYPES: compare_rows = [] - print("-" * 150) + width = 194 if timing else 174 + print("-" * width) print(f"Value dtype: {_dtype_name(value_dtype):<12} | Index dtype: {_dtype_name(index_dtype):<6}") - print("-" * 150) + print("-" * width) + split_header = f"{'ProcGPU':>9} {'Compute':>9} " if timing else "" print( f"{'Op':>5} {'Lay':>4} {'N_rows':>7} {'N_cols':>7} {'NNZ':>10} {'DenseN':>8} {'BN':>4} {'BNNZ':>6} {'Runs':>5} {'Tiles':>5} " - f"{'PyTorch(ms)':>12} {'FlagSparse(ms)':>14} {'cuSPARSE(ms)':>12} {'FS/PT':>8} {'FS/CU':>8} {'PT':>6} {'CU':>6} {'Err(FS)':>11} {'Err(CU)':>12}" + f"{'FS(ms)':>9} {'FS_GPU':>9} {'CPUProc':>9} {split_header}" + f"{'PyTorch':>9} {'cuSPARSE':>9} {'FS/PT':>8} {'FS/CU':>8} {'PT':>6} {'CU':>6} {'Err(FS)':>11} {'Err(CU)':>12}" ) - print("-" * 150) + print("-" * width) combo_reason = None for op_name in op_names: for layout_name in layout_names: for n_rows, n_cols, nnz, n_dense_cols in TEST_CASES: - result = ast_ops.benchmark_spmm_coo_case( + result = _benchmark_spmm_coo_synthetic_policy( n_rows=n_rows, n_cols=n_cols, nnz=nnz, @@ -1464,6 +1958,7 @@ def run_comprehensive_synthetic( compare_routes=(route == "compare"), op=op_name, dense_layout=layout_name, + timing=timing, ) total += 1 params = result["parameters"] @@ -1482,9 +1977,14 @@ def run_comprehensive_synthetic( cusparse_err = None if samples.get("cusparse") is not None: cusparse_err = _scaled_allclose_error(samples["triton"], samples["cusparse"], value_dtype) + split_values = ( + f"{_fmt_ms(perf.get('process_gpu_ms')):>9} {_fmt_ms(perf.get('compute_ms')):>9} " + if timing else "" + ) print( f"{op_name:>5} {layout_name:>4} {n_rows:>7} {n_cols:>7} {nnz:>10} {n_dense_cols:>8} {params['block_n']:>4} {params['block_nnz']:>6} {params['n_row_runs']:>5} {params['required_nnz_tiles']:>5} " - f"{_fmt_ms(perf.get('pytorch_ms')):>12} {_fmt_ms(perf.get('triton_ms')):>14} {_fmt_ms(perf.get('cusparse_ms')):>12} " + f"{_fmt_ms(perf.get('triton_ms')):>9} {_fmt_ms(perf.get('triton_gpu_ms')):>9} {_fmt_ms(perf.get('process_cpu_ms')):>9} {split_values}" + f"{_fmt_ms(perf.get('pytorch_ms')):>9} {_fmt_ms(perf.get('cusparse_ms')):>9} " f"{_fmt_speedup(perf.get('pytorch_ms'), perf.get('triton_ms')):>8} {_fmt_speedup(perf.get('cusparse_ms'), perf.get('triton_ms')):>8} " f"{_fmt_check(triton_ok):>6} {_fmt_check(cusparse_ok):>6} {_fmt_err(triton_err):>11} {_fmt_err(cusparse_err):>12}" ) @@ -1506,7 +2006,7 @@ def run_comprehensive_synthetic( "err_cu_pt": (verify.get("cusparse_max_relative_error") if verify.get("cusparse_match_reference") is not None else None), "err_row_atomic": (parity.get("rowrun_vs_atomic") or {}).get("error_ratio"), }) - print("-" * 150) + print("-" * width) if combo_reason: print(f" cuSPARSE: {combo_reason}") print() @@ -1543,6 +2043,7 @@ def main(): parser.add_argument("--block-n", type=int, default=DEFAULT_BLOCK_N, help="Output column tile override (default: auto from dense-column heuristic)") parser.add_argument("--block-nnz", type=int, default=DEFAULT_BLOCK_NNZ, help="COO nnz tile width override (default: 256)") parser.add_argument("--route", default="rowrun", choices=["rowrun", "atomic", "compare"], help="Native COO route to benchmark/test (default: rowrun)") + parser.add_argument("--timing", action="store_true", help="Add process_gpu_ms/compute_ms split timing columns") parser.add_argument("--warmup", type=int, default=10, help="Warmup runs") parser.add_argument("--iters", type=int, default=50, help="Timing iterations") parser.add_argument("--no-cusparse", action="store_true", help="Skip cuSPARSE baseline") @@ -1584,6 +2085,7 @@ def main(): route=args.route, op_names=op_names, layout_names=layout_names, + timing=args.timing, ) return @@ -1636,6 +2138,7 @@ def main(): index_dtypes=csv_index_dtypes, op_names=csv_op_names, layout_names=csv_layout_names, + timing=args.timing, ) return @@ -1664,8 +2167,9 @@ def main(): route=args.route, op=op_name, layout=layout_name, + timing=args.timing, ) - print_mtx_results(results, value_dtype, index_dtype, route=args.route, layout=layout_name) + print_mtx_results(results, value_dtype, index_dtype, route=args.route, layout=layout_name, timing=args.timing) if args.route == "compare": print_compare_results(results, value_dtype, index_dtype) diff --git a/tests/test_spmm_csr.py b/tests/test_spmm_csr.py index e553e53..2d0c489 100644 --- a/tests/test_spmm_csr.py +++ b/tests/test_spmm_csr.py @@ -42,6 +42,16 @@ DEFAULT_RUN_DTYPE_NAMES = ("float32", "float64") DEFAULT_OP_NAMES = tuple(spmm_ops.SPMM_OP_NAMES.values()) CUSPARSE_DTYPES = (torch.float32, torch.float64, torch.complex64, torch.complex128) +MAIN_CSR_SPMM_ALGORITHMS = { + "csr_base", + "csr_base_accuracy", + "alpha_alg1_tle_opt", + "alpha_alg1_tle_opt2", + "spmm_csr_alg1", + "spmm_csr_alg2", + "spmm_csr_alg2_accuracy", + "spmm_csr_alg2_accuracy_hp", +} PERF_FIELDS = [ "matrix", @@ -429,7 +439,7 @@ def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, for alg in _expand_algs(alg_names, op, dtype): try: resolved = fs.resolve_spmm_csr_algorithm(alg, op, dtype) - if layout == "col" and resolved.name != "csr_base": + if layout == "col" and resolved.name not in MAIN_CSR_SPMM_ALGORITHMS: rows.append( _skip_row( path, @@ -443,7 +453,7 @@ def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, b_stride, torch_ms, cusparse_ms, - "col-major layout is currently supported only by csr_base", + "col-major layout is currently supported only by CSR SpMM main algorithms", timing, cusparse_reason=cusparse_reason, ) diff --git a/tests/test_spmv_coo.py b/tests/test_spmv_coo.py index ceeccb4..5dc6b56 100644 --- a/tests/test_spmv_coo.py +++ b/tests/test_spmv_coo.py @@ -4,6 +4,7 @@ import csv import math import os +import time import torch import flagsparse as fs @@ -230,12 +231,25 @@ def _dense_to_coo(A): COO_SEP = "-" * 200 -COO_HEADER = ( - f"{'Matrix':<28} {'Op':>5} {'Out':>7} {'N_rows':>7} {'N_cols':>7} {'NNZ':>10} " - f"{'Base(ms)':>9} {'Opt(ms)':>9} {'PT(ms)':>9} {'CU(ms)':>9} " - f"{'Opt/Base':>8} {'Opt/PT':>8} {'Opt/CU':>8} " - f"{'Err(Base)':>10} {'Err(Opt)':>10} {'Status':>6}" -) + + +def _coo_header(timing=False): + split = ( + f" {'BasePGPU':>9} {'BaseComp':>9} {'OptPGPU':>9} {'OptComp':>9}" + if timing + else "" + ) + base = ( + f"{'Matrix':<28} {'Op':>5} {'Out':>7} {'N_rows':>7} {'N_cols':>7} {'NNZ':>10} " + f"{'Base(ms)':>9} {'BaseGPU':>9} {'BaseCPU':>9} " + f"{'Opt(ms)':>9} {'OptGPU':>9} {'OptCPU':>9}{split}" + ) + return ( + base + + f" {'PT(ms)':>9} {'CU(ms)':>9} " + f"{'Opt/Base':>8} {'Opt/PT':>8} {'Opt/CU':>8} " + f"{'Err(Base)':>10} {'Err(Opt)':>10} {'Status':>6}" + ) def _spd(num, den): @@ -249,44 +263,136 @@ def _spd(num, den): COO_ATOMIC_WARPS = 4 COO_SEG_BLOCK_INNER = 128 -def _timed_flagsparse_coo( - prepared, + +def _cuda_event_benchmark(op, warmup, iters): + out = None + count = max(1, int(iters)) + for _ in range(max(0, int(warmup))): + out = op() + torch.cuda.synchronize() + e0 = torch.cuda.Event(enable_timing=True) + e1 = torch.cuda.Event(enable_timing=True) + e0.record() + for _ in range(count): + out = op() + e1.record() + torch.cuda.synchronize() + return out, e0.elapsed_time(e1) / count + + +def _run_flagsparse_coo_launch( + launch, + x, +): + x = spmv_coo_mod._validate_x_coo(x, launch) + return spmv_coo_mod._run_spmv_coo_prepared_with_fallback( + launch, + x, + block_size=COO_ATOMIC_BLOCK, + num_warps=COO_ATOMIC_WARPS, + block_inner=COO_SEG_BLOCK_INNER, + ) + + +def _build_flagsparse_coo_launch( data, row, col, - x, shape, sort_by_row, op, +): + return spmv_coo_mod._prepare_spmv_coo_launch_from_raw( + data=data, + row=row, + col=col, + shape=shape, + sort_by_row=sort_by_row, + op=op, + ) + + +def _time_flagsparse_coo_row_run( + data, + row, + col, + x, + shape, + op, warmup, iters, + timing=False, ): - op = op.lower() - if op == "non": - spmv_op = lambda: _run_flagsparse_coo_prepared_non(prepared, x) - else: - spmv_op = lambda: _run_flagsparse_coo_runtime_op( + op = str(op).lower() + + def prepare(): + return _build_flagsparse_coo_launch( data, row, col, - x, shape, - sort_by_row, + True, op, ) - y = spmv_op() - torch.cuda.synchronize() - for _ in range(warmup): - spmv_op() - torch.cuda.synchronize() - e0 = torch.cuda.Event(True) - e1 = torch.cuda.Event(True) - e0.record() - for _ in range(iters): - y = spmv_op() - e1.record() - torch.cuda.synchronize() - return y, e0.elapsed_time(e1) / iters + + def full_op(): + launch = prepare() + return _run_flagsparse_coo_launch(launch, x) + + y, gpu_ms = _cuda_event_benchmark(full_op, warmup, iters) + process_gpu_ms = None + compute_ms = None + total_ms = gpu_ms + if timing: + launch, process_gpu_ms = _cuda_event_benchmark(prepare, warmup, iters) + y, compute_ms = _cuda_event_benchmark( + lambda: _run_flagsparse_coo_launch(launch, x), + warmup, + iters, + ) + total_ms = process_gpu_ms + compute_ms + return { + "out": y, + "ms": total_ms, + "gpu_ms": gpu_ms, + "process_cpu_ms": 0.0, + "process_gpu_ms": process_gpu_ms, + "compute_ms": compute_ms, + } + + +def _time_flagsparse_coo_atomic( + data, + row, + col, + x, + shape, + op, + warmup, + iters, + timing=False, +): + launch = _build_flagsparse_coo_launch( + data, + row, + col, + shape, + False, + str(op).lower(), + ) + y, gpu_ms = _cuda_event_benchmark( + lambda: _run_flagsparse_coo_launch(launch, x), + warmup, + iters, + ) + return { + "out": y, + "ms": gpu_ms, + "gpu_ms": gpu_ms, + "process_cpu_ms": 0.0, + "process_gpu_ms": 0.0 if timing else None, + "compute_ms": gpu_ms if timing else None, + } def _timed_flagsparse_coo_tocsr_runtime( @@ -343,54 +449,13 @@ def _timed_flagsparse_coo_tocsr_prepared( return y, e0.elapsed_time(e1) / iters -def _run_flagsparse_coo_prepared_non( - prepared, - x, -): - return fs.flagsparse_spmv_coo( - x=x, - prepared=prepared, - op="non", - return_time=False, - block_inner=COO_SEG_BLOCK_INNER, - block_size=COO_ATOMIC_BLOCK, - num_warps=COO_ATOMIC_WARPS, - ) - - -def _run_flagsparse_coo_runtime_op( - data, - row, - col, - x, - shape, - sort_by_row, - op, -): - launch = spmv_coo_mod._prepare_spmv_coo_launch_from_raw( - data=data, - row=row, - col=col, - shape=shape, - sort_by_row=sort_by_row, - op=op, - ) - x = spmv_coo_mod._validate_x_coo(x, launch) - return spmv_coo_mod._run_spmv_coo_prepared_with_fallback( - launch, - x, - block_size=COO_ATOMIC_BLOCK, - num_warps=COO_ATOMIC_WARPS, - block_inner=COO_SEG_BLOCK_INNER, - ) - - def run_synthetic( value_dtypes=None, index_dtypes=None, ops=None, warmup=WARMUP, iters=ITERS, + timing=False, ): if not torch.cuda.is_available(): print("CUDA is not available. Please run on a GPU-enabled system.") @@ -416,14 +481,13 @@ def run_synthetic( f"dtype: {_dtype_name(dtype)} index_dtype: {_dtype_name(index_dtype)} op: {op}" ) print(COO_SEP) - print( - "FlagSparse: prepare_spmv_coo + Triton COO SpMV (no CSR). " - "non = compute only; trans/conj = op processing + compute." - ) - print("Base(ms) = row-run (seg) kernel; Opt(ms) = NNZ atomic kernel.") + print("FlagSparse: native COO Triton only (no CSR).") + print("Base(ms)=BaseCPU+BaseGPU; BaseGPU wraps row-run sort + seg_starts + segmented kernel.") + print("Opt(ms)=OptCPU+OptGPU; atomic has no segment/bucket preprocessing.") + print("--timing splits row-run GPU work into BasePGPU + BaseComp; atomic OptPGPU is zero.") print("Speedups use Opt(ms) as the Triton comparison path; Base(ms) is reported separately.") print(COO_SEP) - print(COO_HEADER) + print(_coo_header(timing=timing)) print(COO_SEP) for m, n in TEST_SIZES: A = _random_values((m, n), dtype, device) @@ -442,8 +506,9 @@ def run_synthetic( matrix_name=f"{m}x{n}", warmup=warmup, iters=iters, + timing=timing, ) - _print_coo_result(result) + _print_coo_result(result, timing=timing) print(COO_SEP) print() @@ -459,41 +524,36 @@ def _run_one_coo_case( matrix_name, warmup, iters, + timing=False, ): row = row.to(index_dtype).contiguous() col = col.to(index_dtype).contiguous() x = _random_values((_x_size_for_op(shape, op),), dtype, data.device) atol, rtol = _tol_for_dtype(dtype) - prepared_seg_non = fs.prepare_spmv_coo( - data, row, col, shape, sort_by_row=True, op="non" - ) - prepared_at_non = fs.prepare_spmv_coo( - data, row, col, shape, sort_by_row=False, op="non" - ) - y_base, base_ms = _timed_flagsparse_coo( - prepared_seg_non, + base = _time_flagsparse_coo_row_run( data, row, col, x, shape, - True, op, warmup, iters, + timing=timing, ) - y_opt, opt_ms = _timed_flagsparse_coo( - prepared_at_non, + opt = _time_flagsparse_coo_atomic( data, row, col, x, shape, - False, op, warmup, iters, + timing=timing, ) + y_base = base["out"] + y_opt = opt["out"] y_ref = _pytorch_coo_reference(data, row, col, x, shape, dtype, op=op) err_base = _allclose_error_ratio(y_base, y_ref, atol, rtol) err_opt = _allclose_error_ratio(y_opt, y_ref, atol, rtol) @@ -534,12 +594,20 @@ def _run_one_coo_case( "n_rows": n_rows, "n_cols": n_cols, "nnz": int(data.numel()), - "base_ms": base_ms, - "opt_ms": opt_ms, + "base_ms": base["ms"], + "opt_ms": opt["ms"], + "base_gpu_ms": base["gpu_ms"], + "base_process_cpu_ms": base["process_cpu_ms"], + "base_process_gpu_ms": base["process_gpu_ms"], + "base_compute_ms": base["compute_ms"], + "opt_gpu_ms": opt["gpu_ms"], + "opt_process_cpu_ms": opt["process_cpu_ms"], + "opt_process_gpu_ms": opt["process_gpu_ms"], + "opt_compute_ms": opt["compute_ms"], "cusparse_ms": cu_ms, "pytorch_ms": pt_ms, - "opt_speedup_vs_cusparse": _speedup_ratio(cu_ms, opt_ms), - "opt_speedup_vs_pytorch": _speedup_ratio(pt_ms, opt_ms), + "opt_speedup_vs_cusparse": _speedup_ratio(cu_ms, opt["ms"]), + "opt_speedup_vs_pytorch": _speedup_ratio(pt_ms, opt["ms"]), "pt_status": _status_str(triton_ok_pt, err_pt is not None), "cu_status": _status_str(triton_ok_cu, err_cu is not None), "status": status, @@ -718,14 +786,24 @@ def _run_cupy_runtime_op(data, row, col, x, shape, op): return A_cp @ x -def _print_coo_result(row): +def _print_coo_result(row, timing=False): name = str(row["matrix"])[:27] if len(str(row["matrix"])) > 27: name += "…" + split_text = "" + if timing: + split_text = ( + f" {_fmt_ms(row.get('base_process_gpu_ms')):>9}" + f" {_fmt_ms(row.get('base_compute_ms')):>9}" + f" {_fmt_ms(row.get('opt_process_gpu_ms')):>9}" + f" {_fmt_ms(row.get('opt_compute_ms')):>9}" + ) print( f"{name:<28} {row['op']:>5} {row['out_size']:>7} " f"{row['n_rows']:>7} {row['n_cols']:>7} {row['nnz']:>10} " - f"{_fmt_ms(row['base_ms']):>9} {_fmt_ms(row['opt_ms']):>9} " + f"{_fmt_ms(row['base_ms']):>9} {_fmt_ms(row.get('base_gpu_ms')):>9} {_fmt_ms(row.get('base_process_cpu_ms')):>9} " + f"{_fmt_ms(row['opt_ms']):>9} {_fmt_ms(row.get('opt_gpu_ms')):>9} {_fmt_ms(row.get('opt_process_cpu_ms')):>9}" + f"{split_text} " f"{_fmt_ms(row['pytorch_ms']):>9} {_fmt_ms(row['cusparse_ms']):>9} " f"{_spd(row['base_ms'], row['opt_ms']):>8} " f"{_spd(row['pytorch_ms'], row['opt_ms']):>8} " @@ -874,6 +952,14 @@ def _error_row(path, dtype, index_dtype, op): "nnz": "ERR", "base_ms": None, "opt_ms": None, + "base_gpu_ms": None, + "base_process_cpu_ms": None, + "base_process_gpu_ms": None, + "base_compute_ms": None, + "opt_gpu_ms": None, + "opt_process_cpu_ms": None, + "opt_process_gpu_ms": None, + "opt_compute_ms": None, "cusparse_ms": None, "pytorch_ms": None, "opt_speedup_vs_cusparse": None, @@ -923,6 +1009,7 @@ def run_all_dtypes_coo_csv( ops=None, warmup=WARMUP, iters=ITERS, + timing=False, ): if not torch.cuda.is_available(): print("CUDA is not available.") @@ -936,12 +1023,13 @@ def run_all_dtypes_coo_csv( print("Input: MatrixMarket -> COO. FlagSparse: native COO Triton only (seg + atomic), no CSR.") print("PyTorch = COO sparse.mm; CuPy = COO matvec (coo_matrix @ x, no tocsr).") print( - "Timing policy: non = compute only; trans/conj = op processing + compute. " + "Timing policy: Base/Opt ms = process_cpu_ms + GPU event time. " + "Row-run sort + seg_starts are GPU process; atomic has no process. " "PyTorch/CuPy timings use original dtype." ) print( - f"prepare_spmv_coo once per variant + {warmup} warmup + " - f"{iters} CUDA-event-averaged SpMV per backend." + f"{warmup} warmup + {iters} averaged iterations. " + "--timing splits process_gpu_ms and compute_ms for native COO." ) print("=" * 200) for dtype in value_dtypes: @@ -952,7 +1040,7 @@ def run_all_dtypes_coo_csv( f"Value dtype: {_dtype_name(dtype)} | Index dtype: {_dtype_name(index_dtype)} | op: {op}" ) print(COO_SEP) - print(COO_HEADER) + print(_coo_header(timing=timing)) print(COO_SEP) for path in mtx_paths: try: @@ -970,13 +1058,14 @@ def run_all_dtypes_coo_csv( matrix_name=os.path.basename(path), warmup=warmup, iters=iters, + timing=timing, ) rows_out.append(result) - _print_coo_result(result) + _print_coo_result(result, timing=timing) except Exception as e: row_out = _error_row(path, dtype, index_dtype, op) rows_out.append(row_out) - _print_coo_result(row_out) + _print_coo_result(row_out, timing=timing) print(f" ERROR: {e}") print(COO_SEP) fieldnames = [ @@ -989,7 +1078,15 @@ def run_all_dtypes_coo_csv( "n_cols", "nnz", "base_ms", + "base_gpu_ms", + "base_process_cpu_ms", "opt_ms", + "opt_gpu_ms", + "opt_process_cpu_ms", + "base_process_gpu_ms", + "base_compute_ms", + "opt_process_gpu_ms", + "opt_compute_ms", "cusparse_ms", "pytorch_ms", "opt_speedup_vs_cusparse", @@ -1002,8 +1099,20 @@ def run_all_dtypes_coo_csv( "err_pt", "err_cu", ] + if not timing: + fieldnames = [ + field + for field in fieldnames + if field + not in ( + "base_process_gpu_ms", + "base_compute_ms", + "opt_process_gpu_ms", + "opt_compute_ms", + ) + ] with open(csv_path, "w", newline="", encoding="utf-8") as f: - w = csv.DictWriter(f, fieldnames=fieldnames) + w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") w.writeheader() for r in rows_out: w.writerow(r) @@ -1144,6 +1253,11 @@ def main(): ) parser.add_argument("--warmup", type=int, default=WARMUP) parser.add_argument("--iters", type=int, default=ITERS) + parser.add_argument( + "--timing", + action="store_true", + help="Show/export native COO timing breakdown columns", + ) args = parser.parse_args() value_dtypes = _parse_csv_tokens(args.dtypes, DTYPE_MAP, "--dtypes") index_dtypes = _parse_csv_tokens( @@ -1159,6 +1273,7 @@ def main(): ops=ops, warmup=args.warmup, iters=args.iters, + timing=args.timing, ) return @@ -1182,6 +1297,7 @@ def main(): ops=ops, warmup=args.warmup, iters=args.iters, + timing=args.timing, ) return diff --git a/tests/test_spmv_opt.py b/tests/test_spmv_opt.py index edf4d26..d56816f 100644 --- a/tests/test_spmv_opt.py +++ b/tests/test_spmv_opt.py @@ -91,38 +91,74 @@ def load_mtx_to_csr_torch(file_path, dtype=torch.float32, device=None): return data, indices, indptr, (n_rows, n_cols) -def _timed_spmv(prepared, x, warmup, iters, use_opt): - opt_buckets = None - symbolic_ms = 0.0 +def _cuda_event_benchmark(op, warmup, iters): + out = None count = max(1, int(iters)) - if use_opt: - opt_buckets = spmv_csr_mod._build_spmv_opt_runtime_buckets(prepared) - torch.cuda.synchronize() - t0 = time.perf_counter() - for _ in range(count): - opt_buckets = spmv_csr_mod._build_spmv_opt_runtime_buckets(prepared) - torch.cuda.synchronize() - symbolic_ms = (time.perf_counter() - t0) * 1000.0 / count - - def op(): - return spmv_csr_mod._run_spmv_prepared_with_fallback( - prepared, x, use_opt=use_opt, opt_buckets=opt_buckets - ) - - y = op() torch.cuda.synchronize() for _ in range(warmup): - y = op() + out = op() torch.cuda.synchronize() e0 = torch.cuda.Event(enable_timing=True) e1 = torch.cuda.Event(enable_timing=True) e0.record() for _ in range(count): - y = op() + out = op() e1.record() torch.cuda.synchronize() - compute_ms = e0.elapsed_time(e1) / count - return y, symbolic_ms + compute_ms, symbolic_ms, compute_ms + return out, e0.elapsed_time(e1) / count + + +def _timed_spmv(prepared, x, warmup, iters, use_opt, timing=False): + if not use_opt: + out, gpu_ms = _cuda_event_benchmark( + lambda: spmv_csr_mod._run_spmv_prepared_with_fallback(prepared, x, use_opt=False), + warmup, + iters, + ) + return { + "out": out, + "ms": gpu_ms, + "gpu_ms": gpu_ms, + "process_cpu_ms": 0.0, + "process_gpu_ms": 0.0 if timing else None, + "compute_ms": gpu_ms if timing else None, + "symbolic_ms": 0.0, + } + + def full_op(): + opt_buckets = spmv_csr_mod._build_spmv_opt_runtime_buckets(prepared) + return spmv_csr_mod._run_spmv_prepared_with_fallback( + prepared, x, use_opt=True, opt_buckets=opt_buckets + ) + + out, gpu_ms = _cuda_event_benchmark(full_op, warmup, iters) + process_gpu_ms = None + compute_ms = None + total_ms = gpu_ms + opt_buckets = None + if timing: + opt_buckets, process_gpu_ms = _cuda_event_benchmark( + lambda: spmv_csr_mod._build_spmv_opt_runtime_buckets(prepared), + warmup, + iters, + ) + out, compute_ms = _cuda_event_benchmark( + lambda: spmv_csr_mod._run_spmv_prepared_with_fallback( + prepared, x, use_opt=True, opt_buckets=opt_buckets + ), + warmup, + iters, + ) + total_ms = process_gpu_ms + compute_ms + return { + "out": out, + "ms": total_ms, + "gpu_ms": gpu_ms, + "process_cpu_ms": 0.0, + "process_gpu_ms": process_gpu_ms, + "compute_ms": compute_ms, + "symbolic_ms": (0.0 if process_gpu_ms is None else process_gpu_ms), + } def _timed_pytorch(data, indices, indptr, x, shape, warmup, iters): @@ -198,17 +234,35 @@ def _err(v): return "N/A" if v is None else f"{v:.2e}" -HEADER = ( - f"{'Matrix':<28} {'N_rows':>7} {'N_cols':>7} {'NNZ':>10} " - f"{'Base(ms)':>9} {'Opt(ms)':>9} {'Sym(ms)':>9} {'Comp(ms)':>9} " - f"{'PT(ms)':>9} {'CU(ms)':>9} " - f"{'Opt/Base':>8} {'Opt/PT':>8} {'Opt/CU':>8} " - f"{'Err(Base)':>10} {'Err(Opt)':>10} {'Status':>6}" -) -SEP = "-" * 170 +def _header(timing=False): + split = f"{'OptPGPU':>9} {'OptComp':>9} " if timing else "" + return ( + f"{'Matrix':<28} {'N_rows':>7} {'N_cols':>7} {'NNZ':>10} " + f"{'Base(ms)':>9} {'BaseGPU':>9} {'BaseCPU':>9} " + f"{'Opt(ms)':>9} {'OptGPU':>9} {'OptCPU':>9} {split}" + f"{'PT(ms)':>9} {'CU(ms)':>9} " + f"{'Opt/Base':>8} {'Opt/PT':>8} {'Opt/CU':>8} " + f"{'Err(Base)':>10} {'Err(Opt)':>10} {'Status':>6}" + ) + + +def _sep(timing=False): + return "-" * (210 if timing else 190) -def run_one_mtx(path, dtype, index_dtype, warmup, iters): +def _reference_tolerance(dtype): + if dtype in (torch.float32, torch.complex64): + return 1.3e-6, 1e-3 + if dtype in (torch.float64, torch.complex128): + return 1e-7, 1e-5 + if dtype == torch.float16: + return 1e-3, 2e-3 + if dtype == torch.bfloat16: + return 0.016, 1e-1 + return 1e-6, 1e-5 + + +def run_one_mtx(path, dtype, index_dtype, warmup, iters, timing=False): device = torch.device("cuda") data, indices, indptr, shape = load_mtx_to_csr_torch(path, dtype=dtype, device=device) indices = indices.to(index_dtype) @@ -217,10 +271,7 @@ def run_one_mtx(path, dtype, index_dtype, warmup, iters): x = torch.randn(n_cols, dtype=dtype, device=device) prepared = fs.prepare_spmv_csr(data, indices, indptr, shape) - if dtype == torch.float32: - atol, rtol = 1e-4, 1e-2 - else: - atol, rtol = 1e-12, 1e-10 + atol, rtol = _reference_tolerance(dtype) # ── Reference (float64 accumulation via PyTorch) ── try: @@ -247,14 +298,12 @@ def run_one_mtx(path, dtype, index_dtype, warmup, iters): y_ref = None # ── Baseline (use_opt=False) ── - y_base, base_ms, base_symbolic_ms, base_compute_ms = _timed_spmv( - prepared, x, warmup, iters, use_opt=False - ) + base = _timed_spmv(prepared, x, warmup, iters, use_opt=False, timing=timing) + y_base = base["out"] # ── Optimised (use_opt=True) ── - y_opt, opt_ms, symbolic_ms, compute_ms = _timed_spmv( - prepared, x, warmup, iters, use_opt=True - ) + opt = _timed_spmv(prepared, x, warmup, iters, use_opt=True, timing=timing) + y_opt = opt["out"] # ── PyTorch ── pt_ms = None @@ -286,12 +335,19 @@ def run_one_mtx(path, dtype, index_dtype, warmup, iters): return { "path": path, "shape": shape, "nnz": nnz, - "base_ms": base_ms, "opt_ms": opt_ms, - "base_symbolic_ms": base_symbolic_ms, - "base_compute_ms": base_compute_ms, - "symbolic_ms": symbolic_ms, - "compute_ms": compute_ms, - "op_total_ms": opt_ms, + "base_ms": base["ms"], "opt_ms": opt["ms"], + "base_gpu_ms": base["gpu_ms"], + "opt_gpu_ms": opt["gpu_ms"], + "base_process_cpu_ms": base["process_cpu_ms"], + "opt_process_cpu_ms": opt["process_cpu_ms"], + "base_process_gpu_ms": base["process_gpu_ms"], + "opt_process_gpu_ms": opt["process_gpu_ms"], + "base_compute_ms": base["compute_ms"], + "opt_compute_ms": opt["compute_ms"], + "base_symbolic_ms": base["symbolic_ms"], + "symbolic_ms": opt["symbolic_ms"], + "compute_ms": opt["compute_ms"], + "op_total_ms": opt["ms"], "pt_ms": pt_ms, "cu_ms": cu_ms, "err_base": err_base, "err_opt": err_opt, "base_ok": base_ok, "opt_ok": opt_ok, @@ -299,13 +355,17 @@ def run_one_mtx(path, dtype, index_dtype, warmup, iters): } -def print_row(r): +def print_row(r, timing=False): name = os.path.basename(r["path"])[:27] n_rows, n_cols = r["shape"] + split = ( + f"{_fmt(r.get('opt_process_gpu_ms')):>9} {_fmt(r.get('opt_compute_ms')):>9} " + if timing else "" + ) print( f"{name:<28} {n_rows:>7} {n_cols:>7} {r['nnz']:>10} " - f"{_fmt(r['base_ms']):>9} {_fmt(r['opt_ms']):>9} " - f"{_fmt(r['symbolic_ms']):>9} {_fmt(r['compute_ms']):>9} " + f"{_fmt(r['base_ms']):>9} {_fmt(r.get('base_gpu_ms')):>9} {_fmt(r.get('base_process_cpu_ms')):>9} " + f"{_fmt(r['opt_ms']):>9} {_fmt(r.get('opt_gpu_ms')):>9} {_fmt(r.get('opt_process_cpu_ms')):>9} {split}" f"{_fmt(r['pt_ms']):>9} {_fmt(r['cu_ms']):>9} " f"{_spd(r['base_ms'], r['op_total_ms']):>8} " f"{_spd(r['pt_ms'], r['op_total_ms']):>8} " @@ -314,39 +374,39 @@ def print_row(r): ) -def run_batch(paths, dtype, index_dtype, warmup, iters): +def run_batch(paths, dtype, index_dtype, warmup, iters, timing=False): results = [] for p in paths: try: - r = run_one_mtx(p, dtype, index_dtype, warmup, iters) + r = run_one_mtx(p, dtype, index_dtype, warmup, iters, timing=timing) except Exception as e: print(f" ERROR on {os.path.basename(p)}: {e}") continue results.append(r) - print_row(r) + print_row(r, timing=timing) return results -def run_all_csv(paths, csv_path, warmup, iters, dtype_filter=None): +def run_all_csv(paths, csv_path, warmup, iters, dtype_filter=None, timing=False): all_rows = [] dtypes = VALUE_DTYPES if dtype_filter is None else [dtype_filter] for dtype in dtypes: for idx_dtype in INDEX_DTYPES: dname = str(dtype).replace("torch.", "") iname = str(idx_dtype).replace("torch.", "") - print("=" * 170) + print("=" * (210 if timing else 190)) print(f"Value dtype: {dname} | Index dtype: {iname}") print( - "Base = FlagSparse baseline (fp64-accum for fp32). " - "Opt = FlagSparse CSR-Vector (fp32/fp64 native accum, wide tiles, few launches). " - "Opt(ms) = Sym(ms) CPU wall time + Comp(ms) CUDA event time. " + "Base = prepared baseline kernel. " + "Opt = CSR-Vector with bucket execution-plan data. " + "Base/Opt ms = process_cpu_ms + GPU event time; --timing adds process_gpu_ms/compute_ms. " "Speedup = Base/Opt or Ref/Opt." ) - print(SEP) - print(HEADER) - print(SEP) - results = run_batch(paths, dtype, idx_dtype, warmup, iters) - print(SEP) + print(_sep(timing)) + print(_header(timing)) + print(_sep(timing)) + results = run_batch(paths, dtype, idx_dtype, warmup, iters, timing=timing) + print(_sep(timing)) for r in results: n_rows, n_cols = r["shape"] all_rows.append({ @@ -355,6 +415,14 @@ def run_all_csv(paths, csv_path, warmup, iters, dtype_filter=None): "index_dtype": iname, "n_rows": n_rows, "n_cols": n_cols, "nnz": r["nnz"], "base_ms": r["base_ms"], "opt_ms": r["opt_ms"], + "base_gpu_ms": r["base_gpu_ms"], + "opt_gpu_ms": r["opt_gpu_ms"], + "base_process_cpu_ms": r["base_process_cpu_ms"], + "opt_process_cpu_ms": r["opt_process_cpu_ms"], + "base_process_gpu_ms": r["base_process_gpu_ms"], + "opt_process_gpu_ms": r["opt_process_gpu_ms"], + "base_compute_ms": r["base_compute_ms"], + "opt_compute_ms": r["opt_compute_ms"], "symbolic_ms": r["symbolic_ms"], "compute_ms": r["compute_ms"], "op_total_ms": r["op_total_ms"], @@ -370,11 +438,21 @@ def run_all_csv(paths, csv_path, warmup, iters, dtype_filter=None): fields = [ "matrix", "value_dtype", "index_dtype", "n_rows", "n_cols", "nnz", - "base_ms", "opt_ms", "symbolic_ms", "compute_ms", "op_total_ms", "pt_ms", "cu_ms", + "base_ms", "base_gpu_ms", "base_process_cpu_ms", + "opt_ms", "opt_gpu_ms", "opt_process_cpu_ms", + "symbolic_ms", "compute_ms", "op_total_ms", "pt_ms", "cu_ms", "opt_vs_base", "opt_vs_pt", "opt_vs_cu", "triton_speedup_vs_pytorch", "triton_speedup_vs_cusparse", "err_base", "err_opt", "status", ] + if timing: + insert_at = fields.index("symbolic_ms") + fields[insert_at:insert_at] = [ + "base_process_gpu_ms", + "base_compute_ms", + "opt_process_gpu_ms", + "opt_compute_ms", + ] with open(csv_path, "w", newline="", encoding="utf-8") as f: w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore") w.writeheader() @@ -394,6 +472,7 @@ def main(): choices=["float32", "float64", "all"]) parser.add_argument("--warmup", type=int, default=WARMUP) parser.add_argument("--iters", type=int, default=ITERS) + parser.add_argument("--timing", action="store_true", help="Add process_gpu_ms/compute_ms split timing columns") args = parser.parse_args() paths = [] @@ -413,27 +492,27 @@ def main(): print(f"GPU: {torch.cuda.get_device_name(0)} | Files: {len(paths)} | CSV: {args.csv}") dtype_map = {"float32": torch.float32, "float64": torch.float64} dtype_filter = None if args.dtype == "all" else dtype_map[args.dtype] - run_all_csv(paths, args.csv, args.warmup, args.iters, dtype_filter) + run_all_csv(paths, args.csv, args.warmup, args.iters, dtype_filter, timing=args.timing) return dtype_map = {"float32": torch.float32, "float64": torch.float64} dtypes = VALUE_DTYPES if args.dtype == "all" else [dtype_map[args.dtype]] for dtype in dtypes: dname = str(dtype).replace("torch.", "") - print("=" * 170) + print("=" * (210 if args.timing else 190)) print(f"FLAGSPARSE SpMV Optimisation A/B Test") print(f"GPU: {torch.cuda.get_device_name(0)} | dtype: {dname} | Files: {len(paths)}") print( - "Base = FlagSparse baseline (fp64-accum for fp32). " - "Opt = FlagSparse CSR-Vector (fp32/fp64 native accum, wide tiles, few launches). " - "Opt(ms) = Sym(ms) CPU wall time + Comp(ms) CUDA event time. " + "Base = prepared baseline kernel. " + "Opt = CSR-Vector with bucket execution-plan data. " + "Base/Opt ms = process_cpu_ms + GPU event time; --timing adds process_gpu_ms/compute_ms. " "Speedup = Base/Opt or Ref/Opt." ) - print(SEP) - print(HEADER) - print(SEP) - results = run_batch(paths, dtype, torch.int32, args.warmup, args.iters) - print(SEP) + print(_sep(args.timing)) + print(_header(args.timing)) + print(_sep(args.timing)) + results = run_batch(paths, dtype, torch.int32, args.warmup, args.iters, timing=args.timing) + print(_sep(args.timing)) passed = sum(1 for r in results if r["status"] == "PASS") print(f"Passed: {passed} / {len(results)}") From ada383fbc8d5a5de6e95cadc1077cb8b17569685 Mon Sep 17 00:00:00 2001 From: zyq1105331849 <1105331849@qq.com> Date: Sun, 28 Jun 2026 21:36:49 +0800 Subject: [PATCH 2/3] uniform_timing --- src/flagsparse/sparse_operations/spmm_csr.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/flagsparse/sparse_operations/spmm_csr.py b/src/flagsparse/sparse_operations/spmm_csr.py index e03a5c8..cd4853d 100644 --- a/src/flagsparse/sparse_operations/spmm_csr.py +++ b/src/flagsparse/sparse_operations/spmm_csr.py @@ -1178,8 +1178,10 @@ def _spmm_csr_alg1_build_bucket_descriptors(rows_flat, counts, offsets): def _spmm_csr_alg1_build_process_plan(prepared, *, timing=False): - if prepared.data.dtype not in (torch.float32, torch.float64): - raise TypeError("spmm_csr_alg1 only supports float32 and float64") + if prepared.data.dtype not in (torch.float32, torch.float64, torch.complex64, torch.complex128): + raise TypeError( + "spmm_csr_alg1 only supports float32, float64, complex64, and complex128" + ) device = prepared.data.device row_count = int(prepared.row_lengths.numel()) row_index_dtype = torch.int32 if row_count <= _INDEX_LIMIT_INT32 else torch.int64 @@ -1891,8 +1893,10 @@ def _spmm_csr_alg2_build_bucket_descriptors(rows_flat, counts, offsets, dtype): def _spmm_csr_alg2_build_process_plan(prepared, *, timing=False): - if prepared.data.dtype not in (torch.float32, torch.float64): - raise TypeError("spmm_csr_alg2 only supports float32 and float64") + if prepared.data.dtype not in (torch.float32, torch.float64, torch.complex64, torch.complex128): + raise TypeError( + "spmm_csr_alg2 only supports float32, float64, complex64, and complex128" + ) device = prepared.data.device row_count = int(prepared.row_lengths.numel()) row_index_dtype = torch.int32 if row_count <= _INDEX_LIMIT_INT32 else torch.int64 From 73115bc3929097995a7d4c6bb4c47b69d5c2ba2d Mon Sep 17 00:00:00 2001 From: zyq1105331849 <1105331849@qq.com> Date: Wed, 1 Jul 2026 16:41:00 +0800 Subject: [PATCH 3/3] fix_index --- tests/test_spmm_csr.py | 102 +++++++++++++++++++++++++++++------------ 1 file changed, 73 insertions(+), 29 deletions(-) diff --git a/tests/test_spmm_csr.py b/tests/test_spmm_csr.py index 2d0c489..f6399e3 100644 --- a/tests/test_spmm_csr.py +++ b/tests/test_spmm_csr.py @@ -38,8 +38,13 @@ "complex64": torch.complex64, "complex128": torch.complex128, } +INDEX_DTYPE_MAP = { + "int32": torch.int32, + "int64": torch.int64, +} DEFAULT_DTYPE_NAMES = ("float32", "float64", "complex64", "complex128") DEFAULT_RUN_DTYPE_NAMES = ("float32", "float64") +DEFAULT_INDEX_DTYPE_NAMES = ("int32", "int64") DEFAULT_OP_NAMES = tuple(spmm_ops.SPMM_OP_NAMES.values()) CUSPARSE_DTYPES = (torch.float32, torch.float64, torch.complex64, torch.complex128) MAIN_CSR_SPMM_ALGORITHMS = { @@ -56,6 +61,7 @@ PERF_FIELDS = [ "matrix", "dtype", + "index_dtype", "op", "layout", "alg", @@ -84,6 +90,7 @@ DIAG_FIELDS = [ "matrix", "dtype", + "index_dtype", "op", "layout", "alg", @@ -112,6 +119,7 @@ BEST_FIELDS = [ "matrix", "dtype", + "index_dtype", "op", "layout", "best_alg", @@ -336,6 +344,7 @@ def _time_route(prepared, B, alg, warmup, iters, timing=False, diagnose=False, l def _skip_row( path, dtype, + index_dtype_name, op, layout, alg, @@ -353,6 +362,7 @@ def _skip_row( row = { "matrix": os.path.basename(path), "dtype": _dtype_name(dtype), + "index_dtype": index_dtype_name, "op": op, "layout": layout, "alg": alg, @@ -412,10 +422,25 @@ def _time_cusparse(data, indices, indptr, shape, B, op, warmup, iters, layout="r return None, None, str(exc) -def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, run_cusparse, timing, diagnose): +def run_one_case( + path, + dtype, + index_dtype_name, + index_dtype, + op, + layout, + alg_names, + dense_cols, + warmup, + iters, + run_cusparse, + timing, + diagnose, +): device = torch.device("cuda") data, indices, indptr, shape = load_mtx_to_csr_torch(path, dtype=dtype, device=device) - indices = indices.to(torch.int32) + indices = indices.to(index_dtype) + indptr = indptr.to(index_dtype) n_rows, n_cols = shape b_rows = n_rows if op in ("trans", "conj") else n_cols B = _materialize_dense_layout_for_test( @@ -444,6 +469,7 @@ def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, _skip_row( path, dtype, + index_dtype_name, op, layout, alg, @@ -474,6 +500,7 @@ def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, _skip_row( path, dtype, + index_dtype_name, op, layout, alg, @@ -496,6 +523,7 @@ def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, row = { "matrix": os.path.basename(path), "dtype": _dtype_name(dtype), + "index_dtype": index_dtype_name, "op": op, "layout": layout, "alg": result["alg"], @@ -526,6 +554,7 @@ def run_one_case(path, dtype, op, layout, alg_names, dense_cols, warmup, iters, diag = { "matrix": os.path.basename(path), "dtype": _dtype_name(dtype), + "index_dtype": index_dtype_name, "op": op, "layout": layout, "alg": result["alg"], @@ -542,15 +571,16 @@ def _best_rows(rows): for row in rows: if row.get("status") != "PASS" or row.get("ms") is None: continue - key = (row["matrix"], row["dtype"], row["op"], row["layout"]) + key = (row["matrix"], row["dtype"], row["index_dtype"], row["op"], row["layout"]) groups.setdefault(key, []).append(row) best = [] - for (matrix, dtype, op, layout), group in sorted(groups.items()): + for (matrix, dtype, index_dtype, op, layout), group in sorted(groups.items()): selected = min(group, key=lambda item: item["ms"]) best.append( { "matrix": matrix, "dtype": dtype, + "index_dtype": index_dtype, "op": op, "layout": layout, "best_alg": selected["alg"], @@ -572,7 +602,7 @@ def _write_csv(path, rows, fields): def _print_row(row): print( - f"{row['matrix']:<28} {row['dtype']:<10} {row['op']:<5} {row['layout']:<4} {row['alg']:<10} " + f"{row['matrix']:<28} {row['dtype']:<10} {row['index_dtype']:<5} {row['op']:<5} {row['layout']:<4} {row['alg']:<10} " f"{_fmt(row['ms']):>9} {_fmt(row['gpu_ms']):>9} {_fmt(row['process_cpu_ms']):>9} " f"{_fmt(row['torch_ms']):>9} {_fmt(row['cusparse_ms']):>9} " f"{_fmt(row['torch_vs_alg_speedup'], 2):>9} {_fmt(row['cusparse_vs_alg_speedup'], 2):>9} " @@ -593,6 +623,7 @@ def main(): ), ) parser.add_argument("--op", default="all", help="all or comma-separated ops: non,trans,conj") + parser.add_argument("--index-dtype", default="all", help="int32, int64, or all") parser.add_argument("--layout", default="row", help="row, col, or all") parser.add_argument("--dense-cols", type=int, default=32) parser.add_argument("--warmup", type=int, default=10) @@ -619,6 +650,12 @@ def main(): explicit_names=tuple(DTYPE_MAP), ) op_names = _parse_csv_names(args.op, DEFAULT_OP_NAMES, "--op") + index_dtype_names = _parse_csv_names( + args.index_dtype, + DEFAULT_INDEX_DTYPE_NAMES, + "--index-dtype", + explicit_names=tuple(INDEX_DTYPE_MAP), + ) layout_names = _layout_names(args.layout) alg_names = _parse_algs(args.alg) except ValueError as exc: @@ -629,35 +666,42 @@ def main(): rows = [] diag_rows = [] print( - f"{'Matrix':<28} {'DType':<10} {'Op':<5} {'Lay':<4} {'Alg':<10} " + f"{'Matrix':<28} {'DType':<10} {'Idx':<5} {'Op':<5} {'Lay':<4} {'Alg':<10} " f"{'ms':>9} {'gpu_ms':>9} {'cpu_ms':>9} {'torch':>9} {'cu':>9} " f"{'PT/Alg':>9} {'CU/Alg':>9} {'ErrPT':>10} {'Status':>6}" ) for dtype_name in dtype_names: dtype = DTYPE_MAP[dtype_name] - for op in op_names: - for layout in layout_names: - for path in paths: - try: - case_rows, case_diag = run_one_case( - path, - dtype, - op, - layout, - alg_names, - args.dense_cols, - args.warmup, - args.iters, - not args.no_cusparse, - args.timing, - args.diagnose, - ) - rows.extend(case_rows) - diag_rows.extend(case_diag) - for row in case_rows: - _print_row(row) - except Exception as exc: - print(f" ERROR on {os.path.basename(path)} dtype={dtype_name} op={op} layout={layout}: {exc}") + for index_dtype_name in index_dtype_names: + index_dtype = INDEX_DTYPE_MAP[index_dtype_name] + for op in op_names: + for layout in layout_names: + for path in paths: + try: + case_rows, case_diag = run_one_case( + path, + dtype, + index_dtype_name, + index_dtype, + op, + layout, + alg_names, + args.dense_cols, + args.warmup, + args.iters, + not args.no_cusparse, + args.timing, + args.diagnose, + ) + rows.extend(case_rows) + diag_rows.extend(case_diag) + for row in case_rows: + _print_row(row) + except Exception as exc: + print( + f" ERROR on {os.path.basename(path)} dtype={dtype_name} " + f"index_dtype={index_dtype_name} op={op} layout={layout}: {exc}" + ) if args.csv: csv_path = _normalize_csv_path(args.csv) _write_csv(csv_path, rows, fields)