Fix bug in AOTI lowering

tugsbayasgalan · tugsbayasgalan · commit 002dc270446d · 2025-03-03T13:28:39.000-08:00
ghstack-source-id: 1c81ea4 Pull Request resolved: #148364
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -1161,6 +1161,18 @@ def forward(self, q, k, v):
         path = torch._inductor.aot_compile(ep.module(), inputs)
         aot_model = torch._export.aot_load(path, device=self.device)
         torch.testing.assert_close(m(*inputs), aot_model(*inputs))
+    
+    def test_aoti_constant_tensor(self):
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = torch.ones(4, 4)
+                self.b = torch.ones(4, 4)
+            def forward(self, x):
+                return torch.ops.aten.linear.default(x, self.a, self.b)
+        
+        ep = torch.export.export(Foo(), (torch.ones(4, 4),), strict=False).run_decompositions({})
+        _ = torch._inductor.aoti_compile_and_package(ep)
 
     def test_large_grid(self):
         if self.device != GPU_TYPE:
diff --git a/torch/_inductor/compile_fx.py b/torch/_inductor/compile_fx.py
@@ -2184,6 +2184,22 @@ def bw_compiler(
                     trace_joint=False,
                     decompositions=decompositions,
                 )
+
+                from torch._export.utils import _detect_fake_mode_from_gm
+                fake_mode = _detect_fake_mode_from_gm(gm)
+                # aot_export_module doesn't account for constant tensor attributes 
+                # so we end up having tensors that don't have fake vals attached. 
+                # This can happen when upstream export is non-strict where we 
+                # preserve the original module params/buffers. Once AOTI switches
+                # to ep.run_decompositions() flow to lower to post-autograd opset 
+                # this will go away. 
+                for node in gm.graph.nodes:
+                    if node.op == "get_attr" and "val" not in node.meta:
+                        target = getattr(gm, node.target)
+                        if isinstance(target, torch.Tensor):
+                            node.meta["val"] = fake_mode.from_tensor(target)
+
+
             unlifted_gm = _unlift_graph(model_, gm, graph_signature)
             if "dynamo_flat_name_to_original_fqn" in model_.meta:
                 unlifted_gm.meta["dynamo_flat_name_to_original_fqn"] = model_.meta[
