#pragma once
|
|
#if !defined(CAFFE2_IS_XPLAT_BUILD)
|
#include <ATen/core/function_schema.h>
|
#include <ATen/core/grad_mode.h>
|
#include <ATen/core/op_registration/op_registration.h>
|
#include <torch/csrc/jit/script/function_schema_parser.h>
|
#include <vector>
|
|
namespace caffe2 {
|
namespace detail {
|
|
constexpr const char* PREALLOCATED_OUTPUT_ARGNAME =
|
"_caffe2_preallocated_outputs";
|
|
using _CallCaffe2OpFunc = c10::List<at::Tensor>(
|
const c10::FunctionSchema& schema,
|
std::vector<c10::IValue>&& inputs,
|
c10::List<at::Tensor>&& outputs);
|
|
template <class Caffe2Operator>
|
inline c10::List<at::Tensor> _call_caffe2_op(
|
const c10::FunctionSchema& schema,
|
std::vector<c10::IValue>&& inputs,
|
c10::List<at::Tensor>&& outputs) {
|
Caffe2Operator op(schema, std::move(inputs), std::move(outputs));
|
op.Run();
|
return std::move(op).move_newstyle_outputs();
|
}
|
|
// This function is inline in the hope that compilers optimizing for speed will
|
// inline it into call_caffe2_op_from_c10, allowing call_op to be inlined and
|
// avoiding the function pointer indirection, while compilers optimizing for
|
// binary size will keep it a separate function instead of inlining it into
|
// a template and will reuse the binary code of this function between ops.
|
// We measured and confirmed that binary size off the instagram ios app is
|
// reduced when having _call_caffe2_op_from_c10 separate from the templated
|
// call_caffe2_op_from_c10.
|
inline void _call_caffe2_op_from_c10(
|
c10::Stack* stack,
|
const c10::FunctionSchema& schema,
|
_CallCaffe2OpFunc* call_op) {
|
// precondition: on the stack, there's one IValue for each argument of the
|
// c10 schema. The last argument is an optional tensor list that
|
// (if not ivalue::None) contains a preallocated output tensor for each
|
// operator output.
|
|
// As an invariant, we don't want any autograd gradients to be tracked in
|
// Caffe2 operators.
|
at::NoGradGuard guard;
|
|
AT_ASSERT(
|
schema.arguments().size() != 0 &&
|
schema.arguments().back().type()->isSubtypeOf(
|
OptionalType::create(ListType::ofTensors())));
|
IValue preallocated_outputs = torch::jit::pop(*stack);
|
|
const size_t num_outputs = schema.returns().size();
|
const size_t num_inputs = schema.arguments().size() -
|
1; // -1 because the last argument is the list of preallocated tensors
|
|
c10::List<at::Tensor> outputs;
|
if (preallocated_outputs.isNone()) {
|
// either the schema doesn't support preallocated outputs or it does but
|
// they haven't been passed in. Pass a list of uninitialized tensors to
|
// the caffe2 operator as preallocated outputs.
|
outputs.resize(num_outputs);
|
} else {
|
AT_ASSERT(preallocated_outputs.isTensorList());
|
outputs = std::move(preallocated_outputs).toTensorList();
|
}
|
|
// TODO Avoid vector allocation. One idea would be to keep the std::vector
|
// instances in the cache.
|
std::vector<IValue> inputs = torch::jit::pop(*stack, num_inputs);
|
|
outputs = (*call_op)(schema, std::move(inputs), std::move(outputs));
|
|
for (size_t i = 0; i < outputs.size(); ++i) {
|
torch::jit::push(*stack, outputs.extract(i));
|
}
|
|
// postcondition: All inputs are cleared from the stack, there's now one
|
// IValue for each output which holds the result. This
|
// might reuse one of the preallocated tensors but doesn't have to.
|
}
|
|
template <const c10::FunctionSchema& (*Schema)(), class Caffe2Operator>
|
void call_caffe2_op_from_c10(
|
c10::OperatorKernel* functor,
|
c10::Stack* stack) {
|
_call_caffe2_op_from_c10(stack, Schema(), &_call_caffe2_op<Caffe2Operator>);
|
}
|
|
inline FunctionSchema make_function_schema_for_c10(const char* schema_str) {
|
#if defined(CAFFE2_IS_XPLAT_BUILD)
|
throw std::logic_error("We don't support registering c10 ops on mobile yet because the function schema parser isn't present in the mobile build.");
|
#else
|
c10::FunctionSchema parsed_schema = torch::jit::parseSchema(schema_str);
|
std::vector<c10::Argument> arguments = parsed_schema.arguments();
|
arguments.emplace_back(
|
PREALLOCATED_OUTPUT_ARGNAME,
|
c10::OptionalType::create(c10::ListType::ofTensors()),
|
nullopt,
|
IValue());
|
|
return FunctionSchema(
|
parsed_schema.name(),
|
parsed_schema.overload_name(),
|
std::move(arguments),
|
parsed_schema.returns(),
|
parsed_schema.is_vararg(),
|
parsed_schema.is_varret()
|
);
|
#endif
|
}
|
|
}
|
}
|
|
|
/**
|
* To register a caffe2 operator caffe2::MyOperator with the c10 dispatcher,
|
* call:
|
*
|
* In caffe2/operators/MyOperator.h:
|
*
|
* > C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(C10MyOperator) // C10MyOperator is the name
|
* // used by c10 for this operator
|
*
|
* In caffe2/operators/MyOperator.cc
|
*
|
* > C10_EXPORT_CAFFE2_OP_TO_C10_CPU (
|
* > C10MyOperator,
|
* > "_caffe2::C10MyOperator(Tensor input1, int argument2, float argument3) -> (Tensor output1, Tensor output2)"
|
* > caffe2::MyOperator<caffe2::CPUContext> // This is the caffe2 operator
|
* > // class template
|
* > )
|
*
|
* In caffe2/operators/MyOperator.cu
|
*
|
* > C10_EXPORT_CAFFE2_OP_TO_C10_CUDA(C10MyOperator ,
|
* caffe2::MyOperator<caffe2::CUDAContext>)
|
*
|
* Notes:
|
* - all macros must be defined in the top level namespace, not in namespace
|
* caffe2.
|
* - all operators must call C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10 and
|
* C10_EXPORT_CAFFE2_OP_TO_C10_CPU .
|
* - calling C10_EXPORT_CAFFE2_OP_TO_C10_CUDA is optional and can be omitted i f
|
* you don't want to expose the operator for CUDA operations.
|
* - caffe2 arguments must come after caffe2 inputs, in other words, any tensor
|
* inputs must precede any non-tensor inputs.
|
*
|
* More complex use cases:
|
* - If your operator has a variable number of input tensors, make the first (!)
|
* input an input of type TensorList. There must be no other tensor inputs.
|
*/
|
#define C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(OperatorName) \
|
namespace caffe2 { \
|
namespace _c10_ops { \
|
CAFFE2_API const FunctionSchema& schema_##OperatorName(); \
|
} \
|
}
|
|
#define C10_EXPORT_CAFFE2_OP_TO_C10_CPU( \
|
OperatorName, OperatorSchema, OperatorClass) \
|
/* Register the op schema with the c10 dispatcher */ \
|
namespace caffe2 { \
|
namespace _c10_ops { \
|
C10_EXPORT const FunctionSchema& schema_##OperatorName() { \
|
static const FunctionSchema schema = \
|
::caffe2::detail::make_function_schema_for_c10(OperatorSchema); \
|
return schema; \
|
} \
|
} \
|
} \
|
/* Register call_caffe2_op_from_c10 as a kernel with the c10 dispatcher */ \
|
static auto registry_##OperatorName##_##__COUNTER__ = \
|
::c10::RegisterOperators().op( \
|
::caffe2::_c10_ops::schema_##OperatorName(), \
|
::c10::RegisterOperators::options() \
|
.kernel( \
|
::c10::TensorTypeId::CPUTensorId, \
|
&::caffe2::detail::call_caffe2_op_from_c10< \
|
::caffe2::_c10_ops::schema_##OperatorName, \
|
OperatorClass>));
|
|
#define C10_EXPORT_CAFFE2_OP_TO_C10_CUDA(OperatorName, OperatorClass) \
|
/* Register call_caffe2_op_from_c10 as a kernel with the c10 dispatcher */ \
|
static auto registry_##OperatorName##_##__COUNTER__ = \
|
::c10::RegisterOperators().op( \
|
::caffe2::_c10_ops::schema_##OperatorName(), \
|
::c10::RegisterOperators::options() \
|
.kernel( \
|
::c10::TensorTypeId::CUDATensorId, \
|
&::caffe2::detail::call_caffe2_op_from_c10< \
|
::caffe2::_c10_ops::schema_##OperatorName, \
|
OperatorClass>));
|
|
// You should never manually call the C10_EXPORT_CAFFE2_OP_TO_C10_HIP macro .
|
// The C10_EXPORT_CAFFE2_OP_TO_C10_CUDA macro from above will be automatically
|
// rewritten to C10_EXPORT_CAFFE2_OP_TO_C10_HIP by hipify .
|
#define C10_EXPORT_CAFFE2_OP_TO_C10_HIP(OperatorName, OperatorClass) \
|
/* Register call_caffe2_op_from_c10 as a kernel with the c10 dispatcher */ \
|
static auto registry_##OperatorName##_##__COUNTER__ = \
|
::c10::RegisterOperators().op( \
|
::caffe2::_c10_ops::schema_##OperatorName(), \
|
::c10::RegisterOperators().options() \
|
.kernel( \
|
::c10::TensorTypeId::HIPTensorId, \
|
&::caffe2::detail::call_caffe2_op_from_c10< \
|
::caffe2::_c10_ops::schema_##OperatorName, \
|
OperatorClass>));
|
|
#else
|
// Don't use c10 dispatcher on mobile because of binary size
|
#define C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(OperatorName)
|
#define C10_EXPORT_CAFFE2_OP_TO_C10_CPU(OperatorName, OperatorSchema, OperatorClass)
|
#define C10_EXPORT_CAFFE2_OP_TO_C10_CUDA(OperatorName, OperatorClass)
|
#define C10_EXPORT_CAFFE2_OP_TO_C10_HIP(OperatorName, OperatorClass)
|
#endif
|
