Add file-backed serialization support

backends/aoti/aoti_backend.py

@@ -112,6 +112,23 @@ def codesign_so(cls, so_path: str, compile_specs: List[CompileSpec]) -> None:
         """
         return
 
+    @classmethod
+    def load_weights_blob(
+        cls, blob_path: str, compile_specs: List[CompileSpec]
+    ) -> tuple[Any, str]:
+        """Load an AOTI weights blob and return its data and SHA-256 digest."""
+        with open(blob_path, "rb") as f:
+            blob_data = f.read()
+        os.remove(blob_path)
+        return blob_data, hashlib.sha256(blob_data).hexdigest()
+
+    @classmethod
+    def materialize_weights_blob(
+        cls, paths: Any, compile_specs: List[CompileSpec]
+    ) -> Any:
+        """Materialize backend-specific weight outputs into an AOTI blob."""
+        return paths
+
     @classmethod
     def move_program_to_device(
         cls,
@@ -257,6 +274,8 @@ def preprocess(
                 edge_program_module, tuple(user_input_placeholders), options=options
             )
 
+            paths = cls.materialize_weights_blob(paths, compile_specs)
+
             if len(missing_fallback_kernels) > 0:
                 formatted_kernels = "\n  - ".join(sorted(missing_fallback_kernels))
                 method_name = cls.method_name_from_compile_specs(compile_specs)
@@ -290,9 +309,7 @@ def preprocess(
         with open(so_path, "rb") as f:
             so_data = f.read()
 
-        # Read weights blob
-        with open(blob_path, "rb") as f:
-            blob_data = f.read()
+        blob_data, weights_blob_hash = cls.load_weights_blob(blob_path, compile_specs)
 
         # Create named data store
         named_data_store = NamedDataStore()
@@ -301,7 +318,7 @@ def preprocess(
         # keys (a method-name-only key collides). Runtime recovers them from
         # processed_bytes below.
         so_blob_key = hashlib.sha256(so_data).hexdigest() + "_so_blob"
-        weights_blob_key = hashlib.sha256(blob_data).hexdigest() + "_weights_blob"
+        weights_blob_key = weights_blob_hash + "_weights_blob"
 
         named_data_store.add_named_data(so_blob_key, so_data, 1, None)
         # Determine whether to save named data externally based on backend setting
@@ -314,7 +331,6 @@ def preprocess(
 
         # Clean up the generated files
         os.remove(so_path)
-        os.remove(blob_path)
 
         # Release device memory held by tensors that ``move_to_device_pass``
         # placed on the target device. Default impl is a no-op; concrete

exir/_serialize/_cord.py

@@ -4,10 +4,88 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import hashlib
 import io
+import os
+import shutil
+import tempfile
+import weakref
 from typing import List, Optional, Union
 
 
+class FileBackedData:
+    """A byte buffer that stays on disk until explicitly closed."""
+
+    _COPY_CHUNK_SIZE = 8 * 1024 * 1024
+
+    def __init__(self, path: str, cleanup: bool = False) -> None:
+        self._path = path
+        self._size = os.path.getsize(path)
+        self._sha256: Optional[bytes] = None
+        self._finalizer = (
+            weakref.finalize(self, self._remove, path) if cleanup else None
+        )
+
+    @staticmethod
+    def _remove(path: str) -> None:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+    @classmethod
+    def move_from(cls, path: str) -> "FileBackedData":
+        """Take ownership of ``path`` without loading its contents."""
+        directory = os.path.dirname(path) or "."
+        fd, owned_path = tempfile.mkstemp(
+            prefix=".executorch_", suffix=".data", dir=directory
+        )
+        os.close(fd)
+        try:
+            os.replace(path, owned_path)
+        except Exception:
+            os.remove(owned_path)
+            raise
+        return cls(owned_path, cleanup=True)
+
+    def __len__(self) -> int:
+        return self._size
+
+    def prefix(self, size: int) -> bytes:
+        with open(self._path, "rb") as f:
+            return f.read(size)
+
+    def sha256(self) -> bytes:
+        if self._sha256 is None:
+            digest = hashlib.sha256()
+            with open(self._path, "rb") as f:
+                while chunk := f.read(self._COPY_CHUNK_SIZE):
+                    digest.update(chunk)
+            self._sha256 = digest.digest()
+        return self._sha256
+
+    def to_bytes(self) -> bytes:
+        with open(self._path, "rb") as f:
+            return f.read()
+
+    def write_to_file(self, outfile: io.BufferedIOBase) -> None:
+        with open(self._path, "rb") as f:
+            shutil.copyfileobj(f, outfile, length=self._COPY_CHUNK_SIZE)
+
+    def close(self) -> None:
+        if self._finalizer is not None:
+            self._finalizer()
+
+    def __enter__(self) -> "FileBackedData":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
+
+CordBuffer = Union[bytes, FileBackedData]
+
+
 class Cord:
     """A `bytes`-like sequence of bytes, stored non-contiguously.
 
@@ -16,9 +94,9 @@ class Cord:
     `bytes` or `bytearray` object.
     """
 
-    def __init__(self, data: Optional[Union[bytes, "Cord"]] = None) -> None:
+    def __init__(self, data: Optional[Union[CordBuffer, "Cord"]] = None) -> None:
         """Initialize Cord data structure."""
-        self._buffers: List[bytes] = []
+        self._buffers: List[CordBuffer] = []
         self._byte_size: int = 0
 
         if data is not None:
@@ -30,20 +108,28 @@ def __len__(self):
 
     def __bytes__(self) -> bytes:
         """Return the contents of the Cord as a single `bytes` object."""
-        return b"".join(self._buffers)
+        return b"".join(
+            item if isinstance(item, bytes) else item.to_bytes()
+            for item in self._buffers
+        )
 
-    def append(self, data: Union[bytes, "Cord"]) -> None:
+    def append(self, data: Union[CordBuffer, "Cord"]) -> None:
         """Append a bytes or Cord to the current Cord."""
-        if isinstance(data, bytes):
+        if isinstance(data, (bytes, FileBackedData)):
             self._buffers.append(data)
             self._byte_size += len(data)
         elif isinstance(data, Cord):
             self._buffers.extend(data._buffers)
             self._byte_size += len(data)
         else:
-            raise TypeError(f"Can only append bytes or Cords, received {type(data)}")
+            raise TypeError(
+                f"Can only append bytes, FileBackedData, or Cords, received {type(data)}"
+            )
 
     def write_to_file(self, outfile: io.BufferedIOBase) -> None:
         """Write the Cord to a file."""
         for item in self._buffers:
-            outfile.write(item)
+            if isinstance(item, bytes):
+                outfile.write(item)
+            else:
+                item.write_to_file(outfile)

exir/_serialize/_named_data_store.py

@@ -11,6 +11,7 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
+from executorch.exir._serialize._cord import CordBuffer, FileBackedData
 from executorch.exir._serialize.data_serializer import DataEntry
 from executorch.exir.tensor_layout import TensorLayout
 
@@ -47,7 +48,7 @@ class NamedDataStoreOutput:
             from {filename: {key: DataEntry}}.
     """
 
-    buffers: List[bytes]
+    buffers: List[CordBuffer]
     pte_data: Dict[str, DataEntry]
     external_data: Dict[str, Dict[str, DataEntry]]
 
@@ -68,7 +69,7 @@ class NamedDataStore:
     """
 
     # List of unique blobs.
-    buffers: List[bytes]
+    buffers: List[CordBuffer]
     # Named data stored inside the PTE file. Map of {key: DataEntry}.
     pte_data: Dict[str, DataEntry]
     # Named data stored outside of the PTE file.
@@ -93,17 +94,29 @@ def __init__(self) -> None:
         self.buffer_sha256 = {}
         self.key_to_buffer_idx = {}
 
+    @staticmethod
+    def _sha256(data: CordBuffer) -> bytes:
+        if isinstance(data, FileBackedData):
+            return data.sha256()
+        return hashlib.sha256(data).digest()
+
+    @staticmethod
+    def _prefix(data: CordBuffer, size: int) -> bytes:
+        if isinstance(data, FileBackedData):
+            return data.prefix(size)
+        return data[:size]
+
     def _get_buffer_sha256(self, buffer_idx: int) -> bytes:
         sha = self.buffer_sha256.get(buffer_idx)
         if sha is None:
-            sha = hashlib.sha256(self.buffers[buffer_idx]).digest()
+            sha = self._sha256(self.buffers[buffer_idx])
             self.buffer_sha256[buffer_idx] = sha
         return sha
 
     def _add_named_data_to_map(
         self,
         key: str,
-        data: bytes,
+        data: CordBuffer,
         alignment: int,
         local_key_to_buffer_idx: Dict[str, DataEntry],
         tensor_layout: Optional[TensorLayout] = None,
@@ -127,7 +140,9 @@ def _add_named_data_to_map(
         # Check if the key exists.
         buffer_idx = self.key_to_buffer_idx.get(key, -1)
         if buffer_idx != -1:
-            if data != self.buffers[buffer_idx]:
+            if len(data) != len(self.buffers[buffer_idx]) or self._sha256(
+                data
+            ) != self._get_buffer_sha256(buffer_idx):
                 raise ValueError(
                     f"Duplicate key {key} with different data. "
                     f"Existing data size: {len(self.buffers[buffer_idx])} bytes. "
@@ -136,10 +151,10 @@ def _add_named_data_to_map(
         else:
             # Two-level dedup: cheap fingerprint rejects non-matches fast,
             # SHA-256 confirms matches without full byte comparison.
-            fingerprint = (len(data), data[:32])
+            fingerprint = (len(data), self._prefix(data, 32))
             candidates = self.fingerprint_to_buffer_idx.get(fingerprint)
             if candidates is not None:
-                new_sha = hashlib.sha256(data).digest()
+                new_sha = self._sha256(data)
                 for candidate in candidates:
                     if new_sha == self._get_buffer_sha256(candidate):
                         buffer_idx = candidate
@@ -162,7 +177,7 @@ def _add_named_data_to_map(
     def add_named_data(
         self,
         key: str,
-        data: Union[bytes, torch.Tensor],
+        data: Union[bytes, FileBackedData, torch.Tensor],
         alignment: Optional[int] = 1,
         external_tag: Optional[str] = None,
         tensor_layout: Optional[TensorLayout] = None,
@@ -171,7 +186,8 @@ def add_named_data(
         Adds a named blob to the NamedDataStore.
         Args:
             key (str): key associated with the data.
-            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.
+            data: Bytes, file-backed data, or a torch.Tensor to serialize. If a
+                tensor is passed, its layout is inferred.
             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.
@@ -194,8 +210,10 @@ def add_named_data(
                 )
             tensor_layout = real_tensor_layout
             byte_data = _tensor_to_bytes(data)
-        else:
+        elif isinstance(data, (bytes, FileBackedData)):
             byte_data = data
+        else:
+            raise TypeError(f"Unsupported named data type: {type(data)}")
 
         if external_tag is None:
             self._add_named_data_to_map(

exir/_serialize/data_serializer.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Dict, Optional, Sequence
 
-from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize._cord import Cord, CordBuffer
 from executorch.exir.tensor_layout import TensorLayout
 
 
@@ -36,7 +36,7 @@ class DataPayload:
         key_to_data: a map from unique keys to serializable data.
     """
 
-    buffers: Sequence[bytes]
+    buffers: Sequence[CordBuffer]
     named_data: Dict[str, DataEntry]
 
 

exir/_serialize/test/test_cord.py

@@ -6,9 +6,11 @@
 
 
 import io
+import os
+import tempfile
 import unittest
 
-from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize._cord import Cord, FileBackedData
 
 
 class TestCord(unittest.TestCase):
@@ -61,3 +63,22 @@ def test_cord_write_to_file(self) -> None:
         outfile = io.BytesIO()
         cord.write_to_file(outfile)
         self.assertEqual(b"HelloWorld", outfile.getvalue())
+
+    def test_file_backed_data(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            source_path = os.path.join(directory, "source.bin")
+            with open(source_path, "wb") as f:
+                f.write(b"FileBacked")
+
+            with FileBackedData.move_from(source_path) as data:
+                self.assertFalse(os.path.exists(source_path))
+                self.assertEqual(10, len(data))
+
+                cord = Cord(b"Prefix")
+                cord.append(data)
+                outfile = io.BytesIO()
+                cord.write_to_file(outfile)
+                self.assertEqual(b"PrefixFileBacked", outfile.getvalue())
+                self.assertEqual(b"PrefixFileBacked", bytes(cord))
+
+            self.assertEqual([], os.listdir(directory))