From 180c4543bf025a5d7d58616ed09ca4940a79ba4c Mon Sep 17 00:00:00 2001 From: Shangdi Yu Date: Fri, 24 Jul 2026 18:22:45 -0700 Subject: [PATCH] Cap persistent launch grids at logical PID count --- docs/api/config.md | 8 +++ helion/_compiler/program_id.py | 59 +++++++++++++++++++-- helion/runtime/config.py | 7 +-- test/test_persistent_kernels.py | 93 ++++++++++++++++++++------------- 4 files changed, 125 insertions(+), 42 deletions(-) diff --git a/docs/api/config.md b/docs/api/config.md index b5cee8b49..e286dfdc1 100644 --- a/docs/api/config.md +++ b/docs/api/config.md @@ -103,6 +103,14 @@ Configs are typically discovered automatically through autotuning, but can also - ``"persistent_blocked"``: Persistent kernels with blocked work distribution - ``"persistent_interleaved"``: Persistent kernels with interleaved distribution +.. autoattribute:: Config.num_sm_multiplier + + Controls persistent multi-occupancy. The physical launch grid is capped at + ``min(total logical program IDs, num SMs * num_sm_multiplier)`` so the + autotuner can select an overprovisioned multiplier without launching workers + that cannot own work. Device-dependent bounds retain the SM-derived grid + because their logical program count is unavailable to the host launcher. + .. autoattribute:: Config.l2_groupings Controls reordering of program IDs to improve L2 cache locality. diff --git a/helion/_compiler/program_id.py b/helion/_compiler/program_id.py index b33280ae9..42b59a70f 100644 --- a/helion/_compiler/program_id.py +++ b/helion/_compiler/program_id.py @@ -240,12 +240,26 @@ def codegen_grid(self) -> ast.AST: """Generate grid launch expression for kernel execution.""" raise NotImplementedError + def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST: + """Generate a grid for an explicit aggregate logical PID count.""" + return self.codegen_grid() + + def codegen_uncapped_grid(self) -> ast.AST: + """Generate a grid when the aggregate logical PID count is unknown.""" + return self.codegen_grid() + def total_pids_expr(self, *, is_device: bool) -> str: """Get total PIDs expression for device or host.""" return " * ".join( f"({pid.num_pids_expr(is_device=is_device)})" for pid in self.pid_info ) + def host_total_pids_expr(self) -> str | None: + """Return the host-computable PID count, or None for device bounds.""" + if any(isinstance(pid.numel, str) for pid in self.pid_info): + return None + return self.total_pids_expr(is_device=False) + def setup_persistent_kernel( self, device_function: DeviceFunction, total_pids_expr: str | None = None ) -> list[ast.stmt] | None: @@ -393,6 +407,12 @@ def total_pids_expr(self, *, is_device: bool) -> str: cdivs = [pid.total_pids_expr(is_device=is_device) for pid in self.cases] return " + ".join(cdivs) + def host_total_pids_expr(self) -> str | None: + totals = [case.host_total_pids_expr() for case in self.cases] + if any(total is None for total in totals): + return None + return " + ".join(total for total in totals if total is not None) + def codegen(self, state: CodegenState) -> None: blocks = self._get_cdiv_blocks(state, exclude_last=True) if blocks: @@ -408,8 +428,10 @@ def codegen(self, state: CodegenState) -> None: def codegen_grid(self) -> ast.AST: # Check if any of the pids is a persistent strategy if self.cases[0]._is_persistent(): - # Use SM count grid for persistent kernels - return self.cases[0].codegen_grid() + total_pids_expr = self.host_total_pids_expr() + if total_pids_expr is not None: + return self.cases[0].codegen_grid_for_total(total_pids_expr) + return self.cases[0].codegen_uncapped_grid() # When persistent kernels are not active, use the full grid size host_cdivs = [pid.total_pids_expr(is_device=False) for pid in self.cases] @@ -838,6 +860,14 @@ def codegen_grid(self) -> ast.AST: assert self.parent_strategy is not None return self.parent_strategy.codegen_grid() + def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST: + assert self.parent_strategy is not None + return self.parent_strategy.codegen_grid_for_total(total_pids_expr) + + def codegen_uncapped_grid(self) -> ast.AST: + assert self.parent_strategy is not None + return self.parent_strategy.codegen_uncapped_grid() + def setup_persistent_kernel( self, device_function: DeviceFunction, total_pids_expr: str | None = None ) -> list[ast.stmt] | None: @@ -857,6 +887,10 @@ def total_pids_expr(self, *, is_device: bool) -> str: assert self.parent_strategy is not None return self.parent_strategy.total_pids_expr(is_device=is_device) + def host_total_pids_expr(self) -> str | None: + assert self.parent_strategy is not None + return self.parent_strategy.host_total_pids_expr() + class PersistentProgramIDs(ProgramIDs): """Base class for persistent kernels that use num_sms grid size.""" @@ -911,7 +945,18 @@ def get_device_str(self) -> str: return f"torch.{device!r}" def codegen_grid(self) -> ast.AST: - # Use num_sms * multiplier for persistent kernels (multi-occupancy) + total_pids_expr = self.host_total_pids_expr() + if total_pids_expr is None: + return expr_from_string(f"({self.grid_size_expr},)") + return self.codegen_grid_for_total(total_pids_expr) + + def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST: + # Avoid launching workers that cannot own a logical PID. Keep the + # persistent stride unchanged so overprovisioned configurations still + # map virtual PIDs exactly as before when total work exceeds the grid. + return expr_from_string(f"(min(({total_pids_expr}), ({self.grid_size_expr})),)") + + def codegen_uncapped_grid(self) -> ast.AST: return expr_from_string(f"({self.grid_size_expr},)") def _persistent_setup_statements(self, total_pids_expr: str) -> list[ast.stmt]: @@ -1359,6 +1404,14 @@ def codegen_grid(self) -> ast.AST: grid_work_clusters = self._tcgen05_grid_work_clusters_expr(total_clusters) return expr_from_string(f"({cluster_m}, {cluster_n}, {grid_work_clusters})") + def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST: + # tcgen05 uses its own cluster-aware scheduler and already caps its + # persistent work-cluster grid where required. + return self.codegen_grid() + + def codegen_uncapped_grid(self) -> ast.AST: + return self.codegen_grid() + def _tcgen05_logical_m_coord_expr(self, coord: str) -> str: if self._tcgen05_is_two_cta(): return f"({coord}) // cutlass.Int32({self._tcgen05_cluster_m()})" diff --git a/helion/runtime/config.py b/helion/runtime/config.py index 414f666d0..e304c7550 100644 --- a/helion/runtime/config.py +++ b/helion/runtime/config.py @@ -19,7 +19,7 @@ EvictionPolicyLiteral = Literal["", "first", "last"] LoadCacheModifierLiteral = Literal["", ".cg"] StoreCacheModifierLiteral = Literal["", ".cs", ".wt"] -NumSmMultiplierLiteral = Literal[1, 2, 4, 8] +NumSmMultiplierLiteral = Literal[1, 2, 4, 8, 16, 32, 64, 128] MaxnregLiteral = Literal[32, 64, 128, 256] | None @@ -83,8 +83,9 @@ def __init__( num_warps: Number of warps per block. num_stages: Number of stages for software pipelining. pid_type: Program ID type strategy ("flat", "xyz", "persistent_blocked", "persistent_interleaved"). - num_sm_multiplier: Multiplier for the number of SMs in persistent kernels (1, 2, 4, 8). - Controls multi-occupancy by launching N * num_sms thread blocks instead of just num_sms. + num_sm_multiplier: Multiplier for the number of SMs in persistent kernels. + Controls multi-occupancy with a grid capped at the total logical + program count: min(total_pids, N * num_sms). maxnreg: Maximum number of registers per thread (None, 32, 64, 128, 256). Lower values allow higher occupancy but may hurt performance. Used with persistent kernels to ensure multi-occupancy can be achieved. diff --git a/test/test_persistent_kernels.py b/test/test_persistent_kernels.py index 18fc4de0d..1d8623f49 100644 --- a/test/test_persistent_kernels.py +++ b/test/test_persistent_kernels.py @@ -683,42 +683,23 @@ def test_kernel(x: torch.Tensor) -> torch.Tensor: import re # Look for _launcher(_kernel_name, (grid_size), ...) pattern - # Use a pattern that handles nested parentheses in grid expressions grid_pattern = r"_launcher\([^,]+,\s*\((.+?)\)\s*," flat_grid_match = re.search(grid_pattern, code_flat) - persistent_blocked_grid_match = re.search(grid_pattern, code_persistent_blocked) - persistent_interleaved_grid_match = re.search( - grid_pattern, code_persistent_interleaved - ) self.assertIsNotNone(flat_grid_match, "Could not find grid size in flat code") - self.assertIsNotNone( - persistent_blocked_grid_match, - "Could not find grid size in persistent blocked code", - ) - self.assertIsNotNone( - persistent_interleaved_grid_match, - "Could not find grid size in persistent interleaved code", - ) flat_grid = flat_grid_match.group(1).rstrip(",") # Remove trailing comma - persistent_blocked_grid = persistent_blocked_grid_match.group(1).rstrip(",") - persistent_interleaved_grid = persistent_interleaved_grid_match.group(1).rstrip( - "," - ) # Flat should use the full grid size calculation (ceiling division) self.assertIn("//", flat_grid) - # Persistent kernels should use NUM_SMS - self.assertEqual( - persistent_blocked_grid, - "_NUM_SM", - ) - self.assertEqual( - persistent_interleaved_grid, - "_NUM_SM", + # Persistent kernels cap the SM-derived grid at the logical PID count. + self.assertIn("_launcher(_helion_test_kernel, (min(", code_persistent_blocked) + self.assertIn("_NUM_SM", code_persistent_blocked) + self.assertIn( + "_launcher(_helion_test_kernel, (min(", code_persistent_interleaved ) + self.assertIn("_NUM_SM", code_persistent_interleaved) def test_persistent_loop_variable_names(self): """Test that persistent kernels use correct virtual_pid variable names.""" @@ -1204,6 +1185,39 @@ def data_dependent_kernel( self.assertIn("total_pids", code) self.assertIn("virtual_pid", code) + @skipIfRefEager("Code pattern checking not applicable in ref eager mode") + def test_mixed_static_and_data_dependent_bounds_use_uncapped_grid(self): + """An unknown aggregate PID count must not cap to the static first loop.""" + + @helion.kernel( + autotune_effort="none", + config=helion.Config(pid_type="persistent_interleaved"), + ) + def mixed_bounds_kernel( + x: torch.Tensor, num_elements: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + static_result = x.new_zeros(x.size()) + dynamic_result = x.new_zeros(x.size()) + for tile in hl.tile(32, block_size=32): + static_result[tile] = x[tile] + 1 + for tile in hl.tile(num_elements, block_size=32): + dynamic_result[tile] = x[tile] + 2 + return static_result, dynamic_result + + x = torch.randn([128], device=DEVICE) + num_elements = torch.tensor(64, device=DEVICE) + code, (static_result, dynamic_result) = code_and_output( + mixed_bounds_kernel, (x, num_elements) + ) + + expected_static = torch.zeros_like(x) + expected_static[:32] = x[:32] + 1 + expected_dynamic = torch.zeros_like(x) + expected_dynamic[:64] = x[:64] + 2 + torch.testing.assert_close(static_result, expected_static) + torch.testing.assert_close(dynamic_result, expected_dynamic) + self.assertIn("_launcher(_helion_mixed_bounds_kernel, (_NUM_SM,)", code) + @onlyBackends(["triton"]) class TestNumSmMultiplier(RefEagerTestBase, TestCase): @@ -1221,21 +1235,24 @@ def test_num_sm_multiplier_blocked_grid_size(self): code_m1, result_m1 = code_and_output( add_kernel, args, pid_type="persistent_blocked", num_sm_multiplier=1 ) - self.assertIn("(_NUM_SM,)", code_m1) + self.assertIn("min(", code_m1) + self.assertIn("_NUM_SM", code_m1) self.assertIn("tl.cdiv(total_pids, _NUM_SM)", code_m1) # Test with multiplier=2 code_m2, result_m2 = code_and_output( add_kernel, args, pid_type="persistent_blocked", num_sm_multiplier=2 ) - self.assertIn("(_NUM_SM * 2,)", code_m2) + self.assertIn("min(", code_m2) + self.assertIn("_NUM_SM * 2", code_m2) self.assertIn("tl.cdiv(total_pids, _NUM_SM * 2)", code_m2) # Test with multiplier=4 code_m4, result_m4 = code_and_output( add_kernel, args, pid_type="persistent_blocked", num_sm_multiplier=4 ) - self.assertIn("(_NUM_SM * 4,)", code_m4) + self.assertIn("min(", code_m4) + self.assertIn("_NUM_SM * 4", code_m4) self.assertIn("tl.cdiv(total_pids, _NUM_SM * 4)", code_m4) # All should produce the same result @@ -1256,28 +1273,32 @@ def test_num_sm_multiplier_interleaved_step(self): code_m1, result_m1 = code_and_output( add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=1 ) - self.assertIn("(_NUM_SM,)", code_m1) + self.assertIn("min(", code_m1) + self.assertIn("_NUM_SM", code_m1) self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM", code_m1) # Test with multiplier=2 code_m2, result_m2 = code_and_output( add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=2 ) - self.assertIn("(_NUM_SM * 2,)", code_m2) + self.assertIn("min(", code_m2) + self.assertIn("_NUM_SM * 2", code_m2) self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM * 2", code_m2) - # Test with multiplier=8 - code_m8, result_m8 = code_and_output( - add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=8 + # Test with multiplier=16. The host grid is capped at total_pids while + # the device loop retains the full persistent stride. + code_m16, result_m16 = code_and_output( + add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=16 ) - self.assertIn("(_NUM_SM * 8,)", code_m8) - self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM * 8", code_m8) + self.assertIn("min(", code_m16) + self.assertIn("_NUM_SM * 16", code_m16) + self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM * 16", code_m16) # All should produce the same result expected = args[0] + args[1] torch.testing.assert_close(result_m1, expected) torch.testing.assert_close(result_m2, expected) - torch.testing.assert_close(result_m8, expected) + torch.testing.assert_close(result_m16, expected) @skipIfRefEager("Code pattern checking not applicable in ref eager mode") def test_num_sm_multiplier_matmul_correctness(self):