Auto-select index_dtype

Fixes #1123

stack-info: PR: #1131, branch: jansel/stack/227

Author: jansel

docs/api/kernel.md

@@ -90,6 +90,8 @@ result = bound_static(torch.randn(100, 50))  # Must be exactly [100, 50]
 
 ```{warning}
 Helion shape-specializes kernels by default (`static_shapes=True`) for the best performance. Bound kernels and caches require tensors with the exact same shapes and strides as the examples you compile against. Set `static_shapes=False` if you need the same compiled kernel to serve many shapes.
+With dynamic shapes (`static_shapes=False`), Helion also specializes on if a tensor shape is 0 or 1 and whether a tensor needs 64-bit indexing (more than ``2**31 - 1`` elements).
+This 64-bit indexing specialization can be avoided by setting `index_dtype=torch.int64`.
 ```
 
 ### BoundKernel Methods
@@ -139,6 +141,10 @@ Kernels are automatically cached based on:
 
 By default (`static_shapes=True`), Helion treats shapes and strides as compile-time constants, baking them into generated Triton code for the best performance. To reuse a single compiled kernel across size variations, set `static_shapes=False`, which instead buckets each dimension as `{0, 1, ≥2}` and allows more inputs to share the same cache entry.
 
+```{note}
+Dynamic buckets also track whether any tensor exceeds the ``torch.int32`` indexing limit so that cache entries diverge as soon as large inputs show up. If your deployment regularly touches that regime, pin ``index_dtype=torch.int64`` on the kernel to avoid a cache miss or limit errors.
+```
+
 ```python
 # These create separate cache entries
 tensor_float = torch.randn(100, dtype=torch.float32, device='cuda')

docs/api/settings.md

@@ -81,8 +81,10 @@ def my_kernel(x: torch.Tensor) -> torch.Tensor:
 
 .. autoattribute:: Settings.index_dtype
 
-   The data type used for index variables in generated code. Default is ``torch.int32``.
-   Override via ``HELION_INDEX_DTYPE=int64`` (or any ``torch.<dtype>`` name).
+   The data type used for index variables in generated code. By default Helion auto-selects
+   between ``torch.int32`` and ``torch.int64`` based on whether any input tensor exceeds
+   ``torch.iinfo(torch.int32).max`` elements. Override via ``HELION_INDEX_DTYPE=<dtype>``
+   (or set it to ``auto`` to keep the automatic behavior).
 
 .. autoattribute:: Settings.dot_precision
 
@@ -259,7 +261,7 @@ Built-in values for ``HELION_AUTOTUNER`` include ``"PatternSearch"``, ``"Differe
 | Environment Variable | Maps To | Description |
 |----------------------|---------|-------------|
 | ``TRITON_F32_DEFAULT`` | ``dot_precision`` | Sets default floating-point precision for Triton dot products (``"tf32"``, ``"tf32x3"``, ``"ieee"``). |
-| ``HELION_INDEX_DTYPE`` | ``index_dtype`` | Choose the default index dtype (accepts any ``torch.<dtype>`` name, e.g. ``int64``). |
+| ``HELION_INDEX_DTYPE`` | ``index_dtype`` | Choose the index dtype (accepts any ``torch.<dtype>`` name, e.g. ``int64``), or set to ``auto``/unset to allow Helion to pick ``int32`` vs ``int64`` based on input sizes. |
 | ``HELION_STATIC_SHAPES`` | ``static_shapes`` | Set to ``0``/``false`` to disable global static shape specialization. |
 | ``HELION_PERSISTENT_RESERVED_SMS`` | ``persistent_reserved_sms`` | Reserve this many streaming multiprocessors when launching persistent kernels (``0`` uses all available SMs). |
 | ``HELION_FORCE_AUTOTUNE`` | ``force_autotune`` | Force the autotuner to run even when explicit configs are provided. |

docs/deployment_autotuning.md

@@ -154,8 +154,12 @@ determines when to re-benchmark. Options include:
 - **`static_shapes=False`:** switch to bucketed dynamic shapes. Helion
   reuses results as long as tensor dtypes and device types stay constant.
   Shape changes only trigger a re-selection when a dimension size crosses
-  the buckets `{0, 1, ≥2}`. Use this when you need one compiled kernel to
-  handle many input sizes.
+  the buckets `{0, 1, ≥2}`. Helion also tracks whether any tensor exceeds the
+  `torch.int32` indexing limit (more than ``2**31 - 1`` elements) and will
+  automatically regenerate code with 64-bit indexing in that case. Use this
+  mode when you need one compiled kernel to handle many input sizes, and pin
+  ``@helion.kernel(..., index_dtype=torch.int64)`` if large tensors are the norm
+  so you avoid an extra specialization boundary.
 
 - **Custom keys:** pass `key=` to group calls however you like.
 This custom key is in addition to the above.
@@ -206,10 +210,13 @@ exact shape/stride signature of the example inputs.  The generated code
 has shapes baked in, which often provides a performance boost.
 
 - With `static_shapes=False` it will specialize on the input dtypes,
-device types, and whether each dynamic dimension falls into the 0, 1,
-or ≥2 bucket.  Python types are also specialized.  For dimensions that
-can vary across those buckets, supply representative inputs ≥2 to avoid
-excessive specialization.
+  device types, and whether each dynamic dimension falls into the 0, 1,
+  or ≥2 bucket.  Python types are also specialized.  For dimensions that
+  can vary across those buckets, supply representative inputs ≥2 to avoid
+  excessive specialization.  Just like the autotuning flow above, Helion
+  records whether any tensor crosses the int32 indexing limit when
+  `static_shapes=False`; explicitly set `index_dtype=torch.int64` if your
+  deployment commonly exceeds that threshold to avoid recompilation.
 
 If you need to support multiple input types, bind multiple times with
 representative inputs.

helion/_compiler/aten_lowering.py

@@ -460,9 +460,7 @@ def codegen_iota(ctx: LoweringContext, node: Node) -> object:
     """Generate tl.arange for torch.ops.prims.iota.default operations with automatic power-of-2 padding."""
     start = node.kwargs.get("start", 0)
     step = node.kwargs.get("step", 1)
-    dtype = (
-        node.kwargs.get("dtype") or CompileEnvironment.current().settings.index_dtype
-    )
+    dtype = node.kwargs.get("dtype") or CompileEnvironment.current().index_dtype
     assert isinstance(dtype, torch.dtype)
     (length_arg,) = node.args  # expecting a single argument for length
 

helion/_compiler/compile_environment.py

@@ -73,12 +73,21 @@ class CompileEnvironment:
     No config or codegen specific state should be stored here.
     """
 
-    def __init__(self, device: torch.device, settings: Settings) -> None:
+    def __init__(
+        self,
+        device: torch.device,
+        settings: Settings,
+        *,
+        index_dtype: torch.dtype | None = None,
+    ) -> None:
         from ..autotuner.config_spec import ConfigSpec
 
         super().__init__()
         self.device = device
         self.settings = settings
+        self.index_dtype: torch.dtype = (
+            index_dtype or settings.index_dtype or torch.int32
+        )
         # TODO(jansel): make backend configurable
         self.backend = "triton"
         self.shape_env = ShapeEnv(
@@ -383,7 +392,7 @@ def known_multiple(self, a: sympy.Expr, b: int | torch.SymInt) -> bool:
 
     def triton_index_type(self) -> str:
         """tl.int32 or tl.int64 depending on Settings()"""
-        return triton_type(self.settings.index_dtype)
+        return triton_type(self.index_dtype)
 
     def sympy_debug(self, expr: sympy.Expr) -> str:
         return str(expr.xreplace(self.debug_shape_renames))

helion/_compiler/indexing_strategy.py

@@ -665,7 +665,7 @@ def create(
         env = CompileEnvironment.current()
         dtype = env.triton_index_type()
         if dtype == "tl.int32" and SubscriptIndexing._needs_int64(fake_value):
-            raise exc.IndexOffsetOutOfRangeForInt32(env.settings.index_dtype)
+            raise exc.IndexOffsetOutOfRangeForInt32(env.index_dtype)
 
         def _is_size_one(size: int | torch.SymInt) -> bool:
             return env.known_equal(size, 1)

helion/_compiler/reduction_strategy.py

@@ -332,7 +332,7 @@ def codegen_reduction(
                 )
             else:
                 acc_index = self.fn.new_var(f"{state.fx_node.name}_acc_index", dce=True)
-                index_dtype = CompileEnvironment.current().settings.index_dtype
+                index_dtype = CompileEnvironment.current().index_dtype
                 device_loop.outer_prefix.append(
                     statement_from_string(
                         f"{acc_index} = tl.full({shape}, {torch.iinfo(index_dtype).max!r}, {triton_type(index_dtype)})"

helion/exc.py

@@ -136,6 +136,13 @@ class IndexOffsetOutOfRangeForInt32(BaseError):
     )
 
 
+class InputTensorNumelExceedsIndexType(BaseError):
+    message = (
+        "Kernel index_dtype is {index_dtype}, but input input tensor is too large to fit. "
+        "Use @helion.kernel(index_dtype=torch.int64)."
+    )
+
+
 class DataDependentOutputShapeNotSupported(BaseError):
     message = (
         "{op_desc} is not supported in Helion device loops because it produces "

helion/language/creation_ops.py

@@ -209,7 +209,7 @@ def arange(
     """
     env = CompileEnvironment.current()
     if dtype is None:
-        dtype = env.settings.index_dtype
+        dtype = env.index_dtype
     return torch.arange(
         *args,
         **kwargs,

helion/language/tile_ops.py

@@ -50,7 +50,7 @@ def _(tile: torch.SymInt) -> torch.Tensor:
     assert isinstance(tile, torch.SymInt)
     env = CompileEnvironment.current()
     assert env.get_block_id(tile) is not None
-    return torch.empty([tile], dtype=env.settings.index_dtype, device=env.device)
+    return torch.empty([tile], dtype=env.index_dtype, device=env.device)
 
 
 @_decorators.codegen(tile_index, "triton")