/usr/local/lib64/python3.6/site-packages/torch/include/ATen/core/boxing
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/core/boxing/KernelFunction_impl.h (9784B)
#include
#include
#include
#include
namespace c10 {
inline KernelFunction::KernelFunction()
: functor_(nullptr)
, boxed_kernel_func_(nullptr)
, unboxed_kernel_func_(nullptr)
{}
inline KernelFunction::KernelFunction(std::unique_ptr functor, InternalBoxedKernelFunction* boxed_kernel_func, void* unboxed_kernel_func)
: functor_(std::move(functor))
, boxed_kernel_func_(boxed_kernel_func)
, unboxed_kernel_func_(unboxed_kernel_func)
{}
template
inline void KernelFunction::make_boxed_function(OperatorKernel*, const OperatorHandle& opHandle, DispatchKeySet, Stack* stack) {
// Note that we're dropping the DispatchKeySet argument.
// See Note [Plumbing Keys Through The Dispatcher 2] for details.
func(opHandle, stack);
}
inline bool KernelFunction::isValidUnboxed() const {
return unboxed_kernel_func_ != nullptr;
}
template
inline void KernelFunction::make_boxed_function(OperatorKernel*, const OperatorHandle& opHandle, DispatchKeySet ks, Stack* stack) {
// See Note [Plumbing Keys Through The Dispatcher 2] for details.
func(opHandle, ks, stack);
}
inline bool KernelFunction::isValid() const {
return boxed_kernel_func_ != nullptr;
}
inline bool KernelFunction::isFallthrough() const {
return boxed_kernel_func_ == &fallthrough_kernel;
}
inline void KernelFunction::callBoxed(const OperatorHandle& opHandle, DispatchKeySet dispatchKeySet, Stack* stack) const {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
boxed_kernel_func_ != nullptr,
"Tried to call KernelFunction::callBoxed() on an uninitialized KernelFunction."
);
(*boxed_kernel_func_)(functor_.get(), opHandle, dispatchKeySet, stack);
}
template
inline Return callUnboxedKernelFunction(void* unboxed_kernel_func, OperatorKernel* functor, DispatchKeySet dispatchKeySet, Args&&... args) {
using ActualSignature = Return (OperatorKernel*, DispatchKeySet, Args...);
ActualSignature* func = reinterpret_cast(unboxed_kernel_func);
return (*func)(functor, dispatchKeySet, std::forward(args)...);
}
template
C10_ALWAYS_INLINE Return KernelFunction::call(const OperatorHandle& opHandle, DispatchKeySet dispatchKeySet, Args... args) const {
// note: Args above is intentionally not Args&&. We don't want perfect
// forwarding, which would require Args to be deduced, but instead we
// want callers to explicitly specify the Args.
if (C10_LIKELY(unboxed_kernel_func_ != nullptr)) {
return callUnboxedKernelFunction(unboxed_kernel_func_, functor_.get(), dispatchKeySet, std::forward(args)...);
}
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
boxed_kernel_func_ != nullptr,
"Tried to call KernelFunction::call() on an uninitialized KernelFunction."
);
return impl::BoxedKernelWrapper::call(
boxed_kernel_func_,
functor_.get(),
opHandle,
dispatchKeySet,
std::forward(args)...
);
}
template
inline KernelFunction KernelFunction::makeFromBoxedFunction() {
return KernelFunction(
nullptr, // no functor_ object
&make_boxed_function,
nullptr // no unboxed function pointer
);
}
template
inline KernelFunction KernelFunction::makeFromBoxedFunction() {
return KernelFunction(
nullptr, // no functor_ object
&make_boxed_function,
nullptr // no unboxed function pointer
);
}
inline KernelFunction KernelFunction::makeFallthrough() {
return KernelFunction(
nullptr, // no functor_ object
&fallthrough_kernel,
nullptr // no unboxed function pointer
);
}
inline KernelFunction KernelFunction::makeAmbiguousAutogradOther() {
return KernelFunction(
nullptr, // no functor_ object
&ambiguous_autogradother_kernel,
nullptr // no unboxed function pointer
);
}
inline KernelFunction KernelFunction::makeNamedNotSupported() {
return KernelFunction(
nullptr, // no functor_ object
&named_not_supported_kernel,
nullptr // no unboxed function pointer
);
}
template
inline KernelFunction KernelFunction::makeFromUnboxedFunctor(std::unique_ptr kernelFunctor) {
#ifndef NDEBUG
// This assertion is costly for build time so it's debug-gated.
static_assert(guts::is_functor::value, "Tried to call KernelFunction::makeFromUnboxedFunctor but the argument is not a functor.");
#endif
static_assert(std::is_base_of::value, "Tried to call KernelFunction::makeFromUnboxedFunctor, but the functor doesn't inherit from c10::OperatorKernel. Please have the functor inherit from it.");
return KernelFunction(
std::move(kernelFunctor),
&impl::make_boxed_from_unboxed_functor::call,
reinterpret_cast(&impl::wrap_kernel_functor_unboxed::call)
);
}
template
inline KernelFunction KernelFunction::makeFromBoxedFunctor(std::unique_ptr kernelFunctor) {
static_assert(std::is_base_of::value, "Tried to call KernelFunction::makeFromBoxedFunctor, but the functor doesn't inherit from c10::OperatorKernel. Please have the functor inherit from it.");
return KernelFunction(
std::move(kernelFunctor),
[](OperatorKernel* kernel, const OperatorHandle& op, DispatchKeySet ks, Stack* stack) {
(*static_cast(kernel))(op, ks, stack);
},
nullptr // no unboxed function pointer
);
}
template
inline KernelFunction KernelFunction::makeFromUnboxedFunction(FuncPtr func_ptr) {
static_assert(is_compile_time_function_pointer::value, "Tried to call KernelFunction::makeFromUnboxedFunction with an invalid parameter. It must be a function pointer created with TORCH_FN.");
static_assert(!std::is_same::value, "Tried to call KernelFunction::makeFromUnboxedFunction with a boxed function pointer. Please use KernelFunction::makeFromBoxedFunction instead.");
static_assert(FuncPtr::func_ptr() != nullptr, "Kernel function cannot be nullptr");
#if !defined(C10_MOBILE)
return makeFromUnboxedFunctor::type>(
guts::make_unique_base::type>()
);
#else
// On mobile, we rather want to optimize for binary size than for performance,
// so let's not inline the kernel into the wrapper but use makeFromUnboxedRuntimeFunction
// instead.
return makeFromUnboxedRuntimeFunction(func_ptr.func_ptr());
#endif
}
template
inline KernelFunction KernelFunction::makeFromUnboxedRuntimeFunction(FuncType* func) {
static_assert(guts::is_function_type::value, "Tried to call KernelFunction::makeFromUnboxedRuntimeFunction with a non-function type.");
static_assert(!std::is_same::value, "Tried to call KernelFunction::makeFromUnboxedRuntimeFunction with a boxed function pointer. Please use KernelFunction::makeFromBoxedFunction instead.");
TORCH_INTERNAL_ASSERT(func != nullptr, "Kernel function cannot be nullptr");
return makeFromUnboxedFunctor>>(
guts::make_unique_base>>(func)
);
}
template
inline std::enable_if_t>::value, KernelFunction> KernelFunction::makeFromUnboxedLambda(Lambda&& lambda) {
static_assert(guts::is_functor>::value, "Tried to call KernelFunction::makeFromUnboxedLambda with a non-lambda type.");
#if !defined(C10_MOBILE)
return makeFromUnboxedFunctor>>(
guts::make_unique_base>>(std::forward(lambda))
);
#else
// On mobile, we rather want to optimize for binary size than for performance,
// so let's not inline the kernel into the wrapper but use makeFromUnboxedRuntimeFunction
// instead.
using FuncType = typename guts::infer_function_traits_t>::func_type;
return makeFromUnboxedRuntimeFunction(lambda);
#endif
}
template
inline std::enable_if_t>::value, KernelFunction> KernelFunction::makeFromUnboxedLambda(Lambda&& lambda) {
static_assert(guts::is_functor>::value, "Tried to call KernelFunction::makeFromUnboxedLambda with a non-lambda type.");
return makeFromUnboxedFunctor>>(
guts::make_unique_base>>(std::forward(lambda))
);
}
}