support permute_multi_embedding_function on torch.export (#3897)

Summary:
X-link: facebookresearch/FBGEMM#988

Pull Request resolved: #3897

support fbgemm.permute_multi_embedding_function.default for LPV model
register with separate entry in kernel.yaml for graph mode lowering
added fp16 ref kernel (sitecao)

Reviewed By: StellarrZ

Differential Revision: D71821354

fbshipit-source-id: daaf882f70c102f133d1a64be96386a8acecdde9

Author: zejunh

fbgemm_gpu/fbgemm_gpu/sparse_ops.py

@@ -41,6 +41,9 @@
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_split_cpu"
     )
+    torch.ops.load_library(
+        "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_cpu"
+    )
 
 
 import torch.utils._pytree as pytree
@@ -1122,6 +1125,24 @@ def generic_histogram_binning_calibration_by_feature(
     )
 
 
+def permute_multi_embedding_function_impl_abstract(
+    pooled_embs: List[Tensor],
+    permutes: Tensor,
+    in_shapes: Tensor,
+    out_shapes: Tensor,
+    out_lengths: List[int],
+    reverse: bool = False,
+) -> List[Tensor]:
+    out_dtype = pooled_embs[0].dtype
+    bs = pooled_embs[0].shape[0]
+    torch._check(permutes.shape[1] == 6, lambda: "permutes must have 6 columns")
+
+    output = []
+    for i in range(len(out_lengths)):
+        output.append(torch.empty([bs, out_lengths[i]], dtype=out_dtype))
+    return output
+
+
 def _setup() -> None:
     # pyre-ignore[16]
     _setup.done = getattr(_setup, "done", False)
@@ -1259,6 +1280,10 @@ def impl_autograd(op_name, fn, setup_context: Optional[Callable] = None) -> None
             "fbgemm::generic_histogram_binning_calibration_by_feature",
             generic_histogram_binning_calibration_by_feature,
         )
+        impl_abstract(
+            "fbgemm::permute_multi_embedding_function",
+            permute_multi_embedding_function_impl_abstract,
+        )
         impl_abstract(
             "fbgemm::FloatToHFP8Quantized",
             float_to_hfp8_quantized,

fbgemm_gpu/src/permute_multi_embedding_ops/permute_multi_embedding_ops_cpu.cpp

@@ -344,6 +344,7 @@ std::vector<Tensor> regroup_keyed_tensor_meta(
 } // namespace fbgemm_gpu
 
 TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
+  m.set_python_module("fbgemm_gpu.sparse_ops");
   // register the forward function for internal (autograd) usage
   m.def(
       "permute_multi_embedding_function(Tensor[] pooled_embs, Tensor permutes, Tensor in_shapes, Tensor out_shapes, SymInt[] out_lengths, bool reverse=False) -> Tensor[]");