extend torch.jit._overload to modules

root · root · commit 483786c8fab8 · 2019-08-16T10:16:46.000-07:00
ghstack-source-id: af601e7 Pull Request resolved: #24259
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -13166,31 +13166,26 @@ def test_uses():
 
         self.checkScript(test_uses, ())
 
-    @unittest.skipIf(True, "Removing weak script")
-    def test_overloading(self):
-        @torch._jit_internal.weak_module
+    def test_method_overloading(self):
         class W(torch.nn.Module):
-            __overloads__ = {'forward': ['forward_tuple', 'forward_tensor']}
-
             def __init__(self):
                 super(W, self).__init__()
 
-            @torch._jit_internal.weak_script_method
-            def forward_tuple(self, x):
+            @torch.jit._overload_method  # noqa: F811
+            def forward(self, x):  # noqa: F811
                 # type: (Tuple[Tensor, Tensor]) -> Tensor
-                return x[0] + 5
-
-            def forward(self, x):
-                # manually do argument switching
-                if isinstance(x, tuple):
-                    return self.forward_tuple(x)
-                else:
-                    return self.forward_tensor(x)
+                pass
 
-            @torch._jit_internal.weak_script_method
-            def forward_tensor(self, x):
+            @torch.jit._overload_method  # noqa: F811
+            def forward(self, x):  # noqa: F811
                 # type: (Tensor) -> Tensor
-                return x + 20
+                pass
+
+            def forward(self, x):  # noqa: F811
+                if isinstance(x, Tensor):
+                    return x + 20
+                else:
+                    return x[0] + 5
 
         class S(torch.jit.ScriptModule):
             def __init__(self):
@@ -13201,14 +13196,61 @@ def __init__(self):
             def forward(self, x):
                 return self.weak(x) + self.weak((x, x))
 
-        s = S()
+        s_mod = S()
         x = torch.ones(1)
-        self.assertEqual(s(x), x + 20 + 5 + x)
+        self.assertEqual(s_mod(x), x + 20 + 5 + x)
 
         w = W()
         self.assertEqual(w((x, x)), x + 5)
         self.assertEqual(w((x)), x + 20)
 
+        class Unannotated(torch.nn.Module):
+            def __init__(self):
+                super(Unannotated, self).__init__()
+
+            @torch.jit._overload_method  # noqa: F811
+            def hello(self, x):  # noqa: F811
+                pass
+
+            @torch.jit._overload_method  # noqa: F811
+            def hello(self, x):  # noqa: F811
+                # type: (int) -> (int)
+                pass
+
+            def hello(self, x):  # noqa: F811
+                return x + 3
+
+            def forward(self):
+                return self.hello(1), self.hello(.5)
+
+        w = Unannotated()
+        with self.assertRaisesRegex(Exception, "explicitly add type annotations to overloaded functions"):
+            torch.jit.script(w)
+
+        class CompileOverloadError(torch.nn.Module):
+            def __init__(self):
+                super(CompileOverloadError, self).__init__()
+
+            @torch.jit._overload_method  # noqa: F811
+            def hello(self, x):  # noqa: F811
+                # type: (str) -> (int)
+                pass
+
+            @torch.jit._overload_method  # noqa: F811
+            def hello(self, x):  # noqa: F811
+                # type: (int) -> (int)
+                pass
+
+            def hello(self, x):  # noqa: F811
+                return x + 1
+
+            def forward(self):
+                return self.hello("hi"), self.hello(.5)
+
+        w = CompileOverloadError()
+        with self.assertRaisesRegex(Exception, "but instead found type \'str\'"):
+            torch.jit.script(w)
+
     def test_select_after_chunk(self):
         def foo(x):
             chunked = torch.chunk(x, 1)
diff --git a/torch/_jit_internal.py b/torch/_jit_internal.py
@@ -4,6 +4,8 @@
 circular dependency problems
 """
 
+import threading
+import sys
 import inspect
 import weakref
 import torch._C
@@ -276,6 +278,50 @@ def _get_fn_overloads(qual_name):
 def _clear_fn_overloads(qual_name):
     del _overloaded_fns[qual_name]
 
+def get_class_name(method):
+    current_frame = sys._current_frames()[threading.currentThread().ident]
+
+    # one for the get_class_name call, one for _overload_method call
+    for i in range(2):
+        current_frame = current_frame.f_back
+    class_name = current_frame.f_code.co_name
+    return class_name
+
+# At the the point the decorator is applied to class methods the method
+# has no reference to its owning class. _qualified_name would not include
+# the class it is defined in, so any methods with the same name in the same file
+# would have the same _qualified_name, even if they were defined in different
+# classes. This problem only exists in python 2.
+# We get around this problem by looking at the stack frame and identifying
+# the class name.
+
+# qualified_name => class name => list[overload_functions]
+_overloaded_methods = {}  # noqa: T484
+
+def _overload_method(func):
+    qual_name = _qualified_name(func)
+    global _overloaded_methods
+    class_name_map = _overloaded_methods.get(qual_name, None)
+    if class_name_map is None:
+        class_name_map = {}
+        _overloaded_methods[qual_name] = class_name_map
+
+    class_name = get_class_name(func)
+    method_overloads = class_name_map.get(class_name, None)
+    if method_overloads is None:
+        method_overloads = []
+        class_name_map[class_name] = method_overloads
+
+    method_overloads.append(func)
+    return func
+
+def _get_overloaded_methods(method, class_name):
+    qual_name = _qualified_name(method)
+    class_name_map = _overloaded_methods.get(qual_name, None)
+    if class_name_map is None:
+        return None
+    return class_name_map.get(class_name, None)
+
 try:
     import typing
     from typing import Tuple, List, Dict, Optional
diff --git a/torch/csrc/jit/script/init.cpp b/torch/csrc/jit/script/init.cpp
@@ -750,7 +750,14 @@ void initJitScriptBindings(PyObject* module) {
         auto new_def = implementation_def.withDecl(overload_decl);
         return script_compile_function(name, new_def, defaults, std::move(rcb));
       });
-
+  m.def(
+      "_replace_overloaded_method_decl",
+      [](const Decl& overload_decl,
+         const Def& implementation_def,
+         const std::string& new_name) {
+        checkOverloadDecl(overload_decl, implementation_def.decl());
+        return implementation_def.withDecl(overload_decl).withName(new_name);
+      });
   m.def(
       "_create_function_from_trace",
       [](std::string qualname,
diff --git a/torch/csrc/jit/script/python_sugared_value.cpp b/torch/csrc/jit/script/python_sugared_value.cpp
@@ -218,25 +218,29 @@ std::shared_ptr<SugaredValue> OverloadedMethodValue::call(
   std::vector<NamedValue> new_inputs = inputs.vec();
   new_inputs.insert(new_inputs.begin(), module_);
 
-  for (const std::string& method_name : method_names_) {
-    auto cls = module_->type()->expect<ClassType>();
-    const auto fn = cls->getMethod(method_name);
-    auto match = tryMatchSchema(
-        fn->getSchema(),
-        loc,
-        *caller.graph().get(),
-        c10::nullopt,
-        new_inputs,
-        attributes,
-        &err,
-        true);
-    if (match) {
-      return MethodValue(module_, method_name)
-          .call(loc, caller, inputs, attributes, n_binders);
+  std::stringstream failure_messages;
+  for (bool allow_conversions : {false, true}) {
+    // clear previous error messages
+    failure_messages.str("");
+    for (const std::string& method_name : method_names_) {
+      auto cls = module_->type()->expect<ClassType>();
+      const auto fn = cls->getMethod(method_name);
+      auto match = tryMatchSchema(
+          fn->getSchema(),
+          loc,
+          *caller.graph().get(),
+          c10::nullopt,
+          new_inputs,
+          attributes,
+          &err,
+          allow_conversions);
+      if (match) {
+        return MethodValue(module_, method_name)
+            .call(loc, caller, inputs, attributes, n_binders);
+      }
     }
   }
-  throw ErrorReport(loc) << "Could not find any matching overloads\n"
-                         << err.str();
+  throw ErrorReport(loc) << failure_messages.str();
 }
 
 std::shared_ptr<SugaredValue> OverloadedFunctionValue::call(
diff --git a/torch/csrc/jit/script/python_tree_views.cpp b/torch/csrc/jit/script/python_tree_views.cpp
@@ -150,7 +150,8 @@ void initTreeViewBindings(PyObject* module) {
             const auto& r = name.range();
             return Def::create(r, name, decl, wrap_list(r, std::move(body)));
           }))
-      .def("decl", [](const Def& def) { return def.decl(); });
+      .def("decl", [](const Def& def) { return def.decl(); })
+      .def("name", [](const Def& def) { return def.name(); });
   py::class_<ClassDef, TreeView>(m, "ClassDef")
       .def(py::init([](const Ident& name, std::vector<Stmt> body) {
         const auto& r = name.range();
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -31,7 +31,7 @@
 from collections import OrderedDict, namedtuple
 
 # These are imported so users can access them from the `torch.jit` module
-from torch._jit_internal import Final, _overload  # noqa: F401
+from torch._jit_internal import Final, _overload, _overload_method  # noqa: F401
 from torch._jit_internal import ignore, export  # noqa: F401
 
 if sys.version_info[0] > 2:
@@ -1903,10 +1903,14 @@ def _compile_function_with_overload(qual_name, impl_fn, overload_decl, overload_
     fn = torch._C._jit_script_compile_overload(qual_name, overload_decl, impl_ast, _rcb, overload_defaults)
     return fn
 
-def _get_overload_decl_and_defaults(func):
+def _check_no_signature(func):
     signature = torch.jit.annotations.get_signature(func)
     if signature is None:
-        raise RuntimeError("Must explicitly add type annotations to overloaded functions: {obj}").format(func)
+        qual_name = _qualified_name(func)
+        raise RuntimeError("Must explicitly add type annotations to overloaded functions: {}".format(qual_name))
+
+def _get_overload_decl_and_defaults(func):
+    _check_no_signature(func)
     return (torch.jit.get_jit_def(func).decl(), get_default_args(func))
 
 def _get_overloads(obj):
diff --git a/torch/jit/_recursive.py b/torch/jit/_recursive.py
@@ -117,19 +117,48 @@ def recursive_script(mod):
         if not _jit_internal.is_ignored_fn(mod.forward):
             methods = ('forward',)
     exported = []
+    overloads = []
     for name in dir(mod):
         item = getattr(mod, name)
         if callable(item):
             if _jit_internal.get_torchscript_modifier(item) is _jit_internal.FunctionModifiers.EXPORT:
                 exported.append(name)
+
+            # builtin functions like repr() in python 2 do not have __module__ defined
+            if hasattr(item, "__module__") and item.__module__ is not None:
+                method_overloads = _jit_internal._get_overloaded_methods(item, mod._get_name())
+                if method_overloads is not None:
+                    overloads.append((item, method_overloads))
+
     methods = methods + tuple(exported)
 
+    overload_name_mappings = dict(getattr(mod, "__overloads__", {}))
+    overload_stubs = []
+
+    for orig_fn, overload_fns in overloads:
+        orig_ast = torch.jit.get_jit_def(orig_fn, self_name="ScriptModule")
+        names = list(map(lambda i: orig_ast.name().name + "__" + str(i), range(len(overload_fns))))
+        overload_name_mappings[orig_ast.name().name] = names
+        for overload_fn, name in zip(overload_fns, names):
+            torch.jit._check_no_signature(overload_fn)
+            over_ast = torch.jit.get_jit_def(overload_fn, self_name="ScriptModule")
+            new_ast = torch._C._replace_overloaded_method_decl(over_ast.decl(), orig_ast, name)
+            _rcb = _jit_internal.createResolutionCallbackFromClosure(orig_fn)
+            overload_stubs.append(torch.jit.ScriptMethodStub(_rcb, new_ast, overload_fn))
+
+    mod.__overloads__ = overload_name_mappings
+
+    # we shouldn't directly compile overloaded methods, just its overloads
+    def ignore_overloaded(method_name):
+        return method_name not in overload_name_mappings
+
     def make_stub(method):
         func = get_function_from_type(type(mod), method)
         return torch.jit.script_method(func, _jit_internal.createResolutionCallbackFromClosure(func))
 
-    stubs = list(map(make_stub, methods))
-    return copy_to_script_module(mod, stubs)
+    filtered_methods = filter(ignore_overloaded, methods)
+    stubs = list(map(make_stub, filtered_methods))
+    return copy_to_script_module(mod, overload_stubs + stubs)
 
 
 def create_method_from_fn(module, fn):
