""" NINE SLICES INSTEAD OF A GATHER. The im2col route this programme uses builds its window matrix with one large fancy-index: col = padded[:, idx], 391 million scattered reads a step on the ResNet-20 shape. A scattered read cannot coalesce, so the card spends its time waiting on memory it cannot fetch in blocks, and that is what sets the pace rather than the arithmetic. But a convolution's windows are not scattered. For a FIXED TAP (dy, dx), every output position reads the input at the same offset, so that tap's contribution is a REGULAR STRIDED SLICE of the padded input: padded[:, :, dy : dy + stride*g_out : stride, dx : dx + stride*g_out : stride] which is contiguous in its last axis and evenly strided in the others. Nine such slices assemble the whole window matrix, and the hardware reads each one in blocks. AND THE LAYOUT CAN BE CHOSEN SO NOTHING IS TRANSPOSED. Writing the window matrix as (batch, in-channel x tap, position) and the values as (out- channel, in-channel x tap) makes the layer a batched matrix product whose output is already (batch, out-channel, position), which is exactly the hidden layout this framework uses. The gather version had to transpose twice, once forward and once back. The backward pass gets the same treatment. Its gradient to the input was a second scattered gather over a table of which window slots read each pixel; here it is nine strided ADDS at the same offsets, which is coalesced and deterministic for the same reason the forward is. This changes no mathematics. It is checked against the gather implementation for exact agreement on forward, on the value gradient and on the input gradient, and for bit-identical reruns, before anything is built on it. """ import numpy as np import time try: import cupy as _cp _GPU = _cp.cuda.runtime.getDeviceCount() > 0 except Exception: _GPU = False xp = _cp if _GPU else np DT = np.float32 def to_dev(a, dtype=DT): a = np.asarray(a, dtype=dtype) return xp.asarray(a) if _GPU else a def to_host(a): return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a) class SliceConv: """A folded convolution built from strided slices rather than a gather. Same values, same index meaning, same result. The value index is still (in-channel, out-channel, tap), so every claim this framework makes about partitions carries over untouched.""" def __init__(self, g_in, c_in, k, c_out, stride): self.g_in, self.c_in, self.k = g_in, c_in, k self.c_out, self.stride = c_out, stride self.g_out = g_in // stride self.ins = c_in*g_in*g_in self.out = c_out*self.g_out*self.g_out self.K = c_in*c_out*k*k + 1 self.taps = c_in*k*k self.pad = k//2 self.G = g_in + 2*self.pad self.npos = self.g_out*self.g_out def _slices(self): """Where each tap reads from in the padded input.""" s, g = self.stride, self.g_out for dy in range(self.k): for dx in range(self.k): yield (dy*self.k + dx, slice(dy, dy + s*g, s), slice(dx, dx + s*g, s)) def gather(self, x): """The window matrix as (batch, in-channel x tap, position). Nine strided slice copies. Each is a coalesced read; the fancy index it replaces was 391 million scattered ones on the ResNet-20 shape.""" n = x.shape[0] p = self.pad pd = xp.zeros((n, self.c_in, self.G, self.G), DT) pd[:, :, p:p+self.g_in, p:p+self.g_in] = x.reshape( n, self.c_in, self.g_in, self.g_in) col = xp.empty((n, self.c_in, self.k*self.k, self.npos), DT) for t, sy, sx in self._slices(): col[:, :, t, :] = pd[:, :, sy, sx].reshape(n, self.c_in, self.npos) return col.reshape(n, self.taps, self.npos) def weights_T(self, v): """The values as (out-channel, in-channel x tap). A value index is (in-channel, out-channel, tap), so this is a reshape and one transpose of a vector a few hundred thousand long, done once a step rather than over the whole window matrix.""" return v[:-1].reshape(self.c_in, self.c_out, self.k*self.k) \ .transpose(1, 0, 2).reshape(self.c_out, self.taps) def forward(self, v, x): col = self.gather(x) wt = self.weights_T(v) # (1, c_out, taps) @ (n, taps, positions) -> (n, c_out, positions), # which IS the hidden layout, so nothing is transposed afterwards z = xp.matmul(wt[None], col) return z.reshape(x.shape[0], self.out), col def backward(self, v, col, dz, need_input=True): n = dz.shape[0] d = dz.reshape(n, self.c_out, self.npos) # dW[out, taps] = sum over batch and position dW = xp.matmul(d, col.transpose(0, 2, 1)).sum(0) gv = xp.zeros(self.K, DT) gv[:-1] = dW.reshape(self.c_out, self.c_in, self.k*self.k) \ .transpose(1, 0, 2).reshape(-1) if not need_input: return gv, None # back to the window matrix, then nine strided ADDS rather than a # scatter: the same slices, accumulated in a fixed order dcol = xp.matmul(self.weights_T(v).T[None], d) dcol = dcol.reshape(n, self.c_in, self.k*self.k, self.npos) p = self.pad dpd = xp.zeros((n, self.c_in, self.G, self.G), DT) for t, sy, sx in self._slices(): dpd[:, :, sy, sx] += dcol[:, :, t, :].reshape( n, self.c_in, self.g_out, self.g_out) return gv, dpd[:, :, p:p+self.g_in, p:p+self.g_in].reshape(n, self.ins) # --------------------------------------------------------------- checks class GatherConv: """The implementation this replaces, kept so the new one can be checked against it rather than against a derivation.""" def __init__(self, g_in, c_in, k, c_out, stride): self.g_in, self.c_in, self.k = g_in, c_in, k self.c_out, self.stride = c_out, stride self.g_out = g_in // stride self.ins = c_in*g_in*g_in self.out = c_out*self.g_out*self.g_out self.K = c_in*c_out*k*k + 1 self.taps = c_in*k*k go, gi = self.g_out, g_in oy, ox = np.divmod(np.arange(go*go), go) ci = np.arange(c_in)[:, None, None] dy = np.arange(k)[None, :, None] dx = np.arange(k)[None, None, :] iy = oy[:, None, None, None]*stride - k//2 + dy ix = ox[:, None, None, None]*stride - k//2 + dx ok = (iy >= 0) & (iy < gi) & (ix >= 0) & (ix < gi) flat = (ci*gi*gi + np.clip(iy, 0, gi-1)*gi + np.clip(ix, 0, gi-1)) flat = np.where(ok, flat, self.ins).reshape(go*go, self.taps) self.fwd_idx = to_dev(flat, np.int64) if _GPU else flat.astype(np.int64) readers = {} for pp in range(go*go): for tc in range(self.taps): s = int(flat[pp, tc]) if s < self.ins: readers.setdefault(s, []).append(pp*self.taps + tc) w = max((len(v) for v in readers.values()), default=1) tbl = np.full((self.ins, w), go*go*self.taps, np.int64) for s, v in readers.items(): tbl[s, :len(v)] = v self.bwd_idx = to_dev(tbl, np.int64) if _GPU else tbl def weights(self, v): return v[:-1].reshape(self.c_in, self.c_out, self.k*self.k) \ .transpose(0, 2, 1).reshape(self.taps, self.c_out) def forward(self, v, x): n = x.shape[0] xz = xp.concatenate([x, xp.zeros((n, 1), DT)], 1) col = xz[:, self.fwd_idx].reshape(n*self.g_out**2, self.taps) z = (col @ self.weights(v)).reshape(n, self.g_out**2, self.c_out) return z.transpose(0, 2, 1).reshape(n, self.out), col def backward(self, v, col, dz, need_input=True): n = dz.shape[0] d = dz.reshape(n, self.c_out, self.g_out**2).transpose(0, 2, 1) \ .reshape(n*self.g_out**2, self.c_out) gv = xp.zeros(self.K, DT) gv[:-1] = (col.T @ d).reshape(self.c_in, self.k*self.k, self.c_out) \ .transpose(0, 2, 1).reshape(-1) if not need_input: return gv, None dcol = (d @ self.weights(v).T).reshape(n, -1) pad = xp.concatenate([dcol, xp.zeros((n, 1), DT)], 1) return gv, pad[:, self.bwd_idx].sum(2) def check(shapes=((32, 3, 16, 1), (32, 16, 16, 1), (16, 16, 32, 2), (8, 32, 64, 2), (32, 16, 32, 1)), n=8, reps=3, batches=(8, 64, 256)): rg = np.random.default_rng(0) ok = True print(f" {'shape':>22s} {'forward':>9s} {'values':>9s} {'input':>9s} " f"{'same twice':>11s} " + " ".join(f"{'b='+str(b):>7s}" for b in batches)) print(f" {'':>22s} {'':>9s} {'':>9s} {'':>9s} {'':>11s} " f" how many times faster the slice route is") for (g, ci, co, st) in shapes: A, B = SliceConv(g, ci, 3, co, st), GatherConv(g, ci, 3, co, st) v = rg.normal(size=A.K).astype(np.float32); v[-1] = 0.0 x = rg.normal(size=(n, ci*g*g)).astype(np.float32) vd, xd = to_dev(v), to_dev(x) za, ca = A.forward(vd, xd) zb, cb = B.forward(vd, xd) e1 = float(np.abs(to_host(za) - to_host(zb)).max()) dz = to_dev(rg.normal(size=(n, A.out))) ga, da = A.backward(vd, ca, dz) gb, db = B.backward(vd, cb, dz) sc = max(float(np.abs(to_host(gb)).max()), 1e-9) e2 = float(np.abs(to_host(ga) - to_host(gb)).max())/sc e3 = float(np.abs(to_host(da) - to_host(db)).max()) z2, c2 = A.forward(vd, xd) g2, d2 = A.backward(vd, c2, dz) rep = (to_host(z2).tobytes() == to_host(za).tobytes() and to_host(g2).tobytes() == to_host(ga).tobytes() and to_host(d2).tobytes() == to_host(da).tobytes()) rr = [] for bn in batches: xb = to_dev(rg.normal(size=(bn, ci*g*g))) db = to_dev(rg.normal(size=(bn, A.out))) ts = {} for nm, M in (("slice", A), ("gather", B)): M.forward(vd, xb) # warm the allocator if _GPU: _cp.cuda.Stream.null.synchronize() t0 = time.time() for _ in range(reps): zz, cc = M.forward(vd, xb) M.backward(vd, cc, db) if _GPU: _cp.cuda.Stream.null.synchronize() ts[nm] = (time.time()-t0)/reps rr.append(ts["gather"]/ts["slice"]) print(f" {f'{g}x{g} {ci}->{co} /{st}':>22s} {e1:9.1e} {e2:9.1e} " f"{e3:9.1e} {str(rep):>11s} " + " ".join(f"{r:6.2f}x" for r in rr)) ok = ok and e1 < 2e-4 and e2 < 2e-4 and e3 < 2e-4 and rep return ok if __name__ == "__main__": print("=" * 92) print("NINE SLICES INSTEAD OF A GATHER") print("=" * 92) print(f" backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}\n") good = check() print(f"\n agrees with the gather implementation everywhere and " f"reproduces: {good}") print(f"\n THE BATCH MATTERS AND A FIRST VERSION MEASURED ONLY AT 8,") print(f" where the slice route LOST: nine slice copies forward and nine") print(f" adds backward are eighteen kernel launches against the fancy") print(f" index's one, and at batch 8 there is not enough data for the") print(f" reads to matter. That overhead is FIXED; the data is not. The") print(f" column that decides anything is the rightmost.")