Record inputs at time of tracing, constrain to them for triton fn (#145448)

eellison · pytorchmergebot · commit 7796e308d063 · 2025-01-30T16:54:08.000Z
Record input fake tensors at time of tracing and store them in the node meta. Inductor passes have the possibility of changing strides, so it is safer to record the strides of the inputs at tracing. See, #137979 for more context. We can also extend this to custom ops, and user-visible outputs. If this ends up being compilation time sensitive we can just record strides (and maybe storage offset, per @zou3519) instead of the complete fake tensor. Pull Request resolved: #145448 Approved by: https://github.com/zou3519 ghstack dependencies: #145953
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -393,6 +393,7 @@ def significant_strides_equal(
     """
     Returns true if the strides are equal, ignoring dimensions of size 1 .
     """
+    assert len(shape) == len(strides1) and len(strides1) == len(strides2)
     for dim, s1, s2 in zip(shape, strides1, strides2):
         if V.graph.sizevars.statically_known_leq(dim, 1):  # type: ignore[arg-type]
             continue
