diff --git a/zig/kernels/hip/tiles/stream.hpp b/zig/kernels/hip/tiles/stream.hpp index e17d21ed26..5277e61227 100644 --- a/zig/kernels/hip/tiles/stream.hpp +++ b/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 -__device__ inline void stream_dots(const typename T::pair (&xp)[RG][16], const uint32_t (&w)[Dec::kWords], float (&d)[RG], - std::integer_sequence) { - (([&] { - const typename T::pair q = Dec::template stream_pair(w); +// One column's 16 code pairs of a chunk, made once for every row of the block. +template +__device__ inline void stream_codes(const uint32_t (&w)[Dec::kWords], typename T::pair (&qp)[16], + std::integer_sequence) { + if constexpr (Dec::kWords == 4 && Dec::kGapPairs && !std::is_same_v) { + // 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() == 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(w)), ...); + } } // Epi::kPair: a lane's CB columns are CB / 2 gate columns and the same up columns; out is silu(gate) * up. template -__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(w[c], qp[c], std::make_integer_sequence{}); +#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(xr[r] + static_cast(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(xr[g * RG + r] + static_cast(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(xn, xp[r], std::make_integer_sequence{}); + 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(xn, xp, std::make_integer_sequence{}); #pragma unroll - for (int r = 0; r < RG; ++r) d[r] = 0.f; - stream_dots(xp, w[c], d, std::make_integer_sequence{}); + 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 __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(blockIdx.x) / nby : static_cast(blockIdx.x); + const int by = interleaved ? static_cast(blockIdx.x) % nby : static_cast(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(sides.out[s]); } } - if (!Dec::take_item(a, blockIdx.z) || static_cast(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(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit); + stream_body(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(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit); + stream_body(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit); } else { - stream_body(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit); + stream_body(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit); } } } diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index c1f3e1b506..200b596815 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -31,6 +31,8 @@ test { _ = @import("kernel_tests/affine/rows.zig"); _ = @import("kernel_tests/affine/order.zig"); _ = @import("kernel_tests/affine/round.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"); diff --git a/zig/src/hip/kernel_tests/affine/golden.zig b/zig/src/hip/kernel_tests/affine/golden.zig new file mode 100644 index 0000000000..6754424f80 --- /dev/null +++ b/zig/src/hip/kernel_tests/affine/golden.zig @@ -0,0 +1,119 @@ +//! The 4-bit BF16 stream tile writes the bytes its code pairs made one at a time wrote: a digest per architecture. + +const std = @import("std"); +const base = @import("../rig.zig"); +const affine = @import("../../launches/affine.zig"); +const kernels = @import("../../kernels.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 with stream_pair's code pairs, on each architecture that recorded it. +const pins = [_]struct { arch: []const u8, digest: u64 }{ + .{ .arch = "gfx1151", .digest = 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))); + } + const want = for (pins) |p| { + if (std.mem.eql(u8, p.arch, kernels.arch)) break p.digest; + } else { + std.debug.print("stream golden: no digest pinned for {s}; this run wrote 0x{x:0>16}\n", .{ kernels.arch, h }); + return error.SkipZigTest; + }; + try base.expect(h == want, "stream golden: digest 0x{x:0>16}, pinned 0x{x:0>16}", .{ h, want }); +} diff --git a/zig/src/hip/kernel_tests/affine/layout.zig b/zig/src/hip/kernel_tests/affine/layout.zig new file mode 100644 index 0000000000..0e3d26534b --- /dev/null +++ b/zig/src/hip/kernel_tests/affine/layout.zig @@ -0,0 +1,142 @@ +//! A stream product writes the same bytes with a round's row blocks interleaved on x as with them laid on y. + +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 Rig = @import("product.zig").Rig; +const Kernels = affine.Kernels; + +/// Round widths past one row block: a ragged last block, whole blocks and the widest rounds. +const rounds = [_]usize{ 5, 8, 9, 16, 31, 32 }; +const most_rows = 32; +const n = 1024; +const k_dim = 2048; +const group = 64; + +const Kind = enum { plain, relu2, pair, routed, uneven, group }; + +/// The tables of `experts` stacked products of `rows` columns each, on the device. +const Tables = struct { + words: Buffer, + scale: Buffer, + bias: Buffer, + + fn init(t: *Rig, rng: *data.Rng, rows: usize, bits: c_int) !Tables { + const words_row = k_dim * @as(usize, @intCast(bits)) / 32; + const groups = k_dim / group; + const hw = try data.fill(gpa, u32, rows * words_row, rng, data.makeWord); + defer gpa.free(hw); + const hs = try data.fill(gpa, u16, rows * groups, rng, data.makeScale); + defer gpa.free(hs); + const hb = try data.fill(gpa, u16, rows * groups, rng, data.makeBias); + defer gpa.free(hb); + var words = try Buffer.fromHost(t.base, hw); + errdefer words.free(); + var scale = try Buffer.fromHost(t.base, hs); + errdefer scale.free(); + return .{ .words = words, .scale = scale, .bias = try Buffer.fromHost(t.base, hb) }; + } + + fn deinit(self: *Tables) void { + self.words.free(); + self.scale.free(); + self.bias.free(); + } +}; + +/// One launch of `kind` through `k`: its output goes to `arg.out` (fp32) or `arg.out16` (activation type). +fn launch(t: *Rig, k: *const Kernels, kind: Kind, arg: affine.Arg, sides: []const affine.Side, items: c_int) !void { + const r = &t.base.r; + const s = t.base.stream.handle; + switch (kind) { + .plain => try k.streamGo(r, arg, .{}, arg.n, 1, s), + .relu2 => try k.relu2Go(r, arg, 1, s), + .pair => try k.streamGo(r, arg, .{ .pair_cols = @divExact(arg.n, 2) }, @divExact(arg.n, 2), 1, s), + .routed, .uneven => try k.routed(r, arg, items, s), + .group => try k.groupGo(r, arg, sides, false, s), + } + try t.base.stream.synchronize(); +} + +/// `kind` at every round width, interleaved against the reference layout: the 32-bit words that differ. +fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { + var reference = t.kernels.*; + reference.stream_interleave = false; + const two = kind == .routed or kind == .uneven; + const experts: usize = if (two) 2 else 1; + var tables = try Tables.init(t, rng, experts * n, bits); + defer tables.deinit(); + const hx = try data.fill(gpa, u16, most_rows * k_dim, rng, if (t.fp16) data.makeX16 else data.makeXB); + defer gpa.free(hx); + var x = try Buffer.fromHost(t.base, hx); + defer x.free(); + // a routed round sends each row to both experts, so its output holds two pairs a row + const bytes = @as(usize, if (two) 2 else 1) * most_rows * n * 4; + var outs: [2]Buffer = .{ try Buffer.alloc(t.base, bytes), undefined }; + defer outs[0].free(); + outs[1] = try Buffer.alloc(t.base, bytes); + defer outs[1].free(); + const words_bytes = n * k_dim * @as(usize, @intCast(bits)) / 32 * 4; + const table_bytes = n * (k_dim / group) * 2; + + var diff: usize = 0; + for (rounds) |m| { + // the launch reaches the interleaved grid only past one row block + const block_rows = affine.stream_rows[affine.stream_rows.len - 1]; + try base.expect(m > block_rows, "stream layout: {d} rows fit one row block", .{m}); + // pair 2r of every row goes to the first expert, 2r + 1 to the second (the first half of the rows when uneven) + const rows: c_int = @intCast(m); + const second: usize = if (kind == .uneven) m / 2 else m; + const plan = [_]c_int{ 0, 0, rows, 1, rows, @intCast(second) }; + var members: [2 * most_rows]c_int = undefined; + for (0..m) |i| members[i] = @intCast(2 * i); + for (0..second) |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 .. m + second])); + defer mem.free(); + for (outs, [_]*const Kernels{ t.kernels, &reference }) |out, kern| { + try out.fill8(0xA5, null); + var arg: affine.Arg = .{ + .x = x.ptr, + .words = tables.words.ptr, + .scale = .{ .p = tables.scale.ptr, .kind = 1 }, + .bias = .{ .p = tables.bias.ptr, .kind = 1 }, + .out = if (kind == .relu2 or kind == .pair) 0 else out.ptr, + .out16 = if (kind == .relu2 or kind == .pair) out.ptr else 0, + .m = @intCast(m), + .n = n, + .k = k_dim, + .bits = bits, + .group = group, + .fp16 = @intFromBool(t.fp16), + }; + if (two) arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 2 }; + // three sides of unequal widths over one x, each writing its own stretch of the output + const sides = [_]affine.Side{ + .{ .words = tables.words.ptr, .scale = tables.scale.ptr, .bias = tables.bias.ptr, .n = 512, .out = out.ptr }, + .{ .words = tables.words.ptr + words_bytes / 2, .scale = tables.scale.ptr + table_bytes / 2, .bias = tables.bias.ptr + table_bytes / 2, .n = 256, .out = out.ptr + most_rows * 512 * 4 }, + .{ .words = tables.words.ptr + words_bytes * 3 / 4, .scale = tables.scale.ptr + table_bytes * 3 / 4, .bias = tables.bias.ptr + table_bytes * 3 / 4, .n = 192, .out = out.ptr + most_rows * 768 * 4 }, + }; + try launch(t, kern, kind, arg, &sides, 2); + } + const d = try data.sameOnDevice(outs[0], outs[1], bytes); + if (d != 0) std.debug.print(" stream layout {t} {d}-bit, {d} rows: {d} words differ\n", .{ kind, bits, m, d }); + diff += d; + } + return diff; +} + +test "a stream product writes the bytes of the reference layout with its row blocks interleaved on x" { + const b = try base.Rig.open(1 << 20); + defer b.close(); + var rig = Rig.open(b); + var rng: data.Rng = .{ .state = 0x2545F4914F6CDD1D }; + var diff: usize = 0; + for ([_]Kind{ .plain, .relu2, .pair, .routed, .uneven, .group }) |kind| diff += try compare(&rig, &rng, kind, 4); + diff += try compare(&rig, &rng, .plain, 3); + try base.expect(diff == 0, "stream layout: {d} words differ from the reference layout", .{diff}); +} diff --git a/zig/src/hip/launches/affine.zig b/zig/src/hip/launches/affine.zig index dcd825cf76..96bc1b5d26 100644 --- a/zig/src/hip/launches/affine.zig +++ b/zig/src/hip/launches/affine.zig @@ -15,6 +15,7 @@ const quant = @import("core").quant; const tuning = @import("../tuning/tuning.zig"); const mlx_entries = @import("mlx_entries.zig"); const stream_pick = @import("stream_pick.zig"); +const stream_layout = @import("stream_layout.zig"); const Error = runtime.Error; @@ -43,19 +44,7 @@ pub const Arg = extern struct { out16: u64 = 0, }; -/// quant/mlx_tiles.hpp's StreamSides: up to four products that share x, K, width and group in one launch. -pub const StreamSides = extern struct { - words: [4]u64 = @splat(0), - scale: [4]u64 = @splat(0), - bias: [4]u64 = @splat(0), - out: [4]u64 = @splat(0), - n: [4]c_int = @splat(0), - first: [5]c_int = @splat(0), - count: c_int = 0, - out_half: c_int = 0, - pair_cols: c_int = 0, - limit: f32 = 0, -}; +pub const StreamSides = stream_layout.StreamSides; /// One side of a group: its matrix and where its (rows, n) output goes. pub const Side = struct { words: u64, scale: u64, bias: u64, n: c_int, out: u64 }; @@ -135,6 +124,8 @@ pub const Kernels = struct { relu2_pair: [bit_widths.len][stream_rows.len][stream_cbs.len]?Function, /// Off keeps the decode tiles before the stream tile (the policy's reference switch). stream_on: bool, + /// Off lays a round's stream row blocks on the grid's y: the reference layout the interleaved one must equal. + stream_interleave: bool, block: [bit_widths.len]Function, gemm: [bit_widths.len]Function, wmma_gemm: [bit_widths.len]Function, @@ -157,6 +148,7 @@ pub const Kernels = struct { k.mode = choices.matrix; k.tile = if (choices.gemm) .gemm else .block; k.stream_on = choices.stream; + k.stream_interleave = true; k.reg = Registry.init(tableFor(caps), &mlx_entries.all); k.reg.verify(k.env(), .mlx) catch return error.Invalid; // the GEMM tile's rows of x a lane keeps (tiles/gemm.hpp GemmShape): 8 for BF16, 16 for FP16 @@ -367,7 +359,9 @@ pub const Kernels = struct { })); const table = if (sides.pair_cols > 0) &k.pairs else if (relu2) &k.relu2 else &k.stream; const f = table[b][ri][tile.at].?; - try go(r, f, .{ .x = blocks, .y = cdiv(a.m, stream_rows[ri]), .z = @intCast(items) }, .{ .x = 32 * stream_waves }, s, &args); + // a column block's row blocks run next to each other on x, so the later ones read its weights from cache + const grid = stream_layout.grid(k.stream_interleave, blocks, cdiv(a.m, stream_rows[ri]), items); + try go(r, f, grid, .{ .x = 32 * stream_waves }, s, &args); } /// One row on the row tile: a lane owns a group, a wave reads a column's row in one run. @@ -587,4 +581,5 @@ test "the lane groups of a row tile" { test { _ = mlx_entries; _ = stream_pick; + _ = stream_layout; } diff --git a/zig/src/hip/launches/stream_layout.zig b/zig/src/hip/launches/stream_layout.zig new file mode 100644 index 0000000000..ad5876b7c7 --- /dev/null +++ b/zig/src/hip/launches/stream_layout.zig @@ -0,0 +1,32 @@ +//! The stream tile's launch: quant/mlx_tiles.hpp's StreamSides and where a round's row blocks lie on the grid. +const std = @import("std"); +const hl = @import("../launch.zig"); + +/// quant/mlx_tiles.hpp's StreamSides: up to four products that share x, K, width and group in one launch. +pub const StreamSides = extern struct { + words: [4]u64 = @splat(0), + scale: [4]u64 = @splat(0), + bias: [4]u64 = @splat(0), + out: [4]u64 = @splat(0), + n: [4]c_int = @splat(0), + first: [5]c_int = @splat(0), + count: c_int = 0, + out_half: c_int = 0, + pair_cols: c_int = 0, + limit: f32 = 0, +}; + +/// The stream tile's grid; stream_tile reads a one-row grid over more than one row block as interleaved on x. +pub fn grid(interleave: bool, blocks: c_uint, ny: c_uint, items: c_int) hl.Dim3 { + if (interleave and ny > 1) return .{ .x = blocks * ny, .y = 1, .z = @intCast(items) }; + return .{ .x = blocks, .y = ny, .z = @intCast(items) }; +} + +test "a round's row blocks interleave on x only when there are several and the launch asks" { + const on = grid(true, 30, 4, 2); + try std.testing.expect(on.x == 120 and on.y == 1 and on.z == 2); + const off = grid(false, 30, 4, 2); + try std.testing.expect(off.x == 30 and off.y == 4 and off.z == 2); + const one = grid(true, 30, 1, 1); + try std.testing.expect(one.x == 30 and one.y == 1 and one.z == 1); +}