diff --git a/onnxruntime/core/providers/cpu/tensor/concat.cc b/onnxruntime/core/providers/cpu/tensor/concat.cc index 14c9526623..6ae60bc5aa 100644 --- a/onnxruntime/core/providers/cpu/tensor/concat.cc +++ b/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& 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 diff --git a/onnxruntime/core/providers/cpu/tensor/concat.h b/onnxruntime/core/providers/cpu/tensor/concat.h index e50de54f0f..7752f042a6 100644 --- a/onnxruntime/core/providers/cpu/tensor/concat.h +++ b/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& input_tensors, Prepare& p) const; diff --git a/onnxruntime/core/providers/cuda/tensor/concat.cc b/onnxruntime/core/providers/cuda/tensor/concat.cc index 4a10297957..8792e95895 100644 --- a/onnxruntime/core/providers/cuda/tensor/concat.cc +++ b/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 diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 4b5d71a3bc..16438e4e05 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/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 axes_values = ParseData(axes_proto); + std::vector axes_values = ParseData(axes_proto); std::vector 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(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(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(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(axis))->set_dim_value(total_length); + } + + }); + ONNX_CONTRIB_OPERATOR_SCHEMA(TrainableDropout) .SetDomain(kOnnxDomain) .SinceVersion(9) diff --git a/orttraining/orttraining/test/training_ops/cpu/tensor/concat_op_test.cc b/orttraining/orttraining/test/training_ops/cpu/tensor/concat_op_test.cc new file mode 100644 index 0000000000..26bf993014 --- /dev/null +++ b/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("input1", {1}, {1}); + test.AddInput("input2", {2}, {2, 3}); + test.AddInput("input3", {4}, {4, 5, 6, 7}); + test.AddOutput("concat_result", {7}, {1, 2, 3, 4, 5, 6, 7}); + test.AddOutput("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 dims{4, 1}; + test.AddInput("input1", dims, {11.0f, 21.0f, 31.0f, 41.0f}); + test.AddInput("input2", {4, 2}, {12.0f, 13.0f, 22.0f, 23.0f, 32.0f, 33.0f, 42.0f, 43.0f}); + test.AddInput("input3", dims, {14.0f, 24.0f, 34.0f, 44.0f}); + test.AddOutput("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("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 dims{2, 2, 2}; + test.AddInput("input1", dims, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f}); + test.AddInput("input2", dims, + {9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f}); + test.AddOutput("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("per_input_length", {2}, {2, 2}); + test.Run(); +} + +} // namespace test +} // 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 e567bf1394..bbfdcd41d8 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cpu/tensor/concat.cc b/orttraining/orttraining/training_ops/cpu/tensor/concat.cc new file mode 100644 index 0000000000..d3bf9502ba --- /dev/null +++ b/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 input_tensors; + input_tensors.reserve(input_count); + for (int i = 0; i < input_count; ++i) { + input_tensors.push_back(ctx->Input(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 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(); + 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 diff --git a/orttraining/orttraining/training_ops/cpu/tensor/concat.h b/orttraining/orttraining/training_ops/cpu/tensor/concat.h new file mode 100644 index 0000000000..4dfac38a1f --- /dev/null +++ b/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 diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 478a18c6d6..9444b30776 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/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, BuildKernelCreateInfo, BuildKernelCreateInfo, - + BuildKernelCreateInfo, // Adam BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/tensor/concat.cc b/orttraining/orttraining/training_ops/cuda/tensor/concat.cc new file mode 100644 index 0000000000..3185f3eb6f --- /dev/null +++ b/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 input_tensors; + input_tensors.reserve(input_count); + for (int i = 0; i < input_count; ++i) { + input_tensors.push_back(ctx->Input(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 concat_sizes(input_count); + + CudaAsyncBuffer input_ptr(this, input_count); + gsl::span input_ptr_cpuspan = input_ptr.CpuSpan(); + std::vector 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 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 concat_sizes_gpu(this, concat_sizes); + CudaAsyncBuffer axis_dimension_input_output_mapping_gpu(this, axis_dimension_input_output_mapping); + CudaAsyncBuffer 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(p.output_axis_pitch / p.output_tensor->Shape()[p.axis]); + int block_size_including_axis_dim = static_cast(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(), concat_sizes_gpu.GpuPtr(), input_count * sizeof(int64_t), cudaMemcpyDeviceToDevice)); + + return Status::OK(); +} +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/tensor/concat.h b/orttraining/orttraining/training_ops/cuda/tensor/concat.h new file mode 100644 index 0000000000..f8064f688f --- /dev/null +++ b/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