diff --git a/orttraining/orttraining/core/framework/gradient_graph_builder.h b/orttraining/orttraining/core/framework/gradient_graph_builder.h index 1d5c239b0c..634b2a569c 100644 --- a/orttraining/orttraining/core/framework/gradient_graph_builder.h +++ b/orttraining/orttraining/core/framework/gradient_graph_builder.h @@ -39,7 +39,8 @@ static std::unordered_map> {"Scatter", {1}}, {"OneHot", {0, 1, 2}}, {"Where", {0}}, - {"Range", {0, 1, 2}}}; + {"Range", {0, 1, 2}}, + {"BroadcastGradientArgs", {0, 1}}}; class GradientGraphBuilder { public: diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 2820b09b57..1b835fc080 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -950,7 +950,7 @@ Example 4: updateOutputShape(ctx, 0, TensorShapeProto()); } - if(ctx.getNumOutputs() == 2) { + if (ctx.getNumOutputs() == 2) { propagateElemTypeFromInputToOutput(ctx, 0, 1); if (hasInputShape(ctx, 0)) { propagateShapeFromInputToOutput(ctx, 0, 1); @@ -1167,6 +1167,23 @@ Example 4: propagateShapeAndTypeFromFirstInput(ctx); }); + ONNX_CONTRIB_OPERATOR_SCHEMA(BroadcastGradientArgs) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) + .SetDoc( + "Returns the reduction axes for computing gradients of s0 op s1 with broadcast." + "The ouput axes are deterministic from last to first. " + "Output is an empty vector when no reduction is necessary for the corresponding input.") + .Input(0, "a_shape", "The 1st input shape as Tensor.", "T") + .Input(1, "b_shape", "The 2nd input shape as Tensor.", "T") + .Output(0, "a_axes", "The reduction axes for 1st input, last to first.", "T") + .Output(1, "b_axes", "The reduction axes for 2nd input, last to first.", "T") + .TypeConstraint( + "T", + {"tensor(int64)"}, + "Constrain input and output types to 64-bit integer."); + ONNX_CONTRIB_OPERATOR_SCHEMA(GistBinarizeEncoder) .SetDomain(kMSDomain) .SinceVersion(1) diff --git a/orttraining/orttraining/test/training_ops/cpu/nn/broadcast_grad_args_test.cc b/orttraining/orttraining/test/training_ops/cpu/nn/broadcast_grad_args_test.cc new file mode 100644 index 0000000000..a10c919b10 --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cpu/nn/broadcast_grad_args_test.cc @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace contrib { +namespace test { + +using namespace onnxruntime::test; + +namespace { +constexpr auto k_opset_version = 1; + +void RunBroadcastGradientArgsTest(const char* op, + const std::vector& A_shape_tensor, + const std::vector& B_shape_tensor, + const std::vector& A_axes_expected, + const std::vector& B_axes_expected, + bool fail = false) { + OpTester t{op, k_opset_version, kMSDomain}; + + t.AddInput("a_shape", {static_cast(A_shape_tensor.size())}, A_shape_tensor); + t.AddInput("b_shape", {static_cast(B_shape_tensor.size())}, B_shape_tensor); + + t.AddOutput("a_axes", {static_cast(A_axes_expected.size())}, A_axes_expected); + t.AddOutput("b_axes", {static_cast(B_axes_expected.size())}, B_axes_expected); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + if (fail) + t.Run(OpTester::ExpectResult::kExpectFailure, "", {}, nullptr, &execution_providers, ExecutionMode::ORT_SEQUENTIAL); + else + t.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers, ExecutionMode::ORT_SEQUENTIAL); +} + +} // namespace + +// BroadcastGradientArgs + +TEST(BroadcastGradientArgsTest, Basic) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 16, 1024, 1024}, {1, 1, 1024, 1024}, + {}, {1, 0}); +} + +TEST(BroadcastGradientArgsTest, Basic_both_valid_op) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 16, 1, 1024}, {1, 1, 1024, 1024}, + {2}, {1, 0}); +} + +TEST(BroadcastGradientArgsTest, Basic_no_bcast) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 3, 4, 5}, {2, 3, 4, 5}, + {}, {}); +} + +TEST(BroadcastGradientArgsTest, Basic_B_scalar) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 3, 4, 5}, {}, + {}, {3, 2, 1, 0}); +} + +TEST(BroadcastGradientArgsTest, Basic_B_vector) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 3, 4, 5}, {5}, + {}, {2, 1, 0}); +} + +TEST(BroadcastGradientArgsTest, Basic_A_bcast_different_size) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {4, 5}, {2, 3, 4, 5}, + {1, 0}, {}); +} + +TEST(BroadcastGradientArgsTest, Basic_both_bcast_different_size) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {1, 4, 5}, {2, 3, 1, 1}, + {1, 0}, {3, 2}); +} + +TEST(BroadcastGradientArgsTest, Basic_both_bcast_different_size_2) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {3, 4, 5}, {2, 1, 1, 1}, + {0}, {3, 2, 1}); +} + +TEST(BroadcastGradientArgsTest, Basic_invalid_broadcast) { + RunBroadcastGradientArgsTest("BroadcastGradientArgs", {3, 4, 5}, {2, 1, 6, 1}, + {}, {}, true /*fail*/); +} + +} // namespace test +} // 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 b974d35027..f1f668821b 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -16,6 +16,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, AdamO class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, InPlaceAccumulator); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, ZeroGradient); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Group); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int64_t, BroadcastGradientArgs); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ReduceSumTraining); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, ReduceSumTraining); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int32_t, ReduceSumTraining); @@ -79,6 +80,7 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.cc b/orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.cc new file mode 100644 index 0000000000..aea6b7832d --- /dev/null +++ b/orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.cc @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "orttraining/training_ops/cpu/nn/broadcast_grad_args.h" + +namespace onnxruntime { +namespace contrib { +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + BroadcastGradientArgs, \ + kMSDomain, \ + 1, \ + T, \ + kCpuExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + BroadcastGradientArgs); + +REGISTER_KERNEL_TYPED(int64_t) + +template +Status BroadcastGradientArgs::Compute(OpKernelContext* context) const { + const Tensor* a_shape = context->Input(0); + const Tensor* b_shape = context->Input(1); + const T* A_dims = a_shape->template Data(); + const T* B_dims = b_shape->template Data(); + + const T a_size = a_shape->Shape().Size(); + const T b_size = b_shape->Shape().Size(); + + T ndim = std::max(a_size, b_size); + std::vector a_axes, b_axes; + + T i = a_size - 1; + T j = b_size - 1; + T k = ndim - 1; + + for (; i >= 0 && j >= 0; --k) { + auto A_dim = A_dims[i], + B_dim = B_dims[j]; + + if (A_dim != B_dim) { + if (A_dim == 1) { + a_axes.push_back(k); + } else if (B_dim == 1) { + b_axes.push_back(k); + } else { + TensorShape a(A_dims, a_size); + TensorShape b(B_dims, b_size); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Broadcast is not possible between inputs of shapes: ", + a, " and ", b); + } + } + --i; + --j; + } + + if (i < 0) { + for (; k >= 0; --k) { + a_axes.push_back(k); + } + + } else { + for (; k >= 0; --k) { + b_axes.push_back(k); + } + } + + Tensor* A_axes = context->Output(0, {static_cast(a_axes.size())}); + T* A_axes_data = A_axes->template MutableData(); + std::copy(a_axes.begin(), a_axes.end(), A_axes_data); + Tensor* B_axes = context->Output(1, {static_cast(b_axes.size())}); + T* B_axes_data = B_axes->template MutableData(); + std::copy(b_axes.begin(), b_axes.end(), B_axes_data); + + return Status::OK(); +} + +} // namespace contrib +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.h b/orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.h new file mode 100644 index 0000000000..3dcc7b7a32 --- /dev/null +++ b/orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { +template +class BroadcastGradientArgs final : public OpKernel { + public: + BroadcastGradientArgs(const OpKernelInfo& info) : OpKernel{info} { + } + + Status Compute(OpKernelContext* context) const override; +}; +} // namespace contrib +} // namespace onnxruntime