Add permute_duplicate_pooled_embeddings op for CPU (pytorch#1939)

Summary:
Pull Request resolved: pytorch#1939

This diff builds ontop of the pervious diff and adds support for permute_duplicate_pooled_embeddings for CPU.

# Background
Currently permute_pooled_embs_gpu does not support duplicates in a permutation, this poses a problem with passing the same embeddings to multiple modules. This doc proposes a solution to allow duplicate subsets in the resultant permutation.

# Details
The required implementation of permute_duplicate_pooled_embs_gpu should support a subset being repeated. This is represented by having duplicates in the permute list. This also results in the output list size being greater than the input list.

Input: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Offset_dims: [0,  2,  5,  6, 10]
Permute: [3, 0, 2, 1, 3]

Output:  [6, 7, 8, 9, 0, 1, 5, 2, 3, 4, 6, 7, 8, 9]

Reviewed By: sryap

Differential Revision: D48305145

fbshipit-source-id: a984bebb9f8974015f3d2a4f6a806d7e6d391275

Author: AGZain

fbgemm_gpu/include/fbgemm_gpu/permute_pooled_embedding_ops.h

@@ -11,6 +11,22 @@
 #include <ATen/ATen.h>
 
 namespace fbgemm_gpu {
+
+at::Tensor permute_pooled_embs_cpu_impl(
+    const at::Tensor& pooled_embs, // [B_local][Sum_T_global(D)]
+    const at::Tensor& offset_dim_list,
+    const at::Tensor& permute_list,
+    const at::Tensor& inv_offset_dim_list,
+    const at::Tensor& inv_permute_list,
+    const bool& allow_duplicates);
+
+at::Tensor permute_duplicate_pooled_embs_cpu(
+    const at::Tensor& pooled_embs, // [B_local][Sum_T_global(D)]
+    const at::Tensor& offset_dim_list,
+    const at::Tensor& permute_list,
+    const at::Tensor& inv_offset_dim_list,
+    const at::Tensor& inv_permute_list);
+
 at::Tensor permute_pooled_embs_cpu(
     const at::Tensor& pooled_embs, // [B_local][Sum_T_global(D)]
     const at::Tensor& offset_dim_list,

fbgemm_gpu/src/permute_pooled_embedding_ops/permute_pooled_embedding_ops_gpu.cpp

@@ -22,13 +22,14 @@ using Tensor = at::Tensor;
 
 namespace fbgemm_gpu {
 
-///@ingroup permute-pooled-embs-cpu
-Tensor permute_pooled_embs_cpu(
+///@ingroup permute-pooled-embs-cpu-impl
+Tensor permute_pooled_embs_cpu_impl(
     const Tensor& pooled_embs, // [B_local][Sum_T_global(D)]
     const Tensor& offset_dim_list,
     const Tensor& permute_list,
     const Tensor& inv_offset_dim_list,
-    const Tensor& inv_permute_list) {
+    const Tensor& inv_permute_list,
+    const bool& allow_duplicates) {
   TORCH_CHECK(
       offset_dim_list.scalar_type() == at::ScalarType::Long,
       "offset_dim_list needs to have long/int64 type")
@@ -37,9 +38,10 @@ Tensor permute_pooled_embs_cpu(
       "permute_list needs to have long/int64 type")
   auto permute = permute_list.data_ptr<int64_t>();
   const auto n = permute_list.numel();
+  const auto dims_size = allow_duplicates ? offset_dim_list.numel() : n;
   std::vector<int64_t> dims;
-  dims.reserve(n - 1);
-  for (const auto i : c10::irange(1, n)) {
+  dims.reserve(dims_size - 1);
+  for (const auto i : c10::irange(1, dims_size)) {
     dims.push_back(offset_dim_list[i].item<int64_t>());
   }
   auto ts = pooled_embs.tensor_split(dims, 1);
@@ -51,6 +53,38 @@ Tensor permute_pooled_embs_cpu(
   return at::cat(permuted_ts, 1);
 }
 
+///@ingroup permute-duplicate-pooled-embs-cpu
+at::Tensor permute_pooled_embs_cpu(
+    const at::Tensor& pooled_embs, // [B_local][Sum_T_global(D)]
+    const at::Tensor& offset_dim_list,
+    const at::Tensor& permute_list,
+    const at::Tensor& inv_offset_dim_list,
+    const at::Tensor& inv_permute_list) {
+  return permute_pooled_embs_cpu_impl(
+      pooled_embs,
+      offset_dim_list,
+      permute_list,
+      inv_offset_dim_list,
+      inv_permute_list,
+      false);
+}
+
+///@ingroup permute-pooled-embs-cpu
+at::Tensor permute_duplicate_pooled_embs_cpu(
+    const at::Tensor& pooled_embs, // [B_local][Sum_T_global(D)]
+    const at::Tensor& offset_dim_list,
+    const at::Tensor& permute_list,
+    const at::Tensor& inv_offset_dim_list,
+    const at::Tensor& inv_permute_list) {
+  return permute_pooled_embs_cpu_impl(
+      pooled_embs,
+      offset_dim_list,
+      permute_list,
+      inv_offset_dim_list,
+      inv_permute_list,
+      true);
+}
+
 using torch::autograd::AutogradContext;
 using torch::autograd::Variable;
 using torch::autograd::variable_list;
@@ -201,6 +235,22 @@ Tensor permute_duplicate_pooled_embs_auto_grad_gpu(
       inv_permute_list,
       true);
 }
+
+///@ingroup permute-duplicate-pooled-embs-cpu
+Tensor permute_duplicate_pooled_embs_auto_grad_cpu(
+    const Tensor& pooled_embs,
+    const Tensor& offset_dim_list,
+    const Tensor& permute_list,
+    const Tensor& inv_offset_dim_list,
+    const Tensor& inv_permute_list) {
+  return PermutePooledEmbsFunction::apply(
+      pooled_embs,
+      offset_dim_list,
+      permute_list,
+      inv_offset_dim_list,
+      inv_permute_list,
+      true);
+}
 } // namespace fbgemm_gpu
 
 TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
@@ -228,9 +278,15 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   DISPATCH_TO_CUDA(
       "permute_duplicate_pooled_embs",
       fbgemm_gpu::permute_duplicate_pooled_embs_gpu);
+  DISPATCH_TO_CPU(
+      "permute_duplicate_pooled_embs",
+      fbgemm_gpu::permute_duplicate_pooled_embs_cpu);
   m.def(
       "permute_duplicate_pooled_embs_auto_grad(Tensor pooled_embs, Tensor offset_dim_list, Tensor permute_list, Tensor inv_offset_dim_list, Tensor inv_permute_list) -> Tensor");
   DISPATCH_TO_CUDA(
       "permute_duplicate_pooled_embs_auto_grad",
       fbgemm_gpu::permute_duplicate_pooled_embs_auto_grad_gpu);
+  DISPATCH_TO_CPU(
+      "permute_duplicate_pooled_embs_auto_grad",
+      fbgemm_gpu::permute_duplicate_pooled_embs_auto_grad_cpu);
 }

fbgemm_gpu/test/permute_pooled_embedding_test.py

@@ -209,6 +209,19 @@ def test_duplicate_permutations(self) -> None:
             expected_result,
         )
 
+        input = input.to(device="cpu")
+        result = torch.ops.fbgemm.permute_duplicate_pooled_embs_auto_grad(
+            input,
+            _offset_dim_list.to(device=input.device),
+            _permute.to(device=input.device),
+            _inv_offset_dim_list.to(device=input.device),
+            _inv_permute.to(device=input.device),
+        )
+        self.assertEqual(
+            result.view(16).tolist(),
+            expected_result,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()