Merge branch 'main' into malfet-patch-2

pytorch · Oct 16, 2023 · 83f0f43 · 83f0f43
2 parents ce122a3 + 924f310
commit 83f0f43
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h b/fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h
@@ -45,6 +45,9 @@ at::Tensor asynchronous_inclusive_cumsum_cpu(const at::Tensor& t_in);
 ///@ingroup sparse-data-cuda
 at::Tensor asynchronous_complete_cumsum_meta(const at::Tensor& t_in);
 
+///@ingroup sparse-data-cuda
+at::Tensor asynchronous_exclusive_cumsum_meta(const at::Tensor& t_in);
+
 ///@ingroup sparse-data-cuda
 at::Tensor offsets_range_cuda(const at::Tensor& offsets, int64_t range_size);
 

diff --git a/fbgemm_gpu/src/sparse_ops/sparse_ops_meta.cpp b/fbgemm_gpu/src/sparse_ops/sparse_ops_meta.cpp
@@ -31,6 +31,10 @@ Tensor asynchronous_complete_cumsum_meta(const Tensor& t_in) {
   return output;
 }
 
+Tensor asynchronous_exclusive_cumsum_meta(const Tensor& t_in) {
+  return at::zeros_symint(t_in.sym_sizes(), t_in.options());
+}
+
 namespace {
 
 Tensor pack_segments_forward_meta(
@@ -77,10 +81,6 @@ Tensor asynchronous_inclusive_cumsum_meta(const Tensor& t_in) {
   return at::empty_symint(t_in.sym_sizes(), t_in.options());
 }
 
-Tensor asynchronous_exclusive_cumsum_meta(const Tensor& t_in) {
-  return at::empty_symint(t_in.sym_sizes(), t_in.options());
-}
-
 } // namespace
 
 } // namespace fbgemm_gpu

diff --git a/fbgemm_gpu/test/sparse_ops_test.py b/fbgemm_gpu/test/sparse_ops_test.py
@@ -610,11 +610,11 @@ def test_cumsum(self, n: int, long_index: bool) -> None:
 
         # meta tests
         mx = torch.randint(low=0, high=100, size=(n,)).type(index_dtype).to("meta")
-        # mze = torch.ops.fbgemm.asynchronous_exclusive_cumsum(mx)
+        mze = torch.ops.fbgemm.asynchronous_exclusive_cumsum(mx)
+        self.assertEqual(ze.size(), mze.size())
         # mzi = torch.ops.fbgemm.asynchronous_inclusive_cumsum(mx)
-        mzc = torch.ops.fbgemm.asynchronous_complete_cumsum(mx)
-        # self.assertEqual(ze.size(), mze.size())
         # self.assertEqual(zi.size(), mzi.size())
+        mzc = torch.ops.fbgemm.asynchronous_complete_cumsum(mx)
         self.assertEqual(zc.size(), mzc.size())
 
         if gpu_available: