From aabf47b107a18efae3f7e318358352eb205b7907 Mon Sep 17 00:00:00 2001
From: Yulong Wang <yulongw@microsoft.com>
Date: Tue, 7 Apr 2020 14:44:20 -0700
Subject: [PATCH] Fix Split CUDA implementation for zero sized input (#2942)

* Fix Split CUDA implementation for zero sized input

* resolve comments

* add case

* test case update: split into 2 tensors
---
 .../core/providers/cuda/tensor/split.cc       | 50 ++++++++++---------
 .../providers/cpu/tensor/split_op_test.cc     | 13 ++++-
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/onnxruntime/core/providers/cuda/tensor/split.cc b/onnxruntime/core/providers/cuda/tensor/split.cc
index 7c67efeba2..779cc62f2c 100644
--- a/onnxruntime/core/providers/cuda/tensor/split.cc
+++ b/onnxruntime/core/providers/cuda/tensor/split.cc
@@ -64,34 +64,36 @@ Status Split::ComputeInternal(OpKernelContext* ctx) const {
     }
   }
 
-  output_ptr.CopyToGpu();
+  if (input_tensor->Shape().Size() > 0) {
+    output_ptr.CopyToGpu();
 
-  CudaAsyncBuffer<int64_t> split_sizes_gpu(this, split_sizes);
-  split_sizes_gpu.CopyToGpu();
+    CudaAsyncBuffer<int64_t> split_sizes_gpu(this, split_sizes);
+    split_sizes_gpu.CopyToGpu();
 
-  std::vector<int64_t> split_sizes_range(split_sizes);
-  for (size_t i = 1; i < split_sizes_range.size(); ++i) {
-    split_sizes_range[i] += split_sizes_range[i - 1];
+    std::vector<int64_t> split_sizes_range(split_sizes);
+    for (size_t i = 1; i < split_sizes_range.size(); ++i) {
+      split_sizes_range[i] += split_sizes_range[i - 1];
+    }
+
+    CudaAsyncBuffer<int64_t> split_sizes_range_gpu(this, split_sizes_range);
+    split_sizes_range_gpu.CopyToGpu();
+
+    CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, axis_dimension_input_output_mapping);
+    axis_dimension_input_output_mapping_gpu.CopyToGpu();
+
+    size_t element_size = input_tensor->DataType()->Size();
+    ORT_RETURN_IF_ERROR(SplitImpl(element_size,
+                                  block_size_including_axis_dim,
+                                  block_size_inside_axis_dim,
+                                  split_sizes_gpu.GpuPtr(),
+                                  split_sizes_range_gpu.GpuPtr(),
+                                  axis_dimension_input_output_mapping_gpu.GpuPtr(),
+                                  num_outputs,
+                                  input_data,
+                                  output_ptr.GpuPtr(),
+                                  input_shape.Size()));
   }
 
-  CudaAsyncBuffer<int64_t> split_sizes_range_gpu(this, split_sizes_range);
-  split_sizes_range_gpu.CopyToGpu();
-
-  CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, axis_dimension_input_output_mapping);
-  axis_dimension_input_output_mapping_gpu.CopyToGpu();
-
-  size_t element_size = input_tensor->DataType()->Size();
-  ORT_RETURN_IF_ERROR(SplitImpl(element_size,
-                                block_size_including_axis_dim,
-                                block_size_inside_axis_dim,
-                                split_sizes_gpu.GpuPtr(),
-                                split_sizes_range_gpu.GpuPtr(),
-                                axis_dimension_input_output_mapping_gpu.GpuPtr(),
-                                num_outputs,
-                                input_data,
-                                output_ptr.GpuPtr(),
-                                input_shape.Size()));
-
   return Status::OK();
 }
 
diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc
index 7b1a582e2f..f9f8561144 100644
--- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc
+++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc
@@ -88,11 +88,11 @@ static void SplitTestInt() {
 }
 
 TEST(SplitOperatorTest, Axis0EqualSplitInt32) {
-   SplitTestInt<int32_t>();
+  SplitTestInt<int32_t>();
 }
 
 TEST(SplitOperatorTest, Axis0EqualSplitInt64) {
-   SplitTestInt<int64_t>();
+  SplitTestInt<int64_t>();
 }
 
 TEST(SplitOperatorTest, Axis0EqualSplitString) {
@@ -322,6 +322,15 @@ TEST(SplitOperatorTest, Axis2UnequalSplit) {
   RunTest<float>(axis, splits, input, outputs, false);
 }
 
+TEST(SplitOperatorTest, ZeroSizeInput) {
+  const int64_t axis = -1;
+  std::vector<ShapeAndFloatData> outputs{{{0, 1}, {}}, {{0, 1}, {}}};
+
+  ShapeAndFloatData input = CreateInput({0, 2});
+
+  RunTest<float>(axis, {}, input, outputs, false);
+}
+
 // test a split of a dimension that has leading and trailing dimensions
 TEST(SplitOperatorTest, Axis1SplitMiddleDimensionEqually) {
   const int64_t axis = 1;