QEC permute order caching

torchrec/quant/embedding_modules.py

@@ -93,12 +93,6 @@ def _get_feature_length(feature: KeyedJaggedTensor) -> Tensor:
     return feature.lengths()
 
 
-@torch.fx.wrap
-def _get_kjt_keys(feature: KeyedJaggedTensor) -> List[str]:
-    # this is a fx rule to help with batching hinting jagged sequence tensor coalescing.
-    return feature.keys()
-
-
 def for_each_module_of_type_do(
     module: nn.Module,
     module_types: List[Type[torch.nn.Module]],
@@ -463,6 +457,42 @@ def __init__(
         if register_tbes:
             self.tbes: torch.nn.ModuleList = torch.nn.ModuleList(self._emb_modules)
 
+        self._has_uninitialized_kjt_permute_order: bool = True
+        self._has_features_permute: bool = True
+        self._features_order: List[int] = []
+
+    def _permute_kjt_order(
+        self, features: KeyedJaggedTensor
+    ) -> List[KeyedJaggedTensor]:
+        if self._has_uninitialized_kjt_permute_order:
+            kjt_keys = features.keys()
+            for f in self.feature_names:
+                self._features_order.append(kjt_keys.index(f))
+
+            self.register_buffer(
+                "_features_order_tensor",
+                torch.tensor(
+                    self._features_order,
+                    device=features.device(),
+                    dtype=torch.int32,
+                ),
+                persistent=False,
+            )
+
+            self._features_order = (
+                []
+                if self._features_order == list(range(len(self._features_order)))
+                else self._features_order
+            )
+
+            self._has_uninitialized_kjt_permute_order = False
+
+        if self._features_order:
+            kjt_permute = features.permute(
+                self._features_order, self._features_order_tensor
+            )
+        return kjt_permute.split(self._feature_splits)
+
     def forward(
         self,
         features: KeyedJaggedTensor,
@@ -476,10 +506,7 @@ def forward(
         """
 
         embeddings = []
-        kjt_keys = _get_kjt_keys(features)
-        kjt_permute_order = [kjt_keys.index(k) for k in self._feature_names]
-        kjt_permute = features.permute(kjt_permute_order)
-        kjts_per_key = kjt_permute.split(self._feature_splits)
+        kjts_per_key = self._permute_kjt_order(features)
 
         for i, (emb_op, _) in enumerate(
             zip(self._emb_modules, self._key_to_tables.keys())
@@ -860,6 +887,42 @@ def __init__(  # noqa C901
         if register_tbes:
             self.tbes: torch.nn.ModuleList = torch.nn.ModuleList(self._emb_modules)
 
+        self._has_uninitialized_kjt_permute_order: bool = True
+        self._has_features_permute: bool = True
+        self._features_order: List[int] = []
+
+    def _permute_kjt_order(
+        self, features: KeyedJaggedTensor
+    ) -> List[KeyedJaggedTensor]:
+        if self._has_uninitialized_kjt_permute_order:
+            kjt_keys = features.keys()
+            for f in self.feature_names:
+                self._features_order.append(kjt_keys.index(f))
+
+            self.register_buffer(
+                "_features_order_tensor",
+                torch.tensor(
+                    self._features_order,
+                    device=features.device(),
+                    dtype=torch.int32,
+                ),
+                persistent=False,
+            )
+
+            self._features_order = (
+                []
+                if self._features_order == list(range(len(self._features_order)))
+                else self._features_order
+            )
+
+            self._has_uninitialized_kjt_permute_order = False
+
+        if self._features_order:
+            kjt_permute = features.permute(
+                self._features_order, self._features_order_tensor
+            )
+        return kjt_permute.split(self._feature_splits)
+
     def forward(
         self,
         features: KeyedJaggedTensor,
@@ -873,10 +936,7 @@ def forward(
         """
 
         feature_embeddings: Dict[str, JaggedTensor] = {}
-        kjt_keys = features.keys()
-        kjt_permute_order = [kjt_keys.index(k) for k in self._feature_names]
-        kjt_permute = features.permute(kjt_permute_order)
-        kjts_per_key = kjt_permute.split(self._feature_splits)
+        kjts_per_key = self._permute_kjt_order(features)
         for i, (emb_module, key) in enumerate(
             zip(self._emb_modules, self._key_to_tables.keys())
         ):