Handle size-0 output for split

Summary:
For size-0 output tensors, we do not allocate space for the tensor, therefore a nullptr will be passed in.

Note: This PR only handles the case where we split out a zero size on the splitting dimension. (i.e. split(tensor([8,0]), [4, 4], dim=0) will still fail.)

Differential Revision: D59128217

python/aitemplate/backend/common/split_common.py

@@ -386,7 +386,7 @@
     throw std::runtime_error("input is NULL!");
   }
   for (int i = 0; i < real_num_splits; i++) {
-    if (!outputs[i]) {
+    if (split_sizes[i] && !outputs[i]) {
       throw std::runtime_error("NULL output found at: " + std::to_string(i));
     }
   }

tests/unittest/ops/test_split.py

@@ -204,6 +204,7 @@ def test_split(self):
         self._run_split(input_shape=[2, 0, 4], split_size_or_sections=0, dim=-2)
         self._run_split(input_shape=[2, 0, 4], split_size_or_sections=2, dim=-1)
         self._run_split(input_shape=[2, 0, 7], split_size_or_sections=[2, 3, 2], dim=-1)
+        self._run_split(input_shape=[32, 8], split_size_or_sections=[8, 0, 0], dim=-1)
 
     def test_split_with_mask(self):
         self._run_split(