Rewrite implementation of faithful cpp signatures

This rewrite is as per my comments at #44087 (comment)
I did the rewrite by reverting #44087 and then reimplementing it on top.
You may find it easier to review by diffing against master with only #44087
reverted.

There are two main ideas.

First, we now factor cpp argument processing into two phases operating
on three representations of data:

1. `FunctionSchema` - this is the source from native_functions.yaml
2. `Union[Argument, ThisArgument, TensorOptionsArgument]` - this is
   the arguments after doing some basic semantic analysis to group
   them (for TensorOptions) or identify the this argument (if this
   is a method).  There is only ever one of these per functions.
3. `Union[CppArgument, CppThisArgument, CppTensorOptionsArgument]` -
   this is the arguments after we've elaborated them to C++.  There
   may be multiple of these per actual C++ signature.

You can think of (2) as common processing, whereas (3) bakes in specific
assumptions about whether or not you have a faithful or non-faithful
signature.

Second, we now have CppSignature and CppSignatureGroup representing
the *total* public C++ API signature.  So those dataclasses are what
know how to render definitions/declarations, and you no longer have
to manually type it out in the Functions/TensorMethods codegen.

Here is an exhaustive accounting of the changes.

tools.codegen.api.types

- CppSignature and CppSignatureGroup got moved to tools.codegen.api.types
- Add new CppThisArgument and CppTensorOptionsArguments (modeled off
  of ThisArgument and TensorOptionsArguments) so that we can retain
  high level semantic structure even after elaborating terms with C++
  API information.  Once this is done, we can refine
  CppArgument.argument to no longer contain a ThisArgument (ThisArgument
  is always translated to CppThisArgument.  Note that this doesn't
  apply to TensorOptionsArguments, as those may be expanded or not
  expanded, and so you could get a single CppArgument for 'options')
- Add no_default() functional mutator to easily remove default arguments
  from CppArgument and friends
- Add an explicit_arguments() method to CppArgument and friends to
  extract (flat) argument list that must be explicitly written in the signature.
  This is everything except (Cpp)ThisArgument, and is also convenient
  when you don't care about the extra structure of
  CppTensorOptionsArguments

tools.codegen.api.cpp

- group_arguments is back, and it doesn't send things directly to a
  CppSignatureGroup; instead, it moves us from representation (1) to (2)
  (perhaps it should live in model).  Here I changed my mind from my
  PR comment; I discovered it was not necessary to do classification at
  grouping time, and it was simpler and easier to do it later.
- argument got split into argument_not_this/argument/argument_faithful.
  argument and argument_faithful are obvious enough what they do,
  and I needed argument_not_this as a more refined version of argument
  so that I could get the types to work out on TensorOptionsArguments

tools.codegen.api.dispatcher

- Here we start seeing the payoff.  The old version of this code had a
  "scatter" mode and a "gather" mode.  We don't need that anymore:
  cppargument_exprs is 100% type-directed via the passed in cpp
  arguments.  I am able to write the functions without any reference
  to use_c10_dispatcher

tools.codegen.gen

- Instead of having exprs_str and types_str functions, I moved these to
  live directly on CppSignature, since it seemed pretty logical.
- The actual codegen for TensorMethods/Functions is greatly simplified,
  since (1) all of the heavy lifting is now happening in
  CppSignature(Group) construction, and (2) I don't need to proxy one
  way or another, the new dispatcher translation code is able to handle
  both cases no problem.  There is a little faffing about with ordering
  to reduce the old and new diff which could be removed afterwards.

Here are codegen diffs.  For use_c10_dispatcher: full:

```
+// aten::_cudnn_init_dropout_state(float dropout, bool train, int dropout_seed, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor
 Tensor _cudnn_init_dropout_state(double dropout, bool train, int64_t dropout_seed, const TensorOptions & options) {
-    return _cudnn_init_dropout_state(dropout, train, dropout_seed, optTypeMetaToScalarType(options.dtype_opt()), options.layout_opt(), options.device_opt(), options.pinned_memory_opt());
+    static auto op = c10::Dispatcher::singleton()
+        .findSchemaOrThrow("aten::_cudnn_init_dropout_state", "")
+        .typed<Tensor (double, bool, int64_t, c10::optional<ScalarType>, c10::optional<Layout>, c10::optional<Device>, c10::optional<bool>)>();
+    return op.call(dropout, train, dropout_seed, optTypeMetaToScalarType(options.dtype_opt()), options.layout_opt(), options.device_opt(), options.pinned_memory_opt());
 }
```

Otherwise:

```
+// aten::empty_meta(int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
 Tensor empty_meta(IntArrayRef size, c10::optional<ScalarType> dtype, c10::optional<Layout> layout, c10::optional<Device> device, c10::optional<bool> pin_memory, c10::optional<MemoryFormat> memory_format) {
-    return empty_meta(size, TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory), memory_format);
+    static auto op = c10::Dispatcher::singleton()
+        .findSchemaOrThrow("aten::empty_meta", "")
+        .typed<Tensor (IntArrayRef, const TensorOptions &, c10::optional<MemoryFormat>)>();
+    return op.call(size, TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory), memory_format);
 }
```

Things that I probably did not get right:

- The Union[Argument, TensorOptionsArguments, ThisArgument] and
  the Cpp variants are starting to get a little unwieldy.  Not sure if
  this means I should add a supertype (or at the very least an
  alias); in some cases I do purposely omit one of these from the Union
- Code may not necessarily live in the most logical files.  There isn't
  very much rhyme or reason to it.
- The fields on CppSignature.  They're not very well constrained and
  it will be better if people don't use them directly.
- Disambiguation.  We should do this properly in #44087 and we don't
  need special logic for deleting defaulting for faithful signatures;
  there is a more general story here.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

ghstack-source-id: 5db980cce350fb3f38c75088b7da1b0061b7c717
Pull Request resolved: #45890

tools/codegen/api/cpp.py

@@ -1,9 +1,7 @@
 from tools.codegen.model import *
-from tools.codegen.api.types import TensorOptionsArguments, CppArgument, ThisArgument
+from tools.codegen.api.types import *
 import tools.codegen.local as local
-from typing import Optional, Sequence, Union, Callable, List, Tuple
-import copy
-from dataclasses import dataclass
+from typing import Optional, Sequence, Union, Callable, List
 
 # This file describes the translation of JIT schema to the public C++
 # API, which is what people use when they call functions like at::add.
@@ -197,21 +195,17 @@ def default_expr(d: str, t: Type) -> str:
     return JIT_TO_CPP_DEFAULT.get(d, d)
 
 # Convert an argument into its C++ API form
-def argument(a: Union[Argument, TensorOptionsArguments, ThisArgument]) -> CppArgument:
+
+def argument_not_this(
+    a: Union[Argument, TensorOptionsArguments],
+) -> CppArgument:
     if isinstance(a, Argument):
         return CppArgument(
             type=argument_type(a),
             name=a.name,
             default=default_expr(a.default, a.type) if a.default is not None else None,
             argument=a,
         )
-    elif isinstance(a, ThisArgument):
-        return CppArgument(
-            type=argument_type(a.argument),
-            name="const_cast<Tensor&>(*this)",  # this is an abuse but it's convenient
-            default=None,
-            argument=a,
-        )
     elif isinstance(a, TensorOptionsArguments):
         default = None
         if all(x.default == "None" for x in a.all()):
@@ -227,61 +221,43 @@ def argument(a: Union[Argument, TensorOptionsArguments, ThisArgument]) -> CppArg
     else:
         assert_never(a)
 
-@dataclass(frozen=True)
-class CppSignature:
-    returns: Tuple[Return, ...]
-    arguments: Tuple[Union[Argument, TensorOptionsArguments, ThisArgument], ...]
-
-    def cpp_arguments(self) -> Sequence[CppArgument]:
-        return list(map(argument, self.arguments))
-
-    # Return arguments as a comma separated list, i.e. like they would be in a C++
-    # function signature. Include default values for arguments.
-    def cpp_arguments_str(self, with_defaults: bool) -> str:
-        args_without_this = [argument(a) for a in self.arguments if not isinstance(a, ThisArgument)]
-        if with_defaults:
-            return ', '.join(map(str, args_without_this))
-        else:
-            return ', '.join(map(lambda s: s.str_no_default(), args_without_this))
-
-
-@dataclass(frozen=True)
-class CppSignatureGroup:
-    # arguments contains the arguments for the C++ signature as it is represented
-    # in the JIT schema.
-    signature: CppSignature
-
-    # gathered_signature is an alternative C++ signature in which TensorOptions are
-    # gathered into one TensorOptions object instead of being scattered into
-    # ScalarType, Layout, Device. This is only present for factory operators,
-    # other operators have this set to None. This can be used to generate a
-    # convenience API in the C++ frontend so users can call using TensorOptions objects.
-    gathered_signature: Optional[CppSignature]
-
-    # If it is a factory op, this returns the arguments for the convenience API
-    # that takes TensorOptions. If it is not a factory op and doesn't have
-    # a gathered signature, then this returns the regular signature instead.
-    def signature_prefer_gathered(self) -> CppSignature:
-        if self.gathered_signature is not None:
-            return self.gathered_signature
-        else:
-            return self.signature
+def argument(
+    a: Union[Argument, TensorOptionsArguments, ThisArgument],
+) -> Union[CppSingleArgumentPack, CppThisArgumentPack]:
+    if isinstance(a, ThisArgument):
+        return CppThisArgumentPack(argument=a, type=argument_type(a.argument))
+    else:
+        return CppSingleArgumentPack(argument_not_this(a))
 
+def argument_faithful(
+    a: Union[Argument, TensorOptionsArguments, ThisArgument],
+) -> CppArgumentPack:
+    if isinstance(a, TensorOptionsArguments):
+        return CppTensorOptionsArgumentPack(
+            argument=a,
+            dtype=argument_not_this(a.dtype),
+            layout=argument_not_this(a.layout),
+            device=argument_not_this(a.device),
+            pin_memory=argument_not_this(a.pin_memory),
+        )
+    else:
+        return argument(a)
 
-def signature_group(
-    func: FunctionSchema, *, method: bool = False,
-) -> CppSignatureGroup:
+# NB: this unconditionally groups arguments
+def group_arguments(
+    func: FunctionSchema, *, method: bool
+) -> Sequence[Union[Argument, TensorOptionsArguments, ThisArgument]]:
     args: List[Union[Argument, ThisArgument, TensorOptionsArguments]] = []
+
     args.extend(func.out_arguments)
 
     if method:
         args.extend(ThisArgument(a) if a.name == "self" else a for a in func.arguments)
     else:
         args.extend(func.arguments)
 
-    gathered_args = copy.deepcopy(args)
-
     # group up arguments for tensor options
+
     def pred(name: str, ty: Type) -> Callable[[Argument], bool]:
         return lambda a: a.name == name and a.type in [ty, OptionalType(ty)]
     predicates = [  # order matters
@@ -291,36 +267,22 @@ def pred(name: str, ty: Type) -> Callable[[Argument], bool]:
         pred('pin_memory', Type.parse('bool')),
     ]
 
-    has_tensoroptions_argument = False
     i = 0
     while i < len(func.kwarg_only_arguments):
         # If there is enough space...
         if i <= len(func.kwarg_only_arguments) - len(predicates):
             # And the next len(predicates) arguments look like TensorOptions arguments
             if all(p(a) for p, a in zip(predicates, func.kwarg_only_arguments[i : i + len(predicates)])):
-                has_tensoroptions_argument = True
                 # Group them together as one argument
-                gathered_args.append(TensorOptionsArguments(
+                args.append(TensorOptionsArguments(
                     dtype=func.kwarg_only_arguments[i],
                     layout=func.kwarg_only_arguments[i + 1],
                     device=func.kwarg_only_arguments[i + 2],
                     pin_memory=func.kwarg_only_arguments[i + 3],
                 ))
                 i += len(predicates)
                 continue
-        gathered_args.append(func.kwarg_only_arguments[i])
+        args.append(func.kwarg_only_arguments[i])
         i += 1
 
-    args.extend(func.kwarg_only_arguments)
-
-    if has_tensoroptions_argument:
-        return CppSignatureGroup(
-            signature=CppSignature(arguments=tuple(args), returns=tuple(func.returns)),
-            gathered_signature=CppSignature(arguments=tuple(gathered_args), returns=tuple(func.returns)),
-        )
-    else:
-        assert gathered_args == args
-        return CppSignatureGroup(
-            signature=CppSignature(arguments=tuple(args), returns=tuple(func.returns)),
-            gathered_signature=None,
-        )
+    return args

tools/codegen/api/dispatcher.py

@@ -1,11 +1,10 @@
 from tools.codegen.model import *
 
-from tools.codegen.api.types import CppArgument, DispatcherExpr, TensorOptionsArguments, \
-    DispatcherArgument, ThisArgument, LegacyDispatcherArgument
-from tools.codegen.api import cpp
+from tools.codegen.api.types import *
+import tools.codegen.api.cpp as cpp
 import tools.codegen.api.legacy_dispatcher as legacy_dispatcher
 import tools.codegen.local as local
-from enum import Enum
+
 import itertools
 from typing import Sequence, Optional
 
@@ -75,73 +74,86 @@ def arguments(func: FunctionSchema) -> Sequence[DispatcherArgument]:
             for la in legacy_dispatcher.arguments(func)
         ]
 
-# TODO GATHER is only needed for non-c10-full ops, remove later.
-ProcessTensoroptions = Enum('ProcessTensoroptions', ('GATHER', 'SCATTER', 'PASS_THROUGH'))
-
-
 # Given a set of CppArguments in scope, return a sequence of dispatcher
 # expressions that translate the cpp API into dispatcher API
-def cppargument_exprs(a: CppArgument,
-                      *,
-                      tensor_options: Optional[CppArgument],
-                      process_tensoroptions: ProcessTensoroptions = ProcessTensoroptions.PASS_THROUGH
-                      ) -> Sequence[DispatcherExpr]:
-    if isinstance(a.argument, TensorOptionsArguments):
-        if process_tensoroptions == ProcessTensoroptions.SCATTER:
-            ta = a.argument
+#
+# WARNING: This is unsound if you pass it CppArgument when you were
+# supposed to pass it CppTensorOptionsArguments, it will directly
+# translate device to device, which will give you the wrong signature
+# for dispatcher.  If Argument "knew" that it was part of a
+# TensorOptions that would help us dynamically test for this case
+def cppargument_exprs(
+    a: CppArgumentPack,
+    *, tensor_options: Optional[CppArgument]
+) -> Sequence[DispatcherExpr]:
+    if isinstance(a, CppSingleArgumentPack):
+        if isinstance(a.this.argument, TensorOptionsArguments):
+            if local.use_c10_dispatcher() is UseC10Dispatcher.full:
+                # Scatter
+                ta = a.this.argument
+                name = a.this.name
+                return [
+                    DispatcherExpr(type=argument_type(ta.dtype), expr=f'optTypeMetaToScalarType({name}.dtype_opt())'),
+                    DispatcherExpr(type=argument_type(ta.layout), expr=f'{name}.layout_opt()'),
+                    DispatcherExpr(type=argument_type(ta.device), expr=f'{name}.device_opt()'),
+                    DispatcherExpr(type=argument_type(ta.pin_memory), expr=f'{name}.pinned_memory_opt()'),  # weird discrep
+                ]
+            else:
+                # No-op
+                return [DispatcherExpr(type='const TensorOptions &', expr=a.this.name)]
+        elif isinstance(a.this.argument, Argument):
+            if a.this.name == 'memory_format' and \
+                    tensor_options is not None and \
+                    local.use_c10_dispatcher() is UseC10Dispatcher.full:
+                return [DispatcherExpr(
+                    type=argument_type(a.this.argument),
+                    expr=f'c10::impl::check_tensor_options_and_extract_memory_format({tensor_options.name}, {a.this.name})')
+                ]
+            else:
+                return [DispatcherExpr(type=argument_type(a.this.argument), expr=a.this.name)]
+        else:
+            assert_never(a.this.argument)
+    elif isinstance(a, CppTensorOptionsArgumentPack):
+        if local.use_c10_dispatcher() is UseC10Dispatcher.full:
+            # No-op
             return [
-                DispatcherExpr(type=argument_type(ta.dtype), expr=f'optTypeMetaToScalarType({a.name}.dtype_opt())'),
-                DispatcherExpr(type=argument_type(ta.layout), expr=f'{a.name}.layout_opt()'),
-                DispatcherExpr(type=argument_type(ta.device), expr=f'{a.name}.device_opt()'),
-                DispatcherExpr(type=argument_type(ta.pin_memory), expr=f'{a.name}.pinned_memory_opt()'),  # weird discrep
+                expr
+                for sub_a in a.explicit_arguments()  # NB: don't really care about explicitness here
+                for expr in cppargument_exprs(CppSingleArgumentPack(sub_a), tensor_options=tensor_options)
             ]
-        elif process_tensoroptions == ProcessTensoroptions.GATHER:
-            return [
-                DispatcherExpr(
-                    type='const TensorOptions &',
-                    expr="TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory)")]
         else:
-            assert process_tensoroptions == ProcessTensoroptions.PASS_THROUGH
-            return [DispatcherExpr(type='const TensorOptions &', expr=a.name)]
-    elif isinstance(a.argument, ThisArgument):
-        return [DispatcherExpr(type=argument_type(a.argument.argument), expr=a.name)]
-    elif isinstance(a.argument, Argument):
-        if a.name == 'memory_format' and tensor_options is not None and local.use_c10_dispatcher() is UseC10Dispatcher.full:
+            # Gather
             return [DispatcherExpr(
-                type=argument_type(a.argument),
-                expr=f'c10::impl::check_tensor_options_and_extract_memory_format({tensor_options.name}, {a.name})')
-            ]
-        else:
-            return [DispatcherExpr(type=argument_type(a.argument), expr=a.name)]
+                type='const TensorOptions &',
+                expr=f'TensorOptions().dtype({a.dtype.name}).layout({a.layout.name})'
+                     f'.device({a.device.name}).pinned_memory({a.pin_memory.name})',
+            )]
+    elif isinstance(a, CppThisArgumentPack):
+        return [DispatcherExpr(
+            type=a.type,
+            expr='const_cast<Tensor&>(*this)',
+        )]
     else:
-        assert_never(a.argument)
+        assert_never(a)
 
-def cpparguments_exprs(args: Sequence[CppArgument], process_tensoroptions: ProcessTensoroptions) -> Sequence[DispatcherExpr]:
-    tensor_options = next((a for a in args if isinstance(a.argument, TensorOptionsArguments)), None)
-    return [r for a in args for r in cppargument_exprs(a,
-                                                       tensor_options=tensor_options,
-                                                       process_tensoroptions=process_tensoroptions)]
+def cpparguments_exprs(args: Sequence[CppArgumentPack]) -> Sequence[DispatcherExpr]:
+    tensor_options = next(
+        (a.this for a in args if isinstance(a, CppSingleArgumentPack) and
+            isinstance(a.this.argument, TensorOptionsArguments)),
+        None
+    )
+    return [r for a in args for r in cppargument_exprs(a, tensor_options=tensor_options)]
 
 # I don't think this is entirely sound, but it should be reasonably
 # close
 def legacydispatcherarguments_exprs(args: Sequence[LegacyDispatcherArgument]) -> Sequence[DispatcherExpr]:
-    if local.use_c10_dispatcher() is UseC10Dispatcher.full:
-        process_tensoroptions = ProcessTensoroptions.SCATTER
-    else:
-        process_tensoroptions = ProcessTensoroptions.PASS_THROUGH
-    return cpparguments_exprs([CppArgument(type=a.type,
-                                           name=a.name,
-                                           default=None,
-                                           argument=a.argument) for a in args],
-                              process_tensoroptions=process_tensoroptions)
+    return cpparguments_exprs([
+        CppSingleArgumentPack(CppArgument(type=a.type, name=a.name, default=None, argument=a.argument))
+        for a in args
+    ])
 
 def exprs(args: Sequence[DispatcherArgument]) -> Sequence[DispatcherExpr]:
-    if local.use_c10_dispatcher() is UseC10Dispatcher.full:
-        process_tensoroptions = ProcessTensoroptions.SCATTER
-    else:
-        process_tensoroptions = ProcessTensoroptions.PASS_THROUGH
-    return cpparguments_exprs([CppArgument(type=a.type,
-                                           name=a.name,
-                                           default=None,
-                                           argument=a.argument) for a in args],
-                              process_tensoroptions=process_tensoroptions)
+    return cpparguments_exprs([
+        CppSingleArgumentPack(CppArgument(type=a.type, name=a.name, default=None, argument=a.argument))
+        for a in args
+    ])

tools/codegen/api/legacy_dispatcher.py

@@ -71,6 +71,4 @@ def argument(a: Union[Argument, ThisArgument, TensorOptionsArguments]) -> Legacy
         assert_never(a)
 
 def arguments(func: FunctionSchema) -> Sequence[LegacyDispatcherArgument]:
-    signature_group = cpp.signature_group(func)
-    args = signature_group.signature_prefer_gathered().arguments
-    return list(map(argument, args))
+    return list(map(argument, cpp.group_arguments(func, method=False)))