Add dtype to kt_regroup input (#3250)

Summary:
Pull Request resolved: #3250

Currently ir_kt_regroup always returns a float32, which is incorrect if the KTRegroupAsDict's emb_type is set

Reviewed By: malaybag

Differential Revision: D79422720

fbshipit-source-id: 2e711a7dc186b80226b2631ba94af741e05df4b7

Author: angelayi

torchrec/ir/serializer.py

@@ -23,7 +23,12 @@
 
 from torchrec.ir.types import SerializerInterface
 from torchrec.ir.utils import logging, qualname
-from torchrec.modules.embedding_configs import DataType, EmbeddingBagConfig, PoolingType
+from torchrec.modules.embedding_configs import (
+    data_type_to_dtype,
+    DataType,
+    EmbeddingBagConfig,
+    PoolingType,
+)
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import (
     FeatureProcessor,
@@ -138,7 +143,17 @@ def kt_regroup_meta_forward(
     for i, group in enumerate(op_module._groups):
         out_lengths[i] = sum(lengths_dict[key] for key in group)
     arg_list = [kt.values() for kt in keyed_tensors]
-    outputs = torch.ops.torchrec.ir_kt_regroup(arg_list, batch_size, out_lengths)
+    dtype = (
+        data_type_to_dtype(op_module._emb_dtype)
+        if op_module._emb_dtype
+        else torch.float32
+    )
+    outputs = torch.ops.torchrec.ir_kt_regroup(
+        arg_list,
+        batch_size,
+        out_lengths,
+        dtype=dtype,
+    )
     return dict(zip(op_module._keys, outputs))
 
 

torchrec/ir/tests/test_serializer.py

@@ -23,7 +23,7 @@
     mark_dynamic_kjt,
     qualname,
 )
-from torchrec.modules.embedding_configs import EmbeddingBagConfig
+from torchrec.modules.embedding_configs import data_type_to_dtype, EmbeddingBagConfig
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import (
     PositionWeightedModule,
@@ -822,6 +822,14 @@ def forward(
             preserve_module_call_signature=(tuple(sparse_fqns)),
         )
 
+        for node in ep.graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target == torch.ops.torchrec.ir_kt_regroup.default
+            ):
+                for meta in node.meta["val"]:
+                    self.assertEqual(meta.dtype, data_type_to_dtype(data_type))
+
         unflatten_ep = torch.export.unflatten(ep)
         deserialized = decapsulate_ir_modules(unflatten_ep, JsonSerializer)
         self.assertEqual(deserialized.regroup._emb_dtype, data_type)  # pyre-ignore[16]

torchrec/ir/utils.py

@@ -71,20 +71,26 @@ def ir_emb_lookup_fake(
 
 @torch.library.custom_op("torchrec::ir_kt_regroup", mutates_args={})
 def ir_kt_regroup_impl(
-    tensors: List[Optional[torch.Tensor]], batch_size: int, dims: List[int]
+    tensors: List[Optional[torch.Tensor]],
+    batch_size: int,
+    dims: List[int],
+    dtype: torch.dtype = torch.float32,
 ) -> List[torch.Tensor]:
     device = get_device(tensors)
     logger.info(f"torch.ops.torchrec.ir_kt_regroup -> ({batch_size}, {dims}) {device}")
-    return [torch.empty(batch_size, dim, device=device) for dim in dims]
+    return [torch.empty(batch_size, dim, device=device, dtype=dtype) for dim in dims]
 
 
 @torch.library.register_fake("torchrec::ir_kt_regroup")
 def ir_kt_regroup_fake(
-    tensors: List[Optional[torch.Tensor]], batch_size: int, dims: List[int]
+    tensors: List[Optional[torch.Tensor]],
+    batch_size: int,
+    dims: List[int],
+    dtype: torch.dtype = torch.float32,
 ) -> List[torch.Tensor]:
     device = get_device(tensors)
     logger.info(f"ir_kt_regroup_fake -> ({batch_size}, {dims}) {device}")
-    return [torch.empty(batch_size, dim, device=device) for dim in dims]
+    return [torch.empty(batch_size, dim, device=device, dtype=dtype) for dim in dims]
 
 
 @torch.library.custom_op("torchrec::ir_dynamic_batch_emb_lookup", mutates_args={})