From 88ec95ea967c7f148a526fd33ee1bf0c365763d7 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Fri, 2 Jul 2021 23:04:43 +0800 Subject: [PATCH] Support OrtMemTypeCPUInput for ATenOp/ATenOpGrad (#8116) --- .../core/framework/allocation_planner.cc | 2 +- onnxruntime/core/framework/utils.cc | 22 +++ onnxruntime/core/framework/utils.h | 3 + .../core/optimizer/transformer_memcpy.cc | 10 +- .../core/graph/gradient_builder.cc | 20 +-- .../core/graph/gradient_builder_base.cc | 24 +++ .../core/graph/gradient_builder_base.h | 3 + .../aten_op_executor/aten_op_executor.cc | 56 +------ .../cpu/aten_ops/aten_op_config_test.cc | 101 +++++++++---- .../training_ops/cpu/aten_ops/aten_op.cc | 139 ++++++++++++------ .../training_ops/cpu/aten_ops/aten_op.h | 3 +- .../cpu/aten_ops/aten_op_config.cc | 60 +++++--- .../cpu/aten_ops/aten_op_config.h | 45 +++++- 13 files changed, 320 insertions(+), 168 deletions(-) diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 4f6f21ba45..cd33a96fae 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -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(); diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index fa6c4e58c2..5398ac633b 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -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 diff --git a/onnxruntime/core/framework/utils.h b/onnxruntime/core/framework/utils.h index f91c771dac..150d398d75 100644 --- a/onnxruntime/core/framework/utils.h +++ b/onnxruntime/core/framework/utils.h @@ -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& 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 constexpr ONNXTensorElementDataType GetONNXTensorElementDataType() { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; diff --git a/onnxruntime/core/optimizer/transformer_memcpy.cc b/onnxruntime/core/optimizer/transformer_memcpy.cc index 2a7d7e8230..96edd01cfa 100644 --- a/onnxruntime/core/optimizer/transformer_memcpy.cc +++ b/onnxruntime/core/optimizer/transformer_memcpy.cc @@ -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(); }); diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 9ad11f56a2..1624e66e3c 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -1681,6 +1681,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetMinMaxGradient) { } IMPLEMENT_GRADIENT_BUILDER(GetATenOpGradient) { + std::vector result; const auto& src_attrs = SrcNodeAttributes(); std::vector attrs; ORT_ENFORCE(utils::HasString(src_attrs.at("name"))); @@ -1699,18 +1700,10 @@ IMPLEMENT_GRADIENT_BUILDER(GetATenOpGradient) { std::vector 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(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) { diff --git a/orttraining/orttraining/core/graph/gradient_builder_base.cc b/orttraining/orttraining/core/graph/gradient_builder_base.cc index 0a3b13121f..987fdd4e56 100644 --- a/orttraining/orttraining/core/graph/gradient_builder_base.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_base.cc @@ -293,5 +293,29 @@ std::vector GradientBuilderBase::GetBiasGeluGradNodes( return result; } +ArgDef GradientBuilderBase::HandleATenOpGradInput(const ArgDef& source_arg_def, const std::string& transform_func, + std::vector& 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(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 diff --git a/orttraining/orttraining/core/graph/gradient_builder_base.h b/orttraining/orttraining/core/graph/gradient_builder_base.h index b7fe60ed0b..f765234636 100644 --- a/orttraining/orttraining/core/graph/gradient_builder_base.h +++ b/orttraining/orttraining/core/graph/gradient_builder_base.h @@ -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& output) const; + private: friend class GradientGraphBuilder; diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc index e562ad0c90..91ede87bc5 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc @@ -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 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; - -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> TENSOR_TRANSFORM_FUNCS = { - {"aten::embedding_backward", {{2, embedding_num_weights}}}, - {"aten::unfold_backward", {{1, unfold_input_sizes}}}, -}; - template void SetIValueArguments(const std::vector>& raw_arguments, const std::vector& is_optional_arguments, std::vector& ivalue_arguments) { @@ -120,21 +93,8 @@ std::vector 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(tensor)) : c10::IValue(tensor); - } + arguments[index] = + aten_op.is_optional_arguments[index] ? c10::IValue(c10::optional(tensor)) : c10::IValue(tensor); } SetIValueArguments(int_arguments, aten_op.is_optional_arguments, arguments); diff --git a/orttraining/orttraining/test/training_ops/cpu/aten_ops/aten_op_config_test.cc b/orttraining/orttraining/test/training_ops/cpu/aten_ops/aten_op_config_test.cc index 4076c888b6..e58186b349 100644 --- a/orttraining/orttraining/test/training_ops/cpu/aten_ops/aten_op_config_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/aten_ops/aten_op_config_test.cc @@ -10,11 +10,32 @@ namespace test { using namespace contrib::aten_ops; namespace { +void CompareArgumentConfigs(const std::vector& configs, const std::vector& 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& configs, + const std::vector& 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 weight, float delta) -> Tensor output"; + std::string backward_str = + "at::d_bw(Tensor grad_output, int[] weight.sizes(), int output.size(0), float delta) -> " + "Tensor 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) { diff --git a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.cc b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.cc index 0fb1b5f5f1..37138f2c8b 100644 --- a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.cc +++ b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.cc @@ -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(argument_name, parser, op_config); - int_arguments_.emplace_back(std::make_pair(i, static_cast(int_value))); - } break; - case aten_ops::FLOAT: { - float float_value = GetAttributeValue(argument_name, parser, op_config); - float_arguments_.emplace_back(std::make_pair(i, float_value)); - } break; - case aten_ops::BOOL: { - bool bool_value = GetAttributeValue(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_list = GetArrayAttributeValue(argument_name, parser, op_config); - std::vector long_list; - for (int elem : int_list) { - long_list.emplace_back(static_cast(elem)); - } - int_array_arguments_.emplace_back(std::make_pair(i, long_list)); - } break; - case aten_ops::FLOAT_ARRAY: { - std::vector float_list = GetArrayAttributeValue(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_list = GetArrayAttributeValue(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(argument_name, parser, op_config); + int_arguments_.emplace_back(std::make_pair(i, static_cast(int_value))); + } break; + case aten_ops::FLOAT: { + float float_value = GetAttributeValue(argument_name, parser, op_config); + float_arguments_.emplace_back(std::make_pair(i, float_value)); + } break; + case aten_ops::BOOL: { + bool bool_value = GetAttributeValue(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_list = GetArrayAttributeValue(argument_name, parser, op_config); + std::vector long_list; + for (int elem : int_list) { + long_list.emplace_back(static_cast(elem)); + } + int_array_arguments_.emplace_back(std::make_pair(i, long_list)); + } break; + case aten_ops::FLOAT_ARRAY: { + std::vector float_list = GetArrayAttributeValue(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_list = GetArrayAttributeValue(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 +T GetCpuArgument(OpKernelContext* p_ctx, size_t index) { + const Tensor* tensor = p_ctx->Input(static_cast(index)); + return *tensor->template Data(); +} + +template +std::vector GetCpuArrayArgument(OpKernelContext* p_ctx, size_t index) { + const Tensor* tensor = p_ctx->Input(static_cast(index)); + ORT_ENFORCE(tensor->Shape().NumDimensions() == 1, "Array argument tensor must be a vector tensor."); + size_t length = static_cast(tensor->Shape().Size()); + const T* data = tensor->template Data(); + std::vector result; + result.assign(data, data + length); + return result; +} + Status ATenOpBase::Compute(OpKernelContext* p_ctx) const { auto* p_ctx_internal = static_cast(p_ctx); std::vector> tensor_arguments; + std::vector> int_arguments(int_arguments_); + std::vector> float_arguments(float_arguments_); + std::vector> bool_arguments(bool_arguments_); + std::vector>> int_array_arguments(int_array_arguments_); + std::vector>> float_array_arguments(float_array_arguments_); + std::vector>> 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(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(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(p_ctx, i))); + } break; + case aten_ops::FLOAT: { + float_arguments.emplace_back( + std::make_pair(tensor_argument_indices_[i].first, GetCpuArgument(p_ctx, i))); + } break; + case aten_ops::BOOL: { + bool_arguments.emplace_back(std::make_pair(tensor_argument_indices_[i].first, GetCpuArgument(p_ctx, i))); + } break; + case aten_ops::INT_ARRAY: { + int_array_arguments.emplace_back( + std::make_pair(tensor_argument_indices_[i].first, GetCpuArrayArgument(p_ctx, i))); + } break; + case aten_ops::FLOAT_ARRAY: { + float_array_arguments.emplace_back( + std::make_pair(tensor_argument_indices_[i].first, GetCpuArrayArgument(p_ctx, i))); + } break; + case aten_ops::BOOL_ARRAY: { + bool_array_arguments.emplace_back( + std::make_pair(tensor_argument_indices_[i].first, GetCpuArrayArgument(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(i), dlpack::DlpackToOrtValue(result[i]))); diff --git a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.h b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.h index dedfcab7c3..f4f923f8ae 100644 --- a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.h +++ b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.h @@ -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 tensor_argument_indices_; + std::vector> tensor_argument_indices_; // The size_t value below are the argument indices of the ATen Op. std::vector> int_arguments_; diff --git a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.cc b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.cc index 4f07bc926a..6663f877bc 100644 --- a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.cc +++ b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.cc @@ -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> ATEN_FUNCS = { - {"aten::embedding(Tensor weight, Tensor indices, int padding_idx=-1, bool scale_grad_by_freq=False, bool " - "sparse=False) -> Tensor result", - "aten::embedding_backward(Tensor grad_result, Tensor indices, Tensor weight, int padding_idx=-1, " - "bool scale_grad_by_freq=False, bool sparse=False) -> Tensor grad_weight"}, - {"aten::max_pool2d_with_indices(Tensor self, int[] kernel_size, int[] stride, int[] padding=[0,0], int[] " - "dilation=[1,1], bool ceil_mode=False) -> (Tensor output, Tensor indices)", + {"aten::embedding(Tensor weight, Tensor indices, int padding_idx=-1, bool scale_grad_by_freq=False, " + "bool sparse=False) -> Tensor result", + "aten::embedding_backward(Tensor grad_result, Tensor indices, int weight.size(0), " + "int padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor grad_weight"}, + {"aten::max_pool2d_with_indices(Tensor self, int[] kernel_size, int[] stride, int[] padding=[0,0], " + "int[] dilation=[1,1], bool ceil_mode=False) -> (Tensor output, Tensor indices)", "aten::max_pool2d_with_indices_backward(Tensor grad_output, Tensor self, int[] kernel_size, int[] stride, " "int[] padding=[0,0], int[] dilation=[1,1], bool ceil_mode=False, Tensor indices) -> Tensor grad_self"}, {"aten::unfold(Tensor self, int dimension, int size, int step) -> Tensor output", - "aten::unfold_backward(Tensor grad_output, Tensor self, int dimension, int size, int step) -> Tensor " - "grad_self"}}; + "aten::unfold_backward(Tensor grad_output, int[] self.sizes(), int dimension, int size, int step) -> " + "Tensor 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 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 arguments; std::vector returns; - std::vector> ToArgumentConfigs() { - std::vector> argument_configs; + std::vector ToArgumentConfigs() { + std::vector 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."); } diff --git a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.h b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.h index b1109be7b8..eb94033799 100644 --- a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.h +++ b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op_config.h @@ -5,7 +5,6 @@ #include #include -#include #include 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> forward_argument_configs; - // Backward ATen Op's argument kind, name and if it's optional. - std::vector> backward_argument_configs; + // Forward ATen Op's argument configs. + std::vector forward_argument_configs; + // Backward ATen Op's argument configs. + std::vector backward_argument_configs; // The source config of inputs of com.microsoft::ATenOpGrad. - std::vector> backward_input_source_configs; + std::vector backward_input_source_configs; // The output type infer config of outputs of com.microsoft::ATenOp. std::vector> 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);