diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index df1543761de4..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 == "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 df4aca033a42..c941fb9ea9ec 100644 --- a/tests/python/relax/test_backend_dispatch_sort_scan.py +++ b/tests/python/relax/test_backend_dispatch_sort_scan.py @@ -415,10 +415,12 @@ 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), ], ) -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") @@ -436,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(): @@ -449,6 +453,44 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +@pytest.mark.parametrize("target_kind", ["metal", "webgpu", "cuda"]) +@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 + 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"): + 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=expected_bits + ) + assert_structural_equal(mod["gpu_2d_continuous_cumsum"], expected) + 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."""