| """TIED neural multiplier -- one small cell, iterated (a la neural-raytracing). |
| |
| Multiply is inherently iterative (shift-and-add), so instead of four separate |
| 4x4 atoms we use a SINGLE weight-tied cell applied across the 8 steps -- exactly |
| the raytracing pattern of marching one shared cell rather than stacking many. |
| |
| The tied cell is the conditional-add step of a shift-add multiplier: |
| |
| cell(acc, a, enable) = acc + (a if enable else 0) # 8+8+1 -> 9 bits |
| |
| Its domain is 2^8 * 2^8 * 2 = 131072 -- enumerable, so the cell is verified |
| BIT-EXACT over its whole domain (N/N). The shift and re-assembly between steps |
| are exact wiring, not neural. Looping the one tied cell 8x yields the full |
| unsigned 8x8 -> 16 product; signed uses the exact Baugh-Wooley correction. |
| |
| One verified cell, reused 8 times -- the smallest possible learned core. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from .common import bits_of, int_of, pm, mlp, verify, train |
|
|
|
|
| class NeuralMACStep: |
| """N/N-verified tied cell: acc8 + (a8 if enable else 0) -> 9-bit sum.""" |
|
|
| def __init__(self, h: int = 128, layers: int = 3): |
| self.net = mlp(17, 9, h=h, layers=layers) |
|
|
| def dataset(self) -> tuple[torch.Tensor, torch.Tensor]: |
| X, Y = [], [] |
| for acc in range(256): |
| ab = bits_of(acc, 8) |
| for a in range(256): |
| aa = bits_of(a, 8) |
| for en in (0, 1): |
| X.append(pm(torch.cat([ab, aa, torch.tensor([float(en)])]))) |
| Y.append(bits_of(acc + (a if en else 0), 9)) |
| return torch.stack(X), torch.stack(Y) |
|
|
| def fit(self, steps: int = 5000, lr: float = 2e-3, tag: str = "macstep"): |
| 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 step(self, acc: int, a: int, enable: int) -> int: |
| self.net.eval() |
| x = pm(torch.cat([bits_of(acc & 0xFF, 8), bits_of(a & 0xFF, 8), |
| torch.tensor([float(enable)])])).unsqueeze(0) |
| return int_of((self.net(x)[0] > 0).float()) |
|
|
|
|
| class TiedMul8: |
| """Signed 8x8 -> 16 multiply from ONE tied MAC-step cell, iterated 8x.""" |
|
|
| def __init__(self, h: int = 128, layers: int = 3): |
| self.cell = NeuralMACStep(h=h, layers=layers) |
|
|
| def fit(self, steps: int = 5000, lr: float = 2e-3, tag: str = "macstep"): |
| self.cell.fit(steps=steps, lr=lr, tag=tag) |
| return self |
|
|
| def verify_cell(self) -> tuple[int, int]: |
| return self.cell.verify() |
|
|
| def _umul8(self, a_u: int, b_u: int) -> int: |
| """Unsigned product via the tied shift-add loop (cell reused 8x).""" |
| combined = b_u & 0xFF |
| for _ in range(8): |
| enable = combined & 1 |
| hi = (combined >> 8) & 0xFF |
| s = self.cell.step(hi, a_u, enable) |
| combined = (combined & 0xFF) | (s << 8) |
| combined >>= 1 |
| return combined & 0xFFFF |
|
|
| def mul(self, a: int, b: int) -> int: |
| 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 |
| return prod - 65536 if prod >= 32768 else prod |
|
|
| @torch.no_grad() |
| def verify_unsigned(self) -> tuple[int, int]: |
| ok = 0 |
| for a in range(256): |
| for b in range(256): |
| if self._umul8(a, b) == a * b: |
| ok += 1 |
| return ok, 256 * 256 |
|
|
| @torch.no_grad() |
| def verify(self) -> tuple[int, int]: |
| 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 |
|
|
