Rewrite implementation of faithful cpp signatures

[ezyang]

Stack from ghstack:

• Refactor dispatcher and native to use Signature structure. #45990 Refactor dispatcher and native to use Signature structure.
• Remove unnecessary byte-for-byte compatibility code that is not needed. #45975 Remove unnecessary byte-for-byte compatibility code that is not needed.
• Rename legacy_dispatcher to native. #45974 Rename legacy_dispatcher to native.
• Reorder dispatcher/legacy_dispatcher types #45973 Reorder dispatcher/legacy_dispatcher types
• Add NativeFunctionGroup #45918 Add NativeFunctionGroup
• Rewrite implementation of faithful cpp signatures #45890 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 Add faithful C++ API #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

Differential Revision: D24144035

[ezyang]

cc @peterbell10 you might find this interesting, though in the end I didn't do the advanced disambiguation algorithm

[peterbell10]

Something worth noting, with either of the algorithms proposed in #45666 the overload "disambiguation point" will be the first dtype argument. So layout, device and pin_memory would keep their defaults.

That doesn't actually compile though, because a ScalarType object is implicitly convertible to both TensorOptions and c10::optional<ScalarType>. Since user-defined conversions always have equal priority, there is an ambiguity.

[ezyang]

This is a good point. So do it properly we need to teach codegen about conversions and then consider those when looking for ambiguity.

[peterbell10]

It's complicated because removing the default from layout would remove the ScalarType ambiguity, but also makes it impossible to call with just a c10::optional<ScalarType>.

Two solutions I can think of that wouldn't remove any call signatures:

• make the TensorOptions constructor explicit. If you pass just a ScalarType it selects the faithful overload without ambiguity.
• detect implicit conversions and add new overloads for those types. So, here a new ScalarType overload would be created that just picks one operator and calls it unambiguously. This could be a more general solution, maybe even solving issues like Libtorch tensor.std(0) gets confused about which overload to use (expected vector, not scalar) #40287. However, would be far more compilcated.

For TensorOptions faithful signatures, we could also treat it as one argument and remove all the defaults like you do here.

[ezyang]

So maybe it is best to keep the default stripping here an orthogonal mechanism from disambiguation. The faithful cpp signature is not really intended for users, it's more a convenience for other users who want to be able to uniformly access the API. I will say that this is kind of displeasing though :(

To answer your questions directly:

make the TensorOptions constructor explicit. If you pass just a ScalarType it selects the faithful overload without ambiguity.

This is not great, because the one of the UX affordances of TensorOptions is if you only provide a ScalarType, you can provide it by itself, without having to wrap it in TensorOptions.

detect implicit conversions and add new overloads for those types. So, here a new ScalarType overload would be created that just picks one operator and calls it unambiguously.

Yes I agree this would be complicated.

[codecov]

Codecov Report

Merging #45890 into gh/ezyang/847/base will not change coverage.
The diff coverage is n/a.

@@                 Coverage Diff                 @@
##           gh/ezyang/847/base   #45890   +/-   ##
===================================================
  Coverage               68.28%   68.28%           
===================================================
  Files                     410      410           
  Lines                   53468    53468           
===================================================
  Hits                    36513    36513           
  Misses                  16955    16955           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8de9aa1...c419b70. Read the comment docs.

[smessmer]

I would add a comment here that this removes the implicit this argument. Maybe we should also come up with a better name, since "grouped" vs "ungrouped" doesn't catch the this argument handling, but I don't have an idea right now. It's basically pre-transformation and post-transformation arguments, but that doesn't help with naming. Note that in the same spirit, the CppArgument type returned from cpp_ungrouped_arguments should probably be a different type than the CppArgument returned from cpp_grouped_arguments, since the latter are pre-transformation.

[smessmer]

hm this doesn't specify if it returns the grouped or the ungrouped declaration. I think this points to general issues with the data structure representation. We should probably make that difference and the transformation that is happening here more clear. Maybe have two CppSignature classes, one for grouped and one for ungrouped?

After reading through the code below, it seems that you have one CppSignature class but two instances, one for grouped and one for ungrouped. And in the ungrouped instance, the _cpp_grouped_arguments member holds the ungrouped arguments. This is very confusing.

[ezyang]

This is the correct interpretation. I'm happy to rename the _cpp_grouped_arguments member, the name here is mostly based off of what the types say is possible. I agree it is possible to have two CppSignature classes, one for grouped and another for ungrouped, but this will increase the boilerplate. I can give it a try though if you want.

[smessmer]

ok after some offline discussion, it turned out I misunderstood the terms in your PR. I thought "grouped signature" == "non-faithful signature" because the non-faithful signature groups tensor options arguments into one TensorOptions object...but actually, non-faithful signatures don't even use the grouping logic and just have one CppArgument for each argument (one of them being a TensorOptions). Only for faithful signatures, the difference between a "grouped signature" and "ungrouped signature" comes into play and both always represent the faithful signature, just at different abstraction layers.

[ezyang]

@smessmer I pushed an update at 729aea3 I introduced a new term ArgumentPack for the structured version of CppArgument, and I removed all of the grouped/ungrouped nomenclature (opting for packs instead). I also distinguished CppSingleArgumentPack and CppArgument and that also solved a type safety problem with some APIs.

[smessmer]

you don't need this anymore, right? CppArgument is the low level data structure that was already unpacked.

[smessmer]

yay :) thanks

[smessmer]

Does this work? Looks like you removed the defaults handling - if you don't have a faithful signature (i.e. faithful and non-faithful are identical), then the declaration must have defaults. If faithful and non-faithful differ, then the non-faithful must have defaults and the faithful must not have defaults. Breaking those rules breaks bc.

[ezyang]

This works because I moved the defaults handling into the construction of the signature groups.

        if faithful:
            # Faithful signatures will ungroup arguments into argument
            # packs.
            #
            # After this, manually do overload disambiguation, by
            # dropping defaults from the faithful signature.  In
            # principle, we should be able to do this at some later
            # point in time with other overload disambiguation
            argument_packs = tuple(
                cpp.argument_faithful(a).no_default() for a in arguments
            )   
        else:   
            argument_packs = tuple(
                cpp.argument(a) for a in arguments
            )  

I think handling it at construction time is better, because it preserves the invariant "a CppSignature knows exactly how it should be rendered in C++"

[smessmer]

How did you manage to reduce the 2x2=4 branches down to 2? Last time we talked about this we came to the conclusion that even if you remove the proxies, you still need 4 branches. because you have 2 frontend signatures and 2 dispatcher signatures and you need to map from each to each.

Btw as a side note: Removing the proxies could have side implications because you now store the OperatorHandle singleton twice (once per function) instead of just once in the non-proxy function.

[ezyang]

How did you manage to reduce the 2x2=4 branches down to 2? Last time we talked about this we came to the conclusion that even if you remove the proxies, you still need 4 branches. because you have 2 frontend signatures and 2 dispatcher signatures and you need to map from each to each.

I need 4 branches but it doesn't happen here, it happens at cppargument_exprs, see the "Scatter", "No op", "No op" and "Gather" cases. Because I organized things with argument packs, the decision between the branches can be done entirely in a data directed fashion.

Btw as a side note: Removing the proxies could have side implications because you now store the OperatorHandle singleton twice (once per function) instead of just once in the non-proxy function.

Yes. I decided that besides a mild code size inflation this wouldn't really matter, and simplifying the codegen was more worth it.

[dr-ci]

💊 CI failures summary and remediations

As of commit c419b70 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 5 times.

[bdhirsh]

nit: no need for the f string

[bdhirsh]

never mind, just realized that was fixed in a later commit

[ezyang]

@smessmer All resolved. Are there other parts of the PR which are still conceptually unclear?

[smessmer]

thanks

[bhosmer]

This LGTM - most of the inline commentary is riffing on what would be a fairly sizable refactor, so feel free to disregard - I think it might be a win, but it'd be nontrivial and the current setup is nice. (Mostly I worry about future entropy.)

[bhosmer]

Given the amount of type sniffing that's already being done in the callers, my impulse is to suggest splitting this function into (say) ordinary_argument(Argument) and tensor_options_arguments(TensorOptionsArguments), and coalesce the type sniffing at the caller level. It feels simpler and more precise, e.g. avoiding the bagginess of the assert_never here.

Not a major point, and if the advantage of the current setup becomes clear on further reading I'll try to remember to come back and delete this 😁

[ezyang]

Hmm. Still thinking about the overall comment, but avoiding assert_never is not something you should specifically go for, all it is, is a way of making the typechecker check that the type sniffing (I'll call them pattern matches, because that's the style here) is exhaustive. It's not an antipattern to use it and we should use it as much as possible when we do isinstance style refinements. So for example after your proposed refactor, you will still do an assert_never test in the caller pattern match.

[ezyang]

Re the broader comment, I don't think I really care too much either way. The code is written in its current form because I was trying to minimize the number of top level functions I needed, and so if I write the function to take a Union I only need to write one function (versus two, and also needing to do pattern matching at the caller.)

The beautiful thing about using mypy is that there is almost no precision difference between the two different schemes: the only benefit to the more factored representation is you get to skip an isinstance check when you statically know you only have an Argument or TensorOptionsArgument, but we don't really care about this kind of microoptimization here.

[bhosmer]

Oh of course, it's not about avoiding assert_never per se, more about avoiding partial functions even at the cost of a marginally higher total number of functions (though not at the cost of redundant code). The thought was that since all callers were total over the supertype, cracking this helper avoid the partial function and hence the assert_never. (But I can't say I verified that this factoring would definitely work.)

[ezyang]

avoiding partial functions

There isn't a partial function here! assert_never doesn't mean "this function is partial", it means, "I am doing a case match on a Union and this case match is exhaustive".

[bhosmer]

I know it's not part of this PR, but I'm not sure there's any reason why predicates should be local - pulling it out to a global TENSOR_OPTIONS_PREDS or whatever would call more attention to its significance and also avoid rebuilding it every time thru here

[ezyang]

Not making a judgment call on if it's right or not, but I made it local because I believed there was only one use site, so might as well keep it close and encapsulated.

[bhosmer]

I think my impulse here comes from a relentless compulsion to pull everything that's essentially part of the data model out of the code and into the global scope.

[bhosmer]

Another straight high-level observation so take it fwiw, but here it looks like you could do the thing of "tight data model with convenience projection" by using a dataclass representation that didn't have the space for (e.g.) multiple ThisArguments, misplaced ThisArgument, multiple TensorOptionsArguments instances, etc., and equipping it with a lightweight view method that clients could use to produce a Sequence[...] on demand.

Might be a way to leverage the correct-by-construction goodness of the core data model a bit further downstream. Or it might be overkill :P I haven't looked at what happened to CppSignatureGroup yet so not sure exactly how this observation maps onto the larger design, just noting the thought as I go.

[ezyang]

Yes, I've been thinking about this too. I agree with this proposal. Update: I agree with a limited version of this proposal, which is that grouping should happen during model construction. Extra constraints are more tricky for reasons I listed below.

Some planning notes: group_arguments behaves differently depending on if method is true or false. If we want to make it not do that we will have to unconditionally translate self into ThisArgument (really should call it SelfArgument now), and then delay handling of method kwarg until when we are translating to c++ arguments. This should all work out though.

[bhosmer]

I think you've bought some dynamic complexity at this level by factoring the transformation from CppSignature to Sequence[DispatcherExpr] through a Sequence[CppArgumentPack]. If the builder had access to information CppSignature had (or could have) about the signature as a whole, it could make direct calls to smaller functions defined over CppArgumentPacks subtypes.

I don't know without trying it whether this would be a win, but I mention it bc it feels like the current setup might be falling prey to that thing that happens when trans-sequence structural info gets pushed down into sequence elements, and map/flatMap element handlers wind up needing to reconstruct chunks of the larger structural situation from context.

[ezyang]

I think #45890 (comment) addresses this

[bhosmer]

See there - you can definitely avoid the antipattern I'm mentioning here without sacrificing the expressiveness you need to track self and tensor options that show up in different places. But I wouldn't suggest retroactively tearing up the sidewalk at this point.

[bhosmer]

I'm not sure I get "they're always packing" (even with the analogy 😁 )

[ezyang]

You are probably confused because I meant to write "they're always unpacked"

[bhosmer]

must've been it 😁

[bhosmer]

Just to recap the common theme from a few of the preceding comments - given that representing a signature as (so to speak) (arg | this_arg | tensor_opts_arg)* is really very loose, I'd be tempted to do away with this union and see what the architecture looks like after fixing all the places where it currently appears. Mostly I'd anticipate certain decisions that are currently based on type sniffing would be hoisted into CppSignature and be based on the structure of the signature.

[ezyang]

Ah, there is a reason we can't do this. To restate your proposal, your thinking is that we should have:

args: Sequence[CppArgument]
this_arg: Optional[CppThisArgument]
tensor_opts_arg: Optional[CppTensorOptsArgument]

(and also same structure in the non-cpp model. This prevents us from having to do a pattern match. The reason this doesn't work is that this representation underspecifies the relative locations of all the arguments with each other. There are two major underspecifications:

• The this/self argument isn't always at the beginning of the function. where is the most well-known counterexample:

- func: where.self(Tensor condition, Tensor self, Tensor other) -> Tensor
  use_c10_dispatcher: full
  variants: function, method

• The TensorOptions arguments isn't always at the end of the function. In fact, ALL factory functions that accept memory format violate this condition (because memory format is not considered part of the grouping at the moment):

- func: empty.memory_format(int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor

Now, they do say exceptions shoudn't write the law, and so maybe we should manually work around these exceptions to make your proposal work out. But this was at least some of the motivation to set things up this way, at least initially.

[bhosmer]

Yeah, I realize that's not quite enough - I was thinking something more like

args: Sequence[CppArgument]
this_arg: Optional[Tuple[CppThisArgument, int]]
tensor_opts_arg: Optional[Tuple[CppTensorOptsArgument, int]]

...and this and tensor_opts_arg get spliced into the generated sequence at the specified indexes by as_seq().

[ezyang]

Index splicing is tricky because you need to do the splicing "all in one go" otherwise things get shifted as you splice each one in. It might be better to just have "before" and "after" argument lists.

[facebook-github-bot]

@ezyang merged this pull request in 9079aea.