Add early exit in sparse_async_cumsum ops (#2202)

Summary:
Pull Request resolved: #2202

This diff adds an early exit in sparse_async_cumsum ops. When the size(s) of input tensor are zeor(s) ops return zero tensor.

Reviewed By: jspark1105, sryap, jasonjk-park

Differential Revision: D51999938

fbshipit-source-id: 5c2151a01f547666a1e1742f334953c01dbf6630

fbgemm_gpu/src/sparse_ops/sparse_async_cumsum.cu

@@ -13,6 +13,9 @@ using Tensor = at::Tensor;
 namespace fbgemm_gpu {
 
 DLL_PUBLIC Tensor asynchronous_inclusive_cumsum_gpu(const Tensor& t_in) {
+  if (t_in.numel() == 0) {
+    return at::zeros_like(t_in);
+  }
   TENSOR_ON_CUDA_GPU(t_in);
 
   at::cuda::OptionalCUDAGuard device_guard;
@@ -54,6 +57,9 @@ DLL_PUBLIC Tensor asynchronous_inclusive_cumsum_gpu(const Tensor& t_in) {
 }
 
 DLL_PUBLIC Tensor asynchronous_exclusive_cumsum_gpu(const Tensor& t_in) {
+  if (t_in.numel() == 0) {
+    return at::zeros_like(t_in);
+  }
   TENSOR_ON_CUDA_GPU(t_in);
 
   at::cuda::OptionalCUDAGuard device_guard;
@@ -95,6 +101,9 @@ DLL_PUBLIC Tensor asynchronous_exclusive_cumsum_gpu(const Tensor& t_in) {
 }
 
 DLL_PUBLIC Tensor asynchronous_complete_cumsum_gpu(const Tensor& t_in) {
+  if (t_in.numel() == 0) {
+    return at::zeros({t_in.numel() + 1}, t_in.options());
+  }
   TENSOR_ON_CUDA_GPU(t_in);
 
   at::cuda::OptionalCUDAGuard device_guard;

fbgemm_gpu/test/sparse_ops_test.py

@@ -591,7 +591,7 @@ def test_block_bucketize_sparse_features_long_indices(
             torch.testing.assert_close(new_indices_gpu.cpu(), new_indices_cpu)
 
     @given(
-        n=st.integers(min_value=1, max_value=100),
+        n=st.integers(min_value=0, max_value=100),
         long_index=st.booleans(),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=20, deadline=None)
@@ -655,8 +655,8 @@ def test_cumsum(self, n: int, long_index: bool) -> None:
             )
 
     @given(
-        n=st.integers(min_value=1, max_value=600),
-        b=st.integers(min_value=1, max_value=10),
+        n=st.integers(min_value=0, max_value=600),
+        b=st.integers(min_value=0, max_value=10),
         long_index=st.booleans(),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=20, deadline=None)