From cbe1c80056383810dd141735f035a962b161d007 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 13 Sep 2026 15:26:18 -0400 Subject: [PATCH 1/2] [FIX][Relax] Use 32-bit scan hierarchy thresholds on Metal Generate the cumsum hierarchy for the signed index width supported by Metal, matching the existing WebGPU path. This prevents ForceNarrowIndexToInt32 from rejecting unreachable 64-bit thresholds. Add target-dispatch coverage for Metal, WebGPU, and CUDA, and execute the numerical cumsum test on Metal. --- .../tvm/relax/backend/dispatch_sort_scan.py | 2 +- .../relax/test_backend_dispatch_sort_scan.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index df1543761de4..006252ecf920 100644 --- a/python/tvm/relax/backend/dispatch_sort_scan.py +++ b/python/tvm/relax/backend/dispatch_sort_scan.py @@ -175,7 +175,7 @@ def visit_call_(self, call: relax.Call) -> relax.Expr: kernel = gpu_2d_continuous_cumsum( in_dtype=in_dtype, out_dtype=out_dtype, - index_bits=32 if tgt.kind.name == "webgpu" else 64, + index_bits=32 if tgt.kind.name in ("metal", "webgpu") else 64, ) kernel_name = "gpu_2d_continuous_cumsum" else: diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py index df4aca033a42..8f03dd061ded 100644 --- a/tests/python/relax/test_backend_dispatch_sort_scan.py +++ b/tests/python/relax/test_backend_dispatch_sort_scan.py @@ -415,6 +415,7 @@ def foo(x: R.Tensor((2, 3), "float32", "vulkan")): "target", [ pytest.param("cuda", marks=pytest.mark.gpu), + pytest.param("metal", marks=pytest.mark.gpu), pytest.param({"kind": "vulkan", "supports_int64": True}, marks=pytest.mark.gpu), ], ) @@ -449,6 +450,31 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +@pytest.mark.parametrize("target_kind", ["metal", "webgpu", "cuda"]) +def test_dispatch_cumsum_index_width(target_kind): + """32-bit targets must not generate out-of-range hierarchy thresholds.""" + from tvm.relax.backend.gpu_generic import gpu_2d_continuous_cumsum + + @I.ir_module + class Module: + @R.function + def main(x: R.Tensor(("m", "n"), "float32")): + gv = R.cumsum(x, axis=-1) + return gv + + with tvm.target.Target(target_kind, host="llvm"): + mod = DispatchSortScan()(Module) + + index_bits = 64 if target_kind == "cuda" else 32 + expected = gpu_2d_continuous_cumsum( + in_dtype="float32", out_dtype="float32", index_bits=index_bits + ) + assert_structural_equal(mod["gpu_2d_continuous_cumsum"], expected) + if index_bits == 32: + # This previously failed on Metal with a 2**35 IntImm. + tirx.transform.ForceNarrowIndexToInt32()(mod) + + @pytest.mark.gpu def test_dispatch_cumprod_cuda_large_batch(): """Test that GPU scan supports more batches than CUDA's grid-y limit.""" From f469edaca053522427dee69ff047b74cb90110e7 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 13 Sep 2026 19:23:24 -0400 Subject: [PATCH 2/2] [FIX][Relax] Make the scan index-width policy explicit Restore the default 64-bit continuous cumsum hierarchy for Metal. Add an optional index_bits argument to DispatchSortScan so callers using forced int32 narrowing can request a compatible hierarchy without restricting other Metal pipelines. Preserve the 32-bit WebGPU default and reject explicit 64-bit WebGPU scan budgets. Validate defaults, overrides, invalid budgets, and numerical Metal execution with both index policies. Document that the option neither changes tensor dtypes nor inserts narrowing or runtime bounds checks. --- .../tvm/relax/backend/dispatch_sort_scan.py | 31 +++++++++++++++-- .../relax/test_backend_dispatch_sort_scan.py | 34 ++++++++++++++----- 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index 006252ecf920..7918c92e03c8 100644 --- a/python/tvm/relax/backend/dispatch_sort_scan.py +++ b/python/tvm/relax/backend/dispatch_sort_scan.py @@ -37,9 +37,10 @@ class SortScanDispatcher(BackendDispatcher): calls_to_update: dict[GlobalVar, Target] - def __init__(self, mod): + def __init__(self, mod, index_bits: int | None = None): super().__init__(mod) self.calls_to_update = {} + self.index_bits = index_bits def apply_dlight_gpu_fallback( self, @@ -172,10 +173,15 @@ def visit_call_(self, call: relax.Call) -> relax.Expr: if normalized_axis == len(shape) - 1: outer = reduce(mul, shape_values[:-1], 1) kernel_shape = relax.ShapeExpr([outer, shape[-1]]) + index_bits = self.index_bits + if index_bits is None: + index_bits = 32 if tgt.kind.name == "webgpu" else 64 + if tgt.kind.name == "webgpu" and index_bits != 32: + raise ValueError("WebGPU scan kernels require index_bits=32") kernel = gpu_2d_continuous_cumsum( in_dtype=in_dtype, out_dtype=out_dtype, - index_bits=32 if tgt.kind.name in ("metal", "webgpu") else 64, + index_bits=index_bits, ) kernel_name = "gpu_2d_continuous_cumsum" else: @@ -263,10 +269,29 @@ def allocate_workspace(self, call: relax.Call) -> relax.Var: class DispatchSortScan: """ Pass to dispatch scan and sort operators to platform dependent implementation. + + Parameters + ---------- + index_bits : Optional[int] + Signed index-width budget for the generated continuous GPU cumsum hierarchy. + Must be 32 or 64. By default, use 32 for WebGPU and 64 for other targets. + WebGPU does not support an explicit 64-bit budget. + + Pipelines that subsequently force indices to int32 should request 32 to + avoid generating hierarchy thresholds outside the signed int32 range. + The caller must ensure runtime indices fit the requested width; this + option does not insert runtime bounds checks. + This option does not narrow the generated TIR, change tensor dtypes, or + affect other sort/scan implementations. """ + def __init__(self, index_bits: int | None = None): + if index_bits not in (None, 32, 64): + raise ValueError("index_bits must be either 32 or 64") + self.index_bits = index_bits + def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule: - sort_scan_dispater = SortScanDispatcher(mod) + sort_scan_dispater = SortScanDispatcher(mod, self.index_bits) for gv, func in mod.functions_items(): if isinstance(func, relax.Function): func = sort_scan_dispater.visit_expr(func) diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py index 8f03dd061ded..c941fb9ea9ec 100644 --- a/tests/python/relax/test_backend_dispatch_sort_scan.py +++ b/tests/python/relax/test_backend_dispatch_sort_scan.py @@ -419,7 +419,8 @@ def foo(x: R.Tensor((2, 3), "float32", "vulkan")): pytest.param({"kind": "vulkan", "supports_int64": True}, marks=pytest.mark.gpu), ], ) -def test_dispatch_cumsum_gpu(target): +@pytest.mark.parametrize("index_bits", [None, 32, 64]) +def test_dispatch_cumsum_gpu(target, index_bits): """Test cumsum kernel dispatch and numerical correctness""" if not tvm.testing.device_enabled(target): pytest.skip(f"{target} not enabled") @@ -437,7 +438,9 @@ def main(x: R.Tensor(("m", "n"), "int32")): np_data = np.random.randint(0, 10, size).astype("int32") np_cumsum = np.cumsum(np_data, axis=-1) with tvm.target.Target(target): - mod = DispatchSortScan()(Module) + mod = DispatchSortScan(index_bits=index_bits)(Module) + if index_bits == 32: + mod = tirx.transform.ForceNarrowIndexToInt32()(mod) ex = tvm.compile(mod, target) def run_and_check(): @@ -451,8 +454,9 @@ def run_and_check(): @pytest.mark.parametrize("target_kind", ["metal", "webgpu", "cuda"]) -def test_dispatch_cumsum_index_width(target_kind): - """32-bit targets must not generate out-of-range hierarchy thresholds.""" +@pytest.mark.parametrize("index_bits", [None, 32, 64]) +def test_dispatch_cumsum_index_width(target_kind, index_bits): + """Respect the caller's index budget without restricting Metal's default.""" from tvm.relax.backend.gpu_generic import gpu_2d_continuous_cumsum @I.ir_module @@ -463,18 +467,30 @@ def main(x: R.Tensor(("m", "n"), "float32")): return gv with tvm.target.Target(target_kind, host="llvm"): - mod = DispatchSortScan()(Module) - - index_bits = 64 if target_kind == "cuda" else 32 + if target_kind == "webgpu" and index_bits == 64: + with pytest.raises(ValueError, match="WebGPU scan kernels require index_bits=32"): + DispatchSortScan(index_bits=index_bits)(Module) + return + mod = DispatchSortScan(index_bits=index_bits)(Module) + + expected_bits = ( + index_bits if index_bits is not None else (32 if target_kind == "webgpu" else 64) + ) expected = gpu_2d_continuous_cumsum( - in_dtype="float32", out_dtype="float32", index_bits=index_bits + in_dtype="float32", out_dtype="float32", index_bits=expected_bits ) assert_structural_equal(mod["gpu_2d_continuous_cumsum"], expected) - if index_bits == 32: + if expected_bits == 32: # This previously failed on Metal with a 2**35 IntImm. tirx.transform.ForceNarrowIndexToInt32()(mod) +@pytest.mark.parametrize("index_bits", [0, 16, 128]) +def test_dispatch_cumsum_invalid_index_width(index_bits): + with pytest.raises(ValueError, match="index_bits must be either 32 or 64"): + DispatchSortScan(index_bits=index_bits) + + @pytest.mark.gpu def test_dispatch_cumprod_cuda_large_batch(): """Test that GPU scan supports more batches than CUDA's grid-y limit."""