JaggedTensor permute - less CPU ops (#2786)

Summary:
Pull Request resolved: #2786

`JaggedTensor.permute` could be called with very large `indices` list (a few hundred items) - so calling python properties `self.keys()`, `self.variable_stride_per_key()` and `self.stride_per_key_per_rank()` in the loop over indices start to compound and take noticeable time **on CPU**.

Reviewed By: sarckk

Differential Revision: D70609204

fbshipit-source-id: 257a9a45b204514eef932afcf7df958b194912b6

Author: che-sh

torchrec/sparse/jagged_tensor.py

@@ -2511,22 +2511,23 @@ def permute(
         permuted_stride_per_key_per_rank: List[List[int]] = []
         permuted_length_per_key: List[int] = []
         permuted_length_per_key_sum = 0
+        keys = self._keys
+        variable_stride_per_key = self.variable_stride_per_key()
+        stride_per_key_per_rank = self.stride_per_key_per_rank()
         for index in indices:
-            key = self.keys()[index]
+            key = keys[index]
             permuted_keys.append(key)
             permuted_length_per_key.append(length_per_key[index])
-            if self.variable_stride_per_key():
-                permuted_stride_per_key_per_rank.append(
-                    self.stride_per_key_per_rank()[index]
-                )
+            if variable_stride_per_key:
+                permuted_stride_per_key_per_rank.append(stride_per_key_per_rank[index])
 
         permuted_length_per_key_sum = sum(permuted_length_per_key)
         if not torch.jit.is_scripting() and is_non_strict_exporting():
             torch._check_is_size(permuted_length_per_key_sum)
             torch._check(permuted_length_per_key_sum != -1)
             torch._check(permuted_length_per_key_sum != 0)
 
-        if self.variable_stride_per_key():
+        if variable_stride_per_key:
             length_per_key_tensor = _pin_and_move(
                 torch.tensor(self.length_per_key()), self.device()
             )
@@ -2571,7 +2572,7 @@ def permute(
                 permuted_length_per_key_sum,
             )
         stride_per_key_per_rank = (
-            permuted_stride_per_key_per_rank if self.variable_stride_per_key() else None
+            permuted_stride_per_key_per_rank if variable_stride_per_key else None
         )
         kjt = KeyedJaggedTensor(
             keys=permuted_keys,