Sync from arm/Optimized-Runtimes via hub-sync

9422383 verified 6 days ago

6.4 kB

"""Write the fixed-topology YOLO26n model container consumed by the runtime."""

from __future__ import annotations

import struct
from dataclasses import dataclass

from .model_contract import (
    BOX_COORDINATE_CHANNELS,
    RUNTIME_WEIGHT_NAMES,
    RUNTIME_WEIGHT_SLOT_COUNT,
    SUPPORTED_RESOLUTIONS,
)
from .weight_contract import LoweredWeight

# This module owns the writer's wire-format constants; model_contract.h mirrors
# them in the C++ reader. Native validation exercises the complete writer/reader
# boundary for every supported resolution.
#
# The magic and version let the runtime reject unrelated or incompatible model
# files before interpreting offsets. Both are checked by ModelStore::Load.
_MAGIC = b"Y26MOD01"
_VERSION = 1

# XNNPACK uses 64-byte-aligned static buffers. The header occupies one aligned
# block. Directory position identifies the semantic slot, so each entry needs
# only the payload offset and byte count.
_ALIGNMENT = 64
_HEADER_BYTES = _ALIGNMENT
_ENTRY_BYTES = 16
_HEADER = struct.Struct("<8sIIIIQ")
_ENTRY = struct.Struct("<QQ")
_UINT32_MAX = (1 << 32) - 1
_UINT64_MAX = (1 << 64) - 1
_FP32_BYTES = 4
_CLASSIFICATION_CHANNEL_FLOOR = 64


@dataclass(frozen=True)
class RuntimeModelHeader:
    """Validated header fields from the runtime's serialized ``model.bin``."""

    resolution: int
    class_count: int
    slot_count: int
    payload_offset: int


@dataclass(frozen=True)
class _ModelWeight:
    slot: int
    name: str
    payload: bytes

    def __post_init__(self) -> None:
        if (
            type(self.slot) is not int
            or self.slot < 0
            or not self.name
            or not isinstance(self.payload, bytes)
            or not self.payload
        ):
            raise ValueError("invalid model weight")


def _build_model_weight_table(
    derived_weights: tuple[LoweredWeight, ...],
) -> tuple[_ModelWeight, ...]:
    """Build the runtime's stable table from its closed weight inventory."""
    named = []
    for weight in derived_weights:
        named.append((weight.binding.name, weight.payload))
    named.sort(key=lambda row: row[0])
    names = [name for name, _ in named]
    if len(names) != len(set(names)):
        raise ValueError("semantic model weight slots are not unique")
    if tuple(names) != RUNTIME_WEIGHT_NAMES:
        missing = sorted(set(RUNTIME_WEIGHT_NAMES) - set(names))
        unexpected = sorted(set(names) - set(RUNTIME_WEIGHT_NAMES))
        raise ValueError(
            "lowered weights do not match the fixed runtime inventory: "
            f"missing={missing!r}, unexpected={unexpected!r}"
        )
    weights = tuple(_ModelWeight(slot, name, payload) for slot, (name, payload) in enumerate(named))
    return weights


def _round_up_to_alignment(value: int) -> int:
    """Round a byte offset up to the next 64-byte model-file boundary."""
    return ((value + _ALIGNMENT - 1) // _ALIGNMENT) * _ALIGNMENT


def _class_dependent_weight_bytes_fit(class_count: int) -> bool:
    """Return whether the target runtime can size its largest learned tensor."""
    width = max(_CLASSIFICATION_CHANNEL_FLOOR, class_count)
    return (
        0 < class_count <= _UINT32_MAX - BOX_COORDINATE_CHANNELS
        and class_count <= _UINT64_MAX // _FP32_BYTES // width
    )


def read_runtime_model_header(payload: bytes) -> RuntimeModelHeader:
    """Read and validate the fixed header of an exported runtime ``model.bin``."""
    if len(payload) < _HEADER_BYTES:
        raise ValueError("model file is smaller than its required 64-byte header")
    (
        magic,
        version,
        resolution,
        class_count,
        slot_count,
        payload_offset,
    ) = _HEADER.unpack_from(payload)
    maximum_directory_entries = (len(payload) - _HEADER_BYTES) // _ENTRY_BYTES
    expected_payload_offset = _round_up_to_alignment(_HEADER_BYTES + _ENTRY_BYTES * slot_count)
    if (
        magic != _MAGIC
        or version != _VERSION
        or resolution not in SUPPORTED_RESOLUTIONS
        or not _class_dependent_weight_bytes_fit(class_count)
        or slot_count == 0
        or slot_count != RUNTIME_WEIGHT_SLOT_COUNT
        or slot_count > maximum_directory_entries
        or payload_offset != expected_payload_offset
        or payload[_HEADER.size : _HEADER_BYTES] != bytes(_HEADER_BYTES - _HEADER.size)
    ):
        raise ValueError("model header does not match the fixed runtime format")
    return RuntimeModelHeader(
        resolution=resolution,
        class_count=class_count,
        slot_count=slot_count,
        payload_offset=payload_offset,
    )


def _pack_model_weights(
    weights: tuple[_ModelWeight, ...],
    resolution: int,
    class_count: int,
) -> bytes:
    """Pack validated slot records into the binary model container."""
    if len(weights) != RUNTIME_WEIGHT_SLOT_COUNT or [weight.slot for weight in weights] != list(
        range(RUNTIME_WEIGHT_SLOT_COUNT)
    ):
        raise ValueError("model weight slots must be complete and contiguous")
    if len({weight.name for weight in weights}) != len(weights):
        raise ValueError("model weight names must be unique")
    if (
        resolution not in SUPPORTED_RESOLUTIONS
        or type(class_count) is not int
        or not _class_dependent_weight_bytes_fit(class_count)
    ):
        raise ValueError("invalid model dimensions")

    directory_offset = _HEADER_BYTES
    payload_offset = _round_up_to_alignment(directory_offset + _ENTRY_BYTES * len(weights))
    packed = bytearray(payload_offset)
    for weight in weights:
        while len(packed) % _ALIGNMENT:
            packed.append(0)
        offset = len(packed)
        packed.extend(weight.payload)
        _ENTRY.pack_into(
            packed,
            directory_offset + weight.slot * _ENTRY_BYTES,
            offset,
            len(weight.payload),
        )

    _HEADER.pack_into(
        packed,
        0,
        _MAGIC,
        _VERSION,
        resolution,
        class_count,
        len(weights),
        payload_offset,
    )
    return bytes(packed)


def pack_model(
    lowered_weights: tuple[LoweredWeight, ...],
    resolution: int,
    class_count: int,
) -> bytes:
    """Validate stable weight identities and serialize the runtime model file."""
    return _pack_model_weights(
        _build_model_weight_table(lowered_weights),
        resolution,
        class_count,
    )

	"""Write the fixed-topology YOLO26n model container consumed by the runtime."""

	from __future__ import annotations

	import struct
	from dataclasses import dataclass

	from .model_contract import (
	BOX_COORDINATE_CHANNELS,
	RUNTIME_WEIGHT_NAMES,
	RUNTIME_WEIGHT_SLOT_COUNT,
	SUPPORTED_RESOLUTIONS,
	)
	from .weight_contract import LoweredWeight

	# This module owns the writer's wire-format constants; model_contract.h mirrors
	# them in the C++ reader. Native validation exercises the complete writer/reader
	# boundary for every supported resolution.
	#
	# The magic and version let the runtime reject unrelated or incompatible model
	# files before interpreting offsets. Both are checked by ModelStore::Load.
	_MAGIC = b"Y26MOD01"
	_VERSION = 1

	# XNNPACK uses 64-byte-aligned static buffers. The header occupies one aligned
	# block. Directory position identifies the semantic slot, so each entry needs
	# only the payload offset and byte count.
	_ALIGNMENT = 64
	_HEADER_BYTES = _ALIGNMENT
	_ENTRY_BYTES = 16
	_HEADER = struct.Struct("<8sIIIIQ")
	_ENTRY = struct.Struct("<QQ")
	_UINT32_MAX = (1 << 32) - 1
	_UINT64_MAX = (1 << 64) - 1
	_FP32_BYTES = 4
	_CLASSIFICATION_CHANNEL_FLOOR = 64


	@dataclass(frozen=True)
	class RuntimeModelHeader:
	"""Validated header fields from the runtime's serialized ``model.bin``."""

	resolution: int
	class_count: int
	slot_count: int
	payload_offset: int


	@dataclass(frozen=True)
	class _ModelWeight:
	slot: int
	name: str
	payload: bytes

	def __post_init__(self) -> None:
	if (
	type(self.slot) is not int
	or self.slot < 0
	or not self.name
	or not isinstance(self.payload, bytes)
	or not self.payload
	):
	raise ValueError("invalid model weight")


	def _build_model_weight_table(
	derived_weights: tuple[LoweredWeight, ...],
	) -> tuple[_ModelWeight, ...]:
	"""Build the runtime's stable table from its closed weight inventory."""
	named = []
	for weight in derived_weights:
	named.append((weight.binding.name, weight.payload))
	named.sort(key=lambda row: row[0])
	names = [name for name, _ in named]
	if len(names) != len(set(names)):
	raise ValueError("semantic model weight slots are not unique")
	if tuple(names) != RUNTIME_WEIGHT_NAMES:
	missing = sorted(set(RUNTIME_WEIGHT_NAMES) - set(names))
	unexpected = sorted(set(names) - set(RUNTIME_WEIGHT_NAMES))
	raise ValueError(
	"lowered weights do not match the fixed runtime inventory: "
	f"missing={missing!r}, unexpected={unexpected!r}"
	)
	weights = tuple(_ModelWeight(slot, name, payload) for slot, (name, payload) in enumerate(named))
	return weights


	def _round_up_to_alignment(value: int) -> int:
	"""Round a byte offset up to the next 64-byte model-file boundary."""
	return ((value + _ALIGNMENT - 1) // _ALIGNMENT) * _ALIGNMENT


	def _class_dependent_weight_bytes_fit(class_count: int) -> bool:
	"""Return whether the target runtime can size its largest learned tensor."""
	width = max(_CLASSIFICATION_CHANNEL_FLOOR, class_count)
	return (
	0 < class_count <= _UINT32_MAX - BOX_COORDINATE_CHANNELS
	and class_count <= _UINT64_MAX // _FP32_BYTES // width
	)


	def read_runtime_model_header(payload: bytes) -> RuntimeModelHeader:
	"""Read and validate the fixed header of an exported runtime ``model.bin``."""
	if len(payload) < _HEADER_BYTES:
	raise ValueError("model file is smaller than its required 64-byte header")
	(
	magic,
	version,
	resolution,
	class_count,
	slot_count,
	payload_offset,
	) = _HEADER.unpack_from(payload)
	maximum_directory_entries = (len(payload) - _HEADER_BYTES) // _ENTRY_BYTES
	expected_payload_offset = _round_up_to_alignment(_HEADER_BYTES + _ENTRY_BYTES * slot_count)
	if (
	magic != _MAGIC
	or version != _VERSION
	or resolution not in SUPPORTED_RESOLUTIONS
	or not _class_dependent_weight_bytes_fit(class_count)
	or slot_count == 0
	or slot_count != RUNTIME_WEIGHT_SLOT_COUNT
	or slot_count > maximum_directory_entries
	or payload_offset != expected_payload_offset
	or payload[_HEADER.size : _HEADER_BYTES] != bytes(_HEADER_BYTES - _HEADER.size)
	):
	raise ValueError("model header does not match the fixed runtime format")
	return RuntimeModelHeader(
	resolution=resolution,
	class_count=class_count,
	slot_count=slot_count,
	payload_offset=payload_offset,
	)


	def _pack_model_weights(
	weights: tuple[_ModelWeight, ...],
	resolution: int,
	class_count: int,
	) -> bytes:
	"""Pack validated slot records into the binary model container."""
	if len(weights) != RUNTIME_WEIGHT_SLOT_COUNT or [weight.slot for weight in weights] != list(
	range(RUNTIME_WEIGHT_SLOT_COUNT)
	):
	raise ValueError("model weight slots must be complete and contiguous")
	if len({weight.name for weight in weights}) != len(weights):
	raise ValueError("model weight names must be unique")
	if (
	resolution not in SUPPORTED_RESOLUTIONS
	or type(class_count) is not int
	or not _class_dependent_weight_bytes_fit(class_count)
	):
	raise ValueError("invalid model dimensions")

	directory_offset = _HEADER_BYTES
	payload_offset = _round_up_to_alignment(directory_offset + _ENTRY_BYTES * len(weights))
	packed = bytearray(payload_offset)
	for weight in weights:
	while len(packed) % _ALIGNMENT:
	packed.append(0)
	offset = len(packed)
	packed.extend(weight.payload)
	_ENTRY.pack_into(
	packed,
	directory_offset + weight.slot * _ENTRY_BYTES,
	offset,
	len(weight.payload),
	)

	_HEADER.pack_into(
	packed,
	0,
	_MAGIC,
	_VERSION,
	resolution,
	class_count,
	len(weights),
	payload_offset,
	)
	return bytes(packed)


	def pack_model(
	lowered_weights: tuple[LoweredWeight, ...],
	resolution: int,
	class_count: int,
	) -> bytes:
	"""Validate stable weight identities and serialize the runtime model file."""
	return _pack_model_weights(
	_build_model_weight_table(lowered_weights),
	resolution,
	class_count,
	)