[inductor] Fix var_to_range in IndexPropagation (#130984)

peterbell10 · pytorchmergebot · commit 3c622fbcd3f2 · 2024-07-19T03:08:00.000Z
The current code assumes that indirect variables will be created by the same `IndexPropagation` instance, however that isn't true in the case of masked sub-blocks where we take in variables from the parent block. This fixes the issue by moving the var range information up to the `LoopBody` object where it can be shared by all sub-blocks. Pull Request resolved: #130984 Approved by: https://github.com/lezcano
diff --git a/torch/_inductor/index_propagation.py b/torch/_inductor/index_propagation.py
@@ -189,7 +189,12 @@ class IndexPropagation:
 
     """
 
-    def __init__(self, inner: Any, iter_ranges: Dict[sympy.Symbol, sympy.Expr]):
+    def __init__(
+        self,
+        inner: Any,
+        iter_ranges: Dict[sympy.Symbol, sympy.Expr],
+        indirect_var_ranges: Dict[sympy.Symbol, sympy.Expr],
+    ):
         self._inner = inner
         self.shape_env = V.graph.sizevars.shape_env
 
@@ -199,6 +204,9 @@ def __init__(self, inner: Any, iter_ranges: Dict[sympy.Symbol, sympy.Expr]):
         self.var_to_range = tuple(
             itertools.chain(self.shape_env.var_to_range.items(), var_to_range.items())
         )
+        # NOTE: this is intentionally kept as a reference so the caller can
+        # update it in-place
+        self.indirect_var_ranges = indirect_var_ranges
 
         axioms = []
         for x, s in iter_ranges.items():
@@ -306,10 +314,17 @@ def statically_true(self, e):
               to perform wrap_expr and in CSEProxy.check_bounds to elide upper / lower bounds also
               for indirect_indexing
         """
+        var_to_range = (
+            *self.var_to_range,
+            *(
+                (k, ValueRanges(0, upper_bound(v) - 1))
+                for k, v in self.indirect_var_ranges.items()
+            ),
+        )
         evaluated = self.shape_env._maybe_evaluate_static(
             e,
             axioms=self.axioms,
-            var_to_range=self.var_to_range,
+            var_to_range=var_to_range,
         )
         return bool(evaluated)
 
@@ -351,9 +366,4 @@ def wrap_expr(expr):
         indirect_var = self.fallback(
             "indirect_indexing", (index, size, check), {}
         ).value
-        assert (
-            indirect_var not in self.var_to_range
-        ), f"{indirect_var} should've been created in the fallback."
-        indirect_range = (indirect_var, ValueRanges(0, upper_bound(size) - 1))
-        self.var_to_range = self.var_to_range + (indirect_range,)
         return indirect_var
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -6479,6 +6479,7 @@ def __init__(self, fn, args, var_ranges):
         self.submodules = {"get_index": self.get_index}
         self.subblocks = {}
         self.indirect_vars = []
+        self.indirect_var_ranges: Dict[sympy.Symbol, sympy.Expr] = {}
         self.root_block = LoopBodyBlock(self, fn, args)
         self.indexing = None
 
@@ -6531,7 +6532,9 @@ def add_submodule(self, block, prefix):
 
     def add_indirect(self, size):
         var = sympy_index_symbol_with_prefix(SymT.INDIRECT, len(self.indirect_vars))
+        assert var not in self.indirect_var_ranges
         self.indirect_vars.append(var)
+        self.indirect_var_ranges[var] = size
         return var
 
     def replace_indirect(self, old, new):
@@ -6712,7 +6715,9 @@ def output(result):
             CaptureIndexing(proxy_ops), self.body.var_ranges
         )
         if config.constant_and_index_propagation:
-            handler = IndexPropagation(handler, self.body.var_ranges)
+            handler = IndexPropagation(
+                handler, self.body.var_ranges, self.body.indirect_var_ranges
+            )
 
         with V.set_ops_handler(handler):
             # This indirection is just a cute way to get IndexPropagation to
