Add additional tl.range choices to persistent loop (#287)

Author: jansel

helion/_compiler/program_id.py

@@ -24,6 +24,7 @@ class PIDInfo(NamedTuple):
     pid_var: str
     block_size_var: str
     numel: sympy.Expr
+    block_id: int
 
     def num_pids_expr(self, *, is_device: bool) -> str:
         """Get the number of PIDs expression for device or host."""
@@ -347,11 +348,16 @@ def __init__(self, is_blocked: bool = False) -> None:
             self.block_size_var: str = device_function.new_var("block_size")
             self.start_pid_var: str = device_function.new_var("start_pid")
             self.end_pid_var: str = device_function.new_var("end_pid")
-            self.range_expr: str = f"tl.range({self.start_pid_var}, {self.end_pid_var})"
+            self.range_kwargs: dict[str, str] = {
+                "begin": self.start_pid_var,
+                "end": self.end_pid_var,
+            }
         else:
-            self.range_expr: str = (
-                f"tl.range(tl.program_id(0), {self.total_pids_var}, {NUM_SM_VAR})"
-            )
+            self.range_kwargs: dict[str, str] = {
+                "begin": "tl.program_id(0)",
+                "end": self.total_pids_var,
+                "step": NUM_SM_VAR,
+            }
         if device_function.constexpr_arg(NUM_SM_VAR):
             device = CompileEnvironment.current().device
             device_function.codegen.host_statements.append(
@@ -402,8 +408,18 @@ def setup_persistent_kernel(
                 )
 
         device_function.preamble.extend(setup_statements)
+        # Collect all block IDs from PID info for range configuration
+        pid_block_ids = []
+        for pid_info in self.pid_info:
+            pid_block_ids.append(pid_info.block_id)
+
+        from .tile_strategy import TileStrategy
+
+        range_expr = TileStrategy.get_range_call_str(
+            device_function.config, pid_block_ids, **self.range_kwargs
+        )
         return self._setup_persistent_kernel_and_wrap_body(
-            device_function, self.virtual_pid_var, self.range_expr, total_pids_expr
+            device_function, self.virtual_pid_var, range_expr, total_pids_expr
         )
 
     def _is_persistent(self) -> bool:

helion/_compiler/reduction_strategy.py

@@ -258,7 +258,7 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
             target=create(ast.Name, id=offset_var, ctx=ast.Store()),
             iter=expr_from_string(
                 self.get_range_call_str(
-                    state,
+                    state.config,
                     [self.block_index],
                     begin="0",
                     end=state.sympy_expr(numel),

helion/_compiler/tile_strategy.py

@@ -33,6 +33,7 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    from ..runtime.config import Config
     from .device_function import DeviceFunction
     from .inductor_lowering import CodegenState
 
@@ -125,45 +126,46 @@ def mask_var(self, block_idx: int) -> str | None:
     def block_size_var(self, block_idx: int) -> str | None:
         return self.fn.block_size_var_cache.get((block_idx,))
 
-    def get_tl_range_kwargs(self, state: CodegenState, block_idx: int) -> list[str]:
+    @staticmethod
+    def get_tl_range_kwargs(config: Config, block_idx: int) -> list[str]:
         """Get the range_extra string for loop unroll factor and num_stages based on config."""
         env = CompileEnvironment.current()
         kwargs = []
 
         range_unroll_factor = env.config_spec.range_unroll_factors.config_get(
-            state.config.range_unroll_factors, block_idx, 0
+            config.range_unroll_factors, block_idx, 0
         )
         if range_unroll_factor > 0:
             kwargs.append(f"loop_unroll_factor={range_unroll_factor}")
 
         range_warp_specialize = env.config_spec.range_warp_specialize.config_get(
-            state.config.range_warp_specializes, block_idx, None
+            config.range_warp_specializes, block_idx, None
         )
         if range_warp_specialize is not None:
             kwargs.append(f"warp_specialize={range_warp_specialize}")
 
         range_num_stages = env.config_spec.range_num_stages.config_get(
-            state.config.range_num_stages, block_idx, 0
+            config.range_num_stages, block_idx, 0
         )
         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
+            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}")
 
         range_flatten = env.config_spec.range_flattens.config_get(
-            state.config.range_flattens, block_idx, None
+            config.range_flattens, block_idx, None
         )
         if range_flatten is not None:
             kwargs.append(f"flatten={range_flatten}")
         return kwargs
 
+    @staticmethod
     def get_range_call_str(
-        self,
-        state: CodegenState,
+        config: Config,
         block_ids: list[int],
         *,
         begin: str | None = None,
@@ -173,7 +175,7 @@ def get_range_call_str(
         env = CompileEnvironment.current()
         use_static_range = all(
             env.config_spec.static_ranges.config_get(
-                state.config.static_ranges, block_idx, None
+                config.static_ranges, block_idx, None
             )
             is True
             for block_idx in block_ids
@@ -183,12 +185,12 @@ def get_range_call_str(
         if begin is not None:
             range_args.append(begin)
         range_args.append(end)
-        if step is not None:
-            range_args.append(f"step={step}")
+        if step is not None and step != "1":
+            range_args.append(step)
 
         if use_static_range:
             return f"tl.static_range({', '.join(range_args)})"
-        range_kwargs = self.get_tl_range_kwargs(state, block_ids[0])
+        range_kwargs = TileStrategy.get_tl_range_kwargs(config, block_ids[0])
         return f"tl.range({', '.join(range_args + range_kwargs)})"
 
     def user_size(self, block_index: int) -> sympy.Expr:
@@ -439,7 +441,7 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
             ast.For,
             target=create(ast.Name, id=lid, ctx=ast.Store()),
             iter=expr_from_string(
-                self.get_range_call_str(state, self.block_ids, end=end_var)
+                self.get_range_call_str(state.config, self.block_ids, end=end_var)
             ),
             body=(
                 body := [
@@ -577,7 +579,7 @@ def codegen_grid(self, state: CodegenState) -> DeviceGridState:
             )
             if mask_statement is not None:
                 state.add_statement(mask_statement)
-            pid = PIDInfo(pid_var, block_size_var, numel)
+            pid = PIDInfo(pid_var, block_size_var, numel, block_idx)
             pids.append(pid)
         pids.codegen(state)
         if isinstance(state.device_function.pid, ForEachProgramID):
@@ -656,7 +658,7 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
                 target=create(ast.Name, id=offset_var, ctx=ast.Store()),
                 iter=expr_from_string(
                     self.get_range_call_str(
-                        state,
+                        state.config,
                         [block_idx],
                         begin="begin",
                         end="end",

helion/autotuner/config_fragment.py

@@ -126,10 +126,10 @@ def random(self) -> object:
     def differential_mutation(self, a: object, b: object, c: object) -> object:
         if b == c:
             return a
-        for candidate in random.sample(self.choices, 2):
-            if candidate != a:
-                return candidate
-        return self.random()  # only reachable with duplicate choices
+        choices = [b, c]
+        if a in choices:
+            choices.remove(a)
+        return random.choice(choices)
 
 
 class BooleanFragment(ConfigSpecFragment):

helion/autotuner/config_spec.py

@@ -4,6 +4,7 @@
 import functools
 import operator
 from typing import TYPE_CHECKING
+from typing import cast
 
 from torch._inductor.runtime.runtime_utils import next_power_of_2
 
@@ -95,6 +96,7 @@ class ConfigSpec:
     allowed_pid_types: tuple[PidTypeLiteral, ...] = dataclasses.field(
         default_factory=functools.partial(tuple, VALID_PID_TYPES)
     )
+    grid_block_ids: list[int] = dataclasses.field(default_factory=list)
 
     @staticmethod
     def _valid_indexing_types() -> tuple[IndexingLiteral, ...]:
@@ -165,27 +167,33 @@ def normalize(self, config: helion.Config | dict[str, object]) -> None:
                 name, config.get(name, ()), flatten=flatten
             )
 
-        static_range_block_ids = []
-        for block_id in self.static_ranges.valid_block_ids():
-            use_static_range = self.static_ranges.config_get(
-                config.get(  # pyright: ignore[reportArgumentType]
-                    "static_ranges", ()
-                ),
+        # Disable range_* configs for static ranges
+        static_range_block_ids = [
+            block_id
+            for block_id in self.static_ranges.valid_block_ids()
+            if self.static_ranges.config_get(
+                cast("list[bool]", config.get("static_ranges", [])),
                 block_id,
             )
-            if use_static_range:
-                static_range_block_ids.append(block_id)
-
-        for name, mapping in (
-            ("range_unroll_factors", self.range_unroll_factors),
-            ("range_warp_specializes", self.range_warp_specialize),
-            ("range_num_stages", self.range_num_stages),
-            ("range_multi_buffers", self.range_multi_buffers),
-            ("range_flattens", self.range_flattens),
-        ):
-            config[name] = mapping._reset_config_to_default(
-                name, config.get(name, ()), block_ids=static_range_block_ids
-            )
+        ]
+        if static_range_block_ids:
+            for name, mapping in (
+                ("range_unroll_factors", self.range_unroll_factors),
+                ("range_warp_specializes", self.range_warp_specialize),
+                ("range_num_stages", self.range_num_stages),
+                ("range_multi_buffers", self.range_multi_buffers),
+                ("range_flattens", self.range_flattens),
+            ):
+                config[name] = mapping._reset_config_to_default(
+                    name, config.get(name, ()), block_ids=static_range_block_ids
+                )
+
+        # Only one range_warp_specializes is allowed, take the last one
+        range_warp_specializes = cast(
+            "list[bool | None]", config.get("range_warp_specializes", [])
+        )
+        for i in [j for j, val in enumerate(range_warp_specializes) if val][:-1]:
+            range_warp_specializes[i] = None
 
         for name in (
             "loop_orders",
@@ -218,6 +226,20 @@ def normalize(self, config: helion.Config | dict[str, object]) -> None:
             else:
                 config[name] = values[0]
 
+        # Set default values for grid indices when pid_type is not persistent
+        pid_type = config["pid_type"]
+        if pid_type in ("flat", "xyz") and self.grid_block_ids:
+            for name, mapping in (
+                ("range_unroll_factors", self.range_unroll_factors),
+                ("range_warp_specializes", self.range_warp_specialize),
+                ("range_num_stages", self.range_num_stages),
+                ("range_multi_buffers", self.range_multi_buffers),
+                ("range_flattens", self.range_flattens),
+            ):
+                config[name] = mapping._reset_config_to_default(
+                    name, config.get(name, ()), block_ids=self.grid_block_ids
+                )
+
         # Allow tunable parameter keys in addition to VALID_KEYS
         allowed_keys = VALID_KEYS | {*self.user_defined_tunables.keys()}
         if invalid_keys := ({*config} - allowed_keys):
@@ -263,6 +285,7 @@ def flat_config(self, fn: Callable[[ConfigSpecFragment], object]) -> helion.Conf
         ):
             if not config[name]:
                 config.pop(name)
+        self.normalize(config)
         return helion.Config(**config)