[Fix][Relax] Allow callers to specify scan index-width budget

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)

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."""