Torch package support in dynamo (pytorch#91821)

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#91821
Approved by: https://github.com/suo

Author: voznesenskym

test/dynamo/test_misc.py

@@ -39,6 +39,15 @@ def my_custom_function(x):
     return x + 1
 
 
+class MyPickledModule(torch.nn.Module):
+    def __init__(self, z):
+        super().__init__()
+        self.z = z
+
+    def forward(self, x, y):
+        return x * x * x + y + self.z
+
+
 class MiscTests(torch._dynamo.test_case.TestCase):
     def test_boolarg(self):
         def boolarg(aa, bb, flag):
@@ -3260,6 +3269,31 @@ def fn(x, y):
         res = opt_fn(x, y)
         self.assertTrue(same(ref, res))
 
+    def test_torch_package_working_with_inductor_trace(self):
+        inputs = [torch.randn([2, 2]), torch.randn([2, 2])]
+
+        optimized_model = torch._dynamo.optimize(backend="inductor")(
+            MyPickledModule(torch.randn([2, 2]))
+        )
+        from torch import package
+
+        path = "/tmp/MyPickledModule.pt"
+        package_name = "MyPickledModule"
+        resource_name = "MyPickledModule.pkl"
+
+        model = MyPickledModule(torch.randn([2, 2]))
+
+        with package.PackageExporter(path) as exp:
+            exp.extern("**")
+            exp.save_pickle(package_name, resource_name, model)
+
+        imp = package.PackageImporter(path)
+        loaded_model = imp.load_pickle(package_name, resource_name)
+
+        optimized_loaded_model = torch._dynamo.optimize(backend="inductor")(
+            loaded_model
+        )
+
 
 class CustomFunc1(torch.autograd.Function):
     @staticmethod

torch/_dynamo/guards.py

@@ -97,6 +97,19 @@ def __init__(
         else:
             scope = dict()
         self.scope: Dict[str, object] = scope
+
+        if "__builtins__" not in self.scope:
+            self.scope["__builtins__"] = {}
+        for (
+            name,
+            package_module,
+        ) in torch.package.package_importer._package_imported_modules.items():
+            name = name.replace(">", "_").replace("<", "_").replace(".", "_dot_")
+            # Write the package module into the scope so that we can import it
+            self.scope["__builtins__"][name] = package_module  # type: ignore[index]
+            # Write the demangled name to the scope so that we can use it
+            self.scope[name] = package_module
+
         self.argnames: List[str] = []
         # Code is python expression strings generated for each guard
         self.code: List[str] = []

torch/_dynamo/symbolic_convert.py

@@ -662,8 +662,16 @@ def STORE_GLOBAL(self, inst):
 
     def import_source(self, module_name):
         """Create an alias to a module for use in guards"""
-        value = importlib.import_module(module_name)
-        alias = f"__import_{module_name.replace('.', '_dot_')}"
+        if "torch_package" in module_name:
+            value = torch.package.package_importer._package_imported_modules[
+                module_name
+            ]
+            alias = (
+                module_name.replace(">", "_").replace("<", "_").replace(".", "_dot_")
+            )
+        else:
+            value = importlib.import_module(module_name)
+            alias = f"__import_{module_name.replace('.', '_dot_')}"
         f_globals = self.output.root_globals
         assert alias not in f_globals or f_globals[alias] is value
         f_globals[alias] = value