pytorch

6063ffd0 - Specify dispatch key with kernel (#20821)

Specify dispatch key with kernel (#20821) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/20821 Change registration API. Instead of static auto registry = torch::RegisterOperators() .op("my::op", torch::RegisterOperators::options() .kernel<Kernel>() .dispatchKey(CPUTensorId())); it is now static auto registry = torch::RegisterOperators() .op("my::op", torch::RegisterOperators::options() .kernel<Kernel>(CPUTensorId())); This binds kernel and dispatch key together, allowing them to be separate from other future configuration options like alias analysis or autograd wrappers. The semantic problem behind this is that the dispatch key is a *kernel config parameter* and not an *operator config parameter* while things like autograd wrappers, alias info, and actually the kernel itself are *operator config parameters*. And while previously, the different kind of config parameters have been mixed, this diff now separates them. Before this change, it wouldn't have been well defined if you specified a dispatchKey together with an autogradWrapper or aliasInfo for example. // what is this supposed to do? static auto registry = torch::RegisterOperators() .op("my::op", torch::RegisterOperators::options() .aliasInfo(DEFAULT) .dispatchKey(CPUTensorId())); If we get more kernel config parameters in the future, we could introduce something like this static auto registry = torch::RegisterOperators() .op("my::op", torch::RegisterOperators::options() .kernel<Kernel>(torch::RegisterOperators::kernelOptions() .dispatchKey(CPUTensorId()) .otherConfig()); but that's overkill as long as dispatch keys are the only kernel config parameter, and we can introduce that later without breaking backwards compatibility. A nice side effect of this is that people can register multiple kernels to the same operator in the same `.op()` call: static auto registry = torch::RegisterOperators() .op("my::op", torch::RegisterOperators::options() .kernel<Kernel1>(CPUTensorId()) .kernel<Kernel2>(CUDATensorId())); Reviewed By: dzhulgakov Differential Revision: D15455790 fbshipit-source-id: 1c46bfe676dcacf74cf36bd3f5df3d2c32b8fb11