def _load_aten_op_executor_cpp_extension(verbosity):
from onnxruntime.training.ortmodule.torch_cpp_extensions import aten_op_executor
C.register_aten_op_executor(
str(aten_op_executor.execute_aten_operator_address()))
def load_aten_op_executor_cpp_extension():
_C.register_aten_op_executor(str(is_tensor_argument_address()),
str(execute_aten_operator_address()))
def _load_aten_op_executor_cpp_extension():
from onnxruntime.training.ortmodule.torch_cpp_extensions import aten_op_executor
C.register_aten_op_executor(
str(aten_op_executor.is_tensor_argument_address()),
str(aten_op_executor.execute_aten_operator_address()))
def _load_aten_op_executor_cpp_extension(verbosity, is_rocm_pytorch):
aten_op_executor_cpp_source = """
#include <torch/torch.h>
#include <ATen/DLConvertor.h>
#include <unordered_map>
#include <tuple>
#include <vector>
class ATenOperatorCache {
public:
static ATenOperatorCache& Instance() {
static ATenOperatorCache instance;
return instance;
}
std::shared_ptr<torch::jit::Operator> GetOperator(const std::string& op_name) {
if (ops_.find(op_name) == ops_.end()) {
auto& ops = torch::jit::getAllOperatorsFor(torch::jit::Symbol::fromQualString(op_name));
TORCH_INTERNAL_ASSERT(ops.size() == 1);
ops_[op_name] = ops.front();
}
return ops_.at(op_name);
}
private:
ATenOperatorCache() = default;
std::unordered_map<std::string, std::shared_ptr<torch::jit::Operator>> ops_;
};
// Some arguments of backward operator are not from forward operator's input or output,
// but need some processing. Since we cannot build such processing to ONNX graph for now,
// we are putting such processing code here if needed.
// Take embedding_backward as example:
// weight: embedding_backward(grad, indices, weight.size(0), padding_idx, scale_grad_by_freq, sparse)
// the 3rd argument (index 2) is weight.size(0), we add this processing here.
using TensorTransformFunc = std::function<c10::IValue(const at::Tensor&)>;
static const TensorTransformFunc embedding_num_weights = [](const at::Tensor& tensor) {
return c10::IValue(tensor.size(0));
};
static const std::unordered_map<std::string, std::unordered_map<size_t, TensorTransformFunc>> TENSOR_TRANSFORM_FUNCS = {
{"aten::embedding_backward", {{2, embedding_num_weights}}},
};
template <typename T>
void SetIValueArguments(const std::vector<std::tuple<size_t, T>>& raw_arguments,
std::vector<c10::IValue>& ivalue_arguments) {
for (size_t i = 0; i < raw_arguments.size(); i++) {
size_t index = std::get<0>(raw_arguments[i]);
TORCH_INTERNAL_ASSERT(index < ivalue_arguments.size());
ivalue_arguments[index] = c10::IValue(std::get<1>(raw_arguments[i]));
}
}
// TODO: Add more argument types, such as list type.
std::vector<DLManagedTensor*> ExecuteATenOperator(
const char* op_name, const std::vector<std::tuple<size_t, DLManagedTensor*>>& tensor_arguments,
const std::vector<std::tuple<size_t, int64_t>>& int_arguments,
const std::vector<std::tuple<size_t, float>>& float_arguments,
const std::vector<std::tuple<size_t, bool>>& bool_arguments) {
std::string op_name_str(op_name);
std::shared_ptr<torch::jit::Operator> op = ATenOperatorCache::Instance().GetOperator(op_name_str);
// TODO: need to handle optional argument and arguments with default values.
std::vector<c10::IValue> arguments;
arguments.resize(op->schema().arguments().size());
for (size_t i = 0; i < tensor_arguments.size(); i++) {
size_t index = std::get<0>(tensor_arguments[i]);
at::Tensor tensor = at::fromDLPack(std::get<1>(tensor_arguments[i]));
bool has_transform_func = false;
if (TENSOR_TRANSFORM_FUNCS.find(op_name_str) != TENSOR_TRANSFORM_FUNCS.end()) {
const auto& transform_funcs = TENSOR_TRANSFORM_FUNCS.at(op_name_str);
if (transform_funcs.find(index) != transform_funcs.end()) {
arguments[index] = transform_funcs.at(index)(tensor);
has_transform_func = true;
}
}
if (!has_transform_func) {
arguments[index] = c10::IValue(tensor);
}
}
SetIValueArguments<int64_t>(int_arguments, arguments);
SetIValueArguments<float>(float_arguments, arguments);
SetIValueArguments<bool>(bool_arguments, arguments);
torch::jit::Stack stack;
for (size_t i = 0; i < arguments.size(); i++) {
torch::jit::push(stack, arguments[i]);
}
op->getOperation()(&stack);
// TODO: need to handle multiple-tensor outputs.
at::Tensor output;
torch::jit::pop(stack, output);
std::vector<DLManagedTensor*> result;
result.emplace_back(at::toDLPack(output));
return result;
}
size_t execute_aten_operator_address() { return reinterpret_cast<size_t>(&ExecuteATenOperator); }
"""
aten_op_executor_cpp_extension = load_inline(
name='inline_extension_aten_op_executor',
cpp_sources=[aten_op_executor_cpp_source],
extra_cflags=['-D__HIP_PLATFORM_HCC__=1' if is_rocm_pytorch else ''],
functions=['execute_aten_operator_address'],
verbose=verbosity,
with_cuda=True)
C.register_aten_op_executor(
str(aten_op_executor_cpp_extension.execute_aten_operator_address()))