Speed up named data map + add uint32 support

Differential Revision: D92447071

Pull Request resolved: #17257

Author: metascroy

exir/_serialize/_named_data_store.py

@@ -7,16 +7,32 @@
 # pyre-strict
 
 import hashlib
-
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 import torch
-
 from executorch.exir._serialize.data_serializer import DataEntry
 from executorch.exir.tensor_layout import TensorLayout
 
 
+def _tensor_to_bytes(tensor: torch.Tensor) -> bytes:
+    """Convert tensor to bytes using the fastest method available.
+
+    Uses numpy().tobytes() which is faster than bytes(untyped_storage())
+    for C-contiguous tensors. Falls back to untyped_storage() for
+    non-contiguous tensors (e.g., channels_last) to preserve memory layout.
+    """
+    if not tensor.is_contiguous():
+        # For non-C-contiguous tensors (e.g., channels_last), use untyped_storage
+        # to preserve the actual memory layout
+        return bytes(tensor.untyped_storage())
+    if tensor.dtype == torch.bfloat16:
+        # BFloat16 is not supported by numpy, extract raw bytes via view
+        return tensor.view(torch.uint16).numpy().tobytes()
+    else:
+        return tensor.numpy().tobytes()
+
+
 @dataclass
 class NamedDataStoreOutput:
     """
@@ -169,7 +185,7 @@ def add_named_data(
                     f"Tensor {key} is a torch.Tensor, with tensor_layout {real_tensor_layout}. The provided tensor layout {tensor_layout} does not match."
                 )
             tensor_layout = real_tensor_layout
-            byte_data = bytes(data.untyped_storage())
+            byte_data = _tensor_to_bytes(data)
         else:
             byte_data = data
 

exir/tensor.py

@@ -289,6 +289,7 @@ def memory_format_enum(memory_format: torch.memory_format) -> int:
     torch.bfloat16: ScalarType.BFLOAT16,
     torch.quint4x2: ScalarType.QUINT4x2,
     torch.uint16: ScalarType.UINT16,
+    torch.uint32: ScalarType.UINT32,
 }