[KT.regroup Ops][3/N] implementation of fbgemm op - kt_regroup_permutes

[TroyGarden]

Summary: X-link: https://github.com/pytorch/torchrec/pull/2128

context

• learned from previous benchmark/trace analysis, that the CPU runtime (~2.0 ms) is still comparable with the GPU runtime (~2.2 ms)	Operator	CPU runtime	GPU runtime
  native-pytorch	3.9 ms	3.1 ms
  [prod] permute_pooled_embs	2.1 ms	4.9 ms
  [new] permute_multi_embedding	2.0 ms	2.2 ms

• after a closer look, it takes ~1.1 ms in the meta arguments preparation/calculation, particularly, the _multi_remap_to_groups {F1713121552}
• in order to further improve the CPU runtime performance, we are moving the meta argument preparation into the C++ domain (inside the operator)

details

• the new operator kt_regroup_permutes directly takes input from python/pytorch domain as: 1. feature/key list: List[List[str]] 2. feature/key lengths: List[List[int]] 3. permute feature/key list: List[List[str]].
• use torch.ops.fbgemm.kt_regroup_permutes to do the same work as _kt_regroup_permutes
• in the new op keyed_tensor_regroup, first calls the generate_keyed_tensor_permutes to get the proper arguments
• then call the permute_multi_embedding op for the actual operation.

Differential Revision: D58649553

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D58649553

[netlify[bot]]

❌ Deploy Preview for pytorch-fbgemm-docs failed.

Name	Link
🔨 Latest commit	2034a8085780c2c7b8651795b224609a66d63b68
🔍 Latest deploy log	https://app.netlify.com/sites/pytorch-fbgemm-docs/deploys/6676fc07cbaded0008895cfd

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D58649553