[AOTI] Support fallback ops not in inductor_fallback_ops (#131247)

Summary: For aten ops that are not listed in inductor_fallback_ops, AOTI will use proxy executor to execute them instead of erroring out as missing C shim implementation.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131247
Approved by: https://github.com/angelayi
This commit is contained in:
Bin Bao
2024-07-21 14:04:39 -07:00
committed by PyTorch MergeBot
parent aeca9845a6
commit 3ce6f61416

View File

@ -5807,12 +5807,18 @@ class FallbackKernel(ExternKernelAlloc):
# Aten Fallback Ops
assert isinstance(kernel, torch._ops.OpOverload)
if V.graph.cpp_wrapper:
from torchgen.aoti.fallback_ops import inductor_fallback_ops
if (
config.is_fbcode()
and kernel not in has_c_shim
# C shim v2 is torchgen-ed, which should cover all aten ops.
# If you do hit a missed op, please update gen_aoti_c_shim.py.
and config.c_shim_version == "1"
) or (
config.abi_compatible
and config.c_shim_version == "2"
and str(kernel) not in inductor_fallback_ops
):
log.warning(
"%s is missing a c-shim implementation, using proxy executor as fallback",
@ -5820,8 +5826,7 @@ class FallbackKernel(ExternKernelAlloc):
)
self.use_runtime_dispatch = True
self.set_cpp_kernel(kernel)
else:
self.python_kernel_name = str(kernel)
self.python_kernel_name = f"torch.ops.{kernel}"
elif kernel.namespace == "_quantized": # type: ignore[union-attr]
# Internal Quantized Fallback Ops
assert isinstance(kernel, torch._ops.OpOverload)