Add faithful C++ API

[smessmer]

Stack from ghstack:

• VariableKernel calls into scattered C++ api #44158 VariableKernel calls into scattered C++ api
• Add faithful C++ API #44087 Add faithful C++ API

Each op taking a TensorOptions argument now has an additional overload in the C++ frontend where it takes scattered ScalarType, Layout, Device, bool instead of one TensorOptions argument.

If it is a c10-full op, then the scattered version calls into the dispatcher and the gathered version is a proxy calling into the scattered version.
If it is a non-c10-full op, then the gathered version calls into the dispatcher and the scattered version is a proxy calling into the gathered version.

This should minimize the amount of gathering and scattering needed.

This PR is also a prerequisite to remove the re-gathering of arguments that is currently happening in VariableKernel. Currently, VariableKernels gather arguments into a TensorOptions object
to call into the C++ API. In a PR stacked on top of this, VariableKernel will just directly call into the scattered C++ API introduced here and avoid the gathering step.

Differential Revision: D23492188

[dr-ci]

💊 CI failures summary and remediations

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

---

Commit a086fca was recently pushed. Waiting for builds...

---

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 96 times.

[ezyang]

I'd suggest fixing the schema_order_jit_arguments bug above first; then you'll be able to reuse some functions here.

    # See Note [Byte-for-byte compatibility]
    # NB: NOT actually schema order.  This is almost certainly a BUG.
    schema_order_jit_arguments = list(itertools.chain(f.func.arguments, f.func.out_arguments, f.func.kwarg_only_arguments))

[ezyang]

I think my high level code organization feedback is that the output from api.cpp should be more structured. Instead of having two separate functions which possibly return the same thing and then testing for equality, you should instead have a single function that returns a struct that describes the entire structure of the C++ bindings for this function. This struct would contain a mandatory "schema style" signature, as well as an optional TensorOptions convenience signature.

Furthermore, when binding to the dispatcher, it would be better to have things be more compositional. The chain of transformations that I'd expect to see in the terminal state is TensorOptions signature -> Schema style signature -> Dispatch call. So instead of processing TensorOptions directly, you should get schema-style (cpp) expressions from the tensor options, and then shunt everything through the same processing path. Your current code does an "optimization" of short circuiting the conversion to schema-style when the underlying dispatcher is still legacy, but if you really insist on this optimization, it should be organized in an easy to delete fashion, since we definitely don't want this long term.

[bhosmer]

Couple of API suggestions, plus there's really a lot of new code that I'd love to see factored into helpers, for readability, robustness etc.

[smessmer]

Major rewrite based on feedback. PTAL

[ezyang]

I spent some time working through this diff and trying out some alternative designs, and I think I have a better understanding for which way I'd like the diff to go. I'm hoping the discussion here will help get us more aligned on future new codegen changes.

Don't be afraid to define single-use NamedTuples for very localized situations. In this diff, we have to make a determination from a FunctionSchema whether or not we will need two C++ signatures or one. This determination currently lives signature_group (previously group_arguments) when we scan through the arguments in the schema and work out whether or not there is a set of arguments that group into TensorOptions. So, it seems you worked out that the easiest way to add this functionality without having to reimplement the scanning logic (a good call, by the way) is to modify the signature of signature_group to return a new structure CppBinding (as it was called at the time), that recorded both the gathered and ungathered versions.

The miscommunication that happened during review was that I misunderstood the intent of this struct. What you needed was a struct to report results from signature_group; what I thought the struct was supposed to be was a new abstraction for representing CPP API signatures, similar to how we already had abstractions for CppExpr and CppArg. This means that when I requested certain things (such as recording return types, and being able to render the entirety of the signature) which would make sense for a C++ abstraction, these requests would have lead to awkward implementation, as signature_group would have ended up having to handle a bunch of extra things that didn't really have anything to do with grouping arguments.

When I got to the same position in my variant, I made a separate struct for representing the result of gather_arguments (signature_group in your PR):

class GatherArgumentsResult(NamedTuple):
    arguments: Sequence[Union[Argument, ThisArgument], ...]
    gathered_arguments: Optional[Sequence[Union[Argument, TensorOptionsArguments, ThisArgument], ...]]

modified gather_arguments to return this struct:

def gather_arguments(
    func: FunctionSchema, *, method: bool
) -> GatherArgumentsResult

and then maintained a separate struct hierarchy to represent the actual C++ signature:

class CppSignature:
    # The schema this signature is derived from
    func: FunctionSchema
    # If this signature is a method, this is not None and contains
    # the corresponding ThisArgument for this signature
    method: Optional[ThisArgument]

    cpp_arguments: Tuple[CppArgument, ...]
    cpp_return_type: str

@dataclass(frozen=True)
class CppSignatureGroup:
    func: FunctionSchema
    signature: CppSignature
    gathered_signature: Optional[CppSignature]

No rendering goes on GatherArgumentsResult; it's just a vehicle for gather_arguments to return its results, and the constructor for CppSignatureGroup is what is actually responsible for assembling them into a full-blown concept of a C++ signature. I think this should also resolve your single-responsibility concerns, and it is a lot less weird to store, e.g., the return type in a CppSignature, than it is to store it in GatherArgumentsResult (where it makes no semantic sense.)

Arguably, this is the same as what you had before, minus an extra CppSignature struct. But I think having a CppSignature is definitely worth it. In the old codegen, we manually pull out arguments and returns, and it is already a bit wordy; and as you saw in the first version of your patch, it gets even more wordy when you have two versions of the C++ signature running around. Encapsulating them into a single struct makes sense, and is a refactor that could have been done prior to doing this patch.

CppArgument should accurately reflect the C++ arguments of your function; if they don't, you extend the concept of CppArgument to model groups. Something that seemed a bit smelly to me in early review was the new enum for SCATTER and GATHER and PASS_THROUGH. It seemed to me that this shouldn't be necessary, and I tried to outline the high level design reasoning in #44087 (comment) although you found it difficult to actually solve this problem entirely. I think I know why.

In the old codegen, there is a big coincidence: when mapping from cpp to dispatcher, it is always the case that a cpp argument maps to one or more dispatcher arguments, but never vice versa (never mapping multiple cpp arguments to one dispatcher argument). In your faithful C++ API world this is no longer the case: you may have to map a scattered tensor options (from the faithful API) into a single dispatcher argument (non-faithful API). This is very inconvenient to do with the CppArgument abstraction, because if you actually had separate CppArguments for each scattered argument, then you have to somehow group them back together to map them to a single TensorOptions dispatcher argument. You work around this in your PR by always passing in the gathered CppArguments, and then using process_tensoroptions to disambiguate whether or, no, actually, the binding sites were scattered CppArguments.

The old saying is that all problems in computer science can be solved with another level of indirection. I think this applies here too. If you need CppArguments to be grouped together, add the concept of a group of CppArguments. It would look something like this, in analogy to the existing TensorOptionsArguments struct:

@dataclass(frozen=True)
class CppArgumentGroup:
    def arguments() -> Sequence[CppArgument]:
        raise NotImplementedError

@dataclass(frozen=True)
class CppTensorOptionsArguments:
    argument: TensorOptionsArguments
    dtype: CppArgument
    layout: CppArgument
    device: CppArgument
    pin_memory: CppArgument
    def arguments(self):
        return [self.dtype, self.layout, self.device, self.pin_memory]

Most functions that operate on CppArgument today would now operate on Union[CppArgument, CppArgumentGroup]. In cases where the grouping information is not necessary (e.g., when rendering the declaration/definition signatures in C++), we simply expand the groups into a single CppArgument list. However, when the grouping is necessary (e.g., when computing dispatcher expressions), you then just directly operate on CppArgumentGroup, and your translation problems melt away: you don't need a scatter/gather argument, because you can work out the input convention from the CppArguments you are given as input, and the output convention via is use_c10_dispatcher: full.

This makes it hard to make a mistake, because you always feed in the precise description of what arguments are in scope to expressions (based on what you actually used to make the signature), and that description is what drives the expression process, not some side-band information

The prototype. I pushed an illustrative sketch of some of these ideas to ezyang@5c05730 It doesn't implement CppArgumentGroup, I haven't even type checked it, but I thought relatively carefully about how things should be put together and it lead me to some new observations about of your code which I didn't notice in CR. Let me know what you think.

Also, I'm putting the effort in here to communicate design ideas because I hope you'll be making more changes to the new codegen in the future. If you think that it makes more sense to unblock you by having me finish up this version of the change, or us agreeing to disagree and landing a basic version to be refactored later, please let me know.

[smessmer]

I'm ok with that design. Can you rebase your changes on top of my PR and land them separately? I've already done 3 different versions of this PR based on your feedback and would rather you implemented your vision in a refactoring PR on top instead of me doing a fourth version that again isn't 100% certain to be what you expected.

[ezyang]

@smessmer and I agreed to let me take a whack at refactoring this after landing.

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (gh/smessmer/255/base@ffd1cc9). Click here to learn what that means.
The diff coverage is n/a.

@@                   Coverage Diff                   @@
##             gh/smessmer/255/base   #44087   +/-   ##
=======================================================
  Coverage                        ?   68.56%           
=======================================================
  Files                           ?      405           
  Lines                           ?    51936           
  Branches                        ?        0           
=======================================================
  Hits                            ?    35609           
  Misses                          ?    16327           
  Partials                        ?        0           

---

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 ffd1cc9...a086fca. Read the comment docs.

[facebook-github-bot]

This pull request has been merged in 6e2eee2.