_named_data_store.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

import hashlib
from dataclasses import dataclass
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


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:
    """
    Holds named data for serialization. Note: a DataEntry contains the index into
    `buffers`, the alignment and a tensor layout, if applicable.

    Attributes:
        buffers: A list of unique buffer entries.
        pte_data: Contains data that is stored inside the PTE file. A mapping from
            {key: DataEntry}.
        external_data: Contains data that is stored external to the PTE. A mapping
            from {filename: {key: DataEntry}}.
    """

    buffers: List[CordBuffer]
    pte_data: Dict[str, DataEntry]
    external_data: Dict[str, Dict[str, DataEntry]]


class NamedDataStore:
    """
    NamedDataStore manages the data that delegates want to share. Backends add
    bytes to the store under a unique key. These bytes can be retrieved at
    runtime using the same key with the NamedDataMap.

    Note:
    - Keys are unique in the data store, regardless of whether they are stored
        in the PTE or externally.
    - Multiple keys can point to the same buffer entry.
    - The same data can be added multiple times and all keys will point to one
        buffer. If a duplicate blob is added with a different alignment, the
        lcm of the current and new alignment is taken for that blob.
    """

    # List of unique blobs.
    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.
    # Map of {filename: {key: DataEntry}}.
    external_data: Dict[str, Dict[str, DataEntry]]

    # Fast fingerprint for dedup: (length, first 32 bytes) -> buffer indices.
    fingerprint_to_buffer_idx: Dict[Tuple[int, bytes], List[int]]
    # SHA-256 digest per buffer index, computed lazily on first dedup check.
    buffer_sha256: Dict[int, bytes]
    # Cache of key to buffer idx to detect duplicate key registration.
    key_to_buffer_idx: Dict[str, int]

    def __init__(self) -> None:
        """
        Initializes a new NamedDataStore.
        """
        self.buffers = []
        self.pte_data = {}
        self.external_data = {}
        self.fingerprint_to_buffer_idx = {}
        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 = self._sha256(self.buffers[buffer_idx])
            self.buffer_sha256[buffer_idx] = sha
        return sha

    def _add_named_data_to_map(
        self,
        key: str,
        data: CordBuffer,
        alignment: int,
        local_key_to_buffer_idx: Dict[str, DataEntry],
        tensor_layout: Optional[TensorLayout] = None,
    ) -> None:
        """
        Add data to a map and update the alignment. Ensure that the key-data
        pair is unique.
        - If the key exists, the data must be identical.
        - If multiple unique keys exist for the same data, those keys should
            point to the same buffer.

        Args:
            key (str): key associated with the data.
            data (bytes): Bytes being requested to be serialized.
            alignment (int): alignment for bytes to be serialized with.
            local_key_to_buffer_idx (Dict[str, int]): map to add the data to.
        Raises:
            ValueError: when the key exists in the store, and corresponding data
                is different.
        """
        # Check if the key exists.
        buffer_idx = self.key_to_buffer_idx.get(key, -1)
        if buffer_idx != -1:
            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. "
                    f"New data size: {len(data)} bytes."
                )
        else:
            # Two-level dedup: cheap fingerprint rejects non-matches fast,
            # SHA-256 confirms matches without full byte comparison.
            fingerprint = (len(data), self._prefix(data, 32))
            candidates = self.fingerprint_to_buffer_idx.get(fingerprint)
            if candidates is not None:
                new_sha = self._sha256(data)
                for candidate in candidates:
                    if new_sha == self._get_buffer_sha256(candidate):
                        buffer_idx = candidate
                        break

            if buffer_idx == -1:
                buffer_idx = len(self.buffers)
                self.buffers.append(data)
                self.fingerprint_to_buffer_idx.setdefault(fingerprint, []).append(
                    buffer_idx
                )

            local_key_to_buffer_idx[key] = DataEntry(
                buffer_index=buffer_idx,
                alignment=alignment,
                tensor_layout=tensor_layout,
            )
            self.key_to_buffer_idx[key] = buffer_idx

    def add_named_data(
        self,
        key: str,
        data: Union[bytes, FileBackedData, torch.Tensor],
        alignment: Optional[int] = 1,
        external_tag: Optional[str] = None,
        tensor_layout: Optional[TensorLayout] = None,
    ) -> None:
        """
        Adds a named blob to the NamedDataStore.
        Args:
            key (str): key associated with the data.
            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.
        Raises:
            ValueError: when the key exists in the store, and corresponding data
                is different.
        """

        # Set default alignment.
        if alignment is None:
            alignment = 1
        if alignment <= 0:
            raise ValueError(f"Alignment must be greater than 0, received {alignment}.")

        if isinstance(data, torch.Tensor):
            real_tensor_layout = TensorLayout.from_tensor(data)
            if tensor_layout is not None and not (real_tensor_layout == tensor_layout):
                raise ValueError(
                    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 = _tensor_to_bytes(data)
        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(
                key, byte_data, alignment, self.pte_data, tensor_layout
            )
        else:
            self._add_named_data_to_map(
                key,
                byte_data,
                alignment,
                self.external_data.setdefault(external_tag, {}),
                tensor_layout,
            )

    def externalize_pte_data(
        self,
        max_data_bytes: int,
        tag_prefix: str,
    ) -> None:
        # Keep this generic API defensive because callers can bypass backend
        # option parsing and serialized compile-spec validation.
        if (
            isinstance(max_data_bytes, bool)
            or not isinstance(max_data_bytes, int)
            or max_data_bytes <= 0
        ):
            raise ValueError("external data shard cap must be a positive integer")
        if not tag_prefix:
            raise ValueError("external data tag prefix must be nonempty")

        entries_by_buffer: Dict[int, Dict[str, DataEntry]] = {}
        for key, entry in self.pte_data.items():
            entries_by_buffer.setdefault(entry.buffer_index, {})[key] = entry

        shards: List[Dict[str, DataEntry]] = []
        current_shard: Dict[str, DataEntry] = {}
        current_bytes = 0
        # Prefer stable key order over size-based packing so equivalent stores
        # always produce the same shards.
        ordered_groups = sorted(
            entries_by_buffer.items(), key=lambda item: tuple(sorted(item[1]))
        )
        for buffer_index, entries in ordered_groups:
            buffer_size = len(self.buffers[buffer_index])
            if buffer_size > max_data_bytes:
                raise ValueError(
                    f"buffer {buffer_index} has {buffer_size} bytes and exceeds "
                    f"external data shard cap {max_data_bytes}"
                )
            if current_shard and current_bytes + buffer_size > max_data_bytes:
                shards.append(current_shard)
                current_shard = {}
                current_bytes = 0
            current_shard.update(entries)
            current_bytes += buffer_size
        if current_shard:
            shards.append(current_shard)

        external_data = {
            tag: dict(entries) for tag, entries in self.external_data.items()
        }
        for entries in shards:
            keys = sorted(entries)
            digest = hashlib.sha256("\0".join(keys).encode("utf-8")).hexdigest()
            tag = f"{tag_prefix}_{digest}"
            canonical_entries = {key: entries[key] for key in keys}
            existing = external_data.get(tag)
            if existing is not None and existing != canonical_entries:
                raise ValueError(f"external data tag collision for {tag}")
            external_data[tag] = canonical_entries

        self.external_data = external_data
        self.pte_data = {}

    def get_named_data_store_output(self) -> NamedDataStoreOutput:
        # Clean up empty maps inside self.external_data
        self.external_data = {k: v for k, v in self.external_data.items() if len(v) > 0}
        return NamedDataStoreOutput(self.buffers, self.pte_data, self.external_data)

    def merge_named_data_store(self, other: NamedDataStoreOutput) -> None:
        """
        Merge another NamedDataStore into this one.
        Args:
            other (NamedDataStore): the other NamedDataStore to merge.
        Raises:
            ValueError: when the key exists in both stores, and corresponding
                data is different between them.
        """
        # Merge the pte_data.
        for key, data_entry in other.pte_data.items():
            self.add_named_data(
                key,
                other.buffers[data_entry.buffer_index],
                data_entry.alignment,
                external_tag=None,
                tensor_layout=data_entry.tensor_layout,
            )

        # Merge the external_data.
        for filename, key_to_data_entry in other.external_data.items():
            for key, data_entry in key_to_data_entry.items():
                self.add_named_data(
                    key,
                    other.buffers[data_entry.buffer_index],
                    data_entry.alignment,
                    external_tag=filename,
                    tensor_layout=data_entry.tensor_layout,
                )