vgf_neural_statistics.py

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

# pyre-unsafe

import base64
import binascii
import json
from typing import Any, Dict, List, Sequence, Union

# Functions here are used to take the raw delegate_debug_metadata bytes stored in ETDump,
# check whether those bytes contain VGF neural statistics JSON,
# validate the schema/version, decode base64-encoded binary blobs,
# and return a normal Python dictionary that tooling can consume.

SCHEMA = "executorch.vgf.neural_statistics"
LEGACY_SCHEMA_VERSION = 1
SCHEMA_VERSION = 2
SUPPORTED_SCHEMA_VERSIONS = (LEGACY_SCHEMA_VERSION, SCHEMA_VERSION)

DelegateMetadataBytes = Union[bytes, bytearray, str]


def _to_bytes(metadata: DelegateMetadataBytes) -> bytes:
    if isinstance(metadata, bytes):
        return metadata
    if isinstance(metadata, bytearray):
        return bytes(metadata)
    if isinstance(metadata, str):
        return metadata.encode("utf-8")
    raise TypeError(f"Unsupported delegate metadata type: {type(metadata)}")


def _decode_blob(blob: Dict[str, Any]) -> Dict[str, Any]:
    decoded = dict(blob)

    if not decoded.get("available", False):
        decoded.setdefault("raw_data", b"")
        return decoded

    if decoded.get("encoding") != "base64":
        raise ValueError(
            f"Unsupported VGF neural statistics blob encoding: {decoded.get('encoding')}"
        )

    encoded_data = decoded.get("data", "")
    try:
        decoded["raw_data"] = base64.b64decode(encoded_data, validate=True)
    except (binascii.Error, TypeError) as exc:
        raise ValueError("Malformed base64 data in VGF neural statistics blob") from exc

    return decoded


def _require_dict(payload: Dict[str, Any], key: str) -> Dict[str, Any]:
    value = payload.get(key)
    if not isinstance(value, dict):
        raise ValueError(f"VGF neural statistics metadata is missing {key}")
    return value


def _require_keys(value: Dict[str, Any], section: str, keys: Sequence[str]) -> None:
    missing = [key for key in keys if key not in value]
    if missing:
        raise ValueError(
            f"VGF neural statistics metadata {section} is missing: "
            + ", ".join(missing)
        )


def _validate_v2_interpretation_metadata(payload: Dict[str, Any]) -> None:
    statistics_mode = _require_dict(payload, "statistics_mode")
    _require_keys(
        statistics_mode,
        "statistics_mode",
        ("executorch_index", "vulkan_value", "vulkan_name"),
    )

    target = _require_dict(payload, "target")
    _require_keys(
        target,
        "target",
        (
            "available",
            "device_name",
            "vendor_id",
            "device_id",
            "driver_version",
        ),
    )

    counter_layout = _require_dict(payload, "counter_layout")
    _require_keys(
        counter_layout,
        "counter_layout",
        (
            "available",
            "schema_id",
            "schema_version",
            "endianness",
            "word_type",
            "words_per_task",
            "leading_block_count",
            "reason",
        ),
    )


def parse_vgf_neural_statistics_metadata(
    metadata: DelegateMetadataBytes,
) -> Dict[str, Any]:
    payload = json.loads(_to_bytes(metadata).decode("utf-8"))

    if payload.get("schema") != SCHEMA:
        raise ValueError(f"Not VGF neural statistics metadata: {payload.get('schema')}")

    schema_version = payload.get("schema_version")
    if schema_version not in SUPPORTED_SCHEMA_VERSIONS:
        raise ValueError(
            "Unsupported VGF neural statistics metadata schema version: "
            f"{schema_version}"
        )

    payload = dict(payload)
    if schema_version == SCHEMA_VERSION:
        _validate_v2_interpretation_metadata(payload)
    decoded_segments = []

    for segment in payload.get("segments", []):
        decoded_segment = dict(segment)
        for key in ("debug_database", "statistics_info", "statistics_memory"):
            blob = decoded_segment.get(key)
            if isinstance(blob, dict):
                decoded_segment[key] = _decode_blob(blob)
        decoded_segments.append(decoded_segment)

    payload["segments"] = decoded_segments
    return payload


def parse_vgf_neural_statistics_delegate_metadata(
    delegate_metadata_list: Sequence[DelegateMetadataBytes],
) -> List[Dict[str, Any]]:
    parsed: List[Dict[str, Any]] = []

    for metadata in delegate_metadata_list:
        if metadata is None or metadata == b"" or metadata == "":
            continue

        try:
            metadata_bytes = _to_bytes(metadata)
        except TypeError:
            # Not a valid delegate metadata representation.
            # Treat it as unrelated metadata from another source.
            continue

        try:
            payload = json.loads(metadata_bytes.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            # If the blob appears to be VGF neural statistics metadata but is
            # malformed, surface the error instead of silently dropping it.
            if SCHEMA.encode("utf-8") in metadata_bytes:
                raise ValueError(
                    "Malformed VGF neural statistics delegate metadata"
                ) from exc

            # Otherwise this is generic delegate metadata from another backend.
            continue

        if not isinstance(payload, dict) or payload.get("schema") != SCHEMA:
            # Inspector events can contain delegate metadata from other backends.
            # Ignore only records that are clearly not VGF neural statistics.
            continue

        # From this point onward the record claims to be VGF neural statistics.
        # Do not swallow parse errors: malformed VGF records should be visible.
        parsed.append(parse_vgf_neural_statistics_metadata(metadata_bytes))

    return parsed