[inductor] Enable multilayer reductions with dynamic shapes

test/inductor/test_torchinductor.py

@@ -930,6 +930,15 @@ def fn(a):
 
         self.common(fn, ((torch.rand((10, 3, 352, 352), dtype=torch.float16),)))
 
+    def test_multilayer_prime_size(self):
+        def fn(a):
+            return torch.max(a), torch.sum(a)
+
+        # Requires masked loading for the intermediate reduction
+        sample = torch.full((3999971,), 0, dtype=torch.int64)
+        sample[-1] = 1
+        self.common(fn, (sample,))
+
     def test_expanded_reduction(self):
         if self.device == "cpu":
             raise unittest.SkipTest(

torch/_inductor/ir.py

@@ -595,6 +595,9 @@ def num_splits(
         def _is_static(x):
             return isinstance(x, (int, sympy.Integer))
 
+        reduction_numel_hint = V.graph.sizevars.symbolic_hint(reduction_numel)
+        numel_hint = V.graph.sizevars.symbolic_hint(sympy_product(ranges))
+
         should_split = (
             is_triton(device)
             and reduction_type
@@ -604,9 +607,9 @@ def _is_static(x):
                 "var_unnormalized",
             }
             and config.split_reductions
-            and all(_is_static(r) for r in ranges)
-            and all(_is_static(r) for r in reduction_ranges)
-            and _is_static(reduction_numel)
+            # We don't support unbacked symints
+            and _is_static(reduction_numel_hint)
+            and _is_static(numel_hint)
         )
         if not should_split:
             return ReductionHint.DEFAULT, 1
@@ -689,8 +692,6 @@ def outer_reduction_splits(reduction_numel_hint, numel_hint):
                 rvals_per_thread * split_size
             )
 
-        reduction_numel_hint = V.graph.sizevars.size_hint(reduction_numel)
-        numel_hint = V.graph.sizevars.size_hint(sympy_product(ranges))
         # easy cases
         if numel_hint == 1:
             return ReductionHint.INNER, inner_reduction_splits(
@@ -994,24 +995,14 @@ def create_multilayer(
         """
         reduction_numel = sympy_product(reduction_ranges)
 
-        # TODO(jansel): convert this to dynamic shapes
-        # TODO(jansel): realize the reduction so we can do dynamic indexing
-        reduction_ranges = [
-            sympy.Integer(V.graph.sizevars.evaluate_static_shape(s))
-            for s in reduction_ranges
-        ]
-        reduction_numel = sympy.Integer(
-            V.graph.sizevars.evaluate_static_shape(reduction_numel)
+        need_mask = not V.graph.sizevars.is_expr_static_and_true(
+            sympy.Eq(reduction_numel % split, 0)
         )
 
-        if V.graph.sizevars.size_hint(reduction_numel) % split == 0:
-            need_mask = False
-        else:
-            need_mask = True
-
         split = sympy.Integer(split)
         block_size = FloorDiv(reduction_numel + (split - 1), split)
 
+        # TODO(jansel): realize the reduction so we can do dynamic indexing
         reindex = View.dynamic_reshape_indexer(reduction_ranges, [reduction_numel])
 
         def wrapper_fn(index, reduction_index):

torch/_inductor/sizevars.py

@@ -358,7 +358,8 @@ def evaluate_static_shape(self, left: Expr) -> int:
     def evaluate_static_shapes(self, left: List[Expr]) -> List[int]:
         return [self.evaluate_static_shape(x) for x in left]
 
-    def size_hint(self, expr: Expr) -> int:
+    def symbolic_hint(self, expr: Expr) -> Expr:
+        # Substitute all hints into expr, but leave unbacked symints alone
         if not isinstance(expr, Expr):
             assert isinstance(expr, int)
             return expr
@@ -368,7 +369,10 @@ def size_hint(self, expr: Expr) -> int:
         while any(s.name.startswith("ps") for s in free_symbols):
             expr = sympy_subs(expr, self.inv_precomputed_replacements)
             free_symbols = expr.free_symbols
-        out = sympy_subs(expr, self.var_to_val)
+        return sympy_subs(expr, self.var_to_val)
+
+    def size_hint(self, expr: Expr) -> int:
+        out = self.symbolic_hint(expr)
         try:
             return int(out)
         except Exception: