pytorch
33996118 - [pytorch] refine dispatch keys in native_functions.yaml (1/N) (#45010)

Commit
4 years ago
[pytorch] refine dispatch keys in native_functions.yaml (1/N) (#45010) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/45010 The motivation of this change is to differentiate "backend specific" ops and "generic" ops. "backend specific" ops are those invoking backend specific kernels thus only able to run on certain backends, e.g.: CPU, CUDA. "generic" ops are those not *directly* invoking backend specific kernels. They are usually calling other "backend specific" ops to get things done. Thus, they are also referred to as "composite" ops, or "math" ops (because they are usually pure C++ code constructed from math formula). The other way to see the difference is that: we have to implement new kernels for the "backend specific" ops if we want to run these ops on a new backend. In contrast, "generic"/"composite" ops can run on the new backend if we've added support for all the "backend specific" ops to which they delegate their work. Historically we didn't make a deliberate effort to always populate supported backends to the "dispatch" section for all the "backend specific" ops in native_functions.yaml. So now there are many ops which don't have "dispatch" section but are actually "backend specific" ops. Majority of them are calling "DispatchStub" kernels, which usually only support CPU/CUDA (via TensorIterator) or QuantizedCPU/CUDA. The ultimate goal is to be able to differentiate these two types of ops by looking at the "dispatch" section in native_functions.yaml. This PR leveraged the analysis script on #44963 to populate missing dispatch keys for a set of "backend specific" ops. As the initial step, we only deal with the simplest case: * These ops don't already have dispatch section in native_functions.yaml; * These ops call one or more DispatchStub (thus "backend specific"); * These ops don't call any other aten ops - except for some common ones almost every op calls via framework, e.g. calling aten::eq via Dispatcher::checkSchemaCompatibility. Calling other nontrivial aten ops is a sign of being "composite", so we don't want to deal with this case now; * These ops don't call Tensor::is_quantized() / Tensor::is_sparse() / etc. Some ops call thse Tensor::is_XXX() methods to dispatch to quantized / sparse kernels internally. We don't deal with this case now. Test Plan: Imported from OSS Reviewed By: ezyang Differential Revision: D23803951 Pulled By: ljk53 fbshipit-source-id: aaced7c34427d1ede72380af4513508df366ea16
Author
Parents
Loading