[codegen] split out backend-specific information from NativeFunction in the model by bdhirsh · Pull Request #57361 · pytorch/pytorch

bdhirsh · 2021-04-30T15:45:18Z

Data model change in the codegen, which splits backend-specific information out of NativeFunction

Overview

Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the NativeFunction object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types.

Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the NativeFunction object with the additional external information, by updating the dispatch entry to contain the external backend kernel name and dispatch key.

Instead, this PR tries to split out that information. The NativeFunction class contains all information about an operator from native_functions.yaml that's backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in an external backend yaml just involves updating that index with the new backend.

There were a few places where NativeFunction used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. is_abstract). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends.

This has a few advantages:

We can more easily re-use the existing logic in native_function.py and register_dispatch_key.py for both native and external backends, since they both involve a NativeFunction + a particular backend index
The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend.
There are several of codegen areas that don't care about backend-specific information: mostly the tracing and autograd codegen. We can reason about the codegen there more easily, knowing that backend-specific info is entirely uninvolved.

An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants.

Data Model Changes

Removed 3 classes, which are used by the external codegen:

ExternalBackendFunction
ExternalBackendFunctionsGroup
ExternalBackendMetadata

And added two new ones:

BackendIndex
BackendMetadata

BackendIndex contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op).

BackendMetadata contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured.

Questions

I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out:

Dealing with structured.
This PR separates out the notion of structured into two bits of information:

Does [operator] have a meta() function. This is backend-agnostic, and is represented by the structured property on NativeFunction, same as before. This is used, e.g., to decide what signatures to add to MetaFunctions.h.
Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on BackendMetadata. This is used in most other cases, e.g. in RegisterDispatchKey when we're deciding whether or not to gen a structured or unstructured wrapper.

I also baked is_structured_dispatch_key directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not.

I ended up trying to deal with structured in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make structured backend-dependent. I don't see anything too controversial related to structured in the change, but I tried to call out any areas in the comments

Localizing the fact that external backends follow Dispatcher convention.
Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in native_functions.py by adding a helper, kernel_signature, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need NativeSignature and DispatcherSignature to have matching interfaces. I didn't bother with that in this PR, which is why gen_external_aten_fallbacks.py still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions.

Maybe is_external() shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away.

Thoughts on the BackendIndex / BackendMetadata breakdown.
One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like backend_index.structured(f), which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an Optional[bool] though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though.

Stack from ghstack:

add a boxed CPU fallback kernel #58065 add a boxed CPU fallback kernel
remove xla-specific stuff from codegen (minus CPU fallback) #58064 remove xla-specific stuff from codegen (minus CPU fallback)
generate inplace/out kernels for xla #57510 generate inplace/out kernels for xla
[codegen] split out backend-specific information from NativeFunction in the model #57361 [codegen] split out backend-specific information from NativeFunction in the model
remove bridge API from codegen #55796 remove bridge API from codegen
[external codegen] better yaml error messaging, added explicit error message tests #56597 [external codegen] better yaml error messaging, added explicit error message tests
add _to_cpu() operator #55795 add _to_cpu() operator

Differential Revision: D28474362

…in the model [ghstack-poisoned]

facebook-github-bot · 2021-04-30T15:45:25Z

💊 CI failures summary and remediations

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

1/1 failures introduced in this PR

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

pytorch_xla_linux_bionic_py3_6_clang9_test (1/1)

Step: "Run tests" (full log | diagnosis details | 🔁 rerun)

May 14 21:11:31 ERROR [0.024s]: test_scatter_re...alar_xla_float64 (__main__.TestTorchDeviceTypeXLA)

May 14 21:11:31 AutogradPrivateUse1: registered at /var/lib/jenkins/workspace/torch/csrc/autograd/generated/VariableType_4.cpp:9226 [autograd kernel]
May 14 21:11:31 AutogradPrivateUse2: registered at /var/lib/jenkins/workspace/torch/csrc/autograd/generated/VariableType_4.cpp:9226 [autograd kernel]
May 14 21:11:31 AutogradPrivateUse3: registered at /var/lib/jenkins/workspace/torch/csrc/autograd/generated/VariableType_4.cpp:9226 [autograd kernel]
May 14 21:11:31 Tracer: registered at /var/lib/jenkins/workspace/torch/csrc/autograd/generated/TraceType_4.cpp:9909 [kernel]
May 14 21:11:31 Autocast: fallthrough registered at /var/lib/jenkins/workspace/aten/src/ATen/autocast_mode.cpp:255 [backend fallback]
May 14 21:11:31 Batched: registered at /var/lib/jenkins/workspace/aten/src/ATen/BatchingRegistrations.cpp:1019 [backend fallback]
May 14 21:11:31 VmapMode: fallthrough registered at /var/lib/jenkins/workspace/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
May 14 21:11:31 
May 14 21:11:31 
May 14 21:11:31 ======================================================================
May 14 21:11:31 ERROR [0.024s]: test_scatter_reduce_scalar_xla_float64 (__main__.TestTorchDeviceTypeXLA)
May 14 21:11:31 ----------------------------------------------------------------------
May 14 21:11:31 Traceback (most recent call last):
May 14 21:11:31   File "/opt/conda/lib/python3.6/site-packages/torch/testing/_internal/common_device_type.py", line 297, in instantiated_test
May 14 21:11:31     raise rte
May 14 21:11:31   File "/opt/conda/lib/python3.6/site-packages/torch/testing/_internal/common_device_type.py", line 292, in instantiated_test
May 14 21:11:31     result = test_fn(self, *args)
May 14 21:11:31   File "/var/lib/jenkins/workspace/xla/test/../../test/test_torch.py", line 5548, in test_scatter_reduce_scalar
May 14 21:11:31     input.scatter_(0, index, src, reduce=operation)
May 14 21:11:31 NotImplementedError: Could not run 'aten::_copy_from_and_resize' with arguments from the 'CPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::_copy_from_and_resize' is only available for these backends: [XLA, BackendSelect, Named, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, UNKNOWN_TENSOR_TYPE_ID, AutogradMLC, AutogradHPU, AutogradNestedTensor, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].
May 14 21:11:31

XLA failure

Job pytorch_xla_linux_bionic_py3_6_clang9_test is failing. Please create an issue with title prefixed by [PT_BREAK] in pytorch/xla and link to to this PR. If you have questions, please reach out to @ailzhang / @dlibenzi / @JackCaoG.

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 to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

…in the model ghstack-source-id: ce0971b Pull Request resolved: #57361

…veFunction in the model" [ghstack-poisoned]

…in the model ghstack-source-id: a6ca0f6 Pull Request resolved: #57361

…veFunction in the model" [ghstack-poisoned]

…in the model ghstack-source-id: dc17321 Pull Request resolved: #57361

bdhirsh · 2021-04-30T17:55:12Z

tools/codegen/api/types.py

+
+# Helper functions
+
+def kernel_signature(f: NativeFunction, backend_index: BackendIndex) -> Optional[Union['NativeSignature', 'DispatcherSignature']]:


Mentioned in the PR description, but the goal would be for this function to later abstract away any differences so that RegisterDispatchKey doesn't have to worry about signature differences between external vs. internal backends.

I also kind of wanted this to be in model.py, but that adds a circular dependency since the return type of the function involves NativeSignature/DispatcherSignature.

bdhirsh · 2021-04-30T17:57:53Z

tools/codegen/context.py

+
+# Convenience decorator for functions that explicitly take in a BackendIndex,
+# instead of indirectly taking one in as a closure
+def with_native_function_and_index(func: Callable[[F, BackendIndex], T]) -> Callable[[F, BackendIndex], T]:


😬 indeed

Solve a problem with dynamic scoping, now ya got etc

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that it backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in external backend metadata just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has two advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

bdhirsh · 2021-04-30T18:24:17Z

tools/codegen/model.py

        e.pop('__line__', None)
        assert not e, f"leftover entries: {e}"

+        # Asserts that we can't do in post_init, because they rely on backend-specific info


Instead of moving this assert out of post_init, I could have added some properties on NativeFunction to ensure that post_init has enough information to perform the assert.

I couldn't think of a way to do it that preserves the error message without just saving the list of dispatch keys, so I moved the assert here instead.

bdhirsh · 2021-04-30T18:28:57Z

tools/codegen/dest/register_dispatch_key.py

        if isinstance(f, NativeFunctionsGroup):
+            # Note: We call gen_structured() if the operator is marked structured, regardless of the backend.
+            # gen_structured() has special logic to handle auto-generated kernels.
            if f.structured:


Here's an example of how moving structured to be backend-dependent is "weird". We have special handling to auto-generate kernels for Meta and CompositeExplicitAutograd kernels for structured operators. Because of that, the decision to call gen_structured() should depend on whether the operator is structured, not the backend. Kinda weird- maybe instead I should just move the special handling of those keys directly into this function?

Yeah, this is an interesting modeling question. An alternative to unconditionally calling gen_structured here is to have some preprocessing pass where we stuff the BackendIndex with entries for auto-generated kernels. What is going on here is OK-ish, though, since gen_structured knows how to call back to gen_unstructured if necessary.

Couldn't we decide up here, based on whether the op and the backend are both structured?

I think that this might work, if we mark DispatchKey.Meta and DispatchKey.Composite[Explicit|Implicit]Autograd keys as structured - since structured kernel codegen sometimes tries to autogenerate kernels for those keys if they don't exist. Probably need to try it out and see if that has knock-on effects that break anything else though.

If you're ok with it I'm going to write that down as a TODO and try it in a PR later in the stack, since it looks like this PR and the corresponding changes at pytorch/xla#2915 are finally passing CI.

bdhirsh · 2021-04-30T18:29:56Z

tools/codegen/dest/register_dispatch_key.py


    def gen_class_set_output_body(self, k: SchemaKind) -> str:
-        if self.dispatch_key in [DispatchKey.CUDA, DispatchKey.CompositeExplicitAutograd]:
+        if self.backend_index.dispatch_key in [DispatchKey.CUDA, DispatchKey.CompositeExplicitAutograd]:


we could in theory also store stuff like "needs device guard" in the backend index, but maybe that's overkill. Might be a nice way to localize some of the hardcoded lists of dispatch keys though.

doesn't seem like overkill to me - backend properties now have a place to live! I mean, maybe overkill for this PR 😁

yeah, that's what I meant 😃. Sounds good.

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that it backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in external backend metadata just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has two advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

…in the model ghstack-source-id: 4a0ca55 Pull Request resolved: #57361

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that it backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in external backend metadata just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has two advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

bdhirsh · 2021-04-30T19:19:19Z

tools/codegen/gen_backend_stubs.py

+            'cpp_namespace': cpp_namespace,
+            'dispatch_aten_fallback_declarations': list(concatMap(
+                dest.GenExternalAtenFallback(Target.NAMESPACED_DECLARATION, backend_indices[backend_dispatch_key]),
+                [g for g in grouped_native_functions if not backend_indices[autograd_dispatch_key].has_backend(g)]


Here's an annoying thing about GenExternalAtenFallbacks: It doesn't "just work" when you call it with multiple dispatch keys and concat the results together, since the whole point is that it generates CPU fallback kernels for everything, include operators that don't have a kernel for the specified backend. So calling it twice, once with the XLA key and one with AutogradXLA, results in a bunch of duplicates. And it also does the wrong thing for AutogradXLA, since we don't want to register CPU fallbacks to the AutogradXLA key, just XLA. I think it'll be a little better once I remove all of the non-fallback related logic from GenExternalAtenFallback.

Splitting the codegen into two separate pieces to avoid duplication sounds like the right call.

bdhirsh · 2021-04-30T19:21:48Z

tools/codegen/dest/native_functions.py

@@ -1,96 +1,59 @@
-from typing import List, Union, Set, Any


One thing that's nice about this PR (although it's hard to tell from the git diff) is that this file got a lot shorter. gen_unstructured() and gen_structured() both do less since they only operate on a single (NativeFunction, backend) pair. And gen_structured() already knows that it's been handed a structured backend, so it doesn't need to handle multiple cases.

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that it backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in external backend metadata just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has two advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. I don't see anything *too* controversial related to structured in the change, but I tried to call out any areas in the comments **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. Maybe `is_external()` shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

…in the model ghstack-source-id: 1dddcb4 Pull Request resolved: #57361

ezyang · 2021-05-03T15:40:02Z

tools/codegen/gen_backend_stubs.py

+
+    autograd_key: Optional[DispatchKey] = None
+    if len(supported_autograd) > 0:
+        autograd_key = add_backend_index(supported_autograd, backend, is_autograd=True)


I know it's less efficient, but I probably would have preferred it if this function didn't mutate its argument and just returned a mini backend index to be merged into the main one later. (The general API design concept here is it's OK to use mutation locally, but public APIs should try to avoid mutation as much as possible, because it means that the contract for the function is a bit more complicated).

yep good point, easy enough to perform the mutation outside of that function

tools/codegen/model.py

ezyang · 2021-05-03T15:47:20Z

tools/codegen/model.py

+    # However, external backends like XLA can indendently toggle which ops are structured.
+    structured: bool
+    # Whether or not this op is an in-tree (e.g. CPU/CUDA) or out-of-tree backend (e.g. XLA)
+    external: bool


If you ever get around to making that global backend metadata, external probably belongs there, not here.

...maybe this should be in BackendIndex? (This is not worth holding up this PR for)

Yep, definitely should be in backend index

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that's backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in an external backend yaml just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has a few advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. - There are several of codegen areas that don't care about backend-specific information: mostly the tracing and autograd codegen. We can reason about the codegen there more easily, knowing that backend-specific info is entirely uninvolved. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. I don't see anything *too* controversial related to structured in the change, but I tried to call out any areas in the comments **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. Maybe `is_external()` shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

…in the model ghstack-source-id: e532e92 Pull Request resolved: pytorch#57361

bdhirsh · 2021-05-04T17:38:07Z

Thanks! I think I responded to all of the above but I'll give it one more pass. Also, I accidentally snuck a few more yaml error checking tests into this PR due to a mixup when git rebasing :/ They're mostly just the extra tests from #56834 though, which you already approved.

Looks like this also made the static dispatch tests unhappy- looking into that.

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that's backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in an external backend yaml just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has a few advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. - There are several of codegen areas that don't care about backend-specific information: mostly the tracing and autograd codegen. We can reason about the codegen there more easily, knowing that backend-specific info is entirely uninvolved. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. I don't see anything *too* controversial related to structured in the change, but I tried to call out any areas in the comments **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. Maybe `is_external()` shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

bhosmer

Hey looks good! A few comments/suggestions inline, but per offline discussion I'm not suggesting any extensive restructuring.

bhosmer · 2021-05-12T01:43:23Z

tools/codegen/api/types.py

+def kernel_signature(f: NativeFunction, backend_index: BackendIndex) -> Union['NativeSignature', 'DispatcherSignature']:
+    # Note [External Backends Follow Dispatcher API]
+    # Kernel signatures for in-tree backends follow the "native" API,
+    # while kernels for out-of-tree backends follow the dispatcher API.


I'm wondering how we should treat this design point going forward. Here we're basically institutionalizing it, but given the comment here, do we actually know if it's a pattern we want to encourage for other out of tree backends? If not, we could make it a direct property of BackendIndex, rather than correlate it with external.

And regardless, it might be worth adding a sentence or two here on why the distinction exists - motivation and/or history.

So, my take is that we want to enforce/encourage external backends to all follow the dispatcher API convention if we can: both because (a) any differences between the dispatcher API and the API of the backend kernel results in runtime overhead to convert between the two, and (b) enforcing some uniformity across backends will be easier to manage.

We have a separate native API in-tree, but my understanding is that this just emerged from the fact that our native kernel signatures drifted from the dispatcher API at some point, and introducing a way to translate between the two API's is easier than going through every kernel in the codebase and fixing up its signature.

Compare that to external backends: XLA kernels already follow the dispatcher convention, and it'll be easy to enforce that new backends write kernels that follow the same convention. It seems likely that other backends would follow the same pattern, since we advertise RegistrationDeclarations.yaml as the way to tell which operators a backend has to implement, and the names of the ops in that file all follow dispatcher convention.

But you're right, there's still some chance though that there's a backend out there that already wrote a bunch of kernels and didn't follow the dispatcher naming/schema convention. My take is to allow per-backend kernel API's only if we see that somewhere. Another extension I could imagine we might have to do is if the backend gave all of their kernels custom names that didn't follow any of the API's we currently use- we could probably add an option in the yaml to provide a custom kernel name or something, but for the same reason I don't think we'd want that kind of extension unless there's a need for it.

Let me know if you agree with that argument 😄

Reading over your comment, @bdhirsh, I'm actually imagining the opposite situation, where it is kernels in PyTorch that are easier to migrate (because we have access to them and can easily change them), and external backend functions that might want to move more slowly because we can't just update them whenever we make changes. So actually you could imagine an endless stream of versions of signatures, and backends can migrate to the newer versions by opting in. Or I guess we can force them to rip off the bandaid every version release they migrate onto...

@bdhirsh Yeah that makes perfect sense, and I think you're right that we should proceed as though the native API is legacy - not support it in out-of-tree codebase unless some case comes up that forces the issue. (It might be worth adding a sentence to the comment here giving the high-level color on the native API - otherwise it's pretty easy to read the comment as implying something like "fundamental differences between in tree and out of tree backends have led us to design two different APIs" or whatever.

Right. I'm thinking about the following two things as kind of independent (but maybe that's wrong):

How to ease BC issues for external backends
This applies to stuff like "we added a new default-able param to smooth_l1_loss", or "we renamed an operator from _{op} to {op}". Right now, any naming/schema changes to any native op will break external backends (their master will fail to build) until they make the corresponding change to their kernel. The codegen gives us a hook into making that experience easier, although I'm not too sure what the design for that would look like. I remember talking with Ed a few weeks ago about some notion of explicit versioning, probably per operator, that a backend would need to specify. One way is for backends to specify "XLA backend has a kernel for version 2 of at::add", and to make our codegen aware of how to translate between schema versions. This would put a little more onus on people that modify op schema's in-tree, and doesn't fix all of the BC issues, but it would provide at least some form of shielding backends from BC changes.

The big distinction in my head is that these changes apply on a per-operator basis.

What naming/schema conventions a backend should abide by
In my original comment above, I was thinking more around "how does a backend choose what names to give their kernels" or "do they include packed vs. unpacked TensorOptions", which are conventions that I think should apply to all kernels. So, sort of "what is the default naming/schema convention of external backend kernels", which right now in the codegen is the dispatcher API. We could technically allow external backends to abide by some other convention (like our native API). But I was thinking of that as a different problem compared to dealing with BC-breaking operator changes, which needs to be solved separately.

bhosmer · 2021-05-12T02:12:57Z

tools/codegen/context.py

+
+# Convenience decorator for functions that explicitly take in a BackendIndex,
+# instead of indirectly taking one in as a closure
+def with_native_function_and_index(func: Callable[[F, BackendIndex], T]) -> Callable[[F, BackendIndex], T]:


Solve a problem with dynamic scoping, now ya got etc

bhosmer · 2021-05-12T02:14:53Z

tools/codegen/dest/gen_external_aten_fallbacks.py

-    in_denylist = any([re.match(frx, str(f.native_function.func.name)) for frx in _FN_DENYLIST_REGEX])
+def requires_backend_wrapper(f: NativeFunction, backend_index: BackendIndex) -> bool:
+    requires_lowering = not f.has_composite_kernel and not has_autogenerated_composite_kernel(f)
+    has_xla_lowering = backend_index.has_backend(f)


bhosmer · 2021-05-12T02:21:11Z

tools/codegen/dest/gen_external_aten_fallbacks.py

+    in_denylist = any([re.match(frx, str(f.func.name)) for frx in _FN_DENYLIST_REGEX])
    return not in_denylist and (requires_lowering or has_xla_lowering)

 def xla_tensor_creation_api(


I'm not sure how to read the 'xla' here and elsewhere in this file - is it basically synonymous with "external", or does it really mean XLA and we'll eventually have sibling logic for other out of tree backends?

ah, yeah... just treat "xla" as "external backend". There are probably better ways to slice up the work, but I'm currently going "enable minimum viable version of xla codegen in-tree" -> "expand codegen with new features" -> "make the codegen backend agnostic". So there's a PR further up in this stack to remove references to xla everywhere, etc. Not 100% ready for review but I'll let you know when it is!

bhosmer · 2021-05-12T02:40:59Z

tools/codegen/dest/register_dispatch_key.py

        if isinstance(f, NativeFunctionsGroup):
+            # Note: We call gen_structured() if the operator is marked structured, regardless of the backend.
+            # gen_structured() has special logic to handle auto-generated kernels.
            if f.structured:


Couldn't we decide up here, based on whether the op and the backend are both structured?

bhosmer · 2021-05-12T02:42:44Z

tools/codegen/dest/register_dispatch_key.py


    def gen_class_set_output_body(self, k: SchemaKind) -> str:
-        if self.dispatch_key in [DispatchKey.CUDA, DispatchKey.CompositeExplicitAutograd]:
+        if self.backend_index.dispatch_key in [DispatchKey.CUDA, DispatchKey.CompositeExplicitAutograd]:


doesn't seem like overkill to me - backend properties now have a place to live! I mean, maybe overkill for this PR 😁

bhosmer · 2021-05-12T02:46:18Z

tools/codegen/gen_backend_stubs.py

+        backward_kernels = [f for f in backward_kernels if f is not None]
+        assert len(forward_kernels) == 0 or len(backward_kernels) == 0, \
+            f'Currently, all variants of an op must either be registered to a backend key, or to a backend\'s \
+autograd key. They can not be mix and matched. If this is something you need, feel free to create an issue! \


nit: cannot

bhosmer · 2021-05-12T02:53:40Z

tools/codegen/model.py

+    #
+
+
+# The BackendIndex encodes per-operator information that is potentially different


I think it might clarify to explicitly say up front that a BackendIndex represents a backend - i.e., it's dictionary of per-operator information, rather than information about a single operator. It's clear enough once you get into the code, but my first time through this comment I wasn't actually sure.

bhosmer · 2021-05-12T03:00:31Z

tools/codegen/model.py

+        else:
+            return g.functional
+
+    def has_backend(self, g: Union[NativeFunction, NativeFunctionsGroup]) -> bool:


Shouldn't this be has_kernel or has_function? Not sure the best noun, but has_backend seems backwards - it names the dictionary rather than the thing you're asking the dictionary if it has.

yeah you're absolutely right. I'll go with has_kernel

bhosmer · 2021-05-12T03:02:37Z

tools/codegen/model.py

+        return m is not None
+
+
+    def get(self, g: Union[NativeFunction, NativeFunctionsGroup]) -> Optional[BackendMetadata]:


Also brevity notwithstanding, I'd mildly pitch using the has_<noun> and get_<noun> naming convention here.

…veFunction in the model" Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that's backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in an external backend yaml just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has a few advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. - There are several of codegen areas that don't care about backend-specific information: mostly the tracing and autograd codegen. We can reason about the codegen there more easily, knowing that backend-specific info is entirely uninvolved. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. I don't see anything *too* controversial related to structured in the change, but I tried to call out any areas in the comments **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. Maybe `is_external()` shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. [ghstack-poisoned]

bdhirsh · 2021-05-17T12:46:34Z

@bdhirsh has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot · 2021-05-17T19:26:13Z

@bdhirsh merged this pull request in 9354a68.

…in the model (pytorch#57361) Summary: Pull Request resolved: pytorch#57361 Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that's backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in an external backend yaml just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has a few advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. - There are several of codegen areas that don't care about backend-specific information: mostly the tracing and autograd codegen. We can reason about the codegen there more easily, knowing that backend-specific info is entirely uninvolved. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. I don't see anything *too* controversial related to structured in the change, but I tried to call out any areas in the comments **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. Maybe `is_external()` shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D28474362 Pulled By: bdhirsh fbshipit-source-id: 41a00821acf172467d764cb41e771e096542f661

…in the model (#57361) Summary: Pull Request resolved: #57361 Data model change in the codegen, which splits backend-specific information out of `NativeFunction` ### Overview Currently in the codegen, native_functions.yaml has backend-specific information about each operator that is encoded directly into the data model, in the `NativeFunction` object. That's reasonable, since the native_functions.yaml is the source of truth for information about an operator, and the data model encodes that information into types. Now that external backends can use the codegen though, that information is technically incomplete/inaccurate. In another PR, I tried patching the information on the `NativeFunction` object with the additional external information, by updating the `dispatch` entry to contain the external backend kernel name and dispatch key. Instead, this PR tries to split out that information. The `NativeFunction` class contains all information about an operator from native_functions.yaml that's backend-independent and is known never to change regardless of what extra information backends provide. We also build up a backend "index", which is basically a mapping from [backend] -> [backend-specific-metadata]. Reading in an external backend yaml just involves updating that index with the new backend. There were a few places where `NativeFunction` used the dispatch table directly, that I encoded as properties directly on the NativeFunction object (e.g. `is_abstract`). They were mostly around whether or not the operator has a composite kernel, which isn't something that's going to change for any external backends. This has a few advantages: - We can more easily re-use the existing logic in `native_function.py` and `register_dispatch_key.py` for both native and external backends, since they both involve a NativeFunction + a particular backend index - The data in the data model will be the same regardless of how the codegen is run. Running the codegen with a new external backend doesn't change the data inside of NativeFunction or an existing backend index. It just adds a new index for that backend. - There are several of codegen areas that don't care about backend-specific information: mostly the tracing and autograd codegen. We can reason about the codegen there more easily, knowing that backend-specific info is entirely uninvolved. An alternative to this split would be to augment the NativeFunction objects with external backend information at the time that we create them. So the external codegen could read both native_functions.yaml and the external backend's yaml at the same time, and construct a NativeObject with a full dispatch table (including the XLA entry), and the correct setting of structured (taking into account both yamls). One disadvantage to this approach is that NativeFunction objects now contain different stuff depending on how you ran the codegen, and you have to make sure that any changes to the codegen can properly handle all the different variants. ### Data Model Changes Removed 3 classes, which are used by the external codegen: - ExternalBackendFunction - ExternalBackendFunctionsGroup - ExternalBackendMetadata And added two new ones: - BackendIndex - BackendMetadata `BackendIndex` contains any info that's specific to that backend, plus a mapping from operator names to backend specific metadata about the operator. One example of backend-specific info that's not operator-dependent is the fact that XLA prefers to implement functional kernels instead of out kernels (and so when they eventually mark an op as structured, they're going to mark the functional op and not the out op). `BackendMetadata` contains info specific to an (operator, backend) pair. Right now, that's just (a) the name of the kernel, and (b) whether or not that operator is structured. ### Questions I wanted to get this PR up earlier so I could get feedback, but there are a few things I want to call out: **Dealing with `structured`.** This PR separates out the notion of `structured` into two bits of information: - Does [operator] have a meta() function. This is backend-agnostic, and is represented by the `structured` property on `NativeFunction`, same as before. This is used, e.g., to decide what signatures to add to `MetaFunctions.h`. - Does [operator, backend] have an impl() function. This is backend dependent; even though technically all in-tree backends are forced to write impl() functions for an operator when we port the op to structured in native_functions.yaml, out-of-tree backends can decide to opt in independently. This is represented as a property on `BackendMetadata`. This is used in most other cases, e.g. in `RegisterDispatchKey` when we're deciding whether or not to gen a structured or unstructured wrapper. I also baked `is_structured_dispatch_key` directly into each BackendIndex. So for operators marked "structured" in native_functions.yaml, their corresponding CPU/CUDA BackendIndex entries will be marked structured, and all others (except for potentially external backends) will not. I ended up trying to deal with `structured` in this change since it's technically backend dependent (XLA can opt kernels into structured separately from in-tree ops), but that may have been too ambitious: it's technically not relevant until we actually add support for structured external kernels. If it's not clear that this is the right path for dealing with structured and we want to push that off, I'm fine with backing out the bits of this PR that make `structured` backend-dependent. I don't see anything *too* controversial related to structured in the change, but I tried to call out any areas in the comments **Localizing the fact that external backends follow Dispatcher convention.** Another thing that's sort of backend specific that I didn't totally address in this PR is the fact the fact that in-tree backends follow the Native API while external backends follow the Dispatcher API. I painted over that in `native_functions.py` by adding a helper, `kernel_signature`, that takes in a native function and gives you the "correct" signature for the specified backend- NativeSignature for in-tree backends, and DispatcherSignature for out-of-tree backends. In order to make that fully useable though, we'll need `NativeSignature` and `DispatcherSignature` to have matching interfaces. I didn't bother with that in this PR, which is why `gen_external_aten_fallbacks.py` still has a bunch of direct references to the dispatcher API. Thinking of adding it in a later PR but wanted to see if anyone has other opinions. Maybe `is_external()` shouldn't even be a property on the BackendMetadata, and anything the codegen does that requires asking for that information should just be better abstracted away. **Thoughts on the `BackendIndex` / `BackendMetadata` breakdown.** One thing that's annoying right now is that to query for various pieces of metadata, you call helper functions like `backend_index.structured(f)`, which queries that particular backend and tells you if that specific NativeFunctionGroup is structured for that backend. It has to return an `Optional[bool]` though, since you have to handle the case where that operator doesn't have a kernel for that backend at all. So users of those helpers end up with a bunch of optionals that they need to unpack, even if they know at some point that the result isn't None. I think it would be easier instead to just store the NativeFunction object as a field directly on the BackendMetadata. Curious if there are any other opinions on a better way to model it though. Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D28474362 Pulled By: bdhirsh fbshipit-source-id: 41a00821acf172467d764cb41e771e096542f661

Summary: Pull Request resolved: #58889 fixes #58796 Planning on re-testing locally tomorrow morning to confirm, but this change should fix the non-determinism in the codegen output that was causing `ccache` not to re-use its cached output. I built from the commit referenced in #58796 a few times and ran `diff -Naur` on the codegen output in `build/aten/src/ATen`. After a few tries, `NativeFunctions.h` had a few diffs. The diffs were all related to the ordering of functional/inplace/out variants of a NativeFunctionGroup, which looked non-deterministic. That looks like it's coming from my calling `set()` to filter out duplicate NativeFunction declarations. The earlier version of the codegen also called `set()` to filter out duplicates, but it did so individually for each `NativeFunction` object, before merging the groups (I'm not too sure why this didn't introduce non-determinism before. though). With the refactor from #57361, we're calling `set()` on the declarations from every operator for a given DispatchKey, which is probably what introduced the nondeterminism. Test Plan: Imported from OSS Reviewed By: gchanan Differential Revision: D28675941 Pulled By: bdhirsh fbshipit-source-id: bb66de00aafeeb9720d85e8156ac9f7539aed0d6

Summary: Pull Request resolved: pytorch#58889 fixes pytorch#58796 Planning on re-testing locally tomorrow morning to confirm, but this change should fix the non-determinism in the codegen output that was causing `ccache` not to re-use its cached output. I built from the commit referenced in pytorch#58796 a few times and ran `diff -Naur` on the codegen output in `build/aten/src/ATen`. After a few tries, `NativeFunctions.h` had a few diffs. The diffs were all related to the ordering of functional/inplace/out variants of a NativeFunctionGroup, which looked non-deterministic. That looks like it's coming from my calling `set()` to filter out duplicate NativeFunction declarations. The earlier version of the codegen also called `set()` to filter out duplicates, but it did so individually for each `NativeFunction` object, before merging the groups (I'm not too sure why this didn't introduce non-determinism before. though). With the refactor from pytorch#57361, we're calling `set()` on the declarations from every operator for a given DispatchKey, which is probably what introduced the nondeterminism. Test Plan: Imported from OSS Reviewed By: gchanan Differential Revision: D28675941 Pulled By: bdhirsh fbshipit-source-id: bb66de00aafeeb9720d85e8156ac9f7539aed0d6

[codegen] split out backend-specific information from NativeFunction …

8403086

…in the model [ghstack-poisoned]

bdhirsh requested review from albanD and soulitzer as code owners April 30, 2021 15:45

This was referenced Apr 30, 2021

add _to_cpu() operator #55795

Closed

[external codegen] better yaml error messaging, added explicit error message tests #56597

Closed

remove bridge API from codegen #55796

Closed

facebook-github-bot added the cla signed label Apr 30, 2021

bdhirsh added a commit that referenced this pull request Apr 30, 2021

[codegen] split out backend-specific information from NativeFunction …

a23e27f

…in the model ghstack-source-id: ce0971b Pull Request resolved: #57361

Update on "[codegen] split out backend-specific information from Nati…

c402995

…veFunction in the model" [ghstack-poisoned]

bdhirsh added a commit that referenced this pull request Apr 30, 2021

[codegen] split out backend-specific information from NativeFunction …

f24e3ff

…in the model ghstack-source-id: a6ca0f6 Pull Request resolved: #57361

Update on "[codegen] split out backend-specific information from Nati…

4dcebc2

…veFunction in the model" [ghstack-poisoned]

bdhirsh added a commit that referenced this pull request Apr 30, 2021

[codegen] split out backend-specific information from NativeFunction …

b134c3e

…in the model ghstack-source-id: dc17321 Pull Request resolved: #57361

bdhirsh commented Apr 30, 2021

View reviewed changes

bdhirsh added a commit that referenced this pull request Apr 30, 2021

[codegen] split out backend-specific information from NativeFunction …

3013452

…in the model ghstack-source-id: 4a0ca55 Pull Request resolved: #57361

bdhirsh commented Apr 30, 2021

View reviewed changes

bdhirsh requested a review from ezyang April 30, 2021 20:09

bdhirsh added a commit that referenced this pull request Apr 30, 2021

[codegen] split out backend-specific information from NativeFunction …

2f38f2f

…in the model ghstack-source-id: 1dddcb4 Pull Request resolved: #57361

ezyang requested a review from bhosmer May 3, 2021 15:07

ezyang reviewed May 3, 2021

View reviewed changes

tools/codegen/model.py Outdated Show resolved Hide resolved

ezyang reviewed May 3, 2021

View reviewed changes

ezyang approved these changes May 3, 2021

View reviewed changes

bdhirsh mentioned this pull request May 3, 2021

generate inplace/out kernels for xla #57510

Closed

dgl-intel pushed a commit to dgl-intel/pytorch that referenced this pull request May 3, 2021

[codegen] split out backend-specific information from NativeFunction …

0563224

…in the model ghstack-source-id: e532e92 Pull Request resolved: pytorch#57361

bdhirsh added 2 commits May 4, 2021 13:19

This was referenced May 11, 2021

remove xla-specific stuff from codegen (minus CPU fallback) #58064

Closed

add a boxed CPU fallback kernel #58065

Closed

bhosmer approved these changes May 12, 2021

View reviewed changes

bdhirsh added 2 commits May 12, 2021 07:05

facebook-github-bot closed this in 9354a68 May 17, 2021

facebook-github-bot added the Merged label May 17, 2021

This was referenced May 20, 2021

Cherrypick xla codegen #58717

Closed

[v1.9.0] Release Tracker #58518

Closed

facebook-github-bot deleted the gh/bdhirsh/112/head branch May 21, 2021 14:17

xwang233 mentioned this pull request May 22, 2021

ccache no longer caches pytorch build after PR 57361 #58796

Closed

bdhirsh mentioned this pull request May 25, 2021

Ensure NativeFunctions.h codegen output is deterministic #58889

Closed


		# Helper functions

		def kernel_signature(f: NativeFunction, backend_index: BackendIndex) -> Optional[Union['NativeSignature', 'DispatcherSignature']]:

		#


		# The BackendIndex encodes per-operator information that is potentially different

		return m is not None


		def get(self, g: Union[NativeFunction, NativeFunctionsGroup]) -> Optional[BackendMetadata]:

Conversation

bdhirsh commented Apr 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Data Model Changes

Questions

Uh oh!

facebook-github-bot commented Apr 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

🕵️ 1 new failure recognized by patterns

pytorch_xla_linux_bionic_py3_6_clang9_test (1/1)

XLA failure

Uh oh!

bdhirsh Apr 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdhirsh Apr 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bhosmer May 12, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdhirsh May 12, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdhirsh Apr 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdhirsh commented May 4, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bhosmer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bhosmer May 12, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

bdhirsh commented Apr 30, 2021 •

edited

Loading

facebook-github-bot commented Apr 30, 2021 •

edited

Loading

bdhirsh Apr 30, 2021 •

edited

Loading

bdhirsh Apr 30, 2021 •

edited

Loading

bhosmer May 12, 2021 •

edited

Loading

bdhirsh May 12, 2021 •

edited

Loading

bdhirsh Apr 30, 2021 •

edited

Loading

bdhirsh commented May 4, 2021 •

edited

Loading

bhosmer May 12, 2021 •

edited

Loading

bhosmer May 12, 2021 •

edited

Loading