pytorch
eebe0ee1 - [Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102874)

Commit
1 year ago
[Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102874) Summary: keys and change codegen to take ETKernelIndex We are adding support for dtype and dim order specialized kernel registration. This requires us to reorganize `BackendIndex` (which is a `Dict[DispatchKey, Dict[OperatorName, BackendMetadata]]`) to be `Dict[OperatorName, Dict[ETKernelKey, BackendMetadata]]`. This PR adds new data structures in order to support this change: * `ETKernelKey` to retrieve a certain kernel from the registry. * `ETKernelIndex`, the dictionary from operator name to kernel key to kernel mapping. Note that the codegen logic is not changed yet, we need subsequent diffs to actually generate code for different kernel keys. Test Plan: Added tests Reviewed By: Jack-Khuu Differential Revision: D46407096 Pull Request resolved: https://github.com/pytorch/pytorch/pull/102874 Approved by: https://github.com/Jack-Khuu, https://github.com/kirklandsign
Author
Committer
Parents
Loading