Add torch.tensor to named_data_store

Summary: Allow adding torch.Tensor to named_data_store, and infer the TensorLayout.

Differential Revision: D85992938

Author: lucylq

exir/TARGETS

@@ -86,6 +86,7 @@ runtime.python_library(
     ],
     deps = [
         ":scalar_type",
+        ":tensor",
     ]
 )
 

exir/_serialize/_named_data_store.py

@@ -7,10 +7,11 @@
 # pyre-strict
 
 import hashlib
+
 from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
 
-# from dataclasses import dataclass
-from typing import Dict, List, Optional
+import torch
 
 from executorch.exir._serialize.data_serializer import DataEntry
 from executorch.exir.tensor_layout import TensorLayout
@@ -137,7 +138,7 @@ def _add_named_data_to_map(
     def add_named_data(
         self,
         key: str,
-        data: bytes,
+        data: Union[bytes, torch.Tensor],
         alignment: Optional[int] = 1,
         external_tag: Optional[str] = None,
         tensor_layout: Optional[TensorLayout] = None,
@@ -146,7 +147,7 @@ def add_named_data(
         Adds a named blob to the NamedDataStore.
         Args:
             key (str): key associated with the data.
-            data (bytes): Bytes being requested to be serialized.
+            data (Union[bytes, torch.Tensor]): Union of bytes, or torch.Tensor to serialize. Note: if a tensor is passed, it must have contiguous memory layout. The tensor_layout will be inferred from the tensor and should not be passed in.
             alignment (int): alignment for bytes to be serialized with.
             external (Optional[str]): the external filename that this data is saved to.
             tensor_layout (Optional[TensorLayout]): layout of the tensor, if applicable.
@@ -161,14 +162,24 @@ def add_named_data(
         if alignment <= 0:
             raise ValueError(f"Alignment must be greater than 0, received {alignment}.")
 
+        if isinstance(data, torch.Tensor):
+            if tensor_layout is not None:
+                raise ValueError(
+                    f"Tensor {key} is a torch.Tensor, and also has a tensor_layout that do not match."
+                )
+            tensor_layout = TensorLayout.from_tensor(data)
+            byte_data = bytes(data.untyped_storage())
+        else:
+            byte_data = data
+
         if external_tag is None:
             self._add_named_data_to_map(
-                key, data, alignment, self.pte_data, tensor_layout
+                key, byte_data, alignment, self.pte_data, tensor_layout
             )
         else:
             self._add_named_data_to_map(
                 key,
-                data,
+                byte_data,
                 alignment,
                 self.external_data.setdefault(external_tag, {}),
                 tensor_layout,

exir/_serialize/test/test_named_data_store.py

@@ -8,6 +8,8 @@
 
 import unittest
 
+import torch
+
 from executorch.exir._serialize._named_data_store import NamedDataStore
 from executorch.exir._serialize.data_serializer import DataEntry
 from executorch.exir.scalar_type import ScalarType
@@ -36,6 +38,32 @@ def test_add(self) -> None:
         self.assertEqual(output.external_data["file1"]["key2"], DataEntry(1, 16, None))
         self.assertEqual(output.external_data["file1"]["key3"], DataEntry(2, 16, None))
 
+    def test_add_torch_tensor(self) -> None:
+        store = NamedDataStore()
+        t0 = torch.tensor([[1, 2], [3, 4]], dtype=torch.int)
+        t1 = torch.randn(2, 3, 4, 5).contiguous(memory_format=torch.channels_last)
+
+        store.add_named_data("key0", t0, None, None)
+        store.add_named_data("key1", t1, 16, None)
+
+        output = store.get_named_data_store_output()
+        self.assertEqual(len(output.buffers), 2)
+        self.assertEqual(output.buffers[0], bytes(t0.untyped_storage()))
+        self.assertEqual(output.buffers[1], bytes(t1.untyped_storage()))
+
+        self.assertEqual(len(output.pte_data), 2)
+        self.assertEqual(
+            output.pte_data["key0"],
+            DataEntry(0, 1, TensorLayout(ScalarType.INT, [2, 2], [0, 1])),
+        )
+        self.assertEqual(
+            output.pte_data["key1"],
+            DataEntry(
+                1, 16, TensorLayout(ScalarType.FLOAT, [2, 3, 4, 5], [0, 2, 3, 1])
+            ),
+        )
+        self.assertEqual(len(output.external_data), 0)
+
     def test_add_duplicate_name_and_data(self) -> None:
         store = NamedDataStore()
         store.add_named_data("key", b"data", None, None)

exir/tensor_layout.py

@@ -9,7 +9,10 @@
 from dataclasses import dataclass
 from typing import List
 
+import torch
+
 from executorch.exir.scalar_type import ScalarType
+from executorch.exir.tensor import dim_order_from_stride, scalar_type_enum
 
 
 # Note: keep this in sync with the TensorLayout definition in
@@ -19,3 +22,18 @@ class TensorLayout:
     scalar_type: ScalarType
     sizes: List[int]
     dim_order: List[int]
+
+    @classmethod
+    def from_tensor(cls, tensor: torch.Tensor) -> "TensorLayout":
+        if not (
+            tensor.is_contiguous(memory_format=torch.contiguous_format)
+            or tensor.is_contiguous(memory_format=torch.channels_last)
+        ):
+            raise ValueError(
+                "Tensor is not contiguous. Please call .contiguous() before creating the TensorLayout."
+            )
+        return TensorLayout(
+            scalar_type=scalar_type_enum(tensor.dtype),
+            sizes=list(tensor.shape),
+            dim_order=list(dim_order_from_stride(tensor.stride())),
+        )

extension/flat_tensor/test/test_serialize.py

@@ -10,9 +10,15 @@
 import math
 import unittest
 
-from typing import List, Optional
+from typing import Dict, List, Optional
+
+import torch
 
 from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize._named_data_store import (
+    NamedDataStore,
+    NamedDataStoreOutput,
+)
 
 from executorch.exir._serialize.data_serializer import (
     DataEntry,
@@ -90,6 +96,22 @@ def check_tensor_layout(
         self.assertEqual(expected.sizes, actual.sizes)
         self.assertEqual(expected.dim_order, actual.dim_order)
 
+    def _check_named_data_entries(
+        self, reference: Dict[str, DataEntry], actual: Dict[str, DataEntry]
+    ) -> None:
+        self.assertEqual(reference.keys(), actual.keys())
+        SKIP_FIELDS = {"alignment"}  # Fields to ignore in comparison.
+        for key in reference.keys():
+            ref_entry = reference[key]
+            actual_entry = actual[key]
+            for field in dataclasses.fields(ref_entry):
+                if field.name not in SKIP_FIELDS:
+                    self.assertEqual(
+                        getattr(ref_entry, field.name),
+                        getattr(actual_entry, field.name),
+                        f"Named data record {key}.{field.name} does not match.",
+                    )
+
     def test_serialize(self) -> None:
         config = FlatTensorConfig()
         serializer: DataSerializer = FlatTensorSerializer(config)
@@ -245,19 +267,51 @@ def test_round_trip(self) -> None:
                 f"Buffer at index {i} does not match.",
             )
 
-        self.assertEqual(
-            TEST_DATA_PAYLOAD.named_data.keys(), deserialized_payload.named_data.keys()
+        self._check_named_data_entries(
+            TEST_DATA_PAYLOAD.named_data, deserialized_payload.named_data
         )
 
-        SKIP_FIELDS = {"alignment"}  # Fields to ignore in comparison.
-        for key in TEST_DATA_PAYLOAD.named_data.keys():
-            reference = TEST_DATA_PAYLOAD.named_data[key]
-            actual = deserialized_payload.named_data[key]
+    def test_deserialize_to_named_data_store_output(self) -> None:
+        store = NamedDataStore()
+        external_tag = "model"
+
+        tensor_layout = TensorLayout(ScalarType.FLOAT, [1, 2], [0, 1])
+        store.add_named_data(
+            "key0",
+            b"data0",
+            alignment=1,
+            external_tag=external_tag,
+            tensor_layout=tensor_layout,
+        )
+        store.add_named_data(
+            "key1",
+            torch.tensor([[1, 2], [3, 4]], dtype=torch.float32),
+            alignment=1,
+            external_tag=external_tag,
+        )
 
-            for field in dataclasses.fields(reference):
-                if field.name not in SKIP_FIELDS:
-                    self.assertEqual(
-                        getattr(reference, field.name),
-                        getattr(actual, field.name),
-                        f"Named data record {key}.{field.name} does not match.",
-                    )
+        output = store.get_named_data_store_output()
+        self.assertEqual(len(output.buffers), 2)
+        self.assertEqual(len(output.pte_data), 0)
+        self.assertEqual(len(output.external_data), 1)
+        self.assertEqual(len(output.external_data[external_tag]), 2)
+
+        # Serialize and deserialize.
+        config = FlatTensorConfig()
+        serializer: DataSerializer = FlatTensorSerializer(config)
+        data_payload = DataPayload(
+            buffers=output.buffers, named_data=output.external_data[external_tag]
+        )
+        serialized_data = serializer.serialize(data_payload)
+
+        output2 = serializer.deserialize_to_named_data_store_output(
+            bytes(serialized_data), external_tag
+        )
+
+        self.assertEqual(output.buffers, output2.buffers)
+        self.assertEqual(len(output.pte_data), 0)
+        self.assertEqual(len(output2.pte_data), 0)
+
+        self._check_named_data_entries(
+            output.external_data[external_tag], output2.external_data[external_tag]
+        )