[jit] Fix name collision on load (pytorch#35720)

Summary:
Pull Request resolved: pytorch#35720

When modules are saved, all relevant types are serialized according to
their qualified name with a compilation unit. Since qualified names are
guaranteed to be unique within a compilation unit, this normally works
fine.

On load, all types are registered in a compilation unit owned by the
script::Module. Type names are not unique across compilation units, so
if you load two modules with colliding type names, make them submodules
of yet another module, and save that module, there is the potential of a
name collision. See the added tests for examples if that description is
confusing.

The solution is to unique type names when serializing code by mangling
them if we detect a name collision.

Test Plan: Imported from OSS

Differential Revision: D20749423

Pulled By: suo

fbshipit-source-id: a8827ff1d4a89f3e7964dbbb49b4381863da3e6a

Author: suo

aten/src/ATen/core/jit_type.h

@@ -112,7 +112,7 @@ struct CAFFE2_API Type : std::enable_shared_from_this<Type> {
   //
   // Takes a custom printer that users can pass in to customize the output of
   // this method.
-  std::string python_str(TypePrinter printer = nullptr) const {
+  std::string python_str(TypePrinter printer) const {
     if (printer) {
       // the printer can return nullopt to fall through to the default impl
       if (auto renamed = printer(shared_from_this())) {
@@ -121,6 +121,11 @@ struct CAFFE2_API Type : std::enable_shared_from_this<Type> {
     }
     return python_str_impl(printer);
   }
+  std::string python_str() const {
+    // Overload instead of define a default value for `printer` to help
+    // debuggers out.
+    return python_str(nullptr);
+  }
 
   TypeKind kind() const {
     return kind_;

caffe2/CMakeLists.txt

@@ -467,6 +467,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     ${TORCH_SRC_DIR}/csrc/jit/runtime/jit_exception.cpp
     ${TORCH_SRC_DIR}/csrc/jit/frontend/string_to_type.cpp
     ${TORCH_SRC_DIR}/csrc/jit/serialization/source_range_serialization.cpp
+    ${TORCH_SRC_DIR}/csrc/jit/serialization/type_name_uniquer.cpp
     ${TORCH_SRC_DIR}/csrc/jit/frontend/tracer.cpp
     ${TORCH_SRC_DIR}/csrc/jit/testing/hooks_for_testing.cpp
     ${TORCH_SRC_DIR}/csrc/utils/tensor_flatten.cpp

test/jit/test_save_load.py

@@ -0,0 +1,331 @@
+import os
+import io
+import sys
+
+import torch
+from torch import Tensor
+from typing import NamedTuple
+
+# Make the helper files in test/ importable
+pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+sys.path.append(pytorch_test_dir)
+from torch.testing._internal.jit_utils import JitTestCase, clear_class_registry
+
+if __name__ == "__main__":
+    raise RuntimeError(
+        "This test file is not meant to be run directly, use:\n\n"
+        "\tpython test/test_jit.py TESTNAME\n\n"
+        "instead."
+    )
+
+
+class TestSaveLoad(JitTestCase):
+    def test_different_modules(self):
+        """
+        Exercise the situation where we have the same qualified name
+        in two different CompilationUnits on save/load.
+        """
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super(Foo, self).__init__()
+                self.foo = torch.nn.Linear(2, 2)
+                self.bar = torch.nn.Linear(2, 2)
+
+            def forward(self, x):
+                x = self.foo(x)
+                x = self.bar(x)
+                return x
+
+        first_script_module = torch.jit.script(Foo())
+        first_saved_module = io.BytesIO()
+        torch.jit.save(first_script_module, first_saved_module)
+        first_saved_module.seek(0)
+
+        clear_class_registry()
+
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super(Foo, self).__init__()
+                self.foo = torch.nn.Linear(2, 2)
+
+            def forward(self, x):
+                x = self.foo(x)
+                return x
+
+        second_script_module = torch.jit.script(Foo())
+        second_saved_module = io.BytesIO()
+        torch.jit.save(torch.jit.script(Foo()), second_saved_module)
+        second_saved_module.seek(0)
+
+        clear_class_registry()
+
+        self.assertEqual(
+            first_script_module._c.qualified_name, second_script_module._c.qualified_name
+        )
+
+        class ContainsBoth(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.add_module("second", torch.jit.load(second_saved_module))
+                self.add_module("first", torch.jit.load(first_saved_module))
+
+            def forward(self, x):
+                x = self.first(x)
+                x = self.second(x)
+                return x
+
+        sm = torch.jit.script(ContainsBoth())
+        contains_both = io.BytesIO()
+        torch.jit.save(sm, contains_both)
+        contains_both.seek(0)
+        sm = torch.jit.load(contains_both)
+
+    def test_different_functions(self):
+        """
+        Exercise the situation where we have the same qualified name
+        in two different CompilationUnits on save/load.
+        """
+        def lol(x):
+            return x
+
+        class Foo(torch.nn.Module):
+            def forward(self, x):
+                return lol(x)
+
+        first_script_module = torch.jit.script(Foo())
+        first_saved_module = io.BytesIO()
+        torch.jit.save(first_script_module, first_saved_module)
+        first_saved_module.seek(0)
+
+        clear_class_registry()
+
+        def lol(x):  # noqa: F811
+            return "hello"
+
+        class Foo(torch.nn.Module):
+            def forward(self, x):
+                return lol(x)
+
+        second_script_module = torch.jit.script(Foo())
+        second_saved_module = io.BytesIO()
+        torch.jit.save(torch.jit.script(Foo()), second_saved_module)
+        second_saved_module.seek(0)
+
+        clear_class_registry()
+
+        self.assertEqual(
+            first_script_module._c.qualified_name, second_script_module._c.qualified_name
+        )
+
+        class ContainsBoth(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.add_module("second", torch.jit.load(second_saved_module))
+                self.add_module("first", torch.jit.load(first_saved_module))
+
+            def forward(self, x):
+                x = self.first(x)
+                x = self.second(x)
+                return x
+
+        sm = torch.jit.script(ContainsBoth())
+        contains_both = io.BytesIO()
+        torch.jit.save(sm, contains_both)
+        contains_both.seek(0)
+        sm = torch.jit.load(contains_both)
+
+    def test_different_interfaces(self):
+        """
+        Exercise the situation where we have the same qualified name
+        in two different CompilationUnits on save/load.
+        """
+        @torch.jit.interface
+        class MyInterface(object):
+            def bar(self, x):
+                # type: (Tensor) -> Tensor
+                pass
+
+        @torch.jit.script
+        class ImplementInterface(object):
+            def __init__(self):
+                pass
+
+            def bar(self, x):
+                return x
+
+        class Foo(torch.nn.Module):
+            __annotations__ = {"interface": MyInterface}
+
+            def __init__(self):
+                super().__init__()
+                self.interface = ImplementInterface()
+
+            def forward(self, x):
+                return self.interface.bar(x)
+
+        first_script_module = torch.jit.script(Foo())
+        first_saved_module = io.BytesIO()
+        torch.jit.save(first_script_module, first_saved_module)
+        first_saved_module.seek(0)
+
+        clear_class_registry()
+
+        @torch.jit.interface
+        class MyInterface(object):
+            def not_bar(self, x):
+                # type: (Tensor) -> Tensor
+                pass
+
+        @torch.jit.script  # noqa: F811
+        class ImplementInterface(object):  # noqa: F811
+            def __init__(self):
+                pass
+
+            def not_bar(self, x):
+                return x
+
+        class Foo(torch.nn.Module):
+            __annotations__ = {"interface": MyInterface}
+
+            def __init__(self):
+                super().__init__()
+                self.interface = ImplementInterface()
+
+            def forward(self, x):
+                return self.interface.not_bar(x)
+
+        second_script_module = torch.jit.script(Foo())
+        second_saved_module = io.BytesIO()
+        torch.jit.save(torch.jit.script(Foo()), second_saved_module)
+        second_saved_module.seek(0)
+
+        clear_class_registry()
+
+        self.assertEqual(
+            first_script_module._c.qualified_name, second_script_module._c.qualified_name
+        )
+
+        class ContainsBoth(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.add_module("second", torch.jit.load(second_saved_module))
+                self.add_module("first", torch.jit.load(first_saved_module))
+
+            def forward(self, x):
+                x = self.first(x)
+                x = self.second(x)
+                return x
+
+        sm = torch.jit.script(ContainsBoth())
+        contains_both = io.BytesIO()
+        torch.jit.save(sm, contains_both)
+        contains_both.seek(0)
+        sm = torch.jit.load(contains_both)
+
+    def test_many_collisions(self):
+        class MyCoolNamedTuple(NamedTuple):
+            a: int
+
+        @torch.jit.interface
+        class MyInterface(object):
+            def bar(self, x):
+                # type: (Tensor) -> Tensor
+                pass
+
+        @torch.jit.script
+        class ImplementInterface(object):
+            def __init__(self):
+                pass
+
+            def bar(self, x):
+                return x
+
+        def lol(x):
+            return x
+
+        class Foo(torch.nn.Module):
+            interface: MyInterface
+
+            def __init__(self):
+                super().__init__()
+                self.foo = torch.nn.Linear(2, 2)
+                self.bar = torch.nn.Linear(2, 2)
+                self.interface = ImplementInterface()
+
+            def forward(self, x):
+                x = self.foo(x)
+                x = self.bar(x)
+                x = lol(x)
+                x = self.interface.bar(x)
+
+                return x, MyCoolNamedTuple(a=5)
+
+
+        first_script_module = torch.jit.script(Foo())
+        first_saved_module = io.BytesIO()
+        torch.jit.save(first_script_module, first_saved_module)
+        first_saved_module.seek(0)
+
+        clear_class_registry()
+
+        @torch.jit.interface
+        class MyInterface(object):
+            def not_bar(self, x):
+                # type: (Tensor) -> Tensor
+                pass
+
+        @torch.jit.script   # noqa F811
+        class ImplementInterface(object):  # noqa F811
+            def __init__(self):
+                pass
+
+            def not_bar(self, x):
+                return x
+
+        def lol(x):  # noqa F811
+            return "asdofij"
+
+        class MyCoolNamedTuple(NamedTuple):  # noqa F811
+            a: str
+
+        class Foo(torch.nn.Module):
+            interface: MyInterface
+
+            def __init__(self):
+                super().__init__()
+                self.foo = torch.nn.Linear(2, 2)
+                self.interface = ImplementInterface()
+
+            def forward(self, x):
+                x = self.foo(x)
+                self.interface.not_bar(x)
+                x = lol(x)
+                return x, MyCoolNamedTuple(a="hello")
+
+        second_script_module = torch.jit.script(Foo())
+        second_saved_module = io.BytesIO()
+        torch.jit.save(second_script_module, second_saved_module)
+        second_saved_module.seek(0)
+
+        clear_class_registry()
+
+        self.assertEqual(
+            first_script_module._c.qualified_name, second_script_module._c.qualified_name
+        )
+
+        class ContainsBoth(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.add_module("second", torch.jit.load(second_saved_module))
+                self.add_module("first", torch.jit.load(first_saved_module))
+
+            def forward(self, x):
+                x, named_tuple_1 = self.first(x)
+                x, named_tuple_2 = self.second(x)
+                return len(x + named_tuple_2.a) + named_tuple_1.a
+
+        sm = torch.jit.script(ContainsBoth())
+        contains_both = io.BytesIO()
+        torch.jit.save(sm, contains_both)
+        contains_both.seek(0)
+        sm = torch.jit.load(contains_both)

test/test_jit.py

@@ -24,6 +24,7 @@
 from jit.test_unsupported_ops import TestUnsupportedOps  # noqa: F401
 from jit.test_freezing import TestFreezing  # noqa: F401
 from jit.test_functional_blocks import TestFunctionalBlocks  # noqa: F401
+from jit.test_save_load import TestSaveLoad  # noqa: F401
 
 # Torch
 from torch import Tensor

tools/build_variables.bzl

@@ -87,6 +87,7 @@ libtorch_sources = [
     "torch/csrc/jit/ir/type_hashing.cpp",
     "torch/csrc/jit/serialization/export.cpp",
     "torch/csrc/jit/serialization/export_module.cpp",
+    "torch/csrc/jit/serialization/type_name_uniquer.cpp",
     "torch/csrc/jit/passes/pass_manager.cpp",
     "torch/csrc/jit/serialization/pickler.cpp",
     "torch/csrc/jit/serialization/unpickler.cpp",

torch/csrc/distributed/rpc/utils.cpp

@@ -264,12 +264,10 @@ std::string wireSerialize(
   }
 
   if (!tensors.empty()) {
-    torch::jit::Pickler pickler(
-        [&](const void* buf, size_t sz) -> size_t {
-          metaEntry.append(static_cast<const char*>(buf), sz);
-          return sz;
-        },
-        nullptr);
+    torch::jit::Pickler pickler([&](const void* buf, size_t sz) -> size_t {
+      metaEntry.append(static_cast<const char*>(buf), sz);
+      return sz;
+    });
     pickler.protocol();
     pickler.pushIValue(cloneSparseTensors(tensors));
     pickler.stop();