diff --git a/onnxruntime/core/providers/cpu/tensor/split.cc b/onnxruntime/core/providers/cpu/tensor/split.cc index 39c7621a3e..1f68b7dc52 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.cc +++ b/onnxruntime/core/providers/cpu/tensor/split.cc @@ -16,9 +16,8 @@ ONNX_CPU_OPERATOR_KERNEL( KernelDefBuilder().TypeConstraint("T", std::vector{ DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), - }), + DataTypeImpl::GetTensorType()}), Split); Status SplitBase::PrepareForCompute(const TensorShape& input_shape, @@ -36,8 +35,8 @@ Status SplitBase::PrepareForCompute(const TensorShape& input_shape, before_dims = gsl::narrow(input_shape.SizeToDimension(axis)); after_dims_including_split_axis = gsl::narrow(input_shape.SizeFromDimension(axis)); after_dims_excluding_split = (axis + 1 == num_dimensions) - ? 1 // we multiply by this value so must be 1 not 0 - : gsl::narrow(input_shape.SizeFromDimension(axis + 1)); + ? 1 // we multiply by this value so must be 1 not 0 + : gsl::narrow(input_shape.SizeFromDimension(axis + 1)); if (split_sizes_.empty()) { // equal split based on number of outputs @@ -73,16 +72,25 @@ Status Split::Compute(OpKernelContext* context) const { status = ComputeImpl(*context, input); else if (data_type == DataTypeImpl::GetType()) status = ComputeImpl(*context, input); - else if (data_type == DataTypeImpl::GetType()) { - /* Need to update CopyMatrix to support double... - status = ComputeImpl(*context, input); */ - ORT_NOT_IMPLEMENTED("Split operator does not support double yet"); - } else - ORT_THROW("Invalid data type for Split operator of ", data_type); + else if (data_type == DataTypeImpl::GetType()) + status = ComputeImpl(*context, input); + else + ORT_THROW("Split operator does not support ", data_type, " yet"); return status; } +template +inline void copy_data(const T* src, T* dst, size_t count) { + memcpy(dst, src, count * sizeof(T)); +} + +template<> +inline void copy_data(const std::string* src, std::string* dst, size_t count) { + const std::string* end = src + count; + std::copy(src, end, dst); +} + template Status Split::ComputeImpl(OpKernelContext& context, const Tensor& input) const { auto& input_shape = input.Shape(); @@ -124,7 +132,7 @@ Status Split::ComputeImpl(OpKernelContext& context, const Tensor& input) const { static_cast(output_data), // B split_size * after_dims_excluding_split, // ldb [](const T* src, T* dst, size_t count) { - memcpy(dst, src, count * sizeof(T)); + copy_data(src, dst, count); }); input_offset += split_size * after_dims_excluding_split; // offset by the N data we used in this iteration diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index 964a0d73df..1957f40908 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -7,15 +7,18 @@ namespace onnxruntime { namespace test { -template using ShapeAndData = std::pair, const std::vector>; +template +using ShapeAndData = std::pair, const std::vector>; using ShapeAndFloatData = ShapeAndData; using ShapeAndInt32Data = ShapeAndData; +using ShapeAndStringData = ShapeAndData; using ExpectResult = OpTester::ExpectResult; -template void RunTest(int64_t axis, const std::vector split_sizes, const ShapeAndData& input, - const std::vector>& outputs, bool is_tensorrt_supported = true, - bool expect_failure = false, const std::string& err_msg = {}) { +template +void RunTest(int64_t axis, const std::vector split_sizes, const ShapeAndData& input, + const std::vector>& outputs, bool is_tensorrt_supported = true, + bool expect_failure = false, const std::string& err_msg = {}) { OpTester test("Split"); test.AddAttribute("axis", axis); @@ -40,16 +43,16 @@ template void RunTest(int64_t axis, const std::vector split test.Run(expect_failure ? ExpectResult::kExpectFailure : ExpectResult::kExpectSuccess, err_msg, excluded_providers); } -TEST(SplitOperatorTest, Axis0EqualSplit) { +TEST(SplitOperatorTest, Axis0EqualSplitFloat) { const int64_t axis = 0; std::vector outputs; // input shape and data ShapeAndFloatData input = {{4, 2}, // shape - {1.f, 2.f, - 3.f, 4.f, - 5.f, 6.f, - 7.f, 8.f}}; + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; outputs.push_back({{2, 2}, {1.f, 2.f, @@ -59,7 +62,7 @@ TEST(SplitOperatorTest, Axis0EqualSplit) { {5.f, 6.f, 7.f, 8.f}}); - RunTest(axis, {}, input, outputs, false);//TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM } TEST(SplitOperatorTest, Axis0EqualSplitInt32) { @@ -81,19 +84,41 @@ TEST(SplitOperatorTest, Axis0EqualSplitInt32) { {5, 6, 7, 8}}); - RunTest(axis, {}, input, outputs, false);//TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM } -TEST(SplitOperatorTest, Axis0UnequalSplit) { +TEST(SplitOperatorTest, Axis0EqualSplitString) { + const int64_t axis = 0; + std::vector outputs; + + // input shape and data + ShapeAndStringData input = {{4, 2}, // shape + {"a", "b", + "c", "d", + "e", "f", + "g", "h"}}; + + outputs.push_back({{2, 2}, + {"a", "b", + "c", "d"}}); + + outputs.push_back({{2, 2}, + {"e", "f", + "g", "h"}}); + + RunTest(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM +} + +TEST(SplitOperatorTest, Axis0UnequalSplitFloat) { const int64_t axis = 0; std::vector outputs; // input shape and data ShapeAndFloatData input = {{4, 2}, // shape - {1.f, 2.f, - 3.f, 4.f, - 5.f, 6.f, - 7.f, 8.f}}; + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; std::vector splits{1, 3}; @@ -104,17 +129,40 @@ TEST(SplitOperatorTest, Axis0UnequalSplit) { 5.f, 6.f, 7.f, 8.f}}); - RunTest(axis, splits, input, outputs, false);//TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, splits, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM } -TEST(SplitOperatorTest, Axis1EqualSplit) { +TEST(SplitOperatorTest, Axis0UnequalSplitString) { + const int64_t axis = 0; + std::vector outputs; + + // input shape and data + ShapeAndStringData input = {{4, 2}, // shape + {"a", "b", + "c", "d", + "e", "f", + "g", "h"}}; + + std::vector splits{1, 3}; + + outputs.push_back({{1, 2}, {"a", "b"}}); + + outputs.push_back({{3, 2}, + {"c", "d", + "e", "f", + "g", "h"}}); + + RunTest(axis, splits, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM +} + +TEST(SplitOperatorTest, Axis1EqualSplitFloat) { const int64_t axis = 1; std::vector outputs; // input shape and data ShapeAndFloatData input = {{2, 4}, - {1.f, 2.f, 3.f, 4.f, - 5.f, 6.f, 7.f, 8.f}}; + {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f}}; outputs.push_back({{2, 2}, {1.f, 2.f, @@ -127,14 +175,34 @@ TEST(SplitOperatorTest, Axis1EqualSplit) { RunTest(axis, {}, input, outputs, false); } -TEST(SplitOperatorTest, Axis1UnequalSplit) { +TEST(SplitOperatorTest, Axis1EqualSplitString) { + const int64_t axis = 1; + std::vector outputs; + + // input shape and data + ShapeAndStringData input = {{2, 4}, + {"a", "b", "c", "d", + "e", "f", "g", "h"}}; + + outputs.push_back({{2, 2}, + {"a", "b", + "e", "f"}}); + + outputs.push_back({{2, 2}, + {"c", "d", + "g", "h"}}); + + RunTest(axis, {}, input, outputs, false); +} + +TEST(SplitOperatorTest, Axis1UnequalSplitFloat) { const int64_t axis = 1; std::vector outputs; // input shape and data ShapeAndFloatData input = {{2, 4}, - {1.f, 2.f, 3.f, 4.f, - 5.f, 6.f, 7.f, 8.f}}; + {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f}}; std::vector splits{3, 1}; @@ -149,6 +217,28 @@ TEST(SplitOperatorTest, Axis1UnequalSplit) { RunTest(axis, splits, input, outputs, false); } +TEST(SplitOperatorTest, Axis1UnequalSplitString) { + const int64_t axis = 1; + std::vector outputs; + + // input shape and data + ShapeAndStringData input = {{2, 4}, + {"a", "b", "c", "d", + "e", "f", "g", "h"}}; + + std::vector splits{3, 1}; + + outputs.push_back({{2, 3}, + {"a", "b", "c", + "e", "f", "g"}}); + + outputs.push_back({{2, 1}, + {"d", + "h"}}); + + RunTest(axis, splits, input, outputs, false); +} + ShapeAndFloatData CreateInput(std::vector shape) { auto size = TensorShape(shape).Size(); @@ -280,8 +370,8 @@ TEST(SplitOperatorTest, NegativeAxis) { // input shape and data ShapeAndFloatData input = {{2, 4}, - {1.f, 2.f, 3.f, 4.f, - 5.f, 6.f, 7.f, 8.f}}; + {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f}}; outputs.push_back({{2, 2}, {1.f, 2.f, @@ -300,10 +390,10 @@ TEST(SplitOperatorTest, InvalidAxis) { // input shape and data ShapeAndFloatData input = {{4, 2}, // shape - {1.f, 2.f, - 3.f, 4.f, - 5.f, 6.f, - 7.f, 8.f}}; + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; outputs.push_back({{1}, {0.f}}); @@ -317,17 +407,17 @@ TEST(SplitOperatorTest, SplitAttributeSumTooSmall) { // input shape and data ShapeAndFloatData input = {{4, 2}, // shape - {1.f, 2.f, - 3.f, 4.f, - 5.f, 6.f, - 7.f, 8.f}}; + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; std::vector splits{1, 2}; // should sum to 4 outputs.push_back({{1, 2}, {1.f, 2.f}}); outputs.push_back({{2, 2}, {3.f, 4.f, 5.f, 6.f}}); - RunTest(axis, splits, input, outputs, false, true, "Cannot split using values in 'split' attribute");//TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, splits, input, outputs, false, true, "Cannot split using values in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM } TEST(SplitOperatorTest, InvalidValueInSplitAttribute) { @@ -336,16 +426,16 @@ TEST(SplitOperatorTest, InvalidValueInSplitAttribute) { // input shape and data ShapeAndFloatData input = {{4, 2}, // shape - {1.f, 2.f, - 3.f, 4.f, - 5.f, 6.f, - 7.f, 8.f}}; + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; std::vector splits{1, 0, 3}; // 0 is not valid outputs.push_back({{1, 2}, {1.f, 2.f}}); outputs.push_back({{3, 2}, {3.f, 4.f, 5.f, 6.f, 7.f, 8.f}}); - RunTest(axis, splits, input, outputs, false, true, "Invalid value in 'split' attribute");//TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, splits, input, outputs, false, true, "Invalid value in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM } /*