shortcut for merge_pooled_embedding (#5147)

Summary:
Pull Request resolved: #5147

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

att. When all the input embedding are from the same device, we can just use cat as a short cut. This can avoid unnecessary cross device sync with current impl.

Reviewed By: yyetim

Differential Revision: D87306514

fbshipit-source-id: 71298220bf12b0fba384ce76146824b2bb094e2c

Author: garroud

fbgemm_gpu/src/merge_pooled_embedding_ops/merge_pooled_embedding_ops_gpu.cpp

@@ -688,6 +688,20 @@ Tensor merge_pooled_embeddings(
   at::cuda::OptionalCUDAGuard g;
 
   at::Device out_device = target_device;
+
+  // if target_device is the same as input devices, we can directly call
+  // cat
+  bool is_same_device = true;
+  for (const auto& t : pooled_embeddings) {
+    if (t.device() != target_device) {
+      is_same_device = false;
+      break;
+    }
+  }
+  if (is_same_device) {
+    return at::cat(pooled_embeddings, cat_dim);
+  }
+
   if (target_device.is_cuda()) {
     init_p2p_access();
     g.set_device(target_device);

fbgemm_gpu/test/merge_pooled_embeddings_test.py

@@ -68,6 +68,7 @@ class MergePooledEmbeddingsTest(unittest.TestCase):
         non_default_stream=st.booleans(),
         r=st.randoms(use_true_random=False),
         dim=st.integers(min_value=0, max_value=1),
+        source_from_same_device=st.booleans(),
     )
     # Can instantiate 8 contexts which takes a long time.
     @settings(verbosity=Verbosity.verbose, max_examples=40, deadline=None)
@@ -81,14 +82,19 @@ def test_merge(
         # pyre-fixme[2]: Parameter must be annotated.
         r,
         dim: int,
+        source_from_same_device: bool,
     ) -> None:
         dst_device = r.randint(0, num_gpus - 1)
         torch.cuda.set_device(dst_device)
         ad_ds = [embedding_dimension * ads_tables for _ in range(num_gpus)]
         batch_indices = torch.zeros(num_ads).long().cuda()
         pooled_ad_embeddings = [
-            torch.randn(
-                num_ads, ad_d, dtype=torch.float16, device=torch.device(f"cuda:{i}")
+            (
+                torch.randn(num_ads, ad_d, dtype=torch.float16, device=dst_device)
+                if source_from_same_device
+                else torch.randn(
+                    num_ads, ad_d, dtype=torch.float16, device=torch.device(f"cuda:{i}")
+                )
             )
             for i, ad_d in enumerate(ad_ds)
         ]