Add tl.range disallow_acc_multi_buffer to autotuner

stack-info: PR: #228, branch: jansel/stack/72

Author: jansel

README.md

@@ -199,6 +199,11 @@ Contains one entry per loop dimension, specifying the unroll factor for
 Contains one entry per loop dimension, specifying the number of stages for
 `tl.range()` calls. Values less than 1 omit the `num_stages` parameter.
 
+* **range\_multi\_buffers** (`list[bool | None]`):
+Contains one entry per loop dimension, controlling the `disallow_acc_multi_buffer`
+parameter for `tl.range()` calls. `True` allows multi-buffer (sets `disallow_acc_multi_buffer=False`),
+`False` disallows multi-buffer (sets `disallow_acc_multi_buffer=True`), and `None` omits the parameter.
+
 * **reduction\_loops** (`list[int | None]`):
 Contains one entry per reduction dimension (see
 `examples/softmax.py`). Using `None` triggers a persistent reduction,

helion/_compiler/tile_strategy.py

@@ -140,6 +140,12 @@ def get_tl_range_kwargs(self, state: CodegenState, block_idx: int) -> str:
         if range_num_stages > 0:
             kwargs.append(f"num_stages={range_num_stages}")
 
+        range_multi_buffer = env.config_spec.range_multi_buffers.config_get(
+            state.config.range_multi_buffers, block_idx, None
+        )
+        if range_multi_buffer is not None:
+            kwargs.append(f"disallow_acc_multi_buffer={not range_multi_buffer}")
+
         if kwargs:
             return f", {', '.join(kwargs)}"
         return ""

helion/autotuner/config_spec.py

@@ -38,6 +38,7 @@
         "flatten_loops",
         "range_unroll_factors",
         "range_num_stages",
+        "range_multi_buffers",
         "num_warps",
         "num_stages",
         "use_yz_grid",
@@ -69,6 +70,9 @@ class ConfigSpec:
     range_num_stages: BlockIdSequence[RangeNumStagesSpec] = dataclasses.field(
         default_factory=BlockIdSequence
     )
+    range_multi_buffers: BlockIdSequence[RangeMultiBufferSpec] = dataclasses.field(
+        default_factory=BlockIdSequence
+    )
     user_defined_tunables: dict[str, ConfigSpecFragment] = dataclasses.field(
         default_factory=dict
     )
@@ -80,6 +84,7 @@ def _remove_duplicates(self) -> None:
         self.flatten_loops._remove_duplicates()
         self.range_unroll_factors._remove_duplicates()
         self.range_num_stages._remove_duplicates()
+        self.range_multi_buffers._remove_duplicates()
 
     def normalize(self, config: helion.Config | dict[str, object]) -> None:
         """Normalize the config to match the block_sizes and validate the config."""
@@ -95,6 +100,7 @@ def normalize(self, config: helion.Config | dict[str, object]) -> None:
             "flatten_loop",
             "range_unroll_factor",
             "range_num_stage",
+            "range_multi_buffer",
         ):
             if name in config:
                 names = f"{name}s"
@@ -110,6 +116,7 @@ def normalize(self, config: helion.Config | dict[str, object]) -> None:
             ("reduction_loops", self.reduction_loops, True),
             ("range_unroll_factors", self.range_unroll_factors, True),
             ("range_num_stages", self.range_num_stages, True),
+            ("range_multi_buffers", self.range_multi_buffers, True),
         ]:
             config[name] = mapping._normalize(
                 name, config.get(name, ()), flatten=flatten
@@ -122,6 +129,7 @@ def normalize(self, config: helion.Config | dict[str, object]) -> None:
             "reduction_loops",
             "range_unroll_factors",
             "range_num_stages",
+            "range_multi_buffers",
         ):
             if not config[name]:
                 config.pop(name)
@@ -153,6 +161,7 @@ def flat_config(self, fn: Callable[[ConfigSpecFragment], object]) -> helion.Conf
             "reduction_loops": self.reduction_loops._flat_config(self, fn),
             "range_unroll_factors": self.range_unroll_factors._flat_config(self, fn),
             "range_num_stages": self.range_num_stages._flat_config(self, fn),
+            "range_multi_buffers": self.range_multi_buffers._flat_config(self, fn),
             "num_warps": fn(NumWarpsFragment(1, 32, DEFAULT_NUM_WARPS)),
             "num_stages": fn(IntegerFragment(1, 8, DEFAULT_NUM_STAGES)),
             "indexing": fn(
@@ -181,6 +190,7 @@ def flat_config(self, fn: Callable[[ConfigSpecFragment], object]) -> helion.Conf
             "l2_groupings",
             "range_unroll_factors",
             "range_num_stages",
+            "range_multi_buffers",
         ):
             if not config[name]:
                 config.pop(name)
@@ -350,6 +360,20 @@ def _fill_missing(self) -> int:
         return 0
 
 
+class RangeMultiBufferSpec(_BlockIdItem):
+    def _fragment(self, base: ConfigSpec) -> EnumFragment:
+        return EnumFragment((None, False, True))
+
+    def _normalize(self, name: str, value: object) -> bool | None:
+        if value is not None and not isinstance(value, bool):
+            raise InvalidConfig(f"{name} must be a boolean or None, got {value!r}")
+        return value
+
+    def _fill_missing(self) -> None:
+        """Provide a value when not provided by the user."""
+        return None
+
+
 def _product(seq: Sequence[int]) -> int:
     """Return the product of the elements in the sequence."""
     return functools.reduce(operator.mul, seq, 1)

helion/language/loops.py

@@ -26,6 +26,7 @@
 from ..autotuner.config_spec import FlattenLoopSpec
 from ..autotuner.config_spec import L2GroupingSpec
 from ..autotuner.config_spec import LoopOrderSpec
+from ..autotuner.config_spec import RangeMultiBufferSpec
 from ..autotuner.config_spec import RangeNumStagesSpec
 from ..autotuner.config_spec import RangeUnrollFactorSpec
 from . import _decorators
@@ -249,6 +250,7 @@ def _add_config_choices(
         for block_id in block_ids:
             config_spec.range_unroll_factors.append(RangeUnrollFactorSpec([block_id]))
             config_spec.range_num_stages.append(RangeNumStagesSpec([block_id]))
+            config_spec.range_multi_buffers.append(RangeMultiBufferSpec([block_id]))
 
 
 def _allow_use_yz_grid(config_spec: ConfigSpec, block_ids: list[int]) -> bool:

helion/runtime/config.py

@@ -27,6 +27,7 @@ def __init__(
         reduction_loops: list[int | None] | None = None,
         range_unroll_factors: list[int] | None = None,
         range_num_stages: list[int] | None = None,
+        range_multi_buffers: list[bool | None] | None = None,
         num_warps: int | None = None,
         num_stages: int | None = None,
         use_yz_grid: bool | None = None,
@@ -44,6 +45,7 @@ def __init__(
             reduction_loops: Configures reduction loop behavior.
             range_unroll_factors: Loop unroll factors for tl.range calls.
             range_num_stages: Number of stages for tl.range calls.
+            range_multi_buffers: Controls disallow_acc_multi_buffer for tl.range calls.
             num_warps: Number of warps per block.
             num_stages: Number of stages for software pipelining.
             use_yz_grid: Whether to use yz grid dimensions.
@@ -59,6 +61,7 @@ def __init__(
             "reduction_loops": reduction_loops,
             "range_unroll_factors": range_unroll_factors,
             "range_num_stages": range_num_stages,
+            "range_multi_buffers": range_multi_buffers,
             "num_warps": num_warps,
             "num_stages": num_stages,
             "indexing": indexing,
@@ -152,6 +155,10 @@ def range_unroll_factors(self) -> list[int]:
     def range_num_stages(self) -> list[int]:
         return cast("list[int]", self.config.get("range_num_stages", []))
 
+    @property
+    def range_multi_buffers(self) -> list[bool | None]:
+        return cast("list[bool | None]", self.config.get("range_multi_buffers", []))
+
     @property
     def indexing(self) -> IndexingLiteral:
         return self.config.get("indexing", "pointer")  # type: ignore

pyproject.toml

@@ -58,8 +58,8 @@ select = [
     "TD004", "TRY002", "TRY203", "TRY401", "UP", "W", "YTT",
 ]
 ignore = [
-    "C409", "C419", "COM812", "E501", "ERA001", "FURB189", "G004", "PERF203",
-    "PERF401", "SIM102", "SIM108", "SIM115", "UP035", "UP038",
+    "C409", "C419", "COM812", "E501", "ERA001", "FURB189", "G004", "PERF203", "PERF401",
+    "RET501", "SIM102", "SIM108", "SIM115", "UP035", "UP038",
 ]
 extend-safe-fixes = ["TC", "UP045", "RUF013", "RSE102"]
 preview = true

test/test_autotuner.py

@@ -44,16 +44,16 @@ def test_config_fragment0(self):
         self.assertExpectedInline(
             "\n".join(map(repr, configs)),
             """\
-helion.Config(block_sizes=[16, 16, 16], loop_orders=[[0, 1]], l2_groupings=[1], range_unroll_factors=[0], range_num_stages=[0], num_warps=4, num_stages=3, indexing='pointer')
-helion.Config(block_sizes=[16, 16, 16], loop_orders=[[1, 0]], l2_groupings=[8], range_unroll_factors=[1], range_num_stages=[2], num_warps=2, num_stages=8, indexing='pointer')
-helion.Config(block_sizes=[16, 16, 32], loop_orders=[[1, 0]], l2_groupings=[4], range_unroll_factors=[2], range_num_stages=[1], num_warps=2, num_stages=3, indexing='tensor_descriptor')
-helion.Config(block_sizes=[64, 16, 32], loop_orders=[[1, 0]], l2_groupings=[1], range_unroll_factors=[0], range_num_stages=[0], num_warps=32, num_stages=4, indexing='block_ptr')
-helion.Config(block_sizes=[128, 16, 64], loop_orders=[[0, 1]], l2_groupings=[2], range_unroll_factors=[2], range_num_stages=[4], num_warps=8, num_stages=1, indexing='block_ptr')
-helion.Config(block_sizes=[16, 16, 16], loop_orders=[[1, 0]], l2_groupings=[2], range_unroll_factors=[4], range_num_stages=[4], num_warps=1, num_stages=4, indexing='tensor_descriptor')
-helion.Config(block_sizes=[16, 16, 16], loop_orders=[[0, 1]], l2_groupings=[8], range_unroll_factors=[1], range_num_stages=[4], num_warps=32, num_stages=3, indexing='tensor_descriptor')
-helion.Config(block_sizes=[16, 64, 64], loop_orders=[[0, 1]], l2_groupings=[16], range_unroll_factors=[2], range_num_stages=[0], num_warps=8, num_stages=2, indexing='block_ptr')
-helion.Config(block_sizes=[64, 16, 16], loop_orders=[[0, 1]], l2_groupings=[16], range_unroll_factors=[0], range_num_stages=[2], num_warps=4, num_stages=5, indexing='block_ptr')
-helion.Config(block_sizes=[16, 16, 16], loop_orders=[[1, 0]], l2_groupings=[1], range_unroll_factors=[4], range_num_stages=[1], num_warps=32, num_stages=7, indexing='block_ptr')""",
+helion.Config(block_sizes=[16, 16, 16], loop_orders=[[0, 1]], l2_groupings=[1], range_unroll_factors=[0], range_num_stages=[0], range_multi_buffers=[None], num_warps=4, num_stages=3, indexing='pointer')
+helion.Config(block_sizes=[16, 64, 32], loop_orders=[[1, 0]], l2_groupings=[8], range_unroll_factors=[1], range_num_stages=[2], range_multi_buffers=[None], num_warps=16, num_stages=8, indexing='pointer')
+helion.Config(block_sizes=[64, 16, 128], loop_orders=[[1, 0]], l2_groupings=[4], range_unroll_factors=[2], range_num_stages=[1], range_multi_buffers=[None], num_warps=32, num_stages=3, indexing='tensor_descriptor')
+helion.Config(block_sizes=[64, 16, 32], loop_orders=[[1, 0]], l2_groupings=[1], range_unroll_factors=[0], range_num_stages=[0], range_multi_buffers=[True], num_warps=2, num_stages=5, indexing='block_ptr')
+helion.Config(block_sizes=[16, 16, 16], loop_orders=[[0, 1]], l2_groupings=[4], range_unroll_factors=[4], range_num_stages=[3], range_multi_buffers=[None], num_warps=4, num_stages=7, indexing='tensor_descriptor')
+helion.Config(block_sizes=[16, 64, 32], loop_orders=[[0, 1]], l2_groupings=[2], range_unroll_factors=[4], range_num_stages=[2], range_multi_buffers=[None], num_warps=16, num_stages=3, indexing='block_ptr')
+helion.Config(block_sizes=[16, 16, 16], loop_orders=[[0, 1]], l2_groupings=[4], range_unroll_factors=[4], range_num_stages=[0], range_multi_buffers=[True], num_warps=1, num_stages=6, indexing='pointer')
+helion.Config(block_sizes=[16, 32, 64], loop_orders=[[0, 1]], l2_groupings=[1], range_unroll_factors=[0], range_num_stages=[4], range_multi_buffers=[True], num_warps=32, num_stages=1, indexing='tensor_descriptor')
+helion.Config(block_sizes=[32, 32, 64], loop_orders=[[1, 0]], l2_groupings=[64], range_unroll_factors=[4], range_num_stages=[1], range_multi_buffers=[True], num_warps=4, num_stages=2, indexing='pointer')
+helion.Config(block_sizes=[16, 64, 16], loop_orders=[[1, 0]], l2_groupings=[2], range_unroll_factors=[2], range_num_stages=[3], range_multi_buffers=[True], num_warps=16, num_stages=4, indexing='block_ptr')""",
         )
 
     @patch.object(_compat, "_supports_tensor_descriptor", lambda: True)

test/test_loops.py

@@ -1673,6 +1673,46 @@ def nested_loop_kernel(x: torch.Tensor) -> torch.Tensor:
             "tl.range(0, x_size_1.to(tl.int32), _BLOCK_SIZE_1, num_stages=3)", code3
         )
 
+    def test_range_multi_buffers(self):
+        @helion.kernel()
+        def nested_loop_kernel(x: torch.Tensor) -> torch.Tensor:
+            out = torch.empty_like(x)
+            # Outer loop becomes grid (no tl.range)
+            for tile_outer in hl.tile(x.size(0)):
+                # Inner loop becomes device loop with tl.range
+                for tile_inner in hl.tile(x.size(1)):
+                    out[tile_outer, tile_inner] = x[tile_outer, tile_inner] + 1
+            return out
+
+        # Test configuration validation - that range_multi_buffers works
+        args = (torch.randn([64, 32], device=DEVICE),)
+
+        # Test with range_multi_buffers = [None] (no disallow_acc_multi_buffer for device loop)
+        code_none, result_none = code_and_output(
+            nested_loop_kernel, args, block_sizes=[32, 16], range_multi_buffers=[None]
+        )
+
+        # Test with range_multi_buffers = [True] (disallow_acc_multi_buffer=False for device loop)
+        code_true, result_true = code_and_output(
+            nested_loop_kernel, args, block_sizes=[32, 16], range_multi_buffers=[True]
+        )
+
+        # Test with range_multi_buffers = [False] (disallow_acc_multi_buffer=True for device loop)
+        code_false, result_false = code_and_output(
+            nested_loop_kernel, args, block_sizes=[32, 16], range_multi_buffers=[False]
+        )
+
+        torch.testing.assert_close(result_none, result_true)
+        torch.testing.assert_close(result_none, result_false)
+        torch.testing.assert_close(result_none, args[0] + 1)
+        self.assertNotEqual(code_none, code_true)
+        self.assertNotEqual(code_none, code_false)
+        self.assertNotEqual(code_true, code_false)
+        # Check that disallow_acc_multi_buffer parameter appears in tl.range call
+        self.assertNotIn("disallow_acc_multi_buffer", code_none)
+        self.assertIn("disallow_acc_multi_buffer=False", code_true)
+        self.assertIn("disallow_acc_multi_buffer=True", code_false)
+
 
 if __name__ == "__main__":
     unittest.main()