diff --git a/include/onnxruntime/core/graph/constants.h b/include/onnxruntime/core/graph/constants.h index 0990e09bbf..70c2a1ab54 100644 --- a/include/onnxruntime/core/graph/constants.h +++ b/include/onnxruntime/core/graph/constants.h @@ -13,6 +13,7 @@ constexpr const char* kOnnxDomain = ""; constexpr const char* kOnnxDomainAlias = "ai.onnx"; constexpr const char* kMLDomain = "ai.onnx.ml"; constexpr const char* kMSDomain = "com.microsoft"; +constexpr const char* kPytorchAtenDomain = "org.pytorch.aten"; constexpr const char* kMSExperimentalDomain = "com.microsoft.experimental"; constexpr const char* kMSNchwcDomain = "com.microsoft.nchwc"; constexpr const char* kMSInternalNHWCDomain = "com.ms.internal.nhwc"; diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 2e6e15451b..7c7b78f25c 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -786,10 +786,10 @@ bool IsInputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index) } #ifdef ENABLE_TRAINING - if (node.GetExecutionProviderType() == kCudaExecutionProvider && node.OpType() == "ATenOp" && node.Domain() == kMSDomain) { + if (node.GetExecutionProviderType() == kCudaExecutionProvider && node.OpType() == "ATen" && node.Domain() == kPytorchAtenDomain) { const auto& attrs = node.GetAttributes(); - ORT_ENFORCE(utils::HasString(attrs.at("name"))); - std::string op_name = attrs.at("name").s(); + ORT_ENFORCE(utils::HasString(attrs.at("operator"))); + std::string op_name = attrs.at("operator").s(); std::string overload_name = ""; if (attrs.find("overload_name") != attrs.end() && utils::HasString(attrs.at("overload_name"))) { overload_name = attrs.at("overload_name").s(); diff --git a/onnxruntime/core/providers/cpu/cpu_provider_shared.cc b/onnxruntime/core/providers/cpu/cpu_provider_shared.cc index 20ad4aad06..afce3a29f8 100644 --- a/onnxruntime/core/providers/cpu/cpu_provider_shared.cc +++ b/onnxruntime/core/providers/cpu/cpu_provider_shared.cc @@ -77,7 +77,7 @@ struct ProviderHostCPUImpl : ProviderHostCPU { std::vector& split_sizes) override { return p->SplitBase::PrepareForCompute(input_shape, num_outputs, axis, before_dims, after_dims_including_split_axis, after_dims_excluding_split, split_sizes); } // From cpu/tensor/concatbase.h (direct) - Status ConcatBase__PrepareForCompute(const ConcatBase* p, OpKernelContext* ctx, const ConcatBase_InlinedTensorsVector& input_tensors, Prepare& prepare) override { + Status ConcatBase__PrepareForCompute(const ConcatBase* p, OpKernelContext* ctx, const ConcatBase_InlinedTensorsVector& input_tensors, Prepare& prepare) override { return p->ConcatBase::PrepareForCompute(ctx, reinterpret_cast(input_tensors), prepare); } // GatherElements (direct) @@ -176,7 +176,7 @@ struct ProviderHostCPUImpl : ProviderHostCPU { #endif #ifdef ENABLE_TRAINING - Status ATenOp__Compute(const contrib::ATenOp* p, OpKernelContext* p_ctx) override { return p->ATenOp::Compute(p_ctx); } + Status ATen__Compute(const contrib::ATen* p, OpKernelContext* p_ctx) override { return p->ATen::Compute(p_ctx); } void contrib__record_event_in_tensor(const Tensor& event_id_tensor) override { return contrib::record_event_in_tensor(event_id_tensor); } void contrib__wait_event_in_tensor(const Tensor& event_id_tensor) override { return contrib::wait_event_in_tensor(event_id_tensor); } Status contrib__Group__Compute(const contrib::Group* p, OpKernelContext* context) override { return p->Group::Compute(context); } @@ -189,7 +189,7 @@ struct ProviderHostCPUImpl : ProviderHostCPU { // From aten_op.h (direct) bool contrib__IsATenOperatorExecutorInitialized() override { return contrib::IsATenOperatorExecutorInitialized(); } - Status contrib__ExecuteReduceSumATenOp(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) override { return contrib::ExecuteReduceSumATenOp(p_ctx, axes, keepdims); } + Status contrib__ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) override { return contrib::ExecuteReduceSumATen(p_ctx, axes, keepdims); } #endif #endif }; diff --git a/onnxruntime/core/providers/cpu/cpu_provider_shared.h b/onnxruntime/core/providers/cpu/cpu_provider_shared.h index 5c6b2cd2e5..f3a7822f47 100644 --- a/onnxruntime/core/providers/cpu/cpu_provider_shared.h +++ b/onnxruntime/core/providers/cpu/cpu_provider_shared.h @@ -137,7 +137,7 @@ struct ProviderHostCPU { #endif #ifdef ENABLE_TRAINING - virtual Status ATenOp__Compute(const contrib::ATenOp* p, OpKernelContext* p_ctx) = 0; + virtual Status ATen__Compute(const contrib::ATen* p, OpKernelContext* p_ctx) = 0; virtual void contrib__record_event_in_tensor(const Tensor& event_id_tensor) = 0; virtual void contrib__wait_event_in_tensor(const Tensor& event_id_tensor) = 0; virtual Status contrib__Group__Compute(const contrib::Group* p, OpKernelContext* context) = 0; @@ -150,7 +150,7 @@ struct ProviderHostCPU { // From aten_op.h virtual bool contrib__IsATenOperatorExecutorInitialized() = 0; - virtual Status contrib__ExecuteReduceSumATenOp(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) = 0; + virtual Status contrib__ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) = 0; #endif #endif }; @@ -220,7 +220,7 @@ inline Status PrepareForTrainingCompute(const TensorShape& input_shape, int num_ // From aten_op.h inline bool IsATenOperatorExecutorInitialized() { return g_host_cpu.contrib__IsATenOperatorExecutorInitialized(); } -inline Status ExecuteReduceSumATenOp(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) { return g_host_cpu.contrib__ExecuteReduceSumATenOp(p_ctx, axes, keepdims); } +inline Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) { return g_host_cpu.contrib__ExecuteReduceSumATen(p_ctx, axes, keepdims); } } // namespace contrib #endif // ENABLE_TRAINING #endif // USE_CUDA || USE_ROCM diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 7ffd4bb185..0806a0fd82 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -739,7 +739,7 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe } #ifdef ENABLE_TRAINING - // Use ATenOp for ReduceSum if possible. + // Use ATen for ReduceSum if possible. const TensorShape& input_shape = X->Shape(); if (contrib::IsATenOperatorExecutorInitialized() && cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && !calculate_log_ && !calculate_sqt_ && !log_sum_exp_ && input_shape.Size() > 0) { @@ -747,7 +747,7 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe axes.resize(input_shape.NumDimensions()); std::iota(axes.begin(), axes.end(), 0); } - ORT_RETURN_IF_ERROR(contrib::ExecuteReduceSumATenOp(ctx, axes, keepdims_)); + ORT_RETURN_IF_ERROR(contrib::ExecuteReduceSumATen(ctx, axes, keepdims_)); return Status::OK(); } #endif diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index b36c0c15c8..09b8fab0ea 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/onnxruntime/core/providers/shared_library/provider_api.h @@ -167,7 +167,7 @@ template struct EinsumTypedComputeProcessor; namespace contrib { -class ATenOp; +class ATen; class Group; class PassThrough; class YieldOp; @@ -228,6 +228,7 @@ struct DeleteOnUnloadPtr { constexpr const char* kOnnxDomain = ""; constexpr const char* kMSDomain = "com.microsoft"; +constexpr const char* kPytorchAtenDomain = "org.pytorch.aten"; constexpr const char* kNGraphDomain = "com.intel.ai"; constexpr const char* kCudaExecutionProvider = "CUDAExecutionProvider"; constexpr const char* kDnnlExecutionProvider = "DnnlExecutionProvider"; diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index 392afe355b..03075a728e 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -63,7 +63,7 @@ void operator delete(void* p, size_t /*size*/) noexcept { return Provider_GetHos namespace onnxruntime { #if defined(_MSC_VER) && !defined(__clang__) #pragma warning(push) -// "Global initializer calls a non-constexpr function." +// "Global initializer calls a non-constexpr function." #pragma warning(disable : 26426) #endif ProviderHost* g_host = Provider_GetHost(); @@ -577,7 +577,7 @@ Status Scan<9>::SetupSubgraphExecutionInfo(const SessionState& session_state, co #ifdef ENABLE_TRAINING namespace contrib { -Status ATenOp::Compute(OpKernelContext* p_ctx) const { return g_host_cpu.ATenOp__Compute(this, p_ctx); } +Status ATen::Compute(OpKernelContext* p_ctx) const { return g_host_cpu.ATen__Compute(this, p_ctx); } Status Group::Compute(OpKernelContext* context) const { return g_host_cpu.contrib__Group__Compute(this, context); } Status PassThrough::Compute(OpKernelContext* context) const { return g_host_cpu.contrib__PassThrough__Compute(this, context); } Status YieldOp::Compute(OpKernelContext* context) const { return g_host_cpu.contrib__YieldOp__Compute(this, context); } diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index 27019a7f1d..9c99fb8e0b 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -223,6 +223,7 @@ Status Environment::Initialize(std::unique_ptr logging_ domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSInternalNHWCDomain, 1, 1); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kPytorchAtenDomain, 1, 1); #ifdef USE_DML domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDmlDomain, 1, 1); #endif @@ -272,7 +273,7 @@ Status Environment::Initialize(std::unique_ptr logging_ all_types.insert(all_types.end(), all_tensor_types.begin(), all_tensor_types.end()); all_types.insert(all_types.end(), all_sequence_types.begin(), all_sequence_types.end()); all_types.emplace_back("seq(tensor(bfloat16))"); - all_types.erase(std::remove_if(all_types.begin(), all_types.end(), + all_types.erase(std::remove_if(all_types.begin(), all_types.end(), [](const std::string& s) { return s.find("string") != std::string::npos; }), all_types.end()); return all_types; }(); diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index a728f10ac1..4eff2f22de 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -1898,9 +1898,10 @@ class SymbolicShapeInference: vi = self.known_vi_[node.output[0]] if len(vi.type.tensor_type.shape.dim) == 0: vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - elif node.op_type == 'ATenOp' and node.domain == 'com.microsoft': + elif node.op_type == 'ATen' and node.domain == 'org.pytorch.aten': for attr in node.attribute: - if attr.name == 'name': + # TODO: Is overload_name needed? + if attr.name == 'operator': aten_op_name = attr.s.decode('utf-8') if isinstance(attr.s, bytes) else attr.s if aten_op_name in self.aten_op_dispatcher_: known_aten_op = True diff --git a/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json index 856996d074..95611f5371 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json +++ b/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json @@ -267,4 +267,4 @@ "ZeroGradient com.microsoft CPUExecutionProvider", 3284255990062374928 ] -] \ No newline at end of file +] diff --git a/orttraining/orttraining/core/framework/gradient_graph_builder.cc b/orttraining/orttraining/core/framework/gradient_graph_builder.cc index a28a0aa393..923bf74d5e 100644 --- a/orttraining/orttraining/core/framework/gradient_graph_builder.cc +++ b/orttraining/orttraining/core/framework/gradient_graph_builder.cc @@ -139,7 +139,7 @@ NodeSet GradientGraphBuilder::BFSWithStopGradient(const std::unordered_setOutputEdgesBegin(); edge_it != n->OutputEdgesEnd(); ++edge_it) { const Node& node = edge_it->GetNode(); - const std::unordered_set* edges = GetStopGradientEdges(node); + const std::unordered_set* edges = GetStopGradientEdges(node); if (edges != nullptr && edges->count(edge_it->GetDstArgIndex())) { continue; } @@ -163,7 +163,7 @@ NodeSet GradientGraphBuilder::ReverseBFSWithStopGradient(const NodeSet& nodes) c queue.pop_front(); for (auto edge_it = n->InputEdgesBegin(); edge_it != n->InputEdgesEnd(); ++edge_it) { - const std::unordered_set* edges = GetStopGradientEdges(*n); + const std::unordered_set* edges = GetStopGradientEdges(*n); if (edges != nullptr && edges->count(edge_it->GetDstArgIndex())) { LOGS(logger_, INFO) << "Skip building gradient for input_" << edge_it->GetDstArgIndex() << " of node: " << n->Name(); @@ -203,7 +203,7 @@ Status GradientGraphBuilder::CheckNodeArgsReachable() const { const std::unordered_set* GradientGraphBuilder::GetStopGradientEdges(const Node& node) const { std::string op_type = node.OpType(); - if (op_type == "ATenOp") { + if (op_type == "ATen") { std::string key = GetGradientDefinitionKeyByNode(node); return GradientDefinitionRegistry::Instance().GetStopGradientEdgesForNode(key); } else { diff --git a/orttraining/orttraining/core/graph/gradient_builder_base.cc b/orttraining/orttraining/core/graph/gradient_builder_base.cc index ce70fc38f7..2c6e3e58b4 100644 --- a/orttraining/orttraining/core/graph/gradient_builder_base.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_base.cc @@ -64,7 +64,7 @@ void ComputeBroadcastBackwardAxes( B_dim = B_dims[j].dim_param(); if (A_dim != B_dim) { LOGS_DEFAULT(INFO) << "Gradient building for node " << node_name << ": symbolic dimension expects to match. " << - "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << + "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; } } else if (A_dims[i].has_dim_param() && B_dims[j].has_dim_value()) { @@ -73,8 +73,8 @@ void ComputeBroadcastBackwardAxes( if (B_dim != 1) { LOGS_DEFAULT(INFO) << "Gradient building for node " << node_name << ": symbolic broadcasting expects the B_dimension to be 1. " << - "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << - " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; + "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << + " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; } else { if (B_axes) { B_axes->push_back(gsl::narrow_cast(k)); @@ -86,7 +86,7 @@ void ComputeBroadcastBackwardAxes( if (A_dim != 1) { LOGS_DEFAULT(INFO) << "Gradient building for node " << node_name << ": symbolic broadcasting expects the A_dimension to be 1. " << - "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << + "A_dims:" << ToString(A_dims) << ", B_dims:" << ToString(B_dims) << " This is a relaxing case, and the kernel might run into problem later if A_dims and B_dims turns out not broadcastable."; } else { if (A_axes) { @@ -130,10 +130,10 @@ Status GetShape(const ArgDef& arg_def, std::vector& shape) { std::string GetGradientDefinitionKeyByNode(const Node& node) { std::string op_type = node.OpType(); std::string key = node.Domain() + "::" + op_type; - if (op_type == "ATenOp") { + if (op_type == "ATen") { const auto& attrs = node.GetAttributes(); - ORT_ENFORCE(utils::HasString(attrs.at("name"))); - key = key + "::" + attrs.at("name").s(); + ORT_ENFORCE(utils::HasString(attrs.at("operator"))); + key = key + "::" + attrs.at("operator").s(); std::string overload_name = ""; if (attrs.find("overload_name") != attrs.end() && utils::HasString(attrs.at("overload_name"))) { overload_name = attrs.at("overload_name").s(); diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 013a35d0af..763a3f50f1 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -2100,16 +2100,16 @@ Example 4: .Const("C_MinusHalf", -0.5f, elem_type) .Const("C_alpha", kAlpha, elem_type) .Add(R"( - ErfArg = Mul (X, C_SqrtHalf) - ErfTerm = Erf (ErfArg) - PartialSum = Add (ErfTerm, C_One) - HalfPartialSum = Mul (C_Half, PartialSum) - AlphaX = Mul (X, C_alpha) - MinusHalfX = Mul (C_MinusHalf, X) - ExpArg = Mul (MinusHalfX, X) - ExpTerm = Exp (ExpArg) - Term3 = Mul (AlphaX, ExpTerm) - FullSum = Add (HalfPartialSum, Term3) + ErfArg = Mul (X, C_SqrtHalf) + ErfTerm = Erf (ErfArg) + PartialSum = Add (ErfTerm, C_One) + HalfPartialSum = Mul (C_Half, PartialSum) + AlphaX = Mul (X, C_alpha) + MinusHalfX = Mul (C_MinusHalf, X) + ExpArg = Mul (MinusHalfX, X) + ExpTerm = Exp (ExpArg) + Term3 = Mul (AlphaX, ExpTerm) + FullSum = Add (HalfPartialSum, Term3) dX = Mul (dY, FullSum) )"); @@ -2994,18 +2994,18 @@ Return true if all elements are true and false otherwise. }); #ifdef ENABLE_TRAINING - ONNX_CONTRIB_OPERATOR_SCHEMA(ATenOp) - .SetDomain(kMSDomain) + ONNX_CONTRIB_OPERATOR_SCHEMA(ATen) + .SetDomain(kPytorchAtenDomain) .SinceVersion(1) .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) - .SetDoc("ATenOp") - .Input(0, "inputs", "ATenOp inputs.", "T", OpSchema::Variadic, + .SetDoc("ATen") + .Input(0, "inputs", "ATen Op inputs.", "T", OpSchema::Variadic, /*is_homogeneous*/ false, /*min_arity*/ 1) - .Output(0, "outputs", "ATenOp outputs.", "T", OpSchema::Variadic, + .Output(0, "outputs", "ATen Op outputs.", "T", OpSchema::Variadic, /*is_homogeneous*/ false, /*min_arity*/ 1) - .Attr("name", "Name of ATen operator.", AttributeProto::STRING) + .Attr("operator", "Name of ATen operator.", AttributeProto::STRING) .Attr("overload_name", "Overload name of ATen operator.", AttributeProto::STRING, false) .TypeConstraint("T", OpSchema::all_tensor_types(), "Allow inputs and outputs to be any kind of tensor."); #endif diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py index aca0c632bf..1e9492ef42 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py @@ -3,9 +3,9 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -# Support registration of ATenOp's symbolic (PyTorch exporter overriding) and gradient definition. +# Support registration of ATen op's symbolic (PyTorch exporter overriding) and gradient definition. -# Each gradient definition of ATenOp is a list of node definitions: [node_def1, node_def2, ...]. +# Each gradient definition of ATen op is a list of node definitions: [node_def1, node_def2, ...]. # Each node definition is a tuple: (op_type, inputs, outputs, attributes), while attributes is optional. # 'op_type' is a string or a tuple of two strings with op_name and domain respectively. # If it's string type, then the domain is the default ONNX domain: ''. @@ -82,70 +82,70 @@ def register_gradient(domain, name, *attributes): return gradient_wrapper -# For ATenOp, we need to provide op_name and overload name. -@register_gradient('com.microsoft', 'ATenOp', 'aten::embedding', '') +# For ATen op, we need to provide op_name and overload name. +@register_gradient('org.pytorch.aten', 'ATen', 'aten::embedding', '') def embedding_gradient(): return [ ('Constant', [], ['Const_0'], {'value': {'value': 0, 'dtype': 'int', 'is_tensor': True}}), ('Shape', ['I(0)'], ['Shape_X']), ('Gather', ['Shape_X', 'Const_0'], ['Gather_X_0'], {'axis': {'value': 0, 'dtype': 'int'}}), - (('ATenOp', 'com.microsoft'), ['GO(0)', 'I(1)', 'Gather_X_0', 'I(2)', 'I(3)', 'I(4)'], [ - 'GI(0)'], {'name': {'value': 'aten::embedding_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'I(1)', 'Gather_X_0', 'I(2)', 'I(3)', 'I(4)'], [ + 'GI(0)'], {'operator': {'value': 'aten::embedding_backward', 'dtype': 'string'}}), ] -@register_gradient('com.microsoft', 'ATenOp', 'aten::diagonal', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::diagonal', '') def diagonal_gradient(): return [ ('Shape', ['I(0)'], ['Shape_X']), - (('ATenOp', 'com.microsoft'), ['GO(0)', 'Shape_X', 'I(1)', 'I(2)', 'I(3)'], [ - 'GI(0)'], {'name': {'value': 'aten::diagonal_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'Shape_X', 'I(1)', 'I(2)', 'I(3)'], [ + 'GI(0)'], {'operator': {'value': 'aten::diagonal_backward', 'dtype': 'string'}}), ] -@register_gradient('com.microsoft', 'ATenOp', 'aten::max_pool2d_with_indices', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::max_pool2d_with_indices', '') def max_pool2d_gradient(): return [ - (('ATenOp', 'com.microsoft'), ['GO(0)', 'I(0)', 'I(1)', 'I(2)', 'I(3)', 'I(4)', 'I(5)', 'O(1)'], [ - 'GI(0)'], {'name': {'value': 'aten::max_pool2d_with_indices_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'I(0)', 'I(1)', 'I(2)', 'I(3)', 'I(4)', 'I(5)', 'O(1)'], [ + 'GI(0)'], {'operator': {'value': 'aten::max_pool2d_with_indices_backward', 'dtype': 'string'}}), ] -@register_gradient('com.microsoft', 'ATenOp', 'aten::unfold', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::unfold', '') def unfold_gradient(): return [ ('Shape', ['I(0)'], ['Shape_X']), - (('ATenOp', 'com.microsoft'), ['GO(0)', 'Shape_X', 'I(1)', 'I(2)', 'I(3)'], [ - 'GI(0)'], {'name': {'value': 'aten::unfold_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'Shape_X', 'I(1)', 'I(2)', 'I(3)'], [ + 'GI(0)'], {'operator': {'value': 'aten::unfold_backward', 'dtype': 'string'}}), ] -@register_gradient('com.microsoft', 'ATenOp', 'aten::avg_pool2d', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::avg_pool2d', '') def avg_pool2d_gradient(): return [ - (('ATenOp', 'com.microsoft'), ['GO(0)', 'I(0)', 'I(1)', 'I(2)', 'I(3)', 'I(4)', 'I(5)', 'I(6)'], [ - 'GI(0)'], {'name': {'value': 'aten::avg_pool2d_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'I(0)', 'I(1)', 'I(2)', 'I(3)', 'I(4)', 'I(5)', 'I(6)'], [ + 'GI(0)'], {'operator': {'value': 'aten::avg_pool2d_backward', 'dtype': 'string'}}), ] -@register_gradient('com.microsoft', 'ATenOp', 'aten::_adaptive_avg_pool2d', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::_adaptive_avg_pool2d', '') def adaptive_avg_pool2d_gradient(): return [ - (('ATenOp', 'com.microsoft'), ['GO(0)', 'I(0)'], [ - 'GI(0)'], {'name': {'value': 'aten::_adaptive_avg_pool2d_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'I(0)'], [ + 'GI(0)'], {'operator': {'value': 'aten::_adaptive_avg_pool2d_backward', 'dtype': 'string'}}), ] -CustomGradientRegistry.register_custom_stop_gradient_edges([0], 'com.microsoft', 'ATenOp', 'aten::argmax', '') -CustomGradientRegistry.register_custom_stop_gradient_edges([0], 'com.microsoft', 'ATenOp', 'aten::multinomial', '') +CustomGradientRegistry.register_custom_stop_gradient_edges([0], 'org.pytorch.aten', 'ATen', 'aten::argmax', '') +CustomGradientRegistry.register_custom_stop_gradient_edges([0], 'org.pytorch.aten', 'ATen', 'aten::multinomial', '') -@register_gradient('com.microsoft', 'ATenOp', 'aten::binary_cross_entropy_with_logits', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::binary_cross_entropy_with_logits', '') def binary_cross_entropy_with_logits_gradient(): return [ - (('ATenOp', 'com.microsoft'), ['GO(0)', 'I(0)', 'I(1)', 'I(2)', 'I(3)', 'I(4)'], [ - 'GI(0)'], {'name': {'value': 'aten::binary_cross_entropy_with_logits_backward', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)', 'I(0)', 'I(1)', 'I(2)', 'I(3)', 'I(4)'], [ + 'GI(0)'], {'operator': {'value': 'aten::binary_cross_entropy_with_logits_backward', 'dtype': 'string'}}), ] -@register_gradient('com.microsoft', 'ATenOp', 'aten::numpy_T', '') +@register_gradient('org.pytorch.aten', 'ATen', 'aten::numpy_T', '') def numpy_T_gradient(): return [ - (('ATenOp', 'com.microsoft'), ['GO(0)'], [ - 'GI(0)'], {'name': {'value': 'aten::numpy_T', 'dtype': 'string'}}), + (('ATen', 'org.pytorch.aten'), ['GO(0)'], [ + 'GI(0)'], {'operator': {'value': 'aten::numpy_T', 'dtype': 'string'}}), ] diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py index 5ebeec2fcc..a721fc64a2 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py @@ -64,8 +64,8 @@ def nll_loss(g, self, target, weight, reduction, ignore_index): @register_symbolic('embedding') def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): - output = g.op("com.microsoft::ATenOp", weight, indices, padding_idx, scale_grad_by_freq, sparse, - name_s='aten::embedding') + output = g.op("org.pytorch.aten::ATen", weight, indices, padding_idx, scale_grad_by_freq, sparse, + operator_s='aten::embedding') indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), 'with_sizes'): output_type = weight.type().with_sizes( @@ -75,21 +75,21 @@ def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): @register_symbolic('bitwise_or') def bitwise_or(g, self, other): - return g.op("com.microsoft::ATenOp", self, other, - name_s='aten::bitwise_or', overload_name_s='Tensor') + return g.op("org.pytorch.aten::ATen", self, other, + operator_s='aten::bitwise_or', overload_name_s='Tensor') @register_symbolic('diagonal') def diagonal(g, self, offset, dim1, dim2): - return g.op("com.microsoft::ATenOp", self, offset, dim1, dim2, - name_s='aten::diagonal') + return g.op("org.pytorch.aten::ATen", self, offset, dim1, dim2, + operator_s='aten::diagonal') @register_symbolic('multinomial') def multinomial(g, self, num_samples, replacement=False, generator=None): if generator is not None and not sym_help._is_none(generator): raise RuntimeError("Unsupported: ONNX does not support generator for multinomial") - return g.op("com.microsoft::ATenOp", self, num_samples, replacement, generator, - name_s='aten::multinomial') + return g.op("org.pytorch.aten::ATen", self, num_samples, replacement, generator, + operator_s='aten::multinomial') @register_symbolic('max_pool2d') @@ -97,18 +97,18 @@ def max_pool2d(g, self, kernel_size, stride, padding, dilation, ceil_mode): stride_val = sym_help._maybe_get_const(stride, 'is') if not stride_val: stride = kernel_size - return g.op("com.microsoft::ATenOp", self, kernel_size, stride, padding, dilation, ceil_mode, - name_s='aten::max_pool2d_with_indices', outputs=2)[0] + return g.op("org.pytorch.aten::ATen", self, kernel_size, stride, padding, dilation, ceil_mode, + operator_s='aten::max_pool2d_with_indices', outputs=2)[0] @register_symbolic('unfold') def unfold(g, input, dimension, size, step): - return g.op("com.microsoft::ATenOp", input, dimension, size, step, name_s='aten::unfold') + return g.op("org.pytorch.aten::ATen", input, dimension, size, step, operator_s='aten::unfold') @register_symbolic('argmax') def argmax(g, input, dim, keepdim): - return g.op("com.microsoft::ATenOp", input, dim, keepdim, name_s='aten::argmax') + return g.op("org.pytorch.aten::ATen", input, dim, keepdim, operator_s='aten::argmax') @register_symbolic('avg_pool2d') @@ -116,13 +116,13 @@ def avg_pool2d(g, self, kernel_size, stride, padding, ceil_mode, count_include_p stride_val = sym_help._maybe_get_const(stride, 'is') if not stride_val: stride = kernel_size - return g.op("com.microsoft::ATenOp", self, kernel_size, stride, padding, ceil_mode, - count_include_pad, divisor_override, name_s='aten::avg_pool2d') + return g.op("org.pytorch.aten::ATen", self, kernel_size, stride, padding, ceil_mode, + count_include_pad, divisor_override, operator_s='aten::avg_pool2d') @register_symbolic('adaptive_avg_pool2d') def adaptive_avg_pool2d(g, self, output_size): - return g.op("com.microsoft::ATenOp", self, output_size, name_s='aten::_adaptive_avg_pool2d') + return g.op("org.pytorch.aten::ATen", self, output_size, operator_s='aten::_adaptive_avg_pool2d') @register_symbolic('binary_cross_entropy_with_logits') @@ -131,8 +131,8 @@ def binary_cross_entropy_with_logits(g, self, target, weight, pos_weight, reduct # But current custom_gradient_registry doesn't support such None checking, # So doesn't support non-None weight for now. if weight is None or sym_help._is_none(weight): - return g.op("com.microsoft::ATenOp", self, target, weight, pos_weight, reduction, - name_s='aten::binary_cross_entropy_with_logits') + return g.op("org.pytorch.aten::ATen", self, target, weight, pos_weight, reduction, + operator_s='aten::binary_cross_entropy_with_logits') from torch.onnx.symbolic_opset12 import binary_cross_entropy_with_logits as bce return bce(g, self, target, weight, pos_weight, reduction) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index ede9ffb2ab..7a1f0a07bc 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -179,7 +179,7 @@ class GraphExecutionManager(GraphExecutionInterface): # Re-export will be avoided if _skip_check is enabled. self._original_model_has_changed = False - # Load ATenOp executor extension. + # Load ATen operator executor extension. load_aten_op_executor_cpp_extension() def _get_torch_gpu_allocator_function_addresses(self): diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index 131c8fa80c..23dfea316a 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -343,7 +343,7 @@ def reinitialize_graph_execution_manager(graph_execution_manager): # Re-define the torch allocator graph_execution_manager._get_torch_gpu_allocator_function_addresses() - # Load ATenOp executor extension. + # Load ATen op executor extension. load_aten_op_executor_cpp_extension() def reinitialize_training_manager(training_manager): 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 3b6c9d2f47..d6b84a5425 100644 --- a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.cc +++ b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.cc @@ -10,11 +10,11 @@ namespace onnxruntime { namespace contrib { -ONNX_OPERATOR_KERNEL_EX(ATenOp, kMSDomain, 1, kCpuExecutionProvider, +ONNX_OPERATOR_KERNEL_EX(ATen, kPytorchAtenDomain, 1, kCpuExecutionProvider, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorAndSequenceTensorTypes()), - ATenOp); + ATen); -Status ATenOp::Compute(OpKernelContext* p_ctx) const { +Status ATen::Compute(OpKernelContext* p_ctx) const { auto* p_ctx_internal = static_cast(p_ctx); std::vector dlpacks; for (int i = 0; i < p_ctx_internal->InputCount(); i++) { @@ -39,7 +39,7 @@ bool IsATenOperatorExecutorInitialized() { return aten_ops::ATenOperatorExecutor::Instance().IsInitialized(); } -Status ExecuteReduceSumATenOp(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) { +Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) { ORT_ENFORCE(aten_ops::ATenOperatorExecutor::Instance().IsInitialized() && !axes.empty()); std::vector dlpacks; auto* p_ctx_internal = static_cast(p_ctx); 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 3ab071a271..1a80af45c5 100644 --- a/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.h +++ b/orttraining/orttraining/training_ops/cpu/aten_ops/aten_op.h @@ -8,10 +8,10 @@ namespace onnxruntime { namespace contrib { -class ATenOp : public OpKernel { +class ATen : public OpKernel { public: - ATenOp(const OpKernelInfo& info) : OpKernel(info) { - ORT_THROW_IF_ERROR(info.GetAttr("name", &op_name_)); + ATen(const OpKernelInfo& info) : OpKernel(info) { + ORT_THROW_IF_ERROR(info.GetAttr("operator", &op_name_)); overload_name_ = info.GetAttrOrDefault("overload_name", ""); } @@ -23,7 +23,7 @@ class ATenOp : public OpKernel { }; bool IsATenOperatorExecutorInitialized(); -Status ExecuteReduceSumATenOp(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims); +Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims); } // namespace contrib } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc index 800fd48221..2475ffb9bb 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -98,7 +98,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Recv) class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, RecordEvent); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WaitEvent); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, YieldOp); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, ATenOp); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kPytorchAtenDomain, 1, ATen); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryScalar); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryHistogram); @@ -208,7 +208,7 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/aten_ops/aten_op.cc b/orttraining/orttraining/training_ops/cuda/aten_ops/aten_op.cc index fc07e03412..3a1594d87f 100644 --- a/orttraining/orttraining/training_ops/cuda/aten_ops/aten_op.cc +++ b/orttraining/orttraining/training_ops/cuda/aten_ops/aten_op.cc @@ -8,9 +8,9 @@ namespace onnxruntime { namespace cuda { -ONNX_OPERATOR_KERNEL_EX(ATenOp, kMSDomain, 1, kCudaExecutionProvider, +ONNX_OPERATOR_KERNEL_EX(ATen, kPytorchAtenDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::AllTensorAndSequenceTensorTypes()), - onnxruntime::contrib::ATenOp); + onnxruntime::contrib::ATen); } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 3a0d082d34..5af100a265 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -215,7 +215,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Adas class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, RecordEvent); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, WaitEvent); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, YieldOp); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ATenOp); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kPytorchAtenDomain, 1, ATen); #ifdef ENABLE_TRAINING_TORCH_INTEROP class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, PythonOp); @@ -433,7 +433,7 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, #ifdef ENABLE_TRAINING_TORCH_INTEROP BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc index 7b43f91972..f94f447ac9 100644 --- a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc +++ b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc @@ -191,7 +191,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, Adas class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, RecordEvent); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, WaitEvent); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, YieldOp); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, ATenOp); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kPytorchAtenDomain, 1, ATen); #ifdef ENABLE_TRAINING_TORCH_INTEROP class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, PythonOp); @@ -384,7 +384,7 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, #ifdef ENABLE_TRAINING_TORCH_INTEROP BuildKernelCreateInfo, diff --git a/orttraining/tools/scripts/performance_investigation.py b/orttraining/tools/scripts/performance_investigation.py index b8fb648c18..295eebe2d5 100644 --- a/orttraining/tools/scripts/performance_investigation.py +++ b/orttraining/tools/scripts/performance_investigation.py @@ -27,9 +27,9 @@ def process_file(onnx_file): memcpu_ops.append(f"{node.op_type} {node.name}") if node.op_type == "Cast": cast_ops.append(f"{node.name}") - if node.op_type == "ATenOp": + if node.op_type == "ATen": for attr in node.attribute: - if attr.name == "name": + if attr.name == "operator": aten_ops.append(f"{node.name}: {attr.s.decode('utf-8')}") if node.op_type == "PythonOp": for attr in node.attribute: @@ -51,7 +51,7 @@ def process_file(onnx_file): msgs.append(f"Examine whether {node.name} should be fused with the leading {prev.name} op into BiasSoftmax node.") if aten_ops: - print("ATenOp found:") + print("ATen op found:") for line in aten_ops: print(line) print(10 * '-')