diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index cb54012303..15680d812f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -76,6 +76,12 @@ checked bit for bit against ExLlamaV3's dequantization. Flash Next's optional int8 and int4 KV caches (`families/qwen4_exp/cuda/kvcache.py`) follow the cache quantization scheme of [ExLlamaV3](https://github.com/turboderp-org/exllamav3) `-cq 8` and `-cq 4` (MIT License, Copyright (c) 2025 Turboderp, text below): groups of 32, one fp16 absmax scale per group, the group rotated by a 32-point Hadamard, midpoint-grid codes, `compand_a == 0`. 8-bit stores each code as a signed int8 (`q - 128`). 4-bit stores two unsigned codes per byte, low nibble first (the same bits as ExLlamaV3's little-endian packing, a uint8 tensor rather than their uint32 words). Their dequantizer folds another `1/sqrt(32)` into the scale and applies the unnormalized butterfly on the way out; this cache applies the normalized H32 to the query and to the merged output instead, and leaves the stored codes rotated. Scales match their quantizer bit for bit. Reconstructed values agree within fp16/bf16 rounding (under 0.01 on random groups), not bit for bit. The quantizer and the attention dequant are written for TensorFold and checked against an independent reference of that arithmetic. +The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roce.zig`, `roce_verbs.zig`) adapts +the one-shot all-gather protocol of [b12x](https://github.com/local-inference-lab/b12x)'s RoCEnante +(`b12x/comm/roce/`, PRs #295 and #315): a device-resident epoch, sequence flags after each payload, two slots by +sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, +Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. + ## Vendored code and weights `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified `dflash/model_mlx.py` from diff --git a/build.zig b/build.zig index 054e688280..bd0476b368 100644 --- a/build.zig +++ b/build.zig @@ -5,6 +5,8 @@ const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); +const tp_build = @import("zig/build/tp.zig"); +const sessions_build = @import("zig/build/sessions.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -175,6 +177,8 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); + tp_build.add(b, target, optimize, test_step); + sessions_build.add(b, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/sessions.zig b/zig/build/sessions.zig new file mode 100644 index 0000000000..a5d8ea6f3c --- /dev/null +++ b/zig/build/sessions.zig @@ -0,0 +1,13 @@ +//! Session memory's build: the family-neutral `sessions` module (zig/src/sessions) and its host tests on any OS. + +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/sessions/sessions.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) })); + b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on real files (no GPU)").dependOn(&run.step); + test_step.dependOn(&run.step); +} diff --git a/zig/build/tp.zig b/zig/build/tp.zig new file mode 100644 index 0000000000..f5dacc46e3 --- /dev/null +++ b/zig/build/tp.zig @@ -0,0 +1,59 @@ +//! Tensor parallelism's build: host tests on any OS, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. + +const std = @import("std"); + +/// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL). +fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + cuda.addOptions("kernel_options", options); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); + return cuda; +} + +/// The `tp` module over `cuda` (an engine passes its own runtime module); verbs (zig/src/fabric/verbs.zig) for RoCE. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/tp.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("verbs", b.createModule(.{ .root_source_file = b.path("zig/src/fabric/verbs.zig"), .target = target, .optimize = optimize, .link_libc = true })); + return m; +} + +/// mailbox.cu as a fatbin: sm_120 (RTX PRO 6000), sm_121 (GB10), and compute_90 PTX for anything else. +fn mailboxFatbin(b: *std.Build, nvcc: []const u8) std.Build.LazyPath { + const run = b.addSystemCommand(&.{ nvcc, "-fatbin", "-O3", "-std=c++17", "-gencode=arch=compute_120,code=sm_120", "-gencode=arch=compute_121,code=sm_121", "-gencode=arch=compute_90,code=compute_90", "-o" }); + const out = run.addOutputFileArg("tp_mailbox.fatbin"); + run.addFileArg(b.path("zig/src/fabric/tp/mailbox.cu")); + return out; +} + +/// `pub const mailbox: []const u8`, the fatbin's bytes or empty. +fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { + const wf = b.addWriteFiles(); + if (image) |file| { + _ = wf.addCopyFile(file, "tp_mailbox.fatbin"); + // the driver takes 8-byte aligned images (Module.load refuses others); @embedFile alone promises nothing + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "const bytes align(16) = @embedFile(\"tp_mailbox.fatbin\").*;\npub const mailbox: []const u8 = &bytes;\n") }); + } + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "pub const mailbox: []const u8 = \"\";\n") }); +} + +pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { + const host = b.graph.host; + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); + test_step.dependOn(&tests.step); + if (target.result.os.tag != .linux) return; + const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel"); + const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed"); + const image: ?std.Build.LazyPath = if (prebuilt) |p| b.graph.cwdRelativePath(p) else if (nvcc) |tool| mailboxFatbin(b, tool) else null; + const cuda = runtime(b, target, optimize); + const runner = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + runner.addImport("cuda", cuda); + runner.addImport("tp", module(b, target, optimize, cuda)); + runner.addImport("tp_kernel", kernelModule(b, image)); + const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner }); + b.installArtifact(exe); + b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index f72c13b227..2187c24416 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -216,7 +216,13 @@ pub const Open = struct { }; /// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +pub const Opened = struct { + engine: Engine, + close: *const fn (ctx: *anyopaque) void, + ctx: *anyopaque, + /// A drain's deadline: rounds end at the next boundary, what is left fails with `reason`; null: `close` alone. + halt: ?*const fn (ctx: *anyopaque, reason: []const u8) void = null, +}; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0 }; @@ -381,6 +387,8 @@ pub const LoneHooks = struct { /// The lane core served to the HTTP threads (lane_host.zig). pub const serial_host = @import("serial_host.zig"); pub const LaneHost = @import("lane_host.zig").LaneHost; +/// An opt-in family round planner for the lane host, in place of its `prefill_step` chunks (lane_rounds.zig). +pub const Rounds = @import("lane_host.zig").Rounds; /// Exact prompt reuse between requests, for any family (prompt_cache.zig). pub const prompt_cache = @import("prompt_cache.zig"); @@ -402,6 +410,8 @@ test { _ = @import("serial_host_test.zig"); _ = @import("lane_host.zig"); _ = @import("lane_host_test.zig"); + _ = @import("lane_rounds_test.zig"); + _ = @import("lane_halt_test.zig"); _ = @import("lane_host_reuse_test.zig"); _ = @import("lane_host_warm_test.zig"); _ = @import("prompt_cache.zig"); diff --git a/zig/src/core/lane_halt.zig b/zig/src/core/lane_halt.zig new file mode 100644 index 0000000000..466de561d8 --- /dev/null +++ b/zig/src/core/lane_halt.zig @@ -0,0 +1,34 @@ +//! A stop's drain deadline on the lane host: the rounds end at a boundary and what is left fails with a reason. +const std = @import("std"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Job = LaneHost.Job; + +/// `LaneHost.halt`: the engine thread stops at the next round boundary; what is left fails with `reason`. +pub fn halt(h: *LaneHost, reason: []const u8) void { + h.mutex.lockUncancelable(h.io); + h.closing = true; + h.halt_reason = reason; + h.wake.broadcast(h.io); + h.mutex.unlock(h.io); + if (h.thread) |t| t.join(); + h.thread = null; +} + +/// `halt`: between rounds, every request queued, filling or running fails with `reason` (its stream discarded). +pub fn haltAll(h: *LaneHost, reason: []const u8) void { + h.lock(); + const queued = h.gpa.dupe(*Job, h.queued.items) catch &.{}; + h.queued.clearRetainingCapacity(); + const jobs = h.gpa.dupe(*Job, h.admitted.items) catch &.{}; + h.admitted.clearRetainingCapacity(); + h.filling.clearRetainingCapacity(); // a job still filling is admitted: it fails below, its lane freed + h.unlock(); + if (queued.len + jobs.len > 0) std.log.warn("lane host: stopping at a round boundary: {d} running and {d} queued request(s) end ({s})", .{ jobs.len, queued.len, reason }); + for (jobs) |job| { + if (job.started and !job.stream.finished) h.core.discard(&job.stream); + h.finish(job, .failed, reason); + } + for (queued) |job| h.finish(job, .failed, reason); + h.gpa.free(queued); + h.gpa.free(jobs); +} diff --git a/zig/src/core/lane_halt_test.zig b/zig/src/core/lane_halt_test.zig new file mode 100644 index 0000000000..87ba6710ab --- /dev/null +++ b/zig/src/core/lane_halt_test.zig @@ -0,0 +1,210 @@ +//! The lane host's halt: a stop's drain deadline ends the rounds at a boundary, what is left fails with the reason. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; + +test "halt (a stop's drain deadline): the rounds end at a round boundary, every request queued or running fails with the reason" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + // one prompt prefilling at a time (the next waits queued), 3 rows a round + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: u64 = 0, decoding: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + rounds: u32 = 0, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + const p = self(ctx); + for (p.seqs.items) |x| if (!x.decoding) return .wait; + p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + const p = self(ctx); + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (p.seqs.items) |*x| if (!x.decoding) { + if (x.done < x.n - 1) { + const end = @min(x.done + 3, x.n - 1); + try pieces.append(gpa, .{ .key = x.key, .start = x.done, .end = end, .save = false }); + x.done = end; + } + if (x.done == x.n - 1) { + try finals.append(gpa, x.key); + x.decoding = true; + } + }; + p.rounds += 1; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + message: []const u8 = "", + /// the planner's rounds when the finish arrived (on the engine thread, between rounds) + rounds_at_end: u32 = 0, + plan: *Plan, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| { + b.done = f.reason; + b.message = f.message; + b.rounds_at_end = b.plan.rounds; + }, + else => {}, + } + } + fn count(b: *@This()) usize { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return b.tokens.items.len; + } + }; + var long_prompt: [600]u32 = undefined; + for (&long_prompt, 0..) |*t, i| t.* = @intCast(1 + i % 13); + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; + var boxes: [3]Box = .{ .{ .plan = &pl }, .{ .plan = &pl }, .{ .plan = &pl } }; + defer for (&boxes) |*b| b.tokens.deinit(gpa); + // 1: decoding (no end of its own), 2: prefilling a long prompt in pieces, 3: queued behind 2 + const r1: Request = .{ .prompt = &p1, .max_tokens = 1 << 30 }; + const r2: Request = .{ .prompt = &long_prompt, .max_tokens = 1 << 30 }; + const r3: Request = .{ .prompt = &p1, .max_tokens = 8 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &boxes[0], .event = Box.event }); + while (boxes[0].count() < 4) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + try e.submit(2, &r2, .{ .ctx = &boxes[1], .event = Box.event }); + try e.submit(3, &r3, .{ .ctx = &boxes[2], .event = Box.event }); + while (boxes[0].count() < 20) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + host.halt("server restarting: retry shortly"); + const rounds = pl.rounds; + for (&boxes) |*b| { + try std.testing.expectEqual(@as(?Reason, .failed), b.done); + try std.testing.expectEqualStrings("server restarting: retry shortly", b.message); + // no round was planned after the finishes: the loop stopped at the boundary they were sent at + try std.testing.expectEqual(rounds, b.rounds_at_end); + } + // the decoding stream's tokens are whole rounds' (the fake's own sequence), the prefilling one had none yet + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, &p1); + for (boxes[0].tokens.items) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } + try std.testing.expectEqual(@as(usize, 0), boxes[1].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), boxes[2].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), pl.seqs.items.len); + // nothing runs after it; the stop that follows (a close) finds the thread gone + std.Io.sleep(std.testing.io, .fromMilliseconds(20), .awake) catch {}; + try std.testing.expectEqual(rounds, pl.rounds); + host.stop(); +} + +test "halt fails a request still filling its prompt with the reason and frees its lane; a decoding one fails too" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .prefill_chunks = 2 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.stepped(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + const Chunks = struct { + n: std.atomic.Value(usize) = .init(0), // chunks the second request's pass ran + fn call(ctx: *anyopaque, _: *lanes.Stream, _: usize) void { + const c: *@This() = @ptrCast(@alignCast(ctx)); + _ = c.n.fetchAdd(1, .monotonic); + } + }; + var chunks: Chunks = .{}; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + message: []const u8 = "", + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| { + b.done = f.reason; + b.message = f.message; + }, + else => {}, + } + } + fn count(b: *@This()) usize { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return b.tokens.items.len; + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var boxes: [3]Box = .{ .{}, .{}, .{} }; + defer for (&boxes) |*b| b.tokens.deinit(gpa); + // 1: decoding (no end of its own), 2: filling a prompt that never ends, 3: queued behind both lanes + const r1: Request = .{ .prompt = &p1, .max_tokens = 1 << 30 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 8 }; + const r3: Request = .{ .prompt = &p2, .max_tokens = 8 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &boxes[0], .event = Box.event }); + while (boxes[0].count() < 4) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + // nothing fills now: the next pass takes chunks until the halt + target.prefill_chunks = std.math.maxInt(usize); + target.prefill_hook = Chunks.call; + target.prefill_hook_ctx = &chunks; + try e.submit(2, &r2, .{ .ctx = &boxes[1], .event = Box.event }); + try e.submit(3, &r3, .{ .ctx = &boxes[2], .event = Box.event }); + while (chunks.n.load(.monotonic) < 3) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + host.halt("server restarting: retry shortly"); + for (&boxes) |*b| { + try std.testing.expectEqual(@as(?Reason, .failed), b.done); + try std.testing.expectEqualStrings("server restarting: retry shortly", b.message); + } + try std.testing.expect(boxes[0].tokens.items.len >= 4); + try std.testing.expectEqual(@as(usize, 0), boxes[1].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), boxes[2].tokens.items.len); + // every lane freed: the decoding stream's and the filling one's + try std.testing.expectEqual(@as(usize, 0), target.lanes.count()); + try std.testing.expectEqual(@as(usize, 0), host.filling.items.len); + host.stop(); +} diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 09ff0fa2de..248d4230f2 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -17,12 +17,20 @@ const Reason = api.Reason; const Stats = api.Stats; const SubmitError = api.SubmitError; +pub const Rounds = rounds.Rounds; +pub const Proposers = rounds.Proposers; +const rounds = @import("lane_rounds.zig"); +const lane_stream = @import("lane_stream.zig"); +const lane_halt = @import("lane_halt.zig"); + pub const LaneHost = struct { gpa: Allocator, io: std.Io, core: *lanes.Engine, info_: Info, min_match: i64 = 4, + /// A family's own copy proposer for each request (null: lanes.SuffixLookup at `min_match`) + proposers: ?Proposers = null, mutex: std.Io.Mutex = .init, wake: std.Io.Condition = .init, queued: std.ArrayList(*Job) = .empty, @@ -30,6 +38,8 @@ pub const LaneHost = struct { filling: std.ArrayList(*Job) = .empty, // admitted jobs whose prompt is still going in, a chunk a round, oldest first cancels: std.ArrayList(Id) = .empty, closing: bool = false, + /// `halt`'s reason: what is left finishes failed with it (null: a stop cancels what is left) + halt_reason: ?[]const u8 = null, thread: ?std.Thread = null, decoded: std.ArrayList(status.Mark) = .empty, // tokens a round landed, for the 2 s decode rate prefill_rate: f64 = 0, @@ -47,20 +57,31 @@ pub const LaneHost = struct { learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused lessons: std.ArrayList(Lesson) = .empty, // learn requests waiting for an idle engine lesson_open: bool = false, // the learner holds a begun job (engine thread only) + /// Runs first on the engine thread: a backend whose API state is per thread (a CUDA context) binds it there. + on_thread: ?ThreadInit = null, + /// opt-in: a family round planner (prompt pieces, decode rows in one pass); null: whole or `prefill_step` prompts + rounds: ?Rounds = null, + pieces: std.ArrayList(Rounds.Piece) = .empty, + finals: std.ArrayList(usize) = .empty, + + pub const ThreadInit = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque) anyerror!void }; const Lesson = struct { request: *const api.LearnRequest, sink: api.LearnSink }; - const Job = struct { + pub const Job = struct { host: *LaneHost, id: Id, request: *const Request, sink: Sink, stream: lanes.Stream = undefined, proposer: lanes.SuffixLookup = undefined, + /// the family's proposer (LaneHost.proposers), in place of `proposer` + custom: ?lanes.proposer.Proposer = null, delivered: usize = 0, started: bool = false, prefill_sent: bool = false, // a lone driver's prefilled event went out fill_began: bool = false, // its first prompt chunk ran + submitted: i96 = 0, began: i96 = 0, prefilled: ?i96 = null, entry: ?*pc.Entry = null, // the kept state the backend restores, until its prompt pass reports @@ -68,7 +89,7 @@ pub const LaneHost = struct { kept0: u64 = 0, // the store's kept count when the job looked it up /// The backend's prompt pass stands at a mark: the cache keeps the stream's state there. - fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { + pub fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { const job: *Job = @ptrCast(@alignCast(ptr)); job.reported(); // before a keep can evict the entry the pass restored if (job.host.cache) |store| _ = store.keep(job.request.prompt, at, s, job.request.chunks, job.host.spans(job.request)); @@ -92,10 +113,11 @@ pub const LaneHost = struct { } pub fn start(h: *LaneHost) !void { + if (h.rounds != null and !h.core.backend.pieces()) return error.PiecesUnsupported; // a drafter wrapper runs whole prompts h.thread = try std.Thread.spawn(.{ .stack_size = 16 << 20 }, run, .{h}); } - /// Stops admitting, cancels what is left and joins the engine thread. + /// Stops admitting, cancels what is left and joins the engine thread (after `halt`: the thread is gone already). pub fn stop(h: *LaneHost) void { h.mutex.lockUncancelable(h.io); h.closing = true; @@ -112,8 +134,13 @@ pub const LaneHost = struct { h.cancels.deinit(h.gpa); h.decoded.deinit(h.gpa); h.live_tokens.deinit(h.gpa); + h.pieces.deinit(h.gpa); + h.finals.deinit(h.gpa); } + /// A drain's deadline: the round in flight completes on every rank, so the ranks stop at the same boundary. + pub const halt = lane_halt.halt; + pub fn engine(h: *LaneHost) Engine { return .{ .ctx = h, .vtable = &.{ .info = infoFn, .submit = submitFn, .cancel = cancelFn, .status = status.read, .memory = memoryFn, .keepalive = keepaliveFn, .learn = learn.submit } }; } @@ -134,7 +161,7 @@ pub const LaneHost = struct { fn submitFn(ctx: *anyopaque, id: Id, request: *const Request, sink: Sink) SubmitError!void { const h = self(ctx); const job = h.gpa.create(Job) catch return error.Busy; - job.* = .{ .host = h, .id = id, .request = request, .sink = sink }; + job.* = .{ .host = h, .id = id, .request = request, .sink = sink, .submitted = std.Io.Clock.awake.now(h.io).toNanoseconds() }; h.mutex.lockUncancelable(h.io); defer h.mutex.unlock(h.io); if (h.closing) { @@ -171,7 +198,7 @@ pub const LaneHost = struct { } /// A job's cancel hook for its prompt pass: its id is in `cancels`, or the host is closing (read under the lock). - fn cancelled(ctx: *anyopaque) bool { + pub fn cancelled(ctx: *anyopaque) bool { const job: *Job = @ptrCast(@alignCast(ctx)); job.host.lock(); defer job.host.unlock(); @@ -179,7 +206,7 @@ pub const LaneHost = struct { } /// A background pass's yield hook: a foreground request waits or a stream decodes (read under the lock). - fn yieldsTo(ctx: *anyopaque) bool { + pub fn yieldsTo(ctx: *anyopaque) bool { const job: *Job = @ptrCast(@alignCast(ctx)); job.host.lock(); defer job.host.unlock(); @@ -203,7 +230,7 @@ pub const LaneHost = struct { } /// Sends a stream's new tokens; true once it has finished (its job freed). - fn deliver(h: *LaneHost, job: *Job) bool { + pub fn deliver(h: *LaneHost, job: *Job) bool { h.send(job); if (!job.stream.finished) return false; const reason: Reason = switch (job.stream.reason) { @@ -216,7 +243,8 @@ pub const LaneHost = struct { return true; } - fn finish(h: *LaneHost, job: *Job, reason: Reason, message: []const u8) void { + pub fn finish(h: *LaneHost, job: *Job, reason: Reason, message: []const u8) void { + if (h.rounds) |rd| rd.vtable.left(rd.ctx, @intFromPtr(job)); job.reported(); h.gpa.free(job.marks); const s = &job.stream; @@ -225,13 +253,13 @@ pub const LaneHost = struct { if (job.started) { if (s.grammar) |g| g.deinit(h.gpa); s.deinit(h.gpa); - job.proposer.deinit(); + lane_stream.freeProposer(h, job); } h.gpa.destroy(job); } /// Cancels queued and admitted jobs named since the last round. - fn takeCancels(h: *LaneHost) void { + pub fn takeCancels(h: *LaneHost) void { h.mutex.lockUncancelable(h.io); const ids = h.gpa.dupe(Id, h.cancels.items) catch &.{}; h.cancels.clearRetainingCapacity(); @@ -259,11 +287,11 @@ pub const LaneHost = struct { h.gpa.free(ids); } - fn lock(h: *LaneHost) void { + pub fn lock(h: *LaneHost) void { h.mutex.lockUncancelable(h.io); } - fn unlock(h: *LaneHost) void { + pub fn unlock(h: *LaneHost) void { h.mutex.unlock(h.io); } @@ -288,50 +316,9 @@ pub const LaneHost = struct { return true; }; h.unlock(); - const r = job.request; - var reuse: lanes.stream.Reuse = .{}; - // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass - if (h.cache) |store| if (store.lookupRewind(h.gpa, r.prompt, r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; - job.kept0 = store.counts.kept; - job.marks = l.marks; - reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; - } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); - job.stream = lanes.Stream.init(h.gpa, .{ - .id = "request", - .prompt = r.prompt, - .max_new = r.max_tokens, - .eos = r.eos, - .sampling = r.sampling, - .drafts = r.drafts, - .proposer = job.proposer.proposer(), - .stop_check = if (r.stop) |s| .{ .ptr = s.ctx, .check = s.check } else null, - .cancel_check = .{ .ptr = job, .check = cancelled }, - .yield_check = if (r.background and r.max_tokens == 0) .{ .ptr = job, .check = yieldsTo } else null, - .think_budget = r.think_budget, - .think_close = r.think_close, - .think_end = if (r.think_end) |t| t else -1, - .logprobs = r.logprobs, - .loop_guard = r.loop_guard, - .chunks = r.chunks, - .history_len = r.history_len, - .shared_prefixes = r.shared_prefixes, - .reuse = reuse, - .grammar = if (r.structure) |st| h.constraint(st) catch |e| { - job.proposer.deinit(); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); - } else null, - .decode_spans = h.spans(r), - .decide = r.decide, - }) catch { - job.proposer.deinit(); - return h.drop(job, "out of memory"); - }; - job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; - if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); + if (!lane_stream.open(h, job)) return true; + const began = job.began; + if (job.request.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (h.loneFits(job)) return h.runLone(job, began); if (h.core.fills()) { h.filling.append(h.gpa, job) catch return h.drop(job, "out of memory"); @@ -378,12 +365,12 @@ pub const LaneHost = struct { } /// A reply's grammar from the server's compiler (Info.structures engines get requests that carry one). - fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint { + pub fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint { const c = st.compiler orelse return error.NoGrammarCompiler; return c.constraint(h.gpa, h.io, @fromBackingInt(@intCast(@backingInt(st.kind))), st.text, st.after); } - fn prefilled(h: *LaneHost, job: *Job, began: i96) void { + pub fn prefilled(h: *LaneHost, job: *Job, began: i96) void { const done = std.Io.Clock.awake.now(h.io).toNanoseconds(); h.lock(); if (done > began) h.prefill_rate = @as(f64, @floatFromInt(job.request.prompt.len - job.stream.cached)) / (@as(f64, @floatFromInt(done - began)) / 1e9); @@ -396,7 +383,7 @@ pub const LaneHost = struct { } /// A request's decoded prompt rows, for a backend that keeps their bits (`decoded_rows`); otherwise none. - fn spans(h: *const LaneHost, r: *const Request) []const [2]u32 { + pub fn spans(h: *const LaneHost, r: *const Request) []const [2]u32 { return if (h.decoded_rows) r.decode_spans else &.{}; } @@ -473,7 +460,7 @@ pub const LaneHost = struct { return x.text(x.ctx, e) orelse @errorName(e); } - fn drop(h: *LaneHost, job: *Job, message: []const u8) bool { + pub fn drop(h: *LaneHost, job: *Job, message: []const u8) bool { h.remove(job); if (job.started and !job.stream.finished) h.core.discard(&job.stream); h.finish(job, .failed, message); @@ -483,7 +470,7 @@ pub const LaneHost = struct { /// No memory for another lane yet: the job waits at the queue's head, and admission holds until a lane frees. fn requeue(h: *LaneHost, job: *Job) bool { job.stream.deinit(h.gpa); - job.proposer.deinit(); + lane_stream.freeProposer(h, job); job.started = false; job.fill_began = false; job.entry = null; // looked up again at its next admission @@ -505,14 +492,17 @@ pub const LaneHost = struct { return false; } - /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + /// A job cancelled in its prompt pass, its lane already released (by a `halt`: failed with its reason). + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); - h.finish(job, .cancelled, ""); + h.lock(); + const why = h.halt_reason; + h.unlock(); + h.finish(job, if (why == null) .cancelled else .failed, why orelse ""); return true; } - fn remove(h: *LaneHost, job: *Job) void { + pub fn remove(h: *LaneHost, job: *Job) void { h.unfill(job); h.lock(); defer h.unlock(); @@ -523,6 +513,9 @@ pub const LaneHost = struct { } fn run(h: *LaneHost) void { + if (h.on_thread) |t| t.run(t.ctx) catch |e| std.log.err("lane host: the engine thread's init failed ({t})", .{e}); + // a family planner: prompt pieces and decode rows in one pass; else whole or `prefill_step` chunked prompts + if (h.rounds) |rd| return rounds.run(h, rd); while (true) { h.takeCancels(); while (h.admitOne()) {} @@ -573,7 +566,8 @@ pub const LaneHost = struct { } /// A failed round ends every stream it held, with the backend's error. - fn failAll(h: *LaneHost, message: []const u8) void { + pub fn failAll(h: *LaneHost, message: []const u8) void { + std.log.err("lane host: a round failed ({s}): every admitted stream fails", .{message}); h.lock(); const jobs = h.gpa.dupe(*Job, h.admitted.items) catch &.{}; h.admitted.clearRetainingCapacity(); @@ -586,7 +580,8 @@ pub const LaneHost = struct { h.gpa.free(jobs); } - fn closeAll(h: *LaneHost) void { + pub fn closeAll(h: *LaneHost) void { + if (h.halt_reason) |reason| return lane_halt.haltAll(h, reason); h.lock(); for (h.queued.items) |job| h.cancels.append(h.gpa, job.id) catch {}; for (h.admitted.items) |job| h.cancels.append(h.gpa, job.id) catch {}; @@ -597,4 +592,6 @@ pub const LaneHost = struct { test { _ = @import("lane_host_test.zig"); + _ = @import("lane_rounds_test.zig"); + _ = @import("lane_halt_test.zig"); } diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index d02118e137..4ef15db8f9 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -257,7 +257,7 @@ test "a lane host serves the core's own tokens, in order, and cancels between ro } /// A learner that takes `total` steps over a lesson, then reports it learned in that many. -const FakeLearner = struct { +pub const FakeLearner = struct { sink: ?api.LearnSink = null, steps: u32 = 0, total: u32 = 3, @@ -266,11 +266,11 @@ const FakeLearner = struct { fn emit(l: *FakeLearner, event: api.LearnEvent) void { l.sink.?.event(l.sink.?.ctx, &event); } - fn begin(ctx: *anyopaque, request: *const api.LearnRequest, sink: api.LearnSink) anyerror!void { + pub fn begin(ctx: *anyopaque, request: *const api.LearnRequest, sink: api.LearnSink) anyerror!void { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.* = .{ .sink = sink, .total = l.total, .examples = request.train.len }; } - fn step(ctx: *anyopaque) api.Learner.Step { + pub fn step(ctx: *anyopaque) api.Learner.Step { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.steps += 1; if (l.steps < l.total) return .{ .done = false, .changed = false }; @@ -278,20 +278,20 @@ const FakeLearner = struct { l.emit(.{ .done = .{} }); return .{ .done = true, .changed = true }; } - fn abort(ctx: *anyopaque) void { + pub fn abort(ctx: *anyopaque) void { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.emit(.{ .done = .{ .message = "aborted" } }); } }; /// Learn events as tags, read on the test's thread once `done` arrives. -const LearnBox = struct { +pub const LearnBox = struct { mutex: std.Io.Mutex = .init, tags: std.ArrayList(std.meta.Tag(api.LearnEvent)) = .empty, steps: u32 = 0, done: bool = false, - fn sink(b: *LearnBox) api.LearnSink { + pub fn sink(b: *LearnBox) api.LearnSink { return .{ .ctx = b, .event = event }; } fn event(ctx: *anyopaque, e: *const api.LearnEvent) void { @@ -302,7 +302,7 @@ const LearnBox = struct { if (e.* == .learned) b.steps = e.learned.steps; if (e.* == .done) b.done = true; } - fn wait(b: *LearnBox) void { + pub fn wait(b: *LearnBox) void { while (true) { b.mutex.lockUncancelable(std.testing.io); const d = b.done; diff --git a/zig/src/core/lane_rounds.zig b/zig/src/core/lane_rounds.zig new file mode 100644 index 0000000000..6db9191cca --- /dev/null +++ b/zig/src/core/lane_rounds.zig @@ -0,0 +1,279 @@ +//! A family's round planner under the lane host: prompts in pieces between the other streams' decode rounds. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const host = @import("lane_host.zig"); +const lane_stream = @import("lane_stream.zig"); +const learn = @import("lane_host_learn.zig"); +const status = @import("lane_host_status.zig"); +const LaneHost = host.LaneHost; +const Job = LaneHost.Job; +const Allocator = std.mem.Allocator; + +/// A family's round planner: admissions, then `admitted`, begins, prompt pieces, a decode round, finals, `after`. +pub const Rounds = struct { + ctx: *anyopaque, + vtable: *const VTable, + + pub const Verdict = enum { admit, wait, skip, refuse }; + pub const Piece = struct { key: usize, start: u64, end: u64, save: bool, run: bool = true }; + pub const Ask = struct { key: usize, prompt: []const u32, max_tokens: u32, submitted_s: f64 }; + + pub const VTable = struct { + /// An admission pass starts (a lane is free and requests wait). + begin_admit: *const fn (ctx: *anyopaque) void, + /// The next queued request: start it, wait (it and those behind it), skip past it, or refuse it (`words`). + admit: *const fn (ctx: *anyopaque, ask: Ask, words: *[]const u8) Verdict, + /// The pass is decided: what runs before the admitted streams begin. + admitted: *const fn (ctx: *anyopaque) anyerror!void, + /// Around an admitted request's stream begin (`damaged`: its saved prefix could not be restored). + arm: *const fn (ctx: *anyopaque, key: usize) void, + began: *const fn (ctx: *anyopaque, key: usize, damaged: bool) void, + /// This round's pieces, in order, and the requests whose prompts are in. + plan: *const fn (ctx: *anyopaque, pieces: *std.ArrayList(Piece), finals: *std.ArrayList(usize)) anyerror!void, + /// A request left: finished, cancelled, failed or refused, started or not. + left: *const fn (ctx: *anyopaque, key: usize) void, + /// After the round: the pieces' seconds and the decode round's (finals included). + after: *const fn (ctx: *anyopaque, piece_s: f64, decode_s: f64) void, + }; +}; + +/// Makes and frees a request's copy proposer (`prompt` and `eos` outlive it). +pub const Proposers = struct { + ctx: *anyopaque, + make: *const fn (ctx: *anyopaque, gpa: Allocator, prompt: []const u32, eos: []const u32) anyerror!lanes.proposer.Proposer, + free: *const fn (ctx: *anyopaque, p: lanes.proposer.Proposer) void, +}; + +/// One round under the family's planner; false when it did nothing (nothing to admit, prefill or decode). +fn round(h: *LaneHost, rd: Rounds) bool { + var did = false; + // the admission pass, decided whole before any of it runs + var admits: std.ArrayList(*Job) = .empty; + defer admits.deinit(h.gpa); + var refused: std.ArrayList(struct { job: *Job, why: []const u8 }) = .empty; + defer refused.deinit(h.gpa); + h.lock(); + if (h.queued.items.len > 0 and h.admitted.items.len < h.info_.lanes) { + rd.vtable.begin_admit(rd.ctx); + var free = h.info_.lanes - h.admitted.items.len; + var i: usize = 0; + while (i < h.queued.items.len and free > 0) { + const job = h.queued.items[i]; + var why: []const u8 = ""; + const v = rd.vtable.admit(rd.ctx, .{ .key = @intFromPtr(job), .prompt = job.request.prompt, .max_tokens = job.request.max_tokens, .submitted_s = seconds(job.submitted) }, &why); + switch (v) { + .admit => { + if (h.admitted.append(h.gpa, job)) |_| {} else |_| break; + if (admits.append(h.gpa, job)) |_| {} else |_| { + _ = h.admitted.pop(); + break; + } + _ = h.queued.orderedRemove(i); + free -= 1; + }, + .wait => break, + .skip => i += 1, + .refuse => { + if (refused.append(h.gpa, .{ .job = job, .why = why })) |_| {} else |_| break; + _ = h.queued.orderedRemove(i); + }, + } + } + } + h.unlock(); + for (refused.items) |x| h.finish(x.job, .failed, x.why); + rd.vtable.admitted(rd.ctx) catch |e| { + h.failAll(h.words(e)); + return true; + }; + for (admits.items) |job| { + did = true; + if (!lane_stream.open(h, job)) continue; + rd.vtable.arm(rd.ctx, @intFromPtr(job)); + const b = h.core.beginStream(&job.stream) catch |e| { + _ = h.drop(job, h.words(e)); + continue; + }; + rd.vtable.began(rd.ctx, @intFromPtr(job), b.damaged); + } + // this round's pieces, a decode round of the streams decoding, then the prompts whose rows are in + rd.vtable.plan(rd.ctx, &h.pieces, &h.finals) catch |e| { + h.failAll(h.words(e)); + return true; + }; + const t0 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + if (h.core.cfg.piece_runs and h.core.backend.vtable.pieces != null) { + // one call for the round's pieces (the backend runs several slots' rows in one forward); a failure fails each + var list: [16]lanes.backend.Backend.Piece = undefined; + var jobs: [16]*Job = undefined; + var n: usize = 0; + for (h.pieces.items) |piece| { + if (!piece.run) continue; + did = true; + const job: *Job = @ptrFromInt(piece.key); + if (job.stream.finished) continue; + if (n == list.len) { + _ = h.drop(job, "too many prompt pieces in one round"); + continue; + } + list[n] = .{ .stream = &job.stream, .start = piece.start, .end = piece.end, .save = piece.save }; + jobs[n] = job; + n += 1; + } + if (n > 0) h.core.pieces(list[0..n]) catch |e| { + for (jobs[0..n]) |job| dropPiece(h, job, h.words(e)); + }; + } else for (h.pieces.items) |piece| { + if (!piece.run) continue; + did = true; + const job: *Job = @ptrFromInt(piece.key); + if (job.stream.finished) continue; + h.core.piece(&job.stream, piece.start, piece.end, piece.save) catch |e| { + dropPiece(h, job, h.words(e)); + }; + } + const t1 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + // Config.join_tail: a joined prompt's last token decodes in this round's window; else its own runs after + const join = h.core.cfg.join_tail; + if (join and h.finals.items.len > 0) { + did = true; + runFinals(h); + } + if (h.core.activeCount() > 0) { + did = true; + h.core.step() catch |e| { + h.failAll(@errorName(e)); + return true; + }; + } + if (!join and h.finals.items.len > 0) { + did = true; + runFinals(h); + } + const t2 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + var i: usize = 0; + while (i < h.admitted.items.len) { + const job = h.admitted.items[i]; + if (job.started and h.deliver(job)) { + h.lock(); + _ = h.admitted.orderedRemove(i); + h.unlock(); + } else i += 1; + } + rd.vtable.after(rd.ctx, seconds(t1 - t0), seconds(t2 - t1)); + return did; +} + +/// A job whose piece failed: dropped and taken out of this round's finals (`drop` frees the job). +fn dropPiece(h: *LaneHost, job: *Job, message: []const u8) void { + const key = @intFromPtr(job); + var i: usize = 0; + while (i < h.finals.items.len) { + if (h.finals.items[i] == key) _ = h.finals.orderedRemove(i) else i += 1; + } + _ = h.drop(job, message); +} + +/// The prompts whose rows are in: each stream's last row (Engine.finishStream), then it decodes with the rounds. +fn runFinals(h: *LaneHost) void { + if (h.core.cfg.replay_runs) { + // the finals' shared work first (one forward for all); a failure fails each + var ss: [16]*lanes.Stream = undefined; + var jobs: [16]*Job = undefined; + var n: usize = 0; + for (h.finals.items) |key| { + const job: *Job = @ptrFromInt(key); + if (job.stream.finished or n == ss.len) continue; + ss[n] = &job.stream; + jobs[n] = job; + n += 1; + } + h.core.prepareFinals(ss[0..n]) catch |e| { + for (jobs[0..n]) |job| _ = h.drop(job, h.words(e)); + return; + }; + } + if (h.core.cfg.join_tail and h.core.cfg.join_drafts) { + // the round's finals together: the joined prompts' first drafts in one pass (Engine.finishStreams) + var ss: [16]*lanes.Stream = undefined; + var jobs: [16]*Job = undefined; + var errs: [16]?anyerror = undefined; + var at: usize = 0; + while (at < h.finals.items.len) { + var n: usize = 0; + while (at < h.finals.items.len and n < ss.len) : (at += 1) { + const job: *Job = @ptrFromInt(h.finals.items[at]); + if (job.stream.finished) continue; + ss[n] = &job.stream; + jobs[n] = job; + n += 1; + } + const joined = h.core.finishStreams(ss[0..n], errs[0..n]) catch |e| blk: { + for (errs[0..n]) |*x| x.* = e; + break :blk 0; + }; + // a round's line when its finished prompts' first drafts joined in one pass + if (joined > 0) std.log.info("join drafts: {d} of {d} finished prompts drafted in one pass", .{ joined, n }); + for (jobs[0..n], errs[0..n]) |job, x| { + if (x) |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + } + h.prefilled(job, job.began); + } + } + return; + } + for (h.finals.items) |key| { + const job: *Job = @ptrFromInt(key); + if (job.stream.finished) continue; + h.core.finishStream(&job.stream) catch |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + }; + h.prefilled(job, job.began); + } +} + +/// The engine thread under a round planner; an idle round waits for an arrival, a cancel or a short timeout. +pub fn run(h: *LaneHost, rd: Rounds) void { + while (true) { + h.takeCancels(); + h.lock(); + if (h.closing) { + const left = h.queued.items.len + h.admitted.items.len; + h.unlock(); + if (h.lesson_open) { + h.learner.?.abort(h.learner.?.ctx); + h.lesson_open = false; + } + if (left == 0) return; + h.closeAll(); + continue; + } + const work = h.admitted.items.len > 0 or (h.queued.items.len > 0 and h.admitted.items.len < h.info_.lanes); + if (!work and learn.ready(h)) { + h.unlock(); + learn.step(h); + continue; + } + if (!work) { + if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(100), .clock = .awake } }) catch {}; + h.unlock(); + continue; + } + h.unlock(); + const did = round(h, rd); + status.noteLive(h); + if (!did) { + h.lock(); + if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(20), .clock = .awake } }) catch {}; + h.unlock(); + } + } +} + +fn seconds(ns: i96) f64 { + return @as(f64, @floatFromInt(ns)) / 1e9; +} diff --git a/zig/src/core/lane_rounds_test.zig b/zig/src/core/lane_rounds_test.zig new file mode 100644 index 0000000000..f14f685302 --- /dev/null +++ b/zig/src/core/lane_rounds_test.zig @@ -0,0 +1,392 @@ +//! The lane host under a family round planner: prompt pieces, failed piece runs, learning, refusal. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Memory = api.Memory; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; +const Engine = api.Engine; +const FakeLearner = @import("lane_host_test.zig").FakeLearner; +const LearnBox = @import("lane_host_test.zig").LearnBox; + +test "a round planner: prompts in pieces between the other streams' decode rounds, queued admissions, the same tokens" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + // one prompt prefilling at a time (the next waits queued), 3 rows a round, a final when its rows are in + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: u64 = 0, decoding: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + waits: u32 = 0, + rounds: u32 = 0, + mixed: u32 = 0, // rounds that ran a piece while another stream decoded + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + const p = self(ctx); + for (p.seqs.items) |x| if (!x.decoding) { + p.waits += 1; + return .wait; + }; + p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + const p = self(ctx); + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + var decoding = false; + for (p.seqs.items) |x| decoding = decoding or x.decoding; + for (p.seqs.items) |*x| if (!x.decoding) { + if (x.done < x.n - 1) { + const end = @min(x.done + 3, x.n - 1); + try pieces.append(gpa, .{ .key = x.key, .start = x.done, .end = end, .save = false }); + x.done = end; + if (decoding) p.mixed += 1; + } + if (x.done == x.n - 1) { + try finals.append(gpa, x.key); + x.decoding = true; + } + }; + p.rounds += 1; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) !Reason { + for (0..10_000) |_| { // 10 s + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8, 1, 8, 2, 8, 4, 5, 9 }; + var b1: Box = .{}; + var b2: Box = .{}; + defer b1.tokens.deinit(gpa); + defer b2.tokens.deinit(gpa); + const r1: Request = .{ .prompt = &p1, .max_tokens = 40 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both queued before the first round, so the second waits whatever the scheduler does + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + { + errdefer host.stop(); // a timed-out wait stops the host thread before this frame goes + try std.testing.expectEqual(Reason.length, try b1.wait()); + try std.testing.expectEqual(Reason.length, try b2.wait()); + } + host.stop(); + // the replies are the whole-prompt path's tokens + for ([_]struct { p: []const u32, b: *Box }{ .{ .p = &p1, .b = &b1 }, .{ .p = &p2, .b = &b2 } }) |c| { + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, c.p); + for (c.b.tokens.items) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } + } + try std.testing.expectEqual(@as(usize, 40), b1.tokens.items.len); + try std.testing.expectEqual(@as(usize, 12), b2.tokens.items.len); + // the second prompt waited queued, then ran its pieces while the first decoded + try std.testing.expect(pl.waits > 0 and pl.mixed >= 3); + try std.testing.expectEqual(@as(usize, 3 + 4), target.pieces_run); + try std.testing.expectEqual(@as(usize, 0), pl.seqs.items.len); +} + +test "a failed piece run fails its requests alone: their finals of the same round do not run, the host keeps serving" { + // a run's error drops its jobs, whose finals are listed in the same round; runFinals must not read them + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true, .piece_runs = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_runs = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.runBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + var host = LaneHost.init(std.heap.page_allocator, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + try pieces.append(std.heap.page_allocator, .{ .key = x.key, .start = 0, .end = x.n - 1, .save = false }); + try finals.append(std.heap.page_allocator, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both admitted before the host's first round: one run of their pieces, which fails + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.failed, b2.wait()); + try std.testing.expectEqual(@as(usize, 0), target.joins); // no final of theirs ran + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); + try std.testing.expectEqual(@as(usize, 1), target.runs); +} + +test "under a round planner a learn request still steps while no request waits or runs" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + const Idle = struct { + fn none(_: *anyopaque) void {} + fn admit(_: *anyopaque, _: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + return .wait; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn key(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(_: *anyopaque, _: *std.ArrayList(api.Rounds.Piece), _: *std.ArrayList(usize)) anyerror!void {} + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var ctx: u8 = 0; + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + host.rounds = .{ .ctx = &ctx, .vtable = &.{ .begin_admit = Idle.none, .admit = Idle.admit, .admitted = Idle.admitted, .arm = Idle.key, .began = Idle.began, .plan = Idle.plan, .left = Idle.key, .after = Idle.after } }; + var fake: FakeLearner = .{}; + host.learner = .{ .ctx = &fake, .begin = FakeLearner.begin, .step = FakeLearner.step, .abort = FakeLearner.abort }; + try host.start(); + defer host.stop(); + const example: api.Example = .{ .ids = &.{ 1, 2, 3 }, .start = 2 }; + const learn: api.LearnRequest = .{ .train = &.{example} }; + var box: LearnBox = .{}; + defer box.tags.deinit(gpa); + try host.engine().learn(&learn, box.sink()); + box.wait(); + try std.testing.expectEqualSlices(std.meta.Tag(api.LearnEvent), &.{ .learned, .done }, box.tags.items); + try std.testing.expectEqual(@as(u32, 3), box.steps); +} + +test "a round planner over a backend without prompt pieces is refused at start" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var whole = target.backend(); + var vt = whole.vtable.*; + vt.begin = null; + whole.vtable = &vt; + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, whole, clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + var ctx: u8 = 0; + host.rounds = .{ .ctx = &ctx, .vtable = undefined }; + try std.testing.expectError(error.PiecesUnsupported, host.start()); +} + +test "a failed piece fails its request alone: its final of the same round does not run, the host keeps serving" { + // a job dropped for a failed piece must leave the round's finals: `drop` frees it, runFinals must not read it + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_pieces = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + const pa = std.heap.page_allocator; + var host = LaneHost.init(pa, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + try pieces.append(pa, .{ .key = x.key, .start = 0, .end = x.n - 1, .save = false }); + try finals.append(pa, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both admitted before the host's first round: the first piece fails, the second's request is served + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(usize, 12), b2.tokens); + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); +} diff --git a/zig/src/core/lane_stream.zig b/zig/src/core/lane_stream.zig new file mode 100644 index 0000000000..b213336edb --- /dev/null +++ b/zig/src/core/lane_stream.zig @@ -0,0 +1,60 @@ +//! A request's stream on the lane host: its kept prefix looked up, its drafter made, the stream set up. +const std = @import("std"); +const lanes = @import("lanes"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Job = LaneHost.Job; + +/// A job's drafter freed: the family's (LaneHost.proposers) or the built-in one. +pub fn freeProposer(h: *LaneHost, job: *Job) void { + if (job.custom) |c| h.proposers.?.free(h.proposers.?.ctx, c) else job.proposer.deinit(); + job.custom = null; +} + +/// An admitted job's stream (its kept prefix looked up, its drafter): false when it failed (the job is finished). +pub fn open(h: *LaneHost, job: *Job) bool { + const r = job.request; + var reuse: lanes.stream.Reuse = .{}; + // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass + if (h.cache) |store| if (store.lookupRewind(h.gpa, r.prompt, r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { + job.entry = l.entry; + job.kept0 = store.counts.kept; + job.marks = l.marks; + reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; + } else |_| {}; + if (h.proposers) |ps| { + job.custom = ps.make(ps.ctx, h.gpa, r.prompt, r.eos) catch return !h.drop(job, "the drafter could not start"); + } else job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return !h.drop(job, "the drafter could not start"); + job.stream = lanes.Stream.init(h.gpa, .{ + .id = "request", + .prompt = r.prompt, + .max_new = r.max_tokens, + .eos = r.eos, + .sampling = r.sampling, + .drafts = r.drafts, + .proposer = job.custom orelse job.proposer.proposer(), + .stop_check = if (r.stop) |s| .{ .ptr = s.ctx, .check = s.check } else null, + .cancel_check = .{ .ptr = job, .check = LaneHost.cancelled }, + .yield_check = if (r.background and r.max_tokens == 0) .{ .ptr = job, .check = LaneHost.yieldsTo } else null, + .think_budget = r.think_budget, + .think_close = r.think_close, + .think_end = if (r.think_end) |t| t else -1, + .logprobs = r.logprobs, + .loop_guard = r.loop_guard, + .chunks = r.chunks, + .history_len = r.history_len, + .shared_prefixes = r.shared_prefixes, + .reuse = reuse, + .grammar = if (r.structure) |st| h.constraint(st) catch |e| { + freeProposer(h, job); + return !h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + } else null, + .decode_spans = h.spans(r), + .decide = r.decide, + }) catch { + freeProposer(h, job); + return !h.drop(job, "out of memory"); + }; + job.started = true; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 9d8d9d337f..d8d3451f41 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -100,8 +100,28 @@ pub const Backend = struct { held: ?*const fn (ptr: *anyopaque, s: *Stream, out: []u32) anyerror!void = null, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, + /// A prompt in pieces between rounds (`Rounds`) instead of `prefill`: begin, pieces, finish; null: whole only. + begin: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!Begun = null, + piece: ?*const fn (ptr: *anyopaque, s: *Stream, start: u64, end: u64, save: bool) anyerror!void = null, + finish: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!void = null, + /// `finish` minus the last row (Config.join_tail): true when that token is left for the first round. + join: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// A round's pieces in one call (Config.piece_runs), maybe one forward; a failure fails them all. + pieces: ?*const fn (ptr: *anyopaque, list: []const Piece) anyerror!void = null, + /// Before a round's `finish` / `join` calls (Config.replay_runs): their shared work in one go. + finals: ?*const fn (ptr: *anyopaque, streams: []const *Stream) anyerror!void = null, }; + /// One stream's prompt rows [start, end) of a round (`save`: the prompt's snapshot right after them). + pub const Piece = struct { stream: *Stream, start: u64, end: u64, save: bool }; + + /// `begin`'s outcome: restored prefix tokens, or a damaged saved entry dropped (the prompt runs from 0). + pub const Begun = struct { cached: u64 = 0, damaged: bool = false }; + + pub fn pieces(b: Backend) bool { + return b.vtable.begin != null and b.vtable.piece != null and b.vtable.finish != null; + } + pub fn prefill(b: Backend, s: *Stream) !void { return b.vtable.prefill(b.ptr, s); } diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index 85361bd839..c215c87f86 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -32,6 +32,10 @@ pub const Model = struct { draft_probabilities: bool = false, // the head gives each drafted node's chance of landing draft_streams: bool = false, // the head drafts every stream of a shared round in one batch head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them + join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join) + piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them) + join_drafts: bool = false, // with join_tail: a round's joined prompts' first drafts in one call (Engine.finishStreams) + replay_runs: bool = false, // a round's finals go to the backend together first (Backend.finals) lane_costs: []const Cost = &.{}, // a lone stream's window ms by rows, up to its widest tree (timed at load) lone_copies: bool = true, // a lone stream proposes copies (false: only streams in shared rounds do) }; @@ -55,9 +59,18 @@ pub const Config = struct { fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round draft_streams: bool, head_trees: bool, + /// a prompt in pieces leaves its last token as the pending row of its first shared round's window + join_tail: bool = false, + /// a round's pieces of several streams in one backend call (Backend.pieces), which may run them as one forward + piece_runs: bool = false, + /// with join_tail: a round's joined prompts get their first drafts in one backend call + join_drafts: bool = false, + /// a round's finished prompts to the backend together before their finishes (Backend.finals) + replay_runs: bool = false, head_lanes: u32 = 0, // a lone stream's lanes at most, the planner picking each round's (0: the depth rule's chain) head_depth: u32 = 16, // and their depth at most own_prices: bool = true, // the planner prices each pick from the stream's own measured rounds (seeded from the table) + timed: bool = true, // false: rounds priced from the cost tables alone, no wall clock (a request drafts the same whatever the load or the requests before it) depth_prior: []f64, family_costs: Table, shared_costs: Table, @@ -115,6 +128,10 @@ pub const Config = struct { .node_probabilities = m.draft_probabilities, .draft_streams = m.draft_streams, .head_trees = m.head_trees, + .join_tail = m.join_tail, + .piece_runs = m.piece_runs, + .join_drafts = m.join_drafts, + .replay_runs = m.replay_runs, .depth_prior = depth_prior, .family_costs = family_costs, .shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows), diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 4339af40aa..41c8dfcfec 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -13,6 +13,7 @@ const plan_lanes = @import("plan_lanes.zig"); const fill = @import("fill.zig"); const trail = @import("trail.zig"); const open = @import("open.zig"); +const trim = @import("trim.zig"); const Stream = sm.Stream; const Feed = be.Feed; const LogRow = @import("logprob.zig").Row; @@ -26,11 +27,20 @@ const Result = struct { rows: u32, keep: u32, got: []const u32 }; const Unread = struct { event: usize, handle: u64 }; pub const Engine = struct { + const pieces_ = @import("pieces.zig"); + pub const beginStream = pieces_.beginStream; + pub const piece = pieces_.piece; + pub const pieces = pieces_.pieces; + pub const prepareFinals = pieces_.prepareFinals; + pub const finishStream = pieces_.finishStream; + pub const finishStreams = pieces_.finishStreams; + gpa: Allocator, cfg: *const Config, backend: be.Backend, clock: be.Clock, log: ?*ev.Log = null, + trim: ?trim.Trim = null, // a family's row policy in place of the built-in allocation (lanes/trim.zig) rule: depth.Rule, live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`) drafted: u64 = 0, @@ -218,7 +228,7 @@ pub const Engine = struct { s.priced = null; const a = e.arena.allocator(); var plans = [_]Plan{try win.plan(e, s, copied)}; - if (plans[0].kind == .head and e.cfg.node_probabilities) try win.allocate(e, &plans); + if (plans[0].kind == .head and (e.cfg.node_probabilities or e.trim != null)) try win.allocate(e, &plans); try win.constrain(e, &plans[0]); const plan = plans[0]; const early = e.cfg.family_mtp and s.drafts and plan.kind != .forced and e.cfg.speculate_early and plan.parents == null and s.grammar == null; @@ -240,14 +250,14 @@ pub const Engine = struct { if (early) if (e.backend.vtable.unspeculate) |undo| try undo(e.backend.ptr, s); try e.draftLate(s, plan.position, o.follow, o.path, null); } - if (plan.kind == .head or (plan.kind == .none and e.cfg.plain_guard)) { + if (e.cfg.timed and (plan.kind == .head or (plan.kind == .none and e.cfg.plain_guard))) { const ms = e.clock.elapsedMs(.cost); const first = s.rounds == 1; // a stream's first round carries the prefill-to-decode switch try trail.event(e, &.{ f("ev", str("cost")), f("stream", str(s.id)), f("drafts", int(rows - 1)), f("init", .{ .bool = first }), f("ms", ev.bits(ms)) }); try e.rule.observeCost(@intCast(rows - 1), ms, first, win.who(s)); } } - if (s.graft_room != null and s.rounds > 1) { + if (e.cfg.timed and s.graft_room != null and s.rounds > 1) { // the stream's round time beyond its window and head levels, for the lane planner if (e.cfg.lane_costs.get(@intCast(rows))) |w| { const over = e.clock.elapsedMs(.cost) - w - e.cfg.mtp_step_ms * @as(f64, @floatFromInt(held_levels)); @@ -257,7 +267,7 @@ pub const Engine = struct { // the pick that shaped this round, priced from its own measured rounds next time if (priced) |p| s.measured[p.choice].add(e.clock.elapsedMs(.cost) - p.ms); } - if (plan.kind == .head) { + if (e.cfg.timed and plan.kind == .head) { const ms = e.clock.elapsedMs(.overhead); try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(1)), f("rows", int(rows)), f("ms", ev.bits(ms)) }); try e.rule.observeOverhead(1, rows, ms); @@ -314,9 +324,11 @@ pub const Engine = struct { for (heads.items, budgets) |i, d| try e.draftLate(plans[i].stream, plans[i].position, outcomes[i].follow, outcomes[i].path, d); } try e.backend.keep(windows, paths); - const ms = e.clock.elapsedMs(.overhead); - try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); - try e.rule.observeOverhead(@intCast(plans.len), total, ms); + if (e.cfg.timed) { + const ms = e.clock.elapsedMs(.overhead); + try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); + try e.rule.observeOverhead(@intCast(plans.len), total, ms); + } } /// Commit the target-sampled path (after the thinking budget's cut), each token with its path row's logprob row. @@ -334,6 +346,7 @@ pub const Engine = struct { path = try accept.acceptPath(a, tokens, parents, sampled); for (path[1..]) |r| try committed.append(a, tokens[r]); try committed.append(a, sampled[path[path.len - 1]]); + if (s.proposer) |pr| pr.round(@intCast(rows), @intCast(path.len - 1)); if (rows > 1) { const kept = path.len - 1; e.drafted += rows - 1; @@ -371,6 +384,10 @@ pub const Engine = struct { } } } + if (e.trim) |t| if (p.kind != .forced) { + const chain = parents.len == 0 or accept.isChain(parents); + try t.commit(.{ .stream = s, .drafted = p.kind == .head, .rows = @intCast(rows), .parents = if (chain) null else parents, .path = path, .bonus = sampled[path[path.len - 1]] }); + }; const budget_cut = s.thinkCut(committed.items); const loop_cut = sm.loopCut(s, committed.items); const loop_wins = if (loop_cut) |c| budget_cut == null or c < budget_cut.? else false; @@ -427,7 +444,7 @@ pub const Engine = struct { const s = requests[0].stream; if (s.odds == null) s.odds = shape.Odds.init(shape.start_odds, shape.start_weight); const cap: usize = @intCast(@max(1, @min(@min(shape.max_depth, e.cfg.head_depth), s.draftRoom() - 1))); - const costs = plan_lanes.Costs{ .window = &e.cfg.lane_costs, .head_ms = e.cfg.mtp_step_ms, .over_ms = s.round_over orelse 0, .tree_over_ms = s.tree_over, .measured = if (e.cfg.own_prices) &s.measured else null }; + const costs = plan_lanes.Costs{ .window = &e.cfg.lane_costs, .head_ms = e.cfg.mtp_step_ms, .over_ms = s.round_over orelse 0, .tree_over_ms = s.tree_over, .measured = if (e.cfg.own_prices and e.cfg.timed) &s.measured else null }; // every eighth round one size wider than the best, so the odds of deeper lanes stay measured const p = try plan_lanes.pick(e.gpa, s.odds.?, costs, e.cfg.head_lanes, cap, try e.graftProspect(s), s.rounds % 8 == 4); s.picks[p.choice] += 1; diff --git a/zig/src/core/lanes/engine_test.zig b/zig/src/core/lanes/engine_test.zig index 2100cc3762..5c9043d336 100644 --- a/zig/src/core/lanes/engine_test.zig +++ b/zig/src/core/lanes/engine_test.zig @@ -460,3 +460,124 @@ test "a family that keeps copies to shared rounds proposes none for a lone strea try std.testing.expectEqual(lone_copies, (try win.copyProposal(&engine, &s)).len >= 2); } } + +/// The cases' prompts in pieces, case i opened before round 2 i, others decoding; `join`: Config.join_tail. +fn runPieces(cases: []const Case, join: bool, joins: *usize) ![][]u32 { + var cfg = try model(); + defer cfg.deinit(gpa); + cfg.join_tail = join; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(sm.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |c, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .think_budget = c.think_budget, .think_close = &.{ 90, 91, 92 }, .think_end = 91 }); + defer for (streams) |*s| s.deinit(gpa); + var opened: usize = 0; + var round: usize = 0; + while (opened < streams.len or engine.activeCount() > 0) : (round += 1) { + if (opened < streams.len and round == 2 * opened) { + const s = &streams[opened]; + _ = try engine.beginStream(s); + if (s.prompt_len > 1) try engine.piece(s, 0, s.prompt_len - 1, false); + try engine.finishStream(s); + opened += 1; + } + if (engine.activeCount() > 0) try engine.step(); + } + joins.* = target.joins; + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + return out; +} + +test "a prompt's last token joins its first round's shared window: the same tokens, greedy and sampled" { + const cases = [_]Case{ + .{ .prompt = &.{ 3, 1, 4, 1, 5, 9, 2, 6 }, .max_new = 37 }, + .{ .prompt = &.{ 2, 7, 1, 8, 2, 8 }, .max_new = 29, .sampling = .{ .seed = 5, .temperature = 0.8, .top_k = 7 } }, + .{ .prompt = &.{4}, .max_new = 21 }, + .{ .prompt = &.{ 1, 1, 2, 3, 5 }, .max_new = 25, .drafts = false }, + .{ .prompt = &.{ 6, 6, 6, 1 }, .max_new = 30, .think_budget = 1 }, + }; + var j0: usize = 0; + var j1: usize = 0; + const off = try runPieces(&cases, false, &j0); + defer free(off); + const on = try runPieces(&cases, true, &j1); + defer free(on); + try std.testing.expectEqual(@as(usize, 0), j0); + try std.testing.expectEqual(@as(usize, 3), j1); // not the undrafted stream nor the budget cutting the first token + for (off, on) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // and the whole prompts' tokens + const whole = try run(&cases); + defer free(whole); + for (whole, on) |a, b| try std.testing.expectEqualSlices(u32, a, b); +} + +/// Every prompt in one piece, all finished in one round (`together`: Engine.finishStreams), then rounds until done. +fn runFinals(cases: []const Case, together: bool, calls: *usize, most: *usize) ![][]u32 { + var cfg = try model(); + defer cfg.deinit(gpa); + cfg.join_tail = true; + cfg.join_drafts = together; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(sm.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |c, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .think_budget = c.think_budget, .think_close = &.{ 90, 91, 92 }, .think_end = 91 }); + defer for (streams) |*s| s.deinit(gpa); + const ptrs = try gpa.alloc(*sm.Stream, cases.len); + defer gpa.free(ptrs); + for (streams, ptrs) |*s, *p| { + _ = try engine.beginStream(s); + if (s.prompt_len > 1) try engine.piece(s, 0, s.prompt_len - 1, false); + p.* = s; + } + // the finish phase alone: its draft calls and their largest + target.drafts_calls = 0; + target.drafts_most = 0; + if (together) { + const errs = try gpa.alloc(?anyerror, cases.len); + defer gpa.free(errs); + _ = try engine.finishStreams(ptrs, errs); + for (errs) |x| if (x) |e| return e; + } else for (ptrs) |s| try engine.finishStream(s); + calls.* = target.drafts_calls; + most.* = target.drafts_most; + while (engine.activeCount() > 0) try engine.step(); + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + return out; +} + +test "a round's finished prompts get their first drafts in one call (Python's _windows), the same tokens" { + const cases = [_]Case{ + .{ .prompt = &.{ 3, 1, 4, 1, 5, 9, 2, 6 }, .max_new = 37 }, + .{ .prompt = &.{ 2, 7, 1, 8, 2, 8 }, .max_new = 29, .sampling = .{ .seed = 5, .temperature = 0.8, .top_k = 7 } }, + .{ .prompt = &.{4}, .max_new = 21 }, + .{ .prompt = &.{ 1, 1, 2, 3, 5 }, .max_new = 25, .drafts = false }, + .{ .prompt = &.{ 6, 6, 6, 1 }, .max_new = 30, .think_budget = 1 }, + }; + var c1: usize = 0; + var m1: usize = 0; + var c0: usize = 0; + var m0: usize = 0; + const together = try runFinals(&cases, true, &c1, &m1); + defer free(together); + const apart = try runFinals(&cases, false, &c0, &m0); + defer free(apart); + for (apart, together) |a, b| try std.testing.expectEqualSlices(u32, a, b); + const whole = try run(&cases); + defer free(whole); + for (whole, together) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // three joined prompts in one call, the budget-cut one its own, the undrafted none: 2 calls against 4 + try std.testing.expectEqual(@as(usize, 2), c1); + try std.testing.expectEqual(@as(usize, 3), m1); + try std.testing.expectEqual(@as(usize, 4), c0); + try std.testing.expectEqual(@as(usize, 1), m0); +} diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index dd60ca7f64..50290a9702 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -53,11 +53,18 @@ pub const Fake = struct { answer_cycles: bool = false, prefill_chunks: usize = 0, prefill_count: usize = 0, + fail_pieces: usize = 0, // the next this many `piece` calls fail with error.Unbound prefill_hook: ?*const fn (ctx: *anyopaque, s: *Stream, chunk: usize) void = null, prefill_hook_ctx: ?*anyopaque = null, refuse_sampled: bool = false, // prefill refuses a sampled stream with error.SamplingRefused ranked: bool = false, // the head gives its held chain back (alternatives) memory_lanes: ?usize = null, // prefill past this many live lanes fails with error.LaneMemory (no room yet) + pieces_run: usize = 0, // prompt pieces run (`pieceBackend`) + joins: usize = 0, // prompts whose last token joined a round (`joinBackend`) + drafts_calls: usize = 0, // `draft` calls, and the most requests one took + drafts_most: usize = 0, + fail_runs: usize = 0, // the next this many `pieces` calls fail with error.Unbound (`runBackend`) + runs: usize = 0, // `pieces` calls that ran pub fn deinit(x: *Fake) void { var it = x.lanes.valueIterator(); @@ -81,6 +88,66 @@ pub const Fake = struct { return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow } }; } + /// The backend with prompts in pieces (begin / piece / finish) beside whole prompts. + pub fn pieceBackend(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish } }; + } + + /// `pieceBackend` whose prompts may leave their last token to the stream's first round (Config.join_tail). + pub fn joinBackend(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish, .join = join } }; + } + + /// `joinBackend` that takes a round's pieces in one call (Config.piece_runs): a piece at a time, or a failure. + pub fn runBackend(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish, .join = join, .pieces = pieces } }; + } + + fn pieces(ptr: *anyopaque, list: []const be.Backend.Piece) anyerror!void { + const x = self(ptr); + if (x.fail_runs > 0) { + x.fail_runs -= 1; + return error.Unbound; + } + x.runs += 1; + for (list) |p| try piece(ptr, p.stream, p.start, p.end, p.save); + } + + fn join(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + if (x.lane(s).history.items.len + 1 != s.prompt_len) return error.PositionMismatch; + x.joins += 1; + return true; + } + + fn begin(ptr: *anyopaque, s: *Stream) anyerror!be.Backend.Begun { + const x = self(ptr); + const got = try x.lanes.getOrPut(x.gpa, s); + if (got.found_existing) freeLane(x.gpa, got.value_ptr); + got.value_ptr.* = .{}; + return .{}; + } + + fn piece(ptr: *anyopaque, s: *Stream, start: u64, end: u64, snapshot: bool) anyerror!void { + _ = snapshot; + const x = self(ptr); + if (x.fail_pieces > 0) { + x.fail_pieces -= 1; + return error.Unbound; + } + const h = &x.lane(s).history; + if (h.items.len != start or end + 1 > s.prompt_len) return error.PositionMismatch; + try h.appendSlice(x.gpa, s.prompt()[start..end]); + x.pieces_run += 1; + } + + fn finish(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const h = &x.lane(s).history; + if (h.items.len + 1 != s.prompt_len) return error.PositionMismatch; + try h.append(x.gpa, s.prompt()[s.prompt_len - 1]); + } + fn self(ptr: *anyopaque) *Fake { return @ptrCast(@alignCast(ptr)); } @@ -285,6 +352,8 @@ pub const Fake = struct { fn draft(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { const x = self(ptr); + x.drafts_calls += 1; + x.drafts_most = @max(x.drafts_most, requests.len); for (requests) |r| { const l = x.lane(r.stream); // the head reads the last verify's kept rows (a shared round drafts before its rollback) diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 910f501675..8a1f7213e9 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -14,6 +14,7 @@ pub const cost_rule = @import("cost_rule.zig"); pub const gpu_rule = @import("gpu_rule.zig"); pub const gpu_full = @import("gpu_full.zig"); pub const depth = @import("depth.zig"); +pub const trim = @import("trim.zig"); pub const config = @import("config.zig"); pub const stream = @import("stream.zig"); pub const backend = @import("backend.zig"); @@ -34,6 +35,7 @@ pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; pub const SuffixLookup = proposer.SuffixLookup; pub const LogprobRow = logprob.Row; +pub const Trim = trim.Trim; test { std.testing.refAllDecls(@This()); diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..cd832256ab 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -23,8 +23,16 @@ pub fn add(e: *Engine, s: *Stream) !void { try opened(e, s); } +/// How a stream's first drafts are sized: the depth rule, or the head's whole reach under node probabilities. +pub const Reach = enum { rule, head }; + /// After the prompt pass: the first token drawn and committed, the first drafts asked; rounds from the next. pub fn opened(e: *Engine, s: *Stream) !void { + return openedWith(e, s, .rule); +} + +/// `opened` with its first drafts sized by `reach` (a family planner's prompts in pieces: `.head`). +pub fn openedWith(e: *Engine, s: *Stream, reach: Reach) !void { if (s.isCancelled()) { // cancelled in its last chunk: no first token e.backend.release(s); return error.Cancelled; @@ -40,7 +48,7 @@ pub fn opened(e: *Engine, s: *Stream) !void { var asked: ?u32 = null; if (e.cfg.family_mtp and s.drafts) { // the head reads the prompt's last row and the first token, and drafts the one after it - const d: u32 = @intCast(try e.rule.depth(win.who(s))); + const d = try firstDepth(e, s, reach); asked = d; try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = feed, .rows = null, .start = s.prompt_len, .position = position + 1, .depth = d }}); s.dropHeld(e.gpa); @@ -70,6 +78,13 @@ pub fn opened(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } +/// The first drafts' depth for `reach`. +pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 { + var who = win.who(s); + if (reach == .head and e.cfg.node_probabilities) return @intCast(try e.rule.headDepth(&who, null, win.probe(e))); + return @intCast(try e.rule.depth(who)); +} + /// Whether the backend can run a prompt pass a chunk at a time between rounds (`fillStream`). pub fn fills(e: *const Engine) bool { return e.backend.vtable.prefill_step != null; diff --git a/zig/src/core/lanes/pieces.zig b/zig/src/core/lanes/pieces.zig new file mode 100644 index 0000000000..cd64bb1eb0 --- /dev/null +++ b/zig/src/core/lanes/pieces.zig @@ -0,0 +1,144 @@ +//! Prompts in pieces on the lane core: a stream begun, its prompt rows run in pieces, then its first token. +const std = @import("std"); +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const ev = @import("events.zig"); +const trail = @import("trail.zig"); +const open = @import("open.zig"); +const Engine = @import("engine.zig").Engine; +const Stream = sm.Stream; +const f = ev.f; +const str = trail.str; +const int = trail.int; + +/// A stream whose prompt runs in pieces between rounds: `piece`s, then `finishStream`; a failed begin holds nothing. +pub fn beginStream(e: *Engine, s: *Stream) !be.Backend.Begun { + try trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }); + const begin = e.backend.vtable.begin orelse return error.NoPieces; + const b = try begin(e.backend.ptr, s); + s.cached = @intCast(b.cached); + return b; +} + +/// Prompt rows [start, end) of a begun stream (`save`: the prompt's snapshot right after them). +pub fn piece(e: *Engine, s: *Stream, start: u64, end: u64, save: bool) !void { + const run = e.backend.vtable.piece orelse return error.NoPieces; + return run(e.backend.ptr, s, start, end, save); +} + +/// A round's pieces of several streams in one backend call (Config.piece_runs, Backend.pieces); else one at a time. +pub fn pieces(e: *Engine, list: []const be.Backend.Piece) !void { + if (e.cfg.piece_runs) if (e.backend.vtable.pieces) |run| return run(e.backend.ptr, list); + for (list) |p| try piece(e, p.stream, p.start, p.end, p.save); +} + +/// A round's finished prompts before their `finishStream` calls (Config.replay_runs, Backend.finals). +pub fn prepareFinals(e: *Engine, streams: []const *Stream) !void { + if (!e.cfg.replay_runs or streams.len < 2) return; + const run = e.backend.vtable.finals orelse return; + return run(e.backend.ptr, streams); +} + +/// A begun stream's prompt is in: first token and drafts; a failure is the host's to discard (Cancelled: released). +pub fn finishStream(e: *Engine, s: *Stream) !void { + _ = e.arena.reset(.retain_capacity); + // Config.join_tail: the last token waits for the next round (a forced or budget-cut first token keeps its row) + if (e.cfg.join_tail and e.cfg.family_mtp and s.drafts and s.prompt_len > 0 and s.force.items.len == 0 and !s.cutsNext()) { + if (e.backend.vtable.join) |join| { + if (join(e.backend.ptr, s) catch |x| return released(e, s, x)) return joined(e, s); + return open.openedWith(e, s, .head); + } + } + const fin = e.backend.vtable.finish orelse return error.NoPieces; + fin(e.backend.ptr, s) catch |x| return released(e, s, x); + return open.openedWith(e, s, .head); +} + +/// A round's finished prompts as `finishStream`, joined ones' first drafts in one call; returns how many joined. +pub fn finishStreams(e: *Engine, streams: []const *Stream, errs: []?anyerror) !usize { + _ = e.arena.reset(.retain_capacity); + const a = e.arena.allocator(); + var reqs: std.ArrayList(be.DraftRequest) = .empty; + var joins: std.ArrayList(usize) = .empty; + for (streams, errs, 0..) |s, *err, i| { + err.* = null; + const can = e.cfg.join_tail and e.cfg.family_mtp and s.drafts and s.prompt_len > 0 and s.force.items.len == 0 and !s.cutsNext(); + const join = if (can) e.backend.vtable.join else null; + if (join) |j| { + const ok = j(e.backend.ptr, s) catch |x| { + err.* = released(e, s, x); + continue; + }; + if (ok) { + try reqs.append(a, joinRequest(e, s) catch |x| { + err.* = x; + continue; + }); + try joins.append(a, i); + continue; + } + } else { + const fin = e.backend.vtable.finish orelse return error.NoPieces; + fin(e.backend.ptr, s) catch |x| { + err.* = released(e, s, x); + continue; + }; + } + open.openedWith(e, s, .head) catch |x| { + err.* = x; + }; + } + if (reqs.items.len == 0) return 0; + e.backend.draft(reqs.items) catch |x| { + for (joins.items) |i| errs[i] = x; + return 0; + }; + for (reqs.items) |r| joinedDrafted(e, r) catch |x| { + for (joins.items) |i| if (streams[i] == r.stream) { + errs[i] = x; + }; + }; + return reqs.items.len; +} + +/// A cancelled stream's lane is released here: the host finishes it without the core. +fn released(e: *Engine, s: *Stream, err: anyerror) anyerror { + if (err == error.Cancelled) e.backend.release(s); + return err; +} + +/// A joined prompt (Backend.join): its last token pending, drafts asked now, verified in its first shared round. +fn joined(e: *Engine, s: *Stream) !void { + s.context.shrinkRetainingCapacity(s.prompt_len); + const last = s.context.items[s.prompt_len - 1]; + s.cache_len = s.prompt_len - 1; + s.pending = last; + const position: u64 = s.prompt_len; + const d = try open.firstDepth(e, s, .head); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = last }, .rows = null, .start = s.cache_len, .position = position, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(d)), f("position", int(position)), f("follow", .{ .u32s = &.{last} }), f("rows", .null) }); + try trail.resolve(e); + try e.live.append(e.gpa, s); +} + +/// After its first draft: the joined stream holds it and takes part from the next round. +fn joinedDrafted(e: *Engine, r: be.DraftRequest) !void { + const s = r.stream; + s.dropHeld(e.gpa); + s.next = .{ .count = r.depth }; + try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(r.depth)), f("position", int(r.position)), f("follow", .{ .u32s = &.{s.pending.?} }), f("rows", .null) }); + try trail.resolve(e); + try e.live.append(e.gpa, s); +} + +/// A joined prompt's state and its first draft request (the last token pending at the prompt's last position). +fn joinRequest(e: *Engine, s: *Stream) !be.DraftRequest { + s.context.shrinkRetainingCapacity(s.prompt_len); + const last = s.context.items[s.prompt_len - 1]; + s.cache_len = s.prompt_len - 1; + s.pending = last; + const d = try open.firstDepth(e, s, .head); + return .{ .stream = s, .follow = &.{}, .first = .{ .value = last }, .rows = null, .start = s.cache_len, .position = s.prompt_len, .depth = d }; +} diff --git a/zig/src/core/lanes/proposer.zig b/zig/src/core/lanes/proposer.zig index f49626e74a..3df9a73393 100644 --- a/zig/src/core/lanes/proposer.zig +++ b/zig/src/core/lanes/proposer.zig @@ -11,6 +11,10 @@ pub const Proposer = struct { propose: *const fn (ptr: *anyopaque, context: []const u32, max_draft: i64) anyerror![]const u32, last_match: *const fn (ptr: *anyopaque) i64, observe: ?*const fn (ptr: *anyopaque, proposed: i64, accepted: i64) void = null, + /// The proposer prices its own drafts: `propose`'s result is taken as is, no width, length or match gate. + priced: bool = false, + /// Every non-forced round of the stream, whatever drafted it: the window's rows and its drafts kept. + round: ?*const fn (ptr: *anyopaque, rows: u32, kept: u32) void = null, }; /// Tokens valid until the next `propose`. @@ -25,6 +29,10 @@ pub const Proposer = struct { pub fn observe(p: Proposer, proposed: i64, accepted: i64) void { if (p.vtable.observe) |f| f(p.ptr, proposed, accepted); } + + pub fn round(p: Proposer, rows: u32, kept: u32) void { + if (p.vtable.round) |f| f(p.ptr, rows, kept); + } }; pub const max_ngram = 8; diff --git a/zig/src/core/lanes/trim.zig b/zig/src/core/lanes/trim.zig new file mode 100644 index 0000000000..426577030f --- /dev/null +++ b/zig/src/core/lanes/trim.zig @@ -0,0 +1,72 @@ +//! An optional row policy in place of windows.allocate: a family's own depth rule trims drafts. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const Stream = @import("stream.zig").Stream; + +/// One planned window as the policy sees it. +pub const Window = struct { + stream: *Stream, + drafted: bool, // head drafts the policy may trim; every other window runs whole + rows: u32, // rows as planned, the pending row included + parents: ?[]const i32, // the drafts' parents for a tree (null: a chain) +}; + +/// What a window keeps: its first `count` drafts, or the listed drafts (ascending, each one's parent listed). +pub const Choice = struct { count: u32, nodes: ?[]const u32 = null }; + +/// A committed window: its rows, the accepted path (row 0 first) and the target's choice after the path's last row. +pub const Commit = struct { + stream: *Stream, + drafted: bool, + rows: u32, + parents: ?[]const i32, // each row's parent row (row 0: -1) for a tree; null for a chain + path: []const u32, + bonus: u32, +}; + +pub const Trim = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// Each window's drafts for this round; `alone`: the round's only window. Slices in `out` live in `arena`. + choose: *const fn (ptr: *anyopaque, windows: []const Window, alone: bool, arena: Allocator, out: []Choice) anyerror!void, + /// Every window's commit, in plan order (a shared round's windows one after the other). + commit: *const fn (ptr: *anyopaque, c: Commit) anyerror!void, + }; + + pub fn choose(t: Trim, windows: []const Window, alone: bool, arena: Allocator, out: []Choice) !void { + return t.vtable.choose(t.ptr, windows, alone, arena, out); + } + pub fn commit(t: Trim, c: Commit) !void { + return t.vtable.commit(t.ptr, c); + } +}; + +pub const Sub = struct { tokens: []u32, parents: []i32 }; + +/// The listed drafts of a tree re-indexed among themselves (parents -1: the pending row); arena-owned. +pub fn select(arena: Allocator, tokens: []const u32, parents: []const i32, nodes: []const u32) !Sub { + const at = try arena.alloc(i32, tokens.len); + @memset(at, -1); + const out_t = try arena.alloc(u32, nodes.len); + const out_p = try arena.alloc(i32, nodes.len); + for (nodes, out_t, out_p, 0..) |n, *t, *p, i| { + const q = parents[n]; + if (q >= 0 and at[@intCast(q)] < 0) return error.ParentNotKept; + at[n] = @intCast(i); + t.* = tokens[n]; + p.* = if (q < 0) -1 else at[@intCast(q)]; + } + return .{ .tokens = out_t, .parents = out_p }; +} + +test "select keeps listed drafts with their parents re-indexed" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + // chain 10 11 12, sibling 20 (on the pending row) with child 21 + const s = try select(arena.allocator(), &.{ 10, 11, 12, 20, 21 }, &.{ -1, 0, 1, -1, 3 }, &.{ 0, 3, 4 }); + try std.testing.expectEqualSlices(u32, &.{ 10, 20, 21 }, s.tokens); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 1 }, s.parents); + try std.testing.expectError(error.ParentNotKept, select(arena.allocator(), &.{ 10, 11 }, &.{ -1, 0 }, &.{1})); +} diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index 901cae7b8d..20e401e05e 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -4,6 +4,7 @@ const accept = @import("accept.zig"); const alloc = @import("allocate.zig"); const be = @import("backend.zig"); const depth = @import("depth.zig"); +const trim = @import("trim.zig"); const fill = @import("fill.zig"); const shape = @import("shape.zig"); const ev = @import("events.zig"); @@ -137,6 +138,7 @@ pub fn copyProposal(e: *Engine, s: *Stream) ![]const u32 { const p = s.proposer orelse return &.{}; if (e.cfg.max_copy == 0 or s.force.items.len > 0) return &.{}; if (e.alone and !e.cfg.lone_copies) return &.{}; + if (p.vtable.priced) return p.propose(s.context.items, @min(@as(i64, e.cfg.max_copy), s.draftRoom() - 1)) catch &.{}; const width: u32 = if (e.alone) (s.copy_width orelse e.cfg.first_copy) else e.cfg.first_copy; const copied = p.propose(s.context.items, @min(@as(i64, width), s.draftRoom() - 1)) catch return &.{}; if (copied.len < 2 or p.lastMatch() < e.cfg.enter_match) return &.{}; @@ -160,6 +162,7 @@ pub fn who(s: *Stream) depth.Who { /// Trim draft prefixes by landing probability to the shared width with the most expected tokens a ms. pub fn allocate(e: *Engine, plans: []Plan) !void { + if (e.trim) |t| return trimmed(e, plans, t); const a = e.arena.allocator(); const fixed = try a.alloc(u32, plans.len); const probs = try a.alloc([]const f64, plans.len); @@ -203,6 +206,37 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { try trail.event(e, &.{ f("ev", trail.str("alloc")), f("plans", .{ .list = rows }) }); } +/// A family's row policy decides each head window's drafts (lanes/trim.zig); other windows run whole. +fn trimmed(e: *Engine, plans: []Plan, t: trim.Trim) !void { + const a = e.arena.allocator(); + const asks = try a.alloc(trim.Window, plans.len); + for (asks, plans) |*w, p| w.* = .{ .stream = p.stream, .drafted = p.kind == .head, .rows = @intCast(1 + p.count()), .parents = p.parents }; + const out = try a.alloc(trim.Choice, plans.len); + for (out, plans) |*o, p| o.* = .{ .count = @intCast(p.count()) }; + try t.choose(asks, e.alone and plans.len == 1, a, out); + for (plans, out) |*p, c| { + if (p.kind != .head) continue; + if (c.nodes) |nodes| { + if (p.parents == null or p.held > 0) return error.TrimNodesNeedTree; + const sub = try trim.select(a, p.tokens, p.parents.?, nodes); + p.tokens = sub.tokens; + p.parents = if (accept.isChain(try accept.rowParents(a, 1 + sub.tokens.len, sub.parents))) null else sub.parents; + } else if (c.count == 0) { + p.kind = .none; + p.held = 0; + p.tokens = &.{}; + p.parents = null; + } else if (c.count < p.count()) { + if (p.parents) |parents| { + const tree = try accept.sanitizeTree(a, p.tokens, parents, c.count); + p.tokens = tree.tokens; + p.parents = tree.parents; + } else if (p.held > 0) p.held = c.count else p.tokens = p.tokens[0..c.count]; + } + if (p.count() == 0) p.kind = .none; + } +} + /// The head's chance for each held draft, else chain chances from the stream's per-depth acceptance. pub fn nodeChances(e: *Engine, s: *Stream, n: usize) ![]const f64 { const a = e.arena.allocator(); diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index a608fd6114..dd1bbe1033 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -39,6 +39,8 @@ pub const DeviceAttribute = enum(c_int) { compute_capability_minor = 76, max_shared_memory_per_multiprocessor = 81, max_shared_memory_per_block_optin = 97, + virtual_memory_management_supported = 102, + gpu_direct_rdma_with_cuda_vmm_supported = 110, }; pub const FunctionAttribute = enum(c_int) { @@ -220,3 +222,10 @@ comptime { std.debug.assert(@sizeOf(KernelNodeParams) == 72 and @offsetOf(KernelNodeParams, "params") == 40); std.debug.assert(@sizeOf(ExecUpdateResultInfo) == 24); } + +test "CUDA device attribute matches the RDMA VMM capability ABI" { + // CUDA's CUdevice_attribute assigns 110 to RDMA VMM support; 126 is the memory sync domain count. + try @import("std").testing.expectEqual(@as(c_int, 110), @backingInt( + DeviceAttribute.gpu_direct_rdma_with_cuda_vmm_supported, + )); +} diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig new file mode 100644 index 0000000000..3886c81403 --- /dev/null +++ b/zig/src/cuda/arena.zig @@ -0,0 +1,374 @@ +//! A device sub-allocator under DeviceBuffer.alloc: buffers carved from a few large chunks, not one cuMemAlloc each. + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); + +pub const Error = error{ DriverUnavailable, MissingSymbol, Unsupported, CudaFailed, OutOfMemory }; + +/// Where chunks come from: `chunk` maps `len` bytes (a multiple of `granularity`), `release` gives a chunk back. +pub const Backend = struct { + ctx: *anyopaque, + granularity: usize, + chunk: *const fn (ctx: *anyopaque, len: usize) ?u64, + release: *const fn (ctx: *anyopaque, base: u64, len: usize) void, + name: []const u8, +}; + +pub const Options = struct { + small_max: usize = 4 << 20, // requests up to this many bytes share chunks + small_chunk: usize = 64 << 20, // a shared chunk's bytes, rounded up to the granularity + alignment: usize = 512, // every range starts at a multiple of this +}; + +pub const Stats = struct { small: u64 = 0, large: u64 = 0, reused: u64 = 0, fallback: u64 = 0, chunk_bytes: u64 = 0 }; + +/// Small requests bump 512-byte aligned through shared chunks and are reused by size; large ones get their own chunk. +pub const Arena = struct { + gpa: std.mem.Allocator, + be: Backend, + o: Options, + mutex: std.atomic.Mutex = .unlocked, + small: std.ArrayList([2]u64) = .empty, // shared chunks (base, len) + cursor: u64 = 0, // the current shared chunk's next free byte + limit: u64 = 0, + free_small: std.AutoHashMapUnmanaged(u64, std.ArrayList(u64)) = .empty, // freed small ranges by rounded size + live: std.AutoHashMapUnmanaged(u64, u64) = .empty, // live small ranges: address to rounded size + large: std.AutoHashMapUnmanaged(u64, u64) = .empty, // large chunks: address to mapped size + dead: [2]u32 = .{ 0, 0 }, // removals from `live` and `large` since each was last rehashed + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, be: Backend, o: Options) Arena { + return .{ .gpa = gpa, .be = be, .o = o }; + } + + pub fn deinit(a: *Arena) void { + var it = a.large.iterator(); + while (it.next()) |e| a.be.release(a.be.ctx, e.key_ptr.*, e.value_ptr.*); + for (a.small.items) |c| a.be.release(a.be.ctx, c[0], c[1]); + var f = a.free_small.valueIterator(); + while (f.next()) |l| l.deinit(a.gpa); + a.free_small.deinit(a.gpa); + a.live.deinit(a.gpa); + a.large.deinit(a.gpa); + a.small.deinit(a.gpa); + a.* = undefined; + } + + fn lock(a: *Arena) void { + while (!a.mutex.tryLock()) std.Thread.yield() catch {}; + } + + fn roundUp(x: usize, m: usize) usize { + return std.mem.alignForward(usize, x, m); + } + + /// An address for `len` bytes, or null (the caller then allocates as before). + pub fn alloc(a: *Arena, len: usize) ?u64 { + if (len == 0) return null; + a.lock(); + defer a.mutex.unlock(); + return a.allocLocked(len) catch null; + } + + fn allocLocked(a: *Arena, len: usize) !?u64 { + if (len > a.o.small_max) { + const n = roundUp(len, a.be.granularity); + const p = a.be.chunk(a.be.ctx, n) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, p, n); + try a.large.put(a.gpa, p, n); + a.stats.large += 1; + a.stats.chunk_bytes += n; + return p; + } + const n: u64 = roundUp(len, a.o.alignment); + if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { + try a.live.put(a.gpa, p, n); + a.stats.reused += 1; + return p; + }; + if (a.limit - a.cursor < n) { + const cl = roundUp(@max(a.o.small_chunk, n), a.be.granularity); + const base = a.be.chunk(a.be.ctx, cl) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, base, cl); + try a.small.append(a.gpa, .{ base, cl }); + a.cursor = base; + a.limit = base + cl; + a.stats.chunk_bytes += cl; + } + const p = a.cursor; + try a.live.put(a.gpa, p, n); + a.cursor += n; + a.stats.small += 1; + return p; + } + + /// True when `p` was the arena's (and is now free); false: not ours, the caller frees it as before. + pub fn free(a: *Arena, p: u64) bool { + a.lock(); + defer a.mutex.unlock(); + if (a.large.fetchRemove(p)) |e| { + removed(&a.large, &a.dead[1]); + a.be.release(a.be.ctx, p, e.value); + return true; + } + const e = a.live.fetchRemove(p) orelse return false; + removed(&a.live, &a.dead[0]); + const g = a.free_small.getOrPut(a.gpa, e.value) catch return true; // leaked range, still ours + if (!g.found_existing) g.value_ptr.* = .empty; + g.value_ptr.append(a.gpa, p) catch {}; + return true; + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn removed(map: *std.AutoHashMapUnmanaged(u64, u64), dead: *u32) void { + dead.* += 1; + if (@as(u64, dead.*) * 10 < map.capacity()) return; + map.rehash(std.hash_map.AutoContext(u64){}); + dead.* = 0; + } + + // -- memory.zig's hook ---------------------------------------------------------------------------------------- + + fn hookAlloc(ctx: *anyopaque, len: usize) ?u64 { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.alloc(len); + } + + fn hookFree(ctx: *anyopaque, p: u64) bool { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.free(p); + } + + /// Every later DeviceBuffer.alloc comes from the arena, which must outlive its buffers; only addresses change. + pub fn install(a: *Arena) void { + memory.setHook(.{ .ctx = a, .alloc = hookAlloc, .free = hookFree }); + } +}; + +// -- backends --------------------------------------------------------------------------------------------------------- + +/// One cuMemAlloc a chunk (the driver's own 2 MiB-aligned allocations). +pub const Plain = struct { + d: *const Driver, + + pub fn backend(p: *Plain) Backend { + return .{ .ctx = p, .granularity = 2 << 20, .chunk = chunk, .release = release, .name = "plain" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const p: *Plain = @ptrCast(@alignCast(ctx)); + var q: abi.DevicePtr = 0; + if (p.d.api.cuMemAlloc_v2(&q, len) != abi.success) return null; + return q; + } + + fn release(ctx: *anyopaque, base: u64, _: usize) void { + const p: *Plain = @ptrCast(@alignCast(ctx)); + _ = p.d.api.cuMemFree_v2(base); + } +}; + +/// CUDA's virtual memory management: one reserved virtual range, chunks of the recommended granularity mapped in order. +pub const Vmm = struct { + lib: std.DynLib, + device: abi.Device, + prop: Prop, + granularity: usize, + va: u64 = 0, + va_len: usize = 0, + next: u64 = 0, + api: Api, + + const Api = struct { + reserve: *const fn (*u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + addressFree: *const fn (u64, usize) callconv(.c) abi.Result, + create: *const fn (*u64, usize, *const Prop, c_ulonglong) callconv(.c) abi.Result, + release: *const fn (u64) callconv(.c) abi.Result, + map: *const fn (u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + unmap: *const fn (u64, usize) callconv(.c) abi.Result, + setAccess: *const fn (u64, usize, *const AccessDesc, usize) callconv(.c) abi.Result, + granularity: *const fn (*usize, *const Prop, c_int) callconv(.c) abi.Result, + }; + + pub const Prop = extern struct { // CUmemAllocationProp + type: c_int = 1, // CU_MEM_ALLOCATION_TYPE_PINNED + handle_types: c_int = 0, + loc_type: c_int = 1, // CU_MEM_LOCATION_TYPE_DEVICE + loc_id: c_int = 0, + win32: ?*anyopaque = null, + compression: u8 = 0, + rdma: u8 = 0, + usage: u16 = 0, + reserved: [4]u8 = .{ 0, 0, 0, 0 }, + }; + const AccessDesc = extern struct { loc_type: c_int = 1, loc_id: c_int, flags: c_int = 3 }; // CUmemAccessDesc, read-write + + comptime { + std.debug.assert(@sizeOf(Prop) == 32 and @offsetOf(Prop, "compression") == 24 and @sizeOf(AccessDesc) == 12); + } + + /// `va_bytes` of virtual addresses reserved (cheap: no memory behind them until a chunk maps). + pub fn open(d: *const Driver, device: abi.Device, va_bytes: usize) Error!Vmm { + var vmm_ok: c_int = 0; + if (d.api.cuDeviceGetAttribute(&vmm_ok, .virtual_memory_management_supported, device) != abi.success or vmm_ok == 0) return error.Unsupported; // VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + var api: Api = undefined; + const names = .{ .{ "reserve", "cuMemAddressReserve" }, .{ "addressFree", "cuMemAddressFree" }, .{ "create", "cuMemCreate" }, .{ "release", "cuMemRelease" }, .{ "map", "cuMemMap" }, .{ "unmap", "cuMemUnmap" }, .{ "setAccess", "cuMemSetAccess" }, .{ "granularity", "cuMemGetAllocationGranularity" } }; + inline for (names) |n| @field(api, n[0]) = lib.lookup(@FieldType(Api, n[0]), n[1]) orelse return error.MissingSymbol; + var prop: Prop = .{ .loc_id = device }; + var rdma: c_int = 0; // GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED: keep the memory registrable as cuMemAlloc's is + if (d.api.cuDeviceGetAttribute(&rdma, .gpu_direct_rdma_with_cuda_vmm_supported, device) == abi.success and rdma != 0) prop.rdma = 1; + var g: usize = 0; + if (api.granularity(&g, &prop, 1) != abi.success or g == 0) return error.Unsupported; // RECOMMENDED + g = @max(g, 2 << 20); + const len = std.mem.alignForward(usize, va_bytes, g); + var va: u64 = 0; + if (api.reserve(&va, len, g, 0, 0) != abi.success) return error.CudaFailed; + return .{ .lib = lib, .device = device, .prop = prop, .granularity = g, .va = va, .va_len = len, .next = va, .api = api }; + } + + pub fn close(v: *Vmm) void { + _ = v.api.addressFree(v.va, v.va_len); + v.lib.close(); + } + + pub fn backend(v: *Vmm) Backend { + return .{ .ctx = v, .granularity = v.granularity, .chunk = chunk, .release = release, .name = "vmm" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + if (v.va + v.va_len - v.next < len) return null; + var h: u64 = 0; + if (v.api.create(&h, len, &v.prop, 0) != abi.success) return null; + defer _ = v.api.release(h); // the mapping holds the memory until it is unmapped + const at = v.next; + if (v.api.map(at, len, 0, h, 0) != abi.success) return null; + const desc: AccessDesc = .{ .loc_id = v.device }; + if (v.api.setAccess(at, len, &desc, 1) != abi.success) { + _ = v.api.unmap(at, len); + return null; + } + v.next += len; + return at; + } + + fn release(ctx: *anyopaque, base: u64, len: usize) void { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + _ = v.api.unmap(base, len); // the virtual range is not reused: the reservation is sized for the process + } +}; + +// -- tests (a fake backend: addresses only) --------------------------------------------------------------------------- + +const Fake = struct { + next: u64 = 1 << 40, + chunks: u32 = 0, + released: u32 = 0, + fail: bool = false, + + fn backend(f: *Fake, g: usize) Backend { + return .{ .ctx = f, .granularity = g, .chunk = chunk, .release = release, .name = "fake" }; + } + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail) return null; + const p = f.next; + f.next += len; + f.chunks += 1; + return p; + } + fn release(ctx: *anyopaque, _: u64, _: usize) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.released += 1; + } +}; + +test "arena: small requests pack 512-aligned into shared chunks, are reused by size, large ones get their own" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + const p1 = a.alloc(20480).?; // a norm weight + const p2 = a.alloc(983040).?; // an mHC fn16 + const p3 = a.alloc(100).?; + try t.expectEqual(p1 + 20480, p2); + try t.expectEqual(p2 + 983040, p3); + for ([_]u64{ p1, p2, p3 }) |p| try t.expectEqual(@as(u64, 0), p % 512); + try t.expectEqual(@as(u32, 1), f.chunks); + // a freed range comes back for the same rounded size + try t.expect(a.free(p3)); + try t.expectEqual(p3, a.alloc(17).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + // three more 1 MiB ranges fit the 4 MiB chunk, the fourth opens a new shared chunk + for (0..3) |_| _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 1), f.chunks); + _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 2), f.chunks); + // large: its own chunk, rounded to the granularity, released when freed + const big = a.alloc(5 << 20).?; + try t.expectEqual(@as(u32, 3), f.chunks); + try t.expectEqual(@as(u64, 6 << 20), a.large.get(big).?); + try t.expect(a.free(big)); + try t.expectEqual(@as(u32, 1), f.released); + // not ours + try t.expect(!a.free(12345)); + // a backend that cannot map: null, the caller allocates as before + f.fail = true; + try t.expectEqual(@as(?u64, null), a.alloc(64 << 20)); + try t.expectEqual(@as(u64, 1), a.stats.fallback); +} + +test "arena: under DeviceBuffer.alloc through memory.setHook" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{}); + defer a.deinit(); + a.install(); + defer memory.setHook(null); + const h = memory.currentHook().?; + const p = h.alloc(h.ctx, 4096).?; + try t.expect(h.free(h.ctx, p)); + try t.expect(!h.free(h.ctx, 7)); +} + +/// Slots of `m` that hold neither a key nor a tombstone. +fn freeSlots(m: *const std.AutoHashMapUnmanaged(u64, u64)) usize { + const meta = m.metadata orelse return 0; + var n: usize = 0; + for (meta[0..m.capacity()]) |x| n += @intFromBool(x.isFree()); + return n; +} + +test "arena: a long run of allocations and frees leaves both address maps free slots" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + var held: std.ArrayList(u64) = .empty; + defer held.deinit(t.allocator); + var rng = std.Random.DefaultPrng.init(3); + for (0..20000) |_| { + // a serving process's buffers come and go (prompt workspaces, grown buffers): small and large alike + if (held.items.len > 0 and rng.random().boolean()) { + const k = rng.random().uintLessThan(usize, held.items.len); + try t.expect(a.free(held.swapRemove(k))); + } else { + const big = rng.random().uintLessThan(u32, 4) == 0; + try held.append(t.allocator, a.alloc(if (big) 3 << 20 else 512 * (1 + rng.random().uintLessThan(u64, 64))).?); + } + for ([_]*const std.AutoHashMapUnmanaged(u64, u64){ &a.live, &a.large }) |m| { + if (m.capacity() == 0) continue; + try t.expect(freeSlots(m) * 10 >= m.capacity()); + } + } +} diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..88ee9050dc 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -8,6 +8,10 @@ pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDevi pub const Driver = struct { lib: std.DynLib, api: abi.Api, + /// cuFuncGetName (CUDA 12.3+; optional: an older driver opens without it): a failed launch names its kernel + func_name: ?FuncGetName = null, + + pub const FuncGetName = *const fn (*?[*:0]const u8, abi.Function) callconv(.c) abi.Result; pub fn open() Error!Driver { return openPath("libcuda.so.1"); @@ -25,7 +29,7 @@ pub const Driver = struct { return error.MissingSymbol; }; } - const d: Driver = .{ .lib = lib, .api = api }; + const d: Driver = .{ .lib = lib, .api = api, .func_name = lib.lookup(FuncGetName, "cuFuncGetName") }; try d.check(api.cuInit(0), "cuInit"); return d; } @@ -38,18 +42,33 @@ pub const Driver = struct { pub fn check(self: *const Driver, res: abi.Result, what: []const u8) Error!void { if (res == abi.success) return; if (res == abi.error_not_ready) return error.NotReady; - std.log.err("{s}: {s} ({d}) {s}", .{ what, self.errorName(res), res, self.errorText(res) }); + std.log.err("{s}: {s} ({d}) {s}{s}", .{ what, self.errorName(res), res, self.errorText(res), shutdownNote(res) }); + return errorOf(res); + } + + /// The error a failed call returns (`check` without its log line, for callers that log more). + pub fn errorOf(res: abi.Result) Error { return switch (res) { 2 => error.OutOfDeviceMemory, abi.error_not_found => error.NotFound, + abi.error_not_ready => error.NotReady, else => error.CudaFailed, }; } + /// The driver's name for the code, else knownName's: during the exit's teardown cuGetErrorName fails too. pub fn errorName(self: *const Driver, res: abi.Result) []const u8 { var s: ?[*:0]const u8 = null; - if (self.api.cuGetErrorName(res, &s) != abi.success) return "CUDA_ERROR_UNKNOWN_CODE"; - return if (s) |p| std.mem.span(p) else "CUDA_ERROR_UNKNOWN_CODE"; + if (self.api.cuGetErrorName(res, &s) == abi.success) if (s) |p| return std.mem.span(p); + return knownName(res) orelse "CUDA_ERROR_UNKNOWN_CODE"; + } + + /// A kernel's symbol name (cuFuncGetName), or "?" when the driver cannot say. + pub fn functionName(self: *const Driver, f: abi.Function) []const u8 { + const get = self.func_name orelse return "?"; + var s: ?[*:0]const u8 = null; + if (get(&s, f) != abi.success) return "?"; + return if (s) |p| std.mem.span(p) else "?"; } pub fn errorText(self: *const Driver, res: abi.Result) []const u8 { @@ -82,3 +101,43 @@ pub const Driver = struct { return @intCast(10 * major + minor); } }; + +/// CUresult names (cuda.h) for when cuGetErrorName cannot answer. +pub fn knownName(res: abi.Result) ?[]const u8 { + return switch (res) { + 1 => "CUDA_ERROR_INVALID_VALUE", + 2 => "CUDA_ERROR_OUT_OF_MEMORY", + 3 => "CUDA_ERROR_NOT_INITIALIZED", + 4 => "CUDA_ERROR_DEINITIALIZED", + 101 => "CUDA_ERROR_INVALID_DEVICE", + 200 => "CUDA_ERROR_INVALID_IMAGE", + 201 => "CUDA_ERROR_INVALID_CONTEXT", + 209 => "CUDA_ERROR_NO_BINARY_FOR_GPU", + 400 => "CUDA_ERROR_INVALID_HANDLE", + 500 => "CUDA_ERROR_NOT_FOUND", + 600 => "CUDA_ERROR_NOT_READY", + 700 => "CUDA_ERROR_ILLEGAL_ADDRESS", + 701 => "CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES", + 702 => "CUDA_ERROR_LAUNCH_TIMEOUT", + 709 => "CUDA_ERROR_CONTEXT_IS_DESTROYED", + 710 => "CUDA_ERROR_ASSERT", + 714 => "CUDA_ERROR_HARDWARE_STACK_ERROR", + 715 => "CUDA_ERROR_ILLEGAL_INSTRUCTION", + 716 => "CUDA_ERROR_MISALIGNED_ADDRESS", + 719 => "CUDA_ERROR_LAUNCH_FAILED", + 999 => "CUDA_ERROR_UNKNOWN", + else => null, + }; +} + +/// CUDA_ERROR_DEINITIALIZED: the process is exiting (a stop signal, then exit) while a thread still calls the driver. +pub fn shutdownNote(res: abi.Result) []const u8 { + return if (res == 4) " (the CUDA driver is shutting down: the process is exiting)" else ""; +} + +test "CUresult names without the driver: DEINITIALIZED is 4, unknown codes stay unknown" { + try std.testing.expectEqualStrings("CUDA_ERROR_DEINITIALIZED", knownName(4).?); + try std.testing.expectEqualStrings("CUDA_ERROR_ILLEGAL_ADDRESS", knownName(700).?); + try std.testing.expect(knownName(12345) == null); + try std.testing.expect(shutdownNote(4).len > 0 and shutdownNote(1).len == 0); +} diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig new file mode 100644 index 0000000000..693f8ce25a --- /dev/null +++ b/zig/src/cuda/graph_cache.zig @@ -0,0 +1,401 @@ +//! Keyed CUDA graphs for any family: one captured executable per key, launched again when the key comes back. + +const std = @import("std"); +const graph = @import("graph.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +/// One step's launches on `stream`; it must not synchronize, allocate device memory or read the device on the host. +pub const Body = struct { + ctx: *anyopaque, + run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, +}; + +/// Every rank's verdict AND-ed on a miss (hits never call it), so ranks capture, evict and fall back together. +pub const Agree = struct { + ctx: *anyopaque, + all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, +}; + +/// Capture a body into an executable graph, launch it, free it: a fake engine runs the policy in host tests. +pub const Engine = struct { + ctx: *anyopaque, + capture: *const fn (ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque, + launch: *const fn (ctx: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void, + free: *const fn (ctx: *anyopaque, exec: *anyopaque) void, +}; + +pub const Settings = struct { + on: bool = true, // off: every step runs eagerly + max: u32 = 48, // graphs held at most + min_keep: u32 = 8, // under the memory floor the cache evicts down to this many, then stops capturing + hold: bool = false, // under the floor a miss runs eagerly with no eviction: held graphs cost next to no memory +}; + +pub const Outcome = enum { replayed, captured, eager }; + +pub fn Cache(comptime Key: type) type { + return struct { + const Self = @This(); + const Entry = struct { exec: *anyopaque, used: u64 }; + + gpa: std.mem.Allocator, + engine: Engine, + settings: Settings, + agree: ?Agree = null, + room: ?*const fn () bool = null, // false when this rank is below its memory floor; null: always room + entries: std.AutoArrayHashMapUnmanaged(Key, Entry) = .empty, + cap: u32, + tick: u64 = 0, + fingerprint: ?u64 = null, + failed: bool = false, // a capture failed on some rank: eager from here on + stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, + + pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { + var s = settings; + s.min_keep = @min(s.min_keep, s.max); // the floor never holds more than the configured max + return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; + } + + pub fn deinit(c: *Self) void { + c.dropAll(); + c.entries.deinit(c.gpa); + } + + pub fn count(c: *const Self) usize { + return c.entries.count(); + } + + pub fn has(c: *const Self, key: Key) bool { + return c.entries.contains(key); + } + + /// Every graph freed (a fingerprint change, a shutdown). + pub fn dropAll(c: *Self) void { + for (c.entries.values()) |e| c.engine.free(c.engine.ctx, e.exec); + c.stats.dropped += c.entries.count(); + c.entries.clearRetainingCapacity(); + } + + fn agreed(c: *Self, ok: bool) !bool { + const a = c.agree orelse return ok; + return a.all(a.ctx, ok); + } + + fn evictOldest(c: *Self) void { + var at: ?usize = null; + for (c.entries.values(), 0..) |e, i| { + if (at == null or e.used < c.entries.values()[at.?].used) at = i; + } + const i = at orelse return; + c.engine.free(c.engine.ctx, c.entries.values()[i].exec); + c.entries.swapRemoveAt(i); + c.stats.evicted += 1; + } + + fn eager(c: *Self, stream: Stream, body: Body) !Outcome { + try body.run(body.ctx, stream); + c.stats.eager += 1; + return .eager; + } + + /// One step: the key's graph, captured first on a miss; a new `fingerprint` (same step on all ranks) drops all. + pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { + if (!c.settings.on or c.failed or c.settings.max == 0) return c.eager(stream, body); + if (c.fingerprint != null and c.fingerprint.? != fingerprint) c.dropAll(); + c.fingerprint = fingerprint; + c.tick += 1; + if (c.entries.getPtr(key)) |e| { + e.used = c.tick; + try c.engine.launch(c.engine.ctx, e.exec, stream); + c.stats.replayed += 1; + return .replayed; + } + // a miss: room on every rank (a host map that cannot grow is none), else shed a quarter and lower the cap + const reserved = if (c.entries.ensureUnusedCapacity(c.gpa, 1)) |_| true else |_| false; + const room_here = reserved and (if (c.room) |f| f() else true); + if (!try c.agreed(room_here)) { + if (c.settings.hold) { + if (!quiet and c.stats.floor_hits == 0) std.log.warn("graph cache: under the memory floor on {s} rank: this key eager, {d} graphs held (hold)", .{ if (room_here) "another" else "this", c.entries.count() }); + c.stats.floor_hits += 1; + return c.eager(stream, body); + } + if (c.entries.count() > c.settings.min_keep) { + const drop = @max(1, c.entries.count() / 4); + for (0..drop) |_| if (c.entries.count() > c.settings.min_keep) c.evictOldest(); + } + const was = c.cap; + c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); // min_keep <= max (init) + c.stats.floor_hits += 1; + if (!quiet and c.cap != was) std.log.warn("graph cache: under the memory floor on {s} rank: cap {d} -> {d} ({d} held, {d} evicted, {d} eager so far)", .{ if (room_here) "another" else "this", was, c.cap, c.entries.count(), c.stats.evicted, c.stats.eager + 1 }); + return c.eager(stream, body); + } + if (c.cap < c.settings.max) { + if (!quiet) std.log.info("graph cache: room again on every rank: cap {d} -> {d}", .{ c.cap, c.settings.max }); + c.cap = c.settings.max; + } + while (c.entries.count() >= c.cap) c.evictOldest(); + const exec = c.engine.capture(c.engine.ctx, stream, body) catch |err| blk: { + // the only trace of why graphs went off (the step then runs eagerly, which may fail the same way) + if (!quiet) std.log.warn("graph capture failed: {s} (key {any}, {d} graphs held)", .{ @errorName(err), key, c.entries.count() }); + break :blk null; + }; + var owned = exec; // freed here until the map holds it, so a failed agreement cannot leak it + defer if (owned) |x| c.engine.free(c.engine.ctx, x); + if (!try c.agreed(exec != null)) { + c.failed = true; + c.dropAll(); + return c.eager(stream, body); + } + c.entries.putAssumeCapacityNoClobber(key, .{ .exec = exec.?, .used = c.tick }); + owned = null; + try c.engine.launch(c.engine.ctx, exec.?, stream); + c.stats.captured += 1; + return .captured; + } + }; +} + +/// The driver's engine: thread-local stream capture (other threads' CUDA calls stay out), instantiate, upload, launch. +pub const CudaEngine = struct { + d: *const Driver, + + pub fn engine(e: *CudaEngine) Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + + const Held = struct { exec: graph.Exec }; + + fn capture(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const e: *CudaEngine = @ptrCast(@alignCast(ctx)); + try graph.beginCapture(stream, .thread_local); + body.run(body.ctx, stream) catch |err| { + // end the capture so the stream is usable again; the partial graph is dropped + if (graph.endCapture(stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try graph.endCapture(stream); + defer g.deinit(); + const h = try std.heap.c_allocator.create(Held); + errdefer std.heap.c_allocator.destroy(h); + h.exec = try g.instantiate(); + errdefer h.exec.deinit(); + try h.exec.upload(stream); + _ = e; + return h; + } + + fn launch(_: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void { + const h: *Held = @ptrCast(@alignCast(exec)); + try h.exec.launchOn(stream); + } + + fn free(_: *anyopaque, exec: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(exec)); + h.exec.deinit(); + std.heap.c_allocator.destroy(h); + } +}; + +/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares, on one rank only. +pub fn fingerprintOf(words: []const u64) u64 { + var h: u64 = 0xcbf29ce484222325; + for (words) |w| { + var x = w; + for (0..8) |_| { + h = (h ^ (x & 0xff)) *% 0x100000001b3; + x >>= 8; + } + } + return h; +} + +// -- host tests: the policy with a fake engine -------------------------------------------------------------------- + +const Fake = struct { + captures: u32 = 0, + launches: u32 = 0, + live: u32 = 0, + fail_capture: bool = false, + slots: [64]u8 = undefined, + + fn engine(f: *Fake) Engine { + return .{ .ctx = f, .capture = cap, .launch = lau, .free = fre }; + } + fn cap(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail_capture) return error.CaptureFailed; + _ = stream; + _ = body; // a capture records the body; nothing runs + f.captures += 1; + f.live += 1; + return &f.slots[f.captures % f.slots.len]; + } + fn lau(ctx: *anyopaque, _: *anyopaque, _: Stream) anyerror!void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.launches += 1; + } + fn fre(ctx: *anyopaque, _: *anyopaque) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.live -= 1; + } +}; + +const Ran = struct { + n: u32 = 0, + fn body(r: *Ran) Body { + return .{ .ctx = r, .run = go }; + } + fn go(ctx: *anyopaque, _: Stream) anyerror!void { + const r: *Ran = @ptrCast(@alignCast(ctx)); + r.n += 1; + } +}; + +const no_stream: Stream = .{ .d = undefined, .handle = null }; + +test "graph cache: capture on first use, replay after, LRU at the cap" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 2 }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.replayed, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.captured, try c.run(2, no_stream, 7, r.body())); + _ = try c.run(1, no_stream, 7, r.body()); // 1 is now the most recent + try std.testing.expectEqual(Outcome.captured, try c.run(3, no_stream, 7, r.body())); // evicts 2 + try std.testing.expect(c.has(1) and c.has(3) and !c.has(2)); + try std.testing.expectEqual(@as(u32, 0), r.n); // the body never ran eagerly + try std.testing.expectEqual(@as(u32, 2), f.live); + try std.testing.expectEqual(@as(u32, 5), f.launches); +} + +test "graph cache: a new fingerprint drops every graph" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + defer c.deinit(); + _ = try c.run(1, no_stream, 7, r.body()); + _ = try c.run(2, no_stream, 7, r.body()); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 8, r.body())); + try std.testing.expectEqual(@as(usize, 1), c.count()); + try std.testing.expectEqual(@as(u32, 1), f.live); +} + +const Peer = struct { + theirs: []const bool, // the other rank's verdicts, in call order + at: usize = 0, + fn agree(p: *Peer) Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const t = p.theirs[p.at]; + p.at += 1; + return ok and t; + } +}; + +test "graph cache: a capture failing on another rank turns graphs off here too" { + var f: Fake = .{}; + var r: Ran = .{}; + // per miss: room, then the capture's verdict; the peer's capture fails on the second key + var peer: Peer = .{ .theirs = &.{ true, true, true, false } }; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + c.agree = peer.agree(); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(2, no_stream, 7, r.body())); + try std.testing.expect(c.failed); + try std.testing.expectEqual(@as(u32, 0), f.live); // this rank's good graphs are freed as well + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 2), r.n); +} + +var room_flag = true; +fn roomFn() bool { + return room_flag; +} + +test "graph cache: below the memory floor the least recent quarter goes, the cap follows, the step runs eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 6), c.count()); + try std.testing.expectEqual(@as(u32, 6), c.cap); + try std.testing.expect(!c.has(0) and !c.has(1) and c.has(7)); + room_flag = true; + // room again: the cap is back to max, the miss captures without evicting + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 7), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + try std.testing.expect(c.has(7) and c.has(2)); + try std.testing.expectEqual(@as(u64, 1), c.stats.floor_hits); +} + +test "graph cache: a dip under the floor does not leave later keys churning once memory is back" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 16, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..12) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; // a long prompt's workspace held + _ = try c.run(200, no_stream, 7, r.body()); + room_flag = true; // freed: the decode after it meets 12 new keys (a new context bucket), then replays them + const before = c.stats.evicted; + for (0..12) |k| _ = try c.run(@intCast(300 + k), no_stream, 7, r.body()); + for (0..12) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(300 + k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 16), c.cap); + // only what the cap of 16 forces: 9 left after the dip + 12 new = 21 -> 5 evictions, none during the replays + try std.testing.expectEqual(before + 5, c.stats.evicted); +} + +test "graph cache: hold - below the floor a miss runs eagerly, nothing is evicted, the held keys replay" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2, .hold = true }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(101, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 8), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + for (0..8) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u64, 0), c.stats.evicted); + try std.testing.expectEqual(@as(u64, 2), c.stats.floor_hits); + room_flag = true; + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 9), c.count()); +} + +test "graph cache: off runs every step eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .on = false }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: fingerprint" { + try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); +} + +test { + _ = @import("graph_cache_regression_test.zig"); +} diff --git a/zig/src/cuda/graph_cache_regression_test.zig b/zig/src/cuda/graph_cache_regression_test.zig new file mode 100644 index 0000000000..7f2da1063a --- /dev/null +++ b/zig/src/cuda/graph_cache_regression_test.zig @@ -0,0 +1,181 @@ +//! Failure ownership, rank verdict counts and bounded-cap regressions for the graph cache. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const expect = std.testing.expect; +const equal = std.testing.expectEqual; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + launches: usize = 0, + bodies: usize = 0, + handles: [32]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(p: *anyopaque, _: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.launches += 1; + } + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + std.debug.assert(e.live > 0); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(p: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.bodies += 1; + } +}; + +const Peer = struct { + verdicts: []const bool, + at: usize = 0, + throw_at: ?usize = null, + fn agree(p: *Peer) cache.Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) !bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const at = p.at; + p.at += 1; + if (p.throw_at == at) return error.CollectiveFailed; + if (at >= p.verdicts.len) return error.UnexpectedCollective; + return ok and p.verdicts[at]; + } +}; + +test "graph cache regression: capture agreement throw frees the not-yet-inserted executable" { + var e: Engine = .{}; + // per miss: room, then the capture's verdict + var p: Peer = .{ .verdicts = &.{ true, true }, .throw_at = 1 }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try std.testing.expectError(error.CollectiveFailed, c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.captures); + try equal(@as(usize, 0), e.live); + try equal(@as(usize, 0), c.count()); + try equal(@as(usize, 2), p.at); +} + +test "graph cache regression: local map OOM declines before any capture and can recover" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try expect(failing.has_induced_failure); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), e.live); + try expect(!c.failed); + failing.fail_index = std.math.maxInt(usize); + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.live); +} + +test "graph cache regression: peer map OOM makes a locally prepared rank eager before capture" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{false} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 1), e.bodies); +} + +test "graph cache regression: local OOM still contributes the room verdict" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var p: Peer = .{ .verdicts = &.{true} }; + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); +} + +test "graph cache regression: hits make no collective; a fingerprint change shared by every rank recaptures" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{ true, true, true, true, true, true } }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + _ = try c.run(1, no_stream, 7, e.body()); + _ = try c.run(2, no_stream, 7, e.body()); + try equal(@as(usize, 4), p.at); + for (0..4) |_| try equal(cache.Outcome.replayed, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 4), p.at); + // a new layout generation on every rank: both graphs go, the next miss agrees twice as any other + try equal(cache.Outcome.captured, try c.run(1, no_stream, 8, e.body())); + try equal(@as(usize, 6), p.at); + try equal(@as(u64, 2), c.stats.dropped); + try equal(@as(usize, 1), c.count()); + try equal(@as(usize, 1), e.live); +} + +var room = true; +fn available() bool { + return room; +} + +test "graph cache regression: zero maximum remains eager across pressure recovery, with no collective" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 0 }); + defer c.deinit(); + c.room = available; + c.agree = p.agree(); + for ([_]bool{ false, true, false, true }) |ok| { + room = ok; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + try equal(@as(usize, 0), c.count()); + } + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), p.at); +} + +test "graph cache regression: default floor cannot enlarge a two-entry maximum" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2 }); + defer c.deinit(); + c.room = available; + try equal(@as(u32, 2), c.settings.min_keep); + room = false; + _ = try c.run(1, no_stream, 7, e.body()); + room = true; + for (0..8) |key| { + _ = try c.run(@intCast(key), no_stream, 7, e.body()); + try expect(c.count() <= 2 and c.cap <= 2 and e.live <= 2); + } +} + +test "graph cache regression: zero keep floor recovers after reducing the cap to zero" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2, .min_keep = 0 }); + defer c.deinit(); + c.room = available; + room = false; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + room = true; + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 2), c.cap); +} diff --git a/zig/src/cuda/launch.zig b/zig/src/cuda/launch.zig index 2472042638..c80684f291 100644 --- a/zig/src/cuda/launch.zig +++ b/zig/src/cuda/launch.zig @@ -4,7 +4,8 @@ const std = @import("std"); const abi = @import("abi.zig"); const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; -const Error = @import("driver.zig").Error; +const driver = @import("driver.zig"); +const Error = driver.Error; pub const Dim3 = abi.Dim3; @@ -80,7 +81,7 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void const g = cfg.grid; const b = cfg.block; if (cfg.plain()) { - return d.check(d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); + return launched(f, cfg, d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); } var attrs: [4]abi.LaunchAttribute = undefined; const n = fillAttributes(cfg, &attrs); @@ -96,7 +97,36 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void .attrs = &attrs, .num_attrs = @intCast(n), }; - try d.check(d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); + try launched(f, cfg, d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); +} + +/// A launch's result; a failure names the kernel and its configuration with the code. +fn launched(f: Function, cfg: Config, res: abi.Result, what: []const u8) Error!void { + if (res == abi.success) return; + const d = f.d; + var buf: [160]u8 = undefined; + std.log.err("{s}: {s} ({d}) {s}; {s}{s}", .{ what, d.errorName(res), res, d.errorText(res), describe(&buf, d.functionName(f.handle), cfg), driver.shutdownNote(res) }); + return driver.Driver.errorOf(res); +} + +/// "kernel , grid XxYxZ, block XxYxZ, shared N B" (+ cluster / PDL / cooperative when set). +pub fn describe(buf: []u8, name: []const u8, cfg: Config) []const u8 { + const g = cfg.grid; + const b = cfg.block; + var w = std.Io.Writer.fixed(buf); + w.print("kernel {s}, grid {d}x{d}x{d}, block {d}x{d}x{d}, shared {d} B", .{ name, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared }) catch return w.buffered(); + if (cfg.cluster) |c| w.print(", cluster {d}x{d}x{d}", .{ c.x, c.y, c.z }) catch return w.buffered(); + if (cfg.pdl) w.writeAll(", pdl") catch return w.buffered(); + if (cfg.cooperative) w.writeAll(", cooperative") catch return w.buffered(); + return w.buffered(); +} + +test "a failed launch's description: the kernel, its grid, block and shared bytes, bounded" { + var buf: [160]u8 = undefined; + const cfg: Config = .{ .grid = .{ .x = 70000, .y = 2, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 }, .shared = 49152, .pdl = true }; + try std.testing.expectEqualStrings("kernel _stream_pf, grid 70000x2x1, block 128x1x1, shared 49152 B, pdl", describe(&buf, "_stream_pf", cfg)); + var small: [24]u8 = undefined; + try std.testing.expect(describe(&small, "a_very_long_kernel_symbol_name", cfg).len <= small.len); } fn attribute(id: abi.LaunchAttributeId) abi.LaunchAttribute { diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig index 2c469bae2c..6d802980c7 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -49,16 +49,37 @@ fn lockCounts() void { while (!counts_mutex.tryLock()) std.Thread.yield() catch {}; } +/// A sub-allocator under DeviceBuffer (arena.zig): `alloc` gives an address or null for cuMemAlloc, `free` owns it. +pub const Hook = struct { + ctx: *anyopaque, + alloc: *const fn (ctx: *anyopaque, len: usize) ?abi.DevicePtr, + free: *const fn (ctx: *anyopaque, p: abi.DevicePtr) bool, // true when the address was the hook's +}; + +var alloc_hook: ?Hook = null; + +/// Installs, or with null removes, the sub-allocator: once at boot, before the buffers it should hold. +pub fn setHook(h: ?Hook) void { + alloc_hook = h; +} + +pub fn currentHook() ?Hook { + return alloc_hook; +} + pub const DeviceBuffer = struct { d: *const Driver, ptr: abi.DevicePtr, len: usize, borrowed: bool = false, // a span of memory another value owns (a carveout's): free leaves it alone - /// `len` bytes, 256-byte aligned by the driver; zero bytes allocate nothing and hold address 0. + /// `len` bytes, 256-byte aligned (512 from an installed arena); zero bytes allocate nothing and hold address 0. pub fn alloc(d: *const Driver, len: usize) Error!DeviceBuffer { var p: abi.DevicePtr = 0; - if (len > 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + if (len > 0) { + if (alloc_hook) |h| p = h.alloc(h.ctx, len) orelse 0; + if (p == 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + } held(&device_bytes, len); return .{ .d = d, .ptr = p, .len = len }; } @@ -73,7 +94,10 @@ pub const DeviceBuffer = struct { pub fn free(self: *DeviceBuffer) void { if (!self.borrowed) { - if (self.ptr != 0) _ = self.d.api.cuMemFree_v2(self.ptr); + if (self.ptr != 0) { + const mine = if (alloc_hook) |h| h.free(h.ctx, self.ptr) else false; + if (!mine) _ = self.d.api.cuMemFree_v2(self.ptr); + } freed(&device_bytes, self.len); } self.* = undefined; @@ -209,6 +233,28 @@ test "freeing a borrowed span leaves the driver and the counts alone" { try std.testing.expectEqual(before.device, usage(false).device); } +test "a buffer the hook owns is freed by the hook, never by the driver" { + const Owner = struct { + var frees: u32 = 0; + fn alloc(_: *anyopaque, _: usize) ?abi.DevicePtr { + return null; + } + fn free(_: *anyopaque, p: abi.DevicePtr) bool { + frees += 1; + return p == 0x2000; + } + }; + var ctx: u8 = 0; + setHook(.{ .ctx = &ctx, .alloc = Owner.alloc, .free = Owner.free }); + defer setHook(null); + held(&device_bytes, 4096); + var b: DeviceBuffer = .{ .d = undefined, .ptr = 0x2000, .len = 4096 }; + b.free(); + var c: DeviceBuffer = .{ .d = undefined, .ptr = 0x3000, .len = 64, .borrowed = true }; + c.free(); + try std.testing.expectEqual(@as(u32, 1), Owner.frees); +} + test "peak reset preserves an allocation between its read and store" { const before = usage(true); held(&device_bytes, 100); diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..85f88bace8 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -12,6 +12,8 @@ pub const RedOp = enum(c_int) { sum = 0, prod = 1, max = 2, min = 3, avg = 4 }; pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; +pub const in_progress: Result = 7; // ncclInProgress: a non-blocking communicator's call is still running + const R = Result; const D = abi.DevicePtr; @@ -25,6 +27,11 @@ pub const Api = struct { ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, ncclBroadcast: *const fn (D, D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclSend: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclRecv: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + /// Frees the communicator after ending its pending work; safe from another thread while a kernel waits. + ncclCommAbort: *const fn (Comm) callconv(.c) R, + ncclCommGetAsyncError: *const fn (Comm, *Result) callconv(.c) R, ncclGroupStart: *const fn () callconv(.c) R, ncclGroupEnd: *const fn () callconv(.c) R, }; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index ebbb12b935..cf7e74d559 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -6,6 +6,7 @@ pub const Error = @import("driver.zig").Error; pub const Context = @import("context.zig").Context; pub const DeviceBuffer = @import("memory.zig").DeviceBuffer; pub const HostBuffer = @import("memory.zig").HostBuffer; +pub const memory = @import("memory.zig"); pub const usage = @import("memory.zig").usage; pub const Usage = @import("memory.zig").Usage; pub const carveout = @import("carveout.zig"); @@ -32,6 +33,8 @@ pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); pub const experts = @import("experts.zig"); pub const modelopt = @import("modelopt.zig"); +pub const graph_cache = @import("graph_cache.zig"); +pub const arena = @import("arena.zig"); test { _ = @import("memory.zig"); @@ -47,4 +50,6 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); + _ = graph_cache; + _ = arena; } diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md new file mode 100644 index 0000000000..7ad00ef7d5 --- /dev/null +++ b/zig/src/fabric/tp/README.md @@ -0,0 +1,72 @@ +# Tensor parallelism across processes + +`zig/src/fabric/tp` runs tensor parallelism with one GPU per process, either on one host or on two hosts over RoCE. +`TF_COMM_BACKEND` picks the transport for small exchanges: `nccl` (the default), `mailbox` (two GPUs on one host) +or `roce` (two hosts). NCCL carries everything larger, so it must work on its own as well. + +## Building `tf-tp-test` + +The `mailbox` and `roce` backends need the mailbox kernel (`mailbox.cu`) built into the binary: + +```sh +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-nvcc=/usr/local/cuda/bin/nvcc +# or embed a fatbin built elsewhere: +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-fatbin=/abs/path/tp_mailbox.fatbin +``` + +A build without `-Dtp-nvcc` or `-Dtp-fatbin` still runs `nccl` and `roce-host`, which needs no GPU. +In that build the `mailbox` and `roce` suites fail at load with `NoKernel` (an engine falls back to NCCL instead). +The `INFO` line `tf-tp-test` prints at start shows the embedded kernel's size; `0 B` means none was built in. + +## Choosing the RDMA devices + +| Variable | Used by | Needed when | +| --- | --- | --- | +| `TF_TP_ROCE_HCA` | the RoCE wire | the host has more than one active RoCE port that is not cabled to the peer | +| `NCCL_IB_HCA` | NCCL | the same as above; NCCL otherwise picks its own devices | +| `NCCL_SOCKET_IFNAME` | NCCL bootstrap | the host has several IP interfaces and NCCL picks one the peer cannot reach | +| `TF_TP_MASTER` | the TCP bootstrap | always on two hosts: rank 0's address as the other rank reaches it | + +`TF_TP_ROCE_HCA` is `name[:gid],name[:gid]`: up to two devices, and optionally each one's RoCE v2 GID index. +Without it the rank picks the devices itself, at most `TF_TP_ROCE_HCAS` (default 2). +A host cabled to more than one neighbour should always set it explicitly. +The two hosts may list their devices in a different order and use different names. +Each device pairs with the peer device on the same IPv4 /24. + +`NCCL_IB_HCA` should name the same devices without the `:gid` suffix. +In NCCL, `:N` after a device name selects port `N`, not a GID index. + +`run_hosts.sh` sets `NCCL_IB_HCA` from `TF_TP_ROCE_HCA` when only that is set. +When neither is set, it first runs `tf-tp-test hca` on both ranks: the RoCE pick, made with the peer over the bootstrap. +The same rule picks the wire's devices, and each rank prints its own as `NCCL_IB_HCA=name,...`. +`hca` needs no GPU, and it exits 1 when no port shares a /24 with the peer, so no host falls back to its first ports. +The runner logs each rank's pick and stops with that reason when `hca` fails. +Both hosts must use the same mode: if one names its ports and the other names none, only the latter runs `hca`. +It finds no peer and stops after `TF_TP_PLAN_TIMEOUT` (default 300 s) with that reason. + +Two DGX Spark nodes, with the second ConnectX-7 port of each cabled to the other (both PCIe functions of that +port), and the first port cabled to a third node: + +```sh +# on both hosts; rank 0's link address is 192.0.2.1 +export TF_TP_ROCE_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_SOCKET_IFNAME=enp1s0f1np1 +zig/src/fabric/tp/run_hosts.sh 0 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank0 # host A +zig/src/fabric/tp/run_hosts.sh 1 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank1 # host B +``` + +`ibv_devices` lists the RDMA devices; `ibv_devinfo -d NAME` shows a port's state and link layer. +`/sys/class/infiniband/NAME/ports/1/gids/N` and `gid_attrs/ndevs/N` show each GID and its netdev. +`ulimit -l` must allow the RoCE region to be pinned (`unlimited` is simplest). + +## RoCE limits and the failed-transport check + +No RDMA write is larger than 1 MiB, so on `TF_COMM_BACKEND=roce` a `TF_TP_MAILBOX_MAX_KB` above 1024 fails at boot. +`TF_TP_ROCE_CHUNKED_KB` (default 0, off) sends variable-length gathers past the slot over RoCE in slot-sized ops. +It needs a kernel built from this tree's `mailbox.cu`. +At most 2 MiB is posted and not yet completed at any time. +A rank whose transport fails moves its queue pairs to the error state and keeps its region until the process exits. + +`run_transport.sh MASTER_ADDR PEER_SSH OUTDIR` checks that failure path on two hosts with `tf-tp-test failfast V +transport` and stops at the first kernel log line or HCA error counter that grows; run it with no arguments for help. diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig new file mode 100644 index 0000000000..1486e3a8f7 --- /dev/null +++ b/zig/src/fabric/tp/bootstrap.zig @@ -0,0 +1,262 @@ +//! Rank discovery over TCP (a star on rank 0): small host blobs all-gathered, settings agreed, sockets kept for fate. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; + +pub const max_ranks = 64; +/// Largest blob a rank contributes to one `allGather`. +pub const max_blob = 64 * 1024; + +const magic: u32 = 0x31425054; // "TPB1" +const version: u32 = 2; + +pub const Error = error{ Protocol, Mismatch, TooLarge } || sock.Error; + +/// `role`: 0 the control connection (bootstrap, then the fate channel), 1 the plan link. +const Hello = extern struct { magic: u32, version: u32, world: u32, rank: u32, role: u32 }; + +pub const Bootstrap = struct { + rank: u32, + world: u32, + /// Rank 0: the socket of rank i at index i (index 0 unused); other ranks: rank 0's socket at index 0. + fds: [max_ranks]sock.Fd = @splat(-1), + /// The plan link's sockets, laid out as `fds`. + plan_fds: [max_ranks]sock.Fd = @splat(-1), + /// Each call's limit for the peers' bytes (the connect timeout). + timeout_ns: u64, + + fn deadline(b: *const Bootstrap) u64 { + return sock.nowNs() + b.timeout_ns; + } + + /// Rank 0 listens on `link_port` and accepts every other rank; the others connect to `master` (connect timeout). + pub fn open(cfg: Config) Error!Bootstrap { + if (cfg.world > max_ranks or cfg.rank >= cfg.world) return error.Protocol; + var b: Bootstrap = .{ .rank = cfg.rank, .world = cfg.world, .timeout_ns = cfg.connect_timeout_ns }; + errdefer b.close(); + if (cfg.world == 1) return b; + if (cfg.rank == 0) { + const l = try sock.listen(cfg.link_port); + defer sock.close(l); + try b.acceptAll(l); + } else { + const until = b.deadline(); + for ([_]u32{ 0, 1 }) |role| { + const fd = try sock.connect(cfg.master, cfg.link_port, until); + if (role == 0) b.fds[0] = fd else b.plan_fds[0] = fd; + const h: Hello = .{ .magic = magic, .version = version, .world = cfg.world, .rank = cfg.rank, .role = role }; + try sock.sendAll(fd, std.mem.asBytes(&h)); + } + } + return b; + } + + /// Rank 0 over an already listening socket (tests: an ephemeral port). + pub fn acceptOn(l: sock.Fd, world: u32, timeout_ns: u64) Error!Bootstrap { + var b: Bootstrap = .{ .rank = 0, .world = world, .timeout_ns = timeout_ns }; + errdefer b.close(); + try b.acceptAll(l); + return b; + } + + fn acceptAll(b: *Bootstrap, l: sock.Fd) Error!void { + // every other rank's control and plan connections, in any order + var joined: u32 = 0; + const until = b.deadline(); + while (joined < 2 * (b.world - 1)) { + const fd = try sock.accept(l, until); + var h: Hello = undefined; + sock.recvExact(fd, std.mem.asBytes(&h), until) catch { + sock.close(fd); + continue; + }; + if (h.magic != magic or h.version != version) { + log.warn("tp bootstrap: a connection that is not a rank (magic {x}); ignored", .{h.magic}); + sock.close(fd); + continue; + } + const table = if (h.role == 0) &b.fds else &b.plan_fds; + if (h.world != b.world or h.rank == 0 or h.rank >= b.world or h.role > 1 or table[h.rank] >= 0) { + log.warn("tp bootstrap: rank {d} of {d} joined a world of {d}", .{ h.rank, h.world, b.world }); + sock.close(fd); + return error.Mismatch; + } + table[h.rank] = fd; + joined += 1; + } + } + + /// Takes the control sockets (the fate channel keeps them); `close` then leaves them open. + pub fn release(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.fds; + b.fds = @splat(-1); + return out; + } + + /// Takes the plan link's sockets. + pub fn releasePlan(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.plan_fds; + b.plan_fds = @splat(-1); + return out; + } + + pub fn close(b: *Bootstrap) void { + for ([_]*[max_ranks]sock.Fd{ &b.fds, &b.plan_fds }) |table| for (table) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } + + /// Every rank's `mine` in rank order into `out`, returned as slices of `out`; every rank calls it equally often. + pub fn allGather(b: *Bootstrap, mine: []const u8, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + if (mine.len > max_blob) return error.TooLarge; + std.debug.assert(blobs.len >= b.world); + if (b.world == 1) { + if (out.len < mine.len) return error.TooLarge; + @memcpy(out[0..mine.len], mine); + blobs[0] = out[0..mine.len]; + return blobs[0..1]; + } + if (b.rank != 0) { + try sendFrame(b.fds[0], mine); + return b.readAll(b.fds[0], out, blobs); + } + // rank 0: everyone's blob, then the whole table to everyone + var used: usize = 0; + for (0..b.world) |i| { + const len: u32 = if (i == 0) @intCast(mine.len) else blk: { + var l: u32 = undefined; + try sock.recvExact(b.fds[i], std.mem.asBytes(&l), b.deadline()); + break :blk l; + }; + if (len > max_blob or used + len > out.len) return error.TooLarge; + if (i == 0) @memcpy(out[used..][0..len], mine) else try sock.recvExact(b.fds[i], out[used..][0..len], b.deadline()); + blobs[i] = out[used..][0..len]; + used += len; + } + for (1..b.world) |i| { + for (blobs[0..b.world]) |blob| try sendFrame(b.fds[i], blob); + } + return blobs[0..b.world]; + } + + fn readAll(b: *Bootstrap, fd: sock.Fd, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + var used: usize = 0; + for (0..b.world) |i| { + var l: u32 = undefined; + try sock.recvExact(fd, std.mem.asBytes(&l), b.deadline()); + if (l > max_blob or used + l > out.len) return error.TooLarge; + try sock.recvExact(fd, out[used..][0..l], b.deadline()); + blobs[i] = out[used..][0..l]; + used += l; + } + return blobs[0..b.world]; + } + + /// Rank 0's `bytes` on every rank (into `out` on the others; `bytes` is ignored there). + pub fn broadcast(b: *Bootstrap, bytes: []const u8, out: []u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [4096]u8 = undefined; + if (out.len * b.world > scratch.len) return error.TooLarge; + const all = try b.allGather(if (b.rank == 0) bytes else &.{}, &scratch, &blobs); + if (all[0].len != out.len) return error.Protocol; + @memcpy(out, all[0]); + } + + /// Every rank's `bytes` must be equal (settings that change bytes on the wire); `what` names them in the log. + pub fn agree(b: *Bootstrap, bytes: []const u8, what: []const u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [8192]u8 = undefined; + const all = try b.allGather(bytes, &scratch, &blobs); + for (all, 0..) |blob, i| if (!std.mem.eql(u8, blob, all[0])) { + log.warn("tp bootstrap: rank {d}'s {s} differs from rank 0's; both ranks refuse to start", .{ i, what }); + return error.Mismatch; + }; + } +}; + +fn sendFrame(fd: sock.Fd, bytes: []const u8) sock.Error!void { + const l: u32 = @intCast(bytes.len); + try sock.sendAll(fd, std.mem.asBytes(&l)); + try sock.sendAll(fd, bytes); +} + +/// An ephemeral listening port for tests. +pub fn testListener() !struct { fd: sock.Fd, port: u16 } { + const l = try sock.listen(0); + var addr: std.posix.sockaddr.in = undefined; + var len: std.posix.socklen_t = @sizeOf(std.posix.sockaddr.in); + _ = std.c.getsockname(l, @ptrCast(&addr), &len); + return .{ .fd = l, .port = std.mem.bigToNative(u16, addr.port) }; +} + +fn testRank(cfg: Config, out: *[3][32]u8, err: *?anyerror) void { + var b = Bootstrap.open(cfg) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var id: [32]u8 = undefined; + b.broadcast(&@as([32]u8, @splat(0xAB)), &id) catch |e| { + err.* = e; + return; + }; + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const mine: [3]u8 = @splat(@intCast(cfg.rank)); + const all = b.allGather(mine[0 .. cfg.rank + 1], &buf, &blobs) catch |e| { + err.* = e; + return; + }; + for (all, 0..) |blob, i| @memcpy(out[i][0..blob.len], blob); + b.agree("same", "settings") catch |e| { + err.* = e; + }; +} + +test "three ranks gather in rank order and agree" { + const l = try testListener(); + defer sock.close(l.fd); + var outs: [3][3][32]u8 = undefined; + var errs: [3]?anyerror = @splat(null); + var threads: [2]std.Thread = undefined; + for (0..2) |i| threads[i] = try std.Thread.spawn(.{}, testRank, .{ Config{ .world = 3, .rank = @intCast(i + 1), .link_port = l.port, .connect_timeout_ns = 5 * std.time.ns_per_s }, &outs[i + 1], &errs[i + 1] }); + var b = try Bootstrap.acceptOn(l.fd, 3, 5 * std.time.ns_per_s); + defer b.close(); + var id: [32]u8 = undefined; + try b.broadcast(&@as([32]u8, @splat(0xAB)), &id); + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const all = try b.allGather(&.{0}, &buf, &blobs); + try b.agree("same", "settings"); + for (threads) |t| t.join(); + for (errs) |e| try std.testing.expect(e == null); + try std.testing.expectEqual(@as(usize, 3), all.len); + try std.testing.expectEqualSlices(u8, &.{ 2, 2, 2 }, all[2]); + try std.testing.expectEqualSlices(u8, &.{ 1, 1 }, outs[2][1][0..2]); +} + +test "settings that differ stop every rank" { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, err: *?anyerror) void { + var b = Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = port, .connect_timeout_ns = 5 * std.time.ns_per_s }) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree("knobs-b", "settings") catch |e| { + err.* = e; + }; + } + }; + var err: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, &err }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + try std.testing.expectError(error.Mismatch, b.agree("knobs-a", "settings")); + t.join(); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), err); +} diff --git a/zig/src/fabric/tp/collective.zig b/zig/src/fabric/tp/collective.zig new file mode 100644 index 0000000000..d2334dc74a --- /dev/null +++ b/zig/src/fabric/tp/collective.zig @@ -0,0 +1,136 @@ +//! The tensor-parallel communicator every backend implements: stream-ordered, graph-capturable, rank-order sums. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const DevicePtr = cuda.abi.DevicePtr; +pub const Stream = cuda.abi.Stream; + +pub const Error = error{ Aborted, Unsupported, Invalid, PeerFailed, Timeout, BackendFailed }; + +/// Element types the collectives move (`reduce` takes the float ones and the 32/64-bit integers). +pub const DType = enum(u8) { + u8, + i32, + i64, + f16, + bf16, + f32, + f64, + + pub fn size(t: DType) usize { + return switch (t) { + .u8 => 1, + .f16, .bf16 => 2, + .i32, .f32 => 4, + .i64, .f64 => 8, + }; + } +}; + +pub const Op = enum(u8) { sum, max, min }; + +/// Which implementation is behind a `Collective` (logs, bench labels). +pub const Kind = enum { nccl, mailbox, hybrid, roce, host }; + +pub const Collective = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + kind: Kind, + rank: *const fn (ptr: *anyopaque) u32, + world: *const fn (ptr: *anyopaque) u32, + /// recv[count] <- the element-wise reduction of every rank's send[count], added in rank order. + all_reduce: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void, + /// recv[world * count] <- every rank's send[count], rank 0's first (NCCL's layout). + all_gather: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, stream: Stream) Error!void, + /// Two ranks trade: this rank's send[count] lands in `peer`'s recv[count] and the peer's in ours (one group). + exchange: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + send: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + recv: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + /// Host: returns once every rank has called it and this rank's earlier collectives have finished. + barrier: *const fn (ptr: *anyopaque) Error!void, + /// Host: a failure the transport recorded since the last check (a peer that timed out, an NCCL async error). + check: *const fn (ptr: *anyopaque) Error!void, + /// Any thread: end pending and later work at once; safe while another thread is inside a call. + abort: *const fn (ptr: *anyopaque) void, + /// recv[world][stride]: rank r's first lens[r] bytes (device int32 [world]) land at recv + r * stride. + all_gather_v: ?*const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void = null, + /// Whether `all_gather_v` moves only the lengths for shards of `bytes` here (rank-local). Null: never. + var_fits: ?*const fn (ptr: *anyopaque, bytes: usize) bool = null, + /// Host: whether every rank's transport takes device lengths, agreed once per communicator. Null: false. + var_agreed: ?*const fn (ptr: *anyopaque) Error!bool = null, + }; + + pub fn kind(c: Collective) Kind { + return c.vtable.kind; + } + + pub fn rank(c: Collective) u32 { + return c.vtable.rank(c.ptr); + } + + pub fn world(c: Collective) u32 { + return c.vtable.world(c.ptr); + } + + pub fn allReduce(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void { + return c.vtable.all_reduce(c.ptr, src, dst, count, t, op, stream); + } + + pub fn allGather(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: Stream) Error!void { + return c.vtable.all_gather(c.ptr, src, dst, count, t, stream); + } + + /// `allGather` of `stride` elements a rank where only rank r's first lens[r] bytes matter (`VTable.all_gather_v`). + pub fn allGatherV(c: Collective, src: DevicePtr, dst: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void { + if (c.vtable.all_gather_v) |f| return f(c.ptr, src, dst, stride, lens, t, stream); + return c.vtable.all_gather(c.ptr, src, dst, stride, t, stream); + } + + /// Whether `allGatherV` of `bytes`-byte strides moves only the lengths here (callers agree first: `varAgreed`). + pub fn varFits(c: Collective, bytes: usize) bool { + const f = c.vtable.var_fits orelse return false; + return f(c.ptr, bytes); + } + + /// Every rank's transport takes device lengths: one agreement a communicator, called alike outside any graph. + pub fn varAgreed(c: Collective) Error!bool { + const f = c.vtable.var_agreed orelse return false; + return f(c.ptr); + } + + pub fn exchange(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.exchange(c.ptr, src, dst, count, t, peer, stream); + } + + pub fn send(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.send(c.ptr, buf, count, t, peer, stream); + } + + pub fn recv(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.recv(c.ptr, buf, count, t, peer, stream); + } + + pub fn barrier(c: Collective) Error!void { + return c.vtable.barrier(c.ptr); + } + + pub fn check(c: Collective) Error!void { + return c.vtable.check(c.ptr); + } + + pub fn abort(c: Collective) void { + c.vtable.abort(c.ptr); + } +}; + +/// Refuses a peer that is this rank or out of range. +pub fn checkPeer(rank: u32, world: u32, peer: u32) Error!void { + if (peer == rank or peer >= world) return error.Invalid; +} + +test "dtype sizes" { + try std.testing.expectEqual(@as(usize, 2), DType.bf16.size()); + try std.testing.expectEqual(@as(usize, 8), DType.i64.size()); +} diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig new file mode 100644 index 0000000000..442f82ed0e --- /dev/null +++ b/zig/src/fabric/tp/config.zig @@ -0,0 +1,256 @@ +//! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. +const std = @import("std"); +const log = @import("log.zig"); + +pub const Backend = enum { + /// NCCL for everything (the default). + nccl, + /// Exchanges of at most `mailbox_max_bytes` a rank through the one-shot mailbox kernel, NCCL above that. + mailbox, + /// The same split between two hosts: the small exchanges over RoCE (doorbell + RDMA proxy), NCCL above. + roce, +}; + +/// What `roce` does when its setup or load-time probe fails on any rank (TF_TP_ROCE_FALLBACK). +pub const Fallback = enum { nccl, @"error" }; + +pub const Error = error{BadSetting}; + +/// RoCE's per-write cap (roce_guard.max_write): on the roce backend a slot past it fails boot, never clamped. +pub const max_rdma_write: usize = @import("roce_guard.zig").max_write; + +/// TF_TP_PLAN_PIN: the CPU the plan threads pin to at their first plan (auto: the highest cpu_capacity). +pub const PlanPin = union(enum) { auto, cpu: u32 }; + +pub const Config = struct { + world: u32 = 2, + rank: u32 = 0, + /// Rank 0's address as the other ranks reach it. + master: []const u8 = "localhost", + /// The rendezvous port (TF_TP_PORT); the link listens on `link_port`. + port: u16 = 29500, + /// Rank 0 listens here: TF_TP_PLAN_PORT, else `port` + 7. + link_port: u16 = 29507, + /// How long rank 0 waits for the others and they keep retrying the connect (TF_TP_PLAN_TIMEOUT). + connect_timeout_ns: u64 = 300 * std.time.ns_per_s, + /// CUDA device ordinal this rank uses. + device: u32 = 0, + backend: Backend = .nccl, + mailbox_max_bytes: usize = 256 * 1024, + /// TF_TP_ROCE_CHUNKED_KB: `allGatherV` strides up to this go RoCE in ops of at most a slot (0: off). + roce_chunked_bytes: usize = 0, + /// A mailbox wait for the peer gives up after this long and poisons the transport (fail-stop). + mailbox_timeout_ns: u64 = 120 * std.time.ns_per_s, + failfast: bool = true, + /// After a failure: time for the HTTP threads to answer before exit 70. + grace_ns: u64 = std.time.ns_per_s, + /// Rank 1's memory report period on the fate channel (also how fast it notices rank 0 hanging up). + report_ns: u64 = std.time.ns_per_s / 2, + nccl_lib: []const u8 = "libnccl.so.2", + /// `name[:gid],...` RDMA devices for `roce` (null: active RoCE v2 ports on an IPv4 /24 the peer has, roce_pick). + roce_hca: ?[]const u8 = null, + /// Stripe over at most this many HCAs (auto detection only). + roce_hcas: u8 = 2, + roce_tc: u8 = 0, + /// TF_TP_ROCE_FAST as the RoCE kernel's `opts` word: bit 0 on, bits 8-15 pollers, bits 16-31 stagger ns (local). + roce_fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest RoCE grid, a power of two 1..16; launch shape only, the bytes are the same. + roce_blocks: u32 = 16, + /// Pin the proxy thread (it busy-spins one core between exchanges). + roce_cpu: ?u32 = null, + /// `shm` runs the RoCE protocol on one host without RDMA (tests); `verbs` is the real wire. + roce_wire: enum { verbs, shm } = .verbs, + roce_fallback: Fallback = .nccl, + /// TF_TP_PLAN_LINK=rdma: the follower polls the plan socket this long before blocking (TF_TP_PLAN_SPIN_US). + plan_spin_ns: u64 = 0, + plan_pin: ?PlanPin = null, + + /// Reads every setting from `get` (the environment in `fromEnv`); unset names keep the defaults. + pub fn parse(get: *const fn ([]const u8) ?[]const u8) Error!Config { + var c: Config = .{}; + c.world = try int(u32, get, "TF_TP_WORLD", c.world, 1, 64); + c.rank = try int(u32, get, "TF_TP_RANK", c.rank, 0, c.world - 1); + if (get("TF_TP_MASTER")) |m| c.master = m; + c.port = try int(u16, get, "TF_TP_PORT", c.port, 1, 65535 - 8); + c.link_port = try int(u16, get, "TF_TP_PLAN_PORT", c.port + 7, 1, 65535); + c.connect_timeout_ns = try seconds(get, "TF_TP_PLAN_TIMEOUT", 300, 0.1, 86400); + c.device = try int(u32, get, "TF_TP_DEVICE", c.device, 0, 63); + const backend = "TF_COMM_BACKEND"; + if (get(backend)) |b| c.backend = std.meta.stringToEnum(Backend, std.mem.trim(u8, b, " ")) orelse { + log.warn("{s}={s}: expected nccl, mailbox or roce", .{ backend, b }); + return error.BadSetting; + }; + c.mailbox_max_bytes = 1024 * try int(usize, get, "TF_TP_MAILBOX_MAX_KB", 256, 4, 16 * 1024); + if (c.backend == .roce and c.mailbox_max_bytes > max_rdma_write) { + log.warn("TF_TP_MAILBOX_MAX_KB={d}: a RoCE write is capped at {d} KiB; refusing to start", .{ c.mailbox_max_bytes / 1024, max_rdma_write / 1024 }); + return error.BadSetting; + } + c.roce_chunked_bytes = 1024 * try int(usize, get, "TF_TP_ROCE_CHUNKED_KB", 0, 0, 64 * 1024); + c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); + if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); + c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); + c.report_ns = try seconds(get, "TF_TP_PEER_MEM_S", 0.5, 0.01, 60); + if (get("TF_NCCL_LIB")) |p| c.nccl_lib = p; + c.roce_hca = get("TF_TP_ROCE_HCA"); + c.roce_hcas = try int(u8, get, "TF_TP_ROCE_HCAS", 2, 1, 2); + c.roce_tc = try int(u8, get, "TF_TP_ROCE_TC", try int(u8, get, "NCCL_IB_TC", 0, 0, 255), 0, 255); + if (try int(u8, get, "TF_TP_ROCE_FAST", 0, 0, 1) == 1) { + const pollers = try int(u32, get, "TF_TP_ROCE_POLLERS", 1, 1, 8); + const stagger = try int(u32, get, "TF_TP_ROCE_STAGGER_NS", 150, 0, 65535); + c.roce_fast = 1 | (pollers << 8) | (stagger << 16); + } + c.roce_blocks = try int(u32, get, "TF_TP_ROCE_BLOCKS", 16, 1, 16); + if (c.roce_blocks & (c.roce_blocks - 1) != 0) { + log.warn("TF_TP_ROCE_BLOCKS={d}: expected a power of two", .{c.roce_blocks}); + return error.BadSetting; + } + if (get("TF_TP_ROCE_CPU")) |_| c.roce_cpu = try int(u32, get, "TF_TP_ROCE_CPU", 0, 0, 4095); + if (get("TF_TP_ROCE_WIRE")) |w| c.roce_wire = std.meta.stringToEnum(@TypeOf(c.roce_wire), w) orelse { + log.warn("TF_TP_ROCE_WIRE={s}: expected verbs or shm", .{w}); + return error.BadSetting; + }; + const fallback = "TF_TP_ROCE_FALLBACK"; + if (get(fallback)) |f| c.roce_fallback = std.meta.stringToEnum(Fallback, f) orelse { + log.warn("{s}={s}: expected nccl or error", .{ fallback, f }); + return error.BadSetting; + }; + if (get("TF_TP_PLAN_LINK")) |l| { + const mode = std.mem.trim(u8, l, " "); + if (std.mem.eql(u8, mode, "rdma")) { + c.plan_spin_ns = 1000 * try int(u64, get, "TF_TP_PLAN_SPIN_US", 5000, 0, 10_000_000); + } else if (!std.mem.eql(u8, mode, "tcp") and !std.mem.eql(u8, mode, "nccl")) { + log.warn("TF_TP_PLAN_LINK={s}: expected nccl, tcp or rdma", .{l}); + return error.BadSetting; + } + } + if (get("TF_TP_PLAN_PIN")) |raw| { + const v = std.mem.trim(u8, raw, " "); + if (std.mem.eql(u8, v, "auto")) { + c.plan_pin = .auto; + } else if (!std.mem.eql(u8, v, "-1") and !std.mem.eql(u8, v, "off")) { + c.plan_pin = .{ .cpu = try int(u32, get, "TF_TP_PLAN_PIN", 0, 0, 4095) }; + } + } + return c; + } + + pub fn fromEnv() Error!Config { + return parse(&envGet); + } + + /// The settings both ranks must share (everything but rank, device and timeouts), as bytes to compare. + pub fn digest(c: Config) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(std.mem.asBytes(&c.world)); + h.update(std.mem.asBytes(&c.backend)); + h.update(std.mem.asBytes(&c.mailbox_max_bytes)); + h.update(std.mem.asBytes(&c.roce_chunked_bytes)); + h.update(std.mem.asBytes(&c.failfast)); + h.update(std.mem.asBytes(&c.roce_wire)); + h.update(std.mem.asBytes(&c.roce_fallback)); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; + +fn envGet(name: []const u8) ?[]const u8 { + var buf: [64]u8 = undefined; + if (name.len >= buf.len) return null; + @memcpy(buf[0..name.len], name); + buf[name.len] = 0; + const v = std.c.getenv(buf[0..name.len :0]) orelse return null; + const s = std.mem.span(v); + return if (s.len == 0) null else s; +} + +fn int(comptime T: type, get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: T, lo: T, hi: T) Error!T { + const raw = get(name) orelse return default; + const v = std.fmt.parseInt(T, std.mem.trim(u8, raw, " "), 10) catch { + log.warn("{s}={s}: expected an integer", .{ name, raw }); + return error.BadSetting; + }; + if (v < lo or v > hi) { + log.warn("{s}={d}: expected {d}..{d}", .{ name, v, lo, hi }); + return error.BadSetting; + } + return v; +} + +fn seconds(get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: f64, lo: f64, hi: f64) Error!u64 { + const v = if (get(name)) |raw| std.fmt.parseFloat(f64, std.mem.trim(u8, raw, " ")) catch { + log.warn("{s}={s}: expected seconds", .{ name, raw }); + return error.BadSetting; + } else default; + if (!(v >= lo and v <= hi)) { + log.warn("{s}={d}: expected {d}..{d} s", .{ name, v, lo, hi }); + return error.BadSetting; + } + return @intFromFloat(v * std.time.ns_per_s); +} + +const Fake = struct { + var pairs: []const [2][]const u8 = &.{}; + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "defaults: the plan link 7 ports past the rendezvous, NCCL, fail-fast on" { + Fake.pairs = &.{ .{ "TF_TP_RANK", "1" }, .{ "TF_TP_PORT", "29600" } }; + const c = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u32, 1), c.rank); + try std.testing.expectEqual(@as(u16, 29607), c.link_port); + try std.testing.expectEqual(Backend.nccl, c.backend); + try std.testing.expectEqual(std.time.ns_per_s, c.grace_ns); + try std.testing.expectEqual(std.time.ns_per_s / 2, c.report_ns); +} + +test "bad settings are refused and the digest ignores the rank" { + Fake.pairs = &.{.{ "TF_TP_RANK", "2" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{.{ "TF_COMM_BACKEND", "rdma" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_ROCE_HCA", "rocep1s0f1,roceP2p1s0f1" }, .{ "NCCL_IB_TC", "106" } }; + const r = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, r.backend); + try std.testing.expectEqual(@as(u8, 106), r.roce_tc); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", r.roce_hca.?); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "64" }, .{ "TF_TP_FAILFAST_GRACE_S", "0.25" } }; + const a = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(usize, 64 * 1024), a.mailbox_max_bytes); + try std.testing.expectEqual(std.time.ns_per_s / 4, a.grace_ns); + var b = a; + b.rank = 1; + try std.testing.expectEqualSlices(u8, &a.digest(), &b.digest()); + b.mailbox_max_bytes = 1024; + try std.testing.expect(!std.mem.eql(u8, &a.digest(), &b.digest())); +} + +test "the plan link's spin and pin, the RoCE grid cap" { + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" }, .{ "TF_TP_ROCE_FALLBACK", "error" }, .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_PIN", "auto" } }; + const p = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, p.backend); + try std.testing.expectEqual(@as(usize, 1024 * 1024), p.mailbox_max_bytes); + try std.testing.expectEqual(Fallback.@"error", p.roce_fallback); + try std.testing.expectEqual(@as(u64, 5000 * 1000), p.plan_spin_ns); + try std.testing.expectEqual(PlanPin.auto, p.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "tcp" }, .{ "TF_TP_PLAN_PIN", "7" } }; // tcp blocks at once; an explicit CPU + const o = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 0), o.plan_spin_ns); + try std.testing.expectEqual(PlanPin{ .cpu = 7 }, o.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_SPIN_US", "200" }, .{ "TF_TP_PLAN_PIN", "-1" } }; + const s = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 200 * 1000), s.plan_spin_ns); + try std.testing.expect(s.plan_pin == null); + Fake.pairs = &.{.{ "TF_TP_PLAN_LINK", "udp" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; // the RoCE grid cap: 16 unset, powers of two only + try std.testing.expectEqual(@as(u32, 16), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "4" }}; + try std.testing.expectEqual(@as(u32, 4), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "12" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; +} diff --git a/zig/src/fabric/tp/fate.zig b/zig/src/fabric/tp/fate.zig new file mode 100644 index 0000000000..5084c4648b --- /dev/null +++ b/zig/src/fabric/tp/fate.zig @@ -0,0 +1,261 @@ +//! Fail fast on every rank: a failure or a vanished rank ends all ranks with exit 70 about a second later. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const posix = std.posix; +/// Linux POLLRDHUP: the peer shut its side down (std has no name for it). +const pollrdhup: i16 = 0x2000; + +pub const exit_code: u8 = 70; +const magic: u32 = 0x46463431; +const kind_mem: u32 = 1; +const kind_fail: u32 = 2; +/// rank 0's clean stop: the hang-up that follows is no failure (sent before rank 0 closes its sockets) +const kind_stop: u32 = 3; +const Frame = extern struct { magic: u32, kind: u32, value: i64 }; + +pub const max_ranks = 64; +pub const max_hooks = 8; + +/// A callback with its context (an abort of a collective, rank 0's request failing); `reason` lives until exit. +pub const Hook = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, reason: []const u8) void }; + +pub const Fate = struct { + rank: u32, + world: u32, + fds: [max_ranks]sock.Fd, + grace_ns: u64, + report_ns: u64, + /// This rank's usable memory in bytes for the MEM report (default: MemAvailable). + usable: *const fn () u64 = memAvailable, + /// The process exit; tests replace it (it may return there). + exit_fn: *const fn (u8) void = exitNow, + hooks: [max_hooks]Hook = undefined, + n_hooks: std.atomic.Value(u32) = .init(0), + dead: std.atomic.Value(bool) = .init(false), + closing: std.atomic.Value(bool) = .init(false), + /// Set once the exit was asked for (tests read it; the process is gone in serving). + exited: std.atomic.Value(bool) = .init(false), + reason_buf: [512]u8 = undefined, + reason_len: usize = 0, + send_lock: std.atomic.Value(bool) = .init(false), + /// Rank 0: each rank's last report (bytes) and its monotonic time. + peer_mem: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + peer_mem_at: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + wake: [2]sock.Fd = .{ -1, -1 }, + thread: ?std.Thread = null, + + /// The channel over the bootstrap's sockets (`Bootstrap.release`); `start` runs its thread. + pub fn init(cfg: Config, fds: [max_ranks]sock.Fd) Fate { + return .{ .rank = cfg.rank, .world = cfg.world, .fds = fds, .grace_ns = cfg.grace_ns, .report_ns = cfg.report_ns }; + } + + pub fn start(f: *Fate) !void { + var p: [2]c_int = undefined; + if (std.c.pipe2(&p, .{ .CLOEXEC = true }) != 0) return error.SocketFailed; + f.wake = p; + f.thread = if (f.rank == 0) try std.Thread.spawn(.{}, watch0, .{f}) else try std.Thread.spawn(.{}, watchK, .{f}); + } + + /// Runs `hook` when the ranks go down; collectives register their abort here. + pub fn onFail(f: *Fate, hook: Hook) void { + const i = f.n_hooks.load(.acquire); + std.debug.assert(i < max_hooks); + f.hooks[i] = hook; + f.n_hooks.store(i + 1, .release); + } + + /// This rank failed (`where`: what it was doing): every rank goes down. Returns only in tests. + pub fn fatal(f: *Fate, where: []const u8, err: anyerror) void { + var buf: [256]u8 = undefined; + const msg = std.fmt.bufPrint(&buf, "rank {d} failed in {s}: {t}", .{ f.rank, where, err }) catch "rank failed"; + f.die(msg, true); + } + + /// This rank failed while booting: every rank goes down, and the thread is joined before the caller frees this. + pub fn failBoot(f: *Fate, where: []const u8, err: anyerror) void { + f.fatal(where, err); + f.close(); + } + + /// A normal stop is coming (rank 0 before it tells the others; the others when told): later hang-ups are fine. + pub fn expectStop(f: *Fate) void { + f.closing.store(true, .release); + } + + /// Rank 0's clean stop of every rank: a STOP frame to each ahead of its close, so followers end quietly. + pub fn announceStop(f: *Fate) void { + if (f.closing.swap(true, .acq_rel)) return; + if (f.rank != 0) return; + const h: Frame = .{ .magic = magic, .kind = kind_stop, .value = 0 }; + f.lock(); + defer f.unlock(); + for (f.fds[1..f.world]) |fd| if (fd >= 0) { + _ = sock.trySend(fd, std.mem.asBytes(&h)); + }; + } + + /// A normal stop: the thread ends and the sockets close. + pub fn close(f: *Fate) void { + f.closing.store(true, .release); + if (f.wake[1] >= 0) _ = std.c.write(f.wake[1], "x", 1); + if (f.thread) |t| t.join(); + f.thread = null; + for (&f.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + for (f.wake) |fd| if (fd >= 0) sock.close(fd); + f.wake = .{ -1, -1 }; + } + + /// Rank 0: rank `r`'s last reported usable bytes and the age of the report, or null before its first. + pub fn peerMem(f: *const Fate, r: u32) ?struct { bytes: u64, age_ns: u64 } { + const at = f.peer_mem_at[r].load(.acquire); + if (at == 0) return null; + return .{ .bytes = f.peer_mem[r].load(.acquire), .age_ns = sock.nowNs() -| at }; + } + + pub fn reason(f: *const Fate) []const u8 { + return f.reason_buf[0..f.reason_len]; + } + + fn lock(f: *Fate) void { + while (f.send_lock.swap(true, .acquire)) std.atomic.spinLoopHint(); + } + + fn unlock(f: *Fate) void { + f.send_lock.store(false, .release); + } + + fn die(f: *Fate, why: []const u8, own: bool) void { + if (f.dead.swap(true, .acq_rel)) { + // the first caller exits the process; a second (the fate thread or a round thread) waits for it + while (!f.exited.load(.acquire)) sock.sleepNs(10 * std.time.ns_per_ms); + return; + } + const n = @min(why.len, f.reason_buf.len); + @memcpy(f.reason_buf[0..n], why[0..n]); + f.reason_len = n; + const r = f.reason(); + log.warn("[tensorfold] tp fail-fast: rank {d}: {s}; exiting ({d}) in {d} ms so every rank restarts", .{ f.rank, r, exit_code, f.grace_ns / std.time.ns_per_ms }); + // aborts first: kernels waiting on the dead peer end now; a hook that blocks cannot hold the exit + if (std.Thread.spawn(.{}, runHooks, .{ f, r })) |t| t.detach() else |_| runHooks(f, r); + if (!f.closing.load(.acquire)) { + const head: Frame = .{ .magic = magic, .kind = kind_fail, .value = @intCast(r.len) }; + f.lock(); + for (f.fds[0..f.world]) |fd| if (fd >= 0) { + if (own or f.rank == 0) { + _ = sock.trySend(fd, std.mem.asBytes(&head)); + _ = sock.trySend(fd, r); + } + sock.hangup(fd); + }; + f.unlock(); + } + sock.sleepNs(f.grace_ns); + f.exited.store(true, .release); + f.exit_fn(exit_code); + } + + fn runHooks(f: *Fate, r: []const u8) void { + const n = f.n_hooks.load(.acquire); + for (f.hooks[0..n]) |h| h.call(h.ctx, r); + } + + /// Reads one frame from `fd`: a MEM report, a FAIL (the ranks go down), or the end (the peer is gone). + fn readFrame(f: *Fate, peer: u32, fd: sock.Fd) bool { + var h: Frame = undefined; + const deadline = sock.nowNs() + std.time.ns_per_s; + sock.recvExact(fd, std.mem.asBytes(&h), deadline) catch { + if (f.closing.load(.acquire)) return false; + var buf: [128]u8 = undefined; + f.die(std.fmt.bufPrint(&buf, "rank {d} closed the fate channel (it failed, exited or was killed)", .{peer}) catch "peer gone", false); + return false; + }; + if (h.magic != magic) { + f.die("a frame that is not the fate channel's", false); + return false; + } + if (h.kind == kind_stop) { + // rank 0 stops every rank: its hang-up is next, and no failure + f.closing.store(true, .release); + return false; + } + if (h.kind == kind_mem) { + f.peer_mem[peer].store(@intCast(@max(h.value, 0)), .release); + f.peer_mem_at[peer].store(sock.nowNs(), .release); + return true; + } + var text: [512]u8 = undefined; + const len: usize = @intCast(std.math.clamp(h.value, 0, text.len)); + sock.recvExact(fd, text[0..len], deadline) catch {}; + f.die(if (len > 0) text[0..len] else "a peer failed", false); + return false; + } + + fn watch0(f: *Fate) void { + var p: [max_ranks + 1]std.c.pollfd = undefined; + var who: [max_ranks + 1]u32 = undefined; + var n: usize = 0; + for (f.fds[1..f.world], 1..) |fd, r| if (fd >= 0) { + p[n] = .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }; + who[n] = @intCast(r); + n += 1; + }; + p[n] = .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }; + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + if (std.c.poll(&p, @intCast(n + 1), -1) <= 0) continue; + for (p[0..n], who[0..n]) |*q, r| if (q.revents != 0) { + if (!f.readFrame(r, q.fd)) return; + }; + } + } + + fn watchK(f: *Fate) void { + const fd = f.fds[0]; + var p = [_]std.c.pollfd{ + .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }, + .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }, + }; + const ms: c_int = @intCast(@max(f.report_ns / std.time.ns_per_ms, 1)); + f.report(fd); + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + const ready = std.c.poll(&p, 2, ms); + if (ready == 0) { + f.report(fd); + continue; + } + if (ready > 0 and p[0].revents != 0 and !f.readFrame(0, fd)) return; + } + } + + /// One MEM frame, never blocking (a full buffer drops it: the next comes one period later). + fn report(f: *Fate, fd: sock.Fd) void { + const h: Frame = .{ .magic = magic, .kind = kind_mem, .value = @intCast(f.usable()) }; + f.lock(); + defer f.unlock(); + _ = sock.trySend(fd, std.mem.asBytes(&h)); + } +}; + +fn exitNow(code: u8) void { + std.c._exit(code); +} + +/// MemAvailable from /proc/meminfo in bytes (0 when unreadable). +pub fn memAvailable() u64 { + const fd = std.c.open("/proc/meminfo", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return 0; + defer _ = std.c.close(fd); + var buf: [4096]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return 0; + const text = buf[0..@intCast(n)]; + const key = "MemAvailable:"; + const at = std.mem.indexOf(u8, text, key) orelse return 0; + var it = std.mem.tokenizeAny(u8, text[at + key.len ..], " \n"); + const kb = std.fmt.parseInt(u64, it.next() orelse return 0, 10) catch return 0; + return kb * 1024; +} diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig new file mode 100644 index 0000000000..728f4bfdf3 --- /dev/null +++ b/zig/src/fabric/tp/fate_test.zig @@ -0,0 +1,241 @@ +//! The fate channel between two ranks over real TCP: a failure or a vanished rank ends both, a stop ends neither. +const std = @import("std"); +const sock = @import("sock.zig"); +const fate = @import("fate.zig"); +const bootstrap = @import("bootstrap.zig"); +const Config = @import("config.zig").Config; +const Fate = fate.Fate; + +var exits: std.atomic.Value(u32) = .init(0); +var hooks: std.atomic.Value(u32) = .init(0); + +fn recordExit(code: u8) void { + std.debug.assert(code == fate.exit_code); + _ = exits.fetchAdd(1, .acq_rel); +} + +fn hook(_: *anyopaque, reason: []const u8) void { + std.debug.assert(reason.len > 0); + _ = hooks.fetchAdd(1, .acq_rel); +} + +const grace = 50 * std.time.ns_per_ms; + +/// Two connected ranks' fate channels (rank 0 first), started, with recorded exits and hooks. +fn pair(fates: *[2]Fate) !void { + return pairAt(.{ &fates[0], &fates[1] }); +} + +fn pairAt(fates: [2]*Fate) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var fds0: [fate.max_ranks]sock.Fd = @splat(-1); + var fds1: [fate.max_ranks]sock.Fd = @splat(-1); + fds0[1] = b; + fds1[0] = a; + for (fates, [2][fate.max_ranks]sock.Fd{ fds0, fds1 }, 0..) |f, fds, r| { + f.* = Fate.init(.{ .world = 2, .rank = @intCast(r), .grace_ns = grace, .report_ns = 20 * std.time.ns_per_ms }, fds); + f.exit_fn = recordExit; + f.usable = fakeUsable; + f.onFail(.{ .ctx = f, .call = hook }); + } + for (fates) |f| try f.start(); +} + +fn fakeUsable() u64 { + return 123 << 20; +} + +fn waitExits(want: u32) !u64 { + const t0 = sock.nowNs(); + while (exits.load(.acquire) < want) { + if (sock.nowNs() - t0 > 5 * std.time.ns_per_s) return error.Timeout; + sock.sleepNs(std.time.ns_per_ms); + } + return sock.nowNs() - t0; +} + +fn reset() void { + exits.store(0, .release); + hooks.store(0, .release); +} + +test "rank 1's failure ends both ranks with its reason" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + const mem = fates[0].peerMem(1) orelse return error.NoReport; + try std.testing.expectEqual(@as(u64, 123 << 20), mem.bytes); + fates[1].fatal("prefill", error.OutOfMemory); + const took = try waitExits(2); + try std.testing.expect(took < grace + 500 * std.time.ns_per_ms); + try std.testing.expectEqual(@as(u32, 2), hooks.load(.acquire)); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 failed in prefill: OutOfMemory") != null); +} + +test "rank 0's failure ends rank 1" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + fates[0].fatal("window", error.Unexpected); + _ = try waitExits(2); + try std.testing.expect(std.mem.indexOf(u8, fates[1].reason(), "rank 0 failed in window") != null); +} + +test "a rank that vanishes without a word ends the other" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + // rank 1 killed: its socket closes with no FAIL frame (its own fate thread ends quietly first) + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); +} + +test "a normal stop ends nobody" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + fates[0].expectStop(); + fates[1].expectStop(); + fates[1].close(); + fates[0].close(); + sock.sleepNs(2 * grace); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); +} + +test "a rank failing at boot ends both with its reason, and its channel is down before its owner frees it" { + // a follower's boot failed after its session came up; its Fate is freed while the fate thread could still read it + reset(); + // rank 1's Fate where a Model keeps it: on the heap, its pages unmapped when freed (a stale read faults) + const pa = std.heap.page_allocator; + var f0: Fate = undefined; + const f1 = try pa.create(Fate); + try pairAt(.{ &f0, f1 }); + defer f0.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + f1.failBoot("boot", error.RopeTooShort); + try std.testing.expect(f1.thread == null); + pa.destroy(f1); + _ = try waitExits(2); + // several report periods with rank 1's Fate gone: its thread would have reported from freed memory + sock.sleepNs(5 * 20 * std.time.ns_per_ms); + try std.testing.expect(std.mem.indexOf(u8, f0.reason(), "rank 1 failed in boot: RopeTooShort") != null); +} + +test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order the ranks close in" { + // rank 0 ended its run and closed; the follower, still tearing down, must not see the hang-up as a failure + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + sock.sleepNs(60 * std.time.ns_per_ms); + fates[0].announceStop(); + fates[0].close(); // rank 0 gone before rank 1 even started its own stop + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + fates[1].close(); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +// -- a server's drain on two ranks: the plan link and the fate channel in the order a family's stop uses them -- + +const planlink = @import("planlink.zig"); + +/// Plan words a family's rank 0 sends: a round's operation, and a stop word that ends a follower's loop. +const op_round: i64 = 1; +const op_stop: i64 = 4; + +fn linkPair(links: *[2]planlink.PlanLink) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [planlink.max_ranks]sock.Fd = @splat(-1); + var f1: [planlink.max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + links[0] = planlink.PlanLink.init(0, 2, f0); + links[1] = planlink.PlanLink.init(1, 2, f1); +} + +/// Rank 1's loop: each round as it comes, until the stop word, or until the link ends. +const Follower = struct { + link: *planlink.PlanLink, + rounds: u32 = 0, + ended: enum { running, stopped, peer_closed, failed } = .running, + + fn run(f: *Follower) void { + var buf: [16]i64 = undefined; + while (true) { + const msg = f.link.recv(&buf) catch |e| { + f.ended = if (e == error.PeerClosed) .peer_closed else .failed; + return; + }; + if (msg.len == 0) continue; + if (msg[0] == op_stop) { + f.ended = .stopped; + return; + } + if (msg[0] == op_round) f.rounds += 1; + } + } +}; + +test "a drain on rank 0, then its clean stop at a round boundary: rank 1 ends its loop there and nobody exits 70" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + var follower: Follower = .{ .link = &links[1] }; + const t = try std.Thread.spawn(.{}, Follower.run, .{&follower}); + // serving, then draining (rank 0 refuses new work, rounds go on; the fate channel is untouched), past a grace + var sent: u32 = 0; + for (0..30) |_| { + try links[0].send(&.{ op_round, sent }); + sent += 1; + sock.sleepNs(5 * std.time.ns_per_ms); + } + try std.testing.expect(!fates[0].closing.load(.acquire)); + // the stop's order: the stop word first, so the follower ends its loop before the fate STOP and the hang-up + try links[0].send(&.{op_stop}); + fates[0].announceStop(); + links[0].close(); + fates[0].close(); + t.join(); + // rank 1's loop returned: its model closes and the process exits 0 + fates[1].close(); + links[1].close(); + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + try std.testing.expectEqual(.stopped, follower.ended); + try std.testing.expectEqual(sent, follower.rounds); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +test "a drain on rank 0 does not disarm fail-fast: rank 1 killed mid-drain still ends rank 0 (exit 70)" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + defer links[0].close(); + for (0..5) |i| try links[0].send(&.{ op_round, @intCast(i) }); + // rank 0 is draining (no stop announced); rank 1 is killed: its sockets close with no word + links[1].close(); + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); + try std.testing.expectEqual(@as(u32, 1), exits.load(.acquire)); +} diff --git a/zig/src/fabric/tp/gpu_bench.zig b/zig/src/fabric/tp/gpu_bench.zig new file mode 100644 index 0000000000..99831e6d77 --- /dev/null +++ b/zig/src/fabric/tp/gpu_bench.zig @@ -0,0 +1,82 @@ +//! Latency and bandwidth of each backend's collectives, eager and graphed; RESULT lines are JSON, busbw as nccl-tests. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Ctx = @import("gpu_checks.zig").Ctx; +const Collective = tp.Collective; + +pub const Kind = enum { all_reduce, all_gather }; + +fn issue(c: *Ctx, comm: Collective, kind: Kind, bytes: usize) !void { + const count = bytes / 2; + switch (kind) { + .all_reduce => try comm.allReduce(c.in.ptr, c.out.ptr, count, .bf16, .sum, c.stream.handle), + .all_gather => try comm.allGather(c.in.ptr, c.out.ptr, count, .bf16, c.stream.handle), + } +} + +/// GPU microseconds a call, `iters` calls issued back to back on the stream. +fn eager(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, iters: u32) !f64 { + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + for (0..10) |_| try issue(c, comm, kind, bytes); + try c.stream.synchronize(); + try comm.barrier(); + try e0.record(c.stream); + for (0..iters) |_| try issue(c, comm, kind, bytes); + try e1.record(c.stream); + try e1.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(iters)); +} + +/// GPU microseconds a call inside one graph of `per_graph` calls, replayed `replays` times. +fn graphed(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, per_graph: u32, replays: u32) !f64 { + try cuda.graph.beginCapture(c.stream, .thread_local); + const err: ?tp.collective.Error = blk: { + for (0..per_graph) |_| issue(c, comm, kind, bytes) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (err) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + try exec.upload(c.stream); + try exec.launchOn(c.stream); + try c.stream.synchronize(); + try comm.barrier(); + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + try e0.record(c.stream); + for (0..replays) |_| try exec.launchOn(c.stream); + try e1.record(c.stream); + try e1.synchronize(); + try comm.check(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(per_graph * replays)); +} + +/// Every power of two from `lo` to `hi` bytes a rank, both kinds, eager and graphed. +pub fn sweep(c: *Ctx, comm: Collective, label: []const u8, lo: usize, hi: usize) !void { + for ([_]Kind{ .all_reduce, .all_gather }) |kind| { + var bytes = lo; + while (bytes <= hi) : (bytes *= 2) { + const per_graph: u32 = if (bytes >= 1 << 20) 20 else 100; + const e = try eager(c, comm, kind, bytes, if (bytes >= 1 << 20) 50 else 200); + const g = try graphed(c, comm, kind, bytes, per_graph, 5); + const moved: f64 = @floatFromInt(if (kind == .all_gather) 2 * bytes else bytes); + const factor: f64 = if (kind == .all_gather) 0.5 else 1.0; + const busbw = moved / (g * 1e-6) * factor / 1e9; + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"{t}\",\"bytes\":{d},\"eager_us\":{d:.2},\"graph_us\":{d:.2},\"busbw_GBps\":{d:.2}}}\n", .{ label, kind, bytes, e, g, busbw }); + } + } +} + +/// One decode window's exchanges: `n` all-gathers of one row's bf16 partials (`row_bytes` a rank) in one graph. +pub fn decodeWindow(c: *Ctx, comm: Collective, label: []const u8, n: u32, row_bytes: usize) !void { + const per = try graphed(c, comm, .all_gather, row_bytes, n, 20); + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"decode_window\",\"exchanges\":{d},\"bytes\":{d},\"per_exchange_us\":{d:.2},\"window_us\":{d:.1}}}\n", .{ label, n, row_bytes, per, per * @as(f64, @floatFromInt(n)) }); +} diff --git a/zig/src/fabric/tp/gpu_checks.zig b/zig/src/fabric/tp/gpu_checks.zig new file mode 100644 index 0000000000..5dc4ed6d4c --- /dev/null +++ b/zig/src/fabric/tp/gpu_checks.zig @@ -0,0 +1,188 @@ +//! Every collective on real GPUs against the host reference, eager and replayed from a graph. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Collective = tp.Collective; +const DType = tp.DType; +const Op = tp.Op; +const reference = tp.reference; + +pub const Ctx = struct { + d: *const cuda.Driver, + gpa: std.mem.Allocator, + stream: cuda.Stream, + rank: u32, + /// Device buffers sized for the largest case: input, output (world x), a spare. + in: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + /// Host scratch: this rank's input, the peer's, the expected and the downloaded output. + mine: []u8, + peer: []u8, + want: []u8, + got: []u8, + failures: u32 = 0, + passes: u32 = 0, + + pub fn init(d: *const cuda.Driver, gpa: std.mem.Allocator, rank: u32, max_bytes: usize) !Ctx { + return .{ + .d = d, + .gpa = gpa, + .stream = try cuda.Stream.init(d, true), + .rank = rank, + .in = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .out = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .mine = try gpa.alloc(u8, max_bytes), + .peer = try gpa.alloc(u8, max_bytes), + .want = try gpa.alloc(u8, 2 * max_bytes), + .got = try gpa.alloc(u8, 2 * max_bytes), + }; + } + + pub fn deinit(c: *Ctx) void { + c.stream.deinit(); + c.in.free(); + c.out.free(); + for ([_][]u8{ c.mine, c.peer, c.want, c.got }) |b| c.gpa.free(b); + } + + fn verdict(c: *Ctx, ok: bool, comptime fmt: []const u8, args: anytype) void { + if (ok) c.passes += 1 else c.failures += 1; + std.debug.print("{s} " ++ fmt ++ "\n", .{if (ok) "PASS" else "FAIL"} ++ args); + } + + /// Both ranks' inputs for case `seed`: this rank's into `mine`, the peer's into `peer` (n bytes each). + fn inputs(c: *Ctx, n: usize, t: DType, seed: u64) void { + reference.fill(c.mine[0..n], t, c.rank, seed); + reference.fill(c.peer[0..n], t, 1 - c.rank, seed); + } + + fn rankInput(c: *Ctx, r: u32, n: usize) []const u8 { + return if (r == c.rank) c.mine[0..n] else c.peer[0..n]; + } + + fn compare(c: *Ctx, comm: Collective, label: []const u8, what: []const u8, t: DType, n: usize, out_off: usize, want: []const u8) !void { + try c.stream.synchronize(); + comm.check() catch |e| { + c.verdict(false, "{s} {s} {t} {d} B: transport error {t}", .{ label, what, t, n, e }); + return; + }; + try c.out.download(out_off, c.got[0..want.len]); + const ok = std.mem.eql(u8, c.got[0..want.len], want); + var first: usize = 0; + if (!ok) first = std.mem.indexOfDiff(u8, c.got[0..want.len], want) orelse 0; + c.verdict(ok, "{s} {s} {t} {d} B a rank{s}", .{ label, what, t, n, if (ok) "" else " (first difference at a byte shown below)" }); + if (!ok) std.debug.print(" byte {d}: got {x} want {x}\n", .{ first, c.got[first], want[first] }); + } + + /// all-reduce (out of place, in place, misaligned by `skew` bytes), all-gather (out of place, in place), exchange. + pub fn collectives(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const reduce_types = [_]DType{ .f32, .bf16, .f16, .i32 }; + for (reduce_types, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + const count = n / t.size(); + c.inputs(n, t, seed + ti); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, t, .sum); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum", t, n, 0, c.want[0..n]); + // in place, at an offset of one element (not 16-byte aligned for 2-byte types) + const skew = t.size(); + try c.out.uploadAsync(skew, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.out.ptr + skew, c.out.ptr + skew, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum in place, skewed", t, n, skew, c.want[0..n]); + } + for ([_]Op{ .max, .min }) |op| { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed + 11); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .f32, op); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, n / 4, .f32, op, c.stream.handle); + try c.compare(comm, label, if (op == .max) "all-reduce max" else "all-reduce min", .f32, n, 0, c.want[0..n]); + } + for ([_]DType{ .u8, .bf16, .f32 }, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + c.inputs(n, t, seed + 20 + ti); + @memcpy(c.want[0..n], c.rankInput(0, n)); + @memcpy(c.want[n..][0..n], c.rankInput(1, n)); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allGather(c.in.ptr, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather", t, n, 0, c.want[0 .. 2 * n]); + try c.out.fill8(0, c.stream.handle); + try c.out.uploadAsync(c.rank * n, c.mine[0..n], c.stream.handle); + try comm.allGather(c.out.ptr + c.rank * n, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather in place", t, n, 0, c.want[0 .. 2 * n]); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.exchange(c.in.ptr, c.out.ptr, n / t.size(), t, 1 - c.rank, c.stream.handle); + try c.compare(comm, label, "exchange", t, n, 0, c.peer[0..n]); + } + } + + /// send then recv one way, then the other way (NCCL behind every backend). + pub fn pointToPoint(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + for (0..2) |dir| { + const sender: u32 = @intCast(dir); + if (c.rank == sender) { + try comm.send(c.in.ptr, n / 4, .f32, 1 - c.rank, c.stream.handle); + try c.stream.synchronize(); + } else { + try comm.recv(c.out.ptr, n / 4, .f32, sender, c.stream.handle); + try c.compare(comm, label, "send/recv", .f32, n, 0, c.peer[0..n]); + } + } + } + + /// A captured chain replayed with fresh inputs, eager calls between replays (the mailbox epoch spans both). + pub fn graphs(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, replays: u32) !void { + const n = n_bytes - n_bytes % 16; + const half = n / 2; // bf16 elements + try cuda.graph.beginCapture(c.stream, .thread_local); + // x1 = sum(x0) [bf16], x2 = gather(x1), x3 = sum(x2) [f32 view], x4 = exchange(x3) + const x0 = c.in.ptr; + const x1 = c.out.ptr; + const x2 = c.in.ptr + n; + const x3 = c.out.ptr + n; + const x4 = c.out.ptr + 3 * n; + const ops: ?tp.collective.Error = blk: { + comm.allReduce(x0, x1, half, .bf16, .sum, c.stream.handle) catch |e| break :blk e; + comm.allGather(x1, x2, half, .bf16, c.stream.handle) catch |e| break :blk e; + comm.allReduce(x2, x3, n / 2, .f32, .sum, c.stream.handle) catch |e| break :blk e; + comm.exchange(x3, x4, n / 2, .f32, 1 - c.rank, c.stream.handle) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (ops) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + const nodes = try g.nodeCount(); + const want1 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want1); + const want3 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want3); + for (0..replays) |k| { + c.inputs(n, .bf16, 1000 + k); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + reference.reduce(want1[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + @memcpy(want1[n..][0..n], want1[0..n]); // the gather of equal sums: twice the same + reference.reduce(want3[0 .. 2 * n], &.{ want1[0 .. 2 * n], want1[0 .. 2 * n] }, .f32, .sum); + try exec.launchOn(c.stream); + try c.compare(comm, label, "graph replay: sum, gather, sum, exchange", .bf16, n, n, want3[0 .. 2 * n]); + try c.compare(comm, label, "graph replay: exchange of the sum", .f32, 2 * n, 3 * n, want3[0 .. 2 * n]); + // an eager call between replays + try c.collectivesQuick(comm, label, k); + } + std.debug.print(" ({s}: {d} graph nodes, {d} replays)\n", .{ label, nodes, replays }); + } + + fn collectivesQuick(c: *Ctx, comm: Collective, label: []const u8, k: usize) !void { + const n = 4096; + c.inputs(n, .bf16, 5000 + k); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + try c.in.uploadAsync(5 * n, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr + 5 * n, c.out.ptr + 5 * n, n / 2, .bf16, .sum, c.stream.handle); + try c.compare(comm, label, "eager between replays", .bf16, n, 5 * n, c.want[0..n]); + } +}; diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig new file mode 100644 index 0000000000..82fc1edd5b --- /dev/null +++ b/zig/src/fabric/tp/gpu_main.zig @@ -0,0 +1,152 @@ +//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel`: one rank of TP=2. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const tp_kernel = @import("tp_kernel"); +const checks = @import("gpu_checks.zig"); +const bench = @import("gpu_bench.zig"); + +const max_bytes = (4 << 20) + 4096; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n", .{}); + return 2; + } + // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first + if (std.mem.eql(u8, args[1], "kernel")) { + std.debug.print("mailbox kernel {d} B\n", .{tp_kernel.mailbox.len}); + return if (tp_kernel.mailbox.len > 0) 0 else 1; + } + const cfg = try tp.Config.fromEnv(); + if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + if (std.mem.eql(u8, args[1], "hca")) return hca(cfg, init.io); + if (std.mem.eql(u8, args[1], "failfast") and (!cfg.failfast or args.len > 5) and !hostOnlyAllowed(cfg)) { + std.debug.print("REFUSED: failfast with TF_TP_FAILFAST=0 or a hold is host-only: TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire\n", .{}); + return 2; + } + if (cfg.backend != .nccl and tp_kernel.mailbox.len == 0) { + std.debug.print("FAIL TF_COMM_BACKEND={t} needs the mailbox kernel: rebuild with -Dtp-nvcc or -Dtp-fatbin\n", .{cfg.backend}); + return 2; + } + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, @intCast(cfg.device)); + defer ctx.deinit(); + var name_buf: [128]u8 = undefined; + const t0 = tp.sock.nowNs(); + var s: tp.Session = undefined; + try s.open(&driver, cfg, tp_kernel.mailbox); + const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; + std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); + if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); + if (std.mem.eql(u8, args[1], "failfast") and args.len >= 4 and std.mem.eql(u8, args[3], "transport")) { + const n = if (args.len > 4) try std.fmt.parseInt(u32, args[4], 10) else 50; + const hold = if (args.len > 5) try std.fmt.parseInt(u32, args[5], 10) else 0; + return @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), n, hold); + } + if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); + std.debug.print("unknown command {s}\n", .{args[1]}); + return 2; +} + +/// `hca` (host only, no GPU): both ranks' RoCE pick over the bootstrap; this rank's `NCCL_IB_HCA=...` line on stdout. +fn hca(cfg: tp.Config, io: std.Io) !u8 { + var b = tp.Bootstrap.open(cfg) catch |e| { + const s = @as(f64, @floatFromInt(cfg.connect_timeout_ns)) / std.time.ns_per_s; + switch (e) { + // a peer that never runs `hca` (it names its ports, or is not up) leaves this rank waiting until here + error.Timeout => std.debug.print("FAIL hca rank {d}: no peer joined within TF_TP_PLAN_TIMEOUT ({d:.1} s): is its runner up, and leaving TF_TP_ROCE_HCA and NCCL_IB_HCA unset as this one does?\n", .{ cfg.rank, s }), + error.AddressInUse => std.debug.print("FAIL hca rank {d}: port {d} (TF_TP_PORT + 7) is in use on this host\n", .{ cfg.rank, cfg.link_port }), + error.Closed => std.debug.print("FAIL hca rank {d}: the peer closed the bootstrap connection\n", .{cfg.rank}), + else => std.debug.print("FAIL hca rank {d}: bootstrap: {t}\n", .{ cfg.rank, e }), + } + return 1; + }; + defer b.close(); + var specs: [tp.roce_verbs.max_hcas]tp.roce_verbs.HcaSpec = undefined; + const pick = tp.roce_pick.select(cfg, &b, &specs) catch |e| { + std.debug.print("FAIL hca rank {d}: {t} while picking the ports with the peer\n", .{ cfg.rank, e }); + return 1; + }; + var buf: [tp.roce_pick.list_bytes + 16]u8 = undefined; + const line = tp.roce_pick.hcaLine(pick, &specs, &buf) catch |e| { + const why = switch (e) { + error.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host", + error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's", + }; + std.debug.print("FAIL hca rank {d}: {s}: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA\n", .{ cfg.rank, why }); + return 1; + }; + var names: [tp.roce_pick.list_bytes]u8 = undefined; + const how = if (pick.explicit) "TF_TP_ROCE_HCA" else "on a /24 the peer shares"; + std.debug.print("INFO hca rank {d}: {s} ({s})\n", .{ cfg.rank, try tp.roce_pick.joinNames(specs[0..pick.n], true, &names), how }); + try std.Io.File.stdout().writeStreamingAll(io, line); + return 0; +} + +fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, max_bytes); + defer c.deinit(); + var backends: [2]struct { comm: tp.Collective, label: []const u8 } = undefined; + var n_backends: usize = 1; + backends[0] = .{ .comm = s.nccl.iface(), .label = "nccl" }; + if (s.comm().kind() != .nccl) { + backends[1] = .{ .comm = s.comm(), .label = if (s.roce != null) "roce" else "mailbox" }; + n_backends = 2; + } + const sizes = [_]usize{ 6, 1000, 4096, 65536, 256 * 1024, 256 * 1024 + 16, 1 << 20, 4 << 20 }; + for (backends[0..n_backends]) |b| { + for (sizes, 0..) |n, i| try c.collectives(b.comm, b.label, n, 100 * i); + try c.pointToPoint(b.comm, b.label, 64 * 1024, 7); + for ([_]usize{ 4096, 65536, 1 << 20 }) |n| try c.graphs(b.comm, b.label, n, 4); + try b.comm.barrier(); + } + for (backends[0..n_backends]) |b| { + try bench.sweep(&c, b.comm, b.label, 4096, 4 << 20); + try bench.decodeWindow(&c, b.comm, b.label, 86, 8192); + } + // the stop: rank 0 marks it expected first, then both meet and close + if (s.fate) |*f| f.expectStop(); + try s.comm().barrier(); + std.debug.print("SUMMARY rank {d}: {d} passed, {d} failed\n", .{ cfg.rank, c.passes, c.failures }); + s.close(); + return if (c.failures == 0) 0 else 1; +} + +/// The failed-transport check's host-only forms (fail-fast off, a hold): TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire. +fn hostOnlyAllowed(cfg: tp.Config) bool { + const flag = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_TP_ALLOW_HOST_ONLY") orelse ""), "1"); + return flag and (cfg.backend != .roce or cfg.roce_wire == .shm); +} + +pub fn realNow() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.REALTIME, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) / 1e9; +} + +/// Exchanges until the victim fails; the survivor is left waiting inside one, which only the fate channel ends. +fn failfast(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, mode: []const u8) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, 1 << 20); + defer c.deinit(); + const comm = s.comm(); + var k: u32 = 0; + while (true) : (k += 1) { + if (k == 50 and cfg.rank == victim) { + std.debug.print("INJECT {s} rank {d} at {d:.6}\n", .{ mode, cfg.rank, realNow() }); + if (std.mem.eql(u8, mode, "kill")) _ = std.c.raise(.KILL); + s.fatal("window", error.Injected); + return 3; // not reached: the process exits 70 + } + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + c.stream.synchronize() catch |e| { + std.debug.print("rank {d}: stream error {t} after the peer failed\n", .{ cfg.rank, e }); + }; + comm.check() catch |e| { + std.debug.print("rank {d}: transport reports {t} at exchange {d}; waiting for the fate channel\n", .{ cfg.rank, e, k }); + while (true) tp.sock.sleepNs(std.time.ns_per_s); + }; + } +} diff --git a/zig/src/fabric/tp/gpu_transport.zig b/zig/src/fabric/tp/gpu_transport.zig new file mode 100644 index 0000000000..49cd625844 --- /dev/null +++ b/zig/src/fabric/tp/gpu_transport.zig @@ -0,0 +1,71 @@ +//! `tf-tp-test failfast V transport [N [HOLD_S]]`: the failed-transport path on real HCAs, injected on an idle link. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const checks = @import("gpu_checks.zig"); +const realNow = @import("gpu_main.zig").realNow; + +/// N clean sums, an NCCL meet, the victim's `Roce.inject` once every link is idle, then its close (fail-fast: exit 70). +pub fn run(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, n: u32, hold_s: u32) !u8 { + const r = s.roce orelse { + std.debug.print("failfast transport needs TF_COMM_BACKEND=roce (the session runs {t})\n", .{s.comm().kind()}); + return 2; + }; + var c = try checks.Ctx.init(d, gpa, cfg.rank, 4096); + defer c.deinit(); + const comm = s.comm(); + for (0..n) |_| { + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + try c.stream.synchronize(); + try comm.check(); + } + try s.nccl.iface().barrier(); + std.debug.print("rank {d}: {d} clean exchanges; both ranks met over NCCL\n", .{ cfg.rank, n }); + // the inject only on an idle link: every rank's posts completed, else no rank injects + const t_idle = tp.sock.nowNs(); + while (!r.linkIdle() and tp.sock.nowNs() - t_idle < 2 * std.time.ns_per_s) tp.sock.sleepNs(std.time.ns_per_ms); + var line: [384]u8 = undefined; + std.debug.print("rank {d}: link {s} before the inject: {s}\n", .{ cfg.rank, if (r.linkIdle()) "idle" else "BUSY", r.describeLink(&line) }); + if (try idleRanks(s, &c, r.linkIdle()) != cfg.world) { + std.debug.print("REFUSED rank {d}: a link was not idle; no inject, a clean stop\n", .{cfg.rank}); + if (s.fate) |*f| f.expectStop(); + try s.nccl.iface().barrier(); + s.close(); + return 2; + } + if (cfg.rank == victim) { + std.debug.print("INJECT transport rank {d} at {d:.6}\n", .{ cfg.rank, realNow() }); + r.inject(); + comm.check() catch |e| std.debug.print("rank {d}: transport reports {t}\n", .{ cfg.rank, e }); + } else if (s.fate != null) { + while (true) tp.sock.sleepNs(std.time.ns_per_s); // the fate channel ends this rank (exit 70) + } + const kept0 = tp.roce.Roce.kept.load(.monotonic); + s.close(); // fail-fast on: exit 70 from here + const kept = tp.roce.Roce.kept.load(.monotonic) - kept0; + const want: u32 = @intFromBool(cfg.rank == victim); + const live = cfg.rank != victim or mapped(r.region); + std.debug.print("rank {d}: closed; regions kept {d} (want {d}){s}\n", .{ cfg.rank, kept, want, if (cfg.rank != victim) "" else if (live) ", region still mapped" else ", region NOT mapped" }); + if (hold_s > 0) { + std.debug.print("HOLD rank {d} pid {d} for {d} s\n", .{ cfg.rank, std.c.getpid(), hold_s }); + tp.sock.sleepNs(@as(u64, hold_s) * std.time.ns_per_s); + } + const ok = kept == want and live; + std.debug.print("{s} rank {d}\n", .{ if (ok) "PASS" else "FAIL", cfg.rank }); + return if (ok) 0 else 1; +} + +fn mapped(region: []const u8) bool { + var vec: [1]u8 = undefined; + return std.os.linux.errno(std.os.linux.mincore(@constCast(region.ptr), 4096, &vec)) == .SUCCESS; +} + +/// The ranks whose link was idle, summed over NCCL (one f32 a rank). +fn idleRanks(s: *tp.Session, c: *checks.Ctx, idle: bool) !u32 { + var v: [1]f32 = .{@floatFromInt(@intFromBool(idle))}; + try c.in.uploadAsync(0, std.mem.asBytes(&v), c.stream.handle); + try s.nccl.iface().allReduce(c.in.ptr, c.out.ptr, 1, .f32, .sum, c.stream.handle); + try c.stream.synchronize(); + try c.out.download(0, std.mem.asBytes(&v)); + return @intFromFloat(v[0]); +} diff --git a/zig/src/fabric/tp/host.zig b/zig/src/fabric/tp/host.zig new file mode 100644 index 0000000000..051daea537 --- /dev/null +++ b/zig/src/fabric/tp/host.zig @@ -0,0 +1,293 @@ +//! In-process ranks over host memory: the `Collective` contract without a GPU (desyncs refused, abort wakes waiters). +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub const max_ranks = 8; + +const Call = enum(u8) { none, all_reduce, all_gather, all_gather_v, exchange, barrier }; + +const Post = struct { call: Call = .none, send: DevicePtr = 0, recv: DevicePtr = 0, count: usize = 0, t: DType = .u8, op: Op = .sum, peer: u32 = 0, lens: DevicePtr = 0 }; + +/// One-way message slot from rank `src` to rank `dst` (send / recv). +const Slot = struct { buf: std.atomic.Value(DevicePtr) = .init(0), bytes: usize = 0, taken: std.atomic.Value(bool) = .init(false) }; + +pub const Group = struct { + world: u32, + /// A wait longer than this is `error.Timeout` (a peer that never calls). + timeout_ns: u64 = 10 * std.time.ns_per_s, + posts: [max_ranks]Post = @splat(.{}), + arrived: std.atomic.Value(u32) = .init(0), + generation: std.atomic.Value(u32) = .init(0), + aborted: std.atomic.Value(bool) = .init(false), + /// Ranks (a bit each) whose transport refuses device lengths: their `all_gather_v` moves whole strides (tests). + var_refuse: u32 = 0, + slots: [max_ranks][max_ranks]Slot = @splat(@splat(.{})), + + pub fn init(world: u32) Group { + std.debug.assert(world >= 1 and world <= max_ranks); + return .{ .world = world }; + } + + pub fn rank(g: *Group, r: u32) Rank { + std.debug.assert(r < g.world); + return .{ .group = g, .me = r }; + } + + /// Sense-reversing barrier over every rank; `error.Aborted` once any rank aborts. + fn meet(g: *Group) Error!void { + const gen = g.generation.load(.acquire); + if (g.arrived.fetchAdd(1, .acq_rel) + 1 == g.world) { + g.arrived.store(0, .release); + _ = g.generation.fetchAdd(1, .acq_rel); + return if (g.aborted.load(.acquire)) error.Aborted else {}; + } + const t0 = now(); + while (g.generation.load(.acquire) == gen) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + if (g.aborted.load(.acquire)) return error.Aborted; + } +}; + +fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +fn bytesAt(p: DevicePtr, n: usize) []u8 { + if (n == 0) return &.{}; + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +pub const Rank = struct { + group: *Group, + me: u32, + /// `varAgreed`'s answer, once this communicator agreed. + var_agreed: ?bool = null, + + pub fn iface(self: *Rank) collective.Collective { + return .{ .ptr = self, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .host, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + + fn self_(ptr: *anyopaque) *Rank { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).group.world; + } + + /// Posts this rank's call, waits for all, and refuses calls that do not match rank 0's. + fn enter(self: *Rank, p: Post) Error!void { + const g = self.group; + if (g.aborted.load(.acquire)) return error.Aborted; + g.posts[self.me] = p; + try g.meet(); + const r0 = g.posts[0]; + if (r0.call != p.call or r0.count != p.count or r0.t != p.t or r0.op != p.op) { + abort(self); + return error.Invalid; + } + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_reduce, .send = s, .recv = r, .count = count, .t = t, .op = op }); + const n = count * t.size(); + var ins: [max_ranks][]const u8 = undefined; + for (0..g.world) |i| ins[i] = bytesAt(g.posts[i].send, n); + // the output may alias this rank's input: reduce into scratch first, then meet before writing + var scratch: [1 << 16]u8 = undefined; + if (n > scratch.len) return error.Unsupported; + if (n > 0) reference.reduce(scratch[0..n], ins[0..g.world], t, op); + try g.meet(); + @memcpy(bytesAt(r, n), scratch[0..n]); + try g.meet(); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather, .send = s, .recv = r, .count = count, .t = t }); + const n = count * t.size(); + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| @memcpy(out[i * n ..][0..n], bytesAt(g.posts[i].send, n)); + try g.meet(); + } + + fn lensAt(p: DevicePtr, world: u32) []const i32 { + return @as([*]const i32, @ptrFromInt(p))[0..world]; + } + + /// Copies only each rank's lens[r] bytes (`var_refuse` ranks take whole strides); unequal lengths are refused. + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather_v, .send = s, .recv = r, .count = stride, .t = t, .lens = lens }); + const n = stride * t.size(); + const mine = lensAt(lens, g.world); + for (0..g.world) |i| if (!std.mem.eql(i32, mine, lensAt(g.posts[i].lens, g.world))) { + abort(self); + return error.Invalid; + }; + const whole = g.var_refuse & (@as(u32, 1) << @intCast(self.me)) != 0; + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| { + if (mine[i] < 0 or mine[i] > n) { + abort(self); + return error.Invalid; + } + const len: usize = if (whole) n else @intCast(mine[i]); + @memcpy(out[i * n ..][0..len], bytesAt(g.posts[i].send, len)); + } + try g.meet(); + } + + fn varFits(ptr: *anyopaque, bytes: usize) bool { + const self = self_(ptr); + return bytes > 0 and self.group.var_refuse & (@as(u32, 1) << @intCast(self.me)) == 0; + } + + /// One int all-gather of `varFits(16)` the first time, then the kept answer. + fn varAgreed(ptr: *anyopaque) Error!bool { + const self = self_(ptr); + if (self.var_agreed) |v| return v; + const mine: i32 = @intFromBool(varFits(ptr, 16)); + var all: [max_ranks]i32 = undefined; + try allGather(ptr, @intFromPtr(&mine), @intFromPtr(&all), 1, .i32, null); + self.var_agreed = std.mem.min(i32, all[0..self.group.world]) == 1; + return self.var_agreed.?; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + try self.enter(.{ .call = .exchange, .send = s, .recv = r, .count = count, .t = t, .peer = peer }); + if (g.posts[peer].peer != self.me) { + abort(self); + return error.Invalid; + } + const n = count * t.size(); + @memcpy(bytesAt(r, n), bytesAt(g.posts[peer].send, n)); + try g.meet(); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[self.me][peer]; + slot.bytes = count * t.size(); + slot.taken.store(false, .release); + slot.buf.store(if (buf == 0) 1 else buf, .release); + try waitFor(g, &slot.taken, true); + slot.buf.store(0, .release); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[peer][self.me]; + const t0 = now(); + var src = slot.buf.load(.acquire); + while (src == 0) : (src = slot.buf.load(.acquire)) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + const n = count * t.size(); + if (slot.bytes != n) { + abort(self); + return error.Invalid; + } + @memcpy(bytesAt(buf, n), bytesAt(src, n)); + slot.taken.store(true, .release); + } + + fn waitFor(g: *Group, flag: *std.atomic.Value(bool), want: bool) Error!void { + const t0 = now(); + while (flag.load(.acquire) != want) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + } + + fn barrier(ptr: *anyopaque) Error!void { + const self = self_(ptr); + try self.enter(.{ .call = .barrier }); + try self.group.meet(); + } + + fn check(ptr: *anyopaque) Error!void { + if (self_(ptr).group.aborted.load(.acquire)) return error.Aborted; + } + + fn abort(ptr: *anyopaque) void { + self_(ptr).group.aborted.store(true, .release); + } +}; + +/// Runs `body(rank_collective, rank, ctx)` on `world` threads at once and returns the first error any rank hit. +pub fn run(world: u32, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + return runWith(world, .{}, ctx, body); +} + +pub const Options = struct { + /// `Group.var_refuse`: ranks whose transport takes no device lengths. + var_refuse: u32 = 0, +}; + +pub fn runWith(world: u32, opts: Options, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + var group = Group.init(world); + group.var_refuse = opts.var_refuse; + var ranks: [max_ranks]Rank = undefined; + var errs: [max_ranks]?anyerror = @splat(null); + var threads: [max_ranks]std.Thread = undefined; + const Wrap = struct { + fn go(r: *Rank, c: @TypeOf(ctx), e: *?anyerror) void { + body(r.iface(), r.me, c) catch |err| { + e.* = err; + r.group.aborted.store(true, .release); + }; + } + }; + for (0..world) |i| { + ranks[i] = group.rank(@intCast(i)); + threads[i] = try std.Thread.spawn(.{}, Wrap.go, .{ &ranks[i], ctx, &errs[i] }); + } + for (threads[0..world]) |t| t.join(); + for (errs[0..world]) |e| if (e) |err| return err; +} diff --git a/zig/src/fabric/tp/host_test.zig b/zig/src/fabric/tp/host_test.zig new file mode 100644 index 0000000000..7de32008af --- /dev/null +++ b/zig/src/fabric/tp/host_test.zig @@ -0,0 +1,154 @@ +//! The `Collective` contract on the in-process ranks: every call's bits against the reference, desyncs and aborts. +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const host = @import("host.zig"); +const Collective = collective.Collective; +const DType = collective.DType; + +fn addr(b: []u8) collective.DevicePtr { + return @intFromPtr(b.ptr); +} + +const n = 96; + +const Calls = struct { + fn go(c: Collective, rank: u32, _: void) !void { + const world = c.world(); + inline for (.{ DType.f32, DType.bf16, DType.i32, DType.f16 }) |t| { + const bytes = n * t.size(); + var mine: [n * 8]u8 = undefined; + var out: [n * 8 * host.max_ranks]u8 = undefined; + reference.fill(mine[0..bytes], t, rank, 7); + // all-reduce against the reference sum of every rank's input, in rank order + var ins: [host.max_ranks][n * 8]u8 = undefined; + var views: [host.max_ranks][]const u8 = undefined; + for (0..world) |r| { + reference.fill(ins[r][0..bytes], t, @intCast(r), 7); + views[r] = ins[r][0..bytes]; + } + var want: [n * 8]u8 = undefined; + reference.reduce(want[0..bytes], views[0..world], t, .sum); + try c.allReduce(addr(&mine), addr(&out), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], out[0..bytes]); + // in place + try c.allReduce(addr(&mine), addr(&mine), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], mine[0..bytes]); + // all-gather: rank order + reference.fill(mine[0..bytes], t, rank, 7); + try c.allGather(addr(&mine), addr(&out), n, t, null); + for (0..world) |r| try std.testing.expectEqualSlices(u8, views[r], out[r * bytes ..][0..bytes]); + } + try c.barrier(); + } +}; + +test "every call on two and three in-process ranks equals the reference" { + try host.run(2, {}, Calls.go); + try host.run(3, {}, Calls.go); +} + +const Pairwise = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var mine: [64]u8 = @splat(@intCast(rank + 1)); + var got: [64]u8 = undefined; + const peer = 1 - rank; + try c.exchange(addr(&mine), addr(&got), 64, .u8, peer, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(@intCast(peer + 1))), &got); + // send / recv, ordered so neither waits forever + if (rank == 0) { + try c.send(addr(&mine), 16, .f32, 1, null); + } else { + var buf: [64]u8 = undefined; + try c.recv(addr(&buf), 16, .f32, 0, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(1)), &buf); + } + try std.testing.expectError(error.Invalid, c.exchange(addr(&mine), addr(&got), 64, .u8, rank, null)); + } +}; + +test "exchange, send and recv between two ranks" { + try host.run(2, {}, Pairwise.go); +} + +const Desync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + // the ranks disagree on the count: both see it instead of hanging + const r = c.allGather(addr(&a), addr(&b), if (rank == 0) 16 else 32, .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "a desynchronised call is refused on every rank" { + try host.run(2, {}, Desync.go); +} + +const Aborts = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [16]u8 = @splat(0); + var b: [32]u8 = undefined; + if (rank == 1) { + // rank 1 never calls: it aborts after a while, which must free rank 0's waiting call + @import("sock.zig").sleepNs(50 * std.time.ns_per_ms); + c.abort(); + return; + } + try std.testing.expectError(error.Aborted, c.allGather(addr(&a), addr(&b), 16, .u8, null)); + try std.testing.expectError(error.Aborted, c.check()); + try std.testing.expectError(error.Aborted, c.barrier()); + } +}; + +test "abort frees a rank waiting in a collective" { + try host.run(2, {}, Aborts.go); +} + +const VarGather = struct { + fn go(c: Collective, rank: u32, refuse: u32) !void { + const world = c.world(); + const stride = 200; + const whole = refuse & (@as(u32, 1) << @intCast(rank)) != 0; + try std.testing.expectEqual(!whole, c.varFits(stride)); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); // kept: no second collective + var ins: [host.max_ranks][stride]u8 = undefined; + for (0..world) |r| reference.fill(&ins[r], .u8, @intCast(r), 11); + // lengths of 0, odd, 16 and the whole stride, in bytes; .f32 strides count elements + const cases = [_][3]i32{ .{ 0, 17, 200 }, .{ 200, 0, 3 }, .{ 16, 16, 16 }, .{ 199, 200, 1 } }; + inline for (.{ DType.u8, DType.f32 }) |t| for (cases) |lens| { + var out: [stride * host.max_ranks]u8 = @splat(0xEE); + try c.allGatherV(addr(&ins[rank]), addr(&out), stride / t.size(), @intFromPtr(&lens), t, null); + for (0..world) |r| { + const len: usize = @intCast(lens[r]); + try std.testing.expectEqualSlices(u8, ins[r][0..len], out[r * stride ..][0..len]); + // past a length: the transport's choice (untouched here, or the whole stride where lengths are refused) + const past = out[r * stride + len .. (r + 1) * stride]; + if (whole) try std.testing.expectEqualSlices(u8, ins[r][len..], past) else for (past) |x| try std.testing.expectEqual(@as(u8, 0xEE), x); + } + }; + try c.barrier(); + } +}; + +test "a variable-length all-gather moves each rank's length, and the ranks agree on device lengths once" { + for ([_]u32{ 2, 3 }) |w| { + try host.runWith(w, .{}, @as(u32, 0), VarGather.go); + try host.runWith(w, .{ .var_refuse = 2 }, @as(u32, 2), VarGather.go); + } +} + +const VarDesync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + const lens = [2]i32{ 8, if (rank == 0) 8 else 16 }; + const r = c.allGatherV(addr(&a), addr(&b), 64, @intFromPtr(&lens), .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "lengths that differ between ranks are refused on every rank" { + try host.run(2, {}, VarDesync.go); +} diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig new file mode 100644 index 0000000000..e5f6d72c31 --- /dev/null +++ b/zig/src/fabric/tp/hybrid.zig @@ -0,0 +1,183 @@ +//! Small exchanges through the mailbox (one host) or RoCE (two hosts), the rest through NCCL; ranks always agree. +const std = @import("std"); +const collective = @import("collective.zig"); +const Mailbox = @import("mailbox.zig").Mailbox; +const Roce = @import("roce.zig").Roce; +const Nccl = @import("nccl.zig").Nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +/// The one-kernel path for small shards: both expose the same calls. +pub const Small = union(enum) { + mailbox: *Mailbox, + roce: *Roce, + + fn fits(s: Small, n: usize) bool { + return switch (s) { + inline else => |p| p.fits(n), + }; + } + + fn sum(s: Small, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.sum(src, dst, count, t, stream), + }; + } + + fn gather(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.allGatherBytes(src, dst, n, stream), + }; + } + + fn exchange(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.exchangeBytes(src, dst, n, stream), + }; + } + + fn check(s: Small) Error!void { + return switch (s) { + inline else => |p| p.check(), + }; + } + + fn abort(s: Small) void { + switch (s) { + inline else => |p| p.abort(), + } + } +}; + +/// Whether a variable-length gather of `n`-byte strides one op cannot take goes RoCE in slot-sized ops (needs the ABI). +pub fn chunkRoute(chunked_bytes: usize, n: usize, single: bool, abi: bool) bool { + return abi and !single and n > 0 and n <= chunked_bytes; +} + +pub const Hybrid = struct { + small: Small, + nccl: *Nccl, + /// Largest shard a rank sends through the small path. + max_bytes: usize, + /// Largest stride a variable-length gather takes on RoCE in slot-sized ops (TF_TP_ROCE_CHUNKED_KB; 0: none). + chunked_bytes: usize = 0, + /// Every rank's `varFits(16)`, agreed at `Session.open` (the bootstrap's all-gather): `varAgreed`'s answer. + var_agreed: bool = false, + + pub fn iface(h: *Hybrid) collective.Collective { + return .{ .ptr = h, .vtable = if (h.small == .roce) &roce_vtable else &mailbox_vtable }; + } + + fn vtable(comptime kind: collective.Kind) collective.Collective.VTable { + return .{ + .kind = kind, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + } + + const mailbox_vtable = vtable(.hybrid); + const roce_vtable = vtable(.roce); + + fn self_(ptr: *anyopaque) *Hybrid { + return @ptrCast(@alignCast(ptr)); + } + + fn routed(h: *const Hybrid, n: usize) bool { + return n <= h.max_bytes and h.small.fits(n); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.size; + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (op == .sum and (t == .f32 or t == .bf16) and h.routed(count * t.size())) return h.small.sum(s, r, count, t, stream); + return Nccl.allReduce(h.nccl, s, r, count, t, op, stream); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (h.routed(count * t.size())) return h.small.gather(s, r, count * t.size(), stream); + return Nccl.allGather(h.nccl, s, r, count, t, stream); + } + + /// RoCE moves each rank's device length; the mailbox and NCCL move whole strides (the same bytes where lens says). + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + const n = stride * t.size(); + if (singleFits(h, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + if (varFits(ptr, n)) return h.small.roce.allGatherVChunked(s, r, n, lens, stream); + return allGather(ptr, s, r, stride, t, stream); + } + + fn singleFits(h: *const Hybrid, n: usize) bool { + return switch (h.small) { + .roce => |p| n > 0 and h.routed(n) and p.d != null, + .mailbox => false, + }; + } + + /// RoCE with its device side, in one op or in slot-sized ops (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + const single = singleFits(h, n); + return single or switch (h.small) { + .roce => |p| chunkRoute(h.chunked_bytes, n, single, p.d != null and p.chunk_abi), + .mailbox => false, + }; + } + + fn varAgreed(ptr: *anyopaque) Error!bool { + return self_(ptr).var_agreed; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const h = self_(ptr); + try collective.checkPeer(h.nccl.me, h.nccl.size, peer); + if (h.routed(count * t.size())) return h.small.exchange(s, r, count * t.size(), stream); + return Nccl.exchange(h.nccl, s, r, count, t, peer, stream); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).send(buf, count, t, peer, stream); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).recv(buf, count, t, peer, stream); + } + + fn barrier(ptr: *anyopaque) Error!void { + return Nccl.iface(self_(ptr).nccl).barrier(); + } + + fn check(ptr: *anyopaque) Error!void { + const h = self_(ptr); + try h.small.check(); + try Nccl.iface(h.nccl).check(); + } + + fn abort(ptr: *anyopaque) void { + const h = self_(ptr); + h.small.abort(); + Nccl.iface(h.nccl).abort(); + } +}; diff --git a/zig/src/fabric/tp/log.zig b/zig/src/fabric/tp/log.zig new file mode 100644 index 0000000000..4e3234ab49 --- /dev/null +++ b/zig/src/fabric/tp/log.zig @@ -0,0 +1,15 @@ +//! TP's log lines, silent under test: a test's stderr is the build runner's channel. +const std = @import("std"); +const quiet = @import("builtin").is_test; + +pub fn err(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.err(format, args); +} + +pub fn warn(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.warn(format, args); +} + +pub fn info(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.info(format, args); +} diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu new file mode 100644 index 0000000000..843b4f1af2 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.cu @@ -0,0 +1,564 @@ +// One-shot two-rank exchange through pinned host memory both GPUs map (the TP mailbox; mailbox.zig is its host side). + +#include +#include + +// Protocol adapted from b12x's RoCEnante (Apache-2.0): a device epoch, slots by seq parity, a seq flag after the data. +namespace tftp { + +constexpr int kMaxBlocks = 16; +constexpr int kFlagStride = 32; // words: one 128-byte line per flag +// ctrl words: [0] abort (any rank, any host thread); per rank r at 16 + 8 r: failed, seq, seen, block, why, waited us +constexpr int kCtrlAbort = 0; +constexpr int kCtrlRank = 16; + +enum Mode : uint32_t { kGather = 0, kSumF32 = 1, kSumBf16 = 2, kExchange = 3 }; + +struct MailArgs { + const uint8_t *in; // this rank's shard (device) + uint8_t *out; // gather [2 n], exchange / sums [n] + uint8_t *peer_slots; // the peer's inbox data, as this GPU maps it + uint32_t *peer_flags; // the peer's inbox flags + const uint8_t *my_slots; // this rank's inbox data + const uint32_t *my_flags; // this rank's inbox flags + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison + uint32_t *ctrl; // shared control words + uint64_t nbytes; // bytes a rank + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; +}; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 ld_relaxed_sys_v4(const void *p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u8(const uint8_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u8 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void st_relaxed_sys(uint32_t *p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_sys(uint32_t *p, uint32_t v) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_gpu(uint32_t *p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_gpu(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_sys(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.sys.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_cas_gpu(uint32_t *p, uint32_t cmp, uint32_t v) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.cas.b32 %0, [%1], %2, %3;" : "=r"(old) : "l"(p), "r"(cmp), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ void fence_sc_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } + +__device__ __forceinline__ uint64_t globaltimer() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// dst[i] <- src[i] for i in [0, n), grid-stride, 16-byte vectors when aligned; SysSrc: system-scope loads. +template +__device__ __forceinline__ void copy_bytes(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 v = SysSrc ? ld_relaxed_sys_v4(src + (i << 4)) : reinterpret_cast(src)[i]; + reinterpret_cast(dst)[i] = v; + } + } else { + for (uint64_t i = index; i < n; i += stride) { + dst[i] = static_cast(SysSrc ? ld_relaxed_sys_u8(src + i) : src[i]); + } + } +} + +// TF_TP_ROCE_FAST: copy_bytes with two grid strides' loads issued before their stores (same bytes, one round trip). +template +__device__ __forceinline__ void copy_bytes2(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) != 0u) { + copy_bytes(dst, src, n, index, stride); + return; + } + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += 2u * stride) { + const uint64_t j = i + stride; + uint4 v0, v1 = make_uint4(0u, 0u, 0u, 0u); + if (SysSrc) { + v0 = ld_relaxed_sys_v4(src + (i << 4)); + if (j < packs) v1 = ld_relaxed_sys_v4(src + (j << 4)); + } else { + v0 = reinterpret_cast(src)[i]; + if (j < packs) v1 = reinterpret_cast(src)[j]; + } + reinterpret_cast(dst)[i] = v0; + if (j < packs) reinterpret_cast(dst)[j] = v1; + } +} + +template +__device__ __forceinline__ void copy_any(bool fast, uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + if (fast) { + copy_bytes2(dst, src, n, index, stride); + } else { + copy_bytes(dst, src, n, index, stride); + } +} + +__device__ __forceinline__ uint32_t add_bf16x2(uint32_t x, uint32_t y) { + const float2 a = __bfloat1622float2(*reinterpret_cast(&x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&y)); + const __nv_bfloat162 r = __floats2bfloat162_rn(a.x + b.x, a.y + b.y); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t add_f32(uint32_t x, uint32_t y) { + return __float_as_uint(__uint_as_float(x) + __uint_as_float(y)); +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u16(const uint8_t *p) { + uint16_t v; + asm volatile("ld.relaxed.sys.global.u16 %0, [%1];" : "=h"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint16_t add_bf16(uint16_t x, uint16_t y) { + const __nv_bfloat16 r = __float2bfloat16_rn(__bfloat162float(__ushort_as_bfloat16(x)) + + __bfloat162float(__ushort_as_bfloat16(y))); + return __bfloat16_as_ushort(r); +} + +// out <- mine + peer's, rank 0's first; IEEE addition commutes, so both ranks get the same bits. +template +__device__ __forceinline__ void sum_bytes(uint8_t *out, const uint8_t *mine, const uint8_t *peer, uint64_t n, + uint32_t rank, uint64_t index, uint64_t stride) { + const uint64_t al = reinterpret_cast(out) | reinterpret_cast(mine) | + reinterpret_cast(peer) | n; + if ((al & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 m = reinterpret_cast(mine)[i]; + const uint4 p = ld_relaxed_sys_v4(peer + (i << 4)); + const uint4 a = rank == 0 ? m : p; + const uint4 b = rank == 0 ? p : m; + uint4 r; + r.x = Bf16 ? add_bf16x2(a.x, b.x) : add_f32(a.x, b.x); + r.y = Bf16 ? add_bf16x2(a.y, b.y) : add_f32(a.y, b.y); + r.z = Bf16 ? add_bf16x2(a.z, b.z) : add_f32(a.z, b.z); + r.w = Bf16 ? add_bf16x2(a.w, b.w) : add_f32(a.w, b.w); + reinterpret_cast(out)[i] = r; + } + } else if (Bf16) { + for (uint64_t i = index; i < (n >> 1); i += stride) { + const uint16_t m = reinterpret_cast(mine)[i]; + const uint16_t p = static_cast(ld_relaxed_sys_u16(peer + (i << 1))); + reinterpret_cast(out)[i] = rank == 0 ? add_bf16(m, p) : add_bf16(p, m); + } + } else { + for (uint64_t i = index; i < (n >> 2); i += stride) { + const uint32_t m = reinterpret_cast(mine)[i]; + const uint32_t p = ld_relaxed_sys(reinterpret_cast(peer) + i); + reinterpret_cast(out)[i] = rank == 0 ? add_f32(m, p) : add_f32(p, m); + } + } +} + +// Thread 0 of a block: spin until the peer block's flag holds seq; false on abort (why 1) or timeout (why 2). +__device__ __noinline__ bool wait_flag(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, uint64_t timeout_ns, + uint32_t *seen, uint32_t *why, uint32_t *waited_us) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *seen = v; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + return *seen == seq; + } + } + } +} + +// TF_TP_ROCE_FAST: wait_flag for one of several pollers of a flag; `done` (shared) ends the others once one saw it. +__device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, + uint64_t timeout_ns, uint32_t *seen, uint32_t *why, uint32_t *waited_us, + volatile uint32_t *done) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *done = 1u; + *seen = v; + return true; + } + if (*done != 0u) { + *seen = seq; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + if (*seen == seq) { + *done = 1u; + return true; + } + return *done != 0u; + } + } + } +} + +// One failure record a rank (the first failing wait's): its fields, then its valid word, each fenced at system scope. +__device__ __forceinline__ void fail_record(uint32_t *poison, uint32_t *rec, uint32_t seq, uint32_t seen, uint32_t who, + uint32_t why, uint32_t waited) { + if (atom_cas_gpu(poison, 0u, 1u) == 0u) { + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, who); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } +} + +} // namespace tftp + +using namespace tftp; + +extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + // a poisoned transport does nothing (the host raises on its next check); poison and epoch in one round trip + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const uint64_t flag_at = (slot * kMaxBlocks + blockIdx.x) * kFlagStride; + + // 1. this block's share of the shard into the peer's inbox, then its flag after it (system-scope release) + copy_bytes(a.peer_slots + slot * a.slot_bytes, a.in, a.nbytes, index, stride); + __syncthreads(); + if (tid == 0) { + fence_sc_sys(); + st_release_sys(a.peer_flags + flag_at, seq); + } + // the local shard while the peer's writes travel + if (a.mode == kGather) { + copy_bytes(a.out + static_cast(a.rank) * a.nbytes, a.in, a.nbytes, index, stride); + } + + // 2. the peer's block of the same index (it wrote exactly the bytes this block reads) + if (tid == 0) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.my_flags + flag_at, seq, a.ctrl, a.timeout_ns, &seen, &why, &waited)) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kCtrlRank + 8u * a.rank, seq, seen, blockIdx.x, why, waited); + } + } + __syncthreads(); + + // 3. copy out, gated by this block's own wait + if (s_bad == 0u) { + const uint8_t *inbox = a.my_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_bytes(a.out + static_cast(1u - a.rank) * a.nbytes, inbox, a.nbytes, index, + stride); + break; + case kExchange: + copy_bytes(a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the last block publishes the epoch, unless a wait of this launch failed (later launches stay no-ops) + __syncthreads(); + if (tid == 0) { + // the counter is back at 0 for the next launch (stream order), whatever its grid + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + } + } +} + +// The RoCE variant (roce.zig): blocks stage into the send slot, the last rings the doorbell, a host proxy RDMA-writes. +namespace tftp { + +constexpr int kRoceHcas = 2; +// ctrl words shared with roce.zig's proxy: [0] doorbell seq, [4 + slot] padded bytes, [8] completed, [16..21] failure +constexpr int kRoceSeq = 0; +constexpr int kRoceSlotBytes = 4; +constexpr int kRoceCompleted = 8; +constexpr int kRoceFailed = 16; +constexpr int kRoceAbort = 12; // wait_flag reads ctrl[kCtrlAbort]: the RoCE abort word is passed as that ctrl + +struct RoceArgs { + const uint8_t *in; + uint8_t *out; + uint8_t *send_slots; // this rank's send slots (pinned host, read by the NIC) + const uint8_t *recv_slots; // this rank's receive slots (written by the peer's NIC) + const uint32_t *flags; // this rank's flags [2][kRoceHcas], written by the peer's NIC + uint32_t *ctrl; + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison, [3] stage counter + uint64_t nbytes; + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; + uint32_t n_hca; + uint32_t opts; // TF_TP_ROCE_FAST: bit 0 fast, bits 8-15 pollers a flag, bits 16-31 start stagger ns; 0: classic + // gather only (nullptr: every shard is nbytes): lens[r] = rank r's bytes, padded to 16 on the wire + const int32_t *lens; + uint64_t chunk; // chunked gathers: this op's offset into each shard (at most a slot of each from there; else 0) +}; + +// Rank r's bytes this op: ``nbytes``, or lens[r] capped at nbytes. +__device__ __forceinline__ uint64_t roce_shard(const RoceArgs &a, uint32_t r) { + if (a.lens == nullptr) { + return a.nbytes; + } + const int32_t v = __ldg(a.lens + r); + const uint64_t n = v > 0 ? static_cast(v) : 0u; + return n < a.nbytes ? n : a.nbytes; +} + +// The bytes a copy of an n-byte shard moves: n rounded up to 16 when that stays inside a 16-aligned stride, else n. +__device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t n) { + if (a.lens == nullptr || (a.nbytes & 15u) != 0u) { + return n; + } + return (n + 15u) & ~uint64_t(15u); +} + +// A shard's bytes past `off`, at most a slot (chunked gathers; off 0 and a shard within the slot: the shard). +__device__ __forceinline__ uint64_t roce_part(const RoceArgs &a, uint64_t shard, uint64_t off) { + const uint64_t rest = shard > off ? shard - off : 0u; + return rest < a.slot_bytes ? rest : a.slot_bytes; +} + +} // namespace tftp + +// Present in images whose RoceArgs carry ``chunk`` (roce.zig refuses TF_TP_ROCE_CHUNKED_KB without it); never launched. +extern "C" __global__ void tp_roce_chunk_abi() {} + +extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + __shared__ uint32_t s_done[kRoceHcas]; // fast: some poller of HCA h saw its flag + const bool fast = (a.opts & 1u) != 0u; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + if (tid < static_cast(kRoceHcas)) { + s_done[tid] = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const bool var = a.mode == kGather && a.lens != nullptr; + const uint64_t off = var ? a.chunk : 0u; + const uint64_t mine = var ? roce_part(a, roce_shard(a, a.rank), off) : a.nbytes; + const uint64_t theirs = var ? roce_part(a, roce_shard(a, 1u - a.rank), off) : a.nbytes; + const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + if (off != 0u && mine == 0u && theirs == 0u) { + return; // nothing past this offset on either rank: no op, the epoch stays (both ranks skip alike) + } + if (padded > a.slot_bytes) { // never staged past a slot: a refused wait (why 3), the transport poisoned + if (tid == 0 && blockIdx.x == 0) { + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, 0u, 0u, 3u, 0u); + } + return; + } + + // 1. stage this block's share into the send slot; the last block to finish rings the doorbell + copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in + off, var ? roce_span(a, mine) : a.nbytes, index, + stride); + __syncthreads(); + if (tid == 0 && fast) { + // lean doorbell: release RMW per block arrival, the last acquires, then a release store (no SC system fences) + const bool last = gdim == 1u || atom_add_acq_rel_sys(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } else if (tid == 0) { + fence_sc_sys(); + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + fence_sc_sys(); // every block's staging (seen through the counter) before the doorbell + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } + if (a.mode == kGather) { + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes + off, a.in + off, + var ? roce_span(a, mine) : a.nbytes, index, stride); + } + + // 2. every HCA's flag for this slot; fast: poller k of HCA h is thread 32 k + h, started k x stagger ns late + const uint32_t pollers = (a.opts >> 8) & 0xffu; + const uint32_t stagger_ns = a.opts >> 16; + if (fast) { + const uint32_t lane = tid & 31u; + const uint32_t k = tid >> 5; + if (lane < a.n_hca && k < (pollers == 0u ? 1u : pollers)) { + if (k > 0u && stagger_ns != 0u) { + __nanosleep(k * stagger_ns); + } + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag_shared(a.flags + (slot * kRoceHcas + lane) * kFlagStride, seq, a.ctrl + kRoceAbort, + a.timeout_ns, &seen, &why, &waited, &s_done[lane])) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, lane, why, waited); + } + } + } else if (tid == 0) { + for (uint32_t h = 0; h < a.n_hca; h++) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, + &why, &waited)) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, h, why, waited); + break; + } + } + } + __syncthreads(); + + // 3. copy out of the receive slot the peer's NIC wrote + if (s_bad == 0u) { + const uint8_t *inbox = a.recv_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes + off, inbox, + var ? roce_span(a, theirs) : a.nbytes, index, stride); + break; + case kExchange: + copy_any(fast, a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the epoch, as tp_mailbox; the host also sees the last completed seq + __syncthreads(); + if (tid == 0) { + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + st_relaxed_sys(a.ctrl + kRoceCompleted, seq); + } + } +} diff --git a/zig/src/fabric/tp/mailbox.zig b/zig/src/fabric/tp/mailbox.zig new file mode 100644 index 0000000000..902bb79393 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.zig @@ -0,0 +1,215 @@ +//! The mailbox backend: two-rank exchanges in one kernel through pinned host memory both GPUs map (host side). +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const flag_bytes = 2 * max_blocks * 128; +const ctrl_bytes = 4096; +const ctrl_abort = 0; +const ctrl_rank = 16; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `MailArgs`, field for field. +const MailArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + peer_slots: DevicePtr, + peer_flags: DevicePtr, + my_slots: DevicePtr, + my_flags: DevicePtr, + state: DevicePtr, + ctrl: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, +}; + +/// Driver calls the runtime's table does not carry (page-locking memory this process mapped itself). +pub const HostRegister = struct { + lib: std.DynLib, + register: *const fn (?*anyopaque, usize, c_uint) callconv(.c) cuda.abi.Result, + unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, + + pub fn open() !HostRegister { + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + return .{ + .register = lib.lookup(@FieldType(HostRegister, "register"), "cuMemHostRegister_v2") orelse return error.MissingSymbol, + .unregister = lib.lookup(@FieldType(HostRegister, "unregister"), "cuMemHostUnregister") orelse return error.MissingSymbol, + .lib = lib, + }; + } +}; + +/// Grid for `n` bytes a rank: one block per 16 KiB, a power of two, at most `max_blocks` (both ranks pick the same). +pub fn blocksFor(n: usize) u32 { + const want = std.math.divCeil(usize, @max(n, 1), 16 * 1024) catch 1; + return @intCast(@min(std.math.ceilPowerOfTwoAssert(usize, want), max_blocks)); +} + +pub const Mailbox = struct { + d: *const cuda.Driver, + me: u32, + slot_bytes: usize, + timeout_ns: u64, + region: []align(std.heap.page_size_min) u8, + dev: DevicePtr, + hr: HostRegister, + state: cuda.DeviceBuffer, + module: cuda.Module, + func: cuda.Function, + aborted: std.atomic.Value(bool) = .init(false), + + /// Both ranks call it after the bootstrap (two ranks, one host); `image` is mailbox.cu's fatbin. + pub fn init(d: *const cuda.Driver, b: *Bootstrap, image: []const u8, max_bytes: usize, timeout_ns: u64) !Mailbox { + if (b.world != 2) return error.Unsupported; + if (image.len == 0) return error.NoKernel; + const slot = std.mem.alignForward(usize, max_bytes, 4096); + const len = ctrl_bytes + 2 * (flag_bytes + 2 * slot); + // rank 0 creates and sizes the file before anyone learns its name, the others open it, then rank 0 unlinks it + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + const path = name[0..name_len :0]; + if (b.rank != 0) { + fd = std.c.open(path, .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + try b.agree("mapped", "mailbox file"); + const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + errdefer posix.munmap(region); + try b.agree("open", "mailbox mapping"); + if (b.rank == 0) _ = std.c.unlink(path); + var hr = try HostRegister.open(); + errdefer hr.lib.close(); + try d.check(hr.register(region.ptr, len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + errdefer _ = hr.unregister(region.ptr); + var dev: DevicePtr = 0; + try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); + var state = try cuda.DeviceBuffer.alloc(d, 16); + errdefer state.free(); + try state.fill8(0, null); + var module = try cuda.Module.load(d, image); + errdefer module.unload(); + const func = try module.function("tp_mailbox"); + try b.agree("ready", "mailbox setup"); + return .{ .d = d, .me = b.rank, .slot_bytes = slot, .timeout_ns = timeout_ns, .region = region, .dev = dev, .hr = hr, .state = state, .module = module, .func = func }; + } + + pub fn deinit(m: *Mailbox) void { + m.module.unload(); + m.state.free(); + _ = m.hr.unregister(m.region.ptr); + m.hr.lib.close(); + posix.munmap(m.region); + } + + fn inbox(m: *const Mailbox, r: u32) usize { + return ctrl_bytes + r * (flag_bytes + 2 * m.slot_bytes); + } + + fn ctrlWord(m: *const Mailbox, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(m.region[4 * i ..].ptr)); + } + + /// Whether `n` bytes a rank fit one slot. + pub fn fits(m: *const Mailbox, n: usize) bool { + return n <= m.slot_bytes; + } + + fn launch(m: *Mailbox, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, stream: collective.Stream) Error!void { + if (m.aborted.load(.acquire)) return error.Aborted; + if (n > m.slot_bytes) return error.Unsupported; + if (n == 0) return; + const peer = 1 - m.me; + var a: MailArgs = .{ + .in = in, + .out = out, + .peer_slots = m.dev + m.inbox(peer) + flag_bytes, + .peer_flags = m.dev + m.inbox(peer), + .my_slots = m.dev + m.inbox(m.me) + flag_bytes, + .my_flags = m.dev + m.inbox(m.me), + .state = m.state.ptr, + .ctrl = m.dev, + .nbytes = n, + .slot_bytes = m.slot_bytes, + .timeout_ns = m.timeout_ns, + .rank = m.me, + .mode = @backingInt(mode), + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const r = m.d.api.cuLaunchKernel(m.func.handle, blocksFor(n), 1, 1, threads, 1, 1, 0, stream, ¶ms, null); + m.d.check(r, "cuLaunchKernel tp_mailbox") catch return error.BackendFailed; + } + + pub fn allGatherBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.gather, s, r, n, stream); + } + + pub fn exchangeBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.exchange, s, r, n, stream); + } + + /// Sums only, f32 or bf16 (what the model's partials are); the rest is `error.Unsupported`. + pub fn sum(m: *Mailbox, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return m.launch(mode, s, r, count * t.size(), stream); + } + + /// Ends pending and later exchanges on both ranks (the abort word is shared); this rank's launches stop. + pub fn abort(m: *Mailbox) void { + m.aborted.store(true, .release); + m.ctrlWord(ctrl_abort).store(1, .release); + } + + /// A wait that timed out or was aborted on either rank. + pub fn check(m: *Mailbox) Error!void { + for (0..2) |r| { + const rec = ctrl_rank + 8 * r; + if (m.ctrlWord(rec).load(.acquire) != 0) { + log.warn("tp mailbox: rank {d}'s wait for seq {d} failed ({s} after {d} us; block {d}, flag held {d})", .{ r, m.ctrlWord(rec + 1).load(.acquire), if (m.ctrlWord(rec + 4).load(.acquire) == 1) "aborted" else "timed out", m.ctrlWord(rec + 5).load(.acquire), m.ctrlWord(rec + 3).load(.acquire), m.ctrlWord(rec + 2).load(.acquire) }); + return error.PeerFailed; + } + } + if (m.aborted.load(.acquire) or m.ctrlWord(ctrl_abort).load(.acquire) != 0) return error.Aborted; + } +}; + +test "grids grow with the shard and stay powers of two" { + try std.testing.expectEqual(@as(u32, 1), blocksFor(4096)); + try std.testing.expectEqual(@as(u32, 1), blocksFor(16 * 1024)); + try std.testing.expectEqual(@as(u32, 2), blocksFor(16 * 1024 + 1)); + try std.testing.expectEqual(@as(u32, 4), blocksFor(64 * 1024)); + try std.testing.expectEqual(@as(u32, 16), blocksFor(4 << 20)); + try std.testing.expectEqual(@as(usize, 96), @sizeOf(MailArgs)); +} diff --git a/zig/src/fabric/tp/nccl.zig b/zig/src/fabric/tp/nccl.zig new file mode 100644 index 0000000000..b641839aad --- /dev/null +++ b/zig/src/fabric/tp/nccl.zig @@ -0,0 +1,163 @@ +//! The NCCL backend: one communicator a rank, every collective one enqueue on the caller's stream. +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const nccl = cuda.nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub fn dataType(t: DType) nccl.DataType { + return switch (t) { + .u8 => .u8, + .i32 => .i32, + .i64 => .i64, + .f16 => .f16, + .bf16 => .bf16, + .f32 => .f32, + .f64 => .f64, + }; +} + +pub fn redOp(op: Op) nccl.RedOp { + return switch (op) { + .sum => .sum, + .max => .max, + .min => .min, + }; +} + +pub const Nccl = struct { + lib: *const nccl.Library, + d: *const cuda.Driver, + comm: nccl.Comm, + me: u32, + size: u32, + aborted: std.atomic.Value(bool) = .init(false), + /// The barrier's one-float buffers and its own stream (host path only). + scratch: cuda.DeviceBuffer, + stream: cuda.Stream, + + /// Joins the communicator `id` names (every rank calls it with the same id; it returns once all have). + pub fn init(lib: *const nccl.Library, d: *const cuda.Driver, id: nccl.UniqueId, rank: u32, world: u32) !Nccl { + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + errdefer _ = lib.api.ncclCommDestroy(comm); + var scratch = try cuda.DeviceBuffer.alloc(d, 4 * (world + 1)); + errdefer scratch.free(); + try scratch.fill8(0, null); + return .{ .lib = lib, .d = d, .comm = comm, .me = rank, .size = world, .scratch = scratch, .stream = try cuda.Stream.init(d, true) }; + } + + pub fn deinit(n: *Nccl) void { + _ = n.stream.synchronize() catch {}; + n.stream.deinit(); + n.scratch.free(); + if (!n.aborted.load(.acquire)) _ = n.lib.api.ncclCommDestroy(n.comm); + } + + pub fn iface(n: *Nccl) collective.Collective { + return .{ .ptr = n, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .nccl, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + }; + + fn self_(ptr: *anyopaque) *Nccl { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).size; + } + + fn call(n: *Nccl, r: nccl.Result, what: []const u8) Error!void { + if (r == 0) return; + if (n.aborted.load(.acquire)) return error.Aborted; + n.lib.check(r, what) catch return error.BackendFailed; + } + + fn live(n: *Nccl) Error!void { + if (n.aborted.load(.acquire)) return error.Aborted; + } + + pub fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllReduce(s, r, count, dataType(t), redOp(op), n.comm, stream), "ncclAllReduce"); + } + + pub fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(s, r, count, dataType(t), n.comm, stream), "ncclAllGather"); + } + + pub fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclGroupStart(), "ncclGroupStart"); + const a = n.lib.api.ncclSend(s, count, dataType(t), @intCast(peer), n.comm, stream); + const b = n.lib.api.ncclRecv(r, count, dataType(t), @intCast(peer), n.comm, stream); + const e = n.lib.api.ncclGroupEnd(); + try n.call(a, "ncclSend"); + try n.call(b, "ncclRecv"); + try n.call(e, "ncclGroupEnd"); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclSend(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclSend"); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclRecv(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclRecv"); + } + + /// One float all-gathered on the barrier's own stream, then a sync of it (host path, never captured). + fn barrier(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(n.scratch.ptr, n.scratch.ptr + 4, 1, .f32, n.comm, n.stream.handle), "ncclAllGather"); + n.stream.synchronize() catch return if (n.aborted.load(.acquire)) error.Aborted else error.BackendFailed; + } + + fn check(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + var r: nccl.Result = 0; + try n.call(n.lib.api.ncclCommGetAsyncError(n.comm, &r), "ncclCommGetAsyncError"); + if (r != 0 and r != nccl.in_progress) { + n.lib.check(r, "NCCL async error") catch {}; + return error.PeerFailed; + } + } + + fn abort(ptr: *anyopaque) void { + const n = self_(ptr); + if (n.aborted.swap(true, .acq_rel)) return; + _ = n.lib.api.ncclCommAbort(n.comm); + } +}; diff --git a/zig/src/fabric/tp/planlink.zig b/zig/src/fabric/tp/planlink.zig new file mode 100644 index 0000000000..44dbb61a0c --- /dev/null +++ b/zig/src/fabric/tp/planlink.zig @@ -0,0 +1,467 @@ +//! The round plan from rank 0 to the other ranks over its own TCP connection: a 4-byte length, then int64 words. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const builtin = @import("builtin"); +const config = @import("config.zig"); +const linux = std.os.linux; + +pub const max_ranks = 64; + +/// The CPU a plan thread pinned itself to (TF_TP_PLAN_PIN), -1 until one did; an engine's CPU isolation keeps off it. +pub var pinned_cpu: std.atomic.Value(i32) = .init(-1); +const more: u32 = 0x8000; // MSG_MORE: the length and the ints leave as one segment + +pub const Error = error{ PeerClosed, TooLarge, NotRankZero, NotFollower }; + +/// A batch frame's first word (a plan's op is >= 0): [batch, n, len 1, plan 1, ..., len n, plan n]. +pub const batch: i64 = std.math.minInt(i64); +/// Words rank 0 holds at most (more flushes the held plans first), and the most a batch frame carries. +pub const hold_words: usize = 1024; +pub const batch_words: usize = 2 * hold_words + 2; + +/// Largest plan either side accepts: 64M words (512 MiB), so a garbled length cannot become a 4 GiB allocation. +pub const max_frame_words: usize = 64 << 20; + +comptime { + std.debug.assert(builtin.cpu.arch.endian() == .little); // the ints go out as they are in memory +} + +pub const PlanLink = struct { + rank: u32, + world: u32, + /// Rank 0: rank i's socket at index i; the others: rank 0's at index 0 (the bootstrap's `releasePlan`). + fds: [max_ranks]sock.Fd, + sent: u64 = 0, + received: u64 = 0, + /// TF_TP_PLAN_LINK=rdma (config.plan_spin_ns): the follower polls the socket this long before it blocks. + spin_ns: u64 = 0, + /// TF_TP_PLAN_PIN: the sending / receiving thread pins itself at its first plan; `avoid`: the RoCE proxy's CPU. + pin: ?config.PlanPin = null, + avoid: ?u32 = null, + pinned: bool = false, + /// rank 0: plans held for the next send (`hold`), each as [len, plan...]; `held_n` of them + held: [hold_words]i64 = undefined, + held_len: usize = 0, + held_n: usize = 0, + /// the follower: a batch frame's plans not handed out yet, as [len, plan...] from `inbox_at` + inbox: [batch_words]i64 = undefined, + inbox_len: usize = 0, + inbox_at: usize = 0, + batches: u64 = 0, + + pub fn init(rank: u32, world: u32, fds: [max_ranks]sock.Fd) PlanLink { + return .{ .rank = rank, .world = world, .fds = fds }; + } + + /// Rank 0: one plan to every other rank (host only); held plans go first, in the same frame when they fit. + pub fn send(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (ints.len > max_frame_words) return error.TooLarge; + if (p.held_n > 0) { + if (2 + p.held_len + 1 + ints.len <= batch_words) { + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n + 1); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + frame[2 + p.held_len] = @intCast(ints.len); + @memcpy(frame[3 + p.held_len ..][0..ints.len], ints); + const n = 3 + p.held_len + ints.len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + try p.flush(); + } + return p.sendFrame(ints); + } + + /// Rank 0: a plan the followers run before the next sent one, carried in that plan's frame (`send` or `flush`). + pub fn hold(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (1 + ints.len > hold_words) { + try p.flush(); + return p.sendFrame(ints); + } + if (p.held_len + 1 + ints.len > hold_words) try p.flush(); + p.held[p.held_len] = @intCast(ints.len); + @memcpy(p.held[p.held_len + 1 ..][0..ints.len], ints); + p.held_len += 1 + ints.len; + p.held_n += 1; + } + + /// Rank 0: the held plans now, as one batch frame (nothing held: nothing sent). + pub fn flush(p: *PlanLink) Error!void { + if (p.held_n == 0) return; + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + const n = 2 + p.held_len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + + fn sendFrame(p: *PlanLink, ints: []const i64) Error!void { + p.pinOnce(); + const bytes = std.mem.sliceAsBytes(ints); + const len: u32 = @intCast(bytes.len); + for (p.fds[1..p.world]) |fd| { + if (std.c.send(fd, &len, 4, std.posix.MSG.NOSIGNAL | more) != 4) return error.PeerClosed; + sock.sendAll(fd, bytes) catch return error.PeerClosed; + } + p.sent += 1; + } + + /// The other ranks: the next plan into `out`; `error.PeerClosed` once rank 0 is gone, `error.TooLarge` skips one. + pub fn recv(p: *PlanLink, out: []i64) Error![]i64 { + if (p.inboxNext()) |m| { + if (m.len > out.len) return error.TooLarge; + @memcpy(out[0..m.len], m); + return out[0..m.len]; + } + const n = try p.header(); + if (n > out.len) { + try p.drop(n); + return error.TooLarge; + } + const got = try p.body(out[0..n]); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recv(out); + } + return got; + } + + /// The next plan into `buf`, grown to the frame; the buffer is kept, so steady plans allocate nothing. + pub fn recvGrow(p: *PlanLink, gpa: std.mem.Allocator, buf: *std.ArrayList(i64)) (Error || std.mem.Allocator.Error)![]i64 { + if (p.inboxNext()) |m| { + try buf.resize(gpa, m.len); + @memcpy(buf.items, m); + return buf.items; + } + const n = try p.header(); + if (n > max_frame_words) { + try p.drop(n); + return error.TooLarge; + } + try buf.resize(gpa, n); + const got = try p.body(buf.items); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recvGrow(gpa, buf); + } + return got; + } + + /// The follower: a batch frame's plans into the inbox (checked whole first: a garbled batch ends the link). + fn unpack(p: *PlanLink, frame: []const i64) Error!void { + if (frame.len < 2 or frame.len > batch_words or frame[1] < 1) return error.PeerClosed; + var at: usize = 2; + var k: i64 = 0; + while (k < frame[1]) : (k += 1) { + if (at >= frame.len or frame[at] < 0) return error.PeerClosed; + at += 1 + @as(usize, @intCast(frame[at])); + if (at > frame.len) return error.PeerClosed; + } + if (at != frame.len) return error.PeerClosed; + @memcpy(p.inbox[0 .. frame.len - 2], frame[2..]); + p.inbox_len = frame.len - 2; + p.inbox_at = 0; + p.batches += 1; + } + + /// The follower: the next plan of the last batch frame, if any (a slice of the inbox, valid until the next recv). + fn inboxNext(p: *PlanLink) ?[]const i64 { + if (p.inbox_at >= p.inbox_len) return null; + const n: usize = @intCast(p.inbox[p.inbox_at]); + const m = p.inbox[p.inbox_at + 1 ..][0..n]; + p.inbox_at += 1 + n; + return m; + } + + /// The next frame's length in words. + fn header(p: *PlanLink) Error!usize { + if (p.rank == 0) return error.NotFollower; + p.pinOnce(); + if (p.spin_ns > 0) p.spin(); + var len: u32 = undefined; + sock.recvExact(p.fds[0], std.mem.asBytes(&len), null) catch return error.PeerClosed; + if (len % 8 != 0) return error.PeerClosed; // not a frame boundary: the stream is lost + return len / 8; + } + + fn body(p: *PlanLink, out: []i64) Error![]i64 { + sock.recvExact(p.fds[0], std.mem.sliceAsBytes(out), null) catch return error.PeerClosed; + p.received += 1; + return out; + } + + /// Reads and drops `n` words. + fn drop(p: *PlanLink, n: usize) Error!void { + var left: usize = n * 8; + var scratch: [4096]u8 = undefined; + while (left > 0) : (left -= @min(left, scratch.len)) sock.recvExact(p.fds[0], scratch[0..@min(left, scratch.len)], null) catch return error.PeerClosed; + } + + /// Polls rank 0's socket until it is readable (or hung up) or `spin_ns` passed; the read that follows blocks. + fn spin(p: *PlanLink) void { + if (comptime builtin.os.tag != .linux) return; + var fd = [1]linux.pollfd{.{ .fd = p.fds[0], .events = linux.POLL.IN, .revents = 0 }}; + const end = sock.nowNs() + p.spin_ns; + while (true) { + const rc = linux.poll(&fd, 1, 0); + if (linux.errno(rc) == .SUCCESS and rc > 0) return; + if (sock.nowNs() >= end) return; + std.atomic.spinLoopHint(); + } + } + + /// Pins the calling thread once: an explicit CPU, or auto's highest cpu_capacity. + fn pinOnce(p: *PlanLink) void { + if (p.pinned) return; + p.pinned = true; + if (comptime builtin.os.tag != .linux) return; + const want = p.pin orelse return; + const cpu = switch (want) { + .cpu => |c| c, + .auto => bestCpu(p.avoid) orelse { + log.info("[tensorfold] tp: TF_TP_PLAN_PIN=auto: no cpu_capacity differs among the allowed CPUs; the plan thread stays unpinned", .{}); + return; + }, + }; + pinTo(cpu) catch |e| { + log.warn("[tensorfold] tp: TF_TP_PLAN_PIN: CPU {d} refused ({t}); the plan thread stays unpinned", .{ cpu, e }); + return; + }; + pinned_cpu.store(@intCast(cpu), .release); + log.info("[tensorfold] tp: plan link rank {d} pinned to CPU {d}", .{ p.rank, cpu }); + } + + pub fn close(p: *PlanLink) void { + for (&p.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } +}; + +/// The calling thread onto `cpu` alone. +pub fn pinTo(cpu: u32) !void { + if (cpu >= linux.CPU_SETSIZE) return error.BadCpu; + var set: linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + try linux.sched_setaffinity(0, &set); +} + +/// Auto pin: the allowed CPU (not `avoid`) with the highest cpu_capacity, lowest on ties; null when all are equal. +pub fn bestCpu(avoid: ?u32) ?u32 { + var set: linux.cpu_set_t = undefined; + if (linux.errno(linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set)) != .SUCCESS) return null; + var best: ?u32 = null; + var best_cap: u32 = 0; + var low_cap: u32 = std.math.maxInt(u32); + for (set, 0..) |word, wi| { + var bits = word; + while (bits != 0) : (bits &= bits - 1) { + const cpu: u32 = @intCast(wi * @bitSizeOf(usize) + @ctz(bits)); + if (avoid == cpu) continue; + const c = capacity(cpu) orelse continue; + low_cap = @min(low_cap, c); + if (best == null or c > best_cap) { + best = cpu; + best_cap = c; + } + } + } + return if (best != null and best_cap > low_cap) best else null; +} + +fn capacity(cpu: u32) ?u32 { + var pb: [64]u8 = undefined; + const path = std.fmt.bufPrintSentinel(&pb, "/sys/devices/system/cpu/cpu{d}/cpu_capacity", .{cpu}, 0) catch return null; + const rc = linux.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return null; + const fd: i32 = @intCast(rc); + defer _ = linux.close(fd); + var b: [16]u8 = undefined; + const n = linux.read(fd, &b, b.len); + if (linux.errno(n) != .SUCCESS) return null; + return std.fmt.parseInt(u32, std.mem.trim(u8, b[0..n], " \n"), 10) catch null; +} + +test "plans arrive whole and in order; a closed rank 0 ends the follower" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @as(i64, @intCast(i)) * -3 + 0x7FFF_0000_0000; + try lead.send(&.{ 0x504C414E, 2, 1 }); + try lead.send(&big); + try lead.send(&.{}); + var buf: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 2, 1 }, try follow.recv(&buf)); + try std.testing.expectEqualSlices(i64, &big, try follow.recv(&buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&buf)).len); + try std.testing.expectError(error.NotRankZero, follow.send(&.{1})); + var small: [2]i64 = undefined; + try lead.send(&.{ 1, 2, 3 }); + try std.testing.expectError(error.TooLarge, follow.recv(&small)); + lead.close(); + try std.testing.expectError(error.PeerClosed, follow.recv(&buf)); +} + +test "an admission-sized plan (1.2M words) arrives whole through the growable path" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + const gpa = std.testing.allocator; + const prompt = try gpa.alloc(i64, 1_200_000); + defer gpa.free(prompt); + for (prompt, 0..) |*v, i| v.* = @intCast((i *% 2654435761) % 163840); + // the socket buffer holds far less than 9.6 MB: rank 0 sends from its own thread, as it would + const Send = struct { + fn go(link: *PlanLink, ints: []const i64) void { + link.send(&.{ 0x504C414E, 3, 20 }) catch {}; + link.send(ints) catch {}; + link.send(&.{ 0x504C414E, 3, 21 }) catch {}; + } + }; + const t = try std.Thread.spawn(.{}, Send.go, .{ &lead, prompt }); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(gpa); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 20 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, prompt, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 21 }, try follow.recvGrow(gpa, &buf)); + t.join(); + try std.testing.expect(buf.capacity >= prompt.len); // kept for the next admission +} + +test "TF_TP_PLAN_LINK=rdma: a spinning follower gets every plan, late or early; TF_TP_PLAN_PIN pins the thread" { + if (builtin.os.tag != .linux) return error.SkipZigTest; + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + follow.spin_ns = 2 * std.time.ns_per_ms; + // the CPU this thread may run on first: the pin keeps it allowed, so the test runner is not starved + var set: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set); + var first: u32 = 0; + for (set, 0..) |w, i| if (w != 0) { + first = @intCast(i * @bitSizeOf(usize) + @ctz(w)); + break; + }; + const Follower = struct { + fn run(f: *PlanLink, cpu: u32, out: *[3]i64, ok: *bool) void { + f.pin = .{ .cpu = cpu }; + var buf: [8]i64 = undefined; + const x = (f.recv(&buf) catch return)[0]; // arrives within the spin + const y = (f.recv(&buf) catch return)[0]; // arrives after it: the blocking read + var now: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &now); + var count: usize = 0; + for (now) |w| count += @popCount(w); + out.* = .{ x, y, @intCast(count) }; + ok.* = f.pinned and (now[cpu / @bitSizeOf(usize)] >> @intCast(cpu % @bitSizeOf(usize))) & 1 == 1; + } + }; + var out: [3]i64 = .{ 0, 0, 0 }; + var ok = false; + const t = try std.Thread.spawn(.{}, Follower.run, .{ &follow, first, &out, &ok }); + try lead.send(&.{11}); + var ts: linux.timespec = .{ .sec = 0, .nsec = 20 * std.time.ns_per_ms }; + _ = linux.nanosleep(&ts, null); + try lead.send(&.{22}); + t.join(); + try std.testing.expectEqual([3]i64{ 11, 22, 1 }, out); + try std.testing.expect(ok); + if (bestCpu(null)) |c| try std.testing.expect((set[c / @bitSizeOf(usize)] >> @intCast(c % @bitSizeOf(usize))) & 1 == 1); +} + +test "held plans leave with the next plan, in order, as one frame; a large plan after them flushes them first" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("127.0.0.1", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + const lead = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(lead); + lead.* = PlanLink.init(0, 2, f0); + const follow = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(follow); + follow.* = PlanLink.init(1, 2, f1); + defer follow.close(); + // a round: three keeps held, the window sent: one frame + try lead.hold(&.{ 22, 0, 3 }); + try lead.hold(&.{ 22, 1, 0 }); + try lead.hold(&.{ 22, 2, 5 }); + try lead.send(&.{ 21, 4, 9, 9, 9 }); + try std.testing.expectEqual(@as(u64, 1), lead.sent); + // a held plan, then one too large for a batch: two frames, the held one first + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @intCast(i); + try lead.hold(&.{ 22, 3, 1 }); + try lead.send(&big); + try std.testing.expectEqual(@as(u64, 3), lead.sent); + // flush alone, an empty plan held, a plain send after + try lead.hold(&.{}); + try lead.flush(); + try lead.flush(); // nothing held: nothing sent + try lead.send(&.{7}); + try std.testing.expectEqual(@as(u64, 5), lead.sent); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(std.testing.allocator); + const gpa = std.testing.allocator; + try std.testing.expectEqualSlices(i64, &.{ 22, 0, 3 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 1, 0 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 2, 5 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 21, 4, 9, 9, 9 }, try follow.recvGrow(gpa, &buf)); + var out: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 22, 3, 1 }, try follow.recv(&out)); + try std.testing.expectEqualSlices(i64, &big, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&out)).len); + try std.testing.expectEqualSlices(i64, &.{7}, try follow.recv(&out)); + try std.testing.expectEqual(@as(u64, 3), follow.batches); + try std.testing.expectError(error.NotRankZero, follow.hold(&.{1})); + lead.close(); +} diff --git a/zig/src/fabric/tp/reference.zig b/zig/src/fabric/tp/reference.zig new file mode 100644 index 0000000000..b98966fcc8 --- /dev/null +++ b/zig/src/fabric/tp/reference.zig @@ -0,0 +1,121 @@ +//! Host references for the collectives: reductions in rank order, bf16 / f16 added in fp32 and rounded once. +const std = @import("std"); +const collective = @import("collective.zig"); +const DType = collective.DType; +const Op = collective.Op; + +/// bf16 bits to fp32 (exact). +pub fn bf16ToF32(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// fp32 to bf16 bits, round to nearest even; NaN stays a quiet NaN. +pub fn f32ToBf16(f: f32) u16 { + const u: u32 = @bitCast(f); + if (std.math.isNan(f)) return @intCast((u >> 16) | 0x40); + const round = 0x7FFF + ((u >> 16) & 1); + return @intCast((u +% round) >> 16); +} + +fn combine(comptime T: type, a: T, b: T, op: Op) T { + return switch (op) { + .sum => if (@typeInfo(T) == .int) a +% b else a + b, + .max => @max(a, b), + .min => @min(a, b), + }; +} + +fn reduceTyped(comptime T: type, acc: []u8, src: []const u8, op: Op) void { + const n = acc.len / @sizeOf(T); + for (0..n) |i| { + const a = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const b = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), src[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const r: T = combine(T, @bitCast(a), @bitCast(b), op); + std.mem.writeInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], @bitCast(r), .little); + } +} + +fn reduceHalf(acc: []u8, src: []const u8, op: Op, bf: bool) void { + const n = acc.len / 2; + for (0..n) |i| { + const a = std.mem.readInt(u16, acc[2 * i ..][0..2], .little); + const b = std.mem.readInt(u16, src[2 * i ..][0..2], .little); + const fa: f32 = if (bf) bf16ToF32(a) else @floatCast(@as(f16, @bitCast(a))); + const fb: f32 = if (bf) bf16ToF32(b) else @floatCast(@as(f16, @bitCast(b))); + const r = combine(f32, fa, fb, op); + const out: u16 = if (bf) f32ToBf16(r) else @bitCast(@as(f16, @floatCast(r))); + std.mem.writeInt(u16, acc[2 * i ..][0..2], out, .little); + } +} + +/// acc <- acc (op) src, element-wise; `acc` holds the lower ranks' result, `src` the next rank's input. +pub fn accumulate(acc: []u8, src: []const u8, t: DType, op: Op) void { + std.debug.assert(acc.len == src.len and acc.len % t.size() == 0); + switch (t) { + .u8 => reduceTyped(u8, acc, src, op), + .i32 => reduceTyped(i32, acc, src, op), + .i64 => reduceTyped(i64, acc, src, op), + .f32 => reduceTyped(f32, acc, src, op), + .f64 => reduceTyped(f64, acc, src, op), + .bf16 => reduceHalf(acc, src, op, true), + .f16 => reduceHalf(acc, src, op, false), + } +} + +/// out <- inputs[0] (op) inputs[1] (op) ..., in rank order. +pub fn reduce(out: []u8, inputs: []const []const u8, t: DType, op: Op) void { + @memcpy(out, inputs[0]); + for (inputs[1..]) |in| accumulate(out, in, t, op); +} + +/// Deterministic test data for `rank`'s input of `bytes`: finite values of type `t` that differ per rank and position. +pub fn fill(out: []u8, t: DType, rank: u32, seed: u64) void { + var rng = std.Random.DefaultPrng.init(seed ^ (@as(u64, rank) *% 0x9E3779B97F4A7C15)); + const r = rng.random(); + const n = out.len / t.size(); + for (0..n) |i| { + const x = r.float(f32) * 8.0 - 4.0; + switch (t) { + .u8 => out[i] = r.int(u8), + .i32 => std.mem.writeInt(i32, out[4 * i ..][0..4], r.int(i32), .little), + .i64 => std.mem.writeInt(i64, out[8 * i ..][0..8], r.int(i64), .little), + .f32 => std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x), .little), + .f64 => std.mem.writeInt(u64, out[8 * i ..][0..8], @bitCast(@as(f64, x) * 1.000000123), .little), + .bf16 => std.mem.writeInt(u16, out[2 * i ..][0..2], f32ToBf16(x), .little), + .f16 => std.mem.writeInt(u16, out[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(x))), .little), + } + } +} + +test "bf16 rounding is nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.0)); + // 1 + 2^-8 is exactly halfway between bf16 1.0 and 1.0078125: rounds to the even one (1.0) + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.00390625)); + try std.testing.expectEqual(@as(u16, 0x3F82), f32ToBf16(1.01171875)); + try std.testing.expect(std.math.isNan(bf16ToF32(f32ToBf16(std.math.nan(f32))))); +} + +test "rank-order reduction of bf16 adds in fp32 and rounds once" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(u16, a[0..2], f32ToBf16(1.0), .little); + std.mem.writeInt(u16, b[0..2], f32ToBf16(0.00390625), .little); + std.mem.writeInt(u16, a[2..4], f32ToBf16(-2.5), .little); + std.mem.writeInt(u16, b[2..4], f32ToBf16(3.0), .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .bf16, .sum); + try std.testing.expectEqual(f32ToBf16(1.0), std.mem.readInt(u16, out[0..2], .little)); + try std.testing.expectEqual(f32ToBf16(0.5), std.mem.readInt(u16, out[2..4], .little)); + reduce(&out, &.{ &a, &b }, .bf16, .max); + try std.testing.expectEqual(f32ToBf16(3.0), std.mem.readInt(u16, out[2..4], .little)); +} + +test "integer sums wrap" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(i32, &a, std.math.maxInt(i32), .little); + std.mem.writeInt(i32, &b, 1, .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .i32, .sum); + try std.testing.expectEqual(@as(i32, std.math.minInt(i32)), std.mem.readInt(i32, &out, .little)); +} diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig new file mode 100644 index 0000000000..c9bb72a2c3 --- /dev/null +++ b/zig/src/fabric/tp/roce.zig @@ -0,0 +1,525 @@ +//! The RoCE backend between two hosts: the tp_roce kernel stages and rings a doorbell, a proxy thread moves the slot. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); +const HostRegister = @import("mailbox.zig").HostRegister; +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const page = std.heap.page_size_min; +const flag_stride = 128; +// region: ctrl [4 KiB], flags [2 slots][2 HCAs] at 128 B, send [2][slot], recv [2][slot] (mailbox.cu's tp_roce) +const ctrl_off = 0; +const flag_off = 4096; +const send_off = 8192; +// ctrl words (mailbox.cu): doorbell, padded bytes per slot, completed, abort, proxy failure, a failed wait's record +const w_seq = 0; +const w_slot_bytes = 4; +const w_completed = 8; +const w_abort = 12; +const w_proxy_failed = 13; +const w_failed = 16; +/// Model graphs leave sub-millisecond gaps: the proxy stays hot this many idle polls, then naps 20 us at a time. +const idle_spins: u64 = 20_000_000; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `RoceArgs`, field for field. +const RoceArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + send_slots: DevicePtr, + recv_slots: DevicePtr, + flags: DevicePtr, + ctrl: DevicePtr, + state: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, + n_hca: u32, + /// TF_TP_ROCE_FAST (config.roce_fast; 0: the classic path) + opts: u32 = 0, + /// gather only: device int32 [2] of each rank's bytes (0: every shard is `nbytes`) + lens: DevicePtr = 0, + /// Chunked gathers: this op's byte offset into each rank's shard (it moves at most a slot from there; 0 otherwise). + chunk: u64 = 0, +}; + +pub const WireKind = enum { + /// RDMA writes over the HCAs. + verbs, + /// The proxy copies into the peer's region in shared memory: the same protocol on one host without RDMA (tests). + shm, +}; + +pub const Options = struct { + wire: WireKind = .verbs, + hcas: []const rv.HcaSpec = &.{}, + traffic_class: u8 = 0, + cpu: ?u32 = null, + max_bytes: usize = 256 * 1024, + timeout_ns: u64 = 120 * std.time.ns_per_s, + /// TF_TP_ROCE_FAST: tp_roce's `opts` word (config.roce_fast) + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) + blocks: u32 = 16, + /// The shm wire's stripes, as HCAs: its posts take the verbs wire's split and bounds. + shm_stripes: u32 = 1, + /// Tests: sees every shm post's stripes before the copy. + tap: ?*shm.PostTap = null, +}; + +const Wire = union(WireKind) { + verbs: rv.VerbsWire, + shm: shm.ShmWire, + + fn post(w: *Wire, slot: u32, seq: u32, padded: u32) !void { + switch (w.*) { + .verbs => |*v| try v.post(slot, seq, padded), + .shm => |*s| try s.post(slot, seq, padded), + } + } + + fn reap(w: *Wire) !void { + switch (w.*) { + .verbs => |*v| try v.reap(), + .shm => {}, + } + } + + fn hcas(w: *const Wire) u32 { + return switch (w.*) { + .verbs => |v| v.n_hca, + .shm => |s| s.stripes, + }; + } +}; + +pub const Roce = struct { + me: u32, + slot_bytes: usize, + timeout_ns: u64, + /// TF_TP_ROCE_FAST: the kernel's `opts` word + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid + max_grid: u32 = 16, + /// This rank's region (for shm: half of the mapping in `mapping`). + region: []align(page) u8, + mapping: []align(page) u8, + wire: Wire, + cpu: ?u32, + // device side (absent in host-only use: tests and `tf-tp-test roce-host`) + d: ?*const cuda.Driver = null, + hr: ?HostRegister = null, + dev: DevicePtr = 0, + state: ?cuda.DeviceBuffer = null, + module: ?cuda.Module = null, + func: ?cuda.Function = null, + /// The kernel image takes `RoceArgs.chunk` (it exports tp_roce_chunk_abi): TF_TP_ROCE_CHUNKED_KB needs it. + chunk_abi: bool = false, + // the proxy thread + thread: ?std.Thread = null, + running: std.atomic.Value(bool) = .init(false), + last_seq: u32 = 0, + err_buf: [160]u8 = undefined, + err_len: usize = 0, + // host emulation of the kernel's side (tests) + host_epoch: u32 = 0, + aborted: std.atomic.Value(bool) = .init(false), + + pub fn regionBytes(slot: usize) usize { + return send_off + 4 * slot; + } + + /// Both ranks call it after the bootstrap; any rank failing makes all return `error.RoceUnavailable` (collective). + pub fn init(d: ?*const cuda.Driver, b: *Bootstrap, image: []const u8, o: Options) !*Roce { + if (b.world != 2) return error.Unsupported; + if (o.max_bytes > guard.max_write) return error.WriteTooLarge; // a slot no post could fill in one write + const slot = std.mem.alignForward(usize, o.max_bytes, 4096); + const len = regionBytes(slot); + const r = try std.heap.page_allocator.create(Roce); + errdefer std.heap.page_allocator.destroy(r); + r.* = .{ .me = b.rank, .slot_bytes = slot, .timeout_ns = o.timeout_ns, .fast = o.fast, .max_grid = o.blocks, .region = undefined, .mapping = undefined, .wire = undefined, .cpu = o.cpu }; + // 1. the regions and the wire, each rank alone; every rank learns whether all did + const ok = r.openWire(b, o, len) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, ok, "RoCE setup")) { + if (ok) r.closeWire(); + return error.RoceUnavailable; + } + errdefer r.closeWire(); + // 2. the verbs connection records, through the bootstrap + const connected = r.connectWire(b) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} connect failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, connected, "RoCE connect")) return error.RoceUnavailable; + // 3. the device side + if (d) |drv| { + const dev_ok = r.openDevice(drv, image) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} device setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, dev_ok, "RoCE device setup")) { + if (dev_ok) r.closeDevice(); + return error.RoceUnavailable; + } + } + r.running.store(true, .release); + r.thread = try std.Thread.spawn(.{}, proxyLoop, .{r}); + try b.agree("proxy", "RoCE proxy start"); + return r; + } + + fn agreeAll(b: *Bootstrap, ok: bool, _: []const u8) !bool { + var blobs: [2][]const u8 = undefined; + var buf: [2]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(ok)}, &buf, &blobs); + return all[0][0] == 1 and all[1][0] == 1; + } + + fn openWire(r: *Roce, b: *Bootstrap, o: Options, len: usize) !bool { + switch (o.wire) { + .verbs => { + r.mapping = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + r.region = r.mapping; + errdefer posix.munmap(r.mapping); + try rv.prepareRegion(r.mapping); + r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; + }, + .shm => { + r.mapping = try shm.sharedPair(b, len); + r.region = @alignCast(r.mapping[r.me * len ..][0..len]); + const layout: guard.Layout = .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes, .len = len }; + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .layout = layout, .local = r.region, .stripes = o.shm_stripes, .tap = o.tap } }; + }, + } + return true; + } + + fn connectWire(r: *Roce, b: *Bootstrap) !bool { + switch (r.wire) { + .shm => return true, + .verbs => |*v| { + const mine = v.record(); + var blobs: [2][]const u8 = undefined; + var buf: [2 * @sizeOf(rv.Record)]u8 = undefined; + const all = try b.allGather(std.mem.asBytes(&mine), &buf, &blobs); + var peer: rv.Record = undefined; + if (all[1 - r.me].len != @sizeOf(rv.Record)) return error.BadRecord; + @memcpy(std.mem.asBytes(&peer), all[1 - r.me]); + try v.connect(peer); + var line: [512]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} registered {s}", .{ r.me, v.describe(&line) }); + return true; + }, + } + } + + fn openDevice(r: *Roce, d: *const cuda.Driver, image: []const u8) !bool { + if (image.len == 0) return error.NoKernel; + r.d = d; + r.hr = try HostRegister.open(); + try d.check(r.hr.?.register(r.region.ptr, r.region.len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + try d.check(d.api.cuMemHostGetDevicePointer_v2(&r.dev, r.region.ptr, 0), "cuMemHostGetDevicePointer"); + r.state = try cuda.DeviceBuffer.alloc(d, 16); + try r.state.?.fill8(0, null); + r.module = try cuda.Module.load(d, image); + r.func = try r.module.?.function("tp_roce"); + r.chunk_abi = if (r.module.?.function("tp_roce_chunk_abi")) |_| true else |_| false; + return true; + } + + fn closeDevice(r: *Roce) void { + if (r.module) |*m| m.unload(); + if (r.state) |*s| s.free(); + if (r.hr) |*h| { + _ = h.unregister(r.region.ptr); + h.lib.close(); + } + r.module = null; + r.state = null; + r.hr = null; + } + + fn closeWire(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| v.close(), + .shm => {}, + } + posix.munmap(r.mapping); + } + + /// Regions a failed transport left registered and mapped (the process exit is their only teardown). + pub var kept: std.atomic.Value(u32) = .init(0); + + /// HCAs and GPU quiet, then unregister and unmap; a failed transport's region is kept until exit. + pub fn deinit(r: *Roce) void { + const was_failed = r.failed(); + r.word(w_abort).store(1, .release); + r.running.store(false, .release); + if (r.thread) |t| t.join(); + const quiet = r.quiesce(); + if (was_failed or !quiet) { + const why = if (was_failed) "failed" else "did not quiesce"; + log.warn("[tensorfold] tp roce: rank {d}: the transport {s}; region 0x{x} ({d} B) kept registered until exit", .{ r.me, why, @intFromPtr(r.region.ptr), r.region.len }); + _ = kept.fetchAdd(1, .monotonic); + return; + } + r.closeDevice(); + r.closeWire(); + std.heap.page_allocator.destroy(r); + } + + /// The QPs to error with every posted request reaped, and the GPU's work on the region done. + fn quiesce(r: *Roce) bool { + const wire = switch (r.wire) { + .verbs => |*v| v.quiesce(rv.quiesce_ns), + .shm => true, + }; + const gpu = if (r.d) |d| d.api.cuCtxSynchronize() == cuda.abi.success else true; + var line: [384]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} quiesced ({s}): {s}; GPU sync {s}", .{ r.me, if (wire) "drained" else "NOT drained", r.describeLink(&line), if (gpu) "done" else "FAILED" }); + return wire and gpu; + } + + /// Each HCA's QP state, posts and completions, and what a drain reaped (tf-tp-test's idle check; shm: none). + pub fn describeLink(r: *const Roce, buf: []u8) []const u8 { + return switch (r.wire) { + .verbs => |*v| v.report(buf), + .shm => "shm wire", + }; + } + + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); + } + + fn fail(r: *Roce, comptime fmt: []const u8, args: anytype) void { + const s = std.fmt.bufPrint(&r.err_buf, fmt, args) catch "proxy failure"; + r.err_len = s.len; + log.err("[tensorfold] tp roce proxy: {s}", .{s}); + r.word(w_proxy_failed).store(1, .release); + r.word(w_abort).store(1, .release); // waits on this rank end now instead of at the timeout + r.halt(); + } + + /// The HCAs stop: nothing of this rank's goes out and none of the peer's writes lands while the ranks go down. + fn halt(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| { + const ok = v.halt(); + var line: [384]u8 = undefined; + log.warn("[tensorfold] tp roce: rank {d} halted{s}: {s}", .{ r.me, if (ok) "" else " (a QP refused ERR)", v.report(&line) }); + }, + .shm => {}, + } + } + + /// Every post of this rank's completed (shm: nothing is ever in flight). + pub fn linkIdle(r: *const Roce) bool { + return switch (r.wire) { + .verbs => |*v| v.idle(), + .shm => true, + }; + } + + /// tf-tp-test's `failfast RANK transport`: the proxy's failure path (`fail`), here with no op in flight. + pub fn inject(r: *Roce) void { + r.fail("failure injected by tf-tp-test (no RDMA fault)", .{}); + } + + /// The transport failed or was aborted: a device may still write to its region. + pub fn failed(r: *const Roce) bool { + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return true; + return r.word(w_proxy_failed).load(.acquire) != 0 or r.word(w_failed).load(.acquire) != 0; + } + + /// The proxy loop: spin on the doorbell, post each new sequence (at most two pending), reap while idle, then nap. + fn proxyLoop(r: *Roce) void { + if (r.cpu) |cpu| pin(cpu); + var idle: u64 = 0; + while (r.running.load(.monotonic)) { + if (r.word(w_abort).load(.monotonic) != 0) return; + const seq = r.word(w_seq).load(.acquire); + if (seq == r.last_seq) { + idle += 1; + if (idle % 64 == 0) r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + if (idle >= idle_spins) sock.sleepNs(20 * std.time.ns_per_us) else std.atomic.spinLoopHint(); + continue; + } + idle = 0; + var pending = seq -% r.last_seq; + if (pending > 2) return r.fail("doorbell skipped {d} ops (posted {d}, now {d})", .{ pending, r.last_seq, seq }); + var s = r.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const padded = r.word(w_slot_bytes + (s & 1)).load(.acquire); + if (padded == 0 or padded % 16 != 0 or padded > r.slot_bytes) return r.fail("op {d} of {d} bytes", .{ s, padded }); + r.wire.post(s & 1, s, padded) catch |e| return r.fail("post of op {d}: {t}", .{ s, e }); + r.last_seq = s; + } + r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + } + } + + pub fn fits(r: *const Roce, n: usize) bool { + return n <= r.slot_bytes; + } + + fn launch(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + return r.launchAt(mode, in, out, n, lens, 0, stream); + } + + fn launchAt(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, chunk: u64, stream: collective.Stream) Error!void { + if (r.aborted.load(.acquire)) return error.Aborted; + if (n > r.slot_bytes and (lens == 0 or !r.chunk_abi)) return error.Unsupported; + if (n == 0) return; + const d = r.d orelse return error.Unsupported; + var a: RoceArgs = .{ + .in = in, + .out = out, + .send_slots = r.dev + send_off, + .recv_slots = r.dev + send_off + 2 * r.slot_bytes, + .flags = r.dev + flag_off, + .ctrl = r.dev + ctrl_off, + .state = r.state.?.ptr, + .nbytes = n, + .slot_bytes = r.slot_bytes, + .timeout_ns = r.timeout_ns, + .rank = r.me, + .mode = @backingInt(mode), + .n_hca = r.wire.hcas(), + .opts = r.fast, + .lens = lens, + .chunk = chunk, + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const blocks = @min(@import("mailbox.zig").blocksFor(@min(n, r.slot_bytes)), r.max_grid); + d.check(d.api.cuLaunchKernel(r.func.?.handle, blocks, 1, 1, threads, 1, 1, 0, stream, ¶ms, null), "cuLaunchKernel tp_roce") catch return error.BackendFailed; + } + + pub fn allGatherBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.gather, s, o, n, 0, stream); + } + + /// `allGatherBytes` where rank r posts only lens[r] bytes (device int32 [2], padded to 16, at least 16). + pub fn allGatherV(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (n > r.slot_bytes) return error.Unsupported; // one op: past a slot is allGatherVChunked's + if (lens == 0) return error.Invalid; + return r.launch(.gather, s, o, n, lens, stream); + } + + /// `allGatherV` past the slot: op c moves at most a slot of each shard from c x slot (empty on both ranks: no op). + pub fn allGatherVChunked(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (lens == 0) return error.Invalid; + if (!r.chunk_abi) return error.Unsupported; + var off: usize = 0; + while (off < n) : (off += r.slot_bytes) try r.launchAt(.gather, s, o, n, lens, off, stream); + } + + pub fn exchangeBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.exchange, s, o, n, 0, stream); + } + + pub fn sum(r: *Roce, s: DevicePtr, o: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return r.launch(mode, s, o, count * t.size(), 0, stream); + } + + /// Ends this rank's waits and the proxy (the peer learns through the fate channel and aborts its own). + pub fn abort(r: *Roce) void { + r.aborted.store(true, .release); + r.word(w_abort).store(1, .release); + r.halt(); + } + + pub fn check(r: *Roce) Error!void { + if (r.word(w_failed).load(.acquire) != 0) { + log.warn("tp roce: rank {d}'s wait for seq {d} failed ({s} after {d} us; HCA {d}, flag held {d}); doorbell {d}, posted {d}, completed {d}, abort {d}, proxy failed {d}, flags {d} {d}", .{ r.me, r.word(w_failed + 1).load(.acquire), if (r.word(w_failed + 4).load(.acquire) == 1) "aborted" else "timed out", r.word(w_failed + 5).load(.acquire), r.word(w_failed + 3).load(.acquire), r.word(w_failed + 2).load(.acquire), r.word(w_seq).load(.acquire), r.last_seq, r.word(w_completed).load(.acquire), r.word(w_abort).load(.acquire), r.word(w_proxy_failed).load(.acquire), r.flagWord(0), r.flagWord(1) }); + return error.PeerFailed; + } + if (r.word(w_proxy_failed).load(.acquire) != 0) return error.BackendFailed; + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return error.Aborted; + } + + fn flagWord(r: *const Roce, slot: usize) u32 { + const p: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + slot * 2 * flag_stride ..].ptr)); + return p.load(.acquire); + } + + pub fn proxyError(r: *const Roce) []const u8 { + return r.err_buf[0..r.err_len]; + } + + /// The kernel's side of one op done by the host (tests, `roce-host`): a gather through the same slots and flags. + pub fn hostGather(r: *Roce, in: []const u8, out: []u8) Error!void { + return r.hostGatherV(in, out, null); + } + + /// `hostGather` with variable lengths (`allGatherV`): rank k's first lens[k] bytes land at out + k x n. + pub fn hostGatherV(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32) Error!void { + if (in.len > r.slot_bytes) return error.Invalid; + return r.hostGatherAt(in, out, lens, 0); + } + + /// `allGatherVChunked`'s ops done by the host (the kernel's side of each, as `hostGatherV`). + pub fn hostGatherVChunked(r: *Roce, in: []const u8, out: []u8, lens: [2]u32) Error!void { + var off: usize = 0; + while (off < in.len) : (off += r.slot_bytes) try r.hostGatherAt(in, out, lens, off); + } + + /// One op from shard offset `off`: each rank's bytes past it, at most a slot; an op empty on both ranks is skipped. + fn hostGatherAt(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32, off: usize) Error!void { + const n = in.len; + if (n == 0 or out.len < 2 * n) return error.Invalid; + if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; + const full: [2]usize = if (lens) |l| .{ l[0], l[1] } else .{ n, n }; + const mine: usize = @min(full[r.me] -| off, r.slot_bytes); + const theirs: usize = @min(full[1 - r.me] -| off, r.slot_bytes); + if (off > 0 and mine == 0 and theirs == 0) return; + const seq = r.host_epoch +% 1; + const slot: usize = seq & 1; + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[off..][0..mine]); + r.word(w_slot_bytes + slot).store(@intCast(@max(16, std.mem.alignForward(usize, mine, 16))), .release); + r.word(w_seq).store(seq, .release); + @memcpy(out[r.me * n + off ..][0..mine], in[off..][0..mine]); + const t0 = sock.nowNs(); + for (0..r.wire.hcas()) |h| { + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + (slot * 2 + h) * flag_stride ..].ptr)); + while (flag.load(.acquire) != seq) { + if (r.word(w_abort).load(.monotonic) != 0) return error.Aborted; + if (sock.nowNs() - t0 > r.timeout_ns) return error.Timeout; + std.atomic.spinLoopHint(); + } + } + @memcpy(out[(1 - r.me) * n + off ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + r.host_epoch = seq; + } +}; + +fn pin(cpu: u32) void { + var set: std.os.linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + std.os.linux.sched_setaffinity(0, &set) catch log.warn("tp roce: could not pin the proxy to CPU {d}", .{cpu}); +} diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig new file mode 100644 index 0000000000..c316448ef8 --- /dev/null +++ b/zig/src/fabric/tp/roce_fake.zig @@ -0,0 +1,143 @@ +//! Fake verbs for host tests of VerbsWire: each QP's posted writes counted from the requests, completions that lag. +const std = @import("std"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const rv = @import("roce_verbs.zig"); + +pub const slot: u64 = 1 << 20; +pub const send_off = 8192; + +/// What the HCAs hold (posted, not completed); every count comes from the work requests themselves. +pub const Fake = struct { + pub var region: []const u8 = &.{}; + /// The wire's slot (`Hw.wire`), every data write's bound with the send slots. + pub var slot_bytes: u64 = slot; + pub var lag: u32 = 1; + pub var polls: u64 = 0; + pub var queue: [2][512]u64 = undefined; + pub var head: [2]u32 = .{ 0, 0 }; + pub var len: [2]u32 = .{ 0, 0 }; + pub var inflight: u64 = 0; + pub var max_inflight: u64 = 0; + pub var max_write: u64 = 0; + pub var outside: u64 = 0; + pub var bytes: u64 = 0; + pub var posts: u64 = 0; + + pub fn reset(l: u32) void { + lag = l; + polls = 0; + head = .{ 0, 0 }; + len = .{ 0, 0 }; + inflight = 0; + max_inflight = 0; + max_write = 0; + outside = 0; + bytes = 0; + posts = 0; + } + + pub fn postSend(qp: *abi.Qp, wr: *abi.SendWr, _: *?*abi.SendWr) callconv(.c) c_int { + const q = qp.qp_num; + posts += 1; + var sum: u64 = 0; + var at: ?*abi.SendWr = wr; + while (at) |w| : (at = w.next) { + if (w.send_flags & abi.send_inline == 0) { + const sge = w.sg_list[0]; + const lo = @intFromPtr(region.ptr); + if (sge.addr < lo + send_off or sge.addr + sge.length > lo + send_off + 2 * slot_bytes) outside += 1; + max_write = @max(max_write, sge.length); + sum += sge.length; + } + if (w.send_flags & abi.send_signaled != 0) { + queue[q][(head[q] + len[q]) % 512] = sum; + len[q] += 1; + inflight += sum; + bytes += sum; + max_inflight = @max(max_inflight, inflight); + sum = 0; + } + } + return 0; + } + + /// One completion every `lag` polls of a QP's CQ: the HCA finishes writes far slower than they are posted. + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + polls += 1; + if (polls % lag != 0 or len[q] == 0) return 0; + inflight -= queue[q][head[q]]; + head[q] = (head[q] + 1) % 512; + len[q] -= 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = abi.wc_success; + return 1; + } +}; + +/// QPs that complete nothing on their own: a QP moved to error flushes what it holds, as an HCA reports it. +pub const Halts = struct { + pub var halted: [2]bool = .{ false, false }; + pub var flushed_before_halt: u32 = 0; + pub var flushes: u32 = 0; + + pub fn reset() void { + halted = .{ false, false }; + flushed_before_halt = 0; + flushes = 0; + } + + pub fn modifyQp(qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int) callconv(.c) c_int { + if (attr.qp_state == abi.qps_err and mask == abi.mask.state) halted[qp.qp_num] = true; + return 0; + } + + /// ibv_query_qp as the device answers it: ERR once halted, else RTS. + pub fn queryQp(qp: *abi.Qp, attr: *abi.QpAttr, _: c_int, _: *abi.QpInitAttr) callconv(.c) c_int { + attr.qp_state = if (halted[qp.qp_num]) abi.qps_err else abi.qps_rts; + return 0; + } + + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + if (Fake.len[q] == 0 or !halted[q]) return 0; + if (!halted[0] or !halted[1]) flushed_before_halt += 1; + Fake.inflight -= Fake.queue[q][Fake.head[q]]; + Fake.head[q] = (Fake.head[q] + 1) % 512; + Fake.len[q] -= 1; + flushes += 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = 5; // IBV_WC_WR_FLUSH_ERR + return 1; + } + + /// Points a wire's QPs at these handlers. + pub fn install(hw: *Hw, w: *rv.VerbsWire) void { + w.v.api.ibv_modify_qp = modifyQp; + w.v.query_qp = queryQp; + for (&hw.ctx) |*c| c.ops.poll_cq = pollCq; + } +}; + +/// Two fake HCAs behind one VerbsWire, connected to a peer whose region has this rank's layout. +pub const Hw = struct { + ctx: [2]abi.Context = undefined, + qp: [2]abi.Qp = undefined, + cq: [2]abi.Cq = undefined, + mr: [2]abi.Mr = undefined, + + pub fn wire(hw: *Hw, region: []u8) rv.VerbsWire { + const s = Fake.slot_bytes; + var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * s, .flag_off = 4096, .send_off = send_off, .slot_bytes = s, .region = region }; + for (0..2) |i| { + hw.ctx[i].ops = .{ ._before_poll = @splat(null), .poll_cq = Fake.pollCq, ._before_post = @splat(null), .post_send = Fake.postSend, .post_recv = undefined, ._after = @splat(null) }; + hw.qp[i].qp_num = @intCast(i); + hw.cq[i].handle = @intCast(i); + hw.mr[i].lkey = 7; + w.hcas[i] = .{ .ctx = &hw.ctx[i], .pd = undefined, .mr = &hw.mr[i], .cq = &hw.cq[i], .qp = &hw.qp[i], .gid_index = 0, .gid = @splat(0), .mtu = 0, .lid = 0 }; + } + w.n_hca = 2; + return w; + } +}; diff --git a/zig/src/fabric/tp/roce_guard.zig b/zig/src/fabric/tp/roce_guard.zig new file mode 100644 index 0000000000..8f85247afd --- /dev/null +++ b/zig/src/fabric/tp/roce_guard.zig @@ -0,0 +1,73 @@ +//! Every RoCE post's bounds before the verb: each HCA's stripe inside its slot and the registered region. +const std = @import("std"); + +pub const max_hcas = 2; +pub const flag_stride = 128; + +pub const Error = error{ OutOfRegion, WriteTooLarge }; + +/// The largest RDMA write a post may make, whatever the slot (a 9.6 MB post faulted a DGX Spark's HCA); never raised. +pub const max_write: u32 = 1 << 20; + +/// Both ranks' region layout (roce.zig's): byte offsets and the registered length every post must stay inside. +pub const Layout = struct { send: u64, recv: u64, flags: u64, slot: u64, len: u64 }; + +/// One HCA's share of a post: its data stripe in this rank's send slot and the peer's receive slot, and its flag word. +pub const Stripe = struct { local: u64, remote: u64, len: u32, flag: u64 }; + +/// Slot `slot`'s first `padded` bytes striped over `n_hca` HCAs; error.OutOfRegion before any verb past a slot or MR. +pub fn stripes(l: Layout, n_hca: u32, slot: u32, padded: u32, out: *[max_hcas]Stripe) Error!void { + if (n_hca == 0 or n_hca > max_hcas or slot > 1) return error.OutOfRegion; + if (padded > max_write) return error.WriteTooLarge; + if (padded == 0 or padded % 16 != 0 or padded > l.slot) return error.OutOfRegion; + const packs = padded / 16; + var offset: u64 = 0; + for (out[0..n_hca], 0..) |*o, i| { + var stripe = packs / n_hca; + if (i < packs % n_hca) stripe += 1; + o.* = .{ .local = l.send + slot * l.slot + offset, .remote = l.recv + slot * l.slot + offset, .len = stripe * 16, .flag = l.flags + (@as(u64, slot) * max_hcas + i) * flag_stride }; + if (!within(o.local, o.len, l.send + slot * l.slot, l.slot) or !within(o.remote, o.len, l.recv + slot * l.slot, l.slot)) return error.OutOfRegion; + if (!within(o.flag, 4, l.flags, l.send - l.flags) or !within(o.local, o.len, 0, l.len) or !within(o.remote, o.len, 0, l.len)) return error.OutOfRegion; + offset += o.len; + } +} + +fn within(at: u64, len: u64, lo: u64, size: u64) bool { + return at >= lo and at + len <= lo + size; +} + +/// roce.zig's layout for a 64 KiB slot: ctrl 4 KiB, flags 4 KiB, send [2][slot], recv [2][slot]. +fn testLayout() Layout { + const slot: u64 = 64 << 10; + return .{ .send = 8192, .recv = 8192 + 2 * slot, .flags = 4096, .slot = slot, .len = 8192 + 4 * slot }; +} + +test "a post striped over two HCAs: the stripes tile the slot's first bytes, each flag on its own word" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try stripes(l, 2, 1, 48, &s); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot, .remote = l.recv + l.slot, .len = 32, .flag = l.flags + 2 * flag_stride }, s[0]); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot + 32, .remote = l.recv + l.slot + 32, .len = 16, .flag = l.flags + 3 * flag_stride }, s[1]); + try stripes(l, 2, 1, @intCast(l.slot), &s); + try std.testing.expectEqual(l.recv + 2 * l.slot, s[1].remote + s[1].len); +} + +test "a post past its slot, a third slot, an unpadded or empty post, or no HCA is refused" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, @intCast(l.slot + 16), &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 2, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 24, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 0, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 0, 0, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 3, 0, 16, &s)); + var big = l; + big.slot = 10 << 20; + big.recv = big.send + 2 * big.slot; + big.len = big.send + 4 * big.slot; + try std.testing.expectError(error.WriteTooLarge, stripes(big, 2, 0, max_write + 16, &s)); + try stripes(big, 2, 0, max_write, &s); + var short = l; + short.len -= 16; + try std.testing.expectError(error.OutOfRegion, stripes(short, 2, 1, @intCast(l.slot), &s)); +} diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig new file mode 100644 index 0000000000..1dec07e9ba --- /dev/null +++ b/zig/src/fabric/tp/roce_host.zig @@ -0,0 +1,54 @@ +//! `tf-tp-test roce-host`: the RoCE wire and proxy between two hosts without a GPU, checked bit for bit, then timed. +const std = @import("std"); +const tp = @import("tp"); +const roce = tp.roce; +const rv = tp.roce_verbs; + +const sizes = [_]usize{ 16, 4096, 8192, 16384, 65536, 262144 }; + +pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { + var b = try tp.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var n: usize = 0; + if (cfg.roce_wire == .verbs) n = try tp.roce_pick.resolve(cfg, &b, &specs); + const r = try roce.Roce.init(null, &b, "", .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = 256 * 1024, .timeout_ns = 10 * std.time.ns_per_s }); + defer r.deinit(); + const in = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(in); + const peer = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(peer); + const out = try gpa.alloc(u8, 512 * 1024); + defer gpa.free(out); + // 1. every size from 1 byte to the slot, both slots, bit for bit + var checked: u32 = 0; + for (0..1000) |k| { + const len = 1 + (k * 7919) % (256 * 1024); + tp.reference.fill(in[0..len], .u8, cfg.rank, k); + tp.reference.fill(peer[0..len], .u8, 1 - cfg.rank, k); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + const mine_at = cfg.rank * len; + const peer_at = (1 - cfg.rank) * len; + if (!std.mem.eql(u8, out[mine_at..][0..len], in[0..len]) or !std.mem.eql(u8, out[peer_at..][0..len], peer[0..len])) { + std.debug.print("FAIL roce-host gather {d} of {d} bytes\n", .{ k, len }); + return 1; + } + checked += 1; + } + try r.check(); + std.debug.print("PASS roce-host: {d} gathers of 1 B to 256 KiB bit-exact over {s}\n", .{ checked, @tagName(cfg.roce_wire) }); + // 2. one exchange's host-to-host time by size: median and p90 of 2,000 back-to-back ops + var times: [2000]u64 = undefined; + for (sizes) |len| { + for (×) |*t| { + const t0 = tp.sock.nowNs(); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + t.* = tp.sock.nowNs() - t0; + } + std.mem.sort(u64, ×, {}, std.sort.asc(u64)); + if (cfg.rank == 0) std.debug.print("RESULT {{\"backend\":\"roce-host\",\"wire\":\"{s}\",\"bytes\":{d},\"p50_us\":{d:.2},\"p90_us\":{d:.2}}}\n", .{ @tagName(cfg.roce_wire), len, @as(f64, @floatFromInt(times[1000])) / 1e3, @as(f64, @floatFromInt(times[1800])) / 1e3 }); + } + try r.check(); + try b.agree("done", "roce-host end"); + return 0; +} diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig new file mode 100644 index 0000000000..fd7be53262 --- /dev/null +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -0,0 +1,110 @@ +//! VerbsWire on fake verbs whose completions lag: what it holds in flight and what a failed transport leaves behind. +const std = @import("std"); +const rv = @import("roce_verbs.zig"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Fake = fake.Fake; +const Halts = fake.Halts; +const slot = fake.slot; + +test "chunked posts: the wire's own in-flight count and peak match the HCAs', and reach the 2 MiB cap" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(3); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + for (1..9) |seq| { + try w.post(@intCast(seq & 1), @intCast(seq), slot); + try std.testing.expectEqual(Fake.inflight, w.inflight()); + try std.testing.expect(w.inflight() <= rv.max_inflight); + } + while (w.inflight() > 0) try w.reap(); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + try std.testing.expectEqual(Fake.max_inflight, w.peak.load(.monotonic)); + try std.testing.expectEqual(@as(u64, 2 << 20), w.peak.load(.monotonic)); +} + +test "a failed transport: both QPs go to error before the first flush and every posted write is reaped" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(1); + Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + Halts.install(&hw, &w); + // two 1 MiB ops posted, neither completed: 2 MiB in flight at the HCAs, the link busy + try std.testing.expect(w.idle()); + try w.post(1, 1, slot); + try w.post(0, 2, slot); + try std.testing.expectEqual(@as(u64, 2 * slot), w.inflight()); + try std.testing.expect(!w.idle()); + try std.testing.expect(w.quiesce(std.time.ns_per_s)); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expectEqual(@as(u32, 0), Halts.flushed_before_halt); + try std.testing.expectEqual(@as(u32, 4), Halts.flushes); + try std.testing.expectEqual(@as(u64, 0), w.inflight()); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + // the line tf-tp-test's transport run reads: the device's states and every reap, all of them flushes + var line: [256]u8 = undefined; + try std.testing.expect(w.idle()); + try std.testing.expectEqualStrings("HCA 0 QP ERR; posted 2, completed 2; HCA 1 QP ERR; posted 2, completed 2; reaped 4 (flushed 4), still posted 0", w.report(&line)); + // an aborted transport reports itself failed: its deinit keeps the region + Halts.reset(); + var r: Roce = .{ .me = 0, .slot_bytes = slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + try std.testing.expect(!r.failed()); + r.abort(); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expect(r.failed()); +} + +test "a QP whose inline capacity cannot take the 4-byte flag post is refused" { + try std.testing.expect(!rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 3 })); + try std.testing.expect(rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 4 })); +} + +/// /proc/self/smaps' VmFlags of the mapping that contains `at`. +fn vmFlags(at: usize, buf: []u8) ![]const u8 { + const fd = std.c.open("/proc/self/smaps", .{}); + if (fd < 0) return error.NoSmaps; + defer _ = std.c.close(fd); + var n: usize = 0; + while (n < buf.len) { + const got = std.c.read(fd, buf[n..].ptr, buf.len - n); + if (got <= 0) break; + n += @intCast(got); + } + var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); + var inside = false; + while (lines.next()) |l| { + const sp = std.mem.indexOfScalar(u8, l, ' ') orelse l.len; + if (std.mem.indexOfScalar(u8, l[0..sp], '-')) |dash| { + const lo = std.fmt.parseInt(usize, l[0..dash], 16) catch continue; + const hi = std.fmt.parseInt(usize, l[dash + 1 .. sp], 16) catch continue; + inside = lo <= at and at < hi; + } + if (inside and std.mem.startsWith(u8, l, "VmFlags:")) return l; + } + return error.NotFound; +} + +test "the region is DONTFORK and NOHUGEPAGE before registration" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; + const len = Roce.regionBytes(slot); + const m = try std.posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + defer std.posix.munmap(m); + const buf = try std.testing.allocator.alloc(u8, 4 << 20); + defer std.testing.allocator.free(buf); + // smaps: dc = MADV_DONTFORK, nh = MADV_NOHUGEPAGE + const before = try vmFlags(@intFromPtr(m.ptr), buf); + try std.testing.expect(std.mem.indexOf(u8, before, " dc") == null and std.mem.indexOf(u8, before, " nh") == null); + try rv.prepareRegion(m); + // An interior page stays inside the VMA even when a neighbouring mapping merges with it. + const flags = try vmFlags(@intFromPtr(m.ptr) + 4096, buf); + try std.testing.expect(std.mem.indexOf(u8, flags, " dc") != null); + try std.testing.expect(std.mem.indexOf(u8, flags, " nh") != null); +} diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig new file mode 100644 index 0000000000..b9355c401a --- /dev/null +++ b/zig/src/fabric/tp/roce_pick.zig @@ -0,0 +1,386 @@ +//! Which RoCE devices a rank stripes over: TF_TP_ROCE_HCA, else the usable ports on an IPv4 /24 the peer also has. +const std = @import("std"); +const log = @import("log.zig"); +const rv = @import("roce_verbs.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const Verbs = @import("verbs").Verbs; + +/// Ports a rank lists to its peers (a DGX Spark's ConnectX-7 shows four RDMA devices). +pub const max_ports = 8; +const blob_bytes = 1 + 4 * max_ports; + +/// How many devices, and how: on a /24 the peer shares (`shared`), named by TF_TP_ROCE_HCA (`explicit`), or neither. +pub const Pick = struct { n: usize, shared: bool, explicit: bool = false }; + +fn usable(p: rv.Port) bool { + return p.active and p.ethernet and p.ipv4 != null; +} + +fn subnet(a: [4]u8) u32 { + return @as(u32, a[0]) << 16 | @as(u32, a[1]) << 8 | a[2]; +} + +fn countIn(list: []const [4]u8, net: u32) usize { + var k: usize = 0; + for (list) |a| k += @intFromBool(subnet(a) == net); + return k; +} + +/// Usable ports on a /24 the peers list, by subnet, each /24 as often as both sides have it; else the first usable. +pub fn choose(ports: []const rv.Port, peer: []const [4]u8, out: *[rv.max_hcas]rv.HcaSpec, limit: usize) Pick { + const cap = @min(limit, rv.max_hcas); + var order: [max_ports]usize = undefined; + var m: usize = 0; + for (ports[0..@min(ports.len, max_ports)], 0..) |p, i| if (usable(p)) { + order[m] = i; + m += 1; + }; + var n: usize = 0; + for (order[0..m]) |i| n += @intFromBool(countIn(peer, subnet(ports[i].ipv4.?)) > 0); + if (n == 0) { + for (order[0..@min(m, cap)], 0..) |i, k| out[k] = ports[i].spec; + return .{ .n = @min(m, cap), .shared = false }; + } + // stable insertion sort by subnet (at most max_ports entries) + for (1..m) |a| { + var j = a; + while (j > 0 and subnet(ports[order[j - 1]].ipv4.?) > subnet(ports[order[j]].ipv4.?)) : (j -= 1) { + std.mem.swap(usize, &order[j - 1], &order[j]); + } + } + n = 0; + var taken: [max_ports]u32 = undefined; + for (order[0..m]) |i| { + if (n == cap) break; + const net = subnet(ports[i].ipv4.?); + var mine: usize = 0; + for (taken[0..n]) |t| mine += @intFromBool(t == net); + if (mine >= countIn(peer, net)) continue; + taken[n] = net; + out[n] = ports[i].spec; + n += 1; + } + return .{ .n = n, .shared = true }; +} + +/// `[count][count x IPv4]` of the usable ports, the blob each rank sends. +fn encode(ports: []const rv.Port, out: *[blob_bytes]u8) []const u8 { + var n: u8 = 0; + for (ports[0..@min(ports.len, max_ports)]) |p| if (usable(p)) { + out[1 + 4 * @as(usize, n) ..][0..4].* = p.ipv4.?; + n += 1; + }; + out[0] = n; + return out[0 .. 1 + 4 * @as(usize, n)]; +} + +/// A peer's blob into `out`; a malformed one counts as no ports. +fn decode(blob: []const u8, out: [][4]u8) usize { + if (blob.len == 0 or blob[0] > max_ports or blob.len != 1 + 4 * @as(usize, blob[0]) or out.len < blob[0]) return 0; + for (0..blob[0]) |k| out[k] = blob[1 + 4 * k ..][0..4].*; + return blob[0]; +} + +/// The devices for the verbs wire. A collective: every rank calls it, including ranks that set TF_TP_ROCE_HCA. +pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !usize { + return (try select(cfg, b, out)).n; +} + +/// `resolve` with how the devices were chosen, over this host's ports (a collective, as `resolve`). +pub fn select(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + var ports: [max_ports]rv.Port = undefined; + var n_ports: usize = 0; + if (Verbs.open(null)) |opened| { + var v = opened; + defer v.close(); + n_ports = rv.scanPorts(&v, &ports); + } else |_| {} + return selectFrom(cfg, b, ports[0..n_ports], out); +} + +/// `select` over given ports (the scan's, or a test's fixture ports). +pub fn selectFrom(cfg: Config, b: *bootstrap.Bootstrap, ports: []const rv.Port, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + var mine: [blob_bytes]u8 = undefined; + var buf: [bootstrap.max_ranks * blob_bytes]u8 = undefined; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + const all = try b.allGather(encode(ports, &mine), &buf, &blobs); + if (cfg.roce_hca) |text| return .{ .n = try rv.parseSpecs(text, out), .shared = false, .explicit = true }; + var peer: [bootstrap.max_ranks * max_ports][4]u8 = undefined; + var n_peer: usize = 0; + for (all, 0..) |blob, r| if (r != b.rank) { + n_peer += decode(blob, peer[n_peer..]); + }; + const pick = choose(ports, peer[0..n_peer], out, cfg.roce_hcas); + if (!pick.shared and pick.n > 0) log.warn("tp roce: no RoCE port shares an IPv4 /24 with the peer's; using the first {d} active (set TF_TP_ROCE_HCA)", .{pick.n}); + return pick; +} + +/// Device names joined by commas, with each `:gid` when `gid` (TF_TP_ROCE_HCA's form) or without (NCCL_IB_HCA's). +pub fn joinNames(specs: []const rv.HcaSpec, gid: bool, buf: []u8) error{WriteFailed}![]const u8 { + var w: std.Io.Writer = .fixed(buf); + for (specs, 0..) |s, i| { + try w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }); + if (gid) if (s.gid) |g| try w.print(":{d}", .{g}); + } + return w.buffered(); +} + +/// Room for `joinNames` of every device a rank stripes over (a name is under 64 bytes, a GID index at most 3 digits). +pub const list_bytes = rv.max_hcas * (64 + 4); + +pub const HcaError = error{ + /// no active Ethernet port with an IPv4 RoCE v2 GID on this host + NoRoceDevice, + /// this host's ports share no IPv4 /24 with the peer's: NCCL would not reach the peer on them either + NoSharedSubnet, +}; + +/// `tf-tp-test hca`'s line for run_hosts.sh: `NCCL_IB_HCA=` the picked devices; never the first local ports alone. +pub fn hcaLine(pick: Pick, specs: []const rv.HcaSpec, buf: *[list_bytes + 16]u8) HcaError![]const u8 { + if (pick.n == 0) return error.NoRoceDevice; + if (!pick.shared and !pick.explicit) return error.NoSharedSubnet; + const prefix = "NCCL_IB_HCA="; + @memcpy(buf[0..prefix.len], prefix); + const joined = joinNames(specs[0..pick.n], false, buf[prefix.len .. buf.len - 1]) catch unreachable; // list_bytes holds them + buf[prefix.len + joined.len] = '\n'; + return buf[0 .. prefix.len + joined.len + 1]; +} + +fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { + var p: rv.Port = .{ .active = active, .ethernet = true, .ipv4 = ipv4 }; + p.spec.name_len = @intCast(name.len); + @memcpy(p.spec.name[0..name.len], name); + p.spec.gid = 3; + return p; +} + +fn names(specs: []const rv.HcaSpec, buf: []u8) []const u8 { + return joinNames(specs, false, buf) catch unreachable; +} + +// A host with two neighbours: port f0 cabled to one Spark (both functions on one /24), port f1 to the other. +const a_ports = [_]rv.Port{ + port("rocep1s0f0", true, .{ 203, 0, 113, 1 }), + port("rocep1s0f1", true, .{ 192, 0, 2, 1 }), + port("roceP2p1s0f0", true, .{ 203, 0, 113, 65 }), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 1 }), +}; +// Its peer for this run: only port f1 is cabled, f0 is down. +const b_ports = [_]rv.Port{ + port("rocep1s0f0", false, null), + port("rocep1s0f1", true, .{ 192, 0, 2, 2 }), + port("roceP2p1s0f0", false, null), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 2 }), +}; + +fn listed(ports: []const rv.Port, out: *[max_ports][4]u8) []const [4]u8 { + var blob: [blob_bytes]u8 = undefined; + return out[0..decode(encode(ports, &blob), out)]; +} + +test "a two-neighbour host picks the rails its peer shares, not its first two active ports" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the old rule (no peer list): the first two active ports, one of them cabled to the other neighbour + const old = choose(&a_ports, &.{}, &specs, 2); + try std.testing.expect(!old.shared); + try std.testing.expectEqualStrings("rocep1s0f0,rocep1s0f1", names(specs[0..old.n], &buf)); + var peer_buf: [max_ports][4]u8 = undefined; + const b_addrs = listed(&b_ports, &peer_buf); + try std.testing.expectEqual(@as(usize, 0), countIn(b_addrs, subnet(.{ 203, 0, 113, 1 }))); + // with the peer's addresses: the two functions of the cabled port, on each side + const a = choose(&a_ports, b_addrs, &specs, 2); + try std.testing.expect(a.shared); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..a.n], &buf)); + var a_buf: [max_ports][4]u8 = undefined; + const b = choose(&b_ports, listed(&a_ports, &a_buf), &specs, 2); + try std.testing.expectEqual(@as(usize, 2), b.n); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..b.n], &buf)); + // one rail asked for: the lower /24 on both sides + try std.testing.expectEqualStrings("rocep1s0f1", names(specs[0..choose(&a_ports, b_addrs, &specs, 1).n], &buf)); +} + +test "two rails on one /24 count once per peer port; no shared /24 keeps the list order" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the neighbour cabled to port f0 lists one address on that /24: one rail each, equal counts + const one = [_][4]u8{.{ 203, 0, 113, 2 }}; + const p = choose(&a_ports, &one, &specs, 2); + try std.testing.expectEqual(@as(usize, 1), p.n); + try std.testing.expectEqualStrings("rocep1s0f0", names(specs[0..p.n], &buf)); + const two = [_][4]u8{ .{ 203, 0, 113, 2 }, .{ 203, 0, 113, 66 } }; + try std.testing.expectEqualStrings("rocep1s0f0,roceP2p1s0f0", names(specs[0..choose(&a_ports, &two, &specs, 2).n], &buf)); + // inactive and non-Ethernet ports never count; a peer on other subnets falls back to the first usable ports + var mixed = b_ports; + mixed[1].ethernet = false; + const other = [_][4]u8{.{ 198, 51, 100, 9 }}; + const q = choose(&mixed, &.{.{ 203, 0, 113, 9 }}, &specs, 2); + try std.testing.expect(!q.shared); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..q.n], &buf)); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..choose(&mixed, &other, &specs, 2).n], &buf)); + // a malformed blob is no ports; the detected GID index travels with the pick + var out: [max_ports][4]u8 = undefined; + try std.testing.expectEqual(@as(usize, 0), decode(&.{ 2, 1, 2, 3, 4 }, &out)); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); +} + +test "hca's line: the picked names for NCCL; no device or no shared /24 is refused, an explicit list is taken" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var line: [list_bytes + 16]u8 = undefined; + var peer_buf: [max_ports][4]u8 = undefined; + const shared = choose(&a_ports, listed(&b_ports, &peer_buf), &specs, 2); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(shared, &specs, &line)); + var buf: [list_bytes]u8 = undefined; + try std.testing.expectEqualStrings("rocep1s0f1:3,roceP2p1s0f1:3", try joinNames(specs[0..shared.n], true, &buf)); + // the old fallback (first active ports, no /24 in common) is what NCCL failed on: refused, never printed + try std.testing.expectError(error.NoSharedSubnet, hcaLine(choose(&a_ports, &.{.{ 10, 9, 9, 9 }}, &specs, 2), &specs, &line)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(choose(&.{}, &.{}, &specs, 2), &specs, &line)); + const n = try rv.parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(.{ .n = n, .shared = false, .explicit = true }, &specs, &line)); +} + +/// One RDMA device for `rv.scanPorts` through a stub verbs table: port 1 up or down, Ethernet, mlx5's GIDs 0-3. +const FakeDev = struct { name: [:0]const u8, up: bool, ipv4: [4]u8 = @splat(0) }; + +const ScanFixture = struct { + const vabi = @import("verbs").abi; + var host: []const FakeDev = &.{}; + var list: [max_ports:null]?*vabi.Device = undefined; + + fn handle(i: usize) usize { + return 0x10000 * (i + 1); + } + fn index(p: *const anyopaque) usize { + return @intFromPtr(p) / 0x10000 - 1; + } + fn getList(num: *c_int) callconv(.c) ?[*:null]?*vabi.Device { + for (host, 0..) |_, i| list[i] = @ptrFromInt(handle(i)); + list[host.len] = null; + num.* = @intCast(host.len); + return &list; + } + fn freeList(_: [*:null]?*vabi.Device) callconv(.c) void {} + fn devName(d: *vabi.Device) callconv(.c) ?[*:0]const u8 { + return host[index(d)].name.ptr; + } + fn openDev(d: *vabi.Device) callconv(.c) ?*vabi.Context { + return @ptrFromInt(@intFromPtr(d)); + } + fn closeDev(_: *vabi.Context) callconv(.c) c_int { + return 0; + } + fn queryPort(c: *vabi.Context, _: u8, a: *vabi.PortAttr) callconv(.c) c_int { + a.* = std.mem.zeroes(vabi.PortAttr); + a.state = if (host[index(c)].up) vabi.port_active else 1; + a.link_layer = 2; // Ethernet + a.gid_tbl_len = 4; + return 0; + } + fn queryGid(c: *vabi.Context, _: u8, i: c_int, g: *vabi.Gid) callconv(.c) c_int { + g.* = .{ .raw = @splat(0) }; + if (i < 2) { + g.raw[0] = 0xfe; + g.raw[1] = 0x80; + g.raw[15] = @intCast(index(c) + 1); + } else { + g.raw[10] = 0xff; + g.raw[11] = 0xff; + g.raw[12..16].* = host[index(c)].ipv4; + } + return 0; + } + fn gidEx(_: *vabi.Context, _: u32, i: u32, e: *vabi.GidEntry, _: u32, _: usize) callconv(.c) c_int { + e.* = std.mem.zeroes(vabi.GidEntry); + e.gid_type = if (i % 2 == 1) 2 else 1; // odd indices RoCE v2 + return 0; + } + + /// `rv.scanPorts` over `devs` (names that no real /sys/class/infiniband has, so sysfs never answers for them). + fn scan(devs: []const FakeDev, out: *[max_ports]rv.Port) []const rv.Port { + host = devs; + var v: Verbs = undefined; + v.api.ibv_get_device_list = getList; + v.api.ibv_free_device_list = freeList; + v.api.ibv_get_device_name = devName; + v.api.ibv_open_device = openDev; + v.api.ibv_close_device = closeDev; + v.api.ibv_query_port = queryPort; + v.api.ibv_query_gid = queryGid; + v.gid_ex = gidEx; + return out[0..rv.scanPorts(&v, out)]; + } +}; + +/// #636's reported two GB10 hosts, modelled (no capture): device order and picks reported; addresses, states assumed. +const hosts = [2][4]FakeDev{ .{ + .{ .name = "fx_rocep1s0f0", .up = true, .ipv4 = .{ 192, 0, 2, 1 } }, + .{ .name = "fx_rocep1s0f1", .up = false }, + .{ .name = "fx_roceP2p1s0f0", .up = true, .ipv4 = .{ 198, 51, 100, 1 } }, + .{ .name = "fx_roceP2p1s0f1", .up = false }, +}, .{ + .{ .name = "fx_rocep1s0f0", .up = false }, + .{ .name = "fx_rocep1s0f1", .up = true, .ipv4 = .{ 192, 0, 2, 2 } }, + .{ .name = "fx_roceP2p1s0f0", .up = false }, + .{ .name = "fx_roceP2p1s0f1", .up = true, .ipv4 = .{ 198, 51, 100, 2 } }, +} }; + +test "the port scan reads a host's devices through verbs: state, Ethernet, the IPv4 RoCE v2 GID index" { + var out: [max_ports]rv.Port = undefined; + const a = ScanFixture.scan(&hosts[0], &out); + try std.testing.expectEqual(@as(usize, 4), a.len); + try std.testing.expectEqualStrings("fx_rocep1s0f0", a[0].spec.name[0..a[0].spec.name_len]); + try std.testing.expect(a[0].active and a[0].ethernet and !a[1].active); + try std.testing.expectEqual(@as(?[4]u8, .{ 192, 0, 2, 1 }), a[0].ipv4); + try std.testing.expectEqual(@as(?u8, 3), a[0].spec.gid); // the v2 IPv4 GID, not v1's index 2 + try std.testing.expectEqual(@as(?[4]u8, null), a[1].ipv4); +} + +/// One rank's `selectFrom` over a real bootstrap (rank 0 on `l`), its hca line or error into `out`. +fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, out: *[list_bytes + 16]u8, line: *HcaError![]const u8, err: *?anyerror) void { + var b = (if (l) |fd| bootstrap.Bootstrap.acceptOn(fd, 2, 5 * std.time.ns_per_s) else bootstrap.Bootstrap.open(cfg)) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + const pick = selectFrom(cfg, &b, ports, &specs) catch |e| { + err.* = e; + return; + }; + line.* = hcaLine(pick, &specs, out); +} + +fn hcaPair(cfgs: [2]Config, ports: [2][]const rv.Port, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + var errs: [2]?anyerror = @splat(null); + var c1 = cfgs[1]; + c1.link_port = l.port; + c1.connect_timeout_ns = 5 * std.time.ns_per_s; + const t = try std.Thread.spawn(.{}, hcaRank, .{ c1, null, ports[1], &bufs[1], &lines[1], &errs[1] }); + hcaRank(cfgs[0], l.fd, ports[0], &bufs[0], &lines[0], &errs[0]); + t.join(); + for (errs) |e| if (e) |x| return x; +} + +test "hca over a real bootstrap: both ranks print their peer-linked ports; an explicit rank keeps its list; no ports fail both" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [2][max_ports]rv.Port = undefined; + const ports = [2][]const rv.Port{ ScanFixture.scan(&hosts[0], &scanned[0]), ScanFixture.scan(&hosts[1], &scanned[1]) }; + const auto: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPair(auto, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0,fx_roceP2p1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // TF_TP_ROCE_HCA on one rank: its list as given (no GID for NCCL); it still sends its ports to the other + const one: [2]Config = .{ .{ .rank = 0, .roce_hca = "fx_rocep1s0f0:3" }, .{ .rank = 1 } }; + try hcaPair(one, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // a host without usable ports: it has none, its peer shares no /24 with it, and both refuse + try hcaPair(auto, .{ &.{}, ports[1] }, &lines, &bufs); + try std.testing.expectError(error.NoRoceDevice, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig new file mode 100644 index 0000000000..a284293000 --- /dev/null +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -0,0 +1,120 @@ +//! VerbsWire.post on fake verbs: a post that would leave its slot or region is refused before any verb. +const std = @import("std"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Config = @import("config.zig").Config; +const Fake = fake.Fake; + +test "a post past its slot is refused before any verb; a full slot posts inside the region" { + const small: u64 = 64 << 10; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(small)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = small; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.OutOfRegion, w.post(1, 1, small + 16)); + try std.testing.expectError(error.OutOfRegion, w.post(2, 1, 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(1, 1, small); + try std.testing.expectEqual(@as(u64, 2), Fake.posts); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); +} + +test "no RDMA write past 1 MiB, whatever the slot: a 9.57 MB and a 1 MiB + 16 post refused in a 10 MiB slot" { + const big: u64 = 10 << 20; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(big)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = big; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.WriteTooLarge, w.post(0, 1, 9_568_256)); + try std.testing.expectError(error.WriteTooLarge, w.post(1, 1, (1 << 20) + 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(0, 1, 1 << 20); + try std.testing.expectEqual(@as(u64, 1 << 20), Fake.max_write * 2); +} + +const Env = struct { + var pairs: []const [2][]const u8 = &.{}; + + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "boot refuses a RoCE slot past 1 MiB (TF_TP_MAILBOX_MAX_KB), never clamping it" { + for ([_][]const u8{ "1028", "2048", "16384" }) |kb| { + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", kb } }; + try std.testing.expectError(error.BadSetting, Config.parse(&Env.get)); + } + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" } }; + try std.testing.expectEqual(@as(usize, 1 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); + // the cap is RoCE's: the same-host mailbox keeps its larger slots + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "4096" } }; + try std.testing.expectEqual(@as(usize, 4 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); +} + +test "chunked strides posted back to back: at most 2 MiB in flight over both HCAs and no write past 1 MiB" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = fake.slot; + // gathers of R x 512 x 584 B strides and 584 B rows up to a 128 MiB union, each in <= 1 MiB ops + var strides: [16]u64 = undefined; + var n: usize = 0; + for ([_]u64{ 1, 2, 3, 4, 6, 16, 32, 64 }) |r| { + strides[n] = r * 512 * 584; + n += 1; + } + for ([_]u64{ 1795, 1796, (128 << 20) / (3 * 584) }) |m| { + strides[n] = m * 584; + n += 1; + } + for ([_]u32{ 1, 3, 17 }) |lag| { + Fake.reset(lag); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + var seq: u32 = 0; + var want: u64 = 0; + for (strides[0..n]) |stride| { + var off: u64 = 0; + while (off < stride) : (off += fake.slot) { + seq += 1; + const padded: u32 = @intCast(std.mem.alignForward(u64, @min(fake.slot, stride - off), 16)); + try w.post(seq & 1, seq, padded); + want += padded; + } + } + while (Fake.inflight > 0) try w.reap(); + try std.testing.expectEqual(want, Fake.bytes); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); + // the cap's numbers, not roce_verbs.zig's constants: 1 MiB a write, 2 MiB in flight + try std.testing.expect(Fake.max_write <= 1 << 20); + try std.testing.expect(Fake.max_inflight <= 2 << 20); + // the bound is reached, not avoided: posts outran completions + try std.testing.expect(Fake.max_inflight > 1 << 20); + } +} + +test "Roce.abort (the fail-fast hook) moves both QPs to error: no send goes out and no peer write lands" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = fake.slot; + Fake.reset(1); + fake.Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + fake.Halts.install(&hw, &w); + var r: Roce = .{ .me = 0, .slot_bytes = fake.slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + r.abort(); + try std.testing.expect(fake.Halts.halted[0] and fake.Halts.halted[1]); +} diff --git a/zig/src/fabric/tp/roce_shm.zig b/zig/src/fabric/tp/roce_shm.zig new file mode 100644 index 0000000000..3924be3e96 --- /dev/null +++ b/zig/src/fabric/tp/roce_shm.zig @@ -0,0 +1,63 @@ +//! The RoCE protocol's shm wire: both ranks' regions in one shared mapping, posts through the guard's stripes. +const std = @import("std"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const guard = @import("roce_guard.zig"); +const posix = std.posix; +const page = std.heap.page_size_min; + +/// A test's view of each post: slot, padded bytes and the stripes the wire is about to move. +pub const PostTap = struct { + ctx: *anyopaque, + seen: *const fn (ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void, +}; + +/// The peer's region is mapped here; a post copies each stripe, then stores its flag with release order. +pub const ShmWire = struct { + peer: []u8, + layout: guard.Layout, + local: []u8, + /// Stripes a post makes, as HCAs would: the verbs wire's split, flags and bounds on one host. + stripes: u32 = 1, + tap: ?*PostTap = null, + + pub fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) guard.Error!void { + var plan: [guard.max_hcas]guard.Stripe = undefined; + try guard.stripes(w.layout, w.stripes, slot, padded, &plan); + if (w.tap) |t| t.seen(t.ctx, w.layout, slot, padded, plan[0..w.stripes]); + for (plan[0..w.stripes]) |s| { + @memcpy(w.peer[s.remote..][0..s.len], w.local[s.local..][0..s.len]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[s.flag..].ptr)); + flag.store(seq, .release); + } + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +pub fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + if (b.rank != 0) { + fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + try b.agree("mapped", "RoCE shm wire"); + if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); + return m; +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig new file mode 100644 index 0000000000..ec6bc4d65e --- /dev/null +++ b/zig/src/fabric/tp/roce_test.zig @@ -0,0 +1,173 @@ +//! The RoCE protocol end to end on host threads: two ranks, TCP bootstrap, proxy, shm or verbs (TF_TP_ROCE_TEST_HCA). +const std = @import("std"); +const roce = @import("roce.zig"); +const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const reference = @import("reference.zig"); +const Bootstrap = bootstrap.Bootstrap; + +const Run = struct { + port: u16, + wire: roce.WireKind, + devs: [2][]const u8 = .{ "", "" }, + ops: u32, + timeout_ns: u64 = 10 * std.time.ns_per_s, + /// Rank 1 stops after this many ops (rank 0 must then fail its next wait, by timeout or abort). + stop_after: ?u32 = null, + abort_after_ms: ?u64 = null, + errs: [2]?anyerror = .{ null, null }, + done: [2]u32 = .{ 0, 0 }, + max_bytes: usize = 16 * 1024, + /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. + var_lens: bool = false, + /// Strides up to four slots through `hostGatherVChunked` (variable lengths, slot-sized ops). + chunked: bool = false, + /// The shm wire's stripes (2: the verbs wire's split over two HCAs). + stripes: u32 = 1, + /// Each rank's shm posts as the tap saw them: how many, and how many left their slot or the region. + taps: [2]Tap = .{ .{}, .{} }, +}; + +/// Counts an shm wire's posts and any stripe outside its slot, its flag words or the region. +const Tap = struct { + posts: u32 = 0, + outside: u32 = 0, + tap: shm.PostTap = undefined, + + fn seen(ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void { + const t: *Tap = @ptrCast(@alignCast(ctx)); + t.posts += 1; + var sum: u64 = 0; + for (s) |x| { + sum += x.len; + const in_slot = x.local >= l.send + slot * l.slot and x.local + x.len <= l.send + (slot + 1) * l.slot; + const in_peer = x.remote >= l.recv + slot * l.slot and x.remote + x.len <= l.recv + (slot + 1) * l.slot; + if (!in_slot or !in_peer or x.remote + x.len > l.len or x.flag < l.flags or x.flag + 4 > l.send) t.outside += 1; + } + if (sum != padded) t.outside += 1; + } +}; + +fn rank(run: *Run, me: u32, lfd: sock.Fd) void { + run.errs[me] = body(run, me, lfd); +} + +fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { + var b = (if (me == 0) Bootstrap.acceptOn(lfd, 2, 5 * std.time.ns_per_s) else Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = run.port, .connect_timeout_ns = 5 * std.time.ns_per_s })) catch |e| return e; + defer b.close(); + var spec: [1]rv.HcaSpec = .{.{}}; + if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; + const mine = &run.taps[me]; + mine.tap = .{ .ctx = mine, .seen = Tap.seen }; + const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns, .shm_stripes = run.stripes, .tap = &mine.tap }) catch |e| return e; + defer r.deinit(); + var abort_thread: ?std.Thread = null; + defer if (abort_thread) |t| t.join(); + var in: [64 * 1024]u8 = undefined; + var peer: [64 * 1024]u8 = undefined; + var out: [128 * 1024]u8 = undefined; + for (0..run.ops) |k| { + if (run.stop_after) |s| if (k == s) { + // rank 1 skips op s; rank 0 waits in it until its timeout, or until the abort armed now + if (me == 1) { + sock.sleepNs(300 * std.time.ns_per_ms); + return null; + } + if (run.abort_after_ms) |ms| abort_thread = std.Thread.spawn(.{}, abortLater, .{ r, ms }) catch null; + }; + // sizes from 1 byte to the slot, odd ones included, both slots and wrap-arounds + const n = 1 + (k * 7919) % (if (run.chunked) 4 * run.max_bytes else run.max_bytes); + reference.fill(in[0..n], .u8, me, k); + reference.fill(peer[0..n], .u8, 1 - me, k); + const lens: [2]u32 = if (k % 5 == 0) .{ 0, @intCast(n) } else .{ @intCast((k * 31) % (n + 1)), @intCast((k * 17 + 5) % (n + 1)) }; + const l: [2]usize = if (run.var_lens or run.chunked) .{ lens[0], lens[1] } else .{ n, n }; + if (run.chunked) { + r.hostGatherVChunked(in[0..n], out[0 .. 2 * n], lens) catch |e| return e; + } else r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + const r0 = if (me == 0) in[0..n] else peer[0..n]; + const r1 = if (me == 1) in[0..n] else peer[0..n]; + if (!std.mem.eql(u8, out[0..l[0]], r0[0..l[0]]) or !std.mem.eql(u8, out[n..][0..l[1]], r1[0..l[1]])) return error.WrongBytes; + run.done[me] += 1; + } + r.check() catch |e| return e; + // both ranks finish before either tears down (the peer may still read a flag this rank's wire writes) + b.agree("done", "test end") catch |e| return e; + return null; +} + +fn abortLater(r: *roce.Roce, ms: u64) void { + sock.sleepNs(ms * std.time.ns_per_ms); + r.abort(); +} + +fn twoRanks(run: *Run) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + run.port = l.port; + const t1 = try std.Thread.spawn(.{}, rank, .{ run, 1, l.fd }); + rank(run, 0, l.fd); + t1.join(); +} + +test "the RoCE protocol over the shm wire: 2,000 gathers of every size, bit for bit" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the RoCE protocol with variable lengths over the shm wire: each rank posts its own padded count, empty shards included" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the shm wire striped as over two HCAs: every post inside its slot and the region, bytes and lengths exact" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expectEqual(@as(u32, 2000), t.posts); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + +test "chunked gathers over the shm wire: strides up to four slots in slot-sized ops, tails empty on both ranks skipped" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .chunked = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expect(t.posts > 2000); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + +test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; + try twoRanks(&run); + try std.testing.expectEqual(@as(?anyerror, error.Timeout), run.errs[0]); + try std.testing.expectEqual(@as(u32, 20), run.done[0]); + var run2: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .abort_after_ms = 50 }; + const kept = roce.Roce.kept.load(.monotonic); + try twoRanks(&run2); + // the aborted rank keeps its region until exit, the other frees it + try std.testing.expectEqual(kept + 1, roce.Roce.kept.load(.monotonic)); + try std.testing.expectEqual(@as(?anyerror, error.Aborted), run2.errs[0]); + try std.testing.expectEqual(@as(?anyerror, null), run2.errs[1]); +} + +test "the RoCE protocol over verbs (TF_TP_ROCE_TEST_HCA=dev0,dev1: soft-RoCE or real HCAs)" { + const raw = std.c.getenv("TF_TP_ROCE_TEST_HCA") orelse return error.SkipZigTest; + var it = std.mem.splitScalar(u8, std.mem.span(raw), ','); + var run: Run = .{ .port = 0, .wire = .verbs, .ops = 2000 }; + run.devs = .{ it.next() orelse return error.SkipZigTest, it.next() orelse return error.SkipZigTest }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig new file mode 100644 index 0000000000..6917d292be --- /dev/null +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -0,0 +1,455 @@ +//! The RoCE wire over verbs: one RC queue pair a HCA to the peer, a send slot striped over the HCAs with seq flags. +const std = @import("std"); +const log = @import("log.zig"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const Verbs = verbs.Verbs; +const guard = @import("roce_guard.zig"); +const sock = @import("sock.zig"); + +pub const max_hcas = guard.max_hcas; +/// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). +const send_depth = 256; +const link_ethernet = 2; +const record_magic: u32 = 0x5446_5244 + 2; // "TFRD", bumped for each change of this record's layout + +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord } || guard.Error; + +/// Bytes posted and not yet completed, every HCA together: a post waits for completions past it; never raised. +pub const max_inflight: u64 = 2 * @as(u64, guard.max_write); +/// A post that waits on completions this long fails the proxy (a peer or HCA that stopped answering). +const inflight_wait_ns: u64 = 10 * std.time.ns_per_s; +const ring_len = send_depth / 4; +/// The flag post's bytes: a u32 on the stack under lkey 0, legal only while the QP sends it inline. +pub const flag_bytes: u32 = @sizeOf(u32); +/// The longest `quiesce` waits for the HCAs to flush what was posted. +pub const quiesce_ns: u64 = 5 * std.time.ns_per_s; + +/// Before ibv_reg_mr: no fork shares or copies the region (DONTFORK), its pages stay 4 KiB (NOHUGEPAGE); else no start. +pub fn prepareRegion(region: []align(std.heap.page_size_min) u8) Error!void { + if (comptime @import("builtin").os.tag != .linux) return; // the verbs wire runs on Linux only + const advice = [_]u32{ std.os.linux.MADV.DONTFORK, std.os.linux.MADV.NOHUGEPAGE }; + for (advice) |a| std.posix.madvise(region.ptr, region.len, a) catch |e| { + log.warn("tp roce: madvise({d}) on the {d} B region failed: {t}; refusing to start", .{ a, region.len, e }); + return error.VerbsFailed; + }; +} + +/// Whether a QP's actual inline capacity takes the flag post. +pub fn inlineFits(cap: abi.QpCap) bool { + return cap.max_inline_data >= flag_bytes; +} + +/// One HCA as configured: its device name and RoCE v2 GID index (null: detected). +pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; + +/// What a rank tells its peer to connect (both ranks build it the same way; sent through the bootstrap). +pub const Record = extern struct { + magic: u32 = record_magic, + n_hca: u32 = 0, + region: u64 = 0, + /// The region's slot and registered length: a peer whose layout differs is refused at connect. + slot_bytes: u64 = 0, + region_len: u64 = 0, + rkey: [max_hcas]u32 = @splat(0), + lid: [max_hcas]u32 = @splat(0), + mtu: [max_hcas]u32 = @splat(0), + qpn: [max_hcas]u32 = @splat(0), + gid: [max_hcas][16]u8 = @splat(@splat(0)), +}; + +const Hca = struct { + ctx: *abi.Context, + pd: *abi.Pd, + mr: *abi.Mr, + cq: *abi.Cq, + qp: *abi.Qp, + gid_index: u8, + gid: [16]u8, + mtu: u32, + lid: u16, + outstanding: u32 = 0, + /// Each outstanding post's stripe bytes, oldest at `head` (completions arrive in order). + ring: [ring_len]u32 = undefined, + head: u32 = 0, + inflight: u64 = 0, + /// Posts made and their completions reaped (the proxy writes them, others read): equal = nothing in flight. + posts: std.atomic.Value(u64) = .init(0), + done: std.atomic.Value(u64) = .init(0), +}; + +/// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). +pub fn parseSpecs(text: []const u8, out: *[max_hcas]HcaSpec) !usize { + var n: usize = 0; + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + if (n == max_hcas) return error.TooManyHcas; + var parts = std.mem.splitScalar(u8, std.mem.trim(u8, item, " "), ':'); + const name = parts.next().?; + if (name.len == 0 or name.len >= 64) return error.BadHcaName; + out[n] = .{ .name_len = @intCast(name.len) }; + @memcpy(out[n].name[0..name.len], name); + if (parts.next()) |g| out[n].gid = try std.fmt.parseInt(u8, g, 10); + n += 1; + } + return n; +} + +/// RoCE v2 by sysfs (the kernel's own word for each GID); null when sysfs cannot say. +fn sysfsRoceV2(name: []const u8, index: u32) ?bool { + var path: [256]u8 = undefined; + const p = std.fmt.bufPrintSentinel(&path, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, index }, 0) catch return null; + const fd = std.c.open(p, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return null; + defer _ = std.c.close(fd); + var buf: [32]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return false; + return std.mem.startsWith(u8, buf[0..@intCast(n)], "RoCE v2"); +} + +/// An IPv4-mapped GID (::ffff:a.b.c.d): the address the HCA's netdev carries. +pub fn ipv4Mapped(g: [16]u8) bool { + return std.mem.allEqual(u8, g[0..10], 0) and g[10] == 0xff and g[11] == 0xff; +} + +/// The lowest RoCE v2 GID index that is IPv4-mapped (GID indices move on reboot: never hard-coded). +fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) ?u8 { + var i: u32 = 0; + while (i < @min(@as(u32, @intCast(@max(table, 0))), 255)) : (i += 1) { + const g = v.gid(ctx, 1, i) catch continue; + if (!ipv4Mapped(g.raw)) continue; + const v2 = sysfsRoceV2(name, i) orelse (if (v.gidType(ctx, 1, i)) |t| t == 2 else true); + if (v2) return @intCast(i); + } + return null; +} + +/// A device's port 1 as `roce_pick.choose` sees it: up, Ethernet, and its IPv4 RoCE v2 GID (index and address). +pub const Port = struct { spec: HcaSpec = .{}, active: bool = false, ethernet: bool = false, ipv4: ?[4]u8 = null }; + +/// Port 1 of every device that opens, in the device list's order, at most `out.len`. +pub fn scanPorts(v: *const Verbs, out: []Port) usize { + var count: c_int = 0; + const list = v.api.ibv_get_device_list(&count) orelse return 0; + defer v.api.ibv_free_device_list(list); + var n: usize = 0; + var i: usize = 0; + while (list[i]) |dev| : (i += 1) { + if (n == out.len) break; + const name = std.mem.span(v.api.ibv_get_device_name(dev) orelse continue); + const ctx = v.api.ibv_open_device(dev) orelse continue; + defer _ = v.api.ibv_close_device(ctx); + const port = v.port(ctx, 1) catch continue; + var p: Port = .{ .active = port.state == abi.port_active, .ethernet = port.link_layer == link_ethernet }; + p.spec.name_len = @intCast(@min(name.len, 63)); + @memcpy(p.spec.name[0..p.spec.name_len], name[0..p.spec.name_len]); + if (p.active and p.ethernet) if (detectGid(v, ctx, name, port.gid_tbl_len)) |gi| { + if (v.gid(ctx, 1, gi)) |g| { + p.spec.gid = gi; + p.ipv4 = g.raw[12..16].*; + } else |_| {} + }; + out[n] = p; + n += 1; + } + return n; +} + +pub const VerbsWire = struct { + v: Verbs, + hcas: [max_hcas]Hca = undefined, + n_hca: u32 = 0, + traffic_class: u8, + peer_region: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + /// Offsets in both ranks' regions (the same layout): receive slots and flags. + recv_off: usize, + flag_off: usize, + send_off: usize, + slot_bytes: usize, + region: []u8, + posted: u64 = 0, + /// The most stripe bytes in flight at once (posted, not completed), for tests and reports. + peak: std.atomic.Value(u64) = .init(0), + completed: u64 = 0, + /// Completions `quiesce` reaped, and those of them flushed (not sent: the QP was in error). + reaped: u64 = 0, + flushed: u64 = 0, + + /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. + pub fn open(specs: []const HcaSpec, region: []u8, layout: struct { send: usize, recv: usize, flags: usize, slot: usize }, traffic_class: u8) !VerbsWire { + if (specs.len == 0 or specs.len > max_hcas) return error.NoHca; + var w: VerbsWire = .{ .v = try Verbs.open(null), .traffic_class = traffic_class, .recv_off = layout.recv, .flag_off = layout.flags, .send_off = layout.send, .slot_bytes = layout.slot, .region = region }; + errdefer w.close(); + for (specs) |s| { + try w.openHca(s); + } + return w; + } + + fn openHca(w: *VerbsWire, s: HcaSpec) !void { + const api = w.v.api; + const name = s.name[0..s.name_len]; + const ctx = try w.v.openDevice(name); + errdefer _ = api.ibv_close_device(ctx); + const port = try w.v.port(ctx, 1); + if (port.state != abi.port_active) { + log.warn("tp roce: {s} port 1 is not active", .{name}); + return error.NoHca; + } + const gi = s.gid orelse detectGid(&w.v, ctx, name, port.gid_tbl_len) orelse { + log.warn("tp roce: {s} has no IPv4 RoCE v2 GID", .{name}); + return error.NoHca; + }; + const gid = try w.v.gid(ctx, 1, gi); + const pd = api.ibv_alloc_pd(ctx) orelse return error.VerbsFailed; + errdefer _ = api.ibv_dealloc_pd(pd); + const mr = api.ibv_reg_mr(pd, w.region.ptr, w.region.len, abi.access_local_write | abi.access_remote_write) orelse { + log.warn("tp roce: ibv_reg_mr on {s} failed (memlock limit? ulimit -l unlimited)", .{name}); + return error.VerbsFailed; + }; + errdefer _ = api.ibv_dereg_mr(mr); + const cq = api.ibv_create_cq(ctx, send_depth, null, null, 0) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_cq(cq); + var init: abi.QpInitAttr = .{ .send_cq = cq, .recv_cq = cq, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 }, .qp_type = abi.qpt_rc }; + const qp = api.ibv_create_qp(pd, &init) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_qp(qp); + if (!inlineFits(init.cap)) { + log.warn("tp roce: {s}'s QP sends {d} B inline, the flag post needs {d}; refusing to start", .{ name, init.cap.max_inline_data, flag_bytes }); + return error.VerbsFailed; + } + var a: abi.QpAttr = .{ .qp_state = abi.qps_init, .port_num = 1, .qp_access_flags = abi.access_remote_write }; + if (api.ibv_modify_qp(qp, &a, abi.mask.state | abi.mask.pkey_index | abi.mask.port | abi.mask.access_flags) != 0) return error.VerbsFailed; + w.hcas[w.n_hca] = .{ .ctx = ctx, .pd = pd, .mr = mr, .cq = cq, .qp = qp, .gid_index = gi, .gid = gid.raw, .mtu = @intCast(port.active_mtu), .lid = port.lid }; + w.n_hca += 1; + log.info("[tensorfold] tp roce: {s} GID {d} ({d}.{d}.{d}.{d}), MTU enum {d}", .{ name, gi, gid.raw[12], gid.raw[13], gid.raw[14], gid.raw[15], port.active_mtu }); + } + + pub fn close(w: *VerbsWire) void { + const api = w.v.api; + for (w.hcas[0..w.n_hca]) |h| { + _ = api.ibv_destroy_qp(h.qp); + _ = api.ibv_destroy_cq(h.cq); + _ = api.ibv_dereg_mr(h.mr); + _ = api.ibv_dealloc_pd(h.pd); + _ = api.ibv_close_device(h.ctx); + } + w.n_hca = 0; + w.v.close(); + } + + /// The region and each HCA's keys, this side's and the peer's (a protection error is checked against them). + pub fn describe(w: *const VerbsWire, buf: []u8) []const u8 { + var out: std.Io.Writer = .fixed(buf); + out.print("region 0x{x} len {d}", .{ @intFromPtr(w.region.ptr), w.region.len }) catch return out.buffered(); + for (w.hcas[0..w.n_hca], 0..) |h, i| { + out.print("; HCA {d} lkey 0x{x} rkey 0x{x}, peer region 0x{x} rkey 0x{x}", .{ i, h.mr.lkey, h.mr.rkey, w.peer_region, w.peer_rkey[i] }) catch break; + } + return out.buffered(); + } + + /// The QPs to error: the HCAs stop sending and refuse the peer's writes, and what was posted completes as flushed. + pub fn halt(w: *VerbsWire) bool { + var ok = true; + for (w.hcas[0..w.n_hca]) |h| { + var a: abi.QpAttr = .{ .qp_state = abi.qps_err }; + if (w.v.api.ibv_modify_qp(h.qp, &a, abi.mask.state) != 0) ok = false; + } + return ok; + } + + /// `halt`, then every posted request's completion reaped, flushed or not: no request of the region left at an HCA. + pub fn quiesce(w: *VerbsWire, timeout_ns: u64) bool { + if (!w.halt()) return false; + const t0 = sock.nowNs(); + for (w.hcas[0..w.n_hca]) |*h| while (h.outstanding > 0) { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return false; + for (wc[0..@intCast(n)]) |c| w.flushed += @intFromBool(c.status != abi.wc_success); + w.reaped += @intCast(n); + if (n > 0) w.popped(h, @intCast(n)) else if (sock.nowNs() - t0 > timeout_ns) return false; + }; + return true; + } + + /// Every HCA's posts completed: no request of the region at an HCA (readable beside the proxy). + pub fn idle(w: *const VerbsWire) bool { + for (w.hcas[0..w.n_hca]) |*h| if (h.posts.load(.acquire) != h.done.load(.acquire)) return false; + return true; + } + + /// Each QP's state as its device reports it (ibv_query_qp), then what `quiesce` reaped and what is still posted. + pub fn report(w: *const VerbsWire, buf: []u8) []const u8 { + const names = [_][]const u8{ "RESET", "INIT", "RTR", "RTS", "SQD", "SQE", "ERR" }; + var out: std.Io.Writer = .fixed(buf); + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + var a: abi.QpAttr = .{}; + var init: abi.QpInitAttr = undefined; + const q = w.v.query_qp orelse { + out.print("HCA {d} QP state unknown (no ibv_query_qp); ", .{i}) catch break; + continue; + }; + if (q(h.qp, &a, abi.mask.state, &init) != 0) { + out.print("HCA {d} QP query failed; ", .{i}) catch break; + } else if (a.qp_state >= 0 and a.qp_state < names.len) { + out.print("HCA {d} QP {s}; ", .{ i, names[@intCast(a.qp_state)] }) catch break; + } else out.print("HCA {d} QP state {d}; ", .{ i, a.qp_state }) catch break; + out.print("posted {d}, completed {d}; ", .{ h.posts.load(.acquire), h.done.load(.acquire) }) catch break; + } + var held: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| held += h.outstanding; + out.print("reaped {d} (flushed {d}), still posted {d}", .{ w.reaped, w.flushed, held }) catch {}; + return out.buffered(); + } + + pub fn record(w: *const VerbsWire) Record { + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr), .slot_bytes = w.slot_bytes, .region_len = w.region.len }; + for (w.hcas[0..w.n_hca], 0..) |h, i| { + r.rkey[i] = h.mr.rkey; + r.lid[i] = h.lid; + r.mtu[i] = h.mtu; + r.qpn[i] = h.qp.qp_num; + r.gid[i] = h.gid; + } + return r; + } + + /// The peer HCA for each of ours: the same IPv4 /24 (functions pair by subnet, names may differ), else position. + pub fn pairing(mine: Record, peer: Record) ![max_hcas]u32 { + if (peer.magic != record_magic) return error.BadRecord; + if (peer.n_hca != mine.n_hca or peer.slot_bytes != mine.slot_bytes or peer.region_len != mine.region_len) return error.BadRecord; + var out: [max_hcas]u32 = undefined; + for (0..mine.n_hca) |i| { + out[i] = @intCast(i); + var hits: u32 = 0; + for (0..peer.n_hca) |j| if (std.mem.eql(u8, mine.gid[i][12..15], peer.gid[j][12..15])) { + out[i] = @intCast(j); + hits += 1; + }; + if (hits > 1) out[i] = @intCast(i); // one subnet for both: keep the order + } + return out; + } + + /// INIT to RTR to RTS against the peer's record. + pub fn connect(w: *VerbsWire, peer: Record) !void { + const pairs = try pairing(w.record(), peer); + const api = w.v.api; + const m = abi.mask; + w.peer_region = peer.region; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + const j = pairs[i]; + w.peer_rkey[i] = peer.rkey[j]; + var rtr: abi.QpAttr = .{ .qp_state = abi.qps_rtr, .path_mtu = @intCast(@min(peer.mtu[j], h.mtu)), .dest_qp_num = peer.qpn[j], .rq_psn = 0, .max_dest_rd_atomic = 1, .min_rnr_timer = 12 }; + rtr.ah_attr = .{ .grh = .{ .dgid = .{ .raw = peer.gid[j] }, .sgid_index = h.gid_index, .hop_limit = 64, .traffic_class = w.traffic_class }, .dlid = @intCast(peer.lid[j]), .is_global = 1, .port_num = 1 }; + if (api.ibv_modify_qp(h.qp, &rtr, m.state | m.av | m.path_mtu | m.dest_qpn | m.rq_psn | m.max_dest_rd_atomic | m.min_rnr_timer) != 0) return error.VerbsFailed; + var rts: abi.QpAttr = .{ .qp_state = abi.qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .sq_psn = 0, .max_rd_atomic = 1 }; + if (api.ibv_modify_qp(h.qp, &rts, m.state | m.timeout | m.retry_cnt | m.rnr_retry | m.sq_psn | m.max_qp_rd_atomic) != 0) return error.VerbsFailed; + } + } + + /// Slot `slot`'s first `padded` bytes to the peer, striped over the HCAs, each stripe's `seq` flag on the same QP. + pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { + var plan: [max_hcas]guard.Stripe = undefined; + try guard.stripes(w.regionLayout(), w.n_hca, slot, padded, &plan); + while (w.inflight() + padded > max_inflight) try w.reapWait(); + var seq_copy = seq; + for (w.hcas[0..w.n_hca], plan[0..w.n_hca], 0..) |*h, s, i| { + while (h.outstanding >= ring_len) try w.reapHca(h); + var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + s.flag, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + s.local, .length = s.len, .lkey = h.mr.lkey }; + var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + s.remote, .rkey = w.peer_rkey[i] }; + var bad: ?*abi.SendWr = null; + const first: *abi.SendWr = if (s.len > 0) &data_wr else &flag_wr; + if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.ring[(h.head + h.outstanding) % ring_len] = s.len; + h.outstanding += 1; + h.inflight += s.len; + _ = h.posts.fetchAdd(1, .release); + } + if (w.inflight() > w.peak.load(.monotonic)) w.peak.store(w.inflight(), .monotonic); + w.posted += 1; + } + + /// This rank's layout; the MR covers `region` exactly and the peer's record matched it at connect. + pub fn regionLayout(w: *const VerbsWire) guard.Layout { + return .{ .send = w.send_off, .recv = w.recv_off, .flags = w.flag_off, .slot = w.slot_bytes, .len = w.region.len }; + } + + fn reapHca(w: *VerbsWire, h: *Hca) !void { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return error.Completion; + for (wc[0..@intCast(n)]) |c| if (c.status != abi.wc_success) { + log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); + return error.Completion; + }; + w.popped(h, @intCast(n)); + } + + fn popped(w: *VerbsWire, h: *Hca, n: u32) void { + for (0..n) |_| { + h.inflight -= h.ring[h.head]; + h.head = (h.head + 1) % ring_len; + } + h.outstanding -= n; + w.completed += n; + _ = h.done.fetchAdd(n, .release); + } + + /// Bytes posted and not yet completed over every HCA. + pub fn inflight(w: *const VerbsWire) u64 { + var sum: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| sum += h.inflight; + return sum; + } + + /// Reaps until some post completes (error.Completion when none does within `inflight_wait_ns`). + fn reapWait(w: *VerbsWire) !void { + const before = w.completed; + const t0 = sock.nowNs(); + while (w.completed == before) { + try w.reap(); + if (w.completed == before and sock.nowNs() - t0 > inflight_wait_ns) return error.Completion; + std.atomic.spinLoopHint(); + } + } + + /// Reap every HCA's completions (after each op, and in the proxy's idle loop). + pub fn reap(w: *VerbsWire) !void { + for (w.hcas[0..w.n_hca]) |*h| try w.reapHca(h); + } +}; + +test "HCA lists parse and subnets pair the functions" { + var specs: [max_hcas]HcaSpec = undefined; + const n = try parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqual(@as(usize, 2), n); + try std.testing.expectEqualStrings("rocep1s0f1", specs[0].name[0..specs[0].name_len]); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); + try std.testing.expectEqual(@as(?u8, null), specs[1].gid); + var a: Record = .{ .n_hca = 2 }; + var b: Record = .{ .n_hca = 2 }; + const sub100 = [16]u8{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 192, 168, 100, 10 }; + var sub101 = sub100; + sub101[14] = 101; + a.gid = .{ sub100, sub101 }; + b.gid = .{ sub101, sub100 }; // the peer lists its functions the other way round + b.gid[0][15] = 11; + b.gid[1][15] = 11; + const p = try VerbsWire.pairing(a, b); + try std.testing.expectEqual(@as(u32, 1), p[0]); + try std.testing.expectEqual(@as(u32, 0), p[1]); + try std.testing.expect(ipv4Mapped(sub100)); + b.n_hca = 1; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.n_hca = 2; + b.slot_bytes = 1 << 20; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.slot_bytes = 0; + b.region_len = 4096; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh new file mode 100644 index 0000000000..75f0c7900b --- /dev/null +++ b/zig/src/fabric/tp/run_hosts.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# One rank of the two-host TP=2 gate over RoCE: `run_hosts.sh RANK RANK0_ADDR ./tf-tp-test OUTDIR` on both hosts. +set -u +RANK=${1:?rank}; MASTER=${2:?master address}; BIN=${3:?tf-tp-test}; OUT=${4:?outdir} +mkdir -p "$OUT" +[ -n "${NCCL_SOCKET_IFNAME:-}" ] && export NCCL_SOCKET_IFNAME # else NCCL picks the interface +# NCCL uses the RoCE wire's devices unless NCCL_IB_HCA is set: the names without `:gid` (NCCL reads `:N` as a port) +[ -n "${TF_TP_ROCE_HCA:-}" ] && [ -z "${NCCL_IB_HCA:-}" ] && NCCL_IB_HCA=$(echo "$TF_TP_ROCE_HCA" | sed 's/:[0-9]*//g; s/ //g') +[ -n "${TF_TP_ROCE_HCA:-}" ] && export TF_TP_ROCE_HCA +"$BIN" kernel > /dev/null 2>&1 || { echo "error: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } +ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 +# neither named: NCCL takes the ports the RoCE picker chose with the peer (`tf-tp-test hca`, both ranks), not its own +if [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then + # hca gives up after TF_TP_PLAN_TIMEOUT (default 300 s) with its reason; the cap here only catches a hang past that + CAP=$(awk 'BEGIN { t = ARGV[1] + 0; if (t <= 0) t = 300; printf "%d", int(t) + 31 }' "${TF_TP_PLAN_TIMEOUT:-300}") + echo "hca: rank $RANK picking the ports with the peer (up to $CAP s)" + PICK=$(TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=29700 timeout "$CAP" "$BIN" hca 2> "$OUT/hca.log") + rc=$? + if [ "$rc" -eq 124 ]; then + echo "error: $BIN hca still running after $CAP s: killed ($OUT/hca.log)" >&2 + exit 3 + fi + if [ "$rc" -ne 0 ]; then + echo "error: $BIN hca exited $rc: $(grep -h '^FAIL\|^error' "$OUT/hca.log" | tail -n 1 | grep . || tail -n 1 "$OUT/hca.log") ($OUT/hca.log)" >&2 + exit 3 + fi + if ! echo "$PICK" | grep -Eqx 'NCCL_IB_HCA=[A-Za-z0-9_.-]+(,[A-Za-z0-9_.-]+)*'; then + echo "error: $BIN hca printed '$PICK', not NCCL_IB_HCA=name,... ($OUT/hca.log)" >&2 + exit 3 + fi + NCCL_IB_HCA=${PICK#NCCL_IB_HCA=} + echo "hca: rank $RANK NCCL_IB_HCA=$NCCL_IB_HCA ($(grep -h '^INFO hca' "$OUT/hca.log"))" +fi +[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_HCA +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; env | grep -E '^(TF_TP_ROCE_HCA|NCCL_IB_HCA|NCCL_SOCKET_IFNAME)='; } > "$OUT/host.txt" 2>&1 +PORT=29700 +step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line + local name=$1 backend=$2; shift 2 + PORT=$((PORT + 10)) + ( TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend TF_TP_ROCE_FALLBACK=error \ + timeout 900 "$BIN" "$@"; echo "EXIT rank $RANK rc $? at $(date +%s.%N)" ) > "$OUT/$name.log" 2>&1 + echo "$name: $(grep -h '^SUMMARY\|^PASS roce-host\|^INJECT\|^EXIT' "$OUT/$name.log" | tr '\n' ' ')" +} +step roce-host roce roce-host +step suite-roce roce suite +step suite-nccl nccl suite +for backend in roce nccl; do + for mode in fatal kill; do + for victim in 0 1; do step "ff-$backend-$mode-$victim" "$backend" failfast "$victim" "$mode"; done + done +done diff --git a/zig/src/fabric/tp/run_pair.sh b/zig/src/fabric/tp/run_pair.sh new file mode 100644 index 0000000000..1cf73c6261 --- /dev/null +++ b/zig/src/fabric/tp/run_pair.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# Two ranks of tf-tp-test on one two-GPU host: `[SUITES=..] [FAILFAST=..] run_pair.sh BIN OUTDIR`. +set -u +BIN=${1:?bin}; OUT=${2:?outdir} +mkdir -p "$OUT" +nvidia-smi > "$OUT/smi.txt" 2>&1 +nvidia-smi topo -m > "$OUT/topo.txt" 2>&1 +ldconfig -p | grep -E "libnccl|libcuda" > "$OUT/libs.txt" 2>&1 +PORT=29500 +pair() { # NAME BACKEND ARGS...: both ranks, logs OUT/NAME.rR.log ending with an EXIT line + local name=$1 backend=$2 r; shift 2 + PORT=$((PORT + 10)) + for r in 0 1; do + ( TF_TP_RANK=$r TF_TP_DEVICE=$r TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend NCCL_DEBUG=${NCCL_DEBUG:-WARN} \ + timeout 600 "$BIN" "$@"; echo "EXIT rank $r rc $? at $(date +%s.%N)" ) > "$OUT/$name.r$r.log" 2>&1 & + done + wait +} +# ffsum OUT NAME VICTIM: a fail-fast pair's verdict line (exit codes, exit times after the injection) +ffsum() { + [ -f "$1/$2.r0.log" ] && [ -f "$1/$2.r1.log" ] || { echo "ffsum: $1/$2.r0.log or .r1.log missing" >&2; return 1; } + awk -v name="$2" -v victim="$3" ' + FNR == 1 { r = (FILENAME ~ /\.r1\.log$/) ? 1 : 0 } + !(r in ex) && match($0, /EXIT rank [0-9] rc [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), f, " "); rc[r] = f[5]; at[r] = f[7]; ex[r] = 1 + } + r == victim && !inj && match($0, /INJECT [A-Za-z0-9_]+ rank [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), g, " "); t0 = g[6]; inj = 1 + } + END { + if (!inj || !(0 in ex) || !(1 in ex)) { print name ": FAIL (no inject / exit line)"; exit } + for (i = 0; i < 2; i++) { want[i] = (i == victim && name ~ /kill/) ? 137 : 70; dt[i] = at[i] - t0 } + ok = rc[0] == want[0] && rc[1] == want[1] && dt[1 - victim] < 3.0 + printf "%s: %s rc [%d, %d] (want [%d, %d]), exit after injection: rank0 %.3f s, rank1 %.3f s\n", name, + ok ? "PASS" : "FAIL", rc[0], rc[1], want[0], want[1], dt[0], dt[1] + }' "$1/$2.r0.log" "$1/$2.r1.log" +} +for backend in ${SUITES:-mailbox nccl}; do + pair "suite-$backend" "$backend" suite +done +: > "$OUT/failfast.txt" +for backend in ${FAILFAST:-nccl mailbox}; do + for mode in fatal kill; do + for victim in 0 1; do + name=ff-$backend-$mode-$victim + pair "$name" "$backend" failfast "$victim" "$mode" + ffsum "$OUT" "$name" "$victim" >> "$OUT/failfast.txt" + done + done +done +grep -h "^SUMMARY\|^FAIL" "$OUT"/suite*.log +cat "$OUT/failfast.txt" diff --git a/zig/src/fabric/tp/run_transport.sh b/zig/src/fabric/tp/run_transport.sh new file mode 100755 index 0000000000..68ebcddd9a --- /dev/null +++ b/zig/src/fabric/tp/run_transport.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# The failed-transport check on two hosts (tf-tp-test failfast V transport); `run_transport.sh` alone prints how. +set -u +usage() { + cat << 'TXT' +run_transport.sh MASTER_ADDR PEER_SSH OUTDIR (on rank 0's host, as root; rank 1 runs over ssh) + Two hosts with this script and tf-tp-test (built with the mailbox kernel) at the same path on both. Before the steps, + on both hosts: no GPU process and memlock unlimited. Steps: T0 (suite), T1 (victim 1), T3 (victim 0), fail-fast on; + tf-tp-test injects only when both links are idle. After each step, on both hosts: the kernel log since the step + began and every HCA *err* counter against the step's start. A matching line, a counter that grew, an unexpected rc + or a failure line prints STOP: and ends the run (exit 1), with no later step. +run_transport.sh t2 OUTDIR (one host: TF_TP_ALLOW_HOST_ONLY=1) + T2 (fail-fast off, a hold) runs both ranks on this host over the shm wire; tf-tp-test refuses it on an RDMA wire. +TXT +} +[ $# -ge 2 ] || { usage; exit 2; } +KIT=$(cd "$(dirname "$0")" && pwd) +BIN=$KIT/tf-tp-test +KLOG='EXT4-fs|smmu|iommu|mlx5|AER|nvme|Xid' +BAD='RDMA write of seq|local protection|NOT drained|GPU sync FAILED|REFUSED|refused ERR|panic' +"$BIN" kernel > /dev/null 2>&1 || { echo "STOP: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } + +if [ "${1:-}" = t2 ]; then + OUT=${2:?outdir} + [ "${TF_TP_ALLOW_HOST_ONLY:-}" = 1 ] || { echo "STOP: T2 is host-only: set TF_TP_ALLOW_HOST_ONLY=1" >&2; exit 2; } + mkdir -p "$OUT" + for r in 0 1; do + TF_TP_RANK=$r TF_TP_MASTER=127.0.0.1 TF_TP_PORT=29990 TF_COMM_BACKEND=roce TF_TP_ROCE_WIRE=shm TF_TP_FAILFAST=0 \ + timeout 300 "$BIN" failfast 1 transport 200 5 > "$OUT/T2-r$r.log" 2>&1 & + done + wait + grep -h '^PASS\|^FAIL\|REFUSED\|regions kept' "$OUT"/T2-r*.log + exit 0 +fi + +MASTER=${1:?master address}; PEER=${2:?peer ssh target}; OUT=${3:?outdir} +mkdir -p "$OUT" +stop() { echo "STOP: $*" | tee -a "$OUT/verdict.txt"; exit 1; } + +precheck() { + [ -z "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2> /dev/null)" ] || { echo "a GPU process runs"; return 1; } + [ "$(bash -c 'ulimit -l unlimited 2> /dev/null; ulimit -l')" = unlimited ] || { echo "memlock not unlimited"; return 1; } + echo "ok $(hostname) $(uname -r)" +} +counters() { + find -L /sys/class/infiniband/ -maxdepth 5 -path '*ports*counters*' -name '*err*' 2> /dev/null | sort | while read -r f; do + echo "$f $(cat "$f" 2> /dev/null)" + done +} +klog() { # kernel lines since epoch $1 matching $2 (journalctl must answer: an unreadable log is a STOP) + journalctl -k --since "@$1" --no-pager -o short-iso 2> /dev/null | grep -E "$2" + return "${PIPESTATUS[0]}" +} +probe() { # this host's tf-tp-test QPs and MRs once a second until it exits + local pid="" t + for _ in $(seq 200); do pid=$(pgrep -nx tf-tp-test) && break; sleep 0.1; done + [ -z "$pid" ] && return + for t in $(seq 180); do + kill -0 "$pid" 2> /dev/null || break + echo "== t=$t pid $pid" + rdma resource show qp pid "$pid" 2>&1 + rdma resource show mr pid "$pid" 2>&1 + sleep 1 + done +} +on() { # NODE CMD ARGS...: a command or one of the functions above, here (local) or on the peer + local node=$1 fn=$2 def + shift 2 + if [ "$node" = local ]; then "$fn" "$@"; return; fi + def=$(declare -f "$fn" 2> /dev/null) + ssh -o BatchMode=yes "$PEER" "${def:+$def +}$fn $(printf '%q ' "$@")" +} +grew() { # BEFORE AFTER: counters that went up + awk 'NR == FNR { b[$1] = $2; next } ($1 in b) && $2 > b[$1] { print $1 " " b[$1] " -> " $2 }' "$1" "$2" +} + +for n in local peer; do + v=$(on $n precheck) || stop "$n: $v" + echo "precheck $n: $v" | tee -a "$OUT/verdict.txt" +done +on peer test -x "$BIN" || stop "peer: no $BIN (the kit must sit at the same path on both hosts)" + +PORT=29900 +step() { # NAME WANT_RC -- ARGS... + local name=$1 want=$2 n rc rc0 rc1 hits g + shift 3 + PORT=$((PORT + 10)) + declare -A t + for n in local peer; do + t[$n]=$(on $n date +%s) + on $n counters > "$OUT/$name.$n.err0" + done + on local probe > "$OUT/$name.local.rdma.txt" 2>&1 & + on peer probe > "$OUT/$name.peer.rdma.txt" 2>&1 & + local envs=(TF_TP_MASTER="$MASTER" TF_TP_PORT=$PORT TF_COMM_BACKEND=roce TF_TP_ROCE_FALLBACK=error TF_TP_MAILBOX_MAX_KB=1024 + TF_TP_ROCE_FAST=1 TF_TP_FAILFAST_GRACE_S=15) + ssh -o BatchMode=yes "$PEER" "ulimit -l unlimited && env $(printf '%q ' "${envs[@]}") TF_TP_RANK=1 timeout 300 $(printf '%q ' "$BIN" "$@"); echo \"EXIT rank 1 rc \$? at \$(date +%s.%N)\"" > "$OUT/$name-r1.log" 2>&1 & + ( ulimit -l unlimited && env "${envs[@]}" TF_TP_RANK=0 timeout 300 "$BIN" "$@"; echo "EXIT rank 0 rc $? at $(date +%s.%N)" ) > "$OUT/$name-r0.log" 2>&1 + wait + rc0=$(grep -oP '^EXIT rank 0 rc \K[0-9]+' "$OUT/$name-r0.log") + rc1=$(grep -oP '^EXIT rank 1 rc \K[0-9]+' "$OUT/$name-r1.log") + echo "$name: rank 0 rc ${rc0:-?}, rank 1 rc ${rc1:-?} (want $want)" | tee -a "$OUT/verdict.txt" + for n in local peer; do + on $n counters > "$OUT/$name.$n.err1" + g=$(grew "$OUT/$name.$n.err0" "$OUT/$name.$n.err1") + [ -n "$g" ] && stop "$name $n: HCA error counters grew: $g" + hits=$(on $n klog "${t[$n]}" "$KLOG"); rc=$? + [ "$rc" -ne 0 ] && stop "$name $n: the kernel log is unreadable (journalctl rc $rc)" + [ -n "$hits" ] && { echo "$hits" > "$OUT/$name.$n.klog.txt"; stop "$name $n: kernel lines: $hits"; } + done + hits=$(grep -hE "$BAD" "$OUT/$name-r0.log" "$OUT/$name-r1.log") + [ -n "$hits" ] && stop "$name: $hits" + [ "$rc0" = "$want" ] && [ "$rc1" = "$want" ] || stop "$name: rc $rc0 / $rc1, want $want" + grep -h 'halted\|quiesced\|idle before\|regions kept' "$OUT/$name"-r*.log | sed "s/^/ /" +} +step T0 0 -- suite +step T1 70 -- failfast 1 transport 200 +step T3 70 -- failfast 0 transport 200 +echo "PASS transport check (T0, T1, T3)" | tee -a "$OUT/verdict.txt" diff --git a/zig/src/fabric/tp/sock.zig b/zig/src/fabric/tp/sock.zig new file mode 100644 index 0000000000..b382884034 --- /dev/null +++ b/zig/src/fabric/tp/sock.zig @@ -0,0 +1,169 @@ +//! TCP for the bootstrap and the fate channel: listen, accept and connect with deadlines, whole sends and receives. +const std = @import("std"); +const c = std.c; +const posix = std.posix; + +pub const Fd = c.fd_t; + +pub const Error = error{ SocketFailed, AddressUnresolved, AddressInUse, Timeout, Closed, Interrupted }; + +pub fn nowNs() u64 { + var ts: c.timespec = undefined; + _ = c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +pub fn sleepNs(ns: u64) void { + var ts: c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + while (c.nanosleep(&ts, &ts) != 0) {} +} + +/// TCP_NODELAY (frames are tiny and latency matters) and keepalive (a vanished host fails a read in ~90 s, not never). +fn tune(fd: Fd) void { + const one: c_int = 1; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &one, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, &one, @sizeOf(c_int)); + const idle: c_int = 30; + const intvl: c_int = 10; + const cnt: c_int = 6; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPIDLE, &idle, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPINTVL, &intvl, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPCNT, &cnt, @sizeOf(c_int)); +} + +/// Rank 0's listening socket on every IPv4 address at `port`. +pub fn listen(port: u16) Error!Fd { + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + errdefer _ = c.close(fd); + const one: c_int = 1; + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, &one, @sizeOf(c_int)); + var addr: posix.sockaddr.in = .{ .port = std.mem.nativeToBig(u16, port), .addr = 0 }; + if (c.bind(fd, @ptrCast(&addr), @sizeOf(posix.sockaddr.in)) != 0) return error.AddressInUse; + if (c.listen(fd, 64) != 0) return error.SocketFailed; + return fd; +} + +/// One connection before `deadline_ns` (monotonic). +pub fn accept(fd: Fd, deadline_ns: u64) Error!Fd { + while (true) { + try waitReadable(fd, deadline_ns); + const conn = c.accept(fd, null, null); + if (conn >= 0) { + tune(conn); + return conn; + } + } +} + +/// Connects to `host:port`, retrying every 50 ms (the master may not listen yet) until `deadline_ns`. +pub fn connect(host: []const u8, port: u16, deadline_ns: u64) Error!Fd { + var name: [256]u8 = undefined; + if (host.len >= name.len) return error.AddressUnresolved; + @memcpy(name[0..host.len], host); + name[host.len] = 0; + var service: [8]u8 = undefined; + const svc = std.fmt.bufPrintSentinel(&service, "{d}", .{port}, 0) catch unreachable; + const hints: c.addrinfo = .{ .flags = .{}, .family = posix.AF.INET, .socktype = posix.SOCK.STREAM, .protocol = 0, .addrlen = 0, .addr = null, .canonname = null, .next = null }; + while (true) { + var res: ?*c.addrinfo = null; + if (@backingInt(c.getaddrinfo(name[0..host.len :0], svc, &hints, &res)) == 0) { + defer c.freeaddrinfo(res.?); + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + if (c.connect(fd, res.?.addr.?, res.?.addrlen) == 0) { + tune(fd); + return fd; + } + _ = c.close(fd); + } + if (nowNs() >= deadline_ns) return error.Timeout; + sleepNs(50 * std.time.ns_per_ms); + } +} + +fn waitReadable(fd: Fd, deadline_ns: u64) Error!void { + while (true) { + const t = nowNs(); + if (t >= deadline_ns) return error.Timeout; + const ms: c_int = @intCast(@min((deadline_ns - t) / std.time.ns_per_ms + 1, 1000)); + var p = [_]c.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const n = c.poll(&p, 1, ms); + if (n > 0) return; + } +} + +/// Every byte of `bytes`, or `error.Closed`. +pub fn sendAll(fd: Fd, bytes: []const u8) Error!void { + var off: usize = 0; + while (off < bytes.len) { + const n = c.send(fd, bytes[off..].ptr, bytes.len - off, posix.MSG.NOSIGNAL); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// Fills `out` before `deadline_ns` (null: no deadline), or `error.Closed` at the end of the stream. +pub fn recvExact(fd: Fd, out: []u8, deadline_ns: ?u64) Error!void { + var off: usize = 0; + while (off < out.len) { + if (deadline_ns) |d| try waitReadable(fd, d); + const n = c.recv(fd, out[off..].ptr, out.len - off, 0); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// A non-blocking send of the whole of `bytes` if the socket buffer has room (false: dropped, or the peer is gone). +pub fn trySend(fd: Fd, bytes: []const u8) bool { + const n = c.send(fd, bytes.ptr, bytes.len, posix.MSG.NOSIGNAL | posix.MSG.DONTWAIT); + return n == @as(isize, @intCast(bytes.len)); +} + +pub fn hangup(fd: Fd) void { + _ = c.shutdown(fd, posix.SHUT.RDWR); +} + +pub fn close(fd: Fd) void { + _ = c.close(fd); +} + +/// This host's identity (the kernel's boot id, else the host name): ranks with equal ids share memory. +pub fn hostId(out: *[64]u8) []const u8 { + @memset(out, 0); + const f = c.open("/proc/sys/kernel/random/boot_id", .{ .ACCMODE = .RDONLY }); + if (f >= 0) { + defer _ = c.close(f); + const n = c.read(f, out, out.len); + if (n > 0) return std.mem.trim(u8, out[0..@intCast(n)], " \n"); + } + if (c.gethostname(out, out.len) != 0) return ""; + return std.mem.sliceTo(out, 0); +} + +test "a listener and a client trade bytes and see the hang-up" { + const l = try listen(0); + defer close(l); + var addr: posix.sockaddr.in = undefined; + var len: posix.socklen_t = @sizeOf(posix.sockaddr.in); + try std.testing.expectEqual(@as(c_int, 0), c.getsockname(l, @ptrCast(&addr), &len)); + const port = std.mem.bigToNative(u16, addr.port); + const deadline = nowNs() + 5 * std.time.ns_per_s; + const a = try connect("localhost", port, deadline); + defer close(a); + const b = try accept(l, deadline); + defer close(b); + try sendAll(a, "hello"); + var got: [5]u8 = undefined; + try recvExact(b, &got, deadline); + try std.testing.expectEqualStrings("hello", &got); + try std.testing.expectError(error.Timeout, recvExact(b, &got, nowNs() + 20 * std.time.ns_per_ms)); + hangup(a); + try std.testing.expectError(error.Closed, recvExact(b, &got, deadline)); +} diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig new file mode 100644 index 0000000000..6cc5ceb6ec --- /dev/null +++ b/zig/src/fabric/tp/tp.zig @@ -0,0 +1,217 @@ +//! Tensor parallelism across processes, one GPU each: bootstrap, collectives, fail-fast; `Session.open` starts a rank. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); + +pub const collective = @import("collective.zig"); +pub const reference = @import("reference.zig"); +pub const config = @import("config.zig"); +pub const sock = @import("sock.zig"); +pub const bootstrap = @import("bootstrap.zig"); +pub const fate = @import("fate.zig"); +pub const host = @import("host.zig"); +pub const nccl = @import("nccl.zig"); +pub const mailbox = @import("mailbox.zig"); +pub const hybrid = @import("hybrid.zig"); +pub const roce = @import("roce.zig"); +pub const roce_verbs = @import("roce_verbs.zig"); +pub const roce_pick = @import("roce_pick.zig"); +pub const planlink = @import("planlink.zig"); +pub const Verbs = @import("verbs").Verbs; + +pub const Collective = collective.Collective; +pub const DType = collective.DType; +pub const Op = collective.Op; +pub const Config = config.Config; +pub const Bootstrap = bootstrap.Bootstrap; +pub const Fate = fate.Fate; + +pub const Session = struct { + cfg: Config, + lib: cuda.nccl.Library, + nccl: nccl.Nccl, + mail: ?mailbox.Mailbox = null, + roce: ?*roce.Roce = null, + route: hybrid.Hybrid = undefined, + /// The round plan's link (rank 0 sends, the others receive), on its own TCP connection. + plan: planlink.PlanLink = undefined, + fate: ?Fate = null, + + /// Brings this rank up in place; `kernel_image` is mailbox.cu's fatbin, empty sends those backends to NCCL. + pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { + var b = try Bootstrap.open(cfg); + defer b.close(); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_ROCE_CHUNKED_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); + errdefer lib.close(); + const v = try lib.version(); + try b.agree(std.mem.asBytes(&v), "NCCL version"); + var id: cuda.nccl.UniqueId = undefined; + if (cfg.rank == 0) try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try b.broadcast(&id.internal, &id.internal); + s.* = .{ .cfg = cfg, .lib = lib, .nccl = undefined }; + s.nccl = try nccl.Nccl.init(&s.lib, d, id, cfg.rank, cfg.world); + errdefer s.nccl.deinit(); + switch (cfg.backend) { + .nccl => {}, + .mailbox => if (try mailboxPossible(&b, kernel_image)) { + s.mail = try mailbox.Mailbox.init(d, &b, kernel_image, cfg.mailbox_max_bytes, cfg.mailbox_timeout_ns); + s.route = .{ .small = .{ .mailbox = &s.mail.? }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + } else if (cfg.rank == 0) { + log.warn("[tensorfold] tp: TF_COMM_BACKEND=mailbox needs two ranks on one host and the kernel; every rank uses NCCL", .{}); + }, + .roce => try s.openRoce(d, &b, kernel_image), + } + try s.agreeVar(&b); + s.plan = planlink.PlanLink.init(cfg.rank, cfg.world, b.releasePlan()); + s.plan.spin_ns = cfg.plan_spin_ns; + s.plan.pin = cfg.plan_pin; + s.plan.avoid = cfg.roce_cpu; + if (cfg.failfast) { + s.fate = Fate.init(cfg, b.release()); + s.fate.?.onFail(.{ .ctx = s, .call = abortHook }); + try s.fate.?.start(); + } + log.info("[tensorfold] tp: rank {d} of {d} up ({t}, NCCL {d}, fail-fast {s}, roce fast {s}: {d} poller(s), stagger {d} ns)", .{ cfg.rank, cfg.world, s.comm().kind(), v, if (cfg.failfast) "on" else "off", if (cfg.roce_fast & 1 != 0) "on" else "off", (cfg.roce_fast >> 8) & 0xff, cfg.roce_fast >> 16 }); + } + + /// RoCE for the small exchanges, probed against NCCL at load; on failure every rank obeys TF_TP_ROCE_FALLBACK. + fn openRoce(s: *Session, d: *const cuda.Driver, b: *Bootstrap, image: []const u8) !void { + const cfg = s.cfg; + var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; + var n_specs: usize = 0; + if (cfg.roce_wire == .verbs) n_specs = try roce_pick.resolve(cfg, b, &specs); + const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; + const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); + s.roce = r; + if (cfg.roce_chunked_bytes > 0 and !r.chunk_abi) { + log.warn("TF_TP_ROCE_CHUNKED_KB={d}: the RoCE kernel image has no chunked ABI (tp_roce_chunk_abi); refusing to start", .{cfg.roce_chunked_bytes / 1024}); + r.deinit(); + s.roce = null; + return error.ChunkAbiMissing; + } + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes, .chunked_bytes = cfg.roce_chunked_bytes }; + const good = probe(s, d) catch false; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(good)}, &buf, &blobs); + for (all) |x| if (x[0] != 1) { + r.deinit(); + s.roce = null; + return s.roceFallback(error.ProbeMismatch); + }; + } + + /// `Collective.varAgreed`: whether every rank's transport moves device lengths, agreed once a session. + fn agreeVar(s: *Session, b: *Bootstrap) !void { + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(s.comm().varFits(16))}, &buf, &blobs); + var ok = true; + for (all) |x| ok = ok and x[0] == 1; + if (s.mail != null or s.roce != null) s.route.var_agreed = ok; + } + + fn roceFallback(s: *Session, e: anyerror) !void { + if (s.cfg.roce_fallback == .@"error") return e; + if (s.cfg.rank == 0) log.warn("[tensorfold] tp: RoCE unavailable ({t}); every rank uses NCCL (TF_TP_ROCE_FALLBACK=error refuses instead)", .{e}); + } + + /// One 4 KiB all-gather and one bf16 sum through RoCE and through NCCL: the bytes must be equal. + fn probe(s: *Session, d: *const cuda.Driver) !bool { + const n = 4096; + // input [0, n), gathers [n, 3n) and [3n, 5n), sums [5n, 6n) and [6n, 7n) + var a = try cuda.DeviceBuffer.alloc(d, 7 * n); + defer a.free(); + var stream = try cuda.Stream.init(d, true); + defer stream.deinit(); + var bytes: [n]u8 = undefined; + reference.fill(&bytes, .bf16, s.cfg.rank, 0x7072_6f62); + try a.uploadAsync(0, &bytes, stream.handle); + const fast = s.route.iface(); + const slow = s.nccl.iface(); + try fast.allGather(a.ptr, a.ptr + n, n, .u8, stream.handle); + try slow.allGather(a.ptr, a.ptr + 3 * n, n, .u8, stream.handle); + try fast.allReduce(a.ptr, a.ptr + 5 * n, n / 2, .bf16, .sum, stream.handle); + try slow.allReduce(a.ptr, a.ptr + 6 * n, n / 2, .bf16, .sum, stream.handle); + try stream.synchronize(); + try fast.check(); + var got: [6 * n]u8 = undefined; + try a.download(n, &got); + const gathers = std.mem.eql(u8, got[0 .. 2 * n], got[2 * n .. 4 * n]); + const sums = std.mem.eql(u8, got[4 * n .. 5 * n], got[5 * n .. 6 * n]); + if (!gathers or !sums) log.warn("[tensorfold] tp roce probe: gather {s}, sum {s} against NCCL", .{ if (gathers) "equal" else "DIFFERS", if (sums) "equal" else "DIFFERS" }); + return gathers and sums; + } + + fn mailboxPossible(b: *Bootstrap, image: []const u8) !bool { + if (b.world != 2) return false; + var id_buf: [64]u8 = undefined; + const id = sock.hostId(&id_buf); + var mine: [65]u8 = @splat(0); + mine[0] = @intFromBool(image.len > 0); + @memcpy(mine[1..][0..id.len], id); + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [2 * 65]u8 = undefined; + const all = try b.allGather(&mine, &buf, &blobs); + return all[0][0] == 1 and std.mem.eql(u8, all[0], all[1]); + } + + fn abortHook(ctx: *anyopaque, _: []const u8) void { + const s: *Session = @ptrCast(@alignCast(ctx)); + s.comm().abort(); + } + + /// The communicator the model uses: NCCL, or the hybrid route over the mailbox / RoCE. + pub fn comm(s: *Session) Collective { + if (s.mail != null or s.roce != null) return s.route.iface(); + return s.nccl.iface(); + } + + /// This rank failed inside a round: every rank goes down (exit 70). Without fail-fast it only aborts. + pub fn fatal(s: *Session, where: []const u8, err: anyerror) void { + if (s.fate) |*f| return f.fatal(where, err); + s.comm().abort(); + } + + /// A clean stop announced by rank 0 before it closes: followers end with no fail-fast exit. + pub fn announceStop(s: *Session) void { + if (s.fate) |*f| f.announceStop(); + } + + /// This rank failed after `open`, before serving: the other ranks learn why and go down (fail-fast). + pub fn bootFailed(s: *Session, err: anyerror) void { + if (s.fate) |*f| f.failBoot("boot", err); + s.close(); + } + + /// A normal stop: every rank calls it (rank 0 first marks the stop as expected with `fate.expectStop`). + pub fn close(s: *Session) void { + // a failed transport: every rank goes down (the peers stop posting) before this one lets go of anything + if (s.roce) |r| if (r.failed()) if (s.fate) |*f| f.fatal("transport", error.BackendFailed); + if (s.fate) |*f| f.close(); + s.plan.close(); + if (s.roce) |r| r.deinit(); + if (s.mail) |*m| m.deinit(); + s.nccl.deinit(); + s.lib.close(); + } +}; + +test { + _ = collective; + _ = reference; + _ = config; + _ = sock; + _ = bootstrap; + _ = mailbox; + _ = roce_verbs; + _ = roce_pick; + _ = planlink; + _ = @import("roce_guard.zig"); + _ = @import("roce_test.zig"); + _ = @import("roce_post_test.zig"); + _ = @import("roce_inflight_test.zig"); + _ = @import("host_test.zig"); + _ = @import("fate_test.zig"); +} diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 1da50d386d..8ae58532ef 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,7 +1,7 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); const builtin = @import("builtin"); -const abi = @import("verbs_abi.zig"); +pub const abi = @import("verbs_abi.zig"); pub const Error = error{ LibraryUnavailable, MissingSymbol, NoDeviceList, NoSuchDevice, VerbsFailed, BadDeviceList, NoUsableGid }; pub const max_links = 8; @@ -50,12 +50,15 @@ pub const Api = struct { /// Exported by rdma-core 34+ but not in every header; used only when present. const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, flags: u32, size: usize) callconv(.c) c_int; const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; +/// ibv_query_qp: the QP's state as the device holds it (not every library exports it; used only when present). +pub const QueryQp = *const fn (qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int, init: *abi.QpInitAttr) callconv(.c) c_int; pub const Verbs = struct { lib: std.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, + query_qp: ?QueryQp = null, pub fn open(path: ?[:0]const u8) Error!Verbs { var lib = std.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; @@ -68,6 +71,7 @@ pub const Verbs = struct { var v: Verbs = .{ .lib = lib, .api = api }; v.gid_ex = v.lib.lookup(GidEx, "_ibv_query_gid_ex"); v.in_order = v.lib.lookup(InOrder, "ibv_query_qp_data_in_order"); + v.query_qp = v.lib.lookup(QueryQp, "ibv_query_qp"); return v; } diff --git a/zig/src/sessions/disk.zig b/zig/src/sessions/disk.zig new file mode 100644 index 0000000000..f5ab58f06f --- /dev/null +++ b/zig/src/sessions/disk.zig @@ -0,0 +1,587 @@ +//! The NVMe tier: a file an entry under //rank/, written via .tmp + rename, streamed via the lanes. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const lanes_mod = @import("lanes.zig"); +const Lanes = lanes_mod.Lanes; +const Staging = lanes_mod.Staging; +const Op = lanes_mod.Op; +const up = lanes_mod.up; +const format = @import("format.zig"); +const Layout = format.Layout; +const Head = format.Head; +const Seg = format.Seg; +const PageStore = @import("pagestore.zig").PageStore; +const pool_mod = @import("pool.zig"); +const Family = pool_mod.Family; +const Digest = @import("prefix.zig").Digest; + +pub const none = std.math.maxInt(u32); + +pub const Options = struct { + root: []const u8, + compat: Digest, + rank: u32 = 0, + world: u32 = 1, + /// bytes of files the tier keeps (least recently used files without dependents go first) + budget: u64 = 64 << 30, + /// entries shorter than this are dropped, not parked + min_tokens: u64 = 1024, + /// O_DIRECT where the filesystem takes it (buffered with the page cache dropped otherwise) + direct: bool = true, + /// staging buffers and their size: the host transient of any park or restore + staging: u32 = 16, + chunk: u32 = 8 << 20, + lanes: u32 = 4, + /// delta files a chain may stack before a park writes the entry whole + max_chain: u32 = 8, +}; + +/// A piece of bounded state (rings, carries, taps): small, host bytes. +pub const Blob = struct { + name: []const u8, + bytes: []const u8, +}; + +/// What a park writes: the entry, its tokens from first_page on, and this rank's pages per family, in logical order. +pub const WriteSrc = struct { + key: Digest, + /// the file holding pages [0, first_page) (a delta), or null (a whole entry: first_page 0) + base: ?Digest = null, + tag: u32, + kind: u32 = 0, + page: u32, + pos: u64, + first_page: u32 = 0, + npages: u32, + /// tokens [first_page x page, pos) + ids: []const i32, + store: PageStore, + fence: u64 = 0, + families: []const Family, + /// per family: family-tensor pages of this rank's logical pages [first_page, npages), in logical order + pages: []const []const u32, + blobs: []const Blob = &.{}, +}; + +/// One file of a restore: read logical pages [lo, hi) of it into the slot. +pub const Extent = struct { + key: Digest, + first_page: u32, + lo: u32, + hi: u32, +}; + +/// Where a restore puts rows: per family, the family-tensor page of each of the slot's logical pages. +pub const ReadSink = struct { + store: PageStore, + families: []const Family, + page: u32, + world: u32, + rank: u32, + /// per family: family-tensor pages for logical pages [0, npages) + pages: []const []const u32, +}; + +/// The bounded state a restore returns (owned). +pub const Restored = struct { + blobs: []Blob, + + pub fn deinit(r: *Restored, gpa: std.mem.Allocator) void { + for (r.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(r.blobs); + } +}; + +/// The index record of one file. +pub const Rec = struct { + key: Digest, + base: u32 = none, + tag: u32, + kind: u32, + pos: u64, + first_page: u32, + npages: u32, + size: u64, + /// files stacked on this one + refs: u32 = 0, + /// files under it, this one included + depth: u32 = 1, + older: u32 = none, + newer: u32 = none, + live: bool = true, +}; + +pub const Stats = struct { + writes: u64 = 0, + reads: u64 = 0, + bytes_written: u64 = 0, + bytes_read: u64 = 0, + bad: u64 = 0, + evicted: u64 = 0, +}; + +pub const Disk = struct { + gpa: std.mem.Allocator, + io: Io, + opts: Options, + dir: Io.Dir, + staging: Staging, + lanes: *Lanes, + recs: std.ArrayList(Rec) = .empty, + free_recs: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// least recently used first + oldest: u32 = none, + newest: u32 = none, + used: u64 = 0, + stats: Stats = .{}, + + pub fn open(gpa: std.mem.Allocator, io: Io, opts: Options) !*Disk { + const d = try gpa.create(Disk); + errdefer gpa.destroy(d); + var name_buf: [64]u8 = undefined; + const sub = try std.fmt.bufPrint(&name_buf, "{x}/rank{d}", .{ opts.compat[0..8], opts.rank }); + var root = try Io.Dir.cwd().createDirPathOpen(io, opts.root, .{}); + defer root.close(io); + const dir = try root.createDirPathOpen(io, sub, .{ .open_options = .{ .iterate = true } }); + errdefer dir.close(io); + d.* = .{ .gpa = gpa, .io = io, .opts = opts, .dir = dir, .staging = try Staging.init(gpa, io, @max(2, opts.staging), up(opts.chunk)), .lanes = undefined }; + errdefer d.staging.deinit(); + d.lanes = try Lanes.start(gpa, io, &d.staging, @max(1, opts.lanes)); + return d; + } + + pub fn close(d: *Disk) void { + d.lanes.shutdown(); + d.staging.deinit(); + d.dir.close(d.io); + d.recs.deinit(d.gpa); + d.free_recs.deinit(d.gpa); + d.by_key.deinit(d.gpa); + const gpa = d.gpa; + gpa.destroy(d); + } + + pub fn accepts(d: *const Disk, pos: u64) bool { + return pos >= d.opts.min_tokens; + } + + pub fn get(d: *Disk, key: Digest) ?*Rec { + const i = d.by_key.get(key) orelse return null; + return &d.recs.items[i]; + } + + fn fileName(key: Digest, ext: []const u8, buf: *[48]u8) [:0]const u8 { + const n = (std.fmt.bufPrint(buf[0..47], "{x}.{s}", .{ &key, ext }) catch unreachable).len; + buf[n] = 0; + return buf[0..n :0]; + } + + fn openFile(d: *Disk, name: [:0]const u8, write: bool) !struct { fd: linux.fd_t, direct: bool } { + var flags: linux.O = .{ .ACCMODE = if (write) .WRONLY else .RDONLY, .CLOEXEC = true }; + if (write) { + flags.CREAT = true; + flags.TRUNC = true; + } + if (d.opts.direct) { + var f2 = flags; + f2.DIRECT = true; + const rc = linux.openat(d.dir.handle, name, f2, 0o644); + if (linux.errno(rc) == .SUCCESS) return .{ .fd = @intCast(rc), .direct = true }; + } + const rc = linux.openat(d.dir.handle, name, flags, 0o644); + if (linux.errno(rc) != .SUCCESS) return if (linux.errno(rc) == .NOENT) error.FileNotFound else error.OpenFailed; + return .{ .fd = @intCast(rc), .direct = false }; + } + + /// The layout a WriteSrc gets: segments, chunk sizes (whole pages a chunk for pool segments). + pub fn plan(d: *Disk, src: *const WriteSrc) !Layout { + var segs: std.ArrayList(Seg) = .empty; + defer segs.deinit(d.gpa); + const cap = d.opts.chunk; + try segs.append(d.gpa, .{ .name = format.segName("ids"), .kind = .ids, .family = 0, .nbytes = src.ids.len * 4, .chunk = cap }); + for (src.families, 0..) |f, i| { + const pb: u32 = @intCast(f.pageBytes(src.page)); + const k = @max(1, cap / pb); + if (@as(u64, k) * pb > d.staging.size) return error.PageTooLarge; + try segs.append(d.gpa, .{ .name = format.segName(f.name), .kind = .pool, .family = @intCast(i), .nbytes = @as(u64, src.pages[i].len) * pb, .chunk = k * pb }); + } + for (src.blobs) |b| try segs.append(d.gpa, .{ .name = format.segName(b.name), .kind = .blob, .family = 0, .nbytes = b.bytes.len, .chunk = cap }); + const head: Head = .{ .key = src.key, .base = src.base orelse @splat(0), .compat = d.opts.compat, .tag = src.tag, .kind = src.kind, .page = src.page, .world = d.opts.world, .rank = d.opts.rank, .first_page = src.first_page, .npages = src.npages, .nseg = 0, .ndigest = 0, .pos = src.pos }; + return Layout.build(d.gpa, head, segs.items); + } + + /// Writes an entry's file streamed (the pump copies pages to staging, lanes hash and write); returns its bytes. + pub fn writeFile(d: *Disk, src: *const WriteSrc) !u64 { + var lay = try d.plan(src); + defer lay.deinit(d.gpa); + var tb: [48]u8 = undefined; + var fb: [48]u8 = undefined; + const tmp = fileName(src.key, "tmp", &tb); + const f = try d.openFile(tmp, true); + var ok = false; + defer { + _ = linux.close(f.fd); + if (!ok) _ = linux.unlinkat(d.dir.handle, tmp, 0); + } + try src.store.waitFence(src.fence); + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + for (lay.segs) |*s| { + for (0..s.nchunks) |ci| { + if (op.failed()) |e| return e; + const c = s.chunkAt(@intCast(ci)); + const sb = try d.staging.get(); + errdefer d.staging.put(sb); + const buf = d.staging.buf(sb)[0..c.data]; + const at = @as(u64, ci) * s.chunk; + switch (s.kind) { + .ids => @memcpy(buf, std.mem.sliceAsBytes(src.ids)[at..][0..c.data]), + .blob => @memcpy(buf, blobOf(src, s).bytes[at..][0..c.data]), + .pool => { + const pb: u32 = @intCast(src.families[s.family].pageBytes(src.page)); + const p0: usize = @intCast(at / pb); + try src.store.read(s.family, src.pages[s.family][p0..][0 .. c.data / pb], buf); + }, + _ => unreachable, + } + d.lanes.submit(.{ .kind = .write, .fd = f.fd, .staging = sb, .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + ci], .op = &op }); + } + } + try op.wait(); + const hb = try d.staging.get(); + defer d.staging.put(hb); + if (lay.head_span > d.staging.size) return error.HeaderTooLarge; + lay.encode(d.staging.buf(hb)); + try lanes_mod.pwriteAll(f.fd, d.staging.buf(hb)[0..lay.head_span], 0); + if (linux.errno(linux.fdatasync(f.fd)) != .SUCCESS) return error.SyncFailed; + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + if (linux.errno(linux.renameat(d.dir.handle, tmp, d.dir.handle, fileName(src.key, "tfs", &fb))) != .SUCCESS) return error.RenameFailed; + ok = true; + return lay.head.total; + } + + fn blobOf(src: *const WriteSrc, s: *const Seg) Blob { + for (src.blobs) |b| if (std.mem.eql(u8, b.name, s.nameOf())) return b; + unreachable; + } + + /// A file's header, checked. + pub fn readHead(d: *Disk, fd: linux.fd_t) !Layout { + const hb = try d.staging.get(); + defer d.staging.put(hb); + const buf = d.staging.buf(hb); + try lanes_mod.preadAll(fd, buf[0..4096], 0); + const hn = std.mem.readInt(u32, buf[12..16], .little); + if (up(hn) > buf.len) return error.BadHeader; + if (hn > 4096) try lanes_mod.preadAll(fd, buf[4096..up(hn)], 4096); + return Layout.decode(d.gpa, buf[0..up(hn)]); + } + + /// Reads pages [lo, hi) of one restore file into the sink, each chunk checked first; blobs too when want_blobs. + pub fn readFile(d: *Disk, ext: Extent, sink: *const ReadSink, want_blobs: bool) !?Restored { + var nb: [48]u8 = undefined; + const f = try d.openFile(fileName(ext.key, "tfs", &nb), false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + if (!std.mem.eql(u8, &lay.head.key, &ext.key) or lay.head.first_page != ext.first_page or lay.head.page != sink.page) return error.BadHeader; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS or st.size < lay.head.total) return error.Truncated; + var out: std.ArrayList(Blob) = .empty; + errdefer { + for (out.items) |b| { + d.gpa.free(b.name); + d.gpa.free(b.bytes); + } + out.deinit(d.gpa); + } + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + // (segment, chunk) jobs; the pump keeps at most staging - 1 in flight (one buffer for the header path) + var jobs: std.ArrayList(struct { seg: u32, chunk: u32 }) = .empty; + defer jobs.deinit(d.gpa); + for (lay.segs, 0..) |*s, si| switch (s.kind) { + .pool => { + if (s.family >= sink.families.len) return error.BadHeader; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + if (s.chunk % pb != 0 or !std.mem.eql(u8, s.nameOf(), fam.name)) return error.BadHeader; + const need = ownedIn(fam, sink, ext.first_page, ext.hi) - ownedIn(fam, sink, ext.first_page, ext.lo); + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + if (s.nbytes < (@as(u64, skip) + need) * pb) return error.BadHeader; + const kpp = s.chunk / pb; + const c0 = skip / kpp; + const c1 = if (need == 0) c0 else (skip + need + kpp - 1) / kpp; + for (c0..c1) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + .blob => if (want_blobs) { + const buf = try d.gpa.alloc(u8, s.nbytes); + errdefer d.gpa.free(buf); + const name = try d.gpa.dupe(u8, s.nameOf()); + errdefer d.gpa.free(name); + try out.append(d.gpa, .{ .name = name, .bytes = buf }); + for (0..s.nchunks) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + else => {}, + }; + const window = d.staging.count - 1; + var sent: usize = 0; + var seen: usize = 0; + var bytes: u64 = 0; + while (seen < jobs.items.len) { + while (sent < jobs.items.len and sent - seen < window and op.failed() == null) : (sent += 1) { + const jb = jobs.items[sent]; + const s = &lay.segs[jb.seg]; + const c = s.chunkAt(jb.chunk); + d.lanes.submit(.{ .kind = .read, .fd = f.fd, .staging = try d.staging.get(), .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + jb.chunk], .op = &op, .tag = @intCast(sent) }); + bytes += c.span; + } + const j = op.next() orelse break; + seen += 1; + defer d.staging.put(j.staging); + if (op.failed() != null) continue; + const jb = jobs.items[j.tag]; + const s = &lay.segs[jb.seg]; + const data = d.staging.buf(j.staging)[0..j.data]; + if (s.kind == .blob) { + for (out.items) |b| if (std.mem.eql(u8, b.name, s.nameOf())) @memcpy(@constCast(b.bytes)[@as(u64, jb.chunk) * s.chunk ..][0..data.len], data); + continue; + } + try d.drain(s, jb.chunk, data, ext, sink); + } + try op.wait(); + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + d.stats.bytes_read += bytes; + if (!want_blobs) { + out.deinit(d.gpa); + return null; + } + return .{ .blobs = try out.toOwnedSlice(d.gpa) }; + } + + /// A checked pool chunk's pages that fall in [lo, hi) into the sink's pages. + fn drain(d: *Disk, s: *const Seg, chunk: u32, data: []const u8, ext: Extent, sink: *const ReadSink) !void { + _ = d; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + const kpp = s.chunk / pb; + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + const end = ownedIn(fam, sink, ext.first_page, ext.hi); + const a = @max(chunk * kpp, skip); + const b = @min(chunk * kpp + @as(u32, @intCast(data.len / pb)), end); + if (a >= b) return; + var list: [512]u32 = undefined; + var n: u32 = 0; + var j = a; + while (j < b) : (j += 1) { + list[n] = sink.pages[s.family][logicalOf(fam, sink, ext.first_page, j)]; + n += 1; + if (n == list.len or j + 1 == b) { + const first = j + 1 - n; + try sink.store.write(s.family, list[0..n], data[(first - chunk * kpp) * pb ..][0 .. n * pb]); + n = 0; + } + } + } + + /// Records a written file in the index (round thread, at the plan's point); its base gains a dependent. + pub fn register(d: *Disk, src: *const WriteSrc, size: u64) !u32 { + const id = try d.add(src.key, src.base, src.tag, src.kind, src.pos, src.first_page, src.npages, size); + d.stats.writes += 1; + d.stats.bytes_written += size; + return id; + } + + fn add(d: *Disk, key: Digest, base: ?Digest, tag: u32, kind: u32, pos: u64, first_page: u32, npages: u32, size: u64) !u32 { + var r: Rec = .{ .key = key, .tag = tag, .kind = kind, .pos = pos, .first_page = first_page, .npages = npages, .size = size }; + if (base) |b| { + const bi = d.by_key.get(b) orelse return error.NoBase; + r.base = bi; + r.depth = d.recs.items[bi].depth + 1; + } + const id: u32 = if (d.free_recs.pop()) |i| i else blk: { + try d.recs.append(d.gpa, undefined); + break :blk @intCast(d.recs.items.len - 1); + }; + try d.by_key.put(d.gpa, key, id); + if (r.base != none) d.recs.items[r.base].refs += 1; + d.recs.items[id] = r; + d.link(id); + d.used += size; + return id; + } + + /// A valid file found on restart: its header, size and verified ids (tokens [first_page x page, pos), owned). + pub const Found = struct { + head: Head, + size: u64, + mtime: i64, + ids: []i32, + + fn older(_: void, a: Found, b: Found) bool { + return a.mtime < b.mtime or (a.mtime == b.mtime and a.head.first_page < b.head.first_page); + } + }; + + /// Indexes the directory on restart: drops .tmp and bad files, checks ids, registers bases before their deltas. + pub fn scan(d: *Disk, out: *std.ArrayList(Found)) !void { + var it = d.dir.iterate(); + var nb: [48]u8 = undefined; + while (try it.next(d.io)) |ent| { + if (ent.name.len >= nb.len) continue; + @memcpy(nb[0..ent.name.len], ent.name); + nb[ent.name.len] = 0; + const name = nb[0..ent.name.len :0]; + if (std.mem.endsWith(u8, name, ".tmp")) { + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + } + if (!std.mem.endsWith(u8, name, ".tfs")) continue; + const f = d.found(name) catch { + d.stats.bad += 1; + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + }; + try out.append(d.gpa, f); + } + std.mem.sort(Found, out.items, {}, Found.older); + var i: usize = 0; + while (i < out.items.len) { + const f = out.items[i]; + const base: ?Digest = if (f.head.hasBase()) f.head.base else null; + if (base != null and d.by_key.get(base.?) == null) { + d.stats.bad += 1; + var kb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(f.head.key, "tfs", &kb), 0); + d.gpa.free(f.ids); + _ = out.orderedRemove(i); + continue; + } + _ = try d.add(f.head.key, base, f.head.tag, f.head.kind, f.head.pos, f.head.first_page, f.head.npages, f.size); + i += 1; + } + } + + fn found(d: *Disk, name: [:0]const u8) !Found { + const f = try d.openFile(name, false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + const h = lay.head; + var kb: [48]u8 = undefined; + if (!std.mem.eql(u8, fileName(h.key, "tfs", &kb), name)) return error.BadHeader; + if (!std.mem.eql(u8, &h.compat, &d.opts.compat) or h.rank != d.opts.rank or h.world != d.opts.world) return error.Foreign; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true, .MTIME = true }, &st)) != .SUCCESS or st.size < h.total) return error.Truncated; + const seg = lay.find(.ids, "ids") orelse return error.BadHeader; + if (seg.nbytes % 4 != 0 or seg.nbytes / 4 != h.pos - @as(u64, h.first_page) * h.page) return error.BadHeader; + const ids = try d.gpa.alloc(i32, seg.nbytes / 4); + errdefer d.gpa.free(ids); + const sb = try d.staging.get(); + defer d.staging.put(sb); + for (0..seg.nchunks) |ci| { + const c = seg.chunkAt(@intCast(ci)); + const buf = d.staging.buf(sb); + try lanes_mod.preadAll(f.fd, buf[0..c.span], c.off); + var dg: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(buf[0..c.data], &dg, .{}); + if (!std.mem.eql(u8, &dg, &lay.digests[seg.first_digest + ci])) return error.ChecksumMismatch; + @memcpy(std.mem.sliceAsBytes(ids)[@as(usize, ci) * seg.chunk ..][0..c.data], buf[0..c.data]); + } + return .{ .head = h, .size = st.size, .mtime = st.mtime.sec * std.time.ns_per_s + st.mtime.nsec, .ids = ids }; + } + + fn link(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + r.older = d.newest; + r.newer = none; + if (d.newest != none) d.recs.items[d.newest].newer = id else d.oldest = id; + d.newest = id; + } + + fn unlink(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + if (r.older != none) d.recs.items[r.older].newer = r.newer else d.oldest = r.newer; + if (r.newer != none) d.recs.items[r.newer].older = r.older else d.newest = r.older; + } + + /// Marks a file (and its bases) used now: O(chain). + pub fn touch(d: *Disk, key: Digest) void { + var id = d.by_key.get(key) orelse return; + while (id != none) : (id = d.recs.items[id].base) { + d.unlink(id); + d.link(id); + } + // bases end up newest: walk order keeps the entry itself older than its bases, so trims reach dependents first + } + + /// The files a restore of key reads, base first, with the logical pages each gives. + pub fn extents(d: *Disk, key: Digest, out: *std.ArrayList(Extent)) !void { + out.clearRetainingCapacity(); + var id = d.by_key.get(key) orelse return error.FileNotFound; + var hi = d.recs.items[id].npages; + while (id != none) { + const r = &d.recs.items[id]; + try out.append(d.gpa, .{ .key = r.key, .first_page = r.first_page, .lo = r.first_page, .hi = hi }); + hi = r.first_page; + id = r.base; + } + std.mem.reverse(Extent, out.items); + } + + /// Deletes a file without dependents (error when one stacks on it). + pub fn remove(d: *Disk, key: Digest) !void { + const id = d.by_key.get(key) orelse return; + const r = &d.recs.items[id]; + if (r.refs > 0) return error.HasDependents; + if (r.base != none) d.recs.items[r.base].refs -= 1; + d.unlink(id); + _ = d.by_key.remove(key); + d.used -= r.size; + r.live = false; + var nb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(key, "tfs", &nb), 0); + try d.free_recs.append(d.gpa, id); + } + + /// Deletes least recently used files without dependents while over budget; their keys go to out. + pub fn trim(d: *Disk, out: *std.ArrayList(Digest)) !void { + var id = d.oldest; + while (d.used > d.opts.budget and id != none) { + const next = d.recs.items[id].newer; + if (d.recs.items[id].refs == 0) { + const key = d.recs.items[id].key; + try d.remove(key); + try out.append(d.gpa, key); + d.stats.evicted += 1; + id = d.oldest; // a removed file may free its base + continue; + } + id = next; + } + } +}; + +/// Owned (this rank's) logical pages of family f in [first, k): the file's index of logical page k. +pub fn ownedIn(f: Family, sink: *const ReadSink, first: u32, k: u32) u32 { + if (!(f.split and sink.world > 1)) return k - first; + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, k, sink.world, &by); + return by[sink.rank]; +} + +/// The logical page of a family's j-th stored page of a file that starts at first. +pub fn logicalOf(f: Family, sink: *const ReadSink, first: u32, j: u32) u32 { + if (!(f.split and sink.world > 1)) return first + j; + const w = sink.world; + const r0 = first + (sink.rank + w - first % w) % w; // the first owned logical page at or after first + return r0 + j * w; +} diff --git a/zig/src/sessions/disk_test.zig b/zig/src/sessions/disk_test.zig new file mode 100644 index 0000000000..6ac0525e45 --- /dev/null +++ b/zig/src/sessions/disk_test.zig @@ -0,0 +1,125 @@ +//! The NVMe tier on real files: whole and delta files round-trip bit for bit, split ranks, damage. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk = @import("disk.zig"); + +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Fills a page of every family with bytes that name (family, logical page, byte). +fn fill(h: *HostPool, p: *const Pool, logical: u32, pg: u32, salt: u8) void { + for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pg)) continue; + const v = h.pageOf(@intCast(fi), p.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 131 +% logical *% 7 +% fi *% 29 +% salt); + } +} + +fn familyPages(gpa: std.mem.Allocator, p: *const Pool, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const out = try gpa.alloc([]u32, fams.len); + for (fams, 0..) |f, fi| { + var l: std.ArrayList(u32) = .empty; + for (pages[first..], first..) |pg, k| { + if (f.split and p.world > 1 and owned_only and k % p.world != p.rank) continue; + try l.append(gpa, p.familyPage(f, pg)); + } + out[fi] = try l.toOwnedSlice(gpa); + } + return out; +} + +fn freePages(gpa: std.mem.Allocator, x: [][]u32) void { + for (x) |l| gpa.free(l); + gpa.free(x); +} + +fn roundTrip(world: u32, rank: u32) !void { + const gpa = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(root); + var p = try Pool.init(gpa, .{ .families = &fams, .page = 256 }, 64 * 256, world, rank); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + // small chunks so pool segments span many chunks; few staging buffers so the pump waits on the lanes + const d = try disk.Disk.open(gpa, io, .{ .root = root, .compat = @splat(5), .rank = rank, .world = world, .chunk = 3 * 74_752, .staging = 3, .lanes = 3, .min_tokens = 1 }); + defer d.close(); + const a = try p.newSlot(32 * 256); + const pos: u64 = 9 * 256 + 100; + try a.ensure(pos); + for (a.mapped(), 0..) |pg, k| fill(&h, &p, @intCast(k), pg, 1); + const ids = try gpa.alloc(i32, pos); + defer gpa.free(ids); + for (ids, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + var blob: [5000]u8 = undefined; + for (&blob, 0..) |*x, i| x.* = @truncate(9 - i % 5); + const fp = try familyPages(gpa, &p, a.mapped(), 0, true); + defer freePages(gpa, fp); + const src: disk.WriteSrc = .{ .key = @splat(1), .tag = 4, .page = 256, .pos = pos, .npages = a.len, .ids = ids, .store = h.store(), .families = &fams, .pages = fp, .blobs = &.{.{ .name = "rings", .bytes = &blob }} }; + const size = try d.writeFile(&src); + _ = try d.register(&src, size); + // a newer turn stacks a delta on it: pages [9, 14) only (page 9 was partial in the base) + const pos2: u64 = 13 * 256 + 7; + try a.ensure(pos2); + for (a.mapped()[9..], 9..) |pg, k| fill(&h, &p, @intCast(k), pg, 2); + const ids2 = try gpa.alloc(i32, pos2); + defer gpa.free(ids2); + for (ids2, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + const fp2 = try familyPages(gpa, &p, a.mapped(), 9, true); + defer freePages(gpa, fp2); + const src2: disk.WriteSrc = .{ .key = @splat(2), .base = @splat(1), .tag = 4, .page = 256, .pos = pos2, .first_page = 9, .npages = a.len, .ids = ids2[9 * 256 ..], .store = h.store(), .families = &fams, .pages = fp2, .blobs = &.{.{ .name = "rings", .bytes = blob[0..77] }} }; + const size2 = try d.writeFile(&src2); + _ = try d.register(&src2, size2); + try testing.expect(size2 < size); + try testing.expectError(error.HasDependents, d.remove(@splat(1))); + // restore the delta onto other pages through its chain + const b = try p.newSlot(32 * 256); + try b.ensure(pos2); + var ext: std.ArrayList(disk.Extent) = .empty; + defer ext.deinit(gpa); + try d.extents(@splat(2), &ext); + try testing.expectEqual(@as(usize, 2), ext.items.len); + try testing.expectEqual(@as(u32, 9), ext.items[0].hi); + const tp = try familyPages(gpa, &p, b.mapped(), 0, false); + defer freePages(gpa, tp); + const sink: disk.ReadSink = .{ .store = h.store(), .families = &fams, .page = 256, .world = world, .rank = rank, .pages = tp }; + try testing.expect(try d.readFile(ext.items[0], &sink, false) == null); + var got = (try d.readFile(ext.items[1], &sink, true)).?; + defer got.deinit(gpa); + try testing.expectEqualSlices(u8, blob[0..77], got.blobs[0].bytes); + for (a.mapped(), b.mapped()) |pa, pb| for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pa)) continue; + try testing.expectEqualSlices(u8, h.pageOf(@intCast(fi), p.familyPage(f, pa)), h.pageOf(@intCast(fi), p.familyPage(f, pb))); + }; + // damage one chunk: the read refuses it before or as its rows land + var nb: [48]u8 = undefined; + const key2: [16]u8 = @splat(2); + const name = try std.fmt.bufPrint(&nb, "{x}.tfs", .{&key2}); + const f = try d.dir.openFile(io, name, .{ .mode = .read_write }); + var one = [_]u8{0}; + _ = try f.readPositional(io, &.{&one}, 3 * 4096 + 100); + one[0] ^= 0xff; + try f.writePositionalAll(io, &one, 3 * 4096 + 100); + f.close(io); + try testing.expectError(error.ChecksumMismatch, d.readFile(ext.items[1], &sink, true)); +} + +test "whole and delta files round-trip bit for bit (replicated)" { + try roundTrip(1, 0); +} + +test "split ranks write and restore their own pages (rank 0 and 1 of 2)" { + try roundTrip(2, 0); + try roundTrip(2, 1); +} diff --git a/zig/src/sessions/format.zig b/zig/src/sessions/format.zig new file mode 100644 index 0000000000..2ac391d156 --- /dev/null +++ b/zig/src/sessions/format.zig @@ -0,0 +1,197 @@ +//! A session file (format 2): a header, then segments of 4 KiB-aligned chunks, each with its own SHA-256. + +const std = @import("std"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const up = @import("lanes.zig").up; + +pub const magic = "TFSESS02".*; +pub const version: u32 = 2; + +pub const SegKind = enum(u32) { ids = 1, pool = 2, blob = 3, _ }; + +pub const Head = extern struct { + magic: [8]u8 = magic, + version: u32 = version, + /// bytes of Head + segments + digests (before the pad) + head_bytes: u32 = 0, + /// SHA-256 of those bytes with this field zero + head_sha: [32]u8 = @splat(0), + key: [16]u8, + /// the entry whose files hold logical pages [0, first_page); zero: none + base: [16]u8 = @splat(0), + /// what makes these bytes mean the same thing to a build (family, knobs, layout, split) + compat: [16]u8, + tag: u32, + kind: u32, + page: u32, + world: u32, + rank: u32, + first_page: u32, + npages: u32, + nseg: u32, + ndigest: u32, + _pad: u32 = 0, + pos: u64, + total: u64 = 0, + + pub fn hasBase(h: *const Head) bool { + return !std.mem.allEqual(u8, &h.base, 0); + } +}; + +pub const Seg = extern struct { + name: [40]u8, + kind: SegKind, + /// pool: the family's index in the layout + family: u32, + nbytes: u64, + offset: u64 = 0, + /// bytes a chunk (the last may be shorter) and file bytes a chunk (4 KiB multiple) + chunk: u32, + stride: u32 = 0, + first_digest: u32 = 0, + nchunks: u32 = 0, + + pub fn nameOf(s: *const Seg) []const u8 { + return std.mem.sliceTo(&s.name, 0); + } + + /// Chunk i's data bytes and file span. + pub fn chunkAt(s: *const Seg, i: u32) struct { off: u64, data: u32, span: u32 } { + const data: u32 = @intCast(@min(@as(u64, s.chunk), s.nbytes - @as(u64, i) * s.chunk)); + return .{ .off = s.offset + @as(u64, i) * s.stride, .data = data, .span = @intCast(up(data)) }; + } +}; + +pub const Digest32 = [32]u8; + +pub fn segName(name: []const u8) [40]u8 { + var n: [40]u8 = @splat(0); + @memcpy(n[0..@min(name.len, 39)], name[0..@min(name.len, 39)]); + return n; +} + +/// A file's layout: offsets, chunks and the digest slots, before or after its chunks are written. +pub const Layout = struct { + head: Head, + segs: []Seg, + digests: []Digest32, + /// file bytes of the padded header + head_span: u64, + + /// Places segments (each with name, kind, family, nbytes, chunk) after the header. + pub fn build(gpa: std.mem.Allocator, head: Head, segs_in: []const Seg) !Layout { + const segs = try gpa.dupe(Seg, segs_in); + errdefer gpa.free(segs); + var nd: u32 = 0; + for (segs) |*s| { + if (s.chunk == 0) return error.BadSegment; + s.nchunks = @intCast((s.nbytes + s.chunk - 1) / s.chunk); + s.stride = @intCast(up(s.chunk)); + s.first_digest = nd; + nd += s.nchunks; + } + var h = head; + h.nseg = @intCast(segs.len); + h.ndigest = nd; + h.head_bytes = @intCast(@sizeOf(Head) + segs.len * @sizeOf(Seg) + nd * @sizeOf(Digest32)); + const span = up(h.head_bytes); + var off = span; + for (segs) |*s| { + s.offset = off; + if (s.nchunks > 0) off += @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk); + } + h.total = off; + const d = try gpa.alloc(Digest32, nd); + @memset(d, @splat(0)); + return .{ .head = h, .segs = segs, .digests = d, .head_span = span }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.segs); + gpa.free(l.digests); + l.* = undefined; + } + + /// The header's bytes with its checksum into out (len >= head_span; the pad is zeroed). + pub fn encode(l: *Layout, out: []u8) void { + @memset(out[0..l.head_span], 0); + var h = l.head; + h.head_sha = @splat(0); + const hb = std.mem.asBytes(&h); + const sb = std.mem.sliceAsBytes(l.segs); + const db = std.mem.sliceAsBytes(l.digests); + @memcpy(out[0..hb.len], hb); + @memcpy(out[hb.len..][0..sb.len], sb); + @memcpy(out[hb.len + sb.len ..][0..db.len], db); + var sha: [32]u8 = undefined; + Sha256.hash(out[0..h.head_bytes], &sha, .{}); + @memcpy(out[@offsetOf(Head, "head_sha")..][0..32], &sha); + l.head.head_sha = sha; + } + + /// The header read back from bytes (at least head_bytes of them), every field and offset checked. + pub fn decode(gpa: std.mem.Allocator, bytes: []const u8) !Layout { + if (bytes.len < @sizeOf(Head)) return error.BadHeader; + var h: Head = undefined; + @memcpy(std.mem.asBytes(&h), bytes[0..@sizeOf(Head)]); + if (!std.mem.eql(u8, &h.magic, &magic) or h.version != version) return error.BadHeader; + const want = @sizeOf(Head) + @as(u64, h.nseg) * @sizeOf(Seg) + @as(u64, h.ndigest) * @sizeOf(Digest32); + if (h.head_bytes != want or bytes.len < want) return error.BadHeader; + var check: [32]u8 = undefined; + var sha = Sha256.init(.{}); + sha.update(bytes[0..@offsetOf(Head, "head_sha")]); + sha.update(&@as([32]u8, @splat(0))); + sha.update(bytes[@offsetOf(Head, "head_sha") + 32 .. want]); + sha.final(&check); + if (!std.mem.eql(u8, &check, &h.head_sha)) return error.ChecksumMismatch; + const segs = try gpa.alloc(Seg, h.nseg); + errdefer gpa.free(segs); + @memcpy(std.mem.sliceAsBytes(segs), bytes[@sizeOf(Head)..][0 .. segs.len * @sizeOf(Seg)]); + const d = try gpa.alloc(Digest32, h.ndigest); + errdefer gpa.free(d); + @memcpy(std.mem.sliceAsBytes(d), bytes[@sizeOf(Head) + segs.len * @sizeOf(Seg) ..][0 .. d.len * 32]); + var nd: u64 = 0; + for (segs) |s| { + if (s.chunk == 0 or s.stride != up(s.chunk) or s.first_digest != nd or s.offset % 4096 != 0) return error.BadHeader; + if (s.nchunks != (s.nbytes + s.chunk - 1) / s.chunk) return error.BadHeader; + if (s.nchunks > 0 and s.offset + @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk) > h.total) return error.BadHeader; + nd += s.nchunks; + } + if (nd != h.ndigest) return error.BadHeader; + return .{ .head = h, .segs = segs, .digests = d, .head_span = up(h.head_bytes) }; + } + + pub fn find(l: *const Layout, kind: SegKind, name: []const u8) ?*const Seg { + for (l.segs) |*s| if (s.kind == kind and std.mem.eql(u8, s.nameOf(), name)) return s; + return null; + } +}; + +test "a layout encodes, decodes, and refuses a damaged header" { + const gpa = std.testing.allocator; + const head: Head = .{ .key = @splat(1), .compat = @splat(2), .tag = 3, .kind = 0, .page = 256, .world = 2, .rank = 1, .first_page = 4, .npages = 9, .nseg = 0, .ndigest = 0, .pos = 2200 }; + var l = try Layout.build(gpa, head, &.{ + .{ .name = segName("ids"), .kind = .ids, .family = 0, .nbytes = 4 * (2200 - 1024), .chunk = 1 << 20 }, + .{ .name = segName("comp.20"), .kind = .pool, .family = 0, .nbytes = 3 * 149_504, .chunk = 2 * 149_504 }, + .{ .name = segName("taps"), .kind = .blob, .family = 0, .nbytes = 0, .chunk = 4096 }, + }); + defer l.deinit(gpa); + try std.testing.expectEqual(@as(u32, 1 + 2 + 0), l.head.ndigest); + try std.testing.expectEqual(@as(u64, 4096), l.segs[0].offset); + try std.testing.expectEqual(@as(u64, 4096 + up(4 * 1176)), l.segs[1].offset); + const c1 = l.segs[1].chunkAt(1); + try std.testing.expectEqual(@as(u32, 149_504), c1.data); + try std.testing.expectEqual(l.segs[1].offset + up(2 * 149_504), c1.off); + try std.testing.expectEqual(l.segs[2].offset, l.head.total); + l.digests[1][5] = 9; + var buf: [8192]u8 = undefined; + l.encode(&buf); + var back = try Layout.decode(gpa, &buf); + defer back.deinit(gpa); + try std.testing.expectEqual(l.head.total, back.head.total); + try std.testing.expectEqual(@as(u8, 9), back.digests[1][5]); + try std.testing.expectEqualStrings("comp.20", back.find(.pool, "comp.20").?.nameOf()); + buf[200] ^= 0x40; + try std.testing.expectError(error.ChecksumMismatch, Layout.decode(gpa, &buf)); +} diff --git a/zig/src/sessions/freeset.zig b/zig/src/sessions/freeset.zig new file mode 100644 index 0000000000..b8cce323af --- /dev/null +++ b/zig/src/sessions/freeset.zig @@ -0,0 +1,114 @@ +//! A set of free page indices, lowest first in O(log64 n): a bit a page, then a bit a non-empty word a level up. + +const std = @import("std"); + +const max_levels = 6; // 64^6 pages: more than any pool + +pub const FreeSet = struct { + /// levels[0] holds a bit a page, levels[k] a bit a non-empty word of levels[k - 1] + levels: [max_levels][]u64 = @splat(&.{}), + depth: u8 = 0, + n: u32 = 0, + free: u32 = 0, + + /// All of 0 .. n - 1 free. + pub fn init(gpa: std.mem.Allocator, n: u32) !FreeSet { + var s: FreeSet = .{ .n = n }; + errdefer s.deinit(gpa); + var bits: usize = @max(n, 1); + while (true) { + const words = (bits + 63) / 64; + s.levels[s.depth] = try gpa.alloc(u64, words); + @memset(s.levels[s.depth], 0); + s.depth += 1; + if (words == 1) break; + bits = words; + } + for (0..n) |i| s.setBit(@intCast(i)); + s.free = n; + return s; + } + + pub fn deinit(s: *FreeSet, gpa: std.mem.Allocator) void { + for (s.levels[0..s.depth]) |l| gpa.free(l); + s.* = .{}; + } + + pub fn count(s: *const FreeSet) u32 { + return s.free; + } + + pub fn contains(s: *const FreeSet, i: u32) bool { + return s.levels[0][i / 64] & (@as(u64, 1) << @intCast(i % 64)) != 0; + } + + /// The lowest free index, taken; null when none is free. + pub fn takeLowest(s: *FreeSet) ?u32 { + if (s.free == 0) return null; + var idx: usize = 0; + var lv: usize = s.depth; + while (lv > 0) { + lv -= 1; + const w = s.levels[lv][idx]; + idx = idx * 64 + @ctz(w); + } + const i: u32 = @intCast(idx); + s.clearBit(i); + s.free -= 1; + return i; + } + + /// Takes a given free index (error when it is not free). + pub fn take(s: *FreeSet, i: u32) error{NotFree}!void { + if (i >= s.n or !s.contains(i)) return error.NotFree; + s.clearBit(i); + s.free -= 1; + } + + /// Gives an index back (error when it is already free: a double free). + pub fn give(s: *FreeSet, i: u32) error{DoubleFree}!void { + if (i >= s.n or s.contains(i)) return error.DoubleFree; + s.setBit(i); + s.free += 1; + } + + fn setBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + const was = l[idx / 64]; + l[idx / 64] = was | (@as(u64, 1) << @intCast(idx % 64)); + if (was != 0) return; // the word was non-empty: the levels above already say so + idx /= 64; + } + } + + fn clearBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + l[idx / 64] &= ~(@as(u64, 1) << @intCast(idx % 64)); + if (l[idx / 64] != 0) return; // the word still has a free bit: the levels above stay set + idx /= 64; + } + } +}; + +test "lowest first, give back, double free refused, across levels" { + const gpa = std.testing.allocator; + var s = try FreeSet.init(gpa, 70_000); // three levels + defer s.deinit(gpa); + try std.testing.expectEqual(@as(u8, 3), s.depth); + for (0..5000) |i| try std.testing.expectEqual(@as(?u32, @intCast(i)), s.takeLowest()); + try s.give(4097); + try s.give(17); + try std.testing.expectError(error.DoubleFree, s.give(17)); + try std.testing.expectEqual(@as(?u32, 17), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 4097), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 5000), s.takeLowest()); + try s.take(69_999); + try std.testing.expectError(error.NotFree, s.take(69_999)); + try std.testing.expectEqual(@as(u32, 70_000 - 5001 - 1), s.count()); + while (s.takeLowest()) |_| {} + try std.testing.expectEqual(@as(u32, 0), s.count()); + try s.give(69_999); + try std.testing.expectEqual(@as(?u32, 69_999), s.takeLowest()); +} diff --git a/zig/src/sessions/jobs.zig b/zig/src/sessions/jobs.zig new file mode 100644 index 0000000000..5b445f20e7 --- /dev/null +++ b/zig/src/sessions/jobs.zig @@ -0,0 +1,150 @@ +//! Parks and restores off the round: a tier thread runs each job's file I/O until the plan settles it (Store.settle). + +const std = @import("std"); +const Io = std.Io; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; + +pub const Kind = enum { park, restore }; +pub const State = enum(u8) { queued, running, done, failed }; + +pub const Job = struct { + kind: Kind, + /// the store's entry + entry: u32, + state: std.atomic.Value(State) = .init(.queued), + err: ?anyerror = null, + started_ns: i96 = 0, + elapsed_ns: u64 = 0, + /// park: the file to write (owned buffers below) and its size once written + src: disk_mod.WriteSrc = undefined, + size: u64 = 0, + /// restore: files base first, the sink, the bounded state read + extents: []disk_mod.Extent = &.{}, + sink: disk_mod.ReadSink = undefined, + restored: ?disk_mod.Restored = null, + /// restore: the slot it fills + slot: u32 = 0, + /// pool pages the job holds until settled + held: []u32 = &.{}, + /// owned buffers + ids: []i32 = &.{}, + lists: [][]u32 = &.{}, + blobs: []disk_mod.Blob = &.{}, + + pub fn done(j: *const Job) bool { + const s = j.state.load(.acquire); + return s == .done or s == .failed; + } + + pub fn deinit(j: *Job, gpa: std.mem.Allocator) void { + gpa.free(j.held); + gpa.free(j.ids); + for (j.lists) |l| gpa.free(l); + gpa.free(j.lists); + for (j.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(j.blobs); + gpa.free(j.extents); + if (j.restored) |*r| r.deinit(gpa); + gpa.destroy(j); + } + + fn run(j: *Job, d: *Disk, io: Io) void { + j.state.store(.running, .release); + const t0 = Io.Clock.awake.now(io).toNanoseconds(); + j.err = switch (j.kind) { + .park => if (d.writeFile(&j.src)) |n| blk: { + j.size = n; + break :blk null; + } else |e| e, + .restore => j.restore(d), + }; + j.elapsed_ns = @intCast(Io.Clock.awake.now(io).toNanoseconds() - t0); + j.state.store(if (j.err == null) .done else .failed, .release); + } + + fn restore(j: *Job, d: *Disk) ?anyerror { + for (j.extents, 0..) |e, i| { + const last = i + 1 == j.extents.len; + const r = d.readFile(e, &j.sink, last) catch |err| return err; + if (last) j.restored = r; + } + j.sink.store.publish() catch |err| return err; + return null; + } +}; + +/// The tier thread and its queue (null thread: jobs run inline). +pub const Runner = struct { + gpa: std.mem.Allocator, + io: Io, + disk: *Disk, + thread: ?std.Thread = null, + queue: std.ArrayList(*Job) = .empty, + head: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + + pub fn init(gpa: std.mem.Allocator, io: Io, d: *Disk) Runner { + return .{ .gpa = gpa, .io = io, .disk = d }; + } + + /// Starts the tier thread: jobs then run off the caller's thread. + pub fn startThread(r: *Runner) !void { + if (r.thread != null) return; + r.thread = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, loop, .{r}); + } + + pub fn deinit(r: *Runner) void { + if (r.thread) |t| { + r.mutex.lockUncancelable(r.io); + r.stop = true; + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + t.join(); + } + r.queue.deinit(r.gpa); + } + + pub fn submit(r: *Runner, j: *Job) !void { + if (r.thread == null) return j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (r.head == r.queue.items.len) { + r.queue.clearRetainingCapacity(); + r.head = 0; + } + try r.queue.append(r.gpa, j); + r.cond.signal(r.io); + } + + /// Waits for a job (the plan settles it next). + pub fn wait(r: *Runner, j: *Job) void { + while (!j.done()) { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (!j.done()) r.cond.waitTimeout(r.io, &r.mutex, .{ .duration = .{ .raw = .fromMilliseconds(5), .clock = .awake } }) catch {}; + } + } + + fn loop(r: *Runner) void { + while (true) { + const j = blk: { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + while (r.head == r.queue.items.len and !r.stop) r.cond.waitUncancelable(r.io, &r.mutex); + if (r.head == r.queue.items.len) return; + r.head += 1; + break :blk r.queue.items[r.head - 1]; + }; + j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + } + } +}; diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig new file mode 100644 index 0000000000..38fad8cd89 --- /dev/null +++ b/zig/src/sessions/lanes.zig @@ -0,0 +1,336 @@ +//! The NVMe tier's lane threads, which hash and write (or read and check) 4 KiB chunks, and their staging buffers. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const block = 4096; + +pub fn up(n: u64) u64 { + return (n + block - 1) / block * block; +} + +/// Bounded, aligned staging buffers (allocated once; `release` gives the memory back while idle). +pub const Staging = struct { + io: Io, + size: usize, + count: u32, + bufs: [][]align(block) u8, + free_list: []u32, + free_n: u32, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, io: Io, count: u32, size: usize) !Staging { + if (count == 0 or size % block != 0) return error.BadStaging; + const bufs = try gpa.alloc([]align(block) u8, count); + errdefer gpa.free(bufs); + const fl = try gpa.alloc(u32, count); + errdefer gpa.free(fl); + for (bufs, 0..) |*b, i| { + b.* = &.{}; + fl[i] = @intCast(i); + } + return .{ .io = io, .size = size, .count = count, .bufs = bufs, .free_list = fl, .free_n = count, .gpa = gpa }; + } + + pub fn deinit(s: *Staging) void { + s.release(); + s.gpa.free(s.bufs); + s.gpa.free(s.free_list); + s.* = undefined; + } + + /// A free buffer's index, waiting for one; its memory is mapped on first use. + pub fn get(s: *Staging) !u32 { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + while (s.free_n == 0) s.cond.waitUncancelable(s.io, &s.mutex); + s.free_n -= 1; + const i = s.free_list[s.free_n]; + if (s.bufs[i].len == 0) s.bufs[i] = try s.gpa.alignedAlloc(u8, .fromByteUnits(block), s.size); + return i; + } + + pub fn put(s: *Staging, i: u32) void { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + s.free_list[s.free_n] = i; + s.free_n += 1; + s.cond.signal(s.io); + } + + pub fn buf(s: *Staging, i: u32) []align(block) u8 { + return s.bufs[i]; + } + + /// Unmaps every buffer (all must be free): the memory goes back between operations. + pub fn release(s: *Staging) void { + std.debug.assert(s.free_n == s.count); + for (s.bufs) |*b| if (b.len != 0) { + s.gpa.free(b.*); + b.* = &.{}; + }; + } + + pub fn bytes(s: *const Staging) u64 { + return @as(u64, s.count) * s.size; + } +}; + +pub const Kind = enum { write, read }; + +/// One chunk's I/O: buf[0 .. span] at file offset off (span: the chunk padded to 4 KiB; data: the chunk's own bytes). +pub const Job = struct { + kind: Kind, + fd: linux.fd_t, + staging: u32, + data: u32, + span: u32, + off: u64, + /// write: the chunk's digest goes here; read: the digest it must have + digest: *[32]u8, + op: *Op, + /// the pump's tag for a finished read + tag: u32 = 0, +}; + +/// A park's or restore's chunks in flight: how many are pending, the first error, and finished reads for the pump. +pub const Op = struct { + io: Io, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + pending: u32 = 0, + err: ?anyerror = null, + /// finished reads (their jobs), drained by the pump + done: std.ArrayList(Job) = .empty, + gpa: std.mem.Allocator, + /// write jobs hand their staging back themselves + staging: *Staging, + + pub fn init(gpa: std.mem.Allocator, io: Io, staging: *Staging) !Op { + var o: Op = .{ .io = io, .gpa = gpa, .staging = staging }; + try o.done.ensureTotalCapacity(gpa, staging.count); + return o; + } + + pub fn deinit(o: *Op) void { + o.done.deinit(o.gpa); + } + + fn finish(o: *Op, j: Job, err: ?anyerror) void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + o.pending -= 1; + if (err) |e| { + if (o.err == null) o.err = e; + } + if (j.kind == .read) o.done.appendAssumeCapacity(j) else o.staging.put(j.staging); + o.cond.broadcast(o.io); + } + + /// The next finished read, waiting for one; null when none is pending. + pub fn next(o: *Op) ?Job { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.done.items.len == 0 and o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + return o.done.pop(); + } + + /// Waits for every pending job; returns the first error. + pub fn wait(o: *Op) !void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + if (o.err) |e| return e; + } + + pub fn failed(o: *Op) ?anyerror { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + return o.err; + } +}; + +/// The lane threads and their queue. +pub const Lanes = struct { + io: Io, + gpa: std.mem.Allocator, + staging: *Staging, + threads: []std.Thread, + /// queued jobs: a ring as long as the staging set (every queued job holds a buffer) + ring: []Job, + head: usize = 0, + len: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + /// seconds the lanes spent hashing and in I/O (summed over lanes) + hash_ns: std.atomic.Value(u64) = .init(0), + io_ns: std.atomic.Value(u64) = .init(0), + + pub fn start(gpa: std.mem.Allocator, io: Io, staging: *Staging, n: u32) !*Lanes { + const l = try gpa.create(Lanes); + errdefer gpa.destroy(l); + const ring = try gpa.alloc(Job, staging.count); + errdefer gpa.free(ring); + l.* = .{ .io = io, .gpa = gpa, .staging = staging, .threads = try gpa.alloc(std.Thread, n), .ring = ring }; + errdefer gpa.free(l.threads); + var made: usize = 0; + errdefer { + l.halt(); + for (l.threads[0..made]) |t| t.join(); + } + for (l.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, run, .{l}); + made += 1; + } + return l; + } + + fn halt(l: *Lanes) void { + l.mutex.lockUncancelable(l.io); + l.stop = true; + l.cond.broadcast(l.io); + l.mutex.unlock(l.io); + } + + pub fn shutdown(l: *Lanes) void { + l.halt(); + for (l.threads) |t| t.join(); + l.gpa.free(l.ring); + l.gpa.free(l.threads); + l.gpa.destroy(l); + } + + /// Queues a job (each holds a staging buffer, so the ring never overflows). + pub fn submit(l: *Lanes, j: Job) void { + { + j.op.mutex.lockUncancelable(l.io); + defer j.op.mutex.unlock(l.io); + j.op.pending += 1; + } + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + std.debug.assert(l.len < l.ring.len); + l.ring[(l.head + l.len) % l.ring.len] = j; + l.len += 1; + l.cond.signal(l.io); + } + + fn run(l: *Lanes) void { + while (true) { + const j = blk: { + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + while (l.len == 0 and !l.stop) l.cond.waitUncancelable(l.io, &l.mutex); + if (l.len == 0) return; + const j = l.ring[l.head]; + l.head = (l.head + 1) % l.ring.len; + l.len -= 1; + break :blk j; + }; + j.op.finish(j, l.work(j)); + } + } + + fn work(l: *Lanes, j: Job) ?anyerror { + const b = l.staging.buf(j.staging); + var t0 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + switch (j.kind) { + .write => { + Sha256.hash(b[0..j.data], j.digest, .{}); + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.hash_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + t0 = t1; + @memset(b[j.data..j.span], 0); + pwriteAll(j.fd, b[0..j.span], j.off) catch |e| return e; + _ = l.io_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t0), .monotonic); + }, + .read => { + preadAll(j.fd, b[0..j.span], j.off) catch |e| return e; + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.io_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + var d: [32]u8 = undefined; + Sha256.hash(b[0..j.data], &d, .{}); + _ = l.hash_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t1), .monotonic); + if (!std.mem.eql(u8, &d, j.digest)) return error.ChecksumMismatch; + }, + } + return null; + } +}; + +pub fn pwriteAll(fd: linux.fd_t, data: []const u8, off: u64) !void { + var done: usize = 0; + while (done < data.len) { + const rc = linux.pwrite(fd, data[done..].ptr, data.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + .NOSPC => return error.NoSpaceLeft, + else => return error.WriteFailed, + } + if (rc == 0) return error.WriteFailed; + done += rc; + } +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.ReadFailed, + } + if (rc == 0) return error.ShortRead; + done += rc; + } +} + +test "lanes write, read back and check chunks through bounded staging" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var st = try Staging.init(gpa, io, 3, 2 * block); + defer st.deinit(); + const l = try Lanes.start(gpa, io, &st, 2); + defer l.shutdown(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const f = try tmp.dir.createFile(io, "lanes.bin", .{ .read = true }); + defer f.close(io); + const fd = f.handle; + var digests: [7][32]u8 = undefined; + var op = try Op.init(gpa, io, &st); + defer op.deinit(); + for (0..7) |i| { + const s = try st.get(); + const n: u32 = if (i == 6) 1000 else 2 * block - 7; + for (st.buf(s)[0..n], 0..) |*b, k| b.* = @truncate(k * 31 + i); + l.submit(.{ .kind = .write, .fd = fd, .staging = s, .data = n, .span = @intCast(up(n)), .off = i * 2 * block, .digest = &digests[i], .op = &op }); + } + try op.wait(); + var rop = try Op.init(gpa, io, &st); + defer rop.deinit(); + var seen: u32 = 0; + var sent: u32 = 0; + while (seen < 7) { + while (sent < 7 and sent - seen < 3) : (sent += 1) { + const n: u32 = if (sent == 6) 1000 else 2 * block - 7; + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = n, .span = @intCast(up(n)), .off = sent * 2 * block, .digest = &digests[sent], .op = &rop, .tag = sent }); + } + const j = rop.next().?; + for (st.buf(j.staging)[0..j.data], 0..) |b, k| try std.testing.expectEqual(@as(u8, @truncate(k * 31 + j.tag)), b); + st.put(j.staging); + seen += 1; + } + try rop.wait(); + digests[2][0] ^= 1; // a damaged chunk is refused + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = 2 * block - 7, .span = 2 * block, .off = 2 * 2 * block, .digest = &digests[2], .op = &rop }); + st.put(rop.next().?.staging); + try std.testing.expectError(error.ChecksumMismatch, rop.wait()); +} diff --git a/zig/src/sessions/pagestore.zig b/zig/src/sessions/pagestore.zig new file mode 100644 index 0000000000..0e4792699b --- /dev/null +++ b/zig/src/sessions/pagestore.zig @@ -0,0 +1,129 @@ +//! The bytes behind the pool's pages: a family's tensor of pages, read and written a page run at a time. + +const std = @import("std"); +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; + +pub const PageStore = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// family fam's pages, back to back into dst (pages.len x a page's bytes) + read: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void, + /// src (pages.len x a page's bytes) into family fam's pages + write: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void, + /// one page's rows of family fam into another page (copy on write of a RAM resume's partial page) + copyPage: *const fn (ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void, + /// round thread: a point after every row the compute stream has written so far + fence: *const fn (ptr: *anyopaque) anyerror!u64, + /// before reads: wait until the fence's rows are visible + waitFence: *const fn (ptr: *anyopaque, token: u64) anyerror!void, + /// after writes: make them visible to the compute stream + publish: *const fn (ptr: *anyopaque) anyerror!void, + }; + + pub fn read(s: PageStore, fam: u32, pages: []const u32, dst: []u8) !void { + return s.vtable.read(s.ptr, fam, pages, dst); + } + pub fn write(s: PageStore, fam: u32, pages: []const u32, src: []const u8) !void { + return s.vtable.write(s.ptr, fam, pages, src); + } + pub fn copyPage(s: PageStore, fam: u32, src: u32, dst: u32) !void { + return s.vtable.copyPage(s.ptr, fam, src, dst); + } + pub fn fence(s: PageStore) !u64 { + return s.vtable.fence(s.ptr); + } + pub fn waitFence(s: PageStore, token: u64) !void { + return s.vtable.waitFence(s.ptr, token); + } + pub fn publish(s: PageStore) !void { + return s.vtable.publish(s.ptr); + } +}; + +/// Every family's pages in host memory: [familyTensorPages x page bytes] a family, zeroed (the null page stays zero). +pub const HostPool = struct { + gpa: std.mem.Allocator, + page: u32, + families: []const Family, + tensors: [][]u8, + + pub fn init(gpa: std.mem.Allocator, p: *const Pool) !HostPool { + const t = try gpa.alloc([]u8, p.layout.families.len); + errdefer gpa.free(t); + var made: usize = 0; + errdefer for (t[0..made]) |x| gpa.free(x); + for (p.layout.families, 0..) |f, i| { + t[i] = try gpa.alloc(u8, @as(usize, p.familyTensorPages(f)) * f.pageBytes(p.page)); + @memset(t[i], 0); + made += 1; + } + return .{ .gpa = gpa, .page = p.page, .families = p.layout.families, .tensors = t }; + } + + pub fn deinit(h: *HostPool) void { + for (h.tensors) |x| h.gpa.free(x); + h.gpa.free(h.tensors); + h.* = undefined; + } + + pub fn store(h: *HostPool) PageStore { + return .{ .ptr = h, .vtable = &vtable }; + } + + /// One page of family fam (a view). + pub fn pageOf(h: *HostPool, fam: u32, pg: u32) []u8 { + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + return h.tensors[fam][pg * n ..][0..n]; + } + + const vtable: PageStore.VTable = .{ .read = read, .write = write, .copyPage = copyPage, .fence = fence, .waitFence = waitFence, .publish = publish }; + + fn read(ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (dst.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(dst[i * n ..][0..n], h.pageOf(fam, pg)); + } + + fn write(ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (src.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(h.pageOf(fam, pg), src[i * n ..][0..n]); + } + + fn copyPage(ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + @memcpy(h.pageOf(fam, dst), h.pageOf(fam, src)); + } + + fn fence(_: *anyopaque) anyerror!u64 { + return 0; + } + fn waitFence(_: *anyopaque, _: u64) anyerror!void {} + fn publish(_: *anyopaque) anyerror!void {} +}; + +test "host pool pages round-trip through the interface" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 1, .row_bytes = 8 }, .{ .name = "ik", .ratio = 2, .row_bytes = 4 } }; + var p = try Pool.init(gpa, .{ .families = &fams, .page = 16 }, 64, 1, 0); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + const s = h.store(); + var src: [2 * 128]u8 = undefined; + for (&src, 0..) |*b, i| b.* = @truncate(i * 7 + 3); + try s.write(0, &.{ 3, 1 }, &src); + var dst: [2 * 128]u8 = undefined; + try s.read(0, &.{ 3, 1 }, &dst); + try std.testing.expectEqualSlices(u8, &src, &dst); + try s.copyPage(0, 3, 0); + try std.testing.expectEqualSlices(u8, src[0..128], h.pageOf(0, 0)); + try std.testing.expect(std.mem.allEqual(u8, h.pageOf(0, p.null_page), 0)); + try std.testing.expectError(error.BadLength, s.read(1, &.{0}, &dst)); +} diff --git a/zig/src/sessions/pool.zig b/zig/src/sessions/pool.zig new file mode 100644 index 0000000000..d3e67516af --- /dev/null +++ b/zig/src/sessions/pool.zig @@ -0,0 +1,475 @@ +//! The paged KV pool's host book: lowest-free pages per residue mod W, holder counts, slot page tables, quotas. + +const std = @import("std"); +const FreeSet = @import("freeset.zig").FreeSet; + +pub const Error = error{ PoolExhausted, PastCapacity, SlotNotEmpty, DoubleFree, BadPool, ResidueMismatch } || std.mem.Allocator.Error; + +/// A cache family: rows of `row_bytes` every `ratio` tokens; `split` families shard by logical page over the ranks. +pub const Family = struct { + name: []const u8, + ratio: u32, + row_bytes: u32, + split: bool = false, + + pub fn pageRows(f: Family, page: u32) u32 { + return page / f.ratio; + } + + /// Bytes of one page of this family (a whole page's rows). + pub fn pageBytes(f: Family, page: u32) u64 { + return @as(u64, f.pageRows(page)) * f.row_bytes; + } +}; + +/// The pool's families and page size: what a token costs and how pages are laid out. +pub const Layout = struct { + families: []const Family, + /// tokens a page (a power of two; 256 for every family we serve) + page: u32 = 256, + + /// A rank's bytes a token: the split families count 1 / W. + pub fn bytesPerToken(l: Layout, world: u32) f64 { + var b: f64 = 0; + for (l.families) |f| b += @as(f64, @floatFromInt(f.row_bytes)) / @as(f64, @floatFromInt(f.ratio)) / @as(f64, @floatFromInt(if (f.split and world > 1) world else 1)); + return b; + } + + /// A rank's average bytes a page (split families 1 / W): what the residency budget counts. + pub fn pageBytes(l: Layout, world: u32) u64 { + var b: u64 = 0; + for (l.families) |f| b += f.pageBytes(l.page) / (if (f.split and world > 1) world else 1); + return b; + } +}; + +pub fn pagesFor(tokens: u64, page: u32) u32 { + return @intCast((tokens + page - 1) / page); +} + +/// Logical pages [k0, k1) counted by residue mod `w` (O(w)). +pub fn span(k0: u32, k1: u32, w: u32, out: []u32) void { + for (out[0..w], 0..) |*o, r| { + const rr: u32 = @intCast(r); + const first = k0 + (rr + w - k0 % w) % w; + o.* = if (first >= k1) 0 else (k1 - 1 - first) / w + 1; + } +} + +pub const max_world = 8; + +/// A file whose logical pages [0, pages) match a slot's or entry's rows: a later park writes only past it. +pub const Origin = struct { + key: [16]u8, + pages: u32, +}; + +/// One slot's page table: logical page i -> physical page; the device tables follow through the dirty range. +pub const Slot = struct { + pool: *Pool, + id: u32, + capacity: u64, + max_pages: u32, + /// physical page of each mapped logical page; len is the mapped count + pages: []u32, + len: u32 = 0, + /// pages the admission reserved for this slot's request (null: no admission control) + quota: ?u32 = null, + /// pages taken past the quota, from pages nobody had reserved + over: u32 = 0, + /// table entries changed since the device side last synced them: [lo, hi) + dirty_lo: u32 = std.math.maxInt(u32), + dirty_hi: u32 = 0, + /// the file the leading pages came from (truncation shortens it, release clears it) + origin: ?Origin = null, + + pub fn mapped(s: *const Slot) []const u32 { + return s.pages[0..s.len]; + } + + pub fn tokens(s: *const Slot) u64 { + return @as(u64, s.len) * s.pool.page; + } + + fn mark(s: *Slot, lo: u32, hi: u32) void { + s.dirty_lo = @min(s.dirty_lo, lo); + s.dirty_hi = @max(s.dirty_hi, hi); + } + + /// The table range to upload since the last call (empty when nothing changed): entries past len are the null page. + pub fn takeDirty(s: *Slot) ?struct { lo: u32, hi: u32 } { + if (s.dirty_lo >= s.dirty_hi) return null; + defer { + s.dirty_lo = std.math.maxInt(u32); + s.dirty_hi = 0; + } + return .{ .lo = s.dirty_lo, .hi = s.dirty_hi }; + } + + /// The device table's value at logical page i (the page, or the null page past len). + pub fn tableAt(s: *const Slot, i: u32) u32 { + return if (i < s.len) s.pages[i] else s.pool.null_page; + } + + /// The split families' device table at logical page i: the local page this rank stores, else the discard page. + pub fn localTableAt(s: *const Slot, i: u32) u32 { + return s.pool.localPage(s.tableAt(i)); + } + + fn reservedSpan(s: *const Slot, out: []u32) void { + const q = s.quota orelse 0; + span(s.len, @max(s.len, q), s.pool.world, out); + } + + /// Maps every page holding positions < end (a no-op when they are: the common case). + pub fn ensure(s: *Slot, end: u64) Error!void { + const p = s.pool; + const need = pagesFor(end, p.page); + if (need <= s.len) return; + if (need > s.max_pages) return error.PastCapacity; + const w = p.world; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + if (s.quota) |q| if (need > q) { + // past the reservation: only pages nobody else reserved, residue by residue + var extra: [max_world]u32 = undefined; + span(@max(q, s.len), need, w, &extra); + for (0..w) |r| if (extra[r] > p.free[r].count() -| (p.reserved[r] - before[r])) return error.PoolExhausted; + s.over += need - @max(q, s.len); + }; + var want: [max_world]u32 = undefined; + span(s.len, need, w, &want); + for (0..w) |r| if (want[r] > p.free[r].count()) return error.PoolExhausted; + const lo = s.len; + for (lo..need) |k| { + const pg = p.takeAt(@intCast(k)) catch unreachable; // counted above + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[k] = pg; + } + s.len = need; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(lo, need); + } + + /// Maps shared pages as the first pages of an empty slot (a RAM resume: read only, the slot writes past them). + pub fn adopt(s: *Slot, shared: []const u32) Error!void { + if (s.len != 0) return error.SlotNotEmpty; + if (shared.len > s.max_pages) return error.PastCapacity; + const p = s.pool; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + for (shared, 0..) |pg, k| if (pg >= p.npages or pg % p.world != k % p.world) return error.ResidueMismatch; + for (shared, 0..) |pg, k| { + p.hold(pg, 1, 1); + s.pages[k] = pg; + } + s.len = @intCast(shared.len); + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(0, s.len); + } + + /// Appends one page taken for logical page len (the copy-on-write page of a RAM resume). + pub fn push(s: *Slot, pg: u32) Error!void { + if (s.len >= s.max_pages) return error.PastCapacity; + const p = s.pool; + if (p.world > 1 and pg % p.world != s.len % p.world) return error.ResidueMismatch; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[s.len] = pg; + s.len += 1; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(s.len - 1, s.len); + } + + /// Lets go of the pages past position keep (rows nothing reads again); returns how many. + pub fn truncate(s: *Slot, keep: u64) Error!u32 { + const p = s.pool; + // positions >= keep get written again: a page holding one no longer equals the origin file's + if (s.origin) |*o| o.pages = @min(o.pages, @as(u32, @intCast(keep / p.page))); + const k = pagesFor(keep, p.page); + if (k >= s.len) return 0; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + const n = s.len - k; + for (s.pages[k..s.len]) |pg| try p.release(pg, 1); + const old = s.len; + s.len = k; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(k, old); + return n; + } + + /// Sets the admission's reservation (null: none). + pub fn reserve(s: *Slot, quota: ?u32) void { + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + s.quota = quota; + s.pool.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + s.pool.addReserved(&after); + } + + /// Frees the slot's pages and its reservation. + pub fn releaseAll(s: *Slot) Error!u32 { + s.reserve(null); + s.over = 0; + s.origin = null; + return s.truncate(0); + } +}; + +/// The pool's book: free sets by residue, holders a page, the slots and their reservations. +pub const Pool = struct { + gpa: std.mem.Allocator, + layout: Layout, + npages: u32, + page: u32, + /// log2 of page tokens + shift: u5, + /// ranks the split families are sharded over (1: every rank stores every row) + world: u32, + rank: u32, + /// the page unmapped table entries point at (all zeros: a stray read reads zeros) + null_page: u32, + /// a split family's local page count (npages / W), its local null page and its discard page (other ranks' rows) + local_pages: u32, + local_null: u32, + discard: u32, + free: [max_world]FreeSet = @splat(.{}), + /// pages each residue's slots reserved and have not mapped yet + reserved: [max_world]u32 = @splat(0), + /// holders a page (slots + entries + in-flight jobs) and how many are slots + refs: []u16, + slot_refs: []u16, + /// pages held by something and by no slot (session entries alone): what parking can free + entry_only: u32 = 0, + slots: std.ArrayList(*Slot) = .empty, + + pub fn init(gpa: std.mem.Allocator, layout: Layout, tokens: u64, world: u32, rank: u32) Error!Pool { + if (layout.page == 0 or layout.page & (layout.page - 1) != 0) return error.BadPool; + if (tokens % layout.page != 0 or tokens == 0) return error.BadPool; + if (world == 0 or world > max_world or rank >= world) return error.BadPool; + const npages: u32 = @intCast(tokens / layout.page); + if (npages % world != 0) return error.BadPool; + for (layout.families) |f| if (f.ratio == 0 or layout.page % f.ratio != 0) return error.BadPool; + var p: Pool = .{ + .gpa = gpa, + .layout = layout, + .npages = npages, + .page = layout.page, + .shift = @intCast(std.math.log2_int(u32, layout.page)), + .world = world, + .rank = rank, + .null_page = npages, + .local_pages = npages / world, + .local_null = npages / world, + .discard = npages / world + 1, + .refs = &.{}, + .slot_refs = &.{}, + }; + errdefer p.deinit(); + for (0..world) |r| p.free[r] = try FreeSet.init(gpa, npages / world); + p.refs = try gpa.alloc(u16, npages); + @memset(p.refs, 0); + p.slot_refs = try gpa.alloc(u16, npages); + @memset(p.slot_refs, 0); + return p; + } + + pub fn deinit(p: *Pool) void { + for (p.slots.items) |s| { + p.gpa.free(s.pages); + p.gpa.destroy(s); + } + p.slots.deinit(p.gpa); + for (0..p.world) |r| p.free[r].deinit(p.gpa); + p.gpa.free(p.refs); + p.gpa.free(p.slot_refs); + p.* = undefined; + } + + pub fn newSlot(p: *Pool, capacity: u64) Error!*Slot { + const s = try p.gpa.create(Slot); + errdefer p.gpa.destroy(s); + const mp = pagesFor(capacity, p.page); + s.* = .{ .pool = p, .id = @intCast(p.slots.items.len), .capacity = capacity, .max_pages = mp, .pages = try p.gpa.alloc(u32, mp) }; + errdefer p.gpa.free(s.pages); + try p.slots.append(p.gpa, s); + return s; + } + + /// A free physical page for logical page k (its residue when split), held once by the caller. + pub fn takeAt(p: *Pool, k: u32) Error!u32 { + const r = k % p.world; + const local = p.free[r].takeLowest() orelse return error.PoolExhausted; + const pg = local * p.world + r; + p.hold(pg, 1, 0); // the caller's hold (an entry-like holder until a slot maps it) + return pg; + } + + /// Holders change on page pg: dref total, dslot of them slots; keeps entry_only and gives the page back at zero. + fn hold(p: *Pool, pg: u32, dref: i32, dslot: i32) void { + const r0 = p.refs[pg]; + const s0 = p.slot_refs[pg]; + const r1: u16 = @intCast(@as(i32, r0) + dref); + const s1: u16 = @intCast(@as(i32, s0) + dslot); + p.refs[pg] = r1; + p.slot_refs[pg] = s1; + const was = r0 > 0 and s0 == 0; + const now = r1 > 0 and s1 == 0; + if (was and !now) p.entry_only -= 1; + if (!was and now) p.entry_only += 1; + } + + /// One holder lets go (slot: from a slot's table); the page returns to its free set when nobody holds it. + fn release(p: *Pool, pg: u32, slot: u1) Error!void { + if (p.refs[pg] == 0) return error.DoubleFree; + p.hold(pg, -1, -@as(i32, slot)); + if (p.refs[pg] == 0) p.free[pg % p.world].give(pg / p.world) catch return error.DoubleFree; + } + + /// One more holder (a session entry or an in-flight job) of each page. + pub fn share(p: *Pool, pages: []const u32) void { + for (pages) |pg| p.hold(pg, 1, 0); + } + + /// Those holders let go. + pub fn drop(p: *Pool, pages: []const u32) Error!void { + for (pages) |pg| try p.release(pg, 0); + } + + pub fn holders(p: *const Pool, pg: u32) u16 { + return p.refs[pg]; + } + + fn unreserve(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] -= by[r]; + } + + fn addReserved(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] += by[r]; + } + + /// Free pages no admission reserved (split: W x the scarcer residue, so any logical run fits). + pub fn available(p: *const Pool) u32 { + var low: u32 = std.math.maxInt(u32); + for (0..p.world) |r| low = @min(low, p.free[r].count() -| p.reserved[r]); + return low * p.world; + } + + pub fn freePages(p: *const Pool) u32 { + var n: u32 = 0; + for (0..p.world) |r| n += p.free[r].count(); + return n; + } + + /// Pages an admission needs for a request (prompt + reply + slack, at most the slot's capacity). + pub fn needPages(p: *const Pool, prompt: u64, max_tokens: u64, capacity: u64, slack: u64) u32 { + return pagesFor(@min(prompt + max_tokens + slack, capacity), p.page); + } + + /// This rank stores the split families' rows of page pg. + pub fn owns(p: *const Pool, pg: u32) bool { + return pg != p.null_page and pg % p.world == p.rank; + } + + /// A split family's local page of pg: its own when owned, the local null for the null page, else the discard page. + pub fn localPage(p: *const Pool, pg: u32) u32 { + if (pg == p.null_page) return p.local_null; + return if (pg % p.world == p.rank) pg / p.world else p.discard; + } + + /// The page a family's tensor indexes for physical page pg (split families: localPage). + pub fn familyPage(p: *const Pool, f: Family, pg: u32) u32 { + return if (f.split and p.world > 1) p.localPage(pg) else pg; + } + + /// Pages a family's device tensor holds: npages + null, or (split) local pages + local null + discard page. + pub fn familyTensorPages(p: *const Pool, f: Family) u32 { + return if (f.split and p.world > 1) p.local_pages + 2 else p.npages + 1; + } + + /// Bytes this rank stores of family f for logical pages [first, first + n) of a slot or entry. + pub fn familyBytes(p: *const Pool, f: Family, first: u32, n: u32) u64 { + if (!(f.split and p.world > 1)) return @as(u64, n) * f.pageBytes(p.page); + var by: [max_world]u32 = undefined; + span(first, first + n, p.world, &by); + return @as(u64, by[p.rank]) * f.pageBytes(p.page); + } + + /// Bytes of every family's device tensor on this rank. + pub fn tensorBytes(p: *const Pool) u64 { + var b: u64 = 0; + for (p.layout.families) |f| b += @as(u64, p.familyTensorPages(f)) * f.pageBytes(p.page); + return b; + } + + /// Pages held by session entries alone and their bytes on this rank (the residency budget's measure): O(1). + pub fn entryOnlyBytes(p: *const Pool) u64 { + return @as(u64, p.entry_only) * p.layout.pageBytes(p.world); + } +}; + +test "residues, holders, quotas and the device tables" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 2, .row_bytes = 584, .split = true }, .{ .name = "index_k", .ratio = 2, .row_bytes = 132 } }; + var p = try Pool.init(gpa, .{ .families = &fams }, 16 * 256, 2, 1); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 8 + 2), p.familyTensorPages(fams[0])); + try std.testing.expectEqual(@as(u32, 16 + 1), p.familyTensorPages(fams[1])); + try std.testing.expectApproxEqAbs(@as(f64, 146 + 66), p.layout.bytesPerToken(2), 1e-9); + const a = try p.newSlot(8 * 256); + const b = try p.newSlot(8 * 256); + try a.ensure(3 * 256 + 1); // 4 pages: residues 0,1,0,1 + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3 }, a.mapped()); + try b.ensure(256); + try std.testing.expectEqualSlices(u32, &.{4}, b.mapped()); + try std.testing.expectEqual(p.discard, a.localTableAt(0)); // page 0 is rank 0's + try std.testing.expectEqual(@as(u32, 0), a.localTableAt(1)); // page 1 -> rank 1's local page 0 + try std.testing.expectEqual(p.local_null, a.localTableAt(7)); + const d = a.takeDirty().?; + try std.testing.expectEqual(@as(u32, 0), d.lo); + try std.testing.expectEqual(@as(u32, 4), d.hi); + try std.testing.expect(a.takeDirty() == null); + // an entry shares a's pages; a lets go; the pages stay held by the entry alone + p.share(a.mapped()[0..3]); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + _ = try a.releaseAll(); + try std.testing.expectEqual(@as(u32, 3), p.entry_only); + try std.testing.expectEqual(@as(u32, 16 - 1 - 3), p.freePages()); + try p.drop(&.{ 0, 1, 2 }); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + try std.testing.expectError(error.DoubleFree, p.drop(&.{0})); + // quotas: b reserves 6 pages; available counts W x the scarcer residue + b.reserve(6); + var by: [2]u32 = undefined; + span(1, 6, 2, &by); + try std.testing.expectEqual(@as(u32, 2), by[0]); + try std.testing.expectEqual(@as(u32, 3), by[1]); + try std.testing.expectEqual(@as(u32, 2 * @min(7 - 2, 8 - 3)), p.available()); + try b.ensure(6 * 256); + try std.testing.expectEqual(@as(u32, 0), p.reserved[0] + p.reserved[1]); + for (b.mapped(), 0..) |pg, k| try std.testing.expectEqual(k % 2, pg % 2); + // RAM resume: adopt full pages of the right residues, refuse the wrong ones + const c = try p.newSlot(8 * 256); + try std.testing.expectError(error.ResidueMismatch, c.adopt(&.{b.pages[1]})); + p.share(b.mapped()[0..2]); + try c.adopt(b.mapped()[0..2]); + try c.push(try p.takeAt(2)); + try std.testing.expectEqual(@as(u16, 3), p.holders(b.pages[0])); + try std.testing.expectEqual(@as(u32, 0), c.pages[2] % 2); +} diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig new file mode 100644 index 0000000000..98634bbfd2 --- /dev/null +++ b/zig/src/sessions/prefix.zig @@ -0,0 +1,486 @@ +//! Session entries by prefix: a trie over chained page digests, one node per distinct run of full pages a tag. + +const std = @import("std"); + +pub const Digest = [16]u8; +pub const none = std.math.maxInt(u32); +const Blake = std.crypto.hash.blake2.Blake2b128; + +/// Chained digests of the full pages of ids under a tag: out[j] determines tag and ids[0 .. page (j + 1)]. +pub fn chain(tag: u32, ids: []const i32, page: u32, out: []Digest) void { + var prev = root(tag); + for (out, 0..) |*o, j| { + o.* = step(prev, ids[j * page ..][0..page]); + prev = o.*; + } +} + +/// The digest of a tag's empty prefix. +pub fn root(tag: u32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-root"); + h.update(std.mem.asBytes(&tag)); + return h.finalResult(); +} + +/// The next page's digest. +pub fn step(prev: Digest, page_ids: []const i32) Digest { + var h = Blake.init(.{}); + h.update(&prev); + h.update(std.mem.sliceAsBytes(page_ids)); + return h.finalResult(); +} + +/// An entry's identity: its node's digest (tag and full pages), its length and tail. +pub fn entryKey(node: Digest, pos: u64, tail: []const i32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-entry"); + h.update(&node); + h.update(std.mem.asBytes(&pos)); + h.update(std.mem.sliceAsBytes(tail)); + return h.finalResult(); +} + +pub const Tier = enum { ram, any }; + +pub fn Index(comptime Extra: type) type { + return struct { + const Self = @This(); + + pub const Node = struct { + digest: Digest, + parent: u32, + depth: u32, + first_child: u32 = none, + next_sib: u32 = none, + prev_sib: u32 = none, + first_entry: u32 = none, + /// entries in the subtree (this node's included): all, and those in RAM + sub_all: u32 = 0, + sub_ram: u32 = 0, + /// RAM entries at this node itself + own_ram: u32 = 0, + /// the tokens of the page leading into this node, while sub_ram > 0 + tokens: ?[]i32 = null, + }; + + pub const Entry = struct { + key: Digest, + tag: u32, + pos: u64, + node: u32, + tail: []i32, + ram: bool, + next_at: u32 = none, + prev_at: u32 = none, + live: bool = true, + extra: Extra, + }; + + /// An entry inserted, or the one already there with the same key. + pub const Placed = struct { id: u32, new: bool }; + + gpa: std.mem.Allocator, + page: u32, + nodes: std.ArrayList(Node) = .empty, + free_nodes: std.ArrayList(u32) = .empty, + by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + dead: u32 = 0, // nodes dropped from `by_digest` since it was last rehashed + entries: std.ArrayList(Entry) = .empty, + free_entries: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// page-sized token blocks, reused + blocks: std.ArrayList([]i32) = .empty, + + pub fn init(gpa: std.mem.Allocator, page: u32) Self { + return .{ .gpa = gpa, .page = page }; + } + + pub fn deinit(x: *Self) void { + for (x.nodes.items) |n| if (n.tokens) |t| x.gpa.free(t); + for (x.blocks.items) |b| x.gpa.free(b); + for (x.entries.items) |e| if (e.live) x.gpa.free(e.tail); + x.nodes.deinit(x.gpa); + x.free_nodes.deinit(x.gpa); + x.by_digest.deinit(x.gpa); + x.entries.deinit(x.gpa); + x.free_entries.deinit(x.gpa); + x.by_key.deinit(x.gpa); + x.blocks.deinit(x.gpa); + x.* = undefined; + } + + pub fn get(x: *Self, e: u32) *Entry { + return &x.entries.items[e]; + } + + pub fn node(x: *Self, n: u32) *Node { + return &x.nodes.items[n]; + } + + pub fn lookup(x: *const Self, key: Digest) ?u32 { + return x.by_key.get(key); + } + + fn newNode(x: *Self, digest: Digest, parent: u32, depth: u32) !u32 { + const id: u32 = if (x.free_nodes.pop()) |i| i else blk: { + try x.nodes.append(x.gpa, undefined); + break :blk @intCast(x.nodes.items.len - 1); + }; + x.nodes.items[id] = .{ .digest = digest, .parent = parent, .depth = depth }; + try x.by_digest.put(x.gpa, digest, id); + if (parent != none) { + const p = &x.nodes.items[parent]; + x.nodes.items[id].next_sib = p.first_child; + if (p.first_child != none) x.nodes.items[p.first_child].prev_sib = id; + p.first_child = id; + } + return id; + } + + fn takeBlock(x: *Self) ![]i32 { + if (x.blocks.pop()) |b| return b; + return x.gpa.alloc(i32, x.page); + } + + /// Adds an entry for ids (chain: its full pages' digests), or returns the one with the same key. + pub fn insert(x: *Self, tag: u32, ids: []const i32, digests: []const Digest, ram: bool, extra: Extra) !Placed { + const full: u32 = @intCast(ids.len / x.page); + std.debug.assert(digests.len >= full); + const r = root(tag); + var n = x.by_digest.get(r) orelse try x.newNode(r, none, 0); + for (0..full) |j| { + const d = digests[j]; + n = x.by_digest.get(d) orelse try x.newNode(d, n, @intCast(j + 1)); + } + const r2 = try x.place(tag, n, ids.len, ids[full * x.page ..], extra); + if (r2.new and ram) try x.setRam(r2.id, true, ids); + return r2; + } + + /// Adds a disk entry continuing from node anchor (none: the tag's root); suffix: tokens [depth x page, pos). + pub fn insertFrom(x: *Self, tag: u32, anchor: u32, suffix: []const i32, extra: Extra) !Placed { + var n = if (anchor != none) anchor else blk: { + const r = root(tag); + break :blk x.by_digest.get(r) orelse try x.newNode(r, none, 0); + }; + const d0 = x.nodes.items[n].depth; + const full: u32 = @intCast(suffix.len / x.page); + for (0..full) |j| { + const d = step(x.nodes.items[n].digest, suffix[j * x.page ..][0..x.page]); + n = x.by_digest.get(d) orelse try x.newNode(d, n, d0 + @as(u32, @intCast(j)) + 1); + } + return x.place(tag, n, @as(u64, d0) * x.page + suffix.len, suffix[full * x.page ..], extra); + } + + /// The node at depth on an entry's path (its first `depth` full pages). + pub fn nodeAt(x: *Self, id: u32, depth: u32) ?u32 { + var n = x.entries.items[id].node; + while (n != none and x.nodes.items[n].depth > depth) n = x.nodes.items[n].parent; + return if (n != none and x.nodes.items[n].depth == depth) n else null; + } + + fn place(x: *Self, tag: u32, n: u32, pos: u64, tail_ids: []const i32, extra: Extra) !Placed { + const key = entryKey(x.nodes.items[n].digest, pos, tail_ids); + if (x.by_key.get(key)) |e| return .{ .id = e, .new = false }; + const tail = try x.gpa.dupe(i32, tail_ids); + errdefer x.gpa.free(tail); + const id: u32 = if (x.free_entries.pop()) |i| i else blk: { + try x.entries.append(x.gpa, undefined); + break :blk @intCast(x.entries.items.len - 1); + }; + try x.by_key.put(x.gpa, key, id); + const nd = &x.nodes.items[n]; + x.entries.items[id] = .{ .key = key, .tag = tag, .pos = pos, .node = n, .tail = tail, .ram = false, .next_at = nd.first_entry, .extra = extra }; + if (nd.first_entry != none) x.entries.items[nd.first_entry].prev_at = id; + nd.first_entry = id; + x.addPath(n, 1, 0); + return .{ .id = id, .new = true }; + } + + fn addPath(x: *Self, from: u32, d_all: i32, d_ram: i32) void { + var n = from; + while (n != none) { + const nd = &x.nodes.items[n]; + nd.sub_all = @intCast(@as(i64, nd.sub_all) + d_all); + nd.sub_ram = @intCast(@as(i64, nd.sub_ram) + d_ram); + n = nd.parent; + } + } + + /// Moves an entry in or out of RAM; going in needs its ids (to keep its path's tokens). + pub fn setRam(x: *Self, id: u32, ram: bool, ids: ?[]const i32) !void { + const e = &x.entries.items[id]; + if (e.ram == ram) return; + if (ram) { + const all = ids orelse return error.IdsNeeded; + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.tokens != null) continue; + const b = try x.takeBlock(); + @memcpy(b, all[(nd.depth - 1) * x.page ..][0..x.page]); + nd.tokens = b; + } + } + e.ram = ram; + x.nodes.items[e.node].own_ram = if (ram) x.nodes.items[e.node].own_ram + 1 else x.nodes.items[e.node].own_ram - 1; + x.addPath(e.node, 0, if (ram) 1 else -1); + if (!ram) x.dropTokens(e.node); + } + + fn dropTokens(x: *Self, from: u32) void { + var n = from; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_ram > 0) return; + if (nd.tokens) |t| { + x.blocks.append(x.gpa, t) catch x.gpa.free(t); + nd.tokens = null; + } + } + } + + /// Removes an entry and the nodes no entry needs any more. + pub fn remove(x: *Self, id: u32) void { + x.setRam(id, false, null) catch unreachable; + const e = &x.entries.items[id]; + const n = e.node; + if (e.prev_at != none) x.entries.items[e.prev_at].next_at = e.next_at else x.nodes.items[n].first_entry = e.next_at; + if (e.next_at != none) x.entries.items[e.next_at].prev_at = e.prev_at; + _ = x.by_key.remove(e.key); + x.gpa.free(e.tail); + e.live = false; + x.free_entries.append(x.gpa, id) catch {}; + x.addPath(n, -1, 0); + var m = n; + while (m != none and x.nodes.items[m].sub_all == 0) { + const nd = x.nodes.items[m]; + if (nd.parent != none) { + if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; + if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; + } + if (x.by_digest.remove(nd.digest)) x.dropped(); + x.free_nodes.append(x.gpa, m) catch {}; + m = nd.parent; + } + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn dropped(x: *Self) void { + x.dead += 1; + if (@as(u64, x.dead) * 10 < x.by_digest.capacity()) return; + x.by_digest.rehash(std.hash_map.AutoContext(Digest){}); + x.dead = 0; + } + + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). + fn tailOf(e: *const Entry, prompt: []const i32) bool { + const at: usize = @intCast(e.pos - e.tail.len); + return e.pos < prompt.len and std.mem.eql(i32, e.tail, prompt[at..][0..e.tail.len]); + } + + /// The longest entry of tag whose ids strictly prefix prompt (RAM first on a tie); digests: the prompt's chain. + pub fn find(x: *Self, tag: u32, prompt: []const i32, digests: []const Digest, tier: Tier) ?u32 { + const full: u32 = @intCast(prompt.len / x.page); + const r = x.by_digest.get(root(tag)) orelse return null; + // the deepest node on the prompt's path: each prefix of an entry's path has a node: binary search + var lo: u32 = 0; + var hi: u32 = full; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (x.by_digest.contains(digests[mid - 1])) lo = mid else hi = mid - 1; + } + var n: u32 = if (lo == 0) r else x.by_digest.get(digests[lo - 1]).?; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_all == 0 or (tier == .ram and nd.own_ram == 0)) continue; + var best: u32 = none; + var e = nd.first_entry; + while (e != none) : (e = x.entries.items[e].next_at) { + const en = &x.entries.items[e]; + if (tier == .ram and !en.ram) continue; + if (!tailOf(en, prompt)) continue; + if (best == none or en.pos > x.entries.items[best].pos or (en.pos == x.entries.items[best].pos and en.ram)) best = e; + } + if (best != none) return best; + } + return null; + } + + /// A newer entry of id's chain holds every position it does (tier .ram: in RAM), so id need not be written. + pub fn covered(x: *Self, id: u32, tier: Tier) bool { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == id or (tier == .ram and !o.ram)) continue; + if (o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) return true; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + const n = if (tier == .ram) ch.sub_ram else ch.sub_all; + if (n == 0) continue; + if (e.tail.len == 0) return true; + if (ch.tokens) |t| if (std.mem.startsWith(i32, t, e.tail)) return true; + } + return false; + } + + /// How many RAM entries extend a RAM entry (longer, same tag, its ids their prefix): branch points have more. + pub fn extensions(x: *Self, id: u32) u32 { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var n: u32 = 0; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f != id and o.ram and o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) n += 1; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + if (ch.sub_ram == 0) continue; + if (std.mem.startsWith(i32, ch.tokens.?, e.tail)) n += ch.sub_ram; + } + return n; + } + + /// A RAM entry no other RAM entry extends: a chat's newest state. + pub fn isLeaf(x: *Self, id: u32) bool { + return x.entries.items[id].ram and !x.covered(id, .ram); + } + + /// A RAM leaf's older turns (the RAM entries it extends) into out; returns the chat's recency (newest `used`). + pub fn members(x: *Self, leaf: u32, out: *std.ArrayList(u32), used: *const fn (*const Extra) u64) !u64 { + const l = &x.entries.items[leaf]; + var recency = used(&l.extra); + var below: u32 = none; // the path's node under n: its tokens are the page an entry at n must prefix + var n = l.node; + while (n != none) : ({ + below = n; + n = x.nodes.items[n].parent; + }) { + const nd = &x.nodes.items[n]; + if (nd.own_ram == 0) continue; + const page_ids: []const i32 = if (below == none) l.tail else x.nodes.items[below].tokens.?; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == leaf or !o.ram or o.pos >= l.pos) continue; + if (!std.mem.startsWith(i32, page_ids, o.tail)) continue; + try out.append(x.gpa, f); + recency = @max(recency, used(&o.extra)); + } + } + return recency; + } + + /// A RAM entry's ids, rebuilt from its path's pages and tail, into out (len = pos). + pub fn idsOf(x: *Self, id: u32, out: []i32) !void { + const e = &x.entries.items[id]; + if (out.len != e.pos) return error.BadLength; + const at = e.pos - e.tail.len; + @memcpy(out[at..], e.tail); + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + const t = nd.tokens orelse return error.NotInRam; + @memcpy(out[(nd.depth - 1) * x.page ..][0..x.page], t); + } + } + }; +} + +const T = Index(u64); + +fn usedOf(e: *const u64) u64 { + return e.*; +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + chain(tag, ids, x.page, d[0 .. ids.len / x.page]); + return (try x.insert(tag, ids, d[0 .. ids.len / x.page], ram, used)).id; +} + +test "longest strict prefix, ties, tags and coverage" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + const chat = try seq(gpa, 200, 0); + defer gpa.free(chat); + const t1 = try add(&x, 7, chat[0..40], true, 1); // a prompt's replay point + const t2 = try add(&x, 7, chat[0..72], true, 2); // a turn's end + const t3 = try add(&x, 7, chat[0..128], false, 3); // on disk only, page aligned + _ = try add(&x, 8, chat[0..100], true, 4); // another tag + var d: [16]Digest = undefined; + chain(7, chat[0..40], 16, d[0..2]); + const again = try x.insert(7, chat[0..40], d[0..2], true, 9); + try std.testing.expect(again.id == t1 and !again.new); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t3), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..150], d[0..9], .ram)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..128], d[0..8], .any)); // strict: 128 is not a prefix of 128 + chain(7, chat[0..72], 16, d[0..4]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, chat[0..72], d[0..4], .any)); + var other = try gpa.dupe(i32, chat[0..150]); + defer gpa.free(other); + other[50] = -1; // diverges inside t2's tail page: t1 still prefixes it + chain(7, other, 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, other, d[0..9], .any)); + // t3 extends t2 on disk only: its page's tokens are gone, so a tail there is undecided (the store marks covered_by) + try std.testing.expect(x.covered(t1, .ram) and !x.covered(t2, .any) and !x.covered(t2, .ram) and !x.covered(t3, .any)); + const t4 = try add(&x, 7, chat[0..64], false, 5); // page aligned: any deeper entry covers it + try std.testing.expect(x.covered(t4, .any) and x.covered(t4, .ram)); + x.remove(t4); + try std.testing.expect(x.isLeaf(t2) and !x.isLeaf(t1)); + try std.testing.expectEqual(@as(u32, 1), x.extensions(t1)); + try std.testing.expectEqual(@as(u32, 0), x.extensions(t2)); + var m: std.ArrayList(u32) = .empty; + defer m.deinit(gpa); + try std.testing.expectEqual(@as(u64, 2), try x.members(t2, &m, usedOf)); + try std.testing.expectEqualSlices(u32, &.{t1}, m.items); + const ids = try gpa.alloc(i32, 72); + defer gpa.free(ids); + try x.idsOf(t2, ids); + try std.testing.expectEqualSlices(i32, chat[0..72], ids); + // going to disk drops the path's tokens once no RAM entry needs them; removal prunes the nodes + try x.setRam(t2, false, null); + try x.setRam(t1, false, null); + try std.testing.expectError(error.NotInRam, x.idsOf(t2, ids)); + x.remove(t3); + x.remove(t2); + x.remove(t1); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages +} + +test "a long run of turns saved and dropped leaves the page map free slots" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + var live: std.ArrayList(u32) = .empty; + defer live.deinit(gpa); + var rng = std.Random.DefaultPrng.init(5); + for (0..3000) |i| { + // a chat's turn (its own pages), the oldest dropped once eight are kept: nodes come and go + const ids = try seq(gpa, 16 * 60, @intCast(i)); + defer gpa.free(ids); + try live.append(gpa, try add(&x, 7, ids[0 .. 16 * (1 + rng.random().uintLessThan(usize, 60))], true, i)); + if (live.items.len > 8) x.remove(live.orderedRemove(0)); + const meta = x.by_digest.metadata orelse continue; + var free: usize = 0; + for (meta[0..x.by_digest.capacity()]) |m| free += @intFromBool(m.isFree()); + try std.testing.expect(free * 10 >= x.by_digest.capacity()); + } +} diff --git a/zig/src/sessions/sessions.zig b/zig/src/sessions/sessions.zig new file mode 100644 index 0000000000..7bdc46b386 --- /dev/null +++ b/zig/src/sessions/sessions.zig @@ -0,0 +1,33 @@ +//! Session memory, family-neutral: the KV pool book, page bytes, prefix entries, the RAM tier and the NVMe tier. + +pub const freeset = @import("freeset.zig"); +pub const pool = @import("pool.zig"); +pub const Pool = pool.Pool; +pub const Slot = pool.Slot; +pub const Family = pool.Family; +pub const Layout = pool.Layout; +pub const pagestore = @import("pagestore.zig"); +pub const PageStore = pagestore.PageStore; +pub const HostPool = pagestore.HostPool; +pub const prefix = @import("prefix.zig"); +pub const lanes = @import("lanes.zig"); +pub const format = @import("format.zig"); +pub const disk = @import("disk.zig"); +pub const Disk = disk.Disk; +pub const jobs = @import("jobs.zig"); +pub const store = @import("store.zig"); +pub const Store = store.Store; + +test { + _ = freeset; + _ = pool; + _ = pagestore; + _ = prefix; + _ = lanes; + _ = format; + _ = disk; + _ = jobs; + _ = store; + _ = @import("disk_test.zig"); + _ = @import("store_test.zig"); +} diff --git a/zig/src/sessions/store.zig b/zig/src/sessions/store.zig new file mode 100644 index 0000000000..851ce509aa --- /dev/null +++ b/zig/src/sessions/store.zig @@ -0,0 +1,570 @@ +//! The session store: RAM entries holding pool pages and bounded state, a residency budget, and the NVMe tier. + +const std = @import("std"); +const Io = std.Io; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const PageStore = @import("pagestore.zig").PageStore; +const prefix = @import("prefix.zig"); +const Digest = prefix.Digest; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; +const jobs = @import("jobs.zig"); +const Job = jobs.Job; + +pub const none = prefix.none; + +/// A slot's bounded state (rings, carries, lookback, taps) as the family snapshots it: opaque to the store. +pub const Bounded = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + nbytes: *const fn (ptr: *anyopaque) u64, + /// host copies of its pieces for a park (owned by the caller: name and bytes allocated with gpa) + blobs: *const fn (ptr: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob, + release: *const fn (ptr: *anyopaque) void, + }; +}; + +pub const Policy = struct { + /// bounded state the RAM tier keeps (least recently used entries go past it) + ram_bytes: u64 = 512 << 20, + /// idle chats resident in the pool (0: no limit) and pool bytes held by entries alone (0: no limit) + chats: u32 = 0, + resident_bytes: u64 = 0, + /// park an evicted chat's older turns too (default: drop them, the leaf's file covers them) + park_prefixes: bool = false, + /// drop, never write, an entry a newer entry of its chain covers + skip_covered: bool = true, + /// write a newer turn as a delta on the file its leading pages came from + delta: bool = true, + /// parks and restores on the tier thread (settled by the plan) + async_io: bool = false, +}; + +pub const Extra = struct { + /// RAM: the pool pages it holds, logical [0, npages) + pages: []u32 = &.{}, + bounded: ?Bounded = null, + /// a file with this key is registered on the NVMe tier + file: bool = false, + used: u64 = 0, + origin: ?pool_mod.Origin = null, + /// a newer entry of its chain was parked from RAM: writing this one again is redundant + covered_by: bool = false, + kind: u32 = 0, + /// an unsettled park of it + job: ?*Job = null, + older: u32 = none, + newer: u32 = none, +}; + +const Index = prefix.Index(Extra); + +pub const Hit = struct { id: u32, ram: bool, pos: u64 }; + +pub const Stats = struct { + saves: u64 = 0, + dups: u64 = 0, + hits_ram: u64 = 0, + hits_disk: u64 = 0, + parks: u64 = 0, + delta_parks: u64 = 0, + drops: u64 = 0, + covered_drops: u64 = 0, + chat_evictions: u64 = 0, + park_ns: u64 = 0, + restore_ns: u64 = 0, +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: Io, + pool: *Pool, + pages: PageStore, + index: Index, + disk: ?*Disk, + /// heap-held: the tier thread keeps its address + runner: ?*jobs.Runner = null, + policy: Policy, + clock: u64 = 0, + bounded_total: u64 = 0, + /// RAM entries, least recently used first + oldest: u32 = none, + newest: u32 = none, + ram_count: u32 = 0, + pending: std.ArrayList(*Job) = .empty, + scratch: std.ArrayList(Digest) = .empty, + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, io: Io, pool: *Pool, pages: PageStore, d: ?*Disk, policy: Policy) !Store { + var s: Store = .{ .gpa = gpa, .io = io, .pool = pool, .pages = pages, .index = Index.init(gpa, pool.page), .disk = d, .policy = policy }; + if (d) |dd| { + const r = try gpa.create(jobs.Runner); + r.* = jobs.Runner.init(gpa, io, dd); + s.runner = r; + if (policy.async_io) r.startThread() catch |e| { + gpa.destroy(r); + return e; + }; + } + return s; + } + + pub fn deinit(s: *Store) void { + if (s.runner) |r| { + for (s.pending.items) |j| r.wait(j); + r.deinit(); + s.gpa.destroy(r); + } + for (s.pending.items) |j| j.deinit(s.gpa); + s.pending.deinit(s.gpa); + for (s.index.entries.items) |*e| if (e.live) { + if (e.extra.bounded) |b| b.vtable.release(b.ptr); + s.pool.drop(e.extra.pages) catch {}; + s.gpa.free(e.extra.pages); + }; + s.index.deinit(); + s.scratch.deinit(s.gpa); + s.* = undefined; + } + + fn digests(s: *Store, tag: u32, ids: []const i32) ![]const Digest { + const n = ids.len / s.pool.page; + try s.scratch.resize(s.gpa, n); + prefix.chain(tag, ids, s.pool.page, s.scratch.items); + return s.scratch.items; + } + + pub fn entry(s: *Store, id: u32) *Index.Entry { + return s.index.get(id); + } + + /// The longest stored entry (RAM or NVMe) whose ids strictly prefix prompt, with the same tag. + pub fn find(s: *Store, tag: u32, prompt: []const i32) !?Hit { + const d = try s.digests(tag, prompt); + const id = s.index.find(tag, prompt, d, .any) orelse return null; + const e = s.index.get(id); + return .{ .id = id, .ram = e.ram, .pos = e.pos }; + } + + fn link(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + x.older = s.newest; + x.newer = none; + if (s.newest != none) s.index.get(s.newest).extra.newer = id else s.oldest = id; + s.newest = id; + } + + fn unlink(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + if (x.older != none) s.index.get(x.older).extra.newer = x.newer else s.oldest = x.newer; + if (x.newer != none) s.index.get(x.newer).extra.older = x.older else s.newest = x.older; + } + + fn touch(s: *Store, id: u32) void { + s.clock += 1; + const e = s.index.get(id); + e.extra.used = s.clock; + if (e.ram) { + s.unlink(id); + s.link(id); + } + } + + /// An entry for a slot's state at len(ids), sharing its pages; null for a duplicate (bounded is released then). + pub fn save(s: *Store, slot: *Slot, bounded: Bounded, ids: []const i32, tag: u32, kind: u32) !?u32 { + if (ids.len == 0) return error.EmptyEntry; + const need = pool_mod.pagesFor(ids.len, s.pool.page); + if (slot.len < need) return error.SlotTooShort; + const d = try s.digests(tag, ids); + const ins = try s.index.insert(tag, ids, d, false, .{ .kind = kind }); + errdefer if (ins.new) s.index.remove(ins.id); + const e = s.index.get(ins.id); + if (e.ram) { + bounded.vtable.release(bounded.ptr); + s.touch(ins.id); + s.stats.dups += 1; + return null; + } + const pages = try s.gpa.dupe(u32, slot.mapped()[0..need]); + errdefer s.gpa.free(pages); + try s.index.setRam(ins.id, true, ids); + s.pool.share(pages); + const x = &s.index.get(ins.id).extra; + x.pages = pages; + x.bounded = bounded; + x.kind = kind; + if (slot.origin) |o| x.origin = .{ .key = o.key, .pages = @min(o.pages, @as(u32, @intCast(ids.len / s.pool.page))) }; + s.bounded_total += bounded.vtable.nbytes(bounded.ptr); + s.ram_count += 1; + s.link(ins.id); + s.touch(ins.id); + s.stats.saves += 1; + try s.trimRam(); + _ = try s.enforce(); + return ins.id; + } + + /// A RAM entry into an empty slot: full pages shared, the partial last page copied; returns its bounded state. + pub fn restoreRam(s: *Store, id: u32, slot: *Slot) !Bounded { + const e = s.index.get(id); + if (!e.ram) return error.NotInRam; + const full: u32 = @intCast(e.pos / s.pool.page); + try slot.adopt(e.extra.pages[0..full]); + errdefer _ = slot.releaseAll() catch {}; + if (e.pos % s.pool.page != 0) { + const src = e.extra.pages[full]; + const dst = try s.pool.takeAt(full); + for (s.pool.layout.families, 0..) |f, fi| { + if (f.split and !s.pool.owns(src)) continue; + try s.pages.copyPage(@intCast(fi), s.pool.familyPage(f, src), s.pool.familyPage(f, dst)); + } + try slot.push(dst); + } + if (e.extra.origin) |o| slot.origin = .{ .key = o.key, .pages = @min(o.pages, full) }; + s.touch(id); + s.stats.hits_ram += 1; + return e.extra.bounded.?; + } + + /// Starts an NVMe resume into an empty slot (pages mapped now, filled by the job); settle it before the slot runs. + pub fn beginRestore(s: *Store, id: u32, slot: *Slot) !*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.extra.file) return error.NotOnDisk; + if (slot.len != 0) return error.SlotNotEmpty; + try slot.ensure(e.pos); + const j = try s.gpa.create(Job); + j.* = .{ .kind = .restore, .entry = id, .slot = slot.id }; + errdefer j.deinit(s.gpa); + var ext: std.ArrayList(disk_mod.Extent) = .empty; + errdefer ext.deinit(s.gpa); + try d.extents(e.key, &ext); + j.extents = try ext.toOwnedSlice(s.gpa); + const np = pool_mod.pagesFor(e.pos, s.pool.page); + j.lists = try s.familyLists(slot.mapped()[0..np], 0, false); + j.sink = .{ .store = s.pages, .families = s.pool.layout.families, .page = s.pool.page, .world = s.pool.world, .rank = s.pool.rank, .pages = j.lists }; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Per family: family-tensor pages of logical [first, len(pages)) (owned_only: this rank's of split families). + fn familyLists(s: *Store, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const fams = s.pool.layout.families; + const out = try s.gpa.alloc([]u32, fams.len); + var made: usize = 0; + errdefer { + for (out[0..made]) |l| s.gpa.free(l); + s.gpa.free(out); + } + for (fams, 0..) |f, fi| { + const w = s.pool.world; + const sharded = f.split and w > 1 and owned_only; + var n: usize = pages.len - first; + if (sharded) { + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, @intCast(pages.len), w, &by); + n = by[s.pool.rank]; + } + const l = try s.gpa.alloc(u32, n); + var i: usize = 0; + var k: usize = first; + if (sharded) k = first + (s.pool.rank + w - first % w) % w; + while (k < pages.len) : (k += if (sharded) w else 1) { + l[i] = s.pool.familyPage(f, pages[k]); + i += 1; + } + out[fi] = l; + made += 1; + } + return out; + } + + /// Starts a park of a RAM entry (null: nothing to write, the entry is already settled out of RAM). + pub fn beginPark(s: *Store, id: u32) !?*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.ram or e.extra.job != null) return error.NotParkable; + if (e.extra.file and d.get(e.key) != null) { + try s.leaveRam(id); + return null; + } + const page = s.pool.page; + const npages = pool_mod.pagesFor(e.pos, page); + var base: ?Digest = null; + var first: u32 = 0; + if (s.policy.delta) if (e.extra.origin) |o| if (o.pages > 0) { + var r = d.get(o.key); + while (r) |rec| : (r = if (rec.base == disk_mod.none) null else &d.recs.items[rec.base]) { + if (rec.first_page < o.pages) break; + } + if (r) |rec| if (rec.depth < d.opts.max_chain) { + base = rec.key; + first = o.pages; + }; + }; + const j = try s.gpa.create(Job); + j.* = .{ .kind = .park, .entry = id }; + errdefer j.deinit(s.gpa); + j.ids = try s.gpa.alloc(i32, e.pos); + try s.index.idsOf(id, j.ids); + j.lists = try s.familyLists(e.extra.pages, first, true); + j.blobs = try e.extra.bounded.?.vtable.blobs(e.extra.bounded.?.ptr, s.gpa); + j.src = .{ .key = e.key, .base = base, .tag = e.tag, .kind = e.extra.kind, .page = page, .pos = e.pos, .first_page = first, .npages = npages, .ids = j.ids[first * page ..], .store = s.pages, .fence = try s.pages.fence(), .families = s.pool.layout.families, .pages = j.lists, .blobs = j.blobs }; + if (base) |b| d.get(b).?.refs += 1; // held while the job runs: a trim cannot take the base + e.extra.job = j; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Settles a finished job at the plan's point: a park registers its file, a restore returns the bounded state. + pub fn settle(s: *Store, j: *Job) !?disk_mod.Restored { + s.runner.?.wait(j); + defer { + for (s.pending.items, 0..) |p, i| if (p == j) { + _ = s.pending.orderedRemove(i); + break; + }; + j.deinit(s.gpa); + } + const d = s.disk.?; + switch (j.kind) { + .park => { + if (j.src.base) |b| d.get(b).?.refs -= 1; + const e = s.index.get(j.entry); + e.extra.job = null; + if (j.err) |err| return err; // the entry stays in RAM; the .tmp is gone + _ = try d.register(&j.src, j.size); + e.extra.file = true; + s.stats.parks += 1; + if (j.src.base != null) s.stats.delta_parks += 1; + s.stats.park_ns += j.elapsed_ns; + if (s.policy.skip_covered) { + var m: std.ArrayList(u32) = .empty; + defer m.deinit(s.gpa); + _ = try s.index.members(j.entry, &m, usedOf); + for (m.items) |o| s.index.get(o).extra.covered_by = true; + } + try s.leaveRam(j.entry); + try s.trimDisk(); + return null; + }, + .restore => { + if (j.err) |err| { + // the slot as before the restore: empty, its reservation kept + _ = try s.pool.slots.items[j.slot].truncate(0); + return err; + } + const e = s.index.get(j.entry); + d.touch(e.key); + s.touch(j.entry); + s.pool.slots.items[j.slot].origin = .{ .key = e.key, .pages = @intCast(e.pos / s.pool.page) }; + s.stats.hits_disk += 1; + s.stats.restore_ns += j.elapsed_ns; + const r = j.restored; + j.restored = null; + return r; + }, + } + } + + /// Settles every pending job in submission order (one rank, or the plan's "settle all"). + pub fn settleAll(s: *Store) !void { + while (s.pending.items.len > 0) { + if (try s.settle(s.pending.items[0])) |r| { + var rr = r; + rr.deinit(s.gpa); + } + } + } + + fn usedOf(x: *const Extra) u64 { + return x.used; + } + + /// A RAM entry lets go of its pages and bounded state; it stays indexed if a file holds it. + fn leaveRam(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (!e.ram) return; + try s.pool.drop(e.extra.pages); + s.gpa.free(e.extra.pages); + e.extra.pages = &.{}; + if (e.extra.bounded) |b| { + s.bounded_total -= b.vtable.nbytes(b.ptr); + b.vtable.release(b.ptr); + } + e.extra.bounded = null; + s.unlink(id); + s.ram_count -= 1; + try s.index.setRam(id, false, null); + if (!e.extra.file) s.index.remove(id); + } + + /// An entry whose file could not be read back: the file deleted, the entry gone unless RAM holds it (every rank). + pub fn forget(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (e.extra.job != null) return error.Busy; + if (e.extra.file) if (s.disk) |d| { + d.stats.bad += 1; + // a delta's base stays while a newer file stacks on it (an opt-in delta mode) + d.remove(e.key) catch |err| return if (err == error.HasDependents) {} else err; + e.extra.file = false; + }; + if (!e.ram) s.index.remove(id); + } + + /// Drops an entry from RAM without writing it. + pub fn drop(s: *Store, id: u32) !void { + try s.leaveRam(id); + s.stats.drops += 1; + } + + /// Lets go of a RAM entry: dropped when covered, parked if the tier takes it, else dropped; returns a started park. + pub fn evict(s: *Store, id: u32) !?*Job { + const e = s.index.get(id); + if (e.extra.job != null) return null; + if (s.policy.skip_covered and (e.extra.covered_by or s.index.covered(id, .any))) { + try s.drop(id); + s.stats.covered_drops += 1; + return null; + } + if (s.disk) |d| if (d.accepts(e.pos)) { + const j = try s.beginPark(id) orelse return null; + if (!s.policy.async_io) _ = try s.settle(j); + return if (s.policy.async_io) j else null; + }; + try s.drop(id); + return null; + } + + /// Evicts least recently used RAM entries past the bounded-state budget. + fn trimRam(s: *Store) !void { + var id = s.oldest; + while (s.bounded_total > s.policy.ram_bytes and s.ram_count > 1 and id != none) { + const next = s.index.get(id).extra.newer; + if (s.index.get(id).extra.job == null) _ = try s.evict(id); + id = next; + } + } + + /// Files the NVMe budget deleted leave the index (a RAM entry with that key stays, unwritten again). + fn trimDisk(s: *Store) !void { + const d = s.disk orelse return; + var gone: std.ArrayList(Digest) = .empty; + defer gone.deinit(s.gpa); + try d.trim(&gone); + for (gone.items) |k| if (s.index.lookup(k)) |id| { + s.index.get(id).extra.file = false; + if (!s.index.get(id).ram) s.index.remove(id); + }; + } + + const Chat = struct { leaf: u32, recency: u64 }; + + fn chatLess(_: void, a: Chat, b: Chat) bool { + return a.recency < b.recency or (a.recency == b.recency and a.leaf < b.leaf); + } + + /// A chat is live while a slot maps its leaf's first page. + fn idleLeaf(s: *Store, id: u32) bool { + const e = s.index.get(id); + return e.extra.job == null and s.index.isLeaf(id) and (e.extra.pages.len == 0 or s.pool.slot_refs[e.extra.pages[0]] == 0); + } + + /// Evicts LRU idle chats past the residency budget (leaf parked, unshared older turns dropped); returns how many. + pub fn enforce(s: *Store) !u32 { + const pol = s.policy; + if (pol.chats == 0 and pol.resident_bytes == 0) return 0; + var chats: std.ArrayList(Chat) = .empty; + defer chats.deinit(s.gpa); + var members: std.ArrayList(u32) = .empty; + defer members.deinit(s.gpa); + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) { + if (!s.idleLeaf(id)) continue; + members.clearRetainingCapacity(); + try chats.append(s.gpa, .{ .leaf = id, .recency = try s.index.members(id, &members, usedOf) }); + } + std.mem.sort(Chat, chats.items, {}, chatLess); + var idle: u32 = @intCast(chats.items.len); + var gone: u32 = 0; + for (chats.items) |c| { + const over_n = pol.chats > 0 and idle > pol.chats; + const over_b = pol.resident_bytes > 0 and s.pool.entryOnlyBytes() > pol.resident_bytes; + if (!over_n and !over_b) break; + if (!s.index.get(c.leaf).ram or !s.idleLeaf(c.leaf)) continue; + members.clearRetainingCapacity(); + _ = try s.index.members(c.leaf, &members, usedOf); + _ = try s.evict(c.leaf); + // deepest first: a member more RAM entries extend than this chat's leaf is a branch point, as is all above + const own: u32 = if (s.index.get(c.leaf).ram) 1 else 0; + for (members.items) |m| { + if (s.index.get(m).extra.job != null) continue; + if (s.index.extensions(m) > own) break; + if (pol.park_prefixes) _ = try s.evict(m) else try s.drop(m); + } + idle -= 1; + gone += 1; + s.stats.chat_evictions += 1; + } + return gone; + } + + /// RAM entries to let go of (coldest first, never keep) so need pages are free: [] if they are, null if impossible. + pub fn planSpill(s: *Store, need: u32, keep: []const u32, out: *std.ArrayList(u32)) !bool { + out.clearRetainingCapacity(); + var avail = s.pool.available(); + if (avail >= need) return true; + var left: std.AutoHashMapUnmanaged(u32, u16) = .empty; + defer left.deinit(s.gpa); + var id = s.oldest; + while (id != none and avail < need) : (id = s.index.get(id).extra.newer) { + if (std.mem.indexOfScalar(u32, keep, id) != null or s.index.get(id).extra.job != null) continue; + try out.append(s.gpa, id); + for (s.index.get(id).extra.pages) |pg| { + const g = try left.getOrPut(s.gpa, pg); + if (!g.found_existing) g.value_ptr.* = s.pool.refs[pg]; + g.value_ptr.* -= 1; + if (g.value_ptr.* == 0) avail += 1; // split: a page of the scarcer residue may count for less; the check below decides + } + } + return avail >= need; + } + + /// Rebuilds NVMe entries on restart (ranks then intersect them), deltas on their bases' paths; returns how many. + pub fn reconcile(s: *Store) !u32 { + const d = s.disk orelse return 0; + var found: std.ArrayList(Disk.Found) = .empty; + defer { + for (found.items) |f| s.gpa.free(f.ids); + found.deinit(s.gpa); + } + try d.scan(&found); + for (found.items) |f| { + var anchor: u32 = none; + if (f.head.hasBase()) { + const b = s.index.lookup(f.head.base) orelse return error.NoBase; + anchor = s.index.nodeAt(b, f.head.first_page) orelse return error.BadChain; + } + const r = try s.index.insertFrom(f.head.tag, anchor, f.ids, .{ .file = true, .kind = f.head.kind }); + s.clock += 1; + s.index.get(r.id).extra.used = s.clock; + } + return @intCast(found.items.len); + } + + /// Idle resident chats and the pool bytes held by entries alone. + pub fn resident(s: *Store) struct { chats: u32, bytes: u64 } { + var n: u32 = 0; + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) if (s.idleLeaf(id)) { + n += 1; + }; + return .{ .chats = n, .bytes = s.pool.entryOnlyBytes() }; + } +}; diff --git a/zig/src/sessions/store_test.zig b/zig/src/sessions/store_test.zig new file mode 100644 index 0000000000..86b58a4f59 --- /dev/null +++ b/zig/src/sessions/store_test.zig @@ -0,0 +1,216 @@ +//! The session store on a host pool and a real NVMe tier: chats, evictions, resumes bit for bit, deltas, two ranks. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk_mod = @import("disk.zig"); +const store_mod = @import("store.zig"); +const Store = store_mod.Store; +const Bounded = store_mod.Bounded; + +const page = 64; +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Bounded state in host memory: one blob naming its chat and length. +const HostBounded = struct { + gpa: std.mem.Allocator, + bytes: []u8, + + fn make(gpa: std.mem.Allocator, chat: u32, pos: u64) !Bounded { + const h = try gpa.create(HostBounded); + h.* = .{ .gpa = gpa, .bytes = try gpa.alloc(u8, 3000) }; + for (h.bytes, 0..) |*b, i| b.* = @truncate(i +% chat *% 17 +% pos); + return .{ .ptr = h, .vtable = &vt }; + } + const vt: Bounded.VTable = .{ .nbytes = nbytes, .blobs = blobs, .release = release }; + fn nbytes(p: *anyopaque) u64 { + const h: *HostBounded = @ptrCast(@alignCast(p)); + return h.bytes.len; + } + fn blobs(p: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob { + const h: *HostBounded = @ptrCast(@alignCast(p)); + const out = try gpa.alloc(disk_mod.Blob, 1); + out[0] = .{ .name = try gpa.dupe(u8, "rings"), .bytes = try gpa.dupe(u8, h.bytes) }; + return out; + } + fn release(p: *anyopaque) void { + const h: *HostBounded = @ptrCast(@alignCast(p)); + h.gpa.free(h.bytes); + h.gpa.destroy(h); + } +}; + +/// One rank: its pool, host pages and store, plus what each chat's slot wrote (the reference rows by logical page). +const Rank = struct { + pool: Pool, + host: HostPool, + disk: *disk_mod.Disk, + store: Store, + + fn init(r: *Rank, gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, policy: store_mod.Policy) !void { + r.pool = try Pool.init(gpa, .{ .families = &fams, .page = page }, 96 * page, world, rank); + r.host = try HostPool.init(gpa, &r.pool); + r.disk = try disk_mod.Disk.open(gpa, testing.io, .{ .root = root, .compat = @splat(7), .rank = rank, .world = world, .chunk = 64 << 10, .staging = 4, .lanes = 2, .min_tokens = 1 }); + r.store = try Store.init(gpa, testing.io, &r.pool, r.host.store(), r.disk, policy); + } + + fn deinit(r: *Rank) void { + r.store.deinit(); + r.disk.close(); + r.host.deinit(); + r.pool.deinit(); + } + + /// The rows a chat's slot writes at a logical page: a function of (chat, logical page, family) alone. + fn write(r: *Rank, s: *const Slot, chat: u32, from_page: u32) void { + for (s.mapped()[from_page..], from_page..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi); + }; + } + + /// Every stored row of the slot's pages below pos equals what write() put there. + fn check(r: *Rank, s: *const Slot, chat: u32, pos: u64) !void { + const n = pool_mod.pagesFor(pos, page); + for (s.mapped()[0..n], 0..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + const rows: usize = if (k + 1 == n and pos % page != 0) (pos % page) / f.ratio * f.row_bytes else v.len; + for (v[0..rows], 0..) |b, i| try testing.expectEqual(@as(u8, @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi)), b); + }; + } +}; + +fn chatIds(gpa: std.mem.Allocator, chat: u32, n: usize) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @intCast(chat * 100_000 + i); + return v; +} + +/// Prefill or resume a chat to len(ids) in slot s, save it, free the slot, enforce residency; returns the resume point. +fn turn(gpa: std.mem.Allocator, r: *Rank, s: *Slot, chat: u32, ids: []const i32) ![]const u8 { + var tier: []const u8 = "fresh"; + var from: u32 = 0; + if (try r.store.find(1, ids)) |hit| { + if (hit.ram) { + const b = try r.store.restoreRam(hit.id, s); + try testing.expectEqual(@as(u64, 3000), b.vtable.nbytes(b.ptr)); + tier = "ram"; + } else { + const j = try r.store.beginRestore(hit.id, s); + var got = (try r.store.settle(j)).?; + defer got.deinit(gpa); + try testing.expectEqual(@as(u8, @truncate(chat *% 17 +% hit.pos)), got.blobs[0].bytes[0]); + tier = "disk"; + } + try r.check(s, chat, hit.pos); + from = @intCast(hit.pos / page); + } + try s.ensure(ids.len); + r.write(s, chat, from); + _ = try r.store.save(s, try HostBounded.make(gpa, chat, ids.len), ids, 1, 0); + _ = try s.releaseAll(); + _ = try r.store.enforce(); + try r.store.settleAll(); // the plan's settle point (async parks finish here) + return tier; +} + +/// The scenario on one rank; returns a trace of its decisions (stats and page ids) to compare ranks. +fn scenario(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, async_io: bool, trace: *std.ArrayList(u64)) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1, .async_io = async_io }); + defer r.deinit(); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + const b = try chatIds(gpa, 2, 20 * page); + defer gpa.free(b); + const st = &r.store.stats; + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 1, a[0 .. 5 * page + 10])); + try testing.expectEqualStrings("ram", try turn(gpa, &r, s, 1, a[0 .. 8 * page + 3])); + try testing.expectEqual(@as(u64, 0), st.parks); + // a second idle chat: chat 1 (least recent) parked whole, its older turn dropped (not written) + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 2, b[0 .. 6 * page])); + try testing.expectEqual(@as(u64, 1), st.parks); + try testing.expectEqual(@as(u64, 1), st.drops); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + const whole = r.disk.used; + // chat 1 resumes from NVMe bit for bit; chat 2 goes + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 12 * page + 40])); + try testing.expectEqual(@as(u64, 2), st.parks); + const before = r.disk.used; + // chat 2 resumes; chat 1 goes again, written as a delta on its first file (pages 8..12 of 13) + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 2, b[0 .. 9 * page])); + try testing.expectEqual(@as(u64, 3), st.parks); + try testing.expectEqual(@as(u64, 1), st.delta_parks); + // the delta holds 5 of the entry's 13 pages: less than chat 1's first whole file of 9 pages + try testing.expect(r.disk.used - before < whole); + // chat 1's newest turn restores through its two files, bit for bit + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 15 * page])); + try testing.expectEqual(@as(u64, 3), st.hits_disk); + try trace.appendSlice(gpa, &.{ st.parks, st.delta_parks, st.drops, st.covered_drops, st.hits_ram, st.hits_disk, r.pool.freePages(), r.pool.entry_only }); + var id = r.store.oldest; + while (id != store_mod.none) : (id = r.store.entry(id).extra.newer) for (r.store.entry(id).extra.pages) |pg| try trace.append(gpa, pg); +} + +/// After the scenario: a new process on the same tier finds every file and resumes chat 1 through its delta chain. +fn restart(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1 }); + defer r.deinit(); + try testing.expectEqual(@as(u32, 4), try r.store.reconcile()); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 16 * page])); + try testing.expectEqual(@as(u64, 1), r.store.stats.hits_disk); +} + +fn tmpRoot(gpa: std.mem.Allocator, tmp: *std.testing.TmpDir, sub: []const u8) ![:0]u8 { + try tmp.dir.createDirPath(testing.io, sub); + return tmp.dir.realPathFileAlloc(testing.io, sub, gpa); +} + +test "chats over turns: residency to NVMe, bit-identical resumes, deltas, older turns dropped (sync and async)" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const r0 = try tmpRoot(gpa, &tmp, "sync"); + defer gpa.free(r0); + try scenario(gpa, r0, 1, 0, false, &t0); + try restart(gpa, r0, 1, 0); + const r1 = try tmpRoot(gpa, &tmp, "async"); + defer gpa.free(r1); + try scenario(gpa, r1, 1, 0, true, &t1); + try testing.expectEqualSlices(u64, t0.items, t1.items); // the tier thread changes when bytes move, not what is decided +} + +test "two split ranks make the same decisions on the same calls and each restores its own rows" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const root = try tmpRoot(gpa, &tmp, "tier"); + defer gpa.free(root); + try scenario(gpa, root, 2, 0, false, &t0); + try scenario(gpa, root, 2, 1, true, &t1); + try restart(gpa, root, 2, 0); + try restart(gpa, root, 2, 1); + try testing.expectEqualSlices(u64, t0.items, t1.items); +}