|
|
#pragma once
|
|
|
|
|
|
#include <ATen/core/ivalue.h>
|
|
|
#include <ATen/core/stack.h>
|
|
|
#include <ATen/core/boxing/KernelFunction.h>
|
|
|
#include <ATen/core/dispatch/Dispatcher.h>
|
|
|
#include <c10/util/Metaprogramming.h>
|
|
|
#include <torch/library.h>
|
|
|
|
|
|
namespace at::native {
|
|
|
|
|
|
|
|
|
|
|
|
TORCH_API void cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack, bool error_on_views = false,
|
|
|
c10::DispatchKey cpu_dispatch_key = c10::DispatchKey::CPU);
|
|
|
|
|
|
|
|
|
|
|
|
template<c10::KernelFunction::BoxedKernelFunction* fallback_fn, class Op, bool symint, class ReturnType, class... ParameterTypes>
|
|
|
struct _call_fallback_fn final {};
|
|
|
|
|
|
template<c10::KernelFunction::BoxedKernelFunction* fallback_fn, class Op, bool symint, class ReturnType, class... ParameterTypes>
|
|
|
struct _call_fallback_fn<fallback_fn, Op, symint, ReturnType(ParameterTypes...)> final {
|
|
|
static ReturnType call(typename c10::maybe_keep_symint<symint, ParameterTypes>::type... args) {
|
|
|
auto op = c10::Dispatcher::singleton()
|
|
|
|
|
|
.findSchemaOrThrow((const char*) Op::name, (const char*) Op::overload_name)
|
|
|
|
|
|
.typed<ReturnType (typename c10::maybe_keep_symint<symint, ParameterTypes>::type...)>();
|
|
|
return c10::impl::BoxedKernelWrapper<ReturnType (typename c10::maybe_keep_symint<symint, ParameterTypes>::type...)>::call(
|
|
|
c10::BoxedKernel::makeFromFunction<fallback_fn>(),
|
|
|
op,
|
|
|
c10::DispatchKeySet(),
|
|
|
|
|
|
args...
|
|
|
);
|
|
|
}
|
|
|
};
|
|
|
|
|
|
template<c10::KernelFunction::BoxedKernelFunction* fallback_fn, class Op>
|
|
|
using call_fallback_fn_symint = _call_fallback_fn<fallback_fn, Op, true, typename Op::schema>;
|
|
|
|
|
|
template<c10::KernelFunction::BoxedKernelFunction* fallback_fn, class Op>
|
|
|
using call_fallback_fn = _call_fallback_fn<fallback_fn, Op, false, typename Op::schema>;
|
|
|
|
|
|
}
|
|
|
|