| """Verified INT8 activation / requantize units -- the rest of a quantized layer. |
| |
| A quantized linear layer is: int8 x int8 -> int32 accumulate -> requantize to |
| int8 -> activation. The multiply/accumulate is covered by mul8 + exact int32 |
| sum. These two units cover the tail, each over a FINITE, enumerable domain so |
| they are verified BIT-EXACT (N/N): |
| |
| NeuralReLU8 int8 -> int8, y = max(0, x) domain 256 |
| NeuralRequant16 int16 -> int8, y = sat_int8(x >> shift) domain 65536 |
| |
| The int32 -> int16 narrowing that precedes requantize is an EXACT saturating |
| clamp (integer wiring, not neural). So a full qlinear is exact/verified through |
| every stage. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
|
|
| from .common import bits_of, int_of, pm, mlp, verify, train |
| from . import instrument |
|
|
|
|
| def _s(v: int, bits: int) -> int: |
| """two's-complement raw -> signed.""" |
| m = 1 << (bits - 1) |
| return v - (1 << bits) if v >= m else v |
|
|
|
|
| def sat_int8(x: int) -> int: |
| return -128 if x < -128 else (127 if x > 127 else x) |
|
|
|
|
| class NeuralReLU8: |
| """N/N-verified INT8 ReLU: y = max(0, x).""" |
|
|
| def __init__(self, h: int = 64, layers: int = 2): |
| self.net = mlp(8, 8, h=h, layers=layers) |
|
|
| def dataset(self): |
| X, Y = [], [] |
| for b in range(256): |
| X.append(pm(bits_of(b, 8))) |
| Y.append(bits_of(max(0, _s(b, 8)) & 0xFF, 8)) |
| return torch.stack(X), torch.stack(Y) |
|
|
| def fit(self, steps: int = 3000, lr: float = 2e-3, tag: str = "relu8"): |
| X, Y = self.dataset(); train(self.net, X, Y, steps=steps, lr=lr, tag=tag); return self |
|
|
| def verify(self): |
| X, Y = self.dataset(); return verify(self.net, X, Y) |
|
|
| @torch.no_grad() |
| def relu(self, x: int) -> int: |
| self.net.eval() |
| raw = int_of((self.net(pm(bits_of(x & 0xFF, 8)).unsqueeze(0))[0] > 0).float()) |
| return _s(raw, 8) |
|
|
| @torch.no_grad() |
| def relu_array(self, arr: np.ndarray) -> np.ndarray: |
| """Batched int8 ReLU over an array (one neural forward).""" |
| self.net.eval() |
| a = np.asarray(arr).astype(np.int64).ravel() & 0xFF |
| bits = ((a[:, None] >> np.arange(8)) & 1).astype(np.float32) * 2.0 - 1.0 |
| out = (self.net(torch.from_numpy(bits)) > 0).to(torch.int64).numpy() |
| raw = (out * (1 << np.arange(8))).sum(axis=1) |
| instrument.bump("NeuralReLU8.forward_calls", 1) |
| instrument.bump("NeuralReLU8.elements", a.shape[0]) |
| return np.where(raw >= 128, raw - 256, raw).reshape(np.asarray(arr).shape) |
|
|
|
|
| class NeuralRequant16: |
| """N/N-verified requantize: int16 -> int8, y = sat_int8(x >> shift).""" |
|
|
| def __init__(self, shift: int = 8, h: int = 256, layers: int = 3): |
| self.shift = int(shift) |
| self.net = mlp(16, 8, h=h, layers=layers) |
|
|
| def _ref(self, x_signed: int) -> int: |
| return sat_int8(x_signed >> self.shift) |
|
|
| def dataset(self): |
| X, Y = [], [] |
| for b in range(65536): |
| X.append(pm(bits_of(b, 16))) |
| Y.append(bits_of(self._ref(_s(b, 16)) & 0xFF, 8)) |
| return torch.stack(X), torch.stack(Y) |
|
|
| def fit(self, steps: int = 6000, lr: float = 2e-3, tag: str = "requant16"): |
| X, Y = self.dataset(); train(self.net, X, Y, steps=steps, lr=lr, tag=tag); return self |
|
|
| def verify(self): |
| X, Y = self.dataset(); return verify(self.net, X, Y) |
|
|
| @torch.no_grad() |
| def requant(self, x: int) -> int: |
| self.net.eval() |
| raw = int_of((self.net(pm(bits_of(x & 0xFFFF, 16)).unsqueeze(0))[0] > 0).float()) |
| return _s(raw, 8) |
|
|
| @torch.no_grad() |
| def requant_array(self, arr: np.ndarray) -> np.ndarray: |
| """Batched int16->int8 requantize over an array (one neural forward).""" |
| self.net.eval() |
| a = np.asarray(arr).astype(np.int64).ravel() & 0xFFFF |
| bits = ((a[:, None] >> np.arange(16)) & 1).astype(np.float32) * 2.0 - 1.0 |
| out = (self.net(torch.from_numpy(bits)) > 0).to(torch.int64).numpy() |
| raw = (out * (1 << np.arange(8))).sum(axis=1) |
| instrument.bump("NeuralRequant16.forward_calls", 1) |
| instrument.bump("NeuralRequant16.elements", a.shape[0]) |
| return np.where(raw >= 128, raw - 256, raw).reshape(np.asarray(arr).shape) |
|
|
