Generalize tile_with_offset pass

Author: jansel

helion/_compiler/device_ir.py

@@ -5,7 +5,6 @@
 import contextlib
 import dataclasses
 import functools
-import itertools
 import operator
 import re
 import textwrap
@@ -1216,55 +1215,82 @@ def add_tile_with_offset_metadata(graph_info: GraphInfo) -> None:
     """
     graph = graph_info.graph
     env = CompileEnvironment.current()
-
-    for node in itertools.chain(
-        graph.find_nodes(op="call_function", target=operator.add),
-        graph.find_nodes(op="call_function", target=torch.ops.aten.add.Tensor),
-    ):
-        # Check if this is tile.index + offset pattern
-        # args[0] should be tile_index result, args[1] should be int/SymInt
-        if len(node.args) != 2 and not node.kwargs:
-            continue
-        left_arg, right_arg = node.args
-
-        # Check if left argument is a tile_index call
+    add_targets = (operator.add, torch.ops.aten.add.Tensor)
+    offset_types = (int, torch.SymInt)
+    for node in graph.nodes:
         if (
-            not isinstance(left_arg, torch.fx.Node)
-            or left_arg.op != "call_function"
-            or left_arg.target != hl.tile_index
+            node.op != "call_function"
+            or node.target not in add_targets
+            or node.kwargs
+            or len(node.args) != 2
         ):
             continue
 
-        # Check if right argument is an integer offset
-        # It could be a constant, SymInt node, or another value
-        # We accept int, SymInt, or nodes that represent them
-        offset = None
-        if isinstance(right_arg, (int, torch.SymInt)):
-            offset = right_arg
-        elif isinstance(right_arg, torch.fx.Node):
-            # Check the node's metadata for the value
-            val = right_arg.meta.get("val")
-            if isinstance(val, (int, torch.SymInt)):
-                offset = val
-
-        if offset is None:
-            continue
+        block_id: int | None = None
+        total_offset: int | torch.SymInt = 0
+        valid = True
 
-        # Extract the block_id from the tile_index call
-        tile_arg = left_arg.args[0]
-        block_id = None
-        if isinstance(tile_arg, torch.fx.Node) and isinstance(
-            tile_arg.meta["val"], torch.SymInt
-        ):
-            block_id = env.get_block_id(tile_arg.meta["val"])
+        for arg in node.args:
+            tile_offset_value: int | torch.SymInt | None = None
+            arg_block_id: int | None = None
+
+            if isinstance(arg, torch.fx.Node):
+                meta_tile = arg.meta.get("tile_with_offset")
+                if meta_tile is not None:
+                    arg_block_id = meta_tile.get("block_id")
+                    if arg_block_id is None:
+                        valid = False
+                        break
+                    tile_offset_value = meta_tile.get("offset", 0)
+                elif (
+                    arg.op == "call_function"
+                    and arg.target == hl.tile_index
+                    and arg.args
+                    and isinstance(arg.args[0], torch.fx.Node)
+                ):
+                    tile_val = arg.args[0].meta.get("val")
+                    if isinstance(tile_val, torch.SymInt):
+                        arg_block_id = env.get_block_id(tile_val)
+                        if arg_block_id is None:
+                            valid = False
+                            break
+                        tile_offset_value = 0
+                else:
+                    val = arg.meta.get("val")
+                    if isinstance(val, offset_types):
+                        total_offset = total_offset + val
+                        continue
+
+                if arg_block_id is not None:
+                    if block_id is not None:
+                        valid = False
+                        break
+                    if tile_offset_value is None:
+                        tile_offset_value = 0
+                    block_id = arg_block_id
+                    total_offset = total_offset + tile_offset_value
+                    continue
+
+                val = arg.meta.get("val")
+                if isinstance(val, offset_types):
+                    total_offset = total_offset + val
+                    continue
+
+                valid = False
+                break
+
+            if isinstance(arg, offset_types):
+                total_offset = total_offset + arg
+                continue
+            valid = False
+            break
 
-        if block_id is None:
+        if not valid or block_id is None:
             continue
 
-        # Add metadata to mark this as a tile+offset node
         node.meta["tile_with_offset"] = {
             "block_id": block_id,
-            "offset": offset,
+            "offset": total_offset,
         }
 
 

test/test_autotuner.expected

@@ -2,28 +2,28 @@ This file is automatically generated by assertExpectedJournal calls in test_auto
 Update expected outputs by running tests with the EXPECTTEST_ACCEPT=1 environment variable set.
 
 --- assertExpectedJournal(TestAutotuner.test_config_fragment0)
-helion.Config(block_sizes=[16, 16, 16], indexing='pointer', l2_groupings=[1], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None])
-helion.Config(block_sizes=[32, 128, 64], indexing='block_ptr', l2_groupings=[8], load_eviction_policies=['', ''], loop_orders=[[1, 0]], num_stages=8, num_warps=8, pid_type='persistent_blocked', range_flattens=[None, True], range_multi_buffers=[False, True], range_num_stages=[3, 0], range_unroll_factors=[1, 2], range_warp_specializes=[None, True])
-helion.Config(block_sizes=[16, 16, 16], indexing='tensor_descriptor', l2_groupings=[16], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=7, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 3], range_warp_specializes=[None, False])
-helion.Config(block_sizes=[16, 32, 256], indexing='pointer', l2_groupings=[64], load_eviction_policies=['first', ''], loop_orders=[[1, 0]], num_stages=2, num_warps=16, pid_type='persistent_interleaved', range_flattens=[True, True], range_multi_buffers=[False, None], range_num_stages=[2, 4], range_unroll_factors=[2, 3], range_warp_specializes=[True, None])
-helion.Config(block_sizes=[64, 32, 16], indexing='block_ptr', l2_groupings=[2], load_eviction_policies=['first', 'last'], loop_orders=[[1, 0]], num_stages=2, num_warps=4, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, True], range_num_stages=[0, 4], range_unroll_factors=[0, 1], range_warp_specializes=[None, None])
-helion.Config(block_sizes=[16, 16, 16], indexing='tensor_descriptor', l2_groupings=[32], load_eviction_policies=['last', 'first'], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None, False], range_multi_buffers=[None, None], range_num_stages=[0, 2], range_unroll_factors=[0, 2], range_warp_specializes=[None, False])
-helion.Config(block_sizes=[16, 32, 64], indexing='block_ptr', l2_groupings=[8], load_eviction_policies=['last', 'first'], loop_orders=[[1, 0]], num_stages=5, num_warps=16, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, False], range_num_stages=[0, 3], range_unroll_factors=[0, 3], range_warp_specializes=[None, None])
-helion.Config(block_sizes=[16, 32, 16], indexing='pointer', l2_groupings=[2], load_eviction_policies=['first', 'first'], loop_orders=[[0, 1]], num_stages=8, num_warps=16, pid_type='persistent_interleaved', range_flattens=[False, None], range_multi_buffers=[False, None], range_num_stages=[3, 3], range_unroll_factors=[2, 3], range_warp_specializes=[False, True])
-helion.Config(block_sizes=[256, 16, 16], indexing='pointer', l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=5, num_warps=32, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, False], range_num_stages=[0, 1], range_unroll_factors=[0, 2], range_warp_specializes=[None, True])
-helion.Config(block_sizes=[16, 64, 16], indexing='tensor_descriptor', l2_groupings=[8], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=3, num_warps=32, pid_type='persistent_interleaved', range_flattens=[True, False], range_multi_buffers=[False, None], range_num_stages=[3, 0], range_unroll_factors=[3, 4], range_warp_specializes=[False, True])
+helion.Config(advanced_compiler_configuration=0, block_sizes=[16, 16, 16], indexing='pointer', l2_groupings=[1], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None])
+helion.Config(advanced_compiler_configuration=10, block_sizes=[32, 128, 64], indexing='block_ptr', l2_groupings=[8], load_eviction_policies=['', ''], loop_orders=[[1, 0]], num_stages=8, num_warps=8, pid_type='persistent_blocked', range_flattens=[None, True], range_multi_buffers=[False, True], range_num_stages=[3, 0], range_unroll_factors=[1, 2], range_warp_specializes=[None, True])
+helion.Config(advanced_compiler_configuration=5, block_sizes=[256, 16, 256], indexing='tensor_descriptor', l2_groupings=[8], load_eviction_policies=['', ''], loop_orders=[[1, 0]], num_stages=2, num_warps=32, pid_type='flat', range_flattens=[None, False], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 3], range_warp_specializes=[None, None])
+helion.Config(advanced_compiler_configuration=3, block_sizes=[16, 16, 16], indexing='pointer', l2_groupings=[64], load_eviction_policies=['last', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=2, pid_type='persistent_interleaved', range_flattens=[True, None], range_multi_buffers=[True, True], range_num_stages=[3, 1], range_unroll_factors=[4, 1], range_warp_specializes=[False, True])
+helion.Config(advanced_compiler_configuration=0, block_sizes=[64, 16, 128], indexing='block_ptr', l2_groupings=[4], load_eviction_policies=['', 'last'], loop_orders=[[0, 1]], num_stages=3, num_warps=16, pid_type='flat', range_flattens=[None, False], range_multi_buffers=[None, False], range_num_stages=[0, 2], range_unroll_factors=[0, 0], range_warp_specializes=[None, None])
+helion.Config(advanced_compiler_configuration=9, block_sizes=[16, 16, 16], indexing='pointer', l2_groupings=[4], load_eviction_policies=['first', 'last'], loop_orders=[[1, 0]], num_stages=2, num_warps=8, pid_type='persistent_blocked', range_flattens=[True, False], range_multi_buffers=[True, None], range_num_stages=[0, 0], range_unroll_factors=[4, 1], range_warp_specializes=[True, False])
+helion.Config(advanced_compiler_configuration=5, block_sizes=[16, 16, 16], indexing='block_ptr', l2_groupings=[1], load_eviction_policies=['first', ''], loop_orders=[[1, 0]], num_stages=8, num_warps=8, pid_type='persistent_interleaved', range_flattens=[None, True], range_multi_buffers=[False, False], range_num_stages=[0, 4], range_unroll_factors=[3, 4], range_warp_specializes=[False, False])
+helion.Config(advanced_compiler_configuration=3, block_sizes=[64, 16, 16], indexing='block_ptr', l2_groupings=[32], load_eviction_policies=['', ''], loop_orders=[[1, 0]], num_stages=5, num_warps=4, pid_type='persistent_interleaved', range_flattens=[None, True], range_multi_buffers=[True, False], range_num_stages=[1, 4], range_unroll_factors=[3, 3], range_warp_specializes=[None, False])
+helion.Config(advanced_compiler_configuration=7, block_sizes=[32, 64, 64], indexing='tensor_descriptor', l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=6, num_warps=2, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 2], range_warp_specializes=[None, None])
+helion.Config(advanced_compiler_configuration=2, block_sizes=[16, 16, 16], indexing='block_ptr', l2_groupings=[1], load_eviction_policies=['', 'last'], loop_orders=[[1, 0]], num_stages=4, num_warps=1, pid_type='persistent_interleaved', range_flattens=[None, False], range_multi_buffers=[True, True], range_num_stages=[1, 4], range_unroll_factors=[3, 1], range_warp_specializes=[True, False])
 
 --- assertExpectedJournal(TestAutotuner.test_config_fragment1)
-helion.Config(block_sizes=[8, 16, 16], flatten_loops=[False], indexing='pointer', l2_groupings=[1], load_eviction_policies=['', ''], loop_orders=[[0, 1, 2]], num_stages=2, num_warps=4, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None])
-helion.Config(block_sizes=[1, 64, 64], flatten_loops=[False], indexing='tensor_descriptor', l2_groupings=[4], load_eviction_policies=['first', 'first'], loop_orders=[[1, 2, 0]], num_stages=4, num_warps=8, pid_type='persistent_blocked', range_flattens=[None], range_multi_buffers=[False], range_unroll_factors=[1], range_warp_specializes=[True])
-helion.Config(block_sizes=[2, 8, 512], flatten_loops=[True], indexing='tensor_descriptor', l2_groupings=[8], load_eviction_policies=['first', 'first'], loop_orders=[[2, 0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None])
-helion.Config(block_sizes=[1, 512, 1], flatten_loops=[True], indexing='tensor_descriptor', l2_groupings=[1], load_eviction_policies=['', 'last'], loop_orders=[[0, 2, 1]], num_stages=5, num_warps=2, pid_type='persistent_blocked', range_flattens=[True], range_multi_buffers=[False], range_unroll_factors=[2], range_warp_specializes=[True])
-helion.Config(block_sizes=[1, 4, 256], flatten_loops=[True], indexing='block_ptr', l2_groupings=[8], load_eviction_policies=['last', 'last'], loop_orders=[[1, 0, 2]], num_stages=2, num_warps=32, pid_type='persistent_interleaved', range_flattens=[None], range_multi_buffers=[True], range_unroll_factors=[1], range_warp_specializes=[True])
-helion.Config(block_sizes=[1, 128, 16], flatten_loops=[True], indexing='tensor_descriptor', l2_groupings=[16], load_eviction_policies=['first', 'first'], loop_orders=[[0, 1, 2]], num_stages=1, num_warps=1, pid_type='persistent_blocked', range_flattens=[None], range_multi_buffers=[False], range_unroll_factors=[4], range_warp_specializes=[None])
-helion.Config(block_sizes=[8, 32, 256], flatten_loops=[False], indexing='pointer', l2_groupings=[64], load_eviction_policies=['first', 'last'], loop_orders=[[0, 1, 2]], num_stages=2, num_warps=8, pid_type='persistent_blocked', range_flattens=[False], range_multi_buffers=[True], range_unroll_factors=[4], range_warp_specializes=[None])
-helion.Config(block_sizes=[2, 64, 32], flatten_loops=[False], indexing='block_ptr', l2_groupings=[8], load_eviction_policies=['last', 'first'], loop_orders=[[1, 2, 0]], num_stages=5, num_warps=16, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None])
-helion.Config(block_sizes=[4, 32, 1], flatten_loops=[True], indexing='pointer', l2_groupings=[8], load_eviction_policies=['', 'last'], loop_orders=[[2, 1, 0]], num_stages=8, num_warps=8, pid_type='persistent_blocked', range_flattens=[True], range_multi_buffers=[False], range_unroll_factors=[3], range_warp_specializes=[True])
-helion.Config(block_sizes=[4, 2, 128], flatten_loops=[False], indexing='tensor_descriptor', l2_groupings=[2], load_eviction_policies=['', 'first'], loop_orders=[[1, 2, 0]], num_stages=2, num_warps=4, pid_type='persistent_blocked', range_flattens=[False], range_multi_buffers=[None], range_unroll_factors=[1], range_warp_specializes=[False])
+helion.Config(advanced_compiler_configuration=0, block_sizes=[8, 16, 16], flatten_loops=[False], indexing='pointer', l2_groupings=[1], load_eviction_policies=['', ''], loop_orders=[[0, 1, 2]], num_stages=2, num_warps=4, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None])
+helion.Config(advanced_compiler_configuration=6, block_sizes=[1, 32, 32], flatten_loops=[False], indexing='tensor_descriptor', l2_groupings=[4], load_eviction_policies=['first', 'first'], loop_orders=[[1, 2, 0]], num_stages=4, num_warps=8, pid_type='persistent_blocked', range_flattens=[None], range_multi_buffers=[False], range_unroll_factors=[1], range_warp_specializes=[True])
+helion.Config(advanced_compiler_configuration=10, block_sizes=[1, 32, 1], flatten_loops=[True], indexing='block_ptr', l2_groupings=[16], load_eviction_policies=['last', ''], loop_orders=[[2, 1, 0]], num_stages=4, num_warps=1, pid_type='persistent_blocked', range_flattens=[None], range_multi_buffers=[None], range_unroll_factors=[3], range_warp_specializes=[False])
+helion.Config(advanced_compiler_configuration=2, block_sizes=[1, 8, 16], flatten_loops=[True], indexing='pointer', l2_groupings=[32], load_eviction_policies=['last', 'last'], loop_orders=[[2, 1, 0]], num_stages=7, num_warps=16, pid_type='persistent_interleaved', range_flattens=[False], range_multi_buffers=[None], range_unroll_factors=[3], range_warp_specializes=[True])
+helion.Config(advanced_compiler_configuration=10, block_sizes=[1, 1, 64], flatten_loops=[True], indexing='tensor_descriptor', l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[2, 0, 1]], num_stages=6, num_warps=1, pid_type='persistent_interleaved', range_flattens=[True], range_multi_buffers=[None], range_unroll_factors=[4], range_warp_specializes=[True])
+helion.Config(advanced_compiler_configuration=6, block_sizes=[4, 2, 128], flatten_loops=[True], indexing='block_ptr', l2_groupings=[1], load_eviction_policies=['first', 'first'], loop_orders=[[0, 1, 2]], num_stages=6, num_warps=1, pid_type='persistent_blocked', range_flattens=[True], range_multi_buffers=[True], range_unroll_factors=[0], range_warp_specializes=[True])
+helion.Config(advanced_compiler_configuration=3, block_sizes=[2, 16, 2], flatten_loops=[True], indexing='block_ptr', l2_groupings=[64], load_eviction_policies=['first', 'first'], loop_orders=[[0, 2, 1]], num_stages=4, num_warps=16, pid_type='persistent_blocked', range_flattens=[False], range_multi_buffers=[False], range_unroll_factors=[0], range_warp_specializes=[None])
+helion.Config(advanced_compiler_configuration=5, block_sizes=[4, 128, 16], flatten_loops=[True], indexing='pointer', l2_groupings=[64], load_eviction_policies=['first', ''], loop_orders=[[1, 0, 2]], num_stages=6, num_warps=4, pid_type='persistent_interleaved', range_flattens=[False], range_multi_buffers=[False], range_unroll_factors=[0], range_warp_specializes=[True])
+helion.Config(advanced_compiler_configuration=4, block_sizes=[4, 256, 32], flatten_loops=[False], indexing='block_ptr', l2_groupings=[8], load_eviction_policies=['last', ''], loop_orders=[[2, 1, 0]], num_stages=5, num_warps=2, pid_type='persistent_blocked', range_flattens=[False], range_multi_buffers=[True], range_unroll_factors=[1], range_warp_specializes=[True])
+helion.Config(advanced_compiler_configuration=9, block_sizes=[2, 128, 8], flatten_loops=[True], indexing='pointer', l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[1, 2, 0]], num_stages=5, num_warps=32, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None])
 
 --- assertExpectedJournal(TestAutotuner.test_save_load_config)
 {