File size: 2,150 Bytes
309b968 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | """Native-speed deployment of the verified units, without losing the guarantee.
A verified unit is a finite function. Its neural net is only needed to *prove*
correctness (N/N). For SPEED you materialize the proven function as a lookup
table -- run the net once over its whole (small) domain -- then every later call
is an array index at native memory speed. Because the net is N/N-verified, the
LUT is bit-identical to the net, which is bit-identical to the true op. So:
neural forward == LUT == native integer op (all bit-exact)
That's the "freeze the mesh to its matrix" lesson: verify once (slow, offline),
deploy native (fast). The LUTs are tiny: mul 256x256, requant 65536, relu 256.
"""
from __future__ import annotations
import numpy as np
def build_mul8_lut(mul) -> np.ndarray:
"""[256,256] signed-product table, indexed by unsigned bytes. Net runs once."""
a = np.repeat(np.arange(256), 256)
b = np.tile(np.arange(256), 256)
prod = mul.mul_array(a, b) # verified neural multiply, ONCE
return prod.reshape(256, 256).astype(np.int64)
def build_requant16_lut(rq) -> np.ndarray:
"""[65536] int16->int8 table, indexed by acc & 0xFFFF."""
return rq.requant_array(np.arange(65536)).astype(np.int64)
def build_relu8_lut(relu) -> np.ndarray:
"""[256] int8 ReLU table, indexed by unsigned byte."""
return relu.relu_array(np.arange(256)).astype(np.int64)
class LUTBackend:
"""GEMM via the materialized (verified) multiply table + integer accumulate."""
name = "lut"
def __init__(self, mul):
self.mul_lut = build_mul8_lut(mul)
def available(self):
return True
def gemm(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
au = (A.astype(np.int64) & 0xFF)
bu = (B.astype(np.int64) & 0xFF)
# products via table lookup, then sum over the contraction axis
prod = self.mul_lut[au[:, None, :], bu.T[None, :, :]] # (m, n, k)
from . import instrument
instrument.bump("VerifiedMul(LUT).gemms", 1)
instrument.bump("VerifiedMul(LUT).products", prod.size)
return prod.sum(axis=2)
|