diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 22f6feeab6c..85eb0b1cc00 100644 --- a/backends/aoti/aoti_backend.py +++ b/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 diff --git a/exir/_serialize/_cord.py b/exir/_serialize/_cord.py index b8be3572e16..4177f9088cc 100644 --- a/exir/_serialize/_cord.py +++ b/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) diff --git a/exir/_serialize/_named_data_store.py b/exir/_serialize/_named_data_store.py index d0b18f9d6c2..aabd4e22542 100644 --- a/exir/_serialize/_named_data_store.py +++ b/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( diff --git a/exir/_serialize/data_serializer.py b/exir/_serialize/data_serializer.py index cee34506b66..5d6fa59e757 100644 --- a/exir/_serialize/data_serializer.py +++ b/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] diff --git a/exir/_serialize/test/test_cord.py b/exir/_serialize/test/test_cord.py index d6c60255f5b..85770015da8 100644 --- a/exir/_serialize/test/test_cord.py +++ b/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)) diff --git a/exir/_serialize/test/test_named_data_store.py b/exir/_serialize/test/test_named_data_store.py index 929725dda6d..4dc89002318 100644 --- a/exir/_serialize/test/test_named_data_store.py +++ b/exir/_serialize/test/test_named_data_store.py @@ -8,11 +8,14 @@ import copy import hashlib +import os +import tempfile import unittest from typing import Any, cast import torch +from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir._serialize.data_serializer import DataEntry from executorch.exir.scalar_type import ScalarType @@ -400,3 +403,34 @@ def test_fingerprint_collision_with_dedup(self) -> None: self.assertEqual(output.pte_data["key1"].buffer_index, 0) self.assertEqual(output.pte_data["key2"].buffer_index, 1) self.assertEqual(output.pte_data["key3"].buffer_index, 0) + + def test_file_backed_data_dedup(self) -> None: + with tempfile.TemporaryDirectory() as directory: + paths = [os.path.join(directory, f"data{i}") for i in range(2)] + for path in paths: + with open(path, "wb") as f: + f.write(b"file-backed data") + + with ( + FileBackedData.move_from(paths[0]) as file1, + FileBackedData.move_from(paths[1]) as file2, + ): + store1 = NamedDataStore() + store1.add_named_data("key1", file1, external_tag="model") + + store2 = NamedDataStore() + store2.add_named_data("key2", file2, external_tag="model") + output2 = store2.get_named_data_store_output() + + store1.merge_named_data_store(output2) + output1 = store1.get_named_data_store_output() + self.assertEqual(1, len(output1.buffers)) + self.assertIs(output1.buffers[0], file1) + self.assertIs(output2.buffers[0], file2) + self.assertEqual(2, len(os.listdir(directory))) + + store3 = NamedDataStore() + store3.merge_named_data_store(output2) + self.assertIs(store3.buffers[0], file2) + + self.assertEqual([], os.listdir(directory)) diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index dd8d97d66ac..716040e99e3 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -566,29 +566,35 @@ def lower_all_submodules_to_backend( """ Lower all submodules nodes given in the method_to_submodule_nodes map to backend_id. """ + backend_name_to_subclass = { + subclass.__name__: subclass for subclass in BackendDetails.__subclasses__() + } + if backend_id not in backend_name_to_subclass: + raise NotImplementedError(f"Backend {backend_id} was not found.") + backend_cls = backend_name_to_subclass[backend_id] + + method_to_compile_specs = { + method_name: [node.meta["compile_spec"] for node in call_submodule_nodes] + for method_name, call_submodule_nodes in method_to_submodules_nodes.items() + } + # The created exported program for the submodules are in the call_module node's meta data # We just map the method_to_submodule_nodes directly to the method_to_partitioned_exported_programs method_to_partitioned_program = { method_name: [ # perform deep copy here in case backends change graph inside preprocess method - copy.deepcopy(node.meta["submodule_program"]) - for node in call_submodule_nodes + backend_cls.copy_exported_program_for_preprocess( + node.meta["submodule_program"], compile_spec + ) + for node, compile_spec in zip( + call_submodule_nodes, method_to_compile_specs[method_name] + ) ] for method_name, call_submodule_nodes in method_to_submodules_nodes.items() } - method_to_compile_specs = { - method_name: [node.meta["compile_spec"] for node in call_submodule_nodes] - for method_name, call_submodule_nodes in method_to_submodules_nodes.items() - } - - backend_name_to_subclass = { - subclass.__name__: subclass for subclass in BackendDetails.__subclasses__() - } - if backend_id not in backend_name_to_subclass: - raise NotImplementedError(f"Backend {backend_id} was not found.") method_to_preprocess_result: dict[str, List[PreprocessResult]] = ( - backend_name_to_subclass[backend_id].preprocess_multimethod( + backend_cls.preprocess_multimethod( method_to_partitioned_program, method_to_compile_specs ) ) diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index 9614826c61a..67cf59e0f59 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import copy from abc import ABC, abstractmethod from dataclasses import dataclass @@ -73,6 +74,15 @@ def __init_subclass__(cls, **kwargs): f"(attempted by '{cls.__name__}')." ) + @classmethod + def copy_exported_program_for_preprocess( + cls, + edge_program: ExportedProgram, + compile_specs: List[CompileSpec], + ) -> ExportedProgram: + """Return an isolated program for backend preprocessing.""" + return copy.deepcopy(edge_program) + @staticmethod # all backends need to implement this method @enforcedmethod diff --git a/extension/flat_tensor/test/test_serialize.py b/extension/flat_tensor/test/test_serialize.py index 6ecd6911ac8..57a0e3c38fd 100644 --- a/extension/flat_tensor/test/test_serialize.py +++ b/extension/flat_tensor/test/test_serialize.py @@ -8,6 +8,8 @@ import dataclasses import math +import os +import tempfile import unittest from typing import Dict, List, Optional @@ -280,6 +282,28 @@ def test_round_trip(self) -> None: TEST_DATA_PAYLOAD.named_data, deserialized_payload.named_data ) + def test_file_backed_data_matches_bytes(self) -> None: + from executorch.exir._serialize._cord import FileBackedData + + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "weights.bin") + with open(path, "wb") as f: + f.write(TEST_BUFFER[0]) + file_data = FileBackedData.move_from(path) + payload = DataPayload( + buffers=[file_data], + named_data={"weight": DataEntry(0, 16, None)}, + ) + bytes_payload = DataPayload( + buffers=[TEST_BUFFER[0]], + named_data={"weight": DataEntry(0, 16, None)}, + ) + serializer = FlatTensorSerializer(FlatTensorConfig()) + self.assertEqual( + bytes(serializer.serialize(bytes_payload)), + bytes(serializer.serialize(payload)), + ) + def test_deserialize_to_named_data_store_output(self) -> None: store = NamedDataStore() external_tag = "model"