/usr/local/lib64/python3.6/site-packages/torch/include/ATen/native
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/CPUFallback.h (2404B)
#pragma once
#include
#include
#include
#include
#include
#include
namespace at { namespace native {
// This function implements a boxed fallback to CPU.
// External backends can add their own custom logging on top if it to customize their own CPU fallbacks.
TORCH_API void cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack);
// This is a helper function that backends can use to directly call their boxed CPU fallback
// TODO: update and add a usage example after https://github.com/pytorch/pytorch/pull/58092 lands.
template
struct _call_fallback_fn final {};
template
struct _call_fallback_fn final {
static_assert(std::is_same::return_type>::value,
"Return type mismatch");
static_assert(std::is_same, typename guts::infer_function_traits_t::parameter_types>::value,
"Parameter types mismatch");
static ReturnType call(ParameterTypes... args) {
auto op = c10::Dispatcher::singleton()
// TODO: figure out how to make compiler happy without dynamic casts
.findSchemaOrThrow((const char*) Op::name, (const char*) Op::overload_name)
//.findSchemaOrThrow("a", "b")
.typed();
return c10::impl::BoxedKernelWrapper::call(
c10::KernelFunction::make_boxed_function,
nullptr,
op,
c10::DispatchKeySet(), // we know that the cpu_fallback doesn't use the dispatch keyset.
//std::forward(args...)
// TODO: get std::forward<> to work
args...
);
}
};
template
using call_fallback_fn = _call_fallback_fn;
} // namespace native
} // namespace at