use new op in KTRegroupAsDict module (#2210)

Summary:
Pull Request resolved: #2210

# context
* the new op `permute_multi_embedding` outperforms the original op `permute_pooled_embs_auto_grad`
* this diff makes the move to switch to the new op
* benchmark results: D58907223

# benchmark
* [traces](https://drive.google.com/drive/folders/1v_kD9n1jOkGUmYyix3-dUYiBDE_C3Hiv?usp=drive_link)
* previous prod
 {F1747994738}
* new prod
 {F1747994032}
* metrics
|Operator|GPU runtime|GPU memory|notes|
|---|---|---|---|---|
|**[previous prod] permute_pooled_embs**|4.9 ms|1.5 K|GPU-boudned, does **NOT** allow duplicates, PT2 non-compatible `pin_and_move`|
|**[new prod] permute_multi_embedding**|2.0 ms|1.0 K|both CPU and GPU runtime/memory improved, **ALLOW** duplicates, PT2 friendly|

Reviewed By: dstaay-fb

Differential Revision: D53590566

Author: TroyGarden

torchrec/modules/regroup.py

@@ -9,20 +9,15 @@
 
 #!/usr/bin/env python3
 
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
-from torchrec.sparse.jagged_tensor import (
-    _all_keys_used_once,
-    _desugar_keyed_tensors,
-    _remap_to_groups,
-    KeyedTensor,
-)
+from torchrec.sparse.jagged_tensor import _desugar_keyed_tensors, KeyedTensor
 
 
 @torch.fx.wrap
-def _concat_values(kts: List[KeyedTensor], dim: int) -> torch.Tensor:
-    return torch.cat([kt.values() for kt in kts], dim=dim)
+def _get_kts_values(kts: List[KeyedTensor]) -> List[torch.Tensor]:
+    return [kt.values() for kt in kts]
 
 
 @torch.fx.wrap
@@ -36,11 +31,15 @@ def _permuted_values(
 
 @torch.fx.wrap
 def _build_dict(
-    keys: List[str], values: torch.Tensor, splits: List[int], dim: int
+    keys: List[str],
+    values: Union[torch.Tensor, List[torch.Tensor]],
+    splits: List[int],
+    dim: int,
 ) -> Dict[str, torch.Tensor]:
-    return {
-        key: tensor for key, tensor in zip(keys, torch.split(values, splits, dim=dim))
-    }
+    if isinstance(values, torch.Tensor):
+        return dict(zip(keys, torch.split(values, splits, dim=dim)))
+    else:
+        return dict(zip(keys, values))
 
 
 class KTRegroupAsDict(torch.nn.Module):
@@ -80,23 +79,22 @@ def __init__(self, groups: List[List[str]], keys: List[str]) -> None:
         self._use_fbgemm_regroup: bool = False
         self._splits: List[int] = []
         self._idx_key_pairs: List[Tuple[int, str]] = []
-        self._permute_tensor: Optional[torch.Tensor] = None
-        self._inv_permute_tensor: Optional[torch.Tensor] = None
-        self._offsets_tensor: Optional[torch.Tensor] = None
-        self._inv_offsets_tensor: Optional[torch.Tensor] = None
+        self._permutes: Optional[torch.Tensor] = None
+        self._in_shapes: Optional[torch.Tensor] = None
+        self._out_shapes: Optional[torch.Tensor] = None
+        self._out_lengths: Optional[List[int]] = None
 
     def _init_fbgemm_regroup(self, kts: List[KeyedTensor]) -> None:
         self._use_fbgemm_regroup = True
         keys, lengths, values = _desugar_keyed_tensors(kts)
-        permute, inv_permute, offsets, inv_offsets, splits = _remap_to_groups(
-            keys, lengths, self._groups
+        self._permutes, self._in_shapes, self._out_shapes, self._out_lengths = (
+            torch.ops.fbgemm.kt_regroup_arguments(
+                values[0],
+                keys,
+                lengths,
+                self._groups,
+            )
         )
-        # no need to pin_memory() or to(..., non_blocking=True) since occurs only once
-        self._permute_tensor = permute.to(self.device)
-        self._inv_permute_tensor = inv_permute.to(self.device)
-        self._offsets_tensor = offsets.to(self.device)
-        self._inv_offsets_tensor = inv_offsets.to(self.device)
-        self._splits = splits
 
     def _init_regroup(self, kts: List[KeyedTensor]) -> None:
         lengths = [kt.length_per_key() for kt in kts]
@@ -137,24 +135,23 @@ def forward(self, keyed_tensors: List[KeyedTensor]) -> Dict[str, torch.Tensor]:
             ), "All inputs should have the same key_dim"
             self._dim = keyed_tensors[0].key_dim()
 
-            if _all_keys_used_once(keyed_tensors, self._groups) and self._dim == 1:
+            if self._dim == 1:
                 self._init_fbgemm_regroup(keyed_tensors)
             else:
                 self._init_regroup(keyed_tensors)
             self._is_inited = True
 
         if self._use_fbgemm_regroup:
-            values = _concat_values(keyed_tensors, self._dim)
-            permuted_values = torch.ops.fbgemm.permute_pooled_embs_auto_grad(
+            values = _get_kts_values(keyed_tensors)
+            permuted_values = torch.ops.fbgemm.permute_multi_embedding(
                 values,
-                self._offsets_tensor,
-                self._permute_tensor,
-                self._inv_offsets_tensor,
-                self._inv_permute_tensor,
+                self._permutes,
+                self._in_shapes,
+                self._out_shapes,
+                self._out_lengths,
             )
         else:
             permuted_values = _permuted_values(
                 keyed_tensors, self._idx_key_pairs, self._dim
             )
-
         return _build_dict(self._keys, permuted_values, self._splits, self._dim)