Create the ConcatTraining op (#4595)

* Working changes for ConcatTraining op

* Refactor to move changes to orttraining

* Fix segfault

* Support -ve axis for shape inferencing

* fix build

Co-authored-by: Ethan Tao <ettao@microsoft.com>

onnxruntime/core/providers/cpu/tensor/concat.cc

@@ -21,7 +21,7 @@ ONNX_CPU_OPERATOR_KERNEL(
     KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorTypes()),
     Concat);
 
-// this method will be shared between 'Concat' (CPU and GPU) and 
+// this method will be shared between 'Concat' (CPU and GPU) and
 // 'ConcatFromSequence' ('concat' and 'stack' modes) to validate inputs
 Status ConcatBase::PrepareForCompute(OpKernelContext* ctx,
                                      const std::vector<const Tensor*>& input_tensors,
@@ -169,7 +169,7 @@ Status ConcatBase::ComputeImpl(Prepare& p) const {
     auto input_size = prep.num_elements;
 
     // Copy the data across. For every 'input_axis_pitch' values copied, we move over by the 'output_axis_pitch'
-    // TODO: Optimization possibility: There are cases where we simply need to "merge" raw buffers and this 
+    // TODO: Optimization possibility: There are cases where we simply need to "merge" raw buffers and this
     // could be done without the pointer house-keeping as below. Some scenarios whether this is possible are:
     // 1) Concatenating on input axis = 0
     // 2) Stacking on output axis = 0

onnxruntime/core/providers/cpu/tensor/concat.h

@@ -39,7 +39,7 @@ class ConcatBase {
     }
   }
 
-  // the core method that will be invoked by the 'Concat' (CPU and GPU) 
+  // the core method that will be invoked by the 'Concat' (CPU and GPU)
   // and 'ConcatFromSequence' kernels
   Status PrepareForCompute(OpKernelContext* ctx, const std::vector<const Tensor*>& input_tensors,
                            Prepare& p) const;

onnxruntime/core/providers/cuda/tensor/concat.cc

@@ -80,6 +80,5 @@ Status Concat::ComputeInternal(OpKernelContext* ctx) const {
                                  p.output_num_elements));
   return Status::OK();
 }
-
 }  // namespace cuda
 }  // namespace onnxruntime

orttraining/orttraining/core/graph/training_op_defs.cc

@@ -1107,7 +1107,7 @@ Example 4:
         auto output_shape =
             ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
 
-        std::vector<int64_t> axes_values = ParseData<int64_t>(axes_proto); 
+        std::vector<int64_t> axes_values = ParseData<int64_t>(axes_proto);
         std::vector<int64_t> axes;
         axes.reserve(axes_values.size());
         for (int64_t axis : axes_values) {
@@ -1128,7 +1128,87 @@ Example 4:
           }
         }
       });
-      
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(ConcatTraining)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL)
+      .SetDoc("Concatenate a list of tensors into a single tensor")
+      .Attr("axis", "Which axis to concat on", AttributeProto::INT)
+      .Input(0,
+             "inputs",
+             "List of tensors for concatenation",
+             "T",
+             OpSchema::Variadic)
+      .Output(0, "concat_result", "Concatenated tensor", "T")
+      .Output(1, "per_input_length",
+              "Vector of length of each concatenated "
+              "input along the 'axis' dimension",
+              "Tint")
+      .TypeConstraint(
+          "T",
+          OpSchema::all_tensor_types(),
+          "Constrain output types to any tensor type.")
+      .TypeConstraint(
+          "Tint",
+          {"tensor(int64)"},
+          "Constrain output len types to integer type.")
+      .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+        propagateElemTypeFromInputToOutput(ctx, 0, 0);
+        auto numInputs = ctx.getNumInputs();
+        if (numInputs < 1 ||
+            !hasNInputShapes(ctx, static_cast<int>(numInputs))) {
+          return;
+        }
+
+        auto rank = ctx.getInputType(0)->tensor_type().shape().dim_size();
+
+        auto axisAttr = ctx.getAttribute("axis");
+        if (!axisAttr) {
+          fail_shape_inference("Required attribute axis is missing");
+        }
+        int64_t axis = static_cast<int64_t>(axisAttr->i());
+        axis = HandleNegativeAxis(axis, rank);
+
+        bool all_lengths_known = true;
+        int total_length = 0;
+
+        auto* output_shape =
+            ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+
+        for (int64_t i = 0; i < rank; ++i) {
+          output_shape->add_dim();
+        }
+
+        ONNX_NAMESPACE::TensorShapeProto per_input_len_shape;
+        per_input_len_shape.add_dim()->set_dim_value(numInputs);
+        updateOutputShape(ctx, 1, per_input_len_shape);
+
+        for (size_t i = 0; i < numInputs; i++) {
+          const auto& shape = ctx.getInputType(i)->tensor_type().shape();
+          if (shape.dim_size() != rank)
+            fail_shape_inference("All inputs to Concat must have same rank");
+          for (int j = 0; j < rank; j++) {
+            if (j == axis) {
+              if (shape.dim(j).has_dim_value()) {
+                total_length += static_cast<int>(shape.dim(j).dim_value());
+              } else {
+                all_lengths_known = false;
+              }
+            } else {
+              auto& output_dim = *output_shape->mutable_dim(j);
+              const auto& input_dim = shape.dim(j);
+              mergeInDimensionInfo(input_dim, output_dim, j);
+            }
+          }
+        }
+
+        if (all_lengths_known) {
+          output_shape->mutable_dim(static_cast<int>(axis))->set_dim_value(total_length);
+        }
+
+      });
+
   ONNX_CONTRIB_OPERATOR_SCHEMA(TrainableDropout)
       .SetDomain(kOnnxDomain)
       .SinceVersion(9)

orttraining/orttraining/test/training_ops/cpu/tensor/concat_op_test.cc

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(ConcatTrainingOpTest, Concat1D_int32_negative_axis) {
+  OpTester test("ConcatTraining", 1, kMSDomain);
+  test.AddAttribute("axis", int64_t{-1});
+
+  test.AddInput<int32_t>("input1", {1}, {1});
+  test.AddInput<int32_t>("input2", {2}, {2, 3});
+  test.AddInput<int32_t>("input3", {4}, {4, 5, 6, 7});
+  test.AddOutput<int32_t>("concat_result", {7}, {1, 2, 3, 4, 5, 6, 7});
+  test.AddOutput<int64_t>("per_input_length", {3}, {1, 2, 4});
+  test.Run();
+}
+
+TEST(ConcatTrainingOpTest, Concat2D_float_axis1) {
+  OpTester test("ConcatTraining", 1, kMSDomain);
+  test.AddAttribute("axis", int64_t{1});
+
+  std::vector<int64_t> dims{4, 1};
+  test.AddInput<float>("input1", dims, {11.0f, 21.0f, 31.0f, 41.0f});
+  test.AddInput<float>("input2", {4, 2}, {12.0f, 13.0f, 22.0f, 23.0f, 32.0f, 33.0f, 42.0f, 43.0f});
+  test.AddInput<float>("input3", dims, {14.0f, 24.0f, 34.0f, 44.0f});
+  test.AddOutput<float>("concat_result", {4, 4},
+                        {11.0f, 12.0f, 13.0f, 14.0f,
+                         21.0f, 22.0f, 23.0f, 24.0f,
+                         31.0f, 32.0f, 33.0f, 34.0f,
+                         41.0f, 42.0f, 43.0f, 44.0f});
+  test.AddOutput<int64_t>("per_input_length", {3}, {1, 2, 1});
+  test.Run();
+}
+
+TEST(ConcatTrainingOpTest, Concat3D_same_len) {
+  OpTester test("ConcatTraining", 1, kMSDomain);
+  test.AddAttribute("axis", int64_t{1});
+
+  std::vector<int64_t> dims{2, 2, 2};
+  test.AddInput<float>("input1", dims,
+                       {1.0f, 2.0f,
+                        3.0f, 4.0f,
+
+                        5.0f, 6.0f,
+                        7.0f, 8.0f});
+  test.AddInput<float>("input2", dims,
+                       {9.0f, 10.0f,
+                        11.0f, 12.0f,
+
+                        13.0f, 14.0f,
+                        15.0f, 16.0f});
+  test.AddOutput<float>("concat_result", {2, 4, 2},
+                        {1.0f, 2.0f,
+                         3.0f, 4.0f,
+                         9.0f, 10.0f,
+                         11.0f, 12.0f,
+
+                         5.0f, 6.0f,
+                         7.0f, 8.0f,
+                         13.0f, 14.0f,
+                         15.0f, 16.0f});
+  test.AddOutput<int64_t>("per_input_length", {2}, {2, 2});
+  test.Run();
+}
+
+}  // namespace test
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc

@@ -21,6 +21,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
 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);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int64_t, ReduceSumTraining);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, ConcatTraining);
 
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SoftmaxCrossEntropy);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SoftmaxCrossEntropyGrad);
@@ -110,6 +111,7 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, ReduceSumTraining)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int32_t, ReduceSumTraining)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int64_t, ReduceSumTraining)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, ConcatTraining)>,
 
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SoftmaxCrossEntropy)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SoftmaxCrossEntropyGrad)>,

orttraining/orttraining/training_ops/cpu/tensor/concat.cc

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "orttraining/training_ops/cpu/tensor/concat.h"
+#include "core/providers/common.h"
+#include "core/framework/TensorSeq.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+ONNX_OPERATOR_KERNEL_EX(
+    ConcatTraining,
+    kMSDomain,
+    1,
+    kCpuExecutionProvider,
+    KernelDefBuilder()
+        .TypeConstraint("T", DataTypeImpl::AllTensorTypes()),
+    ConcatTraining);
+
+// core Compute() method for the 'Concat' kernel
+Status ConcatTraining::Compute(OpKernelContext* ctx) const {
+  // Number of input tensors to concatenate
+  auto input_count = Node().InputArgCount().front();
+
+  // Hold pointers to the input tensors to be used in the PrepareForCompute() step
+  std::vector<const Tensor*> input_tensors;
+  input_tensors.reserve(input_count);
+  for (int i = 0; i < input_count; ++i) {
+    input_tensors.push_back(ctx->Input<Tensor>(i));
+  }
+
+  // Validate inputs and prepare some metadata used during actual compute
+  Prepare p;
+  auto status = PrepareForCompute(ctx, input_tensors, p);
+  if (!status.IsOK())
+    return status;
+
+  // Return at this point if output tensor is going to be empty
+  if (p.output_num_elements == 0)
+    return Status::OK();
+
+  // Create output tensor for 'per_input_length'
+  std::vector<int64_t> per_input_length(input_count);
+  for (int i = 0; i < input_count; ++i) {
+    per_input_length[i] = input_tensors[i]->Shape()[p.axis];
+  }
+  Tensor* output_1_tensor = ctx->Output(1, {input_count});
+  int64_t* output_1_tensor_data = output_1_tensor->template MutableData<int64_t>();
+  std::copy(per_input_length.begin(), per_input_length.end(), output_1_tensor_data);
+
+  // Compute values to be placed in the output tensor
+  return ComputeImpl(p);
+}
+
+}  // namespace contrib
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cpu/tensor/concat.h

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/util/math_cpuonly.h"
+#include "core/framework/tensor.h"
+#include "core/providers/cpu/tensor/concat.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+class ConcatTraining : public OpKernel, public ConcatBase {
+ public:
+  ConcatTraining(const OpKernelInfo& info) : OpKernel(info), ConcatBase(info) {}
+
+  Status Compute(OpKernelContext* context) const override;
+};
+
+}  // namespace contrib
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc

@@ -16,6 +16,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, ReduceSumTraining);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int32_t, ReduceSumTraining);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ReduceSumTraining);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ConcatTraining);
+
 // Adam
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_float_float, AdamOptimizer);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float_MLFloat16_float_float, AdamOptimizer);
@@ -132,7 +134,7 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, ReduceSumTraining)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int32_t, ReduceSumTraining)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ReduceSumTraining)>,
-
+    BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ConcatTraining)>,
     // Adam
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_float_float, AdamOptimizer)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float_MLFloat16_float_float, AdamOptimizer)>,

orttraining/orttraining/training_ops/cuda/tensor/concat.cc

@@ -0,0 +1,79 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "orttraining/training_ops/cuda/tensor/concat.h"
+#include "core/providers/cuda/tensor/concat_impl.h"
+
+namespace onnxruntime {
+namespace cuda {
+ONNX_OPERATOR_KERNEL_EX(ConcatTraining,
+                        kMSDomain,
+                        1,
+                        kCudaExecutionProvider,
+                        KernelDefBuilder()
+                            .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()),
+                        ConcatTraining);
+
+Status ConcatTraining::ComputeInternal(OpKernelContext* ctx) const {
+  auto input_count = Node().InputArgCount().front();
+
+  // Hold pointers to the input tensors to be used in the PrepareForCompute() step
+  std::vector<const Tensor*> input_tensors;
+  input_tensors.reserve(input_count);
+  for (int i = 0; i < input_count; ++i) {
+    input_tensors.push_back(ctx->Input<Tensor>(i));
+  }
+
+  Prepare p;
+  ORT_RETURN_IF_ERROR(PrepareForCompute(ctx, input_tensors, p));
+
+  // Return at this point if output tensor is going to be empty
+  if (p.output_num_elements == 0)
+    return Status::OK();
+
+  std::vector<int64_t> concat_sizes(input_count);
+
+  CudaAsyncBuffer<const void*> input_ptr(this, input_count);
+  gsl::span<const void*> input_ptr_cpuspan = input_ptr.CpuSpan();
+  std::vector<int64_t> axis_dimension_input_output_mapping(p.output_tensor->Shape()[p.axis]);
+  int index = 0;
+  for (int i = 0; i < input_count; ++i) {
+    auto input = p.inputs[i];
+    concat_sizes[i] = input.tensor->Shape()[p.axis];
+    input_ptr_cpuspan[i] = input.tensor->DataRaw();
+    for (int j = 0; j < input.tensor->Shape()[p.axis]; ++j) {
+      axis_dimension_input_output_mapping.at(index++) = i;
+    }
+  }
+  std::vector<int64_t> concat_sizes_range(concat_sizes);
+  for (size_t i = 1; i < concat_sizes_range.size(); ++i) {
+    concat_sizes_range[i] += concat_sizes_range[i - 1];
+  }
+
+  CudaAsyncBuffer<int64_t> concat_sizes_gpu(this, concat_sizes);
+  CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, axis_dimension_input_output_mapping);
+  CudaAsyncBuffer<int64_t> concat_sizes_range_gpu(this, concat_sizes_range);
+  concat_sizes_gpu.CopyToGpu();
+  axis_dimension_input_output_mapping_gpu.CopyToGpu();
+  concat_sizes_range_gpu.CopyToGpu();
+  input_ptr.CopyToGpu();
+  int block_size_inside_axis_dim = static_cast<int>(p.output_axis_pitch / p.output_tensor->Shape()[p.axis]);
+  int block_size_including_axis_dim = static_cast<int>(p.output_axis_pitch);
+  auto element_bytes = p.output_tensor->DataType()->Size();
+  ORT_RETURN_IF_ERROR(ConcatImpl(element_bytes,
+                                 block_size_including_axis_dim,
+                                 block_size_inside_axis_dim,
+                                 concat_sizes_gpu.GpuPtr(),
+                                 concat_sizes_range_gpu.GpuPtr(),
+                                 axis_dimension_input_output_mapping_gpu.GpuPtr(),
+                                 p.output_tensor->MutableDataRaw(),
+                                 input_ptr.GpuPtr(),
+                                 p.output_num_elements));
+
+  Tensor* output_1_tensor = ctx->Output(1, {input_count});
+  CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_1_tensor->template MutableData<int64_t>(), concat_sizes_gpu.GpuPtr(), input_count * sizeof(int64_t), cudaMemcpyDeviceToDevice));
+
+  return Status::OK();
+}
+}  // namespace cuda
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/tensor/concat.h

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#pragma once
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cpu/tensor/concat.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+class ConcatTraining final : public CudaKernel, public ConcatBase {
+ public:
+  ConcatTraining(const OpKernelInfo& info) : CudaKernel(info), ConcatBase(info) {}
+  Status ComputeInternal(OpKernelContext* context) const override;
+};
+
+}  // namespace cuda
+}  // namespace onnxruntime