mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-24 19:43:35 +00:00
Support OrtMemTypeCPUInput for ATenOp/ATenOpGrad (#8116)
This commit is contained in:
Vincent Wang
GitHub
parent
b42e7d2c78
commit
88ec95ea96
13 changed files with 320 additions and 168 deletions
|
|
@ -578,7 +578,7 @@ class PlannerImpl {
|
|||
|
||||
const KernelCreateInfo& kernel_create_info = GetKernelCreateInfo(kernel_create_info_map_, node.Index());
|
||||
|
||||
if (kernel_create_info.kernel_def->IsInputOnCpu(input_index))
|
||||
if (utils::IsInputOnCpu(node, &kernel_create_info, input_index))
|
||||
// weights are not output from any node, so it's OK to put its location on CPU provider
|
||||
return execution_providers_.GetDefaultCpuMemoryInfo();
|
||||
return p_provider->GetAllocator(0, OrtMemTypeDefault)->Info();
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@
|
|||
#include "core/framework/TensorSeq.h"
|
||||
#ifdef ENABLE_TRAINING
|
||||
#include "core/framework/orttraining_partial_executor.h"
|
||||
#include "orttraining/training_ops/cpu/aten_ops/aten_op_config.h"
|
||||
#endif
|
||||
|
||||
namespace ONNX_NAMESPACE {
|
||||
|
|
@ -729,5 +730,26 @@ common::Status VerifyInputTensorsAllocatedContiguously(OpKernelContext* context)
|
|||
}
|
||||
#endif
|
||||
|
||||
bool IsInputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index) {
|
||||
if (p_kci && p_kci->kernel_def->IsInputOnCpu(index)) {
|
||||
return true;
|
||||
}
|
||||
|
||||
#ifdef ENABLE_TRAINING
|
||||
if (node.GetExecutionProviderType() == kCudaExecutionProvider &&
|
||||
(node.OpType() == "ATenOp" || node.OpType() == "ATenOpGrad")) {
|
||||
const std::string name = node.GetAttributes().at("name").s();
|
||||
const auto* op_config_ptr = contrib::aten_ops::ATenOperatorConfigs::Instance().GetConfig(name);
|
||||
if (op_config_ptr) {
|
||||
return op_config_ptr->IsInputOnCpu(index, node.OpType() == "ATenOpGrad");
|
||||
}
|
||||
}
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(node);
|
||||
#endif
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
} // namespace utils
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ class KernelRegistryManager;
|
|||
class IExecutionProvider;
|
||||
class Node;
|
||||
class Tensor;
|
||||
struct KernelCreateInfo;
|
||||
|
||||
namespace logging {
|
||||
class Logger;
|
||||
|
|
@ -87,6 +88,8 @@ common::Status ExecuteSubgraph(const SessionState& session_state, const FeedsFet
|
|||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
ExecutionMode execution_mode, const bool& terminate_flag, const logging::Logger& logger);
|
||||
|
||||
bool IsInputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index);
|
||||
|
||||
template <typename T>
|
||||
constexpr ONNXTensorElementDataType GetONNXTensorElementDataType() {
|
||||
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
|
||||
|
|
|
|||
|
|
@ -168,7 +168,7 @@ void TransformerMemcpyImpl::ProcessDefs(onnxruntime::Node& node, const KernelReg
|
|||
|
||||
bool is_implicit_input = false;
|
||||
auto process_inputs =
|
||||
[this, &kci, &initializers_consumed, &is_implicit_input](const onnxruntime::NodeArg& arg, size_t index) {
|
||||
[this, &node, &kci, &initializers_consumed, &is_implicit_input](const onnxruntime::NodeArg& arg, size_t index) {
|
||||
// check if this NodeArg is an initializer defined in current outer graph level
|
||||
const auto* initializer_tensor_proto = GetInitializer(graph_, arg.Name(), true);
|
||||
if (initializer_tensor_proto != nullptr) {
|
||||
|
|
@ -179,7 +179,7 @@ void TransformerMemcpyImpl::ProcessDefs(onnxruntime::Node& node, const KernelReg
|
|||
// flow op (Loop, Scan, If) to do the necessary copy if the input crosses different provider.
|
||||
// PlannerImpl::ComputeUseCounts has matching logic so the allocation plan does the same thing
|
||||
if (!is_implicit_input) {
|
||||
if (kci && kci->kernel_def->IsInputOnCpu(index)) {
|
||||
if (utils::IsInputOnCpu(node, kci, index)) {
|
||||
non_provider_input_defs_.insert(&arg);
|
||||
} else {
|
||||
provider_input_defs_.insert(&arg);
|
||||
|
|
@ -251,7 +251,7 @@ void TransformerMemcpyImpl::BuildDefsMapping(const onnxruntime::NodeArg* arg, co
|
|||
const KernelCreateInfo* kci = nullptr;
|
||||
kernel_registries.SearchKernelRegistry(it, &kci);
|
||||
if (arg_input_index != -1) {
|
||||
if (!kci || !kci->kernel_def->IsInputOnCpu(arg_input_index)) provider_input_nodes_[arg].insert(&it);
|
||||
if (!kci || !utils::IsInputOnCpu(it, kci, arg_input_index)) provider_input_nodes_[arg].insert(&it);
|
||||
}
|
||||
if (arg_output_index != -1) {
|
||||
if (!kci || !kci->kernel_def->IsOutputOnCpu(arg_output_index)) provider_output_nodes_[arg].insert(&it);
|
||||
|
|
@ -338,8 +338,8 @@ bool TransformerMemcpyImpl::ProcessInitializers(const KernelRegistryManager& ker
|
|||
if (kci == nullptr) continue;
|
||||
if (kci->kernel_def == nullptr) continue;
|
||||
onnxruntime::Node::ForEachWithIndex(p_node->InputDefs(),
|
||||
[kci, &dup_replacements](const onnxruntime::NodeArg& arg, size_t index) {
|
||||
if (kci->kernel_def->IsInputOnCpu(index)) dup_replacements.erase(&arg);
|
||||
[kci, &p_node, &dup_replacements](const onnxruntime::NodeArg& arg, size_t index) {
|
||||
if (utils::IsInputOnCpu(*p_node, kci, index)) dup_replacements.erase(&arg);
|
||||
return Status::OK();
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -1681,6 +1681,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetMinMaxGradient) {
|
|||
}
|
||||
|
||||
IMPLEMENT_GRADIENT_BUILDER(GetATenOpGradient) {
|
||||
std::vector<NodeDef> result;
|
||||
const auto& src_attrs = SrcNodeAttributes();
|
||||
std::vector<AttributeProto> attrs;
|
||||
ORT_ENFORCE(utils::HasString(src_attrs.at("name")));
|
||||
|
|
@ -1699,18 +1700,10 @@ IMPLEMENT_GRADIENT_BUILDER(GetATenOpGradient) {
|
|||
std::vector<ArgDef> output_args;
|
||||
|
||||
for (const auto& config : op_config.backward_input_source_configs) {
|
||||
size_t index = config.second;
|
||||
switch (config.first) {
|
||||
case contrib::aten_ops::GRAD_OUTPUT:
|
||||
input_args.emplace_back(GO(index));
|
||||
break;
|
||||
case contrib::aten_ops::FORWARD_INPUT:
|
||||
input_args.emplace_back(I(index));
|
||||
break;
|
||||
case contrib::aten_ops::FORWARD_OUTPUT:
|
||||
input_args.emplace_back(O(index));
|
||||
break;
|
||||
}
|
||||
ArgDef source_arg_def = config.kind == contrib::aten_ops::GRAD_OUTPUT ? GO(config.index)
|
||||
: config.kind == contrib::aten_ops::FORWARD_INPUT ? I(config.index)
|
||||
: O(config.index);
|
||||
input_args.emplace_back(HandleATenOpGradInput(source_arg_def, config.transform_func, result));
|
||||
}
|
||||
|
||||
for (size_t index : op_config.gradient_input_indices) {
|
||||
|
|
@ -1724,7 +1717,8 @@ IMPLEMENT_GRADIENT_BUILDER(GetATenOpGradient) {
|
|||
}
|
||||
|
||||
attrs.emplace_back(MakeAttribute("output_types", grad_output_types));
|
||||
return std::vector<NodeDef>{NodeDef(OpDef{"ATenOpGrad", kMSDomain, 1}, input_args, output_args, attrs)};
|
||||
result.emplace_back(NodeDef(OpDef{"ATenOpGrad", kMSDomain, 1}, input_args, output_args, attrs));
|
||||
return result;
|
||||
}
|
||||
|
||||
IMPLEMENT_GRADIENT_BUILDER(GetPythonOpGradient) {
|
||||
|
|
|
|||
|
|
@ -293,5 +293,29 @@ std::vector<NodeDef> GradientBuilderBase::GetBiasGeluGradNodes(
|
|||
return result;
|
||||
}
|
||||
|
||||
ArgDef GradientBuilderBase::HandleATenOpGradInput(const ArgDef& source_arg_def, const std::string& transform_func,
|
||||
std::vector<NodeDef>& output) const {
|
||||
ArgDef target_arg_def = source_arg_def;
|
||||
if (transform_func == "sizes()") {
|
||||
target_arg_def = IA("Shape_" + source_arg_def.name);
|
||||
output.emplace_back(NodeDef("Shape", {source_arg_def}, {target_arg_def}));
|
||||
} else if (transform_func.find("size(") == 0) {
|
||||
int index = std::stoi(transform_func.substr(5, transform_func.length() - 6));
|
||||
NodeDef index_const_node =
|
||||
ConstantScalarNode(static_cast<int64_t>(index), {1}, Name("Constant_" + std::to_string(index)));
|
||||
ArgDef index_arg_def = index_const_node.output_args[0];
|
||||
output.emplace_back(index_const_node);
|
||||
ArgDef shape_arg_def = IA("Shape_" + source_arg_def.name);
|
||||
target_arg_def = IA("Dim_" + std::to_string(index) + "_" + source_arg_def.name);
|
||||
output.emplace_back(NodeDef("Shape", {source_arg_def}, {shape_arg_def}));
|
||||
output.emplace_back(NodeDef("Gather", {shape_arg_def, index_arg_def}, {target_arg_def},
|
||||
{ONNX_NAMESPACE::MakeAttribute("axis", int64_t(0))}));
|
||||
} else {
|
||||
ORT_ENFORCE(transform_func == "", "Failed to build gradient graph for ", transform_func);
|
||||
}
|
||||
|
||||
return target_arg_def;
|
||||
}
|
||||
|
||||
} // namespace training
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -294,6 +294,9 @@ class GradientBuilderBase {
|
|||
|
||||
const std::string& NodeName() const { return node_->Name(); }
|
||||
|
||||
ArgDef HandleATenOpGradInput(const ArgDef& source_arg_def, const std::string& transform_func,
|
||||
std::vector<NodeDef>& output) const;
|
||||
|
||||
private:
|
||||
friend class GradientGraphBuilder;
|
||||
|
||||
|
|
|
|||
|
|
@ -22,19 +22,13 @@ class ATenOperatorCache {
|
|||
|
||||
const ATenOperator& GetOperator(const std::string& op_name) {
|
||||
if (ops_.find(op_name) == ops_.end()) {
|
||||
auto& ops = torch::jit::getAllOperatorsFor(torch::jit::Symbol::fromQualString(op_name));
|
||||
bool found = false;
|
||||
// Some op name can get multiple ops with different overload names,
|
||||
// we are using the one with empty overload name.
|
||||
c10::OperatorName full_name(op_name, "");
|
||||
auto op = torch::jit::findOperatorFor(full_name);
|
||||
TORCH_INTERNAL_ASSERT(op);
|
||||
ATenOperator aten_op;
|
||||
for (auto op : ops) {
|
||||
// Some op name can get multiple ops with different overload names,
|
||||
// we are using the one without overload name.
|
||||
if (op->schema().overload_name() == "") {
|
||||
aten_op.op = op;
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
TORCH_INTERNAL_ASSERT(found);
|
||||
aten_op.op = op;
|
||||
const auto& schema = aten_op.op->schema();
|
||||
aten_op.argument_size = schema.arguments().size();
|
||||
for (const auto& argument : schema.arguments()) {
|
||||
|
|
@ -54,27 +48,6 @@ class ATenOperatorCache {
|
|||
std::unordered_map<std::string, ATenOperator> 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 TensorTransformFunc unfold_input_sizes = [](const at::Tensor& tensor) {
|
||||
return c10::IValue(tensor.sizes());
|
||||
};
|
||||
|
||||
static const std::unordered_map<std::string, std::unordered_map<size_t, TensorTransformFunc>> TENSOR_TRANSFORM_FUNCS = {
|
||||
{"aten::embedding_backward", {{2, embedding_num_weights}}},
|
||||
{"aten::unfold_backward", {{1, unfold_input_sizes}}},
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void SetIValueArguments(const std::vector<std::pair<size_t, T>>& raw_arguments,
|
||||
const std::vector<bool>& is_optional_arguments, std::vector<c10::IValue>& ivalue_arguments) {
|
||||
|
|
@ -120,21 +93,8 @@ std::vector<DLManagedTensor*> ExecuteATenOperator(
|
|||
for (const auto& tensor_argument : tensor_arguments) {
|
||||
size_t index = tensor_argument.first;
|
||||
at::Tensor tensor = at::fromDLPack(tensor_argument.second);
|
||||
bool has_transform_func = false;
|
||||
|
||||
auto op_it = TENSOR_TRANSFORM_FUNCS.find(op_name_str);
|
||||
if (op_it != TENSOR_TRANSFORM_FUNCS.end()) {
|
||||
auto func_it = op_it->second.find(index);
|
||||
if (func_it != op_it->second.end()) {
|
||||
arguments[index] = func_it->second(tensor);
|
||||
has_transform_func = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!has_transform_func) {
|
||||
arguments[index] =
|
||||
aten_op.is_optional_arguments[index] ? c10::IValue(c10::optional<at::Tensor>(tensor)) : c10::IValue(tensor);
|
||||
}
|
||||
arguments[index] =
|
||||
aten_op.is_optional_arguments[index] ? c10::IValue(c10::optional<at::Tensor>(tensor)) : c10::IValue(tensor);
|
||||
}
|
||||
|
||||
SetIValueArguments<int64_t>(int_arguments, aten_op.is_optional_arguments, arguments);
|
||||
|
|
|
|||
|
|
@ -10,11 +10,32 @@ namespace test {
|
|||
using namespace contrib::aten_ops;
|
||||
|
||||
namespace {
|
||||
void CompareArgumentConfigs(const std::vector<ArgumentConfig>& configs, const std::vector<ArgumentConfig>& others) {
|
||||
EXPECT_TRUE(configs.size() == others.size());
|
||||
for (size_t i = 0; i < configs.size(); i++) {
|
||||
const auto& config = configs[i];
|
||||
const auto& other = others[i];
|
||||
EXPECT_TRUE(config.kind == other.kind && config.name == other.name && config.is_cpu_tensor == other.is_cpu_tensor &&
|
||||
config.is_optional == other.is_optional);
|
||||
}
|
||||
}
|
||||
|
||||
void CompareBackwardInputSourceConfigs(const std::vector<BackwardInputSourceConfig>& configs,
|
||||
const std::vector<BackwardInputSourceConfig>& others) {
|
||||
EXPECT_TRUE(configs.size() == others.size());
|
||||
for (size_t i = 0; i < configs.size(); i++) {
|
||||
const auto& config = configs[i];
|
||||
const auto& other = others[i];
|
||||
EXPECT_TRUE(config.kind == other.kind && config.index == other.index &&
|
||||
config.transform_func == other.transform_func);
|
||||
}
|
||||
}
|
||||
|
||||
void Compare(const ATenOperatorConfig& config, const ATenOperatorConfig& other) {
|
||||
EXPECT_TRUE(config.op_name == other.op_name && config.backward_op_name == other.backward_op_name);
|
||||
EXPECT_TRUE(config.forward_argument_configs == other.forward_argument_configs);
|
||||
EXPECT_TRUE(config.backward_argument_configs == other.backward_argument_configs);
|
||||
EXPECT_TRUE(config.backward_input_source_configs == other.backward_input_source_configs);
|
||||
CompareArgumentConfigs(config.forward_argument_configs, other.forward_argument_configs);
|
||||
CompareArgumentConfigs(config.backward_argument_configs, other.backward_argument_configs);
|
||||
CompareBackwardInputSourceConfigs(config.backward_input_source_configs, other.backward_input_source_configs);
|
||||
EXPECT_TRUE(config.forward_output_type_infer_configs == other.forward_output_type_infer_configs);
|
||||
EXPECT_TRUE(config.gradient_input_indices == other.gradient_input_indices);
|
||||
EXPECT_TRUE(config.default_int_values == other.default_int_values);
|
||||
|
|
@ -37,17 +58,17 @@ TEST(ATenOpConfigTest, ValidATenOpConfig) {
|
|||
ATenOperatorConfig expected;
|
||||
expected.op_name = "at::a";
|
||||
expected.backward_op_name = "at::a_bw";
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(TENSOR, "input", false));
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(TENSOR, "indices", false));
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(INT, "p", false));
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(BOOL, "scale", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(TENSOR, "grad_output", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(TENSOR, "indices", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(INT, "p", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(FLOAT, "d", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(FLOAT, "e", false));
|
||||
expected.backward_input_source_configs.emplace_back(std::make_pair(GRAD_OUTPUT, 0UL));
|
||||
expected.backward_input_source_configs.emplace_back(std::make_pair(FORWARD_INPUT, 1UL));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "input", false, false));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "indices", false, false));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(INT, "p", false, false));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(BOOL, "scale", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "grad_output", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "indices", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(INT, "p", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(FLOAT, "d", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(FLOAT, "e", false, false));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(GRAD_OUTPUT, 0UL, ""));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(FORWARD_INPUT, 1UL, ""));
|
||||
expected.forward_output_type_infer_configs.emplace_back(std::make_pair(PROPAGATE_FROM_INPUT, 0));
|
||||
expected.gradient_input_indices.emplace_back(0UL);
|
||||
expected.default_int_values["p"] = -1;
|
||||
|
|
@ -66,14 +87,14 @@ TEST(ATenOpConfigTest, ValidATenOpConfig) {
|
|||
ATenOperatorConfig expected;
|
||||
expected.op_name = "at::b";
|
||||
expected.backward_op_name = "at::b_bw";
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(TENSOR, "weight", false));
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(TENSOR, "bias", true));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(TENSOR, "grad_r1", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(TENSOR, "grad_r2", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(TENSOR, "r2", false));
|
||||
expected.backward_input_source_configs.emplace_back(std::make_pair(GRAD_OUTPUT, 0UL));
|
||||
expected.backward_input_source_configs.emplace_back(std::make_pair(GRAD_OUTPUT, 1UL));
|
||||
expected.backward_input_source_configs.emplace_back(std::make_pair(FORWARD_OUTPUT, 1UL));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "weight", false, false));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "bias", false, true));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "grad_r1", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "grad_r2", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "r2", false, false));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(GRAD_OUTPUT, 0UL, ""));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(GRAD_OUTPUT, 1UL, ""));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(FORWARD_OUTPUT, 1UL, ""));
|
||||
expected.forward_output_type_infer_configs.emplace_back(std::make_pair(PROPAGATE_FROM_INPUT, 0));
|
||||
expected.forward_output_type_infer_configs.emplace_back(std::make_pair(CONCRETE_TYPE, 1));
|
||||
expected.gradient_input_indices.emplace_back(0UL);
|
||||
|
|
@ -89,12 +110,12 @@ TEST(ATenOpConfigTest, ValidATenOpConfig) {
|
|||
ATenOperatorConfig expected;
|
||||
expected.op_name = "at::c";
|
||||
expected.backward_op_name = "at::c_bw";
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(TENSOR, "input", false));
|
||||
expected.forward_argument_configs.emplace_back(std::make_tuple(INT_ARRAY, "axes", true));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(TENSOR, "grad_output", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(BOOL_ARRAY, "flags", false));
|
||||
expected.backward_argument_configs.emplace_back(std::make_tuple(FLOAT_ARRAY, "es", true));
|
||||
expected.backward_input_source_configs.emplace_back(std::make_pair(GRAD_OUTPUT, 0UL));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "input", false, false));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(INT_ARRAY, "axes", false, true));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "grad_output", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(BOOL_ARRAY, "flags", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(FLOAT_ARRAY, "es", false, true));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(GRAD_OUTPUT, 0UL, ""));
|
||||
expected.forward_output_type_infer_configs.emplace_back(std::make_pair(PROPAGATE_FROM_INPUT, 0));
|
||||
expected.gradient_input_indices.emplace_back(0UL);
|
||||
expected.default_int_array_values["axes"] = {-1, 0, 1};
|
||||
|
|
@ -102,6 +123,30 @@ TEST(ATenOpConfigTest, ValidATenOpConfig) {
|
|||
expected.default_float_array_values["es"] = {};
|
||||
Compare(config, expected);
|
||||
}
|
||||
|
||||
{
|
||||
std::string forward_str = "at::d(Tensor<T> weight, float<CPU> delta) -> Tensor<T> output";
|
||||
std::string backward_str =
|
||||
"at::d_bw(Tensor<T> grad_output, int<CPU>[] weight.sizes(), int<CPU> output.size(0), float<CPU> delta) -> "
|
||||
"Tensor<T> grad_weight";
|
||||
ATenOperatorConfig config = Parse(forward_str, backward_str);
|
||||
ATenOperatorConfig expected;
|
||||
expected.op_name = "at::d";
|
||||
expected.backward_op_name = "at::d_bw";
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "weight", false, false));
|
||||
expected.forward_argument_configs.emplace_back(ArgumentConfig(FLOAT, "delta", true, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(TENSOR, "grad_output", false, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(INT_ARRAY, "weight.sizes()", true, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(INT, "output.size(0)", true, false));
|
||||
expected.backward_argument_configs.emplace_back(ArgumentConfig(FLOAT, "delta", true, false));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(GRAD_OUTPUT, 0UL, ""));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(FORWARD_INPUT, 0UL, "sizes()"));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(FORWARD_OUTPUT, 0UL, "size(0)"));
|
||||
expected.backward_input_source_configs.emplace_back(BackwardInputSourceConfig(FORWARD_INPUT, 1UL, ""));
|
||||
expected.forward_output_type_infer_configs.emplace_back(std::make_pair(PROPAGATE_FROM_INPUT, 0));
|
||||
expected.gradient_input_indices.emplace_back(0UL);
|
||||
Compare(config, expected);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(ATenOpConfigTest, InvalidATenOpConfig) {
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@
|
|||
#include "orttraining/training_ops/cpu/aten_ops/aten_op.h"
|
||||
#include "orttraining/training_ops/cpu/aten_ops/aten_op_attr_parser.h"
|
||||
#include "orttraining/training_ops/cpu/aten_ops/aten_op_executor.h"
|
||||
#include "orttraining/training_ops/cpu/aten_ops/aten_op_config.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
|
|
@ -57,58 +56,112 @@ void ATenOpBase::Init(const OpKernelInfo& info, bool is_backward) {
|
|||
aten_ops::AttributesJsonParser parser(custom_attributes_json);
|
||||
|
||||
for (size_t i = 0; i < argument_configs.size(); i++) {
|
||||
const auto& argument_name = std::get<1>(argument_configs[i]);
|
||||
switch (std::get<0>(argument_configs[i])) {
|
||||
case aten_ops::TENSOR: {
|
||||
tensor_argument_indices_.emplace_back(i);
|
||||
} break;
|
||||
case aten_ops::INT: {
|
||||
// JSON supports INT as 32-bit int, our attribute uses 64-bit int as INT type.
|
||||
int int_value = GetAttributeValue<int>(argument_name, parser, op_config);
|
||||
int_arguments_.emplace_back(std::make_pair(i, static_cast<int64_t>(int_value)));
|
||||
} break;
|
||||
case aten_ops::FLOAT: {
|
||||
float float_value = GetAttributeValue<float>(argument_name, parser, op_config);
|
||||
float_arguments_.emplace_back(std::make_pair(i, float_value));
|
||||
} break;
|
||||
case aten_ops::BOOL: {
|
||||
bool bool_value = GetAttributeValue<bool>(argument_name, parser, op_config);
|
||||
bool_arguments_.emplace_back(std::make_pair(i, bool_value));
|
||||
} break;
|
||||
case aten_ops::INT_ARRAY: {
|
||||
// JSON supports INT as 32-bit int, our attribute uses 64-bit int as INT type.
|
||||
std::vector<int> int_list = GetArrayAttributeValue<int>(argument_name, parser, op_config);
|
||||
std::vector<int64_t> long_list;
|
||||
for (int elem : int_list) {
|
||||
long_list.emplace_back(static_cast<int64_t>(elem));
|
||||
}
|
||||
int_array_arguments_.emplace_back(std::make_pair(i, long_list));
|
||||
} break;
|
||||
case aten_ops::FLOAT_ARRAY: {
|
||||
std::vector<float> float_list = GetArrayAttributeValue<float>(argument_name, parser, op_config);
|
||||
float_array_arguments_.emplace_back(std::make_pair(i, float_list));
|
||||
} break;
|
||||
case aten_ops::BOOL_ARRAY: {
|
||||
std::vector<bool> bool_list = GetArrayAttributeValue<bool>(argument_name, parser, op_config);
|
||||
bool_array_arguments_.emplace_back(std::make_pair(i, bool_list));
|
||||
} break;
|
||||
default:
|
||||
ORT_ENFORCE(false, "Not support for now.");
|
||||
if (argument_configs[i].kind == aten_ops::TENSOR || argument_configs[i].is_cpu_tensor) {
|
||||
tensor_argument_indices_.emplace_back(std::make_pair(i, argument_configs[i].kind));
|
||||
} else {
|
||||
const auto& argument_name = argument_configs[i].name;
|
||||
switch (argument_configs[i].kind) {
|
||||
case aten_ops::INT: {
|
||||
// JSON supports INT as 32-bit int, our attribute uses 64-bit int as INT type.
|
||||
int int_value = GetAttributeValue<int>(argument_name, parser, op_config);
|
||||
int_arguments_.emplace_back(std::make_pair(i, static_cast<int64_t>(int_value)));
|
||||
} break;
|
||||
case aten_ops::FLOAT: {
|
||||
float float_value = GetAttributeValue<float>(argument_name, parser, op_config);
|
||||
float_arguments_.emplace_back(std::make_pair(i, float_value));
|
||||
} break;
|
||||
case aten_ops::BOOL: {
|
||||
bool bool_value = GetAttributeValue<bool>(argument_name, parser, op_config);
|
||||
bool_arguments_.emplace_back(std::make_pair(i, bool_value));
|
||||
} break;
|
||||
case aten_ops::INT_ARRAY: {
|
||||
// JSON supports INT as 32-bit int, our attribute uses 64-bit int as INT type.
|
||||
std::vector<int> int_list = GetArrayAttributeValue<int>(argument_name, parser, op_config);
|
||||
std::vector<int64_t> long_list;
|
||||
for (int elem : int_list) {
|
||||
long_list.emplace_back(static_cast<int64_t>(elem));
|
||||
}
|
||||
int_array_arguments_.emplace_back(std::make_pair(i, long_list));
|
||||
} break;
|
||||
case aten_ops::FLOAT_ARRAY: {
|
||||
std::vector<float> float_list = GetArrayAttributeValue<float>(argument_name, parser, op_config);
|
||||
float_array_arguments_.emplace_back(std::make_pair(i, float_list));
|
||||
} break;
|
||||
case aten_ops::BOOL_ARRAY: {
|
||||
std::vector<bool> bool_list = GetArrayAttributeValue<bool>(argument_name, parser, op_config);
|
||||
bool_array_arguments_.emplace_back(std::make_pair(i, bool_list));
|
||||
} break;
|
||||
default:
|
||||
ORT_ENFORCE(false, "Not support for now.");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T GetCpuArgument(OpKernelContext* p_ctx, size_t index) {
|
||||
const Tensor* tensor = p_ctx->Input<Tensor>(static_cast<int>(index));
|
||||
return *tensor->template Data<T>();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::vector<T> GetCpuArrayArgument(OpKernelContext* p_ctx, size_t index) {
|
||||
const Tensor* tensor = p_ctx->Input<Tensor>(static_cast<int>(index));
|
||||
ORT_ENFORCE(tensor->Shape().NumDimensions() == 1, "Array argument tensor must be a vector tensor.");
|
||||
size_t length = static_cast<size_t>(tensor->Shape().Size());
|
||||
const T* data = tensor->template Data<T>();
|
||||
std::vector<T> result;
|
||||
result.assign(data, data + length);
|
||||
return result;
|
||||
}
|
||||
|
||||
Status ATenOpBase::Compute(OpKernelContext* p_ctx) const {
|
||||
auto* p_ctx_internal = static_cast<OpKernelContextInternal*>(p_ctx);
|
||||
std::vector<std::pair<size_t, DLManagedTensor*>> tensor_arguments;
|
||||
std::vector<std::pair<size_t, int64_t>> int_arguments(int_arguments_);
|
||||
std::vector<std::pair<size_t, float>> float_arguments(float_arguments_);
|
||||
std::vector<std::pair<size_t, bool>> bool_arguments(bool_arguments_);
|
||||
std::vector<std::pair<size_t, std::vector<int64_t>>> int_array_arguments(int_array_arguments_);
|
||||
std::vector<std::pair<size_t, std::vector<float>>> float_array_arguments(float_array_arguments_);
|
||||
std::vector<std::pair<size_t, std::vector<bool>>> bool_array_arguments(bool_array_arguments_);
|
||||
for (size_t i = 0; i < tensor_argument_indices_.size(); i++) {
|
||||
OrtValue ort_value = *p_ctx_internal->GetInputMLValue(static_cast<int>(i));
|
||||
tensor_arguments.emplace_back(std::make_pair(tensor_argument_indices_[i], dlpack::OrtValueToDlpack(ort_value)));
|
||||
switch (tensor_argument_indices_[i].second) {
|
||||
case aten_ops::TENSOR: {
|
||||
OrtValue ort_value = *p_ctx_internal->GetInputMLValue(static_cast<int>(i));
|
||||
tensor_arguments.emplace_back(
|
||||
std::make_pair(tensor_argument_indices_[i].first, dlpack::OrtValueToDlpack(ort_value)));
|
||||
} break;
|
||||
case aten_ops::INT: {
|
||||
int_arguments.emplace_back(
|
||||
std::make_pair(tensor_argument_indices_[i].first, GetCpuArgument<int64_t>(p_ctx, i)));
|
||||
} break;
|
||||
case aten_ops::FLOAT: {
|
||||
float_arguments.emplace_back(
|
||||
std::make_pair(tensor_argument_indices_[i].first, GetCpuArgument<float>(p_ctx, i)));
|
||||
} break;
|
||||
case aten_ops::BOOL: {
|
||||
bool_arguments.emplace_back(std::make_pair(tensor_argument_indices_[i].first, GetCpuArgument<bool>(p_ctx, i)));
|
||||
} break;
|
||||
case aten_ops::INT_ARRAY: {
|
||||
int_array_arguments.emplace_back(
|
||||
std::make_pair(tensor_argument_indices_[i].first, GetCpuArrayArgument<int64_t>(p_ctx, i)));
|
||||
} break;
|
||||
case aten_ops::FLOAT_ARRAY: {
|
||||
float_array_arguments.emplace_back(
|
||||
std::make_pair(tensor_argument_indices_[i].first, GetCpuArrayArgument<float>(p_ctx, i)));
|
||||
} break;
|
||||
case aten_ops::BOOL_ARRAY: {
|
||||
bool_array_arguments.emplace_back(
|
||||
std::make_pair(tensor_argument_indices_[i].first, GetCpuArrayArgument<bool>(p_ctx, i)));
|
||||
} break;
|
||||
default:
|
||||
ORT_ENFORCE(false, "Not support for now.");
|
||||
}
|
||||
}
|
||||
|
||||
auto result = aten_ops::ATenOperatorExecutor::Instance()(op_name_, tensor_arguments, int_arguments_, float_arguments_,
|
||||
bool_arguments_, int_array_arguments_,
|
||||
float_array_arguments_, bool_array_arguments_);
|
||||
auto result = aten_ops::ATenOperatorExecutor::Instance()(op_name_, tensor_arguments, int_arguments, float_arguments,
|
||||
bool_arguments, int_array_arguments, float_array_arguments,
|
||||
bool_array_arguments);
|
||||
|
||||
for (size_t i = 0; i < result.size(); i++) {
|
||||
ORT_RETURN_IF_ERROR(p_ctx_internal->SetOutputMLValue(static_cast<int>(i), dlpack::DlpackToOrtValue(result[i])));
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
#pragma once
|
||||
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "orttraining/training_ops/cpu/aten_ops/aten_op_config.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
|
|
@ -18,7 +19,7 @@ class ATenOpBase : public OpKernel {
|
|||
std::string op_name_;
|
||||
|
||||
// The values in the array are the tensor-type argument indices of Aten Op.
|
||||
std::vector<size_t> tensor_argument_indices_;
|
||||
std::vector<std::pair<size_t, aten_ops::ArgumentKind>> tensor_argument_indices_;
|
||||
|
||||
// The size_t value below are the argument indices of the ATen Op.
|
||||
std::vector<std::pair<size_t, int64_t>> int_arguments_;
|
||||
|
|
|
|||
|
|
@ -14,21 +14,21 @@ namespace aten_ops {
|
|||
// We use regex to parse the strings, to make the parser simple, it requires some special formats
|
||||
// for these function strings, such as spaces in the strings.
|
||||
static const std::vector<std::pair<std::string, std::string>> ATEN_FUNCS = {
|
||||
{"aten::embedding(Tensor<T> weight, Tensor<int64> indices, int padding_idx=-1, bool scale_grad_by_freq=False, bool "
|
||||
"sparse=False) -> Tensor<T> result",
|
||||
"aten::embedding_backward(Tensor<T> grad_result, Tensor<int64> indices, Tensor<T> weight, int padding_idx=-1, "
|
||||
"bool scale_grad_by_freq=False, bool sparse=False) -> Tensor<T> grad_weight"},
|
||||
{"aten::max_pool2d_with_indices(Tensor<T> self, int[] kernel_size, int[] stride, int[] padding=[0,0], int[] "
|
||||
"dilation=[1,1], bool ceil_mode=False) -> (Tensor<T> output, Tensor<int64> indices)",
|
||||
{"aten::embedding(Tensor<T> weight, Tensor<int64> indices, int padding_idx=-1, bool scale_grad_by_freq=False, "
|
||||
"bool sparse=False) -> Tensor<T> result",
|
||||
"aten::embedding_backward(Tensor<T> grad_result, Tensor<int64> indices, int<CPU> weight.size(0), "
|
||||
"int padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor<T> grad_weight"},
|
||||
{"aten::max_pool2d_with_indices(Tensor<T> self, int[] kernel_size, int[] stride, int[] padding=[0,0], "
|
||||
"int[] dilation=[1,1], bool ceil_mode=False) -> (Tensor<T> output, Tensor<int64> indices)",
|
||||
"aten::max_pool2d_with_indices_backward(Tensor<T> grad_output, Tensor<T> self, int[] kernel_size, int[] stride, "
|
||||
"int[] padding=[0,0], int[] dilation=[1,1], bool ceil_mode=False, Tensor<int64> indices) -> Tensor<T> grad_self"},
|
||||
{"aten::unfold(Tensor<T> self, int dimension, int size, int step) -> Tensor<T> output",
|
||||
"aten::unfold_backward(Tensor<T> grad_output, Tensor<T> self, int dimension, int size, int step) -> Tensor<T> "
|
||||
"grad_self"}};
|
||||
"aten::unfold_backward(Tensor<T> grad_output, int<CPU>[] self.sizes(), int dimension, int size, int step) -> "
|
||||
"Tensor<T> grad_self"}};
|
||||
|
||||
const std::regex regex_expr_whole("([a-z0-9:_]+)\\(([A-Za-z0-9_ ,.=+-\\[\\]<>]+)\\) -> \\(?([A-Za-z0-9_ ,<>]+)\\)?");
|
||||
const std::regex regex_expr_whole("([a-z0-9:_]+)\\(([A-Za-z0-9_ ,.=+-\\[\\]\\(\\)<>]+)\\) -> \\(?([A-Za-z0-9_ ,<>]+)\\)?");
|
||||
const std::regex regex_expr_argument(
|
||||
"(Tensor|int|bool|float)(<([A-Za-z0-9_]+)>)?(\\[\\])?(\\?)? ([a-z0-9_]+)(=([TFa-z0-9,.+-\\[\\]]+))?");
|
||||
"(Tensor|int|bool|float)(<([A-Za-z0-9_]+)>)?(\\[\\])?(\\?)? ([a-z0-9_.\\(\\)]+)(=([TFa-z0-9,.+-\\[\\]]+))?");
|
||||
const std::regex regex_expr_comma_space(", ");
|
||||
const std::regex regex_expr_comma(",");
|
||||
// default constructor = end-of-sequence:
|
||||
|
|
@ -37,14 +37,16 @@ const std::regex_token_iterator<std::string::iterator> rend;
|
|||
struct Argument {
|
||||
ArgumentKind type;
|
||||
std::string tensor_elem_type;
|
||||
bool is_cpu_tensor;
|
||||
bool is_optional;
|
||||
std::string name;
|
||||
std::string default_value;
|
||||
|
||||
Argument(ArgumentKind _type, const std::string& _tensor_elem_type, bool _is_optional, const std::string& _name,
|
||||
const std::string& _default_value)
|
||||
Argument(ArgumentKind _type, const std::string& _tensor_elem_type, bool _is_cpu_tensor, bool _is_optional,
|
||||
const std::string& _name, const std::string& _default_value)
|
||||
: type(_type),
|
||||
tensor_elem_type(_tensor_elem_type),
|
||||
is_cpu_tensor(_is_cpu_tensor),
|
||||
is_optional(_is_optional),
|
||||
name(_name),
|
||||
default_value(_default_value) {}
|
||||
|
|
@ -55,10 +57,11 @@ struct Function {
|
|||
std::vector<Argument> arguments;
|
||||
std::vector<Argument> returns;
|
||||
|
||||
std::vector<std::tuple<ArgumentKind, std::string, bool>> ToArgumentConfigs() {
|
||||
std::vector<std::tuple<ArgumentKind, std::string, bool>> argument_configs;
|
||||
std::vector<ArgumentConfig> ToArgumentConfigs() {
|
||||
std::vector<ArgumentConfig> argument_configs;
|
||||
for (const auto& argument : arguments) {
|
||||
argument_configs.emplace_back(std::make_tuple(argument.type, argument.name, argument.is_optional));
|
||||
argument_configs.emplace_back(
|
||||
ArgumentConfig(argument.type, argument.name, argument.is_cpu_tensor, argument.is_optional));
|
||||
}
|
||||
|
||||
return argument_configs;
|
||||
|
|
@ -71,13 +74,18 @@ Argument ParseArgument(const std::string& argument_str) {
|
|||
" is not a vaild argument.");
|
||||
const auto& type_str = sm_argument.str(1);
|
||||
const auto& tensor_elem_type = sm_argument.str(3);
|
||||
bool is_cpu_tensor = tensor_elem_type == "CPU";
|
||||
bool is_array = sm_argument.str(4) == "[]";
|
||||
std::string default_value = sm_argument.str(8);
|
||||
ArgumentKind type;
|
||||
if (type_str == "Tensor") {
|
||||
ORT_ENFORCE(!is_array && tensor_elem_type != "", "Tensor type cannot be an array, and must have element type.");
|
||||
ORT_ENFORCE(!is_array && tensor_elem_type != "" && !is_cpu_tensor,
|
||||
"Tensor type cannot be an array, and must have element type.");
|
||||
type = TENSOR;
|
||||
} else {
|
||||
ORT_ENFORCE(tensor_elem_type == "", "Non-tensor type should not have element type.");
|
||||
ORT_ENFORCE(tensor_elem_type == "" || is_cpu_tensor, "Non-tensor type should not have element type.");
|
||||
// To make it simple, don't support default value for CPU tensor for now.
|
||||
ORT_ENFORCE(!is_cpu_tensor || default_value == "", "CPU tensor doesn't support default value for now.");
|
||||
if (type_str == "int") {
|
||||
type = is_array ? INT_ARRAY : INT;
|
||||
} else if (type_str == "float") {
|
||||
|
|
@ -89,7 +97,7 @@ Argument ParseArgument(const std::string& argument_str) {
|
|||
}
|
||||
}
|
||||
|
||||
return Argument(type, tensor_elem_type, sm_argument.str(5) == "?", sm_argument.str(6), sm_argument.str(8));
|
||||
return Argument(type, tensor_elem_type, is_cpu_tensor, sm_argument.str(5) == "?", sm_argument.str(6), default_value);
|
||||
}
|
||||
|
||||
Function ParseFunction(const std::string& function_str) {
|
||||
|
|
@ -261,22 +269,28 @@ ATenOperatorConfig Parse(const std::string& forward_function_str, const std::str
|
|||
}
|
||||
|
||||
for (const auto& argument : backward_function.arguments) {
|
||||
if (argument.type != TENSOR) {
|
||||
if (argument.type != TENSOR && !argument.is_cpu_tensor) {
|
||||
if (argument.default_value != "" && argument_names.find(argument.name) == argument_names.end()) {
|
||||
AddDefaultValue(argument.name, argument.type, argument.default_value, config);
|
||||
}
|
||||
} else {
|
||||
const auto& name = argument.name;
|
||||
std::string name = argument.name;
|
||||
std::string transform_func = "";
|
||||
size_t pos = name.find('.');
|
||||
if (pos != std::string::npos) {
|
||||
transform_func = name.substr(pos + 1);
|
||||
name = name.substr(0, pos);
|
||||
}
|
||||
if (forward_arguments_name_to_index.find(name) != forward_arguments_name_to_index.end()) {
|
||||
config.backward_input_source_configs.emplace_back(
|
||||
std::make_pair(FORWARD_INPUT, forward_arguments_name_to_index.at(name)));
|
||||
BackwardInputSourceConfig(FORWARD_INPUT, forward_arguments_name_to_index.at(name), transform_func));
|
||||
} else if (forward_returns_name_to_index.find(name) != forward_returns_name_to_index.end()) {
|
||||
config.backward_input_source_configs.emplace_back(
|
||||
std::make_pair(FORWARD_OUTPUT, forward_returns_name_to_index.at(name)));
|
||||
BackwardInputSourceConfig(FORWARD_OUTPUT, forward_returns_name_to_index.at(name), transform_func));
|
||||
} else if (forward_returns_name_to_index.find(name.substr(5UL)) != forward_returns_name_to_index.end()) {
|
||||
// Output gradient has "grad_" prefix.
|
||||
config.backward_input_source_configs.emplace_back(
|
||||
std::make_pair(GRAD_OUTPUT, forward_returns_name_to_index.at(name.substr(5UL))));
|
||||
BackwardInputSourceConfig(GRAD_OUTPUT, forward_returns_name_to_index.at(name.substr(5UL)), transform_func));
|
||||
} else {
|
||||
ORT_ENFORCE(false, "Argument ", name, " is not forward input, output or output gradient.");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@
|
|||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace onnxruntime {
|
||||
|
|
@ -41,16 +40,35 @@ enum ArgumentKind {
|
|||
// TODO: may need more type
|
||||
};
|
||||
|
||||
struct ArgumentConfig {
|
||||
ArgumentKind kind;
|
||||
std::string name;
|
||||
bool is_cpu_tensor;
|
||||
bool is_optional;
|
||||
|
||||
ArgumentConfig(ArgumentKind _kind, const std::string& _name, bool _is_cpu_tensor, bool _is_optional)
|
||||
: kind(_kind), name(_name), is_cpu_tensor(_is_cpu_tensor), is_optional(_is_optional) {}
|
||||
};
|
||||
|
||||
struct BackwardInputSourceConfig {
|
||||
BackwardInputSourceKind kind;
|
||||
size_t index;
|
||||
std::string transform_func;
|
||||
|
||||
BackwardInputSourceConfig(BackwardInputSourceKind _kind, size_t _index, std::string _transform_func)
|
||||
: kind(_kind), index(_index), transform_func(_transform_func) {}
|
||||
};
|
||||
|
||||
// TODO: need to support default attribute value.
|
||||
struct ATenOperatorConfig {
|
||||
std::string op_name;
|
||||
std::string backward_op_name;
|
||||
// Forward ATen Op's argument kind, name and if it's optional.
|
||||
std::vector<std::tuple<ArgumentKind, std::string, bool>> forward_argument_configs;
|
||||
// Backward ATen Op's argument kind, name and if it's optional.
|
||||
std::vector<std::tuple<ArgumentKind, std::string, bool>> backward_argument_configs;
|
||||
// Forward ATen Op's argument configs.
|
||||
std::vector<ArgumentConfig> forward_argument_configs;
|
||||
// Backward ATen Op's argument configs.
|
||||
std::vector<ArgumentConfig> backward_argument_configs;
|
||||
// The source config of inputs of com.microsoft::ATenOpGrad.
|
||||
std::vector<std::pair<BackwardInputSourceKind, size_t>> backward_input_source_configs;
|
||||
std::vector<BackwardInputSourceConfig> backward_input_source_configs;
|
||||
// The output type infer config of outputs of com.microsoft::ATenOp.
|
||||
std::vector<std::pair<OutputTypeInferKind, int>> forward_output_type_infer_configs;
|
||||
// The mapping between com.microsoft::ATenOpGrad's outputs and com.microsoft::ATenOp's inputs,
|
||||
|
|
@ -121,6 +139,21 @@ struct ATenOperatorConfig {
|
|||
|
||||
return has_default_value;
|
||||
}
|
||||
|
||||
bool IsInputOnCpu(size_t input_index, bool is_backward) const {
|
||||
const auto& argument_configs = is_backward ? backward_argument_configs : forward_argument_configs;
|
||||
size_t tensor_input_index = 0;
|
||||
for (const auto& argument_config : argument_configs) {
|
||||
if (argument_config.kind == TENSOR || argument_config.is_cpu_tensor) {
|
||||
if (tensor_input_index == input_index) {
|
||||
return argument_config.is_cpu_tensor;
|
||||
}
|
||||
tensor_input_index++;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
};
|
||||
|
||||
ATenOperatorConfig Parse(const std::string& forward_function_str, const std::string& backward_function_str);
|
||||
|
|
|
|||
Loading…
Reference in a new issue