perf(hip): wide stream rounds reuse cached weights, and a chunk's code pairs are made once

zig/kernels/hip/tiles/stream.hpp

@@ -39,27 +39,53 @@ __device__ inline void stream_x_pairs(const uint32_t (&xn)[16], typename T::pair
      ...);
 }
 
-// One column's chunk against RG rows: each code pair is made once and dotted with every row's x.
-template <class Dec, typename T, int RG, int... I>
-__device__ inline void stream_dots(const typename T::pair (&xp)[RG][16], const uint32_t (&w)[Dec::kWords], float (&d)[RG],
-                                   std::integer_sequence<int, I...>) {
-    (([&] {
-         const typename T::pair q = Dec::template stream_pair<T, I>(w);
+// One column's 16 code pairs of a chunk, made once for every row of the block.
+template <class Dec, typename T, int... I>
+__device__ inline void stream_codes(const uint32_t (&w)[Dec::kWords], typename T::pair (&qp)[16],
+                                    std::integer_sequence<int, I...>) {
+    if constexpr (Dec::kWords == 4 && Dec::kGapPairs && !std::is_same_v<typename T::elem, __half>) {
+        // the byte path below hard-codes the decoder's pairing: pair 4 wd + k is codes 8 wd + k and 8 wd + k + 4
+        static_assert(Dec::gap() == 4);
+        static_assert(((Dec::template lo<I>() == 8 * (I >> 2) + (I & 3)) && ...));
+        // 4-bit BF16 pairs from bytes: a byte's float is exact and its top half is the BF16, as stream_pair makes it
+        // v_cvt_f32_ubyteN reads byte N itself; written out, LLVM would shift and extract first
+        auto f = [](uint32_t v, int b) {
+            uint32_t r;
+            switch (b) {
+                case 0: asm("v_cvt_f32_ubyte0 %0, %1" : "=v"(r) : "v"(v)); break;
+                case 1: asm("v_cvt_f32_ubyte1 %0, %1" : "=v"(r) : "v"(v)); break;
+                case 2: asm("v_cvt_f32_ubyte2 %0, %1" : "=v"(r) : "v"(v)); break;
+                default: asm("v_cvt_f32_ubyte3 %0, %1" : "=v"(r) : "v"(v)); break;
+            }
+            return r;
+        };
+        // the top halves of two floats as one pair, in one v_perm_b32 (stream_pick's selector)
+        auto pack = [](uint32_t lo, uint32_t hi) {
+            return __builtin_bit_cast(typename T::pair, stream_pick(lo, hi, 1));
+        };
 #pragma unroll
-         for (int r = 0; r < RG; ++r) d[r] = T::dot(xp[r][I], q, d[r]);
-     }()),
-     ...);
+        for (int wd = 0; wd < 4; ++wd) {
+            const uint32_t e = w[wd] & 0x0F0F0F0Fu;         // nibbles 0, 2, 4, 6 as bytes 0..3
+            const uint32_t o = (w[wd] >> 4) & 0x0F0F0F0Fu;  // nibbles 1, 3, 5, 7
+            // pair k of a word is nibbles k and k + 4
+            qp[4 * wd + 0] = pack(f(e, 0), f(e, 2));
+            qp[4 * wd + 1] = pack(f(o, 0), f(o, 2));
+            qp[4 * wd + 2] = pack(f(e, 1), f(e, 3));
+            qp[4 * wd + 3] = pack(f(o, 1), f(o, 3));
+        }
+    } else {
+        ((qp[I] = Dec::template stream_pair<T, I>(w)), ...);
+    }
 }
 
 // Epi::kPair: a lane's CB columns are CB / 2 gate columns and the same up columns; out is silu(gate) * up.
 template <class Dec, class Act, class T, class Epi, int R, int CB, bool WIDE>
-__device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, int gshift, int bx, int pair_cols, float limit) {
+__device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, int gshift, int bx, int by, int pair_cols,
+                                   float limit) {
     using pair = typename T::pair;
     constexpr bool PAIR = Epi::kPair;
-    constexpr int RG = R < 4 ? R : 4;  // rows a pass keeps in registers
-    constexpr int NG = R / RG;
     constexpr uint32_t kOne = T::ones;
-    const int row0 = blockIdx.y * R;  // a block's rows: more than 8 rows are 8-row blocks on the grid's y
+    const int row0 = by * R;  // a block's rows: row block `by` of the grid's y, or of x when interleaved
     const int lpc = 1 << lpc_log2;
     const int lane = threadIdx.x & 31;
     const int j = lane & (lpc - 1);
@@ -115,43 +141,37 @@ __device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, in
             sc[c] = live ? terms.widen(sr[c]) : 0.f;
             bi[c] = live ? terms.widen(br[c]) : 0.f;
         }
+        // every column's code pairs are made once, then each row in turn dots them in the order a row alone takes
+        pair qp[CB][16];
+#pragma unroll
+        for (int c = 0; c < CB; ++c) stream_codes<Dec, T>(w[c], qp[c], std::make_integer_sequence<int, 16>{});
+#pragma unroll
+        for (int r = 0; r < R; ++r) {
+            if (row0 + r >= a.m) break;
+            pair xp[16];
+            uint32_t xn[16];
+            const u32x4* src = reinterpret_cast<const u32x4*>(xr[r] + static_cast<long long>(q) * 32);
 #pragma unroll
-        for (int g = 0; g < NG; ++g) {
-            if (row0 + g * RG >= a.m) break;
-            pair xp[RG][16];
-            float sx[RG];
-#pragma unroll
-            for (int r = 0; r < RG; ++r) {
-                uint32_t xn[16];
-                const u32x4* src =
-                    reinterpret_cast<const u32x4*>(xr[g * RG + r] + static_cast<long long>(q) * 32);
-#pragma unroll
-                for (int v = 0; v < 4; ++v) {
-                    const u32x4 t = src[v];
-                    xn[4 * v] = t.x;
-                    xn[4 * v + 1] = t.y;
-                    xn[4 * v + 2] = t.z;
-                    xn[4 * v + 3] = t.w;
-                }
-                float s = 0.f;
-#pragma unroll
-                for (int i = 0; i < 16; ++i) {
-                    s = T::dot(__builtin_bit_cast(pair, xn[i]), __builtin_bit_cast(pair, kOne), s);
-                }
-                sx[r] = s;
-                stream_x_pairs<Dec, T>(xn, xp[r], std::make_integer_sequence<int, 16>{});
+            for (int v = 0; v < 4; ++v) {
+                const u32x4 t = src[v];
+                xn[4 * v] = t.x;
+                xn[4 * v + 1] = t.y;
+                xn[4 * v + 2] = t.z;
+                xn[4 * v + 3] = t.w;
             }
+            float sx = 0.f;
 #pragma unroll
-            for (int c = 0; c < CB; ++c) {
-                float d[RG];
+            for (int i = 0; i < 16; ++i) {
+                sx = T::dot(__builtin_bit_cast(pair, xn[i]), __builtin_bit_cast(pair, kOne), sx);
+            }
+            stream_x_pairs<Dec, T>(xn, xp, std::make_integer_sequence<int, 16>{});
 #pragma unroll
-                for (int r = 0; r < RG; ++r) d[r] = 0.f;
-                stream_dots<Dec, T, RG>(xp, w[c], d, std::make_integer_sequence<int, 16>{});
+            for (int c = 0; c < CB; ++c) {
+                float d = 0.f;
 #pragma unroll
-                for (int r = 0; r < RG; ++r) {
-                    acc[g * RG + r][c] = Dec::fold_scale(acc[g * RG + r][c], d[r], sc[c]);
-                    acc[g * RG + r][c] = Dec::fold_bias(acc[g * RG + r][c], sx[r], bi[c]);
-                }
+                for (int i = 0; i < 16; ++i) d = T::dot(xp[i], qp[c][i], d);
+                acc[r][c] = Dec::fold_scale(acc[r][c], d, sc[c]);
+                acc[r][c] = Dec::fold_bias(acc[r][c], sx, bi[c]);
             }
         }
     };
@@ -213,7 +233,11 @@ __device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, in
 // Item z of the plan (or plain product, or a group's side) at CB columns a lane, R rows; chunk group is q >> gshift.
 template <class Dec, class Act, class T, class Epi, int R, int CB, class Sides>
 __device__ inline void stream_tile(typename Dec::Args& a, const Sides& sides, int lpc_log2, int gshift) {
-    int bx = blockIdx.x;
+    // a one-row grid over more than R rows interleaves (launches/stream_layout.zig grid): row blocks neighbour on x
+    const int nby = (a.m + R - 1) / R;
+    const bool interleaved = gridDim.y == 1 && nby > 1;
+    int bx = interleaved ? static_cast<int>(blockIdx.x) / nby : static_cast<int>(blockIdx.x);
+    const int by = interleaved ? static_cast<int>(blockIdx.x) % nby : static_cast<int>(blockIdx.y);
     if (sides.count > 0) {
         int s = 0;
         while (s + 1 < sides.count && bx >= sides.first[s + 1]) ++s;
@@ -226,15 +250,15 @@ __device__ inline void stream_tile(typename Dec::Args& a, const Sides& sides, in
             a.out = static_cast<float*>(sides.out[s]);
         }
     }
-    if (!Dec::take_item(a, blockIdx.z) || static_cast<int>(blockIdx.y) * R >= a.m) return;
+    if (!Dec::take_item(a, blockIdx.z) || by * R >= a.m) return;
     if constexpr (Dec::kWideBytes == 4) {
-        stream_body<Dec, Act, T, Epi, R, CB, false>(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit);
+        stream_body<Dec, Act, T, Epi, R, CB, false>(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit);
     } else {
         // every chunk is wide-aligned when the words and the row stride are
         if (Dec::wide_aligned(a)) {
-            stream_body<Dec, Act, T, Epi, R, CB, true>(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit);
+            stream_body<Dec, Act, T, Epi, R, CB, true>(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit);
         } else {
-            stream_body<Dec, Act, T, Epi, R, CB, false>(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit);
+            stream_body<Dec, Act, T, Epi, R, CB, false>(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit);
         }
     }
 }

zig/src/hip/kernel_tests.zig

@@ -28,6 +28,8 @@ test {
     _ = @import("kernel_tests/tp.zig");
     _ = @import("kernel_tests/affine/affine.zig");
     _ = @import("kernel_tests/affine/rows.zig");
+    _ = @import("kernel_tests/affine/layout.zig");
+    _ = @import("kernel_tests/affine/golden.zig");
     _ = @import("kernel_tests/graph_split.zig");
     _ = @import("kernel_tests/flashnext.zig");
     _ = @import("kernel_tests/flashnext_attn.zig");

zig/src/hip/kernel_tests/affine/golden.zig

@@ -0,0 +1,110 @@
+//! The 4-bit BF16 stream tile writes the bytes its code pairs made one at a time wrote: a pinned digest of them.
+
+const std = @import("std");
+const base = @import("../rig.zig");
+const affine = @import("../../launches/affine.zig");
+const Buffer = base.Buffer;
+const gpa = base.gpa;
+const data = @import("data.zig");
+const prod = @import("product.zig");
+const Rig = prod.Rig;
+
+/// One row (the 1-row tile), each row block's ragged and whole sizes, and interleaved rounds.
+const rounds = [_]usize{ 1, 2, 3, 4, 5, 8, 16 };
+const most_rows = 16;
+const n = 512;
+const k_dim = 2048;
+const group = 64;
+const bits = 4;
+
+/// The Wyhash of the bytes the tile wrote when its code pairs came from stream_pair; the byte path must reproduce it.
+const want: u64 = 0x24742f9ea12d63cd;
+
+const Kind = enum { plain, pair, routed };
+
+fn digestOf(t: *Rig, rng: *data.Rng, kind: Kind) !u64 {
+    const experts: usize = if (kind == .routed) 2 else 1;
+    const words_row = k_dim * bits / 32;
+    const groups = k_dim / group;
+    const hw = try data.fill(gpa, u32, experts * n * words_row, rng, data.makeWord);
+    defer gpa.free(hw);
+    const hs = try data.fill(gpa, u16, experts * n * groups, rng, data.makeScale);
+    defer gpa.free(hs);
+    const hb = try data.fill(gpa, u16, experts * n * groups, rng, data.makeBias);
+    defer gpa.free(hb);
+    const hx = try data.fill(gpa, u16, most_rows * k_dim, rng, data.makeXB);
+    defer gpa.free(hx);
+    var words = try Buffer.fromHost(t.base, hw);
+    defer words.free();
+    var scale = try Buffer.fromHost(t.base, hs);
+    defer scale.free();
+    var bias = try Buffer.fromHost(t.base, hb);
+    defer bias.free();
+    var x = try Buffer.fromHost(t.base, hx);
+    defer x.free();
+    const bytes = 2 * most_rows * n * 4;
+    var out = try Buffer.alloc(t.base, bytes);
+    defer out.free();
+
+    const r = &t.base.r;
+    const s = t.base.stream.handle;
+    var h: u64 = 0;
+    for (rounds) |m| {
+        // top-2 routing: every row goes to both experts, pair 2r to the first and 2r + 1 to the second
+        const rows: c_int = @intCast(m);
+        const plan = [_]c_int{ 0, 0, rows, 1, rows, rows };
+        var members: [2 * most_rows]c_int = undefined;
+        for (0..m) |i| {
+            members[i] = @intCast(2 * i);
+            members[m + i] = @intCast(2 * i + 1);
+        }
+        var items = try Buffer.fromHost(t.base, @as([]const c_int, &plan));
+        defer items.free();
+        var mem = try Buffer.fromHost(t.base, @as([]const c_int, members[0 .. 2 * m]));
+        defer mem.free();
+        try out.fill8(0xA5, null);
+        var arg: affine.Arg = .{
+            .x = x.ptr,
+            .words = words.ptr,
+            .scale = .{ .p = scale.ptr, .kind = 1 },
+            .bias = .{ .p = bias.ptr, .kind = 1 },
+            .out = if (kind == .pair) 0 else out.ptr,
+            .out16 = if (kind == .pair) out.ptr else 0,
+            .m = rows,
+            .n = n,
+            .k = k_dim,
+            .bits = bits,
+            .group = group,
+            .fp16 = 0,
+        };
+        switch (kind) {
+            .plain => try t.kernels.streamGo(r, arg, .{}, n, 1, s),
+            .pair => try t.kernels.streamGo(r, arg, .{ .pair_cols = n / 2 }, n / 2, 1, s),
+            .routed => {
+                arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 2 };
+                try t.kernels.routed(r, arg, 2, s);
+            },
+        }
+        try t.base.stream.synchronize();
+        const written: usize = switch (kind) {
+            .plain => m * n * 4,
+            .pair => m * (n / 2) * 2,
+            .routed => 2 * m * n * 4,
+        };
+        h = std.hash.Wyhash.hash(h, std.mem.asBytes(&try prod.digestOf(out, written)));
+    }
+    return h;
+}
+
+test "the 4-bit BF16 stream tile writes the pinned bytes of its plain, pair and routed launches" {
+    const b = try base.Rig.open(1 << 20);
+    defer b.close();
+    var rig = Rig.open(b);
+    if (rig.fp16) return error.SkipZigTest; // the FP16 build keeps stream_pair's magic-number pairs
+    var rng: data.Rng = .{ .state = 0x9E3779B97F4A7C15 };
+    var h: u64 = 0;
+    for ([_]Kind{ .plain, .pair, .routed }) |kind| {
+        h = std.hash.Wyhash.hash(h, std.mem.asBytes(&try digestOf(&rig, &rng, kind)));
+    }
+    try base.expect(h == want, "stream golden: digest 0x{x:0>16}, pinned 0x{x:0>16}", .{ h, want });
+}