diff --git a/onnxruntime/core/providers/cpu/tensor/concat.cc b/onnxruntime/core/providers/cpu/tensor/concat.cc index d07b3ed732..5466088553 100644 --- a/onnxruntime/core/providers/cpu/tensor/concat.cc +++ b/onnxruntime/core/providers/cpu/tensor/concat.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cpu/tensor/concat.h" +#include "core/providers/common.h" namespace onnxruntime { @@ -17,6 +18,8 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); const Tensor& inputs_0 = *tensor_pointer; + auto axis = HandleNegativeAxis(axis_, inputs_0.Shape().NumDimensions()); + // Ensure all of the non concatenated axes match each other for (int index = 1; index < input_count; index++) { tensor_pointer = ctx->Input(index); @@ -25,7 +28,7 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep // Ensure all the other axes match auto dimension_count = inputs_0.Shape().NumDimensions(); for (int axis_index = 0; axis_index < dimension_count; axis_index++) { - if (axis_index == axis_) + if (axis_index == axis) continue; ONNXRUNTIME_RETURN_IF_NOT(data_n.Shape()[axis_index] == inputs_0.Shape()[axis_index], "Non concat axis dimensions must match: Axis ", axis_index, " has mismatched dimensions of ", data_n.Shape()[axis_index], " and ", inputs_0.Shape()[axis_index]); } @@ -36,19 +39,19 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep for (int index = 0; index < input_count; index++) { tensor_pointer = ctx->Input(index); if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); - concat_axis_size += tensor_pointer->Shape()[int(axis_)]; + concat_axis_size += tensor_pointer->Shape()[int(axis)]; } // Calculate the shape of the output tensor std::vector dims; for (int dimension_index = 0; dimension_index < inputs_0.Shape().NumDimensions(); dimension_index++) dims.emplace_back(inputs_0.Shape()[dimension_index]); - dims[axis_] = concat_axis_size; + dims[axis] = concat_axis_size; TensorShape outputShape(dims); // The output_axis_pitch is the number of elements to add to move to the next split axis in the output p.output_axis_pitch = 1; - for (auto i = int64_t(dims.size()); i-- > axis_;) + for (auto i = int64_t(dims.size()); i-- > axis;) p.output_axis_pitch *= dims[i]; auto& concat_result = *ctx->Output(0, outputShape); @@ -63,7 +66,7 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep // The input_axis_pitch is the number of elements to add to move to the next split axis in the input int64_t input_axis_pitch = 1; - for (int i = int(data_n.Shape().NumDimensions()); i-- > axis_;) + for (int i = int(data_n.Shape().NumDimensions()); i-- > axis;) input_axis_pitch *= data_n.Shape()[i]; p.inputs.push_back({&data_n, input_axis_pitch}); diff --git a/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc b/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc index 8977cc9fe4..b727db53ae 100644 --- a/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/concat_op_test.cc @@ -29,6 +29,17 @@ TEST(MathOpTest, Concat1D_int32) { test.Run(); } +TEST(MathOpTest, Concat1D_int32_negative_axis) { + OpTester test("Concat"); + 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.Run(); +} + TEST(MathOpTest, Concat1D) { OpTester test("Concat"); test.AddAttribute("axis", int64_t{0}); @@ -103,6 +114,38 @@ TEST(MathOpTest, Concat3D_1) { test.Run(); } +TEST(MathOpTest, Concat3D_1_negative_axis) { + OpTester test("Concat"); + test.AddAttribute("axis", int64_t{-3}); + + std::vector dims{1, 3, 3}; + test.AddInput("input1", dims, + {111.0f, 112.0f, 113.0f, + 121.0f, 122.0f, 123.0f, + 131.0f, 132.0f, 133.0f}); + test.AddInput("input2", dims, + {211.0f, 212.0f, 213.0f, + 221.0f, 222.0f, 223.0f, + 231.0f, 232.0f, 233.0f}); + test.AddInput("input3", dims, + {311.0f, 312.0f, 313.0f, + 321.0f, 322.0f, 323.0f, + 331.0f, 332.0f, 333.0f}); + test.AddOutput("concat_result", {3, 3, 3}, + {111.0f, 112.0f, 113.0f, + 121.0f, 122.0f, 123.0f, + 131.0f, 132.0f, 133.0f, + + 211.0f, 212.0f, 213.0f, + 221.0f, 222.0f, 223.0f, + 231.0f, 232.0f, 233.0f, + + 311.0f, 312.0f, 313.0f, + 321.0f, 322.0f, 323.0f, + 331.0f, 332.0f, 333.0f}); + test.Run(); +} + TEST(MathOpTest, Concat3D_2) { OpTester test("Concat"); test.AddAttribute("axis", int64_t{1});