pytorch/torchgen/executorch
Larry Liu 019c38624c [Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102565)
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.

Differential Revision: [D46206339](https://our.internmc.facebook.com/intern/diff/D46206339/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/102565
Approved by: https://github.com/Jack-Khuu
2023-05-31 09:41:36 +00:00
..
api [Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102565) 2023-05-31 09:41:36 +00:00
__init__.py [torchgen] Introduce Executorch types and signatures (#90781) 2022-12-14 20:13:04 +00:00
model.py [Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102565) 2023-05-31 09:41:36 +00:00