|
| 1 | +import torch |
| 2 | +import torch.nn as nn |
| 3 | +import numpy as np |
| 4 | +import pytfmbs |
| 5 | +import struct |
| 6 | +from .constants import ( |
| 7 | + SRAM_BANK_A_OFFSET, SRAM_BANK_B_OFFSET, SRAM_TILE_STRIDE, |
| 8 | + KERNEL_T_GEMM, HINT_ZERO_SKIP, LANES_PER_TILE, MAX_TILES |
| 9 | +) |
| 10 | + |
| 11 | +def pack_pt5_numpy(trits): |
| 12 | + """ |
| 13 | + Packs trits (-1, 0, 1) into PT-5 format (5 trits per byte). |
| 14 | + trits: numpy array of shape (N,) |
| 15 | + returns: numpy array of uint8 |
| 16 | + """ |
| 17 | + trits = np.asanyarray(trits) |
| 18 | + # Ensure trits are in {-1, 0, 1} |
| 19 | + trits = np.clip(trits, -1, 1).astype(np.int8) |
| 20 | + |
| 21 | + # Pad to multiple of 5 |
| 22 | + padding = (5 - (len(trits) % 5)) % 5 |
| 23 | + if padding > 0: |
| 24 | + trits = np.concatenate([trits, np.zeros(padding, dtype=np.int8)]) |
| 25 | + |
| 26 | + trits_reshaped = trits.reshape(-1, 5) |
| 27 | + powers = 3 ** np.arange(5) |
| 28 | + packed = np.sum((trits_reshaped + 1) * powers, axis=1).astype(np.uint8) |
| 29 | + return packed |
| 30 | + |
| 31 | +def pack_gemv_input(x): |
| 32 | + """ |
| 33 | + Packs a ternary input vector x for GEMV. |
| 34 | + Each element x[d] is duplicated for all lanes to fill a multi-lane frame. |
| 35 | + """ |
| 36 | + x = np.asanyarray(x).astype(np.int8) |
| 37 | + depth = len(x) |
| 38 | + |
| 39 | + # Repeated x: [depth, LANES_PER_TILE] |
| 40 | + x_expanded = np.repeat(x[:, np.newaxis], LANES_PER_TILE, axis=1) |
| 41 | + |
| 42 | + # Reshape to [depth * (LANES_PER_TILE/5), 5] |
| 43 | + x_reshaped = x_expanded.reshape(-1, 5) |
| 44 | + |
| 45 | + powers = 3 ** np.arange(5) |
| 46 | + packed_bytes = np.sum((x_reshaped + 1) * powers, axis=1).astype(np.uint8) |
| 47 | + |
| 48 | + # Now we have depth * 3 bytes. We need to add 1 byte of padding after every 3 bytes for 32-bit alignment. |
| 49 | + packed_32 = np.zeros((depth, 4), dtype=np.uint8) |
| 50 | + packed_32[:, 0:3] = packed_bytes.reshape(depth, 3) |
| 51 | + |
| 52 | + return packed_32.tobytes() |
| 53 | + |
| 54 | +def pack_gemv_weights(w_ternary): |
| 55 | + """ |
| 56 | + Packs ternary weights [out_features, in_features] for T-GEMV. |
| 57 | + Returns a list of bytes, one for each tile (up to MAX_TILES). |
| 58 | + """ |
| 59 | + out_features, in_features = w_ternary.shape |
| 60 | + w_t = w_ternary.T # [in_features, out_features] |
| 61 | + |
| 62 | + tile_data = [] |
| 63 | + for t in range(MAX_TILES): |
| 64 | + tile_out_start = t * LANES_PER_TILE |
| 65 | + tile_out_end = (t + 1) * LANES_PER_TILE |
| 66 | + |
| 67 | + # Extract weights for this tile's lanes |
| 68 | + if tile_out_start < out_features: |
| 69 | + tile_w = w_t[:, tile_out_start:min(tile_out_end, out_features)] |
| 70 | + # Pad lanes to LANES_PER_TILE |
| 71 | + if tile_w.shape[1] < LANES_PER_TILE: |
| 72 | + tile_w = np.pad(tile_w, ((0, 0), (0, LANES_PER_TILE - tile_w.shape[1]))) |
| 73 | + |
| 74 | + # tile_w: [in_features, LANES_PER_TILE] |
| 75 | + tile_w_reshaped = tile_w.reshape(-1, 5) |
| 76 | + powers = 3 ** np.arange(5) |
| 77 | + packed_bytes = np.sum((tile_w_reshaped + 1) * powers, axis=1).astype(np.uint8) |
| 78 | + |
| 79 | + packed_32 = np.zeros((in_features, 4), dtype=np.uint8) |
| 80 | + packed_32[:, 0:3] = packed_bytes.reshape(in_features, 3) |
| 81 | + tile_data.append(packed_32.tobytes()) |
| 82 | + else: |
| 83 | + tile_data.append(None) |
| 84 | + |
| 85 | + return tile_data |
| 86 | + |
| 87 | +class TFMBSLinearFunction(torch.autograd.Function): |
| 88 | + @staticmethod |
| 89 | + def forward(ctx, input, fabric, weight_addr, bias, in_features, out_features, tile_mask): |
| 90 | + # input: [batch, in_features] |
| 91 | + # input must be quantized to {-1, 0, 1} for the fabric |
| 92 | + x_ternary = torch.sign(input).to(torch.int8).cpu().numpy() |
| 93 | + |
| 94 | + batch_size = x_ternary.shape[0] |
| 95 | + outputs = [] |
| 96 | + |
| 97 | + for i in range(batch_size): |
| 98 | + packed_x = pack_gemv_input(x_ternary[i]) |
| 99 | + |
| 100 | + # Load input to Bank B of all active tiles |
| 101 | + for t in range(MAX_TILES): |
| 102 | + if tile_mask & (1 << t): |
| 103 | + fabric.load(SRAM_BANK_B_OFFSET + t * SRAM_TILE_STRIDE, packed_x) |
| 104 | + |
| 105 | + # Run |
| 106 | + tfd = { |
| 107 | + "base_addr": weight_addr, |
| 108 | + "depth": in_features, |
| 109 | + "lane_count": LANES_PER_TILE, |
| 110 | + "tile_mask": tile_mask, |
| 111 | + "exec_hints": KERNEL_T_GEMM | HINT_ZERO_SKIP, |
| 112 | + } |
| 113 | + fabric.run(tfd) |
| 114 | + |
| 115 | + # Results |
| 116 | + all_res = fabric.results(-1) |
| 117 | + outputs.append(torch.tensor(all_res[:out_features], dtype=torch.float32)) |
| 118 | + |
| 119 | + output = torch.stack(outputs).to(input.device) |
| 120 | + if bias is not None: |
| 121 | + output += bias |
| 122 | + return output |
| 123 | + |
| 124 | + @staticmethod |
| 125 | + def backward(ctx, grad_output): |
| 126 | + """ |
| 127 | + Backward pass currently returns the gradient with respect to input. |
| 128 | + Weights are not updated through the fabric in this version (Inference only). |
| 129 | + """ |
| 130 | + # Straight-through estimator or similar could be used for training. |
| 131 | + # For now, we return None for parameters not being trained via Fabric. |
| 132 | + return grad_output, None, None, None, None, None, None |
| 133 | + |
| 134 | +class TFMBSLinear(nn.Module): |
| 135 | + """ |
| 136 | + TFMBS-accelerated Linear layer. |
| 137 | + Automatically quantizes weights and offloads GEMV to Ternary Fabric. |
| 138 | + """ |
| 139 | + def __init__(self, in_features, out_features, fabric=None, bias=True, weight_addr=SRAM_BANK_A_OFFSET): |
| 140 | + super().__init__() |
| 141 | + self.in_features = in_features |
| 142 | + self.out_features = out_features |
| 143 | + self.fabric = fabric or pytfmbs.Fabric() |
| 144 | + |
| 145 | + self.weight = nn.Parameter(torch.empty(out_features, in_features)) |
| 146 | + if bias: |
| 147 | + self.bias = nn.Parameter(torch.empty(out_features)) |
| 148 | + else: |
| 149 | + self.register_parameter('bias', None) |
| 150 | + |
| 151 | + self.reset_parameters() |
| 152 | + self.resident = False |
| 153 | + self.tile_mask = (1 << ((out_features + LANES_PER_TILE - 1) // LANES_PER_TILE)) - 1 |
| 154 | + self.weight_addr = weight_addr |
| 155 | + |
| 156 | + def reset_parameters(self): |
| 157 | + nn.init.kaiming_uniform_(self.weight, a=np.sqrt(5)) |
| 158 | + if self.bias is not None: |
| 159 | + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) |
| 160 | + bound = 1 / np.sqrt(fan_in) |
| 161 | + nn.init.uniform_(self.bias, -bound, bound) |
| 162 | + |
| 163 | + def load_to_fabric(self): |
| 164 | + """Quantizes and moves weights to Fabric memory.""" |
| 165 | + w_ternary = torch.sign(self.weight).to(torch.int8).detach().cpu().numpy() |
| 166 | + tile_data = pack_gemv_weights(w_ternary) |
| 167 | + |
| 168 | + for t, data in enumerate(tile_data): |
| 169 | + if data is not None: |
| 170 | + # Load to Bank A of tile t |
| 171 | + self.fabric.load(SRAM_BANK_A_OFFSET + t * SRAM_TILE_STRIDE, data) |
| 172 | + |
| 173 | + self.resident = True |
| 174 | + |
| 175 | + def forward(self, input): |
| 176 | + if not self.resident: |
| 177 | + self.load_to_fabric() |
| 178 | + |
| 179 | + return TFMBSLinearFunction.apply( |
| 180 | + input, self.fabric, self.weight_addr, self.bias, |
| 181 | + self.in_features, self.out_features, self.tile_mask |
| 182 | + ) |
0 commit comments