Update on "Making ops c10 full: optional out arguments"

smessmer · smessmer · commit 114aa3cf63ba · 2020-12-15T14:50:06.000-08:00
We have some (but very few) ops that take optional out arguments `Tensor(a!)? out`. This PR makes them non-optional mandatory arguments and enables c10-fullness for them. There is only a very small number of ops affected by this. Putting this up for discussion. Alternatives considered: If we keep them optional, we run into lots of issues in the dispatcher. We have to decide what the dispatcher calling convention for this argument type should be. 1) If we keep passing them in as `Tensor&` arguments and return them as `tuple<Tensor&, Tensor&, Tensor&>`, so basically same as currently, then the schema inference check will say "Your kernel function got inferred to have a `Tensor` argument but your native_functions.yaml declaration says `Tensor?`. This is a mismatch, you made an error". We could potentially disable that check, but that would open the door for real mistakes to not be reported anymore in the future. This sounds bad. 2) If we change them to a type that schema inference could differentiate from `Tensor`, say we pass them in as `const optional<Tensor>&` and return them as `tuple<const optional<Tensor>&, const optional<Tensor>&, const optional<Tensor>&>`, then our boxing logic fails because it can't recognize those as out overloads anymore and shortcut the return value as it is doing right now. We might be able to rewrite the boxing logic, but that could be difficult and could easily develop into a rabbit hole of having to clean up `Tensor&` references throughout the system where we use them. Furthermore, having optional out arguments in C++ doesn't really make sense. the C++ API puts them to the front of the argument list, so you can't omit them anyways when calling an op. You would be able to omit them when calling from Python with out kwargs, but not sure if we want that discrepancy between the c++ and python API. Differential Revision: [D25422197](https://our.internmc.facebook.com/intern/diff/D25422197/) [ghstack-poisoned]
diff --git a/c10/util/TypeCast.h b/c10/util/TypeCast.h
@@ -170,3 +170,5 @@ To checked_convert(From f, const char* name) {
 }
 
 }  // namespace c10
+
+// Trigger tests for D25440771. TODO: Remove this line any time you want.
diff --git a/c10/util/complex.h b/c10/util/complex.h
@@ -262,7 +262,7 @@ struct alignas(sizeof(T) * 2) complex {
     return real() || imag();
   }
 
-  constexpr T real() const {
+  C10_HOST_DEVICE constexpr T real() const {
     return real_;
   }
   constexpr void real(T value) {
diff --git a/docs/source/fx.rst b/docs/source/fx.rst
@@ -0,0 +1,30 @@
+.. currentmodule:: torch.fx
+
+torch.fx
+=============
+
+Overview
+--------
+.. automodule:: torch.fx
+
+
+API Reference
+-------------
+
+.. autofunction:: torch.fx.symbolic_trace
+
+.. autoclass:: torch.fx.GraphModule
+  :members:
+
+  .. automethod:: __init__
+
+.. autoclass:: torch.fx.Graph
+  :members:
+
+  .. automethod:: __init__
+
+.. autoclass:: torch.fx.Node
+  :members:
+
+.. autoclass:: torch.fx.Tracer
+  :members:
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -61,6 +61,7 @@ Features described in this documentation are classified by release status:
    torch.distributions <distributions>
    torch.fft <fft>
    futures
+   fx
    torch.hub <hub>
    torch.jit <jit>
    torch.linalg <linalg>
diff --git a/test/run_test.py b/test/run_test.py
@@ -40,6 +40,7 @@
     'distributed/test_distributed_spawn',
     'distributions/test_constraints',
     'distributions/test_distributions',
+    'test_dispatch',
     'test_expecttest',
     'test_foreach',
     'test_indexing',
diff --git a/test/test_dispatch.py b/test/test_dispatch.py
@@ -3,7 +3,6 @@
 
 from collections import namedtuple
 import itertools
-import unittest
 import re
 
 # TODO: Expand the dispatcher API to be a generic API for interfacing with
@@ -256,7 +255,14 @@ def test_def_impl_schema_mismatch(self):
             # m.impl("foo", [](const Tensor & x) { return x })
             lambda m: m.impl_t_t("foo"),
         ], expect_raises=True).state
-        self.assertExpectedInline(state, '''In registration for test::foo: expected schema of operator to be "test::foo(Tensor x, Tensor y) -> (Tensor)" (registered at /dev/null:0), but got inferred schema "(Tensor _0) -> (Tensor _0)" (impl_t_t). The number of arguments is different. 2 vs 1.''')  # noqa
+        self.assertExpectedInline(state, '''\
+Inferred operator schema for a C++ kernel function doesn't match the expected function schema.
+  operator: test::foo
+  expected schema: test::foo(Tensor x, Tensor y) -> (Tensor)
+    registered at /dev/null:0
+  inferred schema: (Tensor _0) -> (Tensor _0)
+    impl_t_t
+  reason: The number of arguments is different. 2 vs 1.''')  # noqa
 
     def test_def_with_inference(self):
         state = self.commute("foo", [
@@ -656,17 +662,19 @@ def test_computed_table_with_cpu_autograd_math_defaultbackend(self):
 AutogradXLA: fn_autograd [autograd kernel]
 ''')
 
-    # Can't do this yet for BC reasons
-    @unittest.expectedFailure
     def test_multiple_def_error(self):
-        state = self.commute("foo", [
+        ops = [
             # m.def("foo(Tensor x, Tensor y) -> Tensor")
             lambda m: m.def_("foo(Tensor x, Tensor y) -> Tensor"),
             # m.def("foo(Tensor x, Tensor y) -> Tensor")
             lambda m: m.def_("foo(Tensor x, Tensor y) -> Tensor"),
-        ], expect_raises=True).state
-        # TODO: fill in the error message here
-        # self.assertExpectedInline(state, '''''')
+        ]
+        self.assertExpectedInline(
+            self.commute("foo", ops, expect_raises=True).state,
+            '''Tried to register an operator (test::foo(Tensor x, Tensor y) -> (Tensor)) with the same name and overload '''
+            '''name multiple times. Each overload's schema should only be registered with a single call to def(). '''
+            '''Duplicate registration: registered at /dev/null:0. Original registration: registered at /dev/null:0'''
+        )
 
     def test_def_with_explicit_alias(self):
         state = self.commute("foo", [
@@ -683,52 +691,22 @@ def test_def_with_explicit_alias(self):
 alias analysis kind: PURE_FUNCTION
 ''')
 
-    # TODO: get rid of this test when multiple defs are wrong
-    def test_multiple_def_schema_mismatch(self):
-        # error message is order dependent
-        ops = [
-            # m.def("foo(Tensor x, Tensor y) -> Tensor")
-            lambda m: m.def_("foo(Tensor x, Tensor y) -> Tensor"),
-            # m.def("foo(Tensor x) -> Tensor")
-            lambda m: m.def_("foo(Tensor x) -> Tensor"),
-        ]
-        self.assertExpectedInline(
-            self.commute("foo", ops, ctor_order=(0, 1), expect_raises=True).state,
-            '''Tried to register multiple operators with the same name and the same overload name but different schemas: test::foo(Tensor x) -> (Tensor) (registered at /dev/null:0) vs test::foo(Tensor x, Tensor y) -> (Tensor) (registered at /dev/null:0)'''  # noqa
-        )
-        self.assertExpectedInline(
-            self.commute("foo", ops, ctor_order=(1, 0), expect_raises=True).state,
-            '''Tried to register multiple operators with the same name and the same overload name but different schemas: test::foo(Tensor x, Tensor y) -> (Tensor) (registered at /dev/null:0) vs test::foo(Tensor x) -> (Tensor) (registered at /dev/null:0)'''  # noqa
-        )
-
     def test_multiple_def_alias_defaulting(self):
-        # TODO: should be an error in both directions soon
         ops = [
             # m.def(torch::schema("foo(Tensor x) -> Tensor",
             #                     c10::AliasAnalysisKind::PURE_FUNCTION))
             lambda m: m.def_("foo(Tensor x) -> Tensor", alias="PURE_FUNCTION"),
             # RegisterOperators().op("foo(Tensor x) -> Tensor")
             lambda m: m.def_legacy("foo(Tensor x) -> Tensor"),
         ]
-        state = self.commute("foo", ops, ctor_order=(0, 1)).state
         self.assertExpectedInline(
-            state,
-            '''\
-name: test::foo
-schema: test::foo(Tensor x) -> (Tensor)
-debug: registered at /dev/null:0
-alias analysis kind: PURE_FUNCTION
-'''
+            self.commute("foo", ops, expect_raises=True).state,
+            '''Tried to register an operator (test::foo(Tensor x) -> (Tensor)) with the same name and overload '''
+            '''name multiple times. Each overload's schema should only be registered with a single call to def(). '''
+            '''Duplicate registration: registered at /dev/null:0. Original registration: registered at /dev/null:0'''
         )
-        # NB: When run with ctor order (1, 0), the destructors are NOT
-        # COMMUTATIVE.  THIS IS A BUG, however we are purposely leaving the bug
-        # in as it is very benign (only leaves us in a bad state during
-        # destruction, when no useful work is being done), will be fixed when we
-        # make alias defaulting a hard error, and is very nontrivial to fix
-        # prior to that.
 
     def test_multiple_def_alias_mismatch(self):
-        # error message is order dependent
         ops = [
             # m.def(torch::schema("foo(Tensor x) -> Tensor",
             #                     c10::AliasAnalysisKind::PURE_FUNCTION))
@@ -738,12 +716,10 @@ def test_multiple_def_alias_mismatch(self):
             lambda m: m.def_("foo(Tensor x) -> Tensor", alias="CONSERVATIVE"),
         ]
         self.assertExpectedInline(
-            self.commute("foo", ops, ctor_order=(0, 1), expect_raises=True).state,
-            '''Tried to define the schema for test::foo with different alias analysis kinds: PURE_FUNCTION (registered at /dev/null:0) vs CONSERVATIVE (registered at /dev/null:0)'''  # noqa
-        )
-        self.assertExpectedInline(
-            self.commute("foo", ops, ctor_order=(1, 0), expect_raises=True).state,
-            '''Tried to define the schema for test::foo with different alias analysis kinds: CONSERVATIVE (registered at /dev/null:0) vs PURE_FUNCTION (registered at /dev/null:0)'''  # noqa
+            self.commute("foo", ops, expect_raises=True).state,
+            '''Tried to register an operator (test::foo(Tensor x) -> (Tensor)) with the same name and overload '''
+            '''name multiple times. Each overload's schema should only be registered with a single call to def(). '''
+            '''Duplicate registration: registered at /dev/null:0. Original registration: registered at /dev/null:0'''  # noqa
         )
 
     def test_multiple_fallback(self):
@@ -754,7 +730,8 @@ def test_multiple_fallback(self):
         except RuntimeError as e:
             self.assertExpectedInline(
                 str(e),
-                '''Tried to register multiple backend fallbacks for the same dispatch key XLA; previous registration registered at /dev/null:0, new registration registered at /dev/null:0'''  # noqa
+                '''Tried to register multiple backend fallbacks for the same dispatch key XLA; previous registration '''
+                '''registered at /dev/null:0, new registration registered at /dev/null:0'''  # noqa
             )
         else:
             self.assertTrue(False)
diff --git a/third_party/tensorpipe b/third_party/tensorpipe
@@ -1 +1 @@
-Subproject commit 5381c57ba923481ffaf7c40f9acc7f164ded887f
+Subproject commit 82a114882e21b176916e2f12a7b566af3d63df71
diff --git a/third_party/tensorpipe.BUILD b/third_party/tensorpipe.BUILD
@@ -93,13 +93,7 @@ TENSORPIPE_HEADERS = glob([
 TENSORPIPE_BASE_SRCS = glob([
     "tensorpipe/*.cc",
     "tensorpipe/channel/*.cc",
-    "tensorpipe/common/address.cc",
-    "tensorpipe/common/epoll_loop.cc",
-    "tensorpipe/common/error.cc",
-    "tensorpipe/common/fd.cc",
-    "tensorpipe/common/ibv.cc",
-    "tensorpipe/common/socket.cc",
-    "tensorpipe/common/system.cc",
+    "tensorpipe/common/*.cc",
     "tensorpipe/core/*.cc",
     "tensorpipe/transport/*.cc",
     "tensorpipe/util/*/*.cc",
@@ -113,10 +107,7 @@ TENSORPIPE_SRCS = TENSORPIPE_BASE_SRCS + glob([
 ])
 
 TENSORPIPE_SRCS_CUDA = TENSORPIPE_SRCS + glob([
-    "tensorpipe/common/cuda_loop.cc",
-    "tensorpipe/channel/cuda_basic/*.cc",
     "tensorpipe/channel/cuda_ipc/*.cc",
-    "tensorpipe/channel/cuda_xth/*.cc",
 ])
 
 cc_library(
diff --git a/torch/fx/graph.py b/torch/fx/graph.py
@@ -415,7 +415,7 @@ def call_method(self,
                     type_expr: Optional[Any] = None) -> Node:
         """
         Insert a ``call_method`` ``Node`` into the ``Graph``. A ``call_method`` node
-        represents a call to a given method on the 0th element of `args.
+        represents a call to a given method on the 0th element of ``args``.
 
         Args:
 
@@ -756,9 +756,9 @@ def lint(self, root : Optional[torch.nn.Module] = None):
         """
         Runs various checks on this Graph to make sure it is well-formed. In
         particular:
-            - Checks Nodes have correct ownership (owned by this graph)
-            - Checks Nodes appear in topological order
-            - If ``root`` is provided, checks that targets exist in ``root``
+        - Checks Nodes have correct ownership (owned by this graph)
+        - Checks Nodes appear in topological order
+        - If ``root`` is provided, checks that targets exist in ``root``
 
         Args:
 
diff --git a/torch/fx/symbolic_trace.py b/torch/fx/symbolic_trace.py
@@ -66,6 +66,7 @@ def create_arg(self, a: Any) -> 'Argument':
             * For a Parameter, emit a ``get_attr`` node referring to that Parameter
             * For a non-Parameter Tensor, store the Tensor away in a special
               attribute referring to that attribute.
+
         This method can be overridden to support more types.
 
         Args:
@@ -126,6 +127,7 @@ def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> boo
         via this parameter.
 
         Args:
+
             m (Module): The module being queried about
             module_qualified_name (str): The path to root of this module. For example,
                 if you have a module hierarchy where submodule ``foo`` contains
@@ -163,7 +165,7 @@ def call_module(self, m: torch.nn.Module, forward: Callable[..., Any], args : Tu
         This method can be overridden to--for example--create nested traced
         GraphModules, or any other behavior you would want while tracing across
         ``Module`` boundaries.
-         ``Module`` boundaries.
+        ``Module`` boundaries.
 
         Args:
 
@@ -186,8 +188,8 @@ def call_module(self, m: torch.nn.Module, forward: Callable[..., Any], args : Tu
     def create_args_for_root(self, root_fn, is_module):
         """
         Create ``placeholder`` nodes corresponding to the signature of the ``root``
-        Module. This method introspects ``root``'s signature and emits those
-        nodes accordingly, also supporting *args and **kwargs.
+        Module. This method introspects root's signature and emits those
+        nodes accordingly, also supporting ``*args`` and ``**kwargs``.
         """
         # In some cases, a function or method has been decorated with a wrapper
         # defined via ``functools.wraps``. In this case, the outer code object

Original file line number	Diff line number	Diff line change
`@@ -170,3 +170,5 @@ To checked_convert(From f, const char* name) {`
`170`	`170`	`}`
`171`	`171`
`172`	`172`	`} // namespace c10`
	`173`	`+`
	`174`	`+// Trigger tests for D25440771. TODO: Remove this line any time you want.`
Original file line number	Diff line number	Diff line change
`@@ -262,7 +262,7 @@ struct alignas(sizeof(T) * 2) complex {`
`262`	`262`	`return real() \|\| imag();`
`263`	`263`	`}`
`264`	`264`
`265`		`- constexpr T real() const {`
	`265`	`+ C10_HOST_DEVICE constexpr T real() const {`
`266`	`266`	`return real_;`
`267`	`267`	`}`
`268`	`268`	`constexpr void real(T value) {`