File size: 5,481 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """Verified neural INT8 multiplier -- the atom of a GPU tensor core / VNNI lane.
A monolithic MLP cannot learn a bit-exact 8x8 multiply (the high-order product
bits are too nonlinear). So -- exactly like the byte-slice / ripple trick the
other projects use for hard functions -- we shrink the *verified atom* to a
4-bit unsigned multiply and compose everything exactly:
atom: NeuralMul4 -- unsigned 4x4 -> 8, domain 16*16 = 256, verified N/N.
8x8: a*b = ah*bh<<8 + (ah*bl + al*bh)<<4 + al*bl (unsigned), exact.
signed: a_s = a_u - 256*a7 ; Baugh-Wooley correction, exact integer glue.
Only the 4x4 multiply is neural (and N/N-proven); the shifts, adds and sign
correction are exact composition. The parallel *throughput* of thousands of such
lanes is NOT here -- that needs real silicon (kernel.py, the GUDA-role path).
"""
from __future__ import annotations
import numpy as np
import torch
from .common import bits_of, int_of, pm, mlp, verify, train
class NeuralMul4:
"""N/N-verified unsigned 4x4 -> 8 multiplier (the finite atom)."""
def __init__(self, h: int = 128, layers: int = 3):
self.net = mlp(8, 8, h=h, layers=layers)
def dataset(self) -> tuple[torch.Tensor, torch.Tensor]:
X, Y = [], []
for a in range(16):
ab = bits_of(a, 4)
for b in range(16):
X.append(pm(torch.cat([ab, bits_of(b, 4)])))
Y.append(bits_of(a * b, 8))
return torch.stack(X), torch.stack(Y)
def fit(self, steps: int = 4000, lr: float = 2e-3, tag: str = "mul4"):
X, Y = self.dataset()
train(self.net, X, Y, steps=steps, lr=lr, tag=tag)
return self
def verify(self) -> tuple[int, int]:
X, Y = self.dataset()
return verify(self.net, X, Y)
@torch.no_grad()
def mul(self, a: int, b: int) -> int:
self.net.eval()
x = pm(torch.cat([bits_of(a & 0xF, 4), bits_of(b & 0xF, 4)])).unsqueeze(0)
return int_of((self.net(x)[0] > 0).float())
@torch.no_grad()
def mul_array(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
"""Batched unsigned 4x4 -> 8 over arrays (one neural forward for all)."""
self.net.eval()
a = (np.asarray(a).astype(np.int64) & 0xF)
b = (np.asarray(b).astype(np.int64) & 0xF)
idx = np.arange(4)
bits_a = (a[:, None] >> idx) & 1
bits_b = (b[:, None] >> idx) & 1
x = np.concatenate([bits_a, bits_b], axis=1).astype(np.float32) * 2.0 - 1.0
out = (self.net(torch.from_numpy(x)) > 0).to(torch.int64).numpy()
from . import instrument
instrument.bump("NeuralMul4.forward_calls", 1)
instrument.bump("NeuralMul4.products", a.shape[0])
return (out * (1 << np.arange(8))).sum(axis=1)
class NeuralMul8:
"""Signed 8x8 -> 16 multiply, composed exactly from the verified 4x4 atom."""
def __init__(self, h: int = 128, layers: int = 3):
self.atom = NeuralMul4(h=h, layers=layers)
def fit(self, steps: int = 4000, lr: float = 2e-3, tag: str = "mul4"):
self.atom.fit(steps=steps, lr=lr, tag=tag)
return self
def verify_atom(self) -> tuple[int, int]:
return self.atom.verify()
def _umul8(self, a_u: int, b_u: int) -> int:
"""Unsigned 8x8 -> 16 via four verified 4x4 sub-products (exact glue)."""
al, ah = a_u & 0xF, (a_u >> 4) & 0xF
bl, bh = b_u & 0xF, (b_u >> 4) & 0xF
ll = self.atom.mul(al, bl)
lh = self.atom.mul(al, bh)
hl = self.atom.mul(ah, bl)
hh = self.atom.mul(ah, bh)
return ll + ((lh + hl) << 4) + (hh << 8)
def mul(self, a: int, b: int) -> int:
"""Signed product via the verified atom + exact Baugh-Wooley correction."""
a_u, b_u = a & 0xFF, b & 0xFF
a7, b7 = (a_u >> 7) & 1, (b_u >> 7) & 1
prod = self._umul8(a_u, b_u) - (a7 * b_u << 8) - (b7 * a_u << 8) + (a7 * b7 << 16)
prod &= 0xFFFF # 16-bit two's complement
return prod - 65536 if prod >= 32768 else prod
@torch.no_grad()
def mul_array(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
"""Batched signed 8x8 -> 16 over arrays, via the verified 4x4 atom.
Four nibble sub-products are batched into ONE neural forward, then the
exact shift/add glue and Baugh-Wooley sign correction are applied in
numpy. Result is bit-exact (the atom is N/N-verified)."""
a = np.asarray(a).astype(np.int64).ravel()
b = np.asarray(b).astype(np.int64).ravel()
au, bu = a & 0xFF, b & 0xFF
al, ah = au & 0xF, (au >> 4) & 0xF
bl, bh = bu & 0xF, (bu >> 4) & 0xF
n = au.shape[0]
pa = np.concatenate([al, al, ah, ah])
pb = np.concatenate([bl, bh, bl, bh])
prod = self.atom.mul_array(pa, pb)
ll, lh, hl, hh = prod[:n], prod[n:2*n], prod[2*n:3*n], prod[3*n:4*n]
u = ll + ((lh + hl) << 4) + (hh << 8)
a7, b7 = (au >> 7) & 1, (bu >> 7) & 1
p = (u - (a7 * bu << 8) - (b7 * au << 8) + (a7 * b7 << 16)) & 0xFFFF
return np.where(p >= 32768, p - 65536, p)
@torch.no_grad()
def verify(self, full: bool = True) -> tuple[int, int]:
"""Exhaustively check the composed signed multiply over all 65536 inputs."""
ok = 0
for a in range(-128, 128):
for b in range(-128, 128):
if self.mul(a, b) == a * b:
ok += 1
return ok, 256 * 256
|