Refine Register fbgemm::sum_reduce_to_one (#5107)

Summary:
Pull Request resolved: #5107

X-link: https://github.com/facebookresearch/FBGEMM/pull/2112

Only support fp16 and bf16 for now.

Reviewed By: domiyy

Differential Revision: D86421539

fbshipit-source-id: ed94dd5236395f2b43d09665e8a645b9d48f1b25

Author: Joy Zhang

fbgemm_gpu/fbgemm_gpu/sparse_ops.py

@@ -1206,6 +1206,16 @@ def all_to_one_device(
     ]
 
 
+def sum_reduce_to_one(
+    input_tensors: list[Tensor],
+    target_device: torch.device,
+) -> Tensor:
+    torch._check(len(input_tensors) > 0, lambda: "reducing no tensor is undefined")
+    # All tensors should have the same shape
+    first_tensor = input_tensors[0]
+    return torch.empty_like(first_tensor, device=torch.device("meta"))
+
+
 def _setup() -> None:
     # pyre-ignore[16]
     _setup.done = getattr(_setup, "done", False)
@@ -1281,6 +1291,7 @@ def impl_autograd(op_name, fn, setup_context: Optional[Callable] = None) -> None
         impl_abstract("fbgemm::segment_sum_csr", segment_sum_csr_abstract)
         impl_abstract("fbgemm::dense_to_jagged_forward", dense_to_jagged_forward)
         impl_abstract("fbgemm::all_to_one_device", all_to_one_device)
+        impl_abstract("fbgemm::sum_reduce_to_one", sum_reduce_to_one)
         impl_abstract(
             "fbgemm::batch_index_select_dim0", batch_index_select_dim0_abstract
         )