mirror of
https://github.com/retroSoC/retroSoC.git
synced 2026-11-04 22:32:36 +00:00
maksyuki
15f4dd555c
- Implemented the KWS engine in the RTL and added necessary components including KWS SRAM client and model loader. - Enhanced the APU ring scheduler testbench to validate KWS descriptors. - Introduced new Python tests for KWS functionality, including model import validation, inference checks, and error handling for mutated models. - Updated existing tests to include KWS components and ensure proper integration within the APU architecture. - Refactored code for clarity and consistency, including adjustments to signal handling and memory management in test cases.
588 lines
21 KiB
Python
588 lines
21 KiB
Python
"""Frozen APUM 1.0 KWS container definitions and independent validators."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import struct
|
|
from dataclasses import dataclass
|
|
from typing import Final
|
|
|
|
APUM_MAGIC: Final = 0x4D555041
|
|
APUM_ABI: Final = 0x00010000
|
|
APUM_BYTES: Final = 32768
|
|
APUM_HEADER_BYTES: Final = 64
|
|
APUM_OPERATOR_OFFSET: Final = 0x0040
|
|
APUM_OPERATOR_COUNT: Final = 12
|
|
APUM_OPERATOR_BYTES: Final = 64
|
|
APUM_TENSOR_OFFSET: Final = 0x0340
|
|
APUM_TENSOR_COUNT: Final = 13
|
|
APUM_TENSOR_BYTES: Final = 32
|
|
APUM_PARAMETER_OFFSET: Final = 0x0500
|
|
APUM_PARAMETER_BYTES: Final = 29072
|
|
APUM_PARAMETER_END: Final = APUM_PARAMETER_OFFSET + APUM_PARAMETER_BYTES
|
|
APUM_PAYLOAD_CRC: Final = 0xB9034B22
|
|
APUM_SHA256: Final = "87861b5f866182969e34917b4bf57cab61cd2245c3ab9b933797209322e3ce63"
|
|
APUM_PARAMS_SHA256: Final = "0e4eb210d247a45747049ef9602bbd7e90868e9cb1c04fb2a02f128db44e241e"
|
|
APUM_MULTIPLIER_SHIFT_SHA256: Final = (
|
|
"222e05821695f3fdd47be682b2cc75e0e735aa106436bf941cdda3ef8cb03b59"
|
|
)
|
|
TFLITE_SHA256: Final = "aeea436800704fce17b17292e4412630ad856e9d777c044c64ef748a880bd0ae"
|
|
CORPUS_SHA256: Final = "57fa600948ef7a19ed5eb0d0e73c7d09fef7d0dbb12bc8833a0a8caad1d6845e"
|
|
|
|
_TENSOR_PROFILE: Final = (
|
|
(1, 49, 10, 1, 0x4000, 490, 0x3F15AF17, 83),
|
|
(1, 25, 5, 64, 0x0000, 8000, 0x3DA13AC8, -128),
|
|
(1, 25, 5, 64, 0x2000, 8000, 0x3DA99AEA, -128),
|
|
(1, 25, 5, 64, 0x0000, 8000, 0x3D75DD81, -128),
|
|
(1, 25, 5, 64, 0x2000, 8000, 0x3D808854, -128),
|
|
(1, 25, 5, 64, 0x0000, 8000, 0x3D19329E, -128),
|
|
(1, 25, 5, 64, 0x2000, 8000, 0x3D3C02C1, -128),
|
|
(1, 25, 5, 64, 0x0000, 8000, 0x3D070C5A, -128),
|
|
(1, 25, 5, 64, 0x2000, 8000, 0x3D3FDF8A, -128),
|
|
(1, 25, 5, 64, 0x0000, 8000, 0x3DA452DB, -128),
|
|
(1, 1, 1, 64, 0x2000, 64, 0x3DA452DB, -128),
|
|
(1, 1, 1, 12, 0x0000, 12, 0x3E142A46, 14),
|
|
(1, 1, 1, 12, 0x2000, 12, 0x3B800000, -128),
|
|
)
|
|
_OPCODES: Final = (1, 2, 1, 2, 1, 2, 1, 2, 1, 3, 4, 5)
|
|
_WEIGHT_BYTES: Final = (2560, 576, 4096, 576, 4096, 576, 4096, 576, 4096, 0, 768, 0)
|
|
_OUTPUT_CHANNELS: Final = (64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 12, 12)
|
|
_KERNELS: Final = (
|
|
(10, 4, 2, 2),
|
|
(3, 3, 1, 1),
|
|
(1, 1, 1, 1),
|
|
(3, 3, 1, 1),
|
|
(1, 1, 1, 1),
|
|
(3, 3, 1, 1),
|
|
(1, 1, 1, 1),
|
|
(3, 3, 1, 1),
|
|
(1, 1, 1, 1),
|
|
(25, 5, 25, 5),
|
|
(1, 1, 1, 1),
|
|
(1, 1, 1, 1),
|
|
)
|
|
_PADDING: Final = (
|
|
(4, 5, 1, 1),
|
|
(1, 1, 1, 1),
|
|
(0, 0, 0, 0),
|
|
(1, 1, 1, 1),
|
|
(0, 0, 0, 0),
|
|
(1, 1, 1, 1),
|
|
(0, 0, 0, 0),
|
|
(1, 1, 1, 1),
|
|
(0, 0, 0, 0),
|
|
(0, 0, 0, 0),
|
|
(0, 0, 0, 0),
|
|
(0, 0, 0, 0),
|
|
)
|
|
|
|
|
|
def crc32_iso_hdlc(data: bytes) -> int:
|
|
"""Calculate CRC-32/ISO-HDLC with the exact APUM byte ordering."""
|
|
crc = 0xFFFFFFFF
|
|
for value in data:
|
|
crc ^= value
|
|
for _ in range(8):
|
|
crc = (crc >> 1) ^ (0xEDB88320 if (crc & 1) != 0 else 0)
|
|
return crc ^ 0xFFFFFFFF
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ApumHeader:
|
|
magic: int
|
|
abi: int
|
|
total_bytes: int
|
|
operator_set: int
|
|
mfcc_rows: int
|
|
mfcc_coefficients: int
|
|
classes: int
|
|
operators: int
|
|
parameter_bytes: int
|
|
scratch_bytes: int
|
|
quantization_profile: int
|
|
crc: int
|
|
model_id: int
|
|
threshold: int
|
|
debounce: int
|
|
tensor_count: int
|
|
layout_id: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ApumOperator:
|
|
opcode: int
|
|
flags: int
|
|
input_tensor: int
|
|
output_tensor: int
|
|
weight_offset: int
|
|
bias_offset: int
|
|
multiplier_offset: int
|
|
shift_offset: int
|
|
weight_bytes: int
|
|
output_channels: int
|
|
kernel_height: int
|
|
kernel_width: int
|
|
stride_height: int
|
|
stride_width: int
|
|
padding: tuple[int, int, int, int]
|
|
input_zero: int
|
|
output_zero: int
|
|
clamp_min: int
|
|
clamp_max: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ApumTensor:
|
|
shape: tuple[int, int, int, int]
|
|
scratch_offset: int
|
|
byte_count: int
|
|
element_type: int
|
|
scale_bits: int
|
|
zero_point: int
|
|
|
|
|
|
def _invalid(message: str) -> ValueError:
|
|
return ValueError(f"invalid APUM 1.0: {message}")
|
|
|
|
|
|
def _interval(offset: int, size: int, label: str) -> tuple[int, int, str]:
|
|
if (offset & 15) != 0 or size <= 0 or offset < APUM_PARAMETER_OFFSET:
|
|
raise _invalid(f"{label} range/alignment")
|
|
end = offset + size
|
|
if end > APUM_PARAMETER_END:
|
|
raise _invalid(f"{label} outside parameter region")
|
|
return offset, end, label
|
|
|
|
|
|
def _parse_header(image: bytes) -> ApumHeader:
|
|
words = struct.unpack_from("<10I", image, 0)
|
|
magic, abi, total, op_set, mfcc, classes, params, scratch, quant, crc = words
|
|
model_id = struct.unpack_from("<Q", image, 0x28)[0]
|
|
defaults, tensor_count, layout_id, reserved = struct.unpack_from("<4I", image, 0x30)
|
|
if reserved != 0:
|
|
raise _invalid("header reserved word is nonzero")
|
|
return ApumHeader(
|
|
magic=magic,
|
|
abi=abi,
|
|
total_bytes=total,
|
|
operator_set=op_set,
|
|
mfcc_rows=mfcc & 0xFFFF,
|
|
mfcc_coefficients=mfcc >> 16,
|
|
classes=classes & 0xFFFF,
|
|
operators=classes >> 16,
|
|
parameter_bytes=params,
|
|
scratch_bytes=scratch,
|
|
quantization_profile=quant,
|
|
crc=crc,
|
|
model_id=model_id,
|
|
threshold=defaults & 0xFF,
|
|
debounce=(defaults >> 8) & 0xFF,
|
|
tensor_count=tensor_count,
|
|
layout_id=layout_id,
|
|
)
|
|
|
|
|
|
def _parse_operator(image: bytes, index: int) -> ApumOperator:
|
|
offset = APUM_OPERATOR_OFFSET + index * APUM_OPERATOR_BYTES
|
|
fields = struct.unpack_from("<HHHHIIIIIIHHHH4BiiiiI", image, offset)
|
|
return ApumOperator(
|
|
opcode=fields[0],
|
|
flags=fields[1],
|
|
input_tensor=fields[2],
|
|
output_tensor=fields[3],
|
|
weight_offset=fields[4],
|
|
bias_offset=fields[5],
|
|
multiplier_offset=fields[6],
|
|
shift_offset=fields[7],
|
|
weight_bytes=fields[8],
|
|
output_channels=fields[9],
|
|
kernel_height=fields[10],
|
|
kernel_width=fields[11],
|
|
stride_height=fields[12],
|
|
stride_width=fields[13],
|
|
padding=(fields[14], fields[15], fields[16], fields[17]),
|
|
input_zero=fields[18],
|
|
output_zero=fields[19],
|
|
clamp_min=fields[20],
|
|
clamp_max=fields[21],
|
|
)
|
|
|
|
|
|
def _parse_tensor(image: bytes, index: int) -> ApumTensor:
|
|
offset = APUM_TENSOR_OFFSET + index * APUM_TENSOR_BYTES
|
|
n, height, width, channels, scratch, count, element, scale, zero, reserved = struct.unpack_from(
|
|
"<HHHHIIIIiI", image, offset
|
|
)
|
|
if reserved != 0:
|
|
raise _invalid(f"tensor {index} reserved word is nonzero")
|
|
return ApumTensor(
|
|
shape=(n, height, width, channels),
|
|
scratch_offset=scratch,
|
|
byte_count=count,
|
|
element_type=element,
|
|
scale_bits=scale,
|
|
zero_point=zero,
|
|
)
|
|
|
|
|
|
def _validate_header(header: ApumHeader, image: bytes) -> None:
|
|
expected = (
|
|
header.magic == APUM_MAGIC,
|
|
header.abi == APUM_ABI,
|
|
header.total_bytes == APUM_BYTES,
|
|
header.operator_set == 1,
|
|
header.mfcc_rows == 49,
|
|
header.mfcc_coefficients == 10,
|
|
header.classes == 12,
|
|
header.operators == APUM_OPERATOR_COUNT,
|
|
header.parameter_bytes == APUM_PARAMETER_BYTES,
|
|
header.scratch_bytes == APUM_BYTES,
|
|
header.quantization_profile == 1,
|
|
header.model_id == 0xCE4F70006843EAAE,
|
|
header.threshold == 128,
|
|
header.debounce == 3,
|
|
header.tensor_count == APUM_TENSOR_COUNT,
|
|
header.layout_id == 1,
|
|
)
|
|
if not all(expected):
|
|
raise _invalid("header profile")
|
|
if image[0x32:0x34] != b"\x00\x00":
|
|
raise _invalid("header defaults reserved bits")
|
|
if crc32_iso_hdlc(image[APUM_HEADER_BYTES:]) != header.crc:
|
|
raise _invalid("payload CRC")
|
|
if header.crc != APUM_PAYLOAD_CRC:
|
|
raise _invalid("non-frozen payload CRC")
|
|
|
|
|
|
def _validate_tensors(image: bytes) -> list[ApumTensor]:
|
|
tensors = [_parse_tensor(image, index) for index in range(APUM_TENSOR_COUNT)]
|
|
for index, tensor in enumerate(tensors):
|
|
n, height, width, channels, scratch, count, scale, zero = _TENSOR_PROFILE[index]
|
|
if (
|
|
tensor.shape != (n, height, width, channels)
|
|
or tensor.scratch_offset != scratch
|
|
or tensor.byte_count != count
|
|
or tensor.element_type != 1
|
|
or tensor.scale_bits != scale
|
|
or tensor.zero_point != zero
|
|
):
|
|
raise _invalid(f"tensor {index} fixed profile")
|
|
if tensor.byte_count != n * height * width * channels:
|
|
raise _invalid(f"tensor {index} byte count")
|
|
if tensor.scratch_offset + tensor.byte_count > APUM_BYTES:
|
|
raise _invalid(f"tensor {index} scratch range")
|
|
return tensors
|
|
|
|
|
|
def _validate_operators(image: bytes, tensors: list[ApumTensor]) -> list[ApumOperator]:
|
|
operators = [_parse_operator(image, index) for index in range(APUM_OPERATOR_COUNT)]
|
|
intervals: list[tuple[int, int, str]] = []
|
|
for index, operator in enumerate(operators):
|
|
if (
|
|
operator.opcode != _OPCODES[index]
|
|
or operator.flags != (1 if index < 9 else 0)
|
|
or operator.input_tensor != index
|
|
or operator.output_tensor != index + 1
|
|
or operator.weight_bytes != _WEIGHT_BYTES[index]
|
|
or operator.output_channels != _OUTPUT_CHANNELS[index]
|
|
or (
|
|
operator.kernel_height,
|
|
operator.kernel_width,
|
|
operator.stride_height,
|
|
operator.stride_width,
|
|
)
|
|
!= _KERNELS[index]
|
|
or operator.padding != _PADDING[index]
|
|
or operator.input_zero != tensors[index].zero_point
|
|
or operator.output_zero != tensors[index + 1].zero_point
|
|
or operator.clamp_min != -128
|
|
or operator.clamp_max != 127
|
|
):
|
|
raise _invalid(f"operator {index} fixed profile")
|
|
if operator.opcode in (3, 5):
|
|
if any(
|
|
value != 0
|
|
for value in (
|
|
operator.weight_offset,
|
|
operator.bias_offset,
|
|
operator.multiplier_offset,
|
|
operator.shift_offset,
|
|
operator.weight_bytes,
|
|
)
|
|
):
|
|
raise _invalid(f"operator {index} parameter fields")
|
|
continue
|
|
if index == 10 and operator.flags != 0:
|
|
raise _invalid("fully connected ReLU flag")
|
|
intervals.extend(
|
|
(
|
|
_interval(
|
|
operator.weight_offset, operator.weight_bytes, f"operator {index} weights"
|
|
),
|
|
_interval(
|
|
operator.bias_offset, operator.output_channels * 4, f"operator {index} biases"
|
|
),
|
|
_interval(
|
|
operator.multiplier_offset,
|
|
operator.output_channels * 4,
|
|
f"operator {index} multipliers",
|
|
),
|
|
_interval(
|
|
operator.shift_offset, operator.output_channels * 4, f"operator {index} shifts"
|
|
),
|
|
)
|
|
)
|
|
intervals.sort()
|
|
if intervals[0][0] != APUM_PARAMETER_OFFSET or intervals[-1][1] != APUM_PARAMETER_END:
|
|
raise _invalid("parameter coverage")
|
|
for previous, current in zip(intervals, intervals[1:]):
|
|
if previous[1] != current[0]:
|
|
raise _invalid("parameter overlap or gap")
|
|
return operators
|
|
|
|
|
|
def parse_apum(image: bytes, *, require_golden_image: bool = True) -> ApumHeader:
|
|
"""Validate the complete immutable APUM 1.0 profile and return its header."""
|
|
if len(image) != APUM_BYTES:
|
|
raise _invalid("image size")
|
|
header = _parse_header(image)
|
|
_validate_header(header, image)
|
|
if image[0x4E0:APUM_PARAMETER_OFFSET] != bytes(APUM_PARAMETER_OFFSET - 0x4E0):
|
|
raise _invalid("metadata padding")
|
|
if image[APUM_PARAMETER_END:] != bytes(APUM_BYTES - APUM_PARAMETER_END):
|
|
raise _invalid("tail padding")
|
|
tensors = _validate_tensors(image)
|
|
_validate_operators(image, tensors)
|
|
if (
|
|
hashlib.sha256(image[APUM_PARAMETER_OFFSET:APUM_PARAMETER_END]).hexdigest()
|
|
!= APUM_PARAMS_SHA256
|
|
):
|
|
raise _invalid("parameter SHA-256")
|
|
if require_golden_image and image_sha256(image) != APUM_SHA256:
|
|
raise _invalid("complete image SHA-256")
|
|
return header
|
|
|
|
|
|
def apum_layout(image: bytes) -> dict[str, object]:
|
|
"""Return the reviewable fixed APUM layout after strict validation."""
|
|
header = parse_apum(image)
|
|
tensors = _validate_tensors(image)
|
|
operators = _validate_operators(image, tensors)
|
|
return {
|
|
"abi": header.abi,
|
|
"image_bytes": len(image),
|
|
"payload_crc": header.crc,
|
|
"image_sha256": image_sha256(image),
|
|
"parameter_sha256": hashlib.sha256(
|
|
image[APUM_PARAMETER_OFFSET:APUM_PARAMETER_END]
|
|
).hexdigest(),
|
|
"operators": [operator.__dict__ for operator in operators],
|
|
"tensors": [tensor.__dict__ for tensor in tensors],
|
|
}
|
|
|
|
|
|
def image_sha256(image: bytes) -> str:
|
|
return hashlib.sha256(image).hexdigest()
|
|
|
|
|
|
def quantize_confidence(softmax: int) -> int:
|
|
"""Apply the frozen signed-INT8 Softmax byte to public unsigned confidence."""
|
|
return max(0, min(255, int(softmax) + 128))
|
|
|
|
|
|
def select_class(scores: list[int]) -> int:
|
|
"""Choose the lowest class ID among maximum signed INT8 scores."""
|
|
if not scores:
|
|
raise ValueError("at least one class score is required")
|
|
return min(range(len(scores)), key=lambda index: (-scores[index], index))
|
|
|
|
|
|
def _signed_byte(value: int) -> int:
|
|
return value - 256 if value >= 128 else value
|
|
|
|
|
|
def _trunc_div(numerator: int, denominator: int) -> int:
|
|
return numerator // denominator if numerator >= 0 else -((-numerator) // denominator)
|
|
|
|
|
|
def _saturating_high_mul(left: int, right: int) -> int:
|
|
if left == -(1 << 31) and right == -(1 << 31):
|
|
return (1 << 31) - 1
|
|
product = left * right
|
|
nudge = (1 << 30) if product >= 0 else 1 - (1 << 30)
|
|
return _trunc_div(product + nudge, 1 << 31)
|
|
|
|
|
|
def _rounding_divide_pot(value: int, exponent: int) -> int:
|
|
if exponent == 0:
|
|
return value
|
|
mask = (1 << exponent) - 1
|
|
remainder = value & mask
|
|
threshold = (mask >> 1) + (1 if value < 0 else 0)
|
|
return (value >> exponent) + (1 if remainder > threshold else 0)
|
|
|
|
|
|
def _multiply_quantized(value: int, multiplier: int, shift: int) -> int:
|
|
left_shift = max(shift, 0)
|
|
right_shift = max(-shift, 0)
|
|
shifted = value << left_shift
|
|
if not -(1 << 31) <= shifted < (1 << 31):
|
|
raise OverflowError("APUM requantization left shift overflow")
|
|
return _rounding_divide_pot(
|
|
_saturating_high_mul(shifted, multiplier), right_shift
|
|
)
|
|
|
|
|
|
def _exp_interval_q31(raw: int) -> int:
|
|
x = raw + (1 << 28)
|
|
x2 = _saturating_high_mul(x, x)
|
|
x3 = _saturating_high_mul(x2, x)
|
|
x4 = _saturating_high_mul(x2, x2)
|
|
series = _rounding_divide_pot(
|
|
_saturating_high_mul(_rounding_divide_pot(x4, 2) + x3, 715827883)
|
|
+ x2,
|
|
1,
|
|
)
|
|
return 1895147668 + _saturating_high_mul(1895147668, x + series)
|
|
|
|
|
|
def _softmax_exp_q31(difference: int) -> int:
|
|
scaled = _saturating_high_mul((-difference) << 24, 1242899200)
|
|
quarter = 1 << 24
|
|
reduced = (scaled & (quarter - 1)) - quarter
|
|
result = _exp_interval_q31(reduced << 5)
|
|
remainder = reduced - scaled
|
|
for exponent, multiplier in (
|
|
(-2, 1672461947),
|
|
(-1, 1302514674),
|
|
(0, 790015084),
|
|
(1, 290630308),
|
|
(2, 39332535),
|
|
(3, 720401),
|
|
(4, 242),
|
|
):
|
|
if (remainder & (1 << (26 + exponent))) != 0:
|
|
result = _saturating_high_mul(result, multiplier)
|
|
return (1 << 31) - 1 if scaled == 0 else result
|
|
|
|
|
|
def _reciprocal_q31(value: int) -> int:
|
|
total = value + ((1 << 31) - 1)
|
|
half_denominator = _trunc_div(total + (1 if total >= 0 else -1), 2)
|
|
estimate = 1515870810 + _saturating_high_mul(
|
|
half_denominator, -1010580540
|
|
)
|
|
for _ in range(3):
|
|
product = _saturating_high_mul(half_denominator, estimate)
|
|
correction = 0x20000000 - product
|
|
estimate += _saturating_high_mul(estimate, correction)
|
|
return max(-(1 << 31), min((1 << 31) - 1, estimate << 1))
|
|
|
|
|
|
def _parameter_i32(image: bytes, offset: int, count: int) -> list[int]:
|
|
return list(struct.unpack_from(f"<{count}i", image, offset))
|
|
|
|
|
|
def _weighted_operator(
|
|
image: bytes, operator: ApumOperator, source: list[int], operator_index: int
|
|
) -> list[int]:
|
|
output = [0] * operator.output_channels
|
|
biases = _parameter_i32(image, operator.bias_offset, operator.output_channels)
|
|
multipliers = _parameter_i32(
|
|
image, operator.multiplier_offset, operator.output_channels
|
|
)
|
|
shifts = _parameter_i32(image, operator.shift_offset, operator.output_channels)
|
|
weights = image[
|
|
operator.weight_offset : operator.weight_offset + operator.weight_bytes
|
|
]
|
|
output_count = 25 * 5 * operator.output_channels if operator_index < 9 else 12
|
|
output = [0] * output_count
|
|
for output_index in range(output_count):
|
|
channel = output_index % operator.output_channels
|
|
spatial = output_index // operator.output_channels
|
|
output_height = spatial // 5
|
|
output_width = spatial % 5
|
|
accumulator = biases[channel]
|
|
if operator_index == 0:
|
|
for term in range(40):
|
|
input_height = output_height * 2 + term // 4 - 4
|
|
input_width = output_width * 2 + term % 4 - 1
|
|
input_value = operator.input_zero
|
|
if 0 <= input_height < 49 and 0 <= input_width < 10:
|
|
input_value = source[input_height * 10 + input_width]
|
|
weight = _signed_byte(weights[channel * 40 + term])
|
|
accumulator += (input_value - operator.input_zero) * weight
|
|
elif operator_index in (1, 3, 5, 7):
|
|
for term in range(9):
|
|
input_height = output_height + term // 3 - 1
|
|
input_width = output_width + term % 3 - 1
|
|
input_value = operator.input_zero
|
|
if 0 <= input_height < 25 and 0 <= input_width < 5:
|
|
input_value = source[
|
|
(input_height * 5 + input_width) * 64 + channel
|
|
]
|
|
weight = _signed_byte(weights[term * 64 + channel])
|
|
accumulator += (input_value - operator.input_zero) * weight
|
|
else:
|
|
for term in range(64):
|
|
input_value = source[term] if operator_index == 10 else source[spatial * 64 + term]
|
|
weight = _signed_byte(weights[channel * 64 + term])
|
|
accumulator += (input_value - operator.input_zero) * weight
|
|
if not -(1 << 31) <= accumulator < (1 << 31):
|
|
raise OverflowError(f"APUM operator {operator_index} accumulator overflow")
|
|
quantized = (
|
|
_multiply_quantized(accumulator, multipliers[channel], shifts[channel])
|
|
+ operator.output_zero
|
|
)
|
|
clamp_min = max(operator.clamp_min, operator.output_zero) if operator.flags & 1 else -128
|
|
output[output_index] = max(clamp_min, min(operator.clamp_max, quantized))
|
|
return output
|
|
|
|
|
|
def _softmax(logits: list[int]) -> list[int]:
|
|
maximum = max(logits)
|
|
exponentials = [
|
|
_softmax_exp_q31(maximum - value) if value - maximum >= -124 else 0
|
|
for value in logits
|
|
]
|
|
accumulation = sum(_rounding_divide_pot(value, 12) for value in exponentials)
|
|
headroom = 32 - accumulation.bit_length()
|
|
shifted_sum = ((accumulation << headroom) & 0xFFFFFFFF) - (1 << 31)
|
|
scale = _reciprocal_q31(shifted_sum)
|
|
num_bits_over_unit = 12 - headroom
|
|
result = []
|
|
for exponential in exponentials:
|
|
probability = _rounding_divide_pot(
|
|
_saturating_high_mul(scale, exponential), num_bits_over_unit + 23
|
|
)
|
|
result.append(max(-128, min(127, probability - 128)))
|
|
return result
|
|
|
|
|
|
def infer_apum(image: bytes, features: bytes) -> list[list[int]]:
|
|
"""Run the frozen scalar INT8 graph and return every published tensor."""
|
|
parse_apum(image)
|
|
if len(features) != 490:
|
|
raise ValueError("APUM inference requires exactly 490 INT8 feature bytes")
|
|
tensors = [_parse_tensor(image, index) for index in range(APUM_TENSOR_COUNT)]
|
|
operators = [_parse_operator(image, index) for index in range(APUM_OPERATOR_COUNT)]
|
|
current = [_signed_byte(value) for value in features]
|
|
layers: list[list[int]] = []
|
|
for operator_index, operator in enumerate(operators):
|
|
if operator_index < 9 or operator_index == 10:
|
|
current = _weighted_operator(image, operator, current, operator_index)
|
|
elif operator_index == 9:
|
|
pooled = []
|
|
for channel in range(64):
|
|
total = sum(current[spatial * 64 + channel] for spatial in range(125))
|
|
pooled.append(
|
|
_trunc_div(total + (62 if total >= 0 else -62), 125)
|
|
)
|
|
current = pooled
|
|
else:
|
|
current = _softmax(current)
|
|
if len(current) != tensors[operator.output_tensor].byte_count:
|
|
raise ValueError(f"APUM operator {operator_index} output size")
|
|
layers.append(current)
|
|
return layers
|