Update on "thread through specialization to compile_fx"

bobrenjc93 · bobrenjc93 · commit 67d9211cbcb9 · 2025-05-01T22:25:16.000-07:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]
diff --git a/torch/_subclasses/meta_utils.py b/torch/_subclasses/meta_utils.py
@@ -882,6 +882,7 @@ def meta_tensor(
         callback_: _MetaTensorCallback[_TensorT],
         source: Optional[Source],
         symbolic_context: Optional[SymbolicContext],
+        specialization=specialization,
     ) -> _TensorT:
         callback: _MetaTensorCallbackOptDevice = functools.partial(
             callback_, device=t.device
@@ -965,6 +966,7 @@ def sym_sizes_strides_storage_offset(
                         [d in t.dynamo_dynamic_indices for d in range(t.ndim)],
                         src,
                         symbolic_context=symbolic_context,
+                        specialization=specialization,
                     )
             else:
                 return (t.size, t.stride, t.storage_offset)
@@ -1924,6 +1926,7 @@ def __call__(
                 callback_,
                 source,
                 symbolic_context,
+                specialization=specialization,
             )
 
         if type(t) is torch.nn.Parameter:
diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
@@ -4142,6 +4142,7 @@ def _create_symbolic_sizes_strides_storage_offset(
         source: Source,
         *,
         symbolic_context: Optional[SymbolicContext] = None,
+        specialization = None,
     ) -> tuple[tuple[IntLikeType, ...], tuple[IntLikeType, ...], IntLikeType,]:
         dim = len(ex_size)
 
@@ -4220,6 +4221,11 @@ def _create_symbolic_sizes_strides_storage_offset(
             )
             for i, (sym, hint) in enumerate(zip(size, ex_size))
         ]
+
+        for i, size in enumerate(sym_sizes):
+            if i in specialization.idxs:
+                expect_true(specialization.lambdas[i](size))
+
         sym_stride = []
         for i, stride_expr in enumerate(stride):
             # NB: Don't duck size the stride; instead use the expression
