[Inductor] Represent tiling as a dict (pytorch#141751)

# Summary

Preparatory refactor for pytorch#137243. This makes it easier to generalize to multi-dimensional reductions.

This diff refactors `self.numels` from a tuple like `(8,16)` to a dict like `{"x": 8, "r": 16}`.

Note: this is based off of pytorch#141738, which enables `tree.is_reduction`. That PR should land first.

# Test plan
The existing CI provides good coverage.

Pull Request resolved: pytorch#141751
Approved by: https://github.com/jansel

Author: blaine-rister

torch/_inductor/codegen/halide.py

@@ -673,10 +673,10 @@ class HalideKernel(SIMDKernel):
 
     def __init__(
         self,
-        *groups,
+        tiling: Dict[str, sympy.Expr],
         **kwargs,
     ) -> None:
-        super().__init__(*groups, **kwargs)
+        super().__init__(tiling, **kwargs)
         # For halide, we just write directly to the body
         self.compute = self.body
         self.loads = self.body

torch/_inductor/codegen/simd.py

@@ -27,6 +27,7 @@
 
 import torch
 import torch._logging
+from torch.fx.immutable_collections import immutable_dict
 from torch.utils._ordered_set import OrderedSet
 from torch.utils._sympy.functions import FloorDiv, Identity, ModularIndexing
 from torch.utils._sympy.symbol import (
@@ -339,7 +340,7 @@ class SIMDKernel(Kernel):
 
     def __init__(
         self,
-        *groups,
+        tiling: Dict[str, sympy.Expr],
         features: SIMDKernelFeatures,
         pid_cache=None,
         override_persistent_reduction=None,
@@ -352,11 +353,13 @@ def __init__(
         self.mutations = features.get_mutations()
         self.body = IndentedBuffer()
         self.indexing_code = IndentedBuffer()
-        self.numels = [V.graph.sizevars.simplify(s) for s in groups]
+        self.numels = {
+            prefix: V.graph.sizevars.simplify(val) for prefix, val in tiling.items()
+        }
         self.range_trees: List[IterationRangesRoot] = []
         self.range_tree_nodes: Dict[sympy.Symbol, IterationRangesEntry] = {}
         self.iter_vars_count = itertools.count()
-        self.inside_reduction = self.numels[-1] != 1
+        self.inside_reduction = self.numels["r"] != 1
         self.cooperative_reduction: bool = (
             override_cooperative_reduction
             if override_cooperative_reduction is not None
@@ -393,7 +396,7 @@ def want_no_x_dim(self):
         return False
 
     def initialize_range_tree(self, pid_cache):
-        no_r_dim = not self.inside_reduction or self.numels[-1] == 1
+        no_r_dim = not self.inside_reduction or self.numels["r"] == 1
 
         prefixes = "zyxr"
         active_prefixes = prefixes[-len(self.numels) :]
@@ -416,7 +419,7 @@ def initialize_range_tree(self, pid_cache):
             self.range_trees.append(
                 IterationRangesRoot(
                     f"{prefix}index",
-                    self.numels[i],
+                    self.numels[prefix],
                     prefix,
                     index,
                     self,
@@ -525,7 +528,7 @@ def disable_reduction(self):
 
         @contextlib.contextmanager
         def ctx():
-            if self.numels[-1] == 1:
+            if self.numels["r"] == 1:
                 assert not self.inside_reduction
                 yield
                 return
@@ -688,7 +691,7 @@ def is_broadcasted(self, index: sympy.Expr):
         simplify = V.graph.sizevars.simplify
         return any(
             simplify(idx_range) != simplify(iter_range)  # type: ignore[arg-type]
-            for idx_range, iter_range in zip(index_numels, self.numels)
+            for idx_range, iter_range in zip(index_numels, self.numels.values())
         )
 
     def index_to_str(self, index: sympy.Expr) -> str:
@@ -855,7 +858,7 @@ def estimate_kernel_num_bytes(self):
         # for the "cat". However, I think it might be a bit overwhelming that
         # we add such complexity only for handling some particular cases for
         # benchmarking.
-        out_numel = V.graph.sizevars.size_hint(sympy_product(self.numels))
+        out_numel = V.graph.sizevars.size_hint(sympy_product(self.numels.values()))
         for i, arg in enumerate(call_args):
             # "buf" may be narrowed. In this case, the number of memory accesses
             # should be estimated based on the reinterpreted layout.
@@ -960,7 +963,7 @@ def warn_mix_layout(self, kernel_name):
     def welford_reduce_fallback(self, dtype, value):
         sum_ = ops.reduction(dtype, dtype, "sum", value)
         self.inside_reduction = False
-        rnumel = ops.index_expr(self.numels[-1], dtype)
+        rnumel = ops.index_expr(self.numels["r"], dtype)
         mean = ops.truediv(sum_, rnumel)
 
         self.inside_reduction = True
@@ -1081,8 +1084,8 @@ def can_fuse(self, node1, node2):
                     config.triton.tiling_prevents_reduction_fusion
                     and not node1.is_template()
                 ):
-                    is_reduction_tiling_valid = self.select_tiling(
-                        node1.get_nodes(), numel1
+                    is_reduction_tiling_valid = tuple(
+                        self.select_tiling(node1.get_nodes(), numel1).values()
                     ) in (
                         (numel1, 1),
                         (numel2, rnumel2, 1),
@@ -1246,11 +1249,11 @@ def can_use_32bit_indexing(
 
     def codegen_node_schedule(self, kernel_features: SIMDKernelFeatures):
         node_schedule = kernel_features.node_schedule
-        tiled_groups = self.select_tiling(
+        tiling = self.select_tiling(
             node_schedule, kernel_features.numel, kernel_features.reduction_numel
         )
         kernels = self.create_kernel_choices(
-            kernel_features, tiled_groups, {"features": kernel_features}
+            kernel_features, [tiling], {"features": kernel_features}
         )
         for kernel in kernels:
             self.codegen_node_schedule_with_kernel(node_schedule, kernel)
@@ -1426,10 +1429,10 @@ def generate_combo_kernel_code(
         for pn, nodes in zip(subkernel_nodes, fused_node_lists):
             _, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group
             node_schedule = self.generate_node_schedule(nodes, numel, rnumel)
-            tiled_groups = self.select_tiling(node_schedule, numel, rnumel)
-            node_schedule_map[pn] = node_schedule, tiled_groups, numel, rnumel
+            tiling = self.select_tiling(node_schedule, numel, rnumel)
+            node_schedule_map[pn] = node_schedule, tiling, numel, rnumel
             subkernel_map[pn] = ComboKernel.create_triton_kernel(
-                *tiled_groups,
+                tiling,
                 features=SIMDKernelFeatures(node_schedule, numel, rnumel),
                 optimize_mask=not mixed_sizes,
             )
@@ -1562,7 +1565,23 @@ def candidate_tilings(node):
         return tilings
 
     @classmethod
-    def select_tiling(cls, node_schedule, numel, reduction_numel=sympy.S.One):
+    def create_tiling(
+        cls, pw_tiling: Sequence[sympy.Expr], reduction_tiling: Sequence[sympy.Expr]
+    ) -> Dict[str, sympy.Expr]:
+        """
+        Create a tiling dict from pointwise and reduction splits.
+        """
+        pw_prefixes = ["z", "y", "x"][-len(pw_tiling) :]
+        reduction_prefixes = ["r"][: len(reduction_tiling)]
+        return immutable_dict(
+            list(zip(pw_prefixes, pw_tiling))
+            + list(zip(reduction_prefixes, reduction_tiling))
+        )
+
+    @classmethod
+    def select_tiling(
+        cls, node_schedule, numel, reduction_numel=sympy.S.One
+    ) -> Dict[str, sympy.Expr]:
         """
         Heuristics to decide how to tile kernels.
         Currently, we tile based on stride-1 dimensions.
@@ -1571,6 +1590,7 @@ def select_tiling(cls, node_schedule, numel, reduction_numel=sympy.S.One):
             `(tile1, tile2, reduction_numel)` s.t. `tile1 * tile2 == numel`
 
         """
+        default_tiling = cls.create_tiling([numel], [reduction_numel])
         if reduction_numel != 1 or config.triton.max_tiles <= 1:
             # TODO(jansel): should we tile reductions?
             # do perf hint here if stride-1 dim is not being reduced
@@ -1579,7 +1599,7 @@ def select_tiling(cls, node_schedule, numel, reduction_numel=sympy.S.One):
                     if len(cls.candidate_tilings(node)) > 0:
                         perf_hint_log.info("reduction over non-contiguous dims")
                         break
-            return (numel, reduction_numel)
+            return default_tiling
 
         seen_names: OrderedSet[str] = OrderedSet()
         candidate_tiles: Counter[Any] = collections.Counter()
@@ -1647,9 +1667,9 @@ def select_tiling(cls, node_schedule, numel, reduction_numel=sympy.S.One):
                 for node in node_schedule
                 if isinstance(node, scheduler.SchedulerNode)
             ):
-                return new_groups
+                return cls.create_tiling(tiled_groups, [reduction_numel])
 
-        return (numel, reduction_numel)
+        return default_tiling
 
     def flush(self):
         pass
@@ -1661,9 +1681,9 @@ def generate_kernel_code_from_nodes(self, nodes, benchmark_kernel=False):
         if not nodes[0].is_template():
             _, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group
             node_schedule = self.generate_node_schedule(nodes, numel, rnumel)
-            tiled_groups = self.select_tiling(node_schedule, numel, rnumel)
+            tiling = self.select_tiling(node_schedule, numel, rnumel)
             kernel = self.kernel_type(
-                *tiled_groups,
+                tiling,
                 features=SIMDKernelFeatures(node_schedule, numel, rnumel),
             )
             self.codegen_node_schedule_with_kernel(node_schedule, kernel)

torch/_inductor/codegen/triton.py

@@ -340,7 +340,7 @@ def remove_dims(it):
         if (
             not V.kernel.inside_reduction
             and len(params.strides) == len(V.kernel.numels) - 1
-            and V.kernel.numels[-1] != 1
+            and V.kernel.numels["r"] != 1
         ):
             # Need to expand rank by 1 to match rank when self.inside_reduction=True
             final_shape.append(sympy.S.One)
@@ -1419,15 +1419,15 @@ class TritonKernel(SIMDKernel):
 
     def __init__(
         self,
-        *groups,
+        tiling: Dict[str, sympy.Expr],
         min_elem_per_thread=0,
         optimize_mask=True,
         fixed_config: Optional[FixedTritonConfig] = None,
         **kwargs,
     ) -> None:
         self.optimize_mask: bool = optimize_mask
         self.fixed_config = fixed_config
-        super().__init__(*groups, **kwargs)
+        super().__init__(tiling, **kwargs)
         self.cse = TritonCSE(self.newvar_prefix, self.suffix)
         self.post_loop_combine: IndentedBuffer = IndentedBuffer()
         self.post_loop_store: IndentedBuffer = IndentedBuffer()
@@ -1463,7 +1463,7 @@ def init_cooperative_reduction(self):
             if tree.grid_dim is not None:
                 tree.grid_dim += 1
 
-        sem_count, _ = self.numels
+        sem_count = self.numels["x"]
         if self.fixed_config:
             sem_count = CeilDiv(sem_count, self.fixed_config["XBLOCK"])
         self.semaphores_name = self.args.semaphores(sem_count)
@@ -2440,7 +2440,7 @@ def codegen_cooperative_reduction_peer_combine(self, result_var, dtype):
         column.  After the barrier, every thread block loads the completed value so that it can compute the final
         value independently.
         """
-        xnumel, rnumel = self.numels
+        xnumel = self.numels["x"]
         mask = "xindex < xnumel" if xnumel != 1 and not self.no_x_dim else None
         expand = "" if self.no_x_dim else "[None,:]"
 
@@ -2946,7 +2946,7 @@ def codegen_kernel(self, name=None):
         code = IndentedBuffer()
 
         size_hints = []
-        for numel in self.numels:
+        for numel in self.numels.values():
             numel_hint = V.graph.sizevars.symbolic_hint(numel)
             if not isinstance(numel_hint, (int, sympy.Integer)):
                 # This default heuristic hint was picked carefully: it is
@@ -3384,7 +3384,7 @@ def iteration_ranges_codegen_header(self, entry, code):
 
 
 class TritonScheduling(SIMDScheduling):
-    kernel_type = TritonKernel
+    kernel_type: Type[Any] = TritonKernel
     backend_features = dict.fromkeys(  # dict for deterministic order
         [
             BackendFeature.FOREACH,
@@ -3642,7 +3642,7 @@ def add_multi_kernel_choices(
                 )
             )
         if optional_cooperative:
-            _, rnumel = kernel.numels
+            rnumel = kernel.numels["r"]
             # for larger sizes non-cooperative gets very slow
             if V.graph.sizevars.statically_known_leq(rnumel, 65536):
                 kernels.append(

torch/_inductor/codegen/triton_combo_kernel.py

@@ -16,6 +16,7 @@
     Union,
 )
 
+import sympy
 from sympy import Integer, Symbol
 
 from .. import config, metrics
@@ -102,7 +103,7 @@ def _default_custom_combo_kernel_horizontal_partition(
             for n in not_reduction
             if not kernel_map[n].inside_reduction
             and len(kernel_map[n].numels) == 2
-            and V.graph.sizevars.size_hint(kernel_map[n].numels[0]) > LARGE_NUMELS
+            and V.graph.sizevars.size_hint(kernel_map[n].numels["x"]) > LARGE_NUMELS
         ]
         if large_pointwise:
             # TODO benchmark the performance when large pointwise nodes combining with others
@@ -216,7 +217,7 @@ def _base_horizontal_partition(
             ndim = len(tiled_groups)
             assert ndim >= 2, f"Combokernel not support tile {tiled_groups}"
             if not mixed_sizes and ndim == 3:
-                y_elem = tiled_groups[0]
+                y_elem = tiled_groups["y"]
                 partition_state = yelem_to_partition_state[y_elem]
                 ComboKernel._update_partition(
                     partition_state, read_write_count, node_info
@@ -463,7 +464,7 @@ def create_sub_kernel(self, triton_kernel: TritonKernel) -> TritonKernel:
 
     @staticmethod
     def create_triton_kernel(
-        *groups: Any,
+        tiling: Dict[str, sympy.Expr],
         features: SIMDKernelFeatures,
         optimize_mask: bool,
     ) -> TritonKernel:
@@ -472,7 +473,7 @@ def create_triton_kernel(
         2) numels except x dimension are the same for each sub kernel.
         """
         return TritonKernel(
-            *groups,
+            tiling,
             features=features,
             pid_cache={"tl.program_id(0)": "pid_offset"},
             optimize_mask=optimize_mask,
@@ -564,7 +565,7 @@ def min_x_blocks_sub_kernel(self, sub_kernel: TritonKernel, num: int) -> None:
     def select_heuristics(self, sub_kernel: TritonKernel) -> Tuple[str, List[int]]:
         size_hints = [
             next_power_of_2(V.graph.sizevars.size_hint(numel))
-            for numel in sub_kernel.numels
+            for numel in sub_kernel.numels.values()
         ]
         if sub_kernel.persistent_reduction:
             assert sub_kernel.inside_reduction

torch/_inductor/codegen/triton_split_scan.py

@@ -1,13 +1,18 @@
 # mypy: allow-untyped-defs
 import functools
+from typing import Dict
+
+import sympy
 
 from torch._inductor import config
 from torch._inductor.codegen.simd import IterationRangesRoot
 from torch._inductor.codegen.triton import triton_compute_type, TritonKernel
 from torch._inductor.runtime.triton_heuristics import split_scan_grid
-from torch._prims_common import prod
 from torch.utils._sympy.functions import CeilDiv
 
+from ..utils import sympy_product
+from .simd import prefix_is_reduction
+
 
 class TritonSplitScanKernel(TritonKernel):
     """Generates a triton kernel that supports ops.scan calls while also splitting
@@ -27,15 +32,15 @@ class TritonSplitScanKernel(TritonKernel):
 
     def __init__(
         self,
-        *groups,
+        tiling: Dict[str, sympy.Expr],
         pid_cache=None,
         fixed_config=None,
         **kwargs,
     ) -> None:
         assert pid_cache is None, "not supported"
         assert fixed_config is None, "not supported"
         super().__init__(
-            *groups,
+            tiling,
             **kwargs,
         )
         self.no_x_dim = True
@@ -54,7 +59,8 @@ def initialize_range_tree(self, pid_cache):
         active_prefixes = prefixes[len(prefixes) - len(self.numels) :]
 
         grid_dims = "rxy"
-        for numel, prefix in zip(self.numels, active_prefixes):
+        for prefix in active_prefixes:
+            numel = self.numels[prefix]
             is_reduction = prefix == "r"
             tensor_dim = 0 if is_reduction else None
             grid_dim = grid_dims.find(prefix)
@@ -99,7 +105,17 @@ def scan(self, dtypes, combine_fn, values):
 
         assert len(self.numels) == 2, "Unexpected tiling"
         min_rblock = config.triton.min_split_scan_rblock
-        max_blocks = prod(self.numels[:-1]) * CeilDiv(self.numels[-1], min_rblock)
+        reduction_numel = sympy_product(
+            numel
+            for prefix, numel in self.numels.items()
+            if prefix_is_reduction(prefix)
+        )
+        pointwise_numel = sympy_product(
+            numel
+            for prefix, numel in self.numels.items()
+            if not prefix_is_reduction(prefix)
+        )
+        max_blocks = pointwise_numel * CeilDiv(reduction_numel, min_rblock)
         nbytes = scratch_nbytes_per_block * max_blocks
         scratch_base, offset = self.args.workspace(nbytes=nbytes, zero_fill=True)
         if offset != 0: