Generalize l2_grouping to support 3+ dimensions (#313)

Author: jansel

helion/_compiler/program_id.py

@@ -266,7 +266,7 @@ def codegen(self, state: CodegenState) -> None:
         # Note: Persistent kernel setup is handled by ForEachProgramID if needed
         assert self.parent_strategy is not None
         parent_pids = self.parent_strategy.pid_info
-        assert len(parent_pids) == 2
+        assert len(parent_pids) >= 2, "L2 grouping requires at least 2 dimensions"
         new_var = state.device_function.new_var
 
         # Use shared_pid_var if we're in a ForEachProgramID context, otherwise use virtual_program_id
@@ -275,26 +275,87 @@ def codegen(self, state: CodegenState) -> None:
         else:
             pid = self.virtual_program_id
 
-        num_pid_m = new_var("num_pid_m")
-        num_pid_n = new_var("num_pid_n")
-        num_pid_in_group = new_var("num_pid_in_group")
-        group_id = new_var("group_id")
-        first_pid_m = new_var("first_pid_m")
-        group_size_m = new_var("group_size_m")
-
-        assignments = [
-            (num_pid_m, parent_pids[0].num_pids_expr(is_device=True)),
-            (num_pid_n, parent_pids[1].num_pids_expr(is_device=True)),
-            (num_pid_in_group, f"{self.group_size} * {num_pid_n}"),
-            (group_id, f"{pid} // {num_pid_in_group}"),
-            (first_pid_m, f"{group_id} * {self.group_size}"),
-            (group_size_m, f"min({num_pid_m} - {first_pid_m}, {self.group_size})"),
-            (
-                parent_pids[0].pid_var,
-                f"{first_pid_m} + (({pid} % {num_pid_in_group}) % {group_size_m})",
-            ),
-            (parent_pids[1].pid_var, f"({pid} % {num_pid_in_group}) // {group_size_m}"),
+        # Apply L2 grouping to the 2 fastest varying dimensions (pid_0, pid_1)
+        # These are always the first 2 dimensions in the PID decomposition
+        num_dims = len(parent_pids)
+        assignments = []
+
+        # Generate size variables for all dimensions (except the last which doesn't need one)
+        num_blocks = []
+        for i in range(num_dims - 1):
+            num_block_var = new_var(f"num_blocks_{i}", dce=True)
+            assignments.append(
+                (num_block_var, parent_pids[i].num_pids_expr(is_device=True))
+            )
+            num_blocks.append(num_block_var)
+
+        # Apply L2 grouping to the 2 fastest varying dimensions (pid_0, pid_1)
+        fastest_m_idx = 0  # pid_0 (fastest varying)
+        fastest_n_idx = 1  # pid_1 (second fastest varying)
+
+        # Extract the 2D portion for the fastest 2 dimensions
+        inner_2d_size = new_var("inner_2d_size", dce=True)
+        inner_2d_pid = new_var("inner_2d_pid", dce=True)
+
+        num_pid_m = new_var("num_pid_m", dce=True)
+        num_pid_n = new_var("num_pid_n", dce=True)
+        num_pid_in_group = new_var("num_pid_in_group", dce=True)
+        group_id = new_var("group_id", dce=True)
+        first_pid_m = new_var("first_pid_m", dce=True)
+        group_size_m = new_var("group_size_m", dce=True)
+
+        # Set up L2 grouping for the fastest 2 dimensions
+        inner_2d_assignments = [
+            (num_pid_m, parent_pids[fastest_m_idx].num_pids_expr(is_device=True)),
+            (num_pid_n, parent_pids[fastest_n_idx].num_pids_expr(is_device=True)),
         ]
+
+        # Only add modulo for 3D+ cases where we need to extract the 2D portion
+        if num_dims > 2:
+            inner_2d_assignments.extend(
+                [
+                    (inner_2d_size, f"{num_pid_m} * {num_pid_n}"),
+                    (
+                        inner_2d_pid,
+                        f"{pid} % {inner_2d_size}",
+                    ),  # Extract fastest 2D portion
+                ]
+            )
+        else:
+            # For 2D case, the entire PID space is the 2D space
+            inner_2d_assignments.append((inner_2d_pid, pid))
+
+        assignments.extend(inner_2d_assignments)
+        assignments.extend(
+            [
+                (num_pid_in_group, f"{self.group_size} * {num_pid_n}"),
+                (group_id, f"{inner_2d_pid} // {num_pid_in_group}"),
+                (first_pid_m, f"{group_id} * {self.group_size}"),
+                (group_size_m, f"min({num_pid_m} - {first_pid_m}, {self.group_size})"),
+                (
+                    parent_pids[fastest_m_idx].pid_var,
+                    f"{first_pid_m} + (({inner_2d_pid} % {num_pid_in_group}) % {group_size_m})",
+                ),
+                (
+                    parent_pids[fastest_n_idx].pid_var,
+                    f"({inner_2d_pid} % {num_pid_in_group}) // {group_size_m}",
+                ),
+            ]
+        )
+
+        # Process remaining dimensions (if any) using standard decomposition
+        for i in range(2, num_dims):
+            expr = pid
+            # Add divisor for all faster dimensions
+            if i > 0:
+                divisor = " * ".join(num_blocks[:i])
+                expr = f"({expr}) // ({divisor})"
+            # Add modulo unless this is the outermost dimension
+            if i + 1 < num_dims:  # Not the outermost dimension
+                expr = f"({expr}) % {num_blocks[i]}"
+
+            assignments.append((parent_pids[i].pid_var, expr))
+
         statements = [
             statement_from_string(f"{var} = {expr}") for var, expr in assignments
         ]

helion/language/loops.py

@@ -386,8 +386,8 @@ def _add_config_choices(
         config_spec.grid_block_ids.extend(
             [x for x in block_ids if x not in existing_ids]
         )
-        if len(block_ids) == 2:
-            # TODO(jansel): support L2 grouping with 3+ dims (and maybe non-grids?)
+        if len(block_ids) >= 2:
+            # L2 grouping now supports 3D+ grids by applying to innermost 2 dimensions
             config_spec.l2_groupings.append(L2GroupingSpec(block_ids))
         if not _allow_use_yz_grid(config_spec, block_ids):
             config_spec.disallow_pid_type("xyz")

test/test_autotuner.expected

@@ -14,16 +14,16 @@ helion.Config(block_sizes=[16, 32, 16], loop_orders=[[1, 0]], l2_groupings=[8],
 helion.Config(block_sizes=[16, 32, 32], loop_orders=[[0, 1]], l2_groupings=[2], range_unroll_factors=[0, 2], range_warp_specializes=[None, None], range_num_stages=[0, 0], range_multi_buffers=[None, None], range_flattens=[None, True], num_warps=2, num_stages=6, indexing='tensor_descriptor', pid_type='flat')
 
 --- assertExpectedJournal(TestAutotuner.test_config_fragment1)
-helion.Config(block_sizes=[8, 16, 16], loop_orders=[[0, 1, 2]], flatten_loops=[False], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=4, num_stages=3, indexing='pointer', pid_type='flat')
-helion.Config(block_sizes=[2, 128, 128], loop_orders=[[1, 2, 0]], flatten_loops=[False], range_unroll_factors=[2], range_warp_specializes=[None], range_num_stages=[4], range_multi_buffers=[False], range_flattens=[None], num_warps=8, num_stages=4, indexing='tensor_descriptor', pid_type='persistent_blocked')
-helion.Config(block_sizes=[2, 16, 4], loop_orders=[[0, 2, 1]], flatten_loops=[True], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=1, num_stages=2, indexing='tensor_descriptor', pid_type='flat')
-helion.Config(block_sizes=[8, 2, 512], loop_orders=[[0, 2, 1]], flatten_loops=[True], range_unroll_factors=[4], range_warp_specializes=[False], range_num_stages=[0], range_multi_buffers=[False], range_flattens=[True], num_warps=8, num_stages=3, indexing='block_ptr', pid_type='persistent_interleaved')
-helion.Config(block_sizes=[1, 16, 32], loop_orders=[[0, 2, 1]], flatten_loops=[False], range_unroll_factors=[4], range_warp_specializes=[True], range_num_stages=[4], range_multi_buffers=[None], range_flattens=[False], num_warps=8, num_stages=4, indexing='tensor_descriptor', pid_type='persistent_interleaved')
-helion.Config(block_sizes=[1, 32, 512], loop_orders=[[0, 2, 1]], flatten_loops=[False], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=2, num_stages=5, indexing='pointer', pid_type='flat')
-helion.Config(block_sizes=[1, 32, 32], loop_orders=[[1, 2, 0]], flatten_loops=[True], range_unroll_factors=[4], range_warp_specializes=[None], range_num_stages=[2], range_multi_buffers=[None], range_flattens=[None], num_warps=16, num_stages=3, indexing='tensor_descriptor', pid_type='persistent_blocked')
-helion.Config(block_sizes=[1, 4, 32], loop_orders=[[1, 0, 2]], flatten_loops=[True], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=16, num_stages=6, indexing='block_ptr', pid_type='flat')
-helion.Config(block_sizes=[4, 16, 1], loop_orders=[[2, 1, 0]], flatten_loops=[True], range_unroll_factors=[2], range_warp_specializes=[None], range_num_stages=[2], range_multi_buffers=[None], range_flattens=[True], num_warps=4, num_stages=8, indexing='block_ptr', pid_type='persistent_interleaved')
-helion.Config(block_sizes=[8, 128, 4], loop_orders=[[1, 0, 2]], flatten_loops=[False], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=2, num_stages=4, indexing='tensor_descriptor', pid_type='flat')
+helion.Config(block_sizes=[8, 16, 16], loop_orders=[[0, 1, 2]], flatten_loops=[False], l2_groupings=[1], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=4, num_stages=3, indexing='pointer', pid_type='flat')
+helion.Config(block_sizes=[2, 128, 128], loop_orders=[[1, 2, 0]], flatten_loops=[False], l2_groupings=[4], range_unroll_factors=[1], range_warp_specializes=[True], range_num_stages=[3], range_multi_buffers=[None], range_flattens=[False], num_warps=16, num_stages=4, indexing='tensor_descriptor', pid_type='persistent_blocked')
+helion.Config(block_sizes=[4, 32, 8], loop_orders=[[0, 2, 1]], flatten_loops=[True], l2_groupings=[8], range_unroll_factors=[4], range_warp_specializes=[False], range_num_stages=[2], range_multi_buffers=[None], range_flattens=[None], num_warps=1, num_stages=4, indexing='block_ptr', pid_type='persistent_blocked')
+helion.Config(block_sizes=[1, 512, 1], loop_orders=[[0, 2, 1]], flatten_loops=[True], l2_groupings=[1], range_unroll_factors=[2], range_warp_specializes=[True], range_num_stages=[3], range_multi_buffers=[True], range_flattens=[None], num_warps=4, num_stages=7, indexing='pointer', pid_type='persistent_interleaved')
+helion.Config(block_sizes=[1, 8, 512], loop_orders=[[1, 0, 2]], flatten_loops=[True], l2_groupings=[8], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=4, num_stages=2, indexing='tensor_descriptor', pid_type='flat')
+helion.Config(block_sizes=[4, 2, 128], loop_orders=[[0, 1, 2]], flatten_loops=[True], l2_groupings=[1], range_unroll_factors=[0], range_warp_specializes=[True], range_num_stages=[0], range_multi_buffers=[True], range_flattens=[False], num_warps=4, num_stages=5, indexing='block_ptr', pid_type='persistent_blocked')
+helion.Config(block_sizes=[2, 16, 2], loop_orders=[[0, 2, 1]], flatten_loops=[True], l2_groupings=[64], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[2], range_multi_buffers=[False], range_flattens=[True], num_warps=16, num_stages=4, indexing='block_ptr', pid_type='persistent_blocked')
+helion.Config(block_sizes=[4, 4, 1], loop_orders=[[1, 2, 0]], flatten_loops=[False], l2_groupings=[16], range_unroll_factors=[2], range_warp_specializes=[False], range_num_stages=[0], range_multi_buffers=[True], range_flattens=[False], num_warps=8, num_stages=5, indexing='tensor_descriptor', pid_type='persistent_blocked')
+helion.Config(block_sizes=[4, 4, 16], loop_orders=[[1, 2, 0]], flatten_loops=[True], l2_groupings=[8], range_unroll_factors=[1], range_warp_specializes=[False], range_num_stages=[1], range_multi_buffers=[True], range_flattens=[True], num_warps=8, num_stages=3, indexing='tensor_descriptor', pid_type='persistent_blocked')
+helion.Config(block_sizes=[4, 8, 8], loop_orders=[[2, 0, 1]], flatten_loops=[False], l2_groupings=[4], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=8, num_stages=5, indexing='tensor_descriptor', pid_type='flat')
 
 --- assertExpectedJournal(TestAutotuner.test_save_load_config)
 {