frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[AOTI] Add more fallback ops (#126720)

Browse Source 
Summary: These ops are either in either unit tests or TorchBench. Fixes https://github.com/pytorch/pytorch/issues/122050

Pull Request resolved: https://github.com/pytorch/pytorch/pull/126720
Approved by: https://github.com/chenyang78
This commit is contained in:
Bin Bao
2024-05-21 11:23:36 -07:00
committed by PyTorch MergeBot
parent 0d17aae242
commit 19cd4484ec
6 changed files with 17 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
torchgen/gen_aoti_c_shim.py
View File

@ -333,6 +333,8 @@ def get_backend_index_for_aoti(
backend_index = backend_indices[
DispatchKey.CompositeExplicitAutogradNonFunctional
]
elif backend_indices[DispatchKey.CompositeImplicitAutograd].has_kernel(func):
backend_index = backend_indices[DispatchKey.CompositeImplicitAutograd]
return backend_index
@ -471,6 +473,7 @@ extern "C" {{
#include <ATen/{str(dispatch_key)}Functions.h>
#include <ATen/CompositeExplicitAutogradFunctions.h>
#include <ATen/CompositeExplicitAutogradNonFunctionalFunctions.h>
#include <ATen/CompositeImplicitAutogradFunctions.h>
#else
{includes}
#endif