Ensure block analysis only matches integer dims and strides

kundaMwiza · kundaMwiza · commit 9f49c4b1b48c · 2025-03-20T10:42:18.000Z
diff --git a/test/inductor/test_torchinductor_strided_blocks.py b/test/inductor/test_torchinductor_strided_blocks.py
@@ -952,6 +952,31 @@ def fn(a):
             rtol=0.06,
         )
 
+    # Integration test to ensure that matched dims & strides from match_mod_div_expr
+    # are nonnegative integers. This test case has the following index
+    # index=(ModularIndexing(xindex, 4, 4)) + 4*(ModularIndexing(xindex, 32, 2))
+    # and the match below is a candidate that is invalid:
+    # match={
+    #   dim_mod4_: 32, dim_mod3_: 2, stride_mod3_: 4, dim_mod2_: 1/16,
+    #   dim_mod1_: 4, stride_mod1_: 1, stride_mod4_: 0, stride_mod2_: 0, stride_mod0_: 0
+    # }
+    # This is now fixed by ensuring that that wild symbols only match nonnegative integers
+    def test_ensure_integral_dims_and_strides(self):
+        def model(data, *args):
+            return torch.nn.functional.unfold(data, *args)
+
+        data = torch.zeros([2, 3, 5, 5], dtype=torch.float16, requires_grad=True)
+        args = [2, 1, 0, 1]
+        run_and_compare(
+            self,
+            model,
+            data,
+            *args,
+            expected_num_triton_kernels=2,
+            expected_num_block_pointers=4,
+            compile_kwargs={"fullgraph": True},
+        )
+
 
 @unittest.skipIf(not TRITON_HAS_CPU, "requires triton CPU backend")
 @config.patch(cpu_backend="triton")
diff --git a/torch/_inductor/codegen/block_analysis.py b/torch/_inductor/codegen/block_analysis.py
@@ -63,7 +63,11 @@ def match_mod_div_block_expr(
         index = cls._preprocess(index)
 
         # Pattern match to find the strides and offset.
-        wild = functools.partial(sympy.Wild, exclude=[index_var])
+        wild = functools.partial(
+            sympy.Wild,
+            exclude=[index_var],
+            properties=[lambda x: x.is_integer and x.is_nonnegative],
+        )
         dims: list[Expr] = [wild(f"dim_mod{idx}") for idx in range(num_dims)]
         strides: list[Expr] = [wild(f"stride_mod{idx}") for idx in range(num_dims)]
 
@@ -167,7 +171,11 @@ def match_affine_block_expr(
         stride.
         """
         index = cls._preprocess(index)
-        stride = sympy.Wild("stride", exclude=[index_var])
+        stride = sympy.Wild(
+            "stride",
+            exclude=[index_var],
+            properties=[lambda x: x.is_integer and x.is_nonnegative],
+        )
         m = index.match(index_var * stride)
         if m is None:
             return None
