From 1be6dc00b243b1c692d5ccc06e52d74143e4773d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:03:35 +0000 Subject: [PATCH 01/46] Windows: a platform layer and host builds that pass - zig/src/platform: `sys`, imported by every module (zig/build/sys.zig wires it). On POSIX it is std.c, std.posix and std.DynLib unchanged; on Windows win_c.zig gives libc's surface over the UCRT and Win32 (integer descriptors, directory descriptors, readdir, pread/pwrite, high-resolution sleep, monotonic clock, /tmp and /proc/self/exe mapped), win_posix.zig the server's sockets over Winsock, console events as SIGINT/SIGTERM, and LoadLibraryW for DynLib. - std.c and std.DynLib calls go through `platform` (mechanical). - build.zig routes Windows targets to the CUDA build. - Server: socket closes and writes are explicit, the key file's mode check is skipped on NTFS (ACLs), descriptor limits are a no-op. - shard_edit: the shard's map is released before the rename, since Windows refuses to replace a mapped file. - Tests that need symlink privilege or mode bits skip on Windows. zig build test-cuda-host passes on x86_64-windows. --- build.zig | 6 +- zig/build/sys.zig | 41 ++ zig/src/core/checkpoint_metal.zig | 23 +- zig/src/core/direct_io.zig | 25 +- zig/src/core/heat.zig | 13 +- zig/src/core/lanes/cost_cache.zig | 5 +- zig/src/core/lanes/grammar.zig | 3 +- zig/src/core/lanes/learned_dirs.zig | 105 ++-- zig/src/core/lanes/learned_disk.zig | 6 +- zig/src/core/learned_fault_test.zig | 41 +- zig/src/core/learned_index_orphan_test.zig | 19 +- zig/src/core/learned_part_test.zig | 45 +- zig/src/core/paced_read.zig | 25 +- zig/src/core/pair.zig | 15 +- zig/src/core/prompt_cache_fake.zig | 11 +- zig/src/core/prompt_cache_learn_test.zig | 29 +- zig/src/core/prompt_imprint.zig | 133 ++-- zig/src/core/prompt_imprint_fault_test.zig | 19 +- zig/src/core/prompt_imprint_test.zig | 31 +- zig/src/core/shard_edit.zig | 57 +- zig/src/core/snapshot_file.zig | 21 +- zig/src/core/template/methods.zig | 15 +- zig/src/cuda/cublaslt.zig | 5 +- zig/src/cuda/driver.zig | 7 +- zig/src/cuda/nccl.zig | 5 +- zig/src/fabric/mailbox.zig | 33 +- zig/src/fabric/mcdma.zig | 3 +- zig/src/fabric/mcdma_abi.zig | 5 +- zig/src/fabric/registration.zig | 45 +- zig/src/fabric/verbs.zig | 7 +- zig/src/fabric/verbs_abi.zig | 11 +- zig/src/fabric/words.zig | 19 +- zig/src/families/flashnext/call_log.zig | 3 +- zig/src/families/flashnext/engine.zig | 19 +- zig/src/families/flashnext/follow.zig | 9 +- zig/src/families/flashnext/host_test.zig | 9 +- zig/src/families/flashnext/learned.zig | 19 +- zig/src/families/flashnext/replay.zig | 9 +- zig/src/families/flashnext/rows_load.zig | 3 +- zig/src/families/flashnext/tp.zig | 41 +- zig/src/families/glm/checks.zig | 15 +- zig/src/families/glm/engine.zig | 51 +- zig/src/families/glm/ep.zig | 23 +- zig/src/families/glm/ep_control.zig | 11 +- zig/src/families/glm/learned_pair.zig | 11 +- zig/src/families/glm/load_plan.zig | 5 +- zig/src/families/glm/mirror.zig | 7 +- zig/src/families/glm/names.zig | 7 +- zig/src/families/glm/page_cache.zig | 35 +- zig/src/families/glm/plain_rounds.zig | 5 +- zig/src/families/glm/pool.zig | 3 +- zig/src/families/glm/prompt.zig | 3 +- zig/src/families/glm/slots.zig | 7 +- zig/src/families/glm/snapshot_file.zig | 29 +- zig/src/families/glm/weights.zig | 1 + zig/src/families/glm/weights_fix.zig | 19 +- zig/src/families/glm/weights_test.zig | 7 +- zig/src/families/nemotron/backend.zig | 3 +- zig/src/families/nemotron/cuda_engine.zig | 3 +- zig/src/families/nemotron/cuda_native.zig | 5 +- zig/src/families/nemotron/cuda_seq_graphs.zig | 3 +- zig/src/families/nemotron/forward.zig | 3 +- zig/src/families/nemotron/gpu_round.zig | 5 +- zig/src/families/nemotron/learned_prompts.zig | 21 +- zig/src/families/nemotron/learner.zig | 3 +- zig/src/families/nemotron/living.zig | 5 +- zig/src/families/nemotron/model.zig | 3 +- zig/src/families/nemotron/pocket.zig | 7 +- zig/src/families/nemotron/pocket_dump.zig | 7 +- zig/src/families/nemotron/whole.zig | 5 +- zig/src/families/qwen27/model.zig | 3 +- zig/src/hip/driver.zig | 5 +- zig/src/hip/runtime.zig | 5 +- zig/src/hip/symbols.zig | 3 +- zig/src/learned_warning_test.zig | 3 +- zig/src/metal/clock.zig | 7 +- zig/src/metal/device.zig | 7 +- zig/src/native/cuda.zig | 3 +- zig/src/native/flashnext_host.zig | 3 +- zig/src/native/metal.zig | 15 +- zig/src/native/qwen27_lanes.zig | 3 +- zig/src/native/qwen27_served.zig | 3 +- zig/src/platform/sys.zig | 47 ++ zig/src/platform/win_c.zig | 577 ++++++++++++++++++ zig/src/platform/win_posix.zig | 279 +++++++++ zig/src/server/auth.zig | 50 +- zig/src/server/http_conn.zig | 3 +- zig/src/server/late_system_test.zig | 21 +- zig/src/server/listener.zig | 3 +- zig/src/server/live.zig | 3 +- zig/src/server/log.zig | 5 +- zig/src/server/request_log.zig | 8 +- zig/src/server/serve.zig | 6 +- zig/src/server/server.zig | 5 +- zig/tests/cuda/pack_compare.zig | 3 +- zig/tests/cuda/train_pack.zig | 3 +- zig/tests/fabric_library.zig | 3 +- zig/tests/flashnext_bench.zig | 3 +- zig/tests/flashnext_reuse.zig | 5 +- zig/tests/flashnext_run.zig | 137 ++--- zig/tests/flashnext_think.zig | 5 +- zig/tests/glm_attn_probe.zig | 7 +- zig/tests/glm_dsa_prompt_check.zig | 5 +- zig/tests/glm_dsa_run.zig | 9 +- zig/tests/glm_rows_speed.zig | 7 +- zig/tests/glm_run.zig | 69 +-- zig/tests/linktest.zig | 9 +- zig/tests/linktest_probe.zig | 3 +- zig/tests/linktest_suite.zig | 3 +- zig/tests/mm_bench.zig | 5 +- zig/tests/moe_bench.zig | 9 +- zig/tests/nemotron_prefill_check.zig | 3 +- zig/tests/pocket_fit_check.zig | 3 +- zig/tests/qwen27/dflash_graph.zig | 13 +- zig/tests/qwen27/dflash_run.zig | 3 +- zig/tests/qwen27/dflash_served.zig | 3 +- zig/tests/qwen27/reuse.zig | 3 +- zig/tests/server/golden_check.zig | 5 +- zig/tests/train_check.zig | 5 +- zig/tests/weight_read_check.zig | 3 +- 120 files changed, 1883 insertions(+), 762 deletions(-) create mode 100644 zig/build/sys.zig create mode 100644 zig/src/platform/sys.zig create mode 100644 zig/src/platform/win_c.zig create mode 100644 zig/src/platform/win_posix.zig diff --git a/build.zig b/build.zig index 08765bd1df..01a2a30bb4 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const sys_build = @import("zig/build/sys.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -153,13 +154,13 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&cost_run.step); switch (target.result.os.tag) { .macos => metalTargets(b, target, optimize, draft_ids, build_options, test_step), - .linux => cuda_build.targets(b, target, optimize, draft_ids, build_options), + .linux, .windows => cuda_build.targets(b, target, optimize, draft_ids, build_options), else => {}, } const xgrammar = grammar_build.option(b); dist_build.targets(b, draft_ids, build_options, release_version, xgrammar); cuda_build.hostTests(b, draft_ids, build_options, test_step, target.result.os.tag != .macos); - if (target.result.os.tag == .linux) { + if (target.result.os.tag == .linux or target.result.os.tag == .windows) { // the engine seam's tests (lane_host.zig and its siblings), which the macOS branch runs from its own api module const lanes_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const api_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes_mod }} }); @@ -189,6 +190,7 @@ 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); + sys_build.wire(b); } /// `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/sys.zig b/zig/build/sys.zig new file mode 100644 index 0000000000..b34be717ed --- /dev/null +++ b/zig/build/sys.zig @@ -0,0 +1,41 @@ +//! `sys` (zig/src/platform/sys.zig) for every module the build compiles: one walk of the step graph after it is +//! declared, so no module's declaration needs to name it. + +const std = @import("std"); + +pub fn wire(b: *std.Build) void { + var seen_steps: std.AutoHashMapUnmanaged(*std.Build.Step, void) = .empty; + var seen_mods: std.AutoHashMapUnmanaged(*std.Build.Module, void) = .empty; + var per_target: std.StringHashMapUnmanaged(*std.Build.Module) = .empty; + var w: Walk = .{ .b = b, .seen_steps = &seen_steps, .seen_mods = &seen_mods, .per_target = &per_target }; + for (b.top_level_steps.values()) |top| w.step(&top.step); +} + +const Walk = struct { + b: *std.Build, + seen_steps: *std.AutoHashMapUnmanaged(*std.Build.Step, void), + seen_mods: *std.AutoHashMapUnmanaged(*std.Build.Module, void), + per_target: *std.StringHashMapUnmanaged(*std.Build.Module), + + fn step(w: *Walk, s: *std.Build.Step) void { + if ((w.seen_steps.getOrPut(w.b.allocator, s) catch @panic("OOM")).found_existing) return; + if (s.cast(std.Build.Step.Compile)) |compile| w.module(compile.root_module, compile.root_module.resolved_target); + for (s.dependencies.items) |dep| w.step(dep); + } + + fn module(w: *Walk, m: *std.Build.Module, inherited: ?std.Build.ResolvedTarget) void { + if ((w.seen_mods.getOrPut(w.b.allocator, m) catch @panic("OOM")).found_existing) return; + const target = m.resolved_target orelse inherited orelse w.b.graph.host; + // only TensorFold's own sources see it; the sys module itself and generated modules do not + if (m.root_source_file) |root| if (root == .src_path and !m.import_table.contains("sys") and !std.mem.endsWith(u8, root.src_path.sub_path, "platform/sys.zig")) + m.addImport("sys", w.sys(target)); + for (m.import_table.values()) |dep| w.module(dep, target); + } + + fn sys(w: *Walk, target: std.Build.ResolvedTarget) *std.Build.Module { + const key = target.result.zigTriple(w.b.allocator) catch @panic("OOM"); + const slot = w.per_target.getOrPut(w.b.allocator, key) catch @panic("OOM"); + if (!slot.found_existing) slot.value_ptr.* = w.b.createModule(.{ .root_source_file = w.b.path("zig/src/platform/sys.zig"), .target = target, .link_libc = true }); + return slot.value_ptr.*; + } +}; diff --git a/zig/src/core/checkpoint_metal.zig b/zig/src/core/checkpoint_metal.zig index aabd70f340..45939c6bb5 100644 --- a/zig/src/core/checkpoint_metal.zig +++ b/zig/src/core/checkpoint_metal.zig @@ -1,5 +1,6 @@ //! The shared safetensors index read into Metal buffers: one shared buffer a shard, each tensor at its data offset. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("safetensors.zig"); @@ -33,7 +34,7 @@ pub const Deferred = struct { file: usize, at: usize, bytes: usize, dtype: DType pub const Later = *const fn (name: []const u8, rank: usize) bool; /// A shard file that deferred tensors stay in, and its address-only views' reservation once one is asked for. -const File = struct { fd: std.c.fd_t, end: usize = 0, views: ?[]align(std.heap.page_size_min) u8 = null }; +const File = struct { fd: platform.c.fd_t, end: usize = 0, views: ?[]align(std.heap.page_size_min) u8 = null }; /// Kept tensors pack at this alignment; reads this large or larger run on 8 threads. pub const alignment = 256; @@ -63,14 +64,14 @@ pub const Checkpoint = struct { /// As addFileSelected; with `later`, the tensors it names stay in the file for `read` and the rest pack into one buffer. pub fn addFileSplit(self: *Checkpoint, device: mtl.Device, path: [:0]const u8, prefix: []const u8, selected: ?[]const u8, later: ?Later) !void { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) { std.log.err("cannot open {s}", .{path}); return error.OpenFailed; } var held = false; defer if (!held) { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var head: [8]u8 = undefined; try readAll(fd, &head, 0); @@ -78,7 +79,7 @@ pub const Checkpoint = struct { const header = try self.allocator.alloc(u8, header_len); defer self.allocator.free(header); try readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); const data_start = 8 + header_len; const data_len: usize = @as(usize, @intCast(end)) - data_start; @@ -114,7 +115,7 @@ pub const Checkpoint = struct { } /// The kept entries packed into one buffer, the `later` ones recorded for `read`; `held` once the file stays open. - fn split(self: *Checkpoint, device: mtl.Device, fd: std.c.fd_t, held: *bool, data_start: usize, entries: st.Header, prefix: []const u8, later: Later) !void { + fn split(self: *Checkpoint, device: mtl.Device, fd: platform.c.fd_t, held: *bool, data_start: usize, entries: st.Header, prefix: []const u8, later: Later) !void { var arena = std.heap.ArenaAllocator.init(self.allocator); defer arena.deinit(); const names = try arena.allocator().alloc([]const u8, entries.count()); @@ -184,7 +185,7 @@ pub const Checkpoint = struct { /// Close the files the deferred tensors live in, once load has read them. pub fn closeFiles(self: *Checkpoint) void { for (self.files.items) |*file| if (file.fd >= 0) { - _ = std.c.close(file.fd); + _ = platform.c.close(file.fd); file.fd = -1; }; } @@ -300,16 +301,16 @@ pub const Stage = struct { } }; -fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readAll(fd: platform.c.fd_t, dest: []u8, at: usize) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } } -const Part = struct { fd: std.c.fd_t, dest: []u8, at: usize, failed: bool = false }; +const Part = struct { fd: platform.c.fd_t, dest: []u8, at: usize, failed: bool = false }; fn readPart(part: *Part) void { readAll(part.fd, part.dest, part.at) catch { @@ -318,7 +319,7 @@ fn readPart(part: *Part) void { } /// pread `dest.len` bytes from `at` on 8 threads (page cache or disk, the copy is the load's main cost). -fn readParallel(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readParallel(fd: platform.c.fd_t, dest: []u8, at: usize) !void { const n_threads = 8; var parts: [n_threads]Part = undefined; var threads: [n_threads]?std.Thread = @splat(null); @@ -335,7 +336,7 @@ fn readParallel(fd: std.c.fd_t, dest: []u8, at: usize) !void { } /// One tensor's bytes: small ones on this thread, large ones on 8. -fn readInto(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readInto(fd: platform.c.fd_t, dest: []u8, at: usize) !void { return if (dest.len >= parallel_bytes) readParallel(fd, dest, at) else readAll(fd, dest, at); } diff --git a/zig/src/core/direct_io.zig b/zig/src/core/direct_io.zig index c06f213388..2f4cc546f6 100644 --- a/zig/src/core/direct_io.zig +++ b/zig/src/core/direct_io.zig @@ -1,6 +1,7 @@ //! Positional reads that skip the page cache (O_DIRECT on Linux): 4 KiB-aligned spans into aligned buffers. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); /// Offset, length and buffer alignment a direct read needs (the logical block size of every NVMe we run on). @@ -20,22 +21,22 @@ pub fn fits(bytes: usize) usize { } pub const File = struct { - fd: std.c.fd_t, + fd: platform.c.fd_t, direct: bool, /// `path` read-only, direct when its file system allows it, else through the page cache. pub fn open(path: [:0]const u8) !File { if (builtin.os.tag == .linux) { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); if (fd >= 0) return .{ .fd = fd, .direct = true }; } - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }); if (fd < 0) return error.FileNotFound; return .{ .fd = fd, .direct = false }; } pub fn close(f: *File) void { - _ = std.c.close(f.fd); + _ = platform.c.close(f.fd); f.* = undefined; } @@ -46,11 +47,11 @@ pub const File = struct { } }; -fn sysPread(_: *anyopaque, fd: std.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { - return std.c.pread(fd, buf, len, @intCast(offset)); +fn sysPread(_: *anyopaque, fd: platform.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { + return platform.c.pread(fd, buf, len, @intCast(offset)); } -fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: *const fn (*anyopaque, std.c.fd_t, [*]u8, usize, u64) isize) ![]u8 { +fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: *const fn (*anyopaque, platform.c.fd_t, [*]u8, usize, u64) isize) ![]u8 { const lo = std.mem.alignBackward(u64, offset, alignment); const want: usize = @intCast(offset + len - lo); const whole = span(offset, len); @@ -62,8 +63,8 @@ fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: const ask = whole - at; const n = pread(ctx, f.fd, buf.ptr + at, ask, lo + @as(u64, @intCast(at))); if (n < 0) { - if (std.c.errno(n) == .INTR) continue; - std.log.err("read of {d} bytes at {d} failed: {t}", .{ ask, lo + at, std.c.errno(n) }); + if (platform.c.errno(n) == .INTR) continue; + std.log.err("read of {d} bytes at {d} failed: {t}", .{ ask, lo + at, platform.c.errno(n) }); return error.ReadFailed; } if (n == 0) return error.EndOfFile; @@ -89,7 +90,7 @@ test "direct reads return the same bytes as buffered reads" { var f = try File.open("/proc/self/exe"); defer f.close(); // The reference uses the page cache, so its unaligned reads do not inherit O_DIRECT. - var plain: File = .{ .fd = std.c.open("/proc/self/exe", .{ .ACCMODE = .RDONLY, .CLOEXEC = true }), .direct = false }; + var plain: File = .{ .fd = platform.c.open("/proc/self/exe", .{ .ACCMODE = .RDONLY, .CLOEXEC = true }), .direct = false }; try std.testing.expect(plain.fd >= 0); defer plain.close(); const buf = try gpa.alignedAlloc(u8, .fromByteUnits(alignment), 1 << 20); @@ -97,7 +98,7 @@ test "direct reads return the same bytes as buffered reads" { const want = try gpa.alloc(u8, 1 << 19); defer gpa.free(want); for ([_][2]usize{ .{ 0, 1 }, .{ 1, 4095 }, .{ 4095, 2 }, .{ 12345, 100000 }, .{ 8192, 1 << 19 } }) |c| { - const n = std.c.pread(plain.fd, want.ptr, c[1], @intCast(c[0])); + const n = platform.c.pread(plain.fd, want.ptr, c[1], @intCast(c[0])); try std.testing.expectEqual(@as(isize, @intCast(c[1])), n); try std.testing.expectEqualSlices(u8, want[0..c[1]], try f.read(buf, c[0], c[1])); } @@ -111,7 +112,7 @@ const Script = struct { calls: usize = 0, bad: bool = false, - fn pread(ctx: *anyopaque, _: std.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { + fn pread(ctx: *anyopaque, _: platform.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { const s: *Script = @ptrCast(@alignCast(ctx)); if (s.calls < s.offs.len) s.offs[s.calls] = offset; s.calls += 1; diff --git a/zig/src/core/heat.zig b/zig/src/core/heat.zig index dc024d653d..eec72a6348 100644 --- a/zig/src/core/heat.zig +++ b/zig/src/core/heat.zig @@ -1,5 +1,6 @@ //! Opt-in wait before a prompt chunk when the hottest thermal zone is above a band. const std = @import("std"); +const platform = @import("sys"); pub const default_root = "/sys/class/thermal"; @@ -129,7 +130,7 @@ fn noteNoZone(root: []const u8) void { } fn envSpan(name: [*:0]const u8) ?[]const u8 { - const p = std.c.getenv(name) orelse return null; + const p = platform.c.getenv(name) orelse return null; const s = std.mem.span(p); if (s.len == 0) return null; return s; @@ -194,13 +195,15 @@ test "the chunk runs under the high band, and a hot reading waits until the low try std.testing.expectEqual(@as(u32, 1), stuck.sleeps); // a cancel ends the wait after the pause it came in } -extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; -extern "c" fn unsetenv(name: [*:0]const u8) i32; +const env = if (@import("builtin").os.tag == .windows) @import("sys").c else struct { + extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; + extern "c" fn unsetenv(name: [*:0]const u8) i32; +}; fn putEnv(name: [*:0]const u8, value: ?[*:0]const u8) void { if (value) |v| { - _ = setenv(name, v, 1); - } else _ = unsetenv(name); + _ = env.setenv(name, v, 1); + } else _ = env.unsetenv(name); } test "TF_HEAT_HIGH and TF_HEAT_LOW set the bands together, and TF_HEAT_ROOT the zones' folder" { diff --git a/zig/src/core/lanes/cost_cache.zig b/zig/src/core/lanes/cost_cache.zig index 1b04fe8fde..279aca3931 100644 --- a/zig/src/core/lanes/cost_cache.zig +++ b/zig/src/core/lanes/cost_cache.zig @@ -1,5 +1,6 @@ //! Timed costs are measured once per engine build, chip and model shape, then read back by later loads. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; const magic = "TFWC2\n"; @@ -53,7 +54,7 @@ fn hashPart(h: *std.crypto.hash.sha2.Sha256, part: []const u8) void { } fn path(a: Allocator, k: [32]u8, suffix: []const u8) !?[]u8 { - const home = std.c.getenv("HOME") orelse return null; + const home = platform.c.getenv("HOME") orelse return null; return try std.fmt.allocPrint(a, "{s}/.cache/tensorfold/window-costs/{s}{s}", .{ std.mem.span(home), std.fmt.bytesToHex(k, .lower), suffix }); } @@ -78,7 +79,7 @@ pub fn save(comptime T: type, a: Allocator, io: std.Io, k: [32]u8, value: T) voi const p = (path(a, k, ".bin") catch return) orelse return; defer a.free(p); std.Io.Dir.cwd().createDirPath(io, std.fs.path.dirname(p) orelse return) catch return; - const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, std.c.getpid() }) catch return; + const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, platform.c.getpid() }) catch return; defer a.free(temp); const bytes = std.mem.concat(a, u8, &.{ magic, std.mem.asBytes(&value) }) catch return; defer a.free(bytes); diff --git a/zig/src/core/lanes/grammar.zig b/zig/src/core/lanes/grammar.zig index 371c2ac4d8..2fe9ae3d87 100644 --- a/zig/src/core/lanes/grammar.zig +++ b/zig/src/core/lanes/grammar.zig @@ -1,6 +1,7 @@ //! Structured output: libtfgrammar (xgrammar, loaded at the first structured request) and a stream's grammar walk. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Allocator = std.mem.Allocator; @@ -37,7 +38,7 @@ fn load(io: std.Io, a: Allocator) ?Fns { library, }; var dl = for (names) |name| { - break std.DynLib.open(name) catch continue; + break platform.DynLib.open(name) catch continue; } else return null; var f: Fns = undefined; const info = @typeInfo(Fns).@"struct"; diff --git a/zig/src/core/lanes/learned_dirs.zig b/zig/src/core/lanes/learned_dirs.zig index 780a997a5c..cdf9648438 100644 --- a/zig/src/core/lanes/learned_dirs.zig +++ b/zig/src/core/lanes/learned_dirs.zig @@ -1,5 +1,6 @@ //! Bounded-name learned directories: both ranks plan the same identity deletions before either writes. const std = @import("std"); +const platform = @import("sys"); pub fn rootOf(dir: []const u8) []const u8 { return std.fs.path.dirname(dir) orelse "."; } @@ -13,23 +14,23 @@ fn isCurrent(keep: []const u8, id: u64) bool { } pub fn next(root: []const u8, keep: []const u8, after: ?u64) !?u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var out: ?u64 = null; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name.len != 16) continue; const id = std.fmt.parseInt(u64, name, 16) catch continue; if (id == std.math.maxInt(u64)) return error.LearnedDirectoryType; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; if (isCurrent(keep, id) or (after != null and id <= after.?)) continue; if (out == null or id < out.?) out = id; } return out; } pub fn unlink(file: [:0]const u8) !void { - const rc = std.c.unlink(file); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlink(file); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } fn stateKey(name: []const u8) ?u64 { @@ -52,38 +53,38 @@ fn identityName(name: []const u8) bool { } fn unlinkAt(fd: c_int, name: [*:0]const u8) !void { - const rc = std.c.unlinkat(fd, name, 0); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlinkat(fd, name, 0); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } /// Only the configured root may be an alias; cleanup never follows an identity or file link. pub fn cleanupParts(root: []const u8) !void { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); - if (ent.type == std.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { + if (ent.type == platform.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { try unlinkAt(fd, @ptrCast(&ent.name)); continue; } if (!identityName(name)) continue; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; - const sub_fd = std.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; + const sub_fd = platform.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (sub_fd < 0) return error.LearnedDirectoryRead; - const sub = std.c.fdopendir(sub_fd) orelse { - _ = std.c.close(sub_fd); + const sub = platform.c.fdopendir(sub_fd) orelse { + _ = platform.c.close(sub_fd); return error.LearnedDirectoryRead; }; { - defer _ = std.c.closedir(sub); - while (std.c.readdir(sub)) |file| { - if (file.type == std.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); + defer _ = platform.c.closedir(sub); + while (platform.c.readdir(sub)) |file| { + if (file.type == platform.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); } } } @@ -91,15 +92,15 @@ pub fn cleanupParts(root: []const u8) !void { /// Payloads have no token metadata to recover a lost index entry; only exact owned, unreferenced names go. pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { - const fd = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; const name = std.mem.sliceTo(&ent.name, 0); const key = stateKey(name) orelse continue; const indexed = for (metas) |x| { @@ -110,29 +111,29 @@ pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { } } pub fn bytes(dir: [:0]const u8) !u64 { - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var buf: [1400]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end < 0) return error.LearnedDirectoryRead; n = try std.math.add(u64, n, @intCast(end)); } return n; } pub fn fileBytes(file: [:0]const u8) !u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.LearnedDirectoryRead; return @intCast(end); } @@ -156,16 +157,16 @@ pub fn total(root: []const u8) !u64 { pub fn used(root: []const u8, id: u64) !u64 { var buf: [1100]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}/used", .{ root, id }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); var n: u64 = 0; const data = std.mem.asBytes(&n); var done: usize = 0; while (done < data.len) { - const count = std.c.read(fd, data.ptr + done, data.len - done); + const count = platform.c.read(fd, data.ptr + done, data.len - done); if (count < 0) { - if (std.c.errno(count) == .INTR) continue; + if (platform.c.errno(count) == .INTR) continue; return error.LearnedDirectoryRead; } if (count == 0) return 0; // A torn recency hint is oldest; it cannot make a learned payload valid. @@ -178,23 +179,23 @@ pub fn removeOther(root: []const u8, keep: []const u8, id: u64) !void { var buf: [1100]u8 = undefined; const dir = try path(&buf, root, id); _ = try bytes(dir); - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var f: [1400]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&f, "{s}/{s}", .{ dir, name }, 0)); } - const rc = std.c.rmdir(dir); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.rmdir(dir); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } test "acknowledged deletion refuses a directory and accepts an absent file" { var buf: [128]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{std.c.getpid()}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{platform.c.getpid()}, 0); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, unlink(dir)); var missing: [160]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&missing, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/lanes/learned_disk.zig b/zig/src/core/lanes/learned_disk.zig index 9ca5fbd4c0..d6e275fae1 100644 --- a/zig/src/core/lanes/learned_disk.zig +++ b/zig/src/core/lanes/learned_disk.zig @@ -1,5 +1,6 @@ //! Reserve learned-state writes above a free-disk floor, with bounded backoff after refusals or write failures. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Count = if (builtin.os.tag == .linux) u64 else u32; const Stat = extern struct { @@ -19,14 +20,15 @@ const Stat = extern struct { extern "c" fn statvfs(path: [*:0]const u8, result: *Stat) c_int; pub fn available(dir: [:0]const u8) ?u64 { + if (builtin.os.tag == .windows) return platform.c.diskAvailable(dir); var stat: Stat = undefined; if (statvfs(dir, &stat) != 0) return null; return std.math.mul(u64, stat.frsize, stat.bavail) catch null; } fn now() ?u64 { - var t: std.c.timespec = undefined; - if (std.c.clock_gettime(std.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; + var t: platform.c.timespec = undefined; + if (platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; return @intCast(t.sec); } diff --git a/zig/src/core/learned_fault_test.zig b/zig/src/core/learned_fault_test.zig index fefd560d99..98d9c9209b 100644 --- a/zig/src/core/learned_fault_test.zig +++ b/zig/src/core/learned_fault_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const Imprint = @import("prompt_imprint.zig").Imprint; const Fake = @import("prompt_cache_fake.zig").Fake; @@ -14,14 +15,14 @@ const Disk = struct { } }; fn root(buf: []u8, suffix: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-learn-fault-{s}-{d}", .{ suffix, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-learn-fault-{s}-{d}", .{ suffix, platform.c.getpid() }); } fn put(dir: []const u8, name: []const u8, count: usize) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.Create; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(count)) != 0) return error.Create; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(count)) != 0) return error.Create; } pub fn learnOne(im: *Imprint, prompt: []const u32) !void { @@ -61,9 +62,9 @@ pub fn recallOne(im: *Imprint, prompt: []const u32) !void { fn appendIndex(dir: []const u8, bytes: []const u8) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/index", .{dir}, 0), .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/index", .{dir}, 0), .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try @import("prompt_imprint.zig").writeAll(fd, bytes); } @@ -125,9 +126,9 @@ test "invalid complete records salvage the valid prefix but unknown versions rem } else |err| try std.testing.expectEqual(error.ImprintRead, err); try std.testing.expectEqual(@as(u64, 104), try im.indexScratchBytes()); var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var actual: [44]u8 = undefined; try std.testing.expect(@import("prompt_imprint.zig").readAt(fd, &actual, 60)); try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&record), &actual); @@ -206,9 +207,9 @@ test "short future-version prefixes preserve index bytes and payloads on open an defer im.deinit(); if (prefix) try learnOne(&im, &first); var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try @import("prompt_imprint.zig").writeAll(fd, std.mem.sliceAsBytes(&future)[0..cut]); var expected: [96]u8 = undefined; const len = (if (prefix) @as(usize, 60) else 0) + cut; @@ -225,13 +226,13 @@ test "short future-version prefixes preserve index bytes and payloads on open an var actual: [96]u8 = undefined; try std.testing.expect(@import("prompt_imprint.zig").readAt(fd, actual[0..len], 0)); try std.testing.expectEqualSlices(u8, expected[0..len], actual[0..len]); - const payload = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/00000000000000af.bin", .{im.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const payload = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/00000000000000af.bin", .{im.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (payload < 0) return error.FuturePayloadRemoved; - _ = std.c.close(payload); + _ = platform.c.close(payload); if (prefix) { - const kept = std.c.open(try Fake.file(&path, im.dir, Imprint.keyOf(first[0..6])), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const kept = platform.c.open(try Fake.file(&path, im.dir, Imprint.keyOf(first[0..6])), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (kept < 0) return error.ValidPrefixPayloadRemoved; - _ = std.c.close(kept); + _ = platform.c.close(kept); } } }; @@ -272,14 +273,14 @@ test "an index replacement failure retains its metadata and removes its partial var idx_buf: [1200]u8 = undefined; const idx = try std.fmt.bufPrintSentinel(&idx_buf, "{s}/index", .{im.dir}, 0); try dirs.unlink(idx); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(idx, 0o700)); - defer _ = std.c.rmdir(idx); // the tree removal goes two levels down + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(idx, 0o700)); + defer _ = platform.c.rmdir(idx); // the tree removal goes two levels down try std.testing.expectError(error.ImprintWrite, im.remove(Imprint.keyOf(&.{ 1, 2 }))); try std.testing.expect(im.has(Imprint.keyOf(&.{ 1, 2 }))); const part = try std.fmt.bufPrintSentinel(&path, "{s}/index.part", .{im.dir}, 0); - const fd = std.c.open(part, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(part, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.PartRemains; } } @@ -315,8 +316,8 @@ test "unlink failure is observable and missing files are idempotently removed" { const r = try root(&buf, "unlink"); var path: [150]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&path, "{s}", .{r}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, dirs.unlink(dir)); var absent: [180]u8 = undefined; try dirs.unlink(try std.fmt.bufPrintSentinel(&absent, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/learned_index_orphan_test.zig b/zig/src/core/learned_index_orphan_test.zig index 7b361f0907..6b210bb469 100644 --- a/zig/src/core/learned_index_orphan_test.zig +++ b/zig/src/core/learned_index_orphan_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const imprint = @import("prompt_imprint.zig"); const Imprint = imprint.Imprint; const Fake = @import("prompt_cache_fake.zig").Fake; @@ -26,15 +27,15 @@ fn learn(im: *Imprint, prompt: []const u32) !void { try std.testing.expect(im.has(Imprint.keyOf(prompt[0..6]))); } fn exists(file: [:0]const u8) bool { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return false; - _ = std.c.close(fd); + _ = platform.c.close(fd); return true; } fn put(file: [:0]const u8, bytes: []const u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(file, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, bytes); } @@ -42,7 +43,7 @@ test "reopen sweeps only unindexed owned payloads left by old post-tear learning const a = std.testing.allocator; for (0..3) |kind| { var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-lost-{d}-{d}", .{ std.c.getpid(), kind }); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-lost-{d}-{d}", .{ platform.c.getpid(), kind }); defer @import("prompt_imprint_test.zig").rmTree(root); const first = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 }; var later = [_]u32{ 11, 12, 13, 14, 15, 16, 17, 18 }; @@ -57,14 +58,14 @@ test "reopen sweeps only unindexed owned payloads left by old post-tear learning later[0] += 10; } var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var original: [240]u8 = undefined; try std.testing.expect(imprint.readAt(fd, &original, 0)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 60)); - if (std.c.lseek(fd, 60, std.c.SEEK.SET) != 60) return error.TestFile; + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 60)); + if (platform.c.lseek(fd, 60, platform.c.SEEK.SET) != 60) return error.TestFile; const key = Imprint.keyOf(&.{ 31, 32 }); const torn = [_]u32{ 0x494d5052, 2, @truncate(key), @truncate(key >> 32), 1, if (kind == 0) 512 else 2, 0, 100, 0 }; if (kind != 1) try imprint.writeAll(fd, std.mem.sliceAsBytes(&torn)); diff --git a/zig/src/core/learned_part_test.zig b/zig/src/core/learned_part_test.zig index c013e5de1b..2f7ab53027 100644 --- a/zig/src/core/learned_part_test.zig +++ b/zig/src/core/learned_part_test.zig @@ -1,22 +1,23 @@ const std = @import("std"); +const platform = @import("sys"); const Imprint = @import("prompt_imprint.zig").Imprint; const life = @import("learned_fault_test.zig"); fn scratch(buf: []u8, tag: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-learn-part-{s}-{d}", .{ tag, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-learn-part-{s}-{d}", .{ tag, platform.c.getpid() }); } fn put(dir: []const u8, name: []const u8, count: usize) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(count)) != 0) return error.TestFile; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(count)) != 0) return error.TestFile; } fn exists(dir: []const u8, name: []const u8) bool { var buf: [1200]u8 = undefined; - const fd = std.c.open(std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0) catch return false, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0) catch return false, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return false; - _ = std.c.close(fd); + _ = platform.c.close(fd); return true; } @@ -95,30 +96,30 @@ test "owned-name links and directories stay untouched and foreign directories ar const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}/target", .{root}, 0); var link_buf: [1200]u8 = undefined; const link = try std.fmt.bufPrintSentinel(&link_buf, "{s}/00000000000000af.bin.part", .{im.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, link)); - defer _ = std.c.unlink(link); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, link)); + defer _ = platform.c.unlink(link); var dir_buf: [1200]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&dir_buf, "{s}/index.part", .{im.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); var foreign_buf: [1200]u8 = undefined; const foreign = try std.fmt.bufPrintSentinel(&foreign_buf, "{s}/foreign", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(foreign, 0o700)); - defer _ = std.c.rmdir(foreign); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(foreign, 0o700)); + defer _ = platform.c.rmdir(foreign); try put(foreign, "index.part", 3); var reopened = try Imprint.open(a, root, 1, 1 << 20); defer reopened.deinit(); var actual: [1200]u8 = undefined; - const n = std.c.readlink(link, &actual, actual.len); + const n = platform.c.readlink(link, &actual, actual.len); try std.testing.expectEqual(@as(isize, @intCast(target.len)), n); try std.testing.expectEqualSlices(u8, target, actual[0..@intCast(n)]); - const directory = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const directory = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); try std.testing.expect(directory >= 0); - _ = std.c.close(directory); + _ = platform.c.close(directory); try std.testing.expect(exists(root, "target")); try std.testing.expect(exists(foreign, "index.part")); var marker_buf: [1400]u8 = undefined; - _ = std.c.unlink(try std.fmt.bufPrintSentinel(&marker_buf, "{s}/index.part", .{foreign}, 0)); + _ = platform.c.unlink(try std.fmt.bufPrintSentinel(&marker_buf, "{s}/index.part", .{foreign}, 0)); } test "an identity-directory symlink refuses open and preserves foreign payloads" { @@ -132,12 +133,12 @@ test "an identity-directory symlink refuses open and preserves foreign payloads" } var target_buf: [1200]u8 = undefined; const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}/foreign", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(target, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(target, 0o700)); try put(target, "00000000000000af.bin", 7); var link_buf: [1200]u8 = undefined; const link = try std.fmt.bufPrintSentinel(&link_buf, "{s}/0000000000000001", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, link)); - defer _ = std.c.unlink(link); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, link)); + defer _ = platform.c.unlink(link); try std.testing.expectError(error.ImprintRead, Imprint.open(a, root, 1, 1 << 20)); try std.testing.expect(exists(target, "00000000000000af.bin")); try std.testing.expect(!exists(target, "used")); @@ -154,11 +155,11 @@ test "an explicitly configured root alias still cleans only its real identity ch try put(seeded.dir, "index.part", 4096); } var alias_buf: [128]u8 = undefined; - const alias = try std.fmt.bufPrintSentinel(&alias_buf, "/tmp/tf-learn-root-alias-{d}", .{std.c.getpid()}, 0); + const alias = try std.fmt.bufPrintSentinel(&alias_buf, "/tmp/tf-learn-root-alias-{d}", .{platform.c.getpid()}, 0); var target_buf: [128]u8 = undefined; const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, alias)); - defer _ = std.c.unlink(alias); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, alias)); + defer _ = platform.c.unlink(alias); var reopened = try Imprint.open(a, alias, 1, 1 << 20); defer reopened.deinit(); try std.testing.expect(!exists(reopened.dir, "index.part")); diff --git a/zig/src/core/paced_read.zig b/zig/src/core/paced_read.zig index b23dd15c4a..0b72289b1f 100644 --- a/zig/src/core/paced_read.zig +++ b/zig/src/core/paced_read.zig @@ -1,5 +1,6 @@ //! Checkpoint reads that keep out of the page cache and wait for free memory: a load never outruns the kernel's reclaim. const std = @import("std"); +const platform = @import("sys"); /// Offsets, lengths and buffers of uncached reads align to the VM page (F_NOCACHE only skips the cache for aligned reads). pub const alignment = 16384; @@ -16,17 +17,17 @@ pub fn stage(gpa: std.mem.Allocator, bytes: usize) !Stage { pub fn floor() u64 { var mem: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; + if (platform.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; return @max(8 << 30, mem / 100 * 3); } pub fn freeBytes() u64 { var pages: u32 = 0; var len: usize = @sizeOf(u32); - if (std.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); + if (platform.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); var size: u64 = 0; var len2: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; + if (platform.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; return @as(u64, pages) * size; } @@ -35,13 +36,13 @@ pub fn waitForRoom(limit_ms: u32) void { const want = floor(); var waited: u32 = 0; while (waited < limit_ms and freeBytes() < want) : (waited += 2) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; + _ = platform.c.nanosleep(&ts, null); } } /// `dest.len` bytes at `at` of `fd` (opened with F_NOCACHE) through page-aligned reads into `buf`, pausing for free memory. -pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { +pub fn read(fd: platform.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { var done: usize = 0; while (done < dest.len) { waitForRoom(10_000); @@ -52,7 +53,7 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { const span = std.mem.alignForward(usize, skip + want, alignment); var got: usize = 0; while (got < skip + want) { // a read past the file's end comes back short once the bytes asked for are in - const n = std.c.pread(fd, buf.ptr + got, span - got, @intCast(lo + got)); + const n = platform.c.pread(fd, buf.ptr + got, span - got, @intCast(lo + got)); if (n <= 0) return error.ShortRead; got += @intCast(n); } @@ -64,17 +65,17 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { test "aligned spans copy the bytes asked for, at any offset and length" { const gpa = std.testing.allocator; var name_buf: [64]u8 = undefined; - const name = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-paced-read-{d}", .{std.c.getpid()}, 0); - const fd = std.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const name = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-paced-read-{d}", .{platform.c.getpid()}, 0); + const fd = platform.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); defer { - _ = std.c.close(fd); - _ = std.c.unlink(name); + _ = platform.c.close(fd); + _ = platform.c.unlink(name); } const data = try gpa.alloc(u8, 3 * alignment + 777); defer gpa.free(data); for (data, 0..) |*b, i| b.* = @truncate(i *% 131 +% 7); - try std.testing.expectEqual(@as(isize, @intCast(data.len)), std.c.write(fd, data.ptr, data.len)); + try std.testing.expectEqual(@as(isize, @intCast(data.len)), platform.c.write(fd, data.ptr, data.len)); const buf = try stage(gpa, 2 * alignment); defer gpa.free(buf); for ([_][2]usize{ .{ 0, 1 }, .{ 5, 100 }, .{ alignment - 3, 10 }, .{ 1, 2 * alignment + 50 }, .{ 3 * alignment, 777 }, .{ 9, data.len - 9 } }) |c| { diff --git a/zig/src/core/pair.zig b/zig/src/core/pair.zig index b7d6daed68..d6fc64cd68 100644 --- a/zig/src/core/pair.zig +++ b/zig/src/core/pair.zig @@ -1,5 +1,6 @@ //! Two Macs' tensor-parallel exchanges in a command buffer over MCDMA, and rank 0's commands to rank 1. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const sources = @import("kernel_sources"); @@ -54,7 +55,7 @@ threadlocal var raised = false; /// This thread at the user-interactive QoS once: busy CPUs must not delay the pair's sends and commands. pub fn raise() void { if (raised) return; - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); raised = true; } @@ -230,16 +231,16 @@ pub const Pair = struct { /// Wait until this Mac's word at `off` reaches `value`: spins, then short sleeps (`bounded`: 60 s at most). fn reach(p: *Pair, off: usize, value: u64, bounded: bool) !void { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); var spins: usize = 0; while (@atomicLoad(u64, p.word64(off), .acquire) < value) { if (p.stop.load(.acquire)) return error.PairStopped; if (p.failed.load(.acquire)) return error.PairLinkFailed; - if (bounded and std.c.mach_absolute_time() - t0 > 60 * 24_000_000) return error.PairPeerSilent; + if (bounded and platform.c.mach_absolute_time() - t0 > 60 * 24_000_000) return error.PairPeerSilent; spins += 1; if (spins > 200_000) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } } @@ -254,8 +255,8 @@ pub const Pair = struct { if (!reached(posted, next_x) or !reached(p.queued.load(.acquire), next_x)) { idle += 1; if (idle > 1_000_000) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); continue; } diff --git a/zig/src/core/prompt_cache_fake.zig b/zig/src/core/prompt_cache_fake.zig index 6f6754d0b6..734ecafae9 100644 --- a/zig/src/core/prompt_cache_fake.zig +++ b/zig/src/core/prompt_cache_fake.zig @@ -1,5 +1,6 @@ //! A host-memory snapshot family shared by the prompt cache's tests and the learned-disk fault tests. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const Snapshots = pc.Snapshots; @@ -96,7 +97,7 @@ pub const Fake = struct { } var path: [512]u8 = undefined; - _ = std.c.unlink(file(&path, dir, key) catch return); + _ = platform.c.unlink(file(&path, dir, key) catch return); } pub fn file(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { return std.fmt.bufPrintSentinel(buf, "{s}/{x:0>16}.bin", .{ dir, key }, 0); @@ -106,17 +107,17 @@ pub const Fake = struct { f.writes += 1; if (f.fail_write) return error.WriteFailed; var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.WriteFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, std.mem.asBytes(@as(*State, @ptrCast(@alignCast(saved))))); } fn readFn(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror!Saved { const f = of(ptr); var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.ReadFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const st = try f.gpa.create(State); errdefer f.gpa.destroy(st); if (!imprint.readAt(fd, std.mem.asBytes(st), 0) or st.at != at) return error.ReadFailed; diff --git a/zig/src/core/prompt_cache_learn_test.zig b/zig/src/core/prompt_cache_learn_test.zig index 73eb857f2a..26f9f3f28e 100644 --- a/zig/src/core/prompt_cache_learn_test.zig +++ b/zig/src/core/prompt_cache_learn_test.zig @@ -1,5 +1,6 @@ //! The prompt cache's learned-state tests: harness states kept on disk, the cap, the disk floor and peers' halves. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const rmTree = @import("prompt_imprint_test.zig").rmTree; @@ -14,7 +15,7 @@ test "a learned harness state outlives its store: a fresh session on a new one r defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{platform.c.getpid()}); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; // shared by every session: a state at 5 reads it all (lookahead 1) const key = imprint.Imprint.keyOf(&harness); @@ -50,7 +51,7 @@ test "past the learned-state cap the least recently used state is forgotten, and defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const one = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -67,8 +68,8 @@ test "past the learned-state cap the least recently used state is forgotten, and try std.testing.expectEqual(fresh(&p2), f.pass(&s, &p2, try s.begin(a, &p2, 7, &.{5}, &.{}, null, &.{}))); try std.testing.expect(im.has(imprint.Imprint.keyOf(&two)) and !im.has(imprint.Imprint.keyOf(&one))); var path: [512]u8 = undefined; - try std.testing.expect(std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too - _ = std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back + try std.testing.expect(platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too + _ = platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back var g: Fake = .{ .gpa = gpa }; var t = Store.init(gpa, g.learned(), rules, 1 << 20); defer t.deinit(); @@ -87,7 +88,7 @@ test "learned states are prompt arithmetic: a request that decodes rows below on defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -138,7 +139,7 @@ test "Store learning preserves the disk floor, backs off writes across keys, and var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 11, 1 << 20); defer im.deinit(); @@ -180,7 +181,7 @@ test "learning reserves the existing index rewrite above the disk floor and resu var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 12, 1 << 20); defer im.deinit(); @@ -213,7 +214,7 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 51, 1 << 20); defer im.deinit(); @@ -242,9 +243,9 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem try std.testing.expectEqual(@as(u64, 0), im.admission.reserved); var file_path: [512]u8 = undefined; const file = try Fake.file(&file_path, im.dir, imprint.Imprint.keyOf(prompt[0..6])); - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.OrphanedHalf; } } @@ -263,16 +264,16 @@ test "under-cap disk pressure selects the same LRU victim before either half is }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 52, 1 << 20); defer im.deinit(); try im.add(imprint.Imprint.keyOf(&.{ 1, 2 }), 1, &.{ 1, 2 }, &.{}, 16); var victim_path: [512]u8 = undefined; - const victim_fd = std.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const victim_fd = platform.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (victim_fd < 0) return error.Create; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(victim_fd, 16)); - _ = std.c.close(victim_fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(victim_fd, 16)); + _ = platform.c.close(victim_fd); im.admission = .{ .floor = 100, .free = Disk.free, .clock = Disk.clock }; const prompt = [_]u32{ 7, 7, 7, 7, 7, 7, 9, 10 }; Disk.space = 100 + 105 + 256 + (6 * 4) + (try im.indexScratchBytes()) - 8; diff --git a/zig/src/core/prompt_imprint.zig b/zig/src/core/prompt_imprint.zig index 8fee9da5ad..024a8a9b6a 100644 --- a/zig/src/core/prompt_imprint.zig +++ b/zig/src/core/prompt_imprint.zig @@ -1,5 +1,6 @@ //! Prompt-cache states on disk (Imprint): a learned harness state outlives the server and serves later fresh sessions. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; /// One state on disk: the tokens it read (position plus lookahead), the chunk starts below it, and its bytes. @@ -74,10 +75,10 @@ pub const Imprint = struct { /// Reserve the current index's size too: LRU eviction rewrites it through a temporary file. pub fn indexScratchBytes(m: *const Imprint) !u64 { var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.ImprintRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.ImprintRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.ImprintRead; return @intCast(end); } @@ -119,10 +120,10 @@ pub const Imprint = struct { pub fn pendingWrite(m: *const Imprint) !?u64 { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/pending", .{m.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) null else error.ImprintRead; - defer _ = std.c.close(fd); - if (std.c.lseek(fd, 0, std.c.SEEK.END) != 48) return error.ImprintRead; + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/pending", .{m.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) null else error.ImprintRead; + defer _ = platform.c.close(fd); + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != 48) return error.ImprintRead; var data: [48]u8 = undefined; if (!readAt(fd, &data, 0)) return error.ImprintRead; var digest: [32]u8 = undefined; @@ -135,19 +136,19 @@ pub const Imprint = struct { var buf: [1100]u8 = undefined; var dest: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&buf, "{s}/pending.part", .{m.dir}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var data: [48]u8 = undefined; std.mem.writeInt(u64, data[0..8], 0x494d5052494e5431, .little); std.mem.writeInt(u64, data[8..16], key, .little); std.crypto.hash.sha2.Sha256.hash(data[0..16], data[16..48], .{}); try writeAll(fd, &data); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; } - if (std.c.rename(tmp, try std.fmt.bufPrintSentinel(&dest, "{s}/pending", .{m.dir}, 0)) != 0) return error.ImprintWrite; + if (platform.c.rename(tmp, try std.fmt.bufPrintSentinel(&dest, "{s}/pending", .{m.dir}, 0)) != 0) return error.ImprintWrite; } pub fn clearPending(m: *const Imprint) !void { var buf: [1100]u8 = undefined; @@ -159,16 +160,16 @@ pub const Imprint = struct { if (m.has(key)) return; const below = cut(starts, at); var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true, .NOFOLLOW = true, .NONBLOCK = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true, .NOFOLLOW = true, .NONBLOCK = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); _ = indexSize(fd) catch return error.ImprintWrite; - var end = std.c.lseek(fd, 0, std.c.SEEK.END); + var end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.ImprintWrite; if (end != m.index_end) { try m.reload(fd); if (m.has(key)) return; - end = std.c.lseek(fd, 0, std.c.SEEK.END); + end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0 or end != m.index_end) return error.ImprintWrite; } const t = try m.gpa.dupe(u32, tokens); @@ -176,9 +177,9 @@ pub const Imprint = struct { const s = try m.gpa.dupe(u32, below); errdefer m.gpa.free(s); try m.metas.ensureUnusedCapacity(m.gpa, 1); - errdefer _ = std.c.ftruncate(fd, end); + errdefer _ = platform.c.ftruncate(fd, end); try record(fd, key, at, tokens, below, bytes); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; m.index_end = @as(u64, @intCast(end)) + HEAD * 4 + 4 * (tokens.len + below.len); m.clock += 1; m.metas.appendAssumeCapacity(.{ .key = key, .at = at, .tokens = t, .starts = s, .bytes = bytes, .used = m.clock }); @@ -190,15 +191,15 @@ pub const Imprint = struct { var path: [1100]u8 = undefined; var part: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&part, "{s}/index.part", .{m.dir}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); for (m.metas.items, 0..) |x, j| if (j != i) try record(fd, x.key, x.at, x.tokens, x.starts, x.bytes); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; } - if (std.c.rename(tmp, try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0)) != 0) return error.ImprintWrite; + if (platform.c.rename(tmp, try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0)) != 0) return error.ImprintWrite; free(m.gpa, m.metas.orderedRemove(i)); m.index_end = 0; for (m.metas.items) |x| m.index_end += HEAD * 4 + 4 * (x.tokens.len + x.starts.len); @@ -207,9 +208,9 @@ pub const Imprint = struct { /// Keep a valid prefix across torn or invalid suffixes; unknown versions and I/O errors refuse learning. fn load(m: *Imprint) !void { var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .NONBLOCK = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) {} else error.ImprintRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .NONBLOCK = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) {} else error.ImprintRead; + defer _ = platform.c.close(fd); try m.readIndex(fd); } @@ -300,20 +301,20 @@ fn indexSize(fd: c_int) !u64 { if (@import("builtin").os.tag == .linux) { const linux = std.os.linux; var sx: linux.Statx = undefined; - if (std.c.statx(fd, "", linux.AT.EMPTY_PATH, .{ .TYPE = true, .SIZE = true }, &sx) != 0) return error.ImprintRead; - if (!sx.mask.TYPE or !sx.mask.SIZE or @as(u32, sx.mode) & std.c.S.IFMT != std.c.S.IFREG) return error.ImprintRead; + if (platform.c.statx(fd, "", linux.AT.EMPTY_PATH, .{ .TYPE = true, .SIZE = true }, &sx) != 0) return error.ImprintRead; + if (!sx.mask.TYPE or !sx.mask.SIZE or @as(u32, sx.mode) & platform.c.S.IFMT != platform.c.S.IFREG) return error.ImprintRead; if (sx.size > @as(u64, std.math.maxInt(i64))) return error.ImprintRead; return sx.size; } else { - var stat: std.c.Stat = undefined; - if (std.c.fstat(fd, &stat) != 0 or stat.mode & std.c.S.IFMT != std.c.S.IFREG or stat.size < 0) return error.ImprintRead; + var stat: platform.c.Stat = undefined; + if (platform.c.fstat(fd, &stat) != 0 or stat.mode & platform.c.S.IFMT != platform.c.S.IFREG or stat.size < 0) return error.ImprintRead; return @intCast(stat.size); } } fn repairTail(fd: c_int, whole: u64, observed: u64) !void { - if (std.c.lseek(fd, 0, std.c.SEEK.END) != observed) return error.ImprintRead; - if (std.c.ftruncate(fd, @intCast(whole)) != 0 or std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != observed) return error.ImprintRead; + if (platform.c.ftruncate(fd, @intCast(whole)) != 0 or platform.c.fsync(fd) != 0) return error.ImprintWrite; } fn ordered(starts: []const u32) bool { @@ -358,9 +359,9 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { var dirs: std.ArrayList(Dir) = .empty; defer dirs.deinit(gpa); var buf: [1100]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return 0) orelse return 0; + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return 0) orelse return 0; var held: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (!identityName(name) or std.mem.endsWith(u8, keep, name)) continue; var sub: [1100]u8 = undefined; @@ -370,7 +371,7 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { held += x.bytes; dirs.append(gpa, x) catch break; } - _ = std.c.closedir(d); + _ = platform.c.closedir(d); std.mem.sort(Dir, dirs.items, {}, struct { fn older(_: void, a: Dir, b: Dir) bool { return a.used < b.used; @@ -381,7 +382,7 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { var sub: [1100]u8 = undefined; const path = std.fmt.bufPrintSentinel(&sub, "{s}/{s}", .{ root, &x.name }, 0) catch continue; _ = files(path, true); - if (std.c.rmdir(path) == 0) held -|= x.bytes; + if (platform.c.rmdir(path) == 0) held -|= x.bytes; } return held; } @@ -389,27 +390,27 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { /// The bytes of the regular files directly in `dir`, each unlinked when `remove`. fn files(dir: []const u8, remove: bool) u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{dir}, 0) catch return 0) orelse return 0; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{dir}, 0) catch return 0) orelse return 0; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; var p: [1400]u8 = undefined; const f = std.fmt.bufPrintSentinel(&p, "{s}/{s}", .{ dir, std.mem.sliceTo(&ent.name, 0) }, 0) catch continue; - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) continue; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end > 0) n += @intCast(end); - if (remove) _ = std.c.unlink(f); + if (remove) _ = platform.c.unlink(f); } return n; } fn readWord(path: [:0]const u8) u64 { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return 0; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var w: u64 = 0; return if (readAt(fd, std.mem.asBytes(&w), 0)) w else 0; } @@ -417,19 +418,19 @@ fn readWord(path: [:0]const u8) u64 { fn writeWord(path: [:0]const u8, w: u64) void { var buf: [1200]u8 = undefined; const part = std.fmt.bufPrintSentinel(&buf, "{s}.part", .{path}, 0) catch return; - const fd = std.c.open(part, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(part, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return; - defer _ = std.c.close(fd); - defer _ = std.c.unlink(part); + defer _ = platform.c.close(fd); + defer _ = platform.c.unlink(part); writeAll(fd, std.mem.asBytes(&w)) catch return; - if (std.c.fsync(fd) != 0) return; - _ = std.c.rename(part, path); + if (platform.c.fsync(fd) != 0) return; + _ = platform.c.rename(part, path); } pub fn writeAll(fd: c_int, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, bytes.len - done); + const n = platform.c.write(fd, bytes.ptr + done, bytes.len - done); if (n <= 0) return error.ImprintWrite; done += @intCast(n); } @@ -438,7 +439,7 @@ pub fn writeAll(fd: c_int, bytes: []const u8) !void { pub fn readAt(fd: c_int, dest: []u8, at: u64) bool { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return false; done += @intCast(n); } @@ -448,13 +449,13 @@ pub fn readAt(fd: c_int, dest: []u8, at: u64) bool { /// This Mac's OS build (sysctl kern.osversion): its Metal compiler can change a kernel's bits. pub fn osBuild(buf: []u8) []const u8 { var n: usize = buf.len; - if (std.c.sysctlbyname("kern.osversion", buf.ptr, &n, null, 0) != 0) return ""; + if (platform.c.sysctlbyname("kern.osversion", buf.ptr, &n, null, 0) != 0) return ""; return std.mem.sliceTo(buf[0..n], 0); } /// Where learned states live unless --learn-dir says: $HOME/.cache/tensorfold/learned. pub fn defaultRoot(a: Allocator) ![]const u8 { - const home = std.c.getenv("HOME") orelse return error.ImprintDir; + const home = platform.c.getenv("HOME") orelse return error.ImprintDir; return std.fmt.allocPrint(a, "{s}/.cache/tensorfold/learned", .{std.mem.span(home)}); } @@ -464,11 +465,11 @@ fn makePath(gpa: Allocator, dir: [:0]const u8) !void { defer gpa.free(tmp); for (tmp[1..], 1..) |ch, i| if (ch == '/') { tmp[i] = 0; - _ = std.c.mkdir(tmp.ptr, 0o700); + _ = platform.c.mkdir(tmp.ptr, 0o700); tmp[i] = '/'; }; - const rc = std.c.mkdir(tmp.ptr, 0o700); - if (rc != 0 and std.c.errno(rc) != .EXIST) return error.ImprintDir; + const rc = platform.c.mkdir(tmp.ptr, 0o700); + if (rc != 0 and platform.c.errno(rc) != .EXIST) return error.ImprintDir; } test { @@ -479,19 +480,19 @@ test { test "a failed tail truncation keeps the bytes and refuses repair" { var buf: [128]u8 = undefined; - const path = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-index-repair-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(path); - const writer = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const path = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-index-repair-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(path); + const writer = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (writer < 0) return error.TestFile; { - defer _ = std.c.close(writer); + defer _ = platform.c.close(writer); try writeAll(writer, "record"); } - const reader = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const reader = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (reader < 0) return error.TestFile; - defer _ = std.c.close(reader); + defer _ = platform.c.close(reader); try std.testing.expectError(error.ImprintWrite, repairTail(reader, 0, 6)); - try std.testing.expectEqual(@as(i64, 6), std.c.lseek(reader, 0, std.c.SEEK.END)); + try std.testing.expectEqual(@as(i64, 6), platform.c.lseek(reader, 0, platform.c.SEEK.END)); var actual: [6]u8 = undefined; try std.testing.expect(readAt(reader, &actual, 0)); try std.testing.expectEqualStrings("record", &actual); diff --git a/zig/src/core/prompt_imprint_fault_test.zig b/zig/src/core/prompt_imprint_fault_test.zig index 418f69e0b5..a1d954bfd2 100644 --- a/zig/src/core/prompt_imprint_fault_test.zig +++ b/zig/src/core/prompt_imprint_fault_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const Imprint = @import("prompt_imprint.zig").Imprint; const imp = @import("prompt_imprint.zig"); const dirs = @import("lanes").learned_dirs; @@ -41,12 +42,12 @@ const Tracked = struct { } }; fn scratch(buf: []u8, suffix: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-imprint-fault-{s}-{d}", .{ suffix, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-imprint-fault-{s}-{d}", .{ suffix, platform.c.getpid() }); } fn word(path: [:0]const u8, n: u64) !void { - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imp.writeAll(fd, std.mem.asBytes(&n)); } test "Imprint failed open releases each allocation once after another identity scan" { @@ -58,8 +59,8 @@ test "Imprint failed open releases each allocation once after another identity s defer old.deinit(); var path: [256]u8 = undefined; const unreadable = try std.fmt.bufPrintSentinel(&path, "{s}/unreadable", .{old.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(unreadable, 0o700)); - defer _ = std.c.rmdir(unreadable); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(unreadable, 0o700)); + defer _ = platform.c.rmdir(unreadable); var tracked: Tracked = .{}; try std.testing.expectError(error.ImprintRead, Imprint.open(tracked.allocator(), root, 2, 1 << 20)); try std.testing.expectEqual(@as(usize, 0), tracked.duplicate_frees); @@ -92,8 +93,8 @@ test "failed atomic stamp replacement preserves the previous complete used word" var path: [256]u8 = undefined; try word(try std.fmt.bufPrintSentinel(&path, "{s}/used", .{m.dir}, 0), 42); const part = try std.fmt.bufPrintSentinel(&path, "{s}/used.part", .{m.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(part, 0o700)); - defer _ = std.c.rmdir(part); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(part, 0o700)); + defer _ = platform.c.rmdir(part); var reopened = try Imprint.open(a, root, 1, 1 << 20); defer reopened.deinit(); try std.testing.expectEqual(@as(u64, 42), try dirs.used(root, 1)); @@ -108,7 +109,7 @@ test "used stamp reads still refuse an actual I/O error" { var path: [256]u8 = undefined; const used = try std.fmt.bufPrintSentinel(&path, "{s}/used", .{m.dir}, 0); try dirs.unlink(used); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(used, 0o700)); - defer _ = std.c.rmdir(used); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(used, 0o700)); + defer _ = platform.c.rmdir(used); try std.testing.expectError(error.LearnedDirectoryRead, dirs.used(root, 1)); } diff --git a/zig/src/core/prompt_imprint_test.zig b/zig/src/core/prompt_imprint_test.zig index 66bdabd15e..f3e856e21b 100644 --- a/zig/src/core/prompt_imprint_test.zig +++ b/zig/src/core/prompt_imprint_test.zig @@ -1,5 +1,6 @@ //! Learned prompt states on disk: reopen, eviction order, removal and the sweep of other identities. const std = @import("std"); +const platform = @import("sys"); const imp = @import("prompt_imprint.zig"); const Imprint = imp.Imprint; @@ -7,38 +8,38 @@ const Imprint = imp.Imprint; pub fn rmTree(root: []const u8) void { var buf: [512]u8 = undefined; const r = std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return; - const d = std.c.opendir(r) orelse return; - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(r) orelse return; + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name[0] == '.') continue; var sub: [512]u8 = undefined; const p = std.fmt.bufPrintSentinel(&sub, "{s}/{s}", .{ root, name }, 0) catch continue; - if (std.c.unlink(p) == 0) continue; - if (std.c.opendir(p)) |s| { - while (std.c.readdir(s)) |f| { + if (platform.c.unlink(p) == 0) continue; + if (platform.c.opendir(p)) |s| { + while (platform.c.readdir(s)) |f| { const fname = std.mem.sliceTo(&f.name, 0); if (fname[0] == '.') continue; var fb: [600]u8 = undefined; - _ = std.c.unlink(std.fmt.bufPrintSentinel(&fb, "{s}/{s}", .{ p, fname }, 0) catch continue); + _ = platform.c.unlink(std.fmt.bufPrintSentinel(&fb, "{s}/{s}", .{ p, fname }, 0) catch continue); } - _ = std.c.closedir(s); + _ = platform.c.closedir(s); } - _ = std.c.rmdir(p); + _ = platform.c.rmdir(p); } - _ = std.c.closedir(d); - _ = std.c.rmdir(r); + _ = platform.c.closedir(d); + _ = platform.c.rmdir(r); } fn scratch(buf: []u8, tag: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-imprint-{s}-{d}", .{ tag, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-imprint-{s}-{d}", .{ tag, platform.c.getpid() }); } /// A file of `n` bytes at `dir`/`name`. fn fill(dir: []const u8, name: []const u8, n: usize) !void { var buf: [512]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var zeros: [256]u8 = @splat(0); var left = n; while (left > 0) : (left -= @min(left, zeros.len)) try imp.writeAll(fd, zeros[0..@min(left, zeros.len)]); @@ -46,8 +47,8 @@ fn fill(dir: []const u8, name: []const u8, n: usize) !void { fn exists(path: []const u8) bool { var buf: [512]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{path}, 0) catch return false) orelse return false; - _ = std.c.closedir(d); + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{path}, 0) catch return false) orelse return false; + _ = platform.c.closedir(d); return true; } diff --git a/zig/src/core/shard_edit.zig b/zig/src/core/shard_edit.zig index b340a40523..dff4f9dd5a 100644 --- a/zig/src/core/shard_edit.zig +++ b/zig/src/core/shard_edit.zig @@ -1,5 +1,6 @@ //! Replaced tensors written into a checkpoint's own shards, index and config: in place when sizes hold, else by rename. const std = @import("std"); +const platform = @import("sys"); const st = @import("safetensors.zig"); const Allocator = std.mem.Allocator; const Io = std.Io; @@ -200,18 +201,18 @@ fn inPlace(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacement) } var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); - const fd = std.c.open(try arena_state.allocator().dupeSentinel(u8, path, 0), .{ .ACCMODE = .WRONLY }); + const fd = platform.c.open(try arena_state.allocator().dupeSentinel(u8, path, 0), .{ .ACCMODE = .WRONLY }); if (fd < 0) return error.ShardWrite; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); for (edits, places[0..edits.len]) |r, at| { var done: usize = 0; while (done < r.bytes.len) { - const n = std.c.pwrite(fd, r.bytes.ptr + done, @min(r.bytes.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pwrite(fd, r.bytes.ptr + done, @min(r.bytes.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.ShardWrite; done += @intCast(n); } } - if (std.c.fsync(fd) != 0) return error.ShardWrite; + if (platform.c.fsync(fd) != 0) return error.ShardWrite; return true; } @@ -221,7 +222,8 @@ pub fn rewrite(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacem defer arena_state.deinit(); const arena = arena_state.allocator(); var file = try st.File.open(gpa, io, path); - defer file.close(io); + var mapped = true; + defer if (mapped) file.close(io); for (edits) |r| if (!file.names.contains(r.name)) return error.MissingTensor; const raw = file.map.memory; const original = try std.json.parseFromSliceLeaky(Value, arena, raw[8..file.data], .{}); @@ -257,7 +259,11 @@ pub fn rewrite(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacem const spaces = try arena.alloc(u8, padded - text.len); @memset(spaces, ' '); try pieces.insertSlice(arena, 0, &.{ &head, text, spaces }); - try writeFile(arena, path, pieces.items); + const tmp = try writeTemp(arena, path, pieces.items); + // the pieces are written; the map goes before the rename, since Windows refuses to replace a mapped file + file.close(io); + mapped = false; + try commit(arena, tmp, path); return @as(i64, @intCast(offset)) - @as(i64, @intCast(before)); } @@ -325,23 +331,32 @@ fn writeJson(arena: Allocator, path: []const u8, v: Value) !void { /// `parts` to `path` through `path`.tmp, synced, then renamed over it; a reader sees the old file or the new one. fn writeFile(arena: Allocator, path: []const u8, parts: []const []const u8) !void { + try commit(arena, try writeTemp(arena, path, parts), path); +} + +/// `parts` written and synced to `path`.tmp, which is returned. +fn writeTemp(arena: Allocator, path: []const u8, parts: []const []const u8) ![:0]const u8 { const tmp = try std.fmt.allocPrintSentinel(arena, "{s}.tmp", .{path}, 0); - const dest = try arena.dupeSentinel(u8, path, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o644)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o644)); if (fd < 0) return error.ShardWrite; - errdefer _ = std.c.unlink(tmp); - { - defer _ = std.c.close(fd); - for (parts) |p| try put(fd, p); - if (std.c.fsync(fd) != 0) return error.ShardWrite; - } - if (std.c.rename(tmp, dest) != 0) return error.ShardWrite; + errdefer _ = platform.c.unlink(tmp); + defer _ = platform.c.close(fd); + for (parts) |p| try put(fd, p); + if (platform.c.fsync(fd) != 0) return error.ShardWrite; + return tmp; +} + +/// `tmp` renamed over `path`; a failed rename removes `tmp`. +fn commit(arena: Allocator, tmp: [:0]const u8, path: []const u8) !void { + errdefer _ = platform.c.unlink(tmp); + const dest = try arena.dupeSentinel(u8, path, 0); + if (platform.c.rename(tmp, dest) != 0) return error.ShardWrite; } fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.ShardWrite; done += @intCast(n); } @@ -518,13 +533,19 @@ test "a folder whose shard is a symlink or a hard link is refused for learning; defer tmp.cleanup(); try tmp.dir.createDirPath(io, "blobs"); for ([_][]const u8{ "config.json", index_name, "model-00001-of-00002.safetensors", "blobs/b2", "notes.txt" }) |name| try tmp.dir.writeFile(io, .{ .sub_path = name, .data = "x" }); - try tmp.dir.symLink(io, "model-00001-of-00002.safetensors", "readme-link", .{}); + // Windows grants symlinks only with Developer Mode or elevation + tmp.dir.symLink(io, "model-00001-of-00002.safetensors", "readme-link", .{}) catch |e| return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; const dir = try std.fs.path.join(arena, &.{ ".zig-cache", "tmp", &tmp.sub_path }); try std.testing.expect((try linked(arena, io, dir)) == null); try tmp.dir.symLink(io, "blobs/b2", "model-00002-of-00002.safetensors", .{}); try std.testing.expectEqualStrings("model-00002-of-00002.safetensors", (try linked(arena, io, dir)).?); try tmp.dir.deleteFile(io, "model-00002-of-00002.safetensors"); - try tmp.dir.hardLink("blobs/b2", tmp.dir, "model-00002-of-00002.safetensors", io, .{}); + if (@import("builtin").os.tag == .windows) { + // std's Io has no hard links on Windows; the platform's link (CreateHardLinkW) is what pull uses there + const old = try std.fs.path.joinZ(arena, &.{ dir, "blobs", "b2" }); + const new = try std.fs.path.joinZ(arena, &.{ dir, "model-00002-of-00002.safetensors" }); + try std.testing.expectEqual(@as(c_int, 0), platform.c.link(old, new)); + } else try tmp.dir.hardLink("blobs/b2", tmp.dir, "model-00002-of-00002.safetensors", io, .{}); try std.testing.expectEqualStrings("model-00002-of-00002.safetensors", (try linked(arena, io, dir)).?); } diff --git a/zig/src/core/snapshot_file.zig b/zig/src/core/snapshot_file.zig index 609916ed99..f56997b862 100644 --- a/zig/src/core/snapshot_file.zig +++ b/zig/src/core/snapshot_file.zig @@ -1,5 +1,6 @@ //! Learned snapshots carry a SHA-256 of their entire header and payload; exact length is required on admission. const std = @import("std"); +const platform = @import("sys"); const Hash = std.crypto.hash.sha2.Sha256; @@ -15,7 +16,7 @@ pub fn write(fd: c_int, header: []const u8, payload: []const u8) !void { pub fn read(fd: c_int, header: []const u8, payload: []u8) !void { const length = std.math.add(usize, header.len, payload.len) catch return error.SnapshotRead; const total = std.math.add(usize, length, 32) catch return error.SnapshotRead; - if (std.c.lseek(fd, 0, std.c.SEEK.END) != total) return error.SnapshotRead; + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != total) return error.SnapshotRead; if (!get(fd, payload, header.len)) return error.SnapshotRead; var stored: [32]u8 = undefined; if (!get(fd, &stored, length)) return error.SnapshotRead; @@ -27,21 +28,21 @@ pub fn read(fd: c_int, header: []const u8, payload: []u8) !void { test "payload bit flips, header changes, truncation, legacy files and extra bytes are refused" { var path: [128]u8 = undefined; - const name = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-snapshot-test-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(name); - const fd = std.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const name = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-snapshot-test-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(name); + const fd = platform.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var output: [4]u8 = undefined; try write(fd, "head", "body"); try read(fd, "head", &output); try std.testing.expectEqualStrings("body", &output); try std.testing.expectError(error.SnapshotRead, read(fd, "dead", &output)); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, 5)); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, 5)); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 8)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 8)); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); - _ = std.c.lseek(fd, 0, std.c.SEEK.SET); + _ = platform.c.lseek(fd, 0, platform.c.SEEK.SET); try write(fd, "head", "body"); try put(fd, "x"); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); @@ -50,7 +51,7 @@ test "payload bit flips, header changes, truncation, legacy files and extra byte fn put(fd: c_int, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30)); + const n = platform.c.write(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -58,7 +59,7 @@ fn put(fd: c_int, bytes: []const u8) !void { fn get(fd: c_int, bytes: []u8, at: usize) bool { var done: usize = 0; while (done < bytes.len) { - const n = std.c.pread(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return false; done += @intCast(n); } diff --git a/zig/src/core/template/methods.zig b/zig/src/core/template/methods.zig index 2ec14e09e1..ecccb1497e 100644 --- a/zig/src/core/template/methods.zig +++ b/zig/src/core/template/methods.zig @@ -1,5 +1,6 @@ //! Python methods templates call on strings, dicts and lists, the loop helpers, and the template globals. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const v = @import("value.zig"); const uni = @import("unicode.zig"); @@ -432,14 +433,22 @@ pub fn callFunc(rt: *Rt, f: v.Func, pos: []const Value, kw: []const Kw) Error!Va } const Tm = extern struct { sec: c_int, min: c_int, hour: c_int, mday: c_int, mon: c_int, year: c_int, wday: c_int, yday: c_int, isdst: c_int, gmtoff: c_long, zone: ?[*:0]const u8 }; -extern "c" fn localtime_r(t: *const std.c.time_t, out: *Tm) ?*Tm; +const localtime_r = if (builtin.os.tag == .windows) localtimeWindows else struct { + extern "c" fn localtime_r(t: *const platform.c.time_t, out: *Tm) ?*Tm; +}.localtime_r; +/// The UCRT's localtime_s fills the nine ISO fields; gmtoff and zone are POSIX's and stay unset. +extern "c" fn _localtime64_s(out: *Tm, t: *const i64) c_int; +fn localtimeWindows(t: *const platform.c.time_t, out: *Tm) ?*Tm { + out.* = std.mem.zeroes(Tm); + return if (_localtime64_s(out, t) == 0) out else null; +} extern "c" fn strftime(buf: [*]u8, max: usize, fmt: [*:0]const u8, tm: *const Tm) usize; /// datetime.now().strftime(): %f, %z and %Z are Python's own (naive time), the rest is the C library's. fn strftimeNow(rt: *Rt, fmt: []const u8) Error![]const u8 { if (!builtin.link_libc) return rt.fail("strftime_now needs libc", .{}); - var tv: std.c.timeval = undefined; - _ = std.c.gettimeofday(&tv, null); + var tv: platform.c.timeval = undefined; + _ = platform.c.gettimeofday(&tv, null); var tm: Tm = undefined; if (localtime_r(&tv.sec, &tm) == null) return rt.fail("localtime failed", .{}); var pre: std.ArrayList(u8) = .empty; diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index cc733b8537..9e1f154172 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -1,6 +1,7 @@ //! cuBLASLt opened at run time: descriptors, layouts, heuristics and matmul, with fp32 compute only (no TF32 or fast modes). const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Status = c_int; @@ -54,7 +55,7 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, /// The CUDA 13 soname first, then the unversioned link name. @@ -69,7 +70,7 @@ pub const Library = struct { } pub fn openPath(path: []const u8) Error!Library { - var lib = std.DynLib.open(path) catch return error.LibraryUnavailable; + var lib = platform.DynLib.open(path) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..1d644f3742 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -1,21 +1,22 @@ //! The CUDA driver opened at run time: libcuda.so.1's versioned entry points in one table, and checked calls. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDeviceMemory, NotReady, NotFound, Invalid }; pub const Driver = struct { - lib: std.DynLib, + lib: platform.DynLib, api: abi.Api, pub fn open() Error!Driver { - return openPath("libcuda.so.1"); + return openPath(platform.cuda_driver); } /// Resolves every field of `abi.Api` by its exact name; a missing symbol refuses the whole driver. pub fn openPath(path: []const u8) Error!Driver { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); var api: abi.Api = undefined; const info = @typeInfo(abi.Api).@"struct"; diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..5bc478c6b5 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -1,6 +1,7 @@ //! NCCL opened at run time for tensor-parallel ranks: communicators and the collectives the lane engines use. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Result = c_int; @@ -30,7 +31,7 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn open() Error!Library { @@ -38,7 +39,7 @@ pub const Library = struct { } pub fn openPath(path: []const u8) Error!Library { - var lib = std.DynLib.open(path) catch return error.LibraryUnavailable; + var lib = platform.DynLib.open(path) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; diff --git a/zig/src/fabric/mailbox.zig b/zig/src/fabric/mailbox.zig index 3010d50af6..49b08f4b4b 100644 --- a/zig/src/fabric/mailbox.zig +++ b/zig/src/fabric/mailbox.zig @@ -1,5 +1,6 @@ //! One mcdma-rpcd mailbox (protocol 1): the mapped halves, a client's ordered call and a service's request and reply. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const layout = @import("layout.zig"); const words = @import("words.zig"); @@ -48,32 +49,32 @@ pub const Mailbox = struct { fn openShm(name: []const u8) Error!Mailbox { var path: [64]u8 = undefined; const z = std.mem.printSentinel(&path, "/mcdma-rpc.{s}", .{name}, 0) catch return error.InvalidName; - const flags: c_int = @bitCast(std.c.O{ .ACCMODE = .RDWR }); - const fd = if (builtin.os.tag == .macos) std.c.shm_open(z.ptr, flags, @as(c_uint, 0)) else std.c.shm_open(z.ptr, flags, 0); + const flags: c_int = @bitCast(platform.c.O{ .ACCMODE = .RDWR }); + const fd = if (builtin.os.tag == .macos) platform.c.shm_open(z.ptr, flags, @as(c_uint, 0)) else platform.c.shm_open(z.ptr, flags, 0); if (fd < 0) return error.NoMailbox; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); return mapFd(fd); } /// The listen end's mailbox: POSIX shared memory on a Mac, the file /dev/shm/mcdma-rpc.NAME on Linux; MCDMA_RPC_BOX_DIR names a test directory. pub fn openListen(name: []const u8) Error!Mailbox { if (!layout.validName(name)) return error.InvalidName; - const env = std.c.getenv("MCDMA_RPC_BOX_DIR"); + const env = platform.c.getenv("MCDMA_RPC_BOX_DIR"); if (env == null and builtin.os.tag == .macos) return openShm(name); const dir = if (env) |d| std.mem.span(d) else "/dev/shm"; var path: [256]u8 = undefined; const z = std.mem.printSentinel(&path, "{s}/mcdma-rpc.{s}", .{ dir, name }, 0) catch return error.InvalidName; - const fd = std.c.open(z.ptr, .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .CLOEXEC = true }); + const fd = platform.c.open(z.ptr, .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .CLOEXEC = true }); if (fd < 0) return error.NoMailbox; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); return mapFd(fd); } /// Map a mailbox the daemon created; it must be this user's and owner-only, as the daemons make them. - fn mapFd(fd: std.c.fd_t) Error!Mailbox { - var st: std.c.Stat = undefined; - if (std.c.fstat(fd, &st) != 0) return error.NoMailbox; - if (st.uid != std.c.getuid() or st.mode & 0o077 != 0) return error.UnsafeMailbox; + fn mapFd(fd: platform.c.fd_t) Error!Mailbox { + var st: platform.c.Stat = undefined; + if (platform.c.fstat(fd, &st) != 0) return error.NoMailbox; + if (st.uid != platform.c.getuid() or st.mode & 0o077 != 0) return error.UnsafeMailbox; const size: usize = @intCast(st.size); if (size < 2 * layout.segment) return error.BadSizes; const mem = std.posix.mmap(null, size, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0) catch return error.NoMailbox; @@ -210,14 +211,14 @@ pub fn allocFake(req: usize, rep: usize) ![]align(page) u8 { } test "a listen mailbox on this Mac is POSIX shared memory, as MCDMA 2.0.0's Mac listen end makes it" { - if (builtin.os.tag != .macos or std.c.getenv("MCDMA_RPC_BOX_DIR") != null) return error.SkipZigTest; + if (builtin.os.tag != .macos or platform.c.getenv("MCDMA_RPC_BOX_DIR") != null) return error.SkipZigTest; const name = "/mcdma-rpc.tf-zig-shm-test"; - const flags: c_int = @bitCast(std.c.O{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true }); - const fd = std.c.shm_open(name, flags, @as(c_uint, 0o600)); + const flags: c_int = @bitCast(platform.c.O{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true }); + const fd = platform.c.shm_open(name, flags, @as(c_uint, 0o600)); try std.testing.expect(fd >= 0); - defer _ = std.c.shm_unlink(name); - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 2 * layout.segment)); + defer _ = platform.c.shm_unlink(name); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 2 * layout.segment)); const mem = try std.posix.mmap(null, 2 * layout.segment, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); std.mem.writeInt(u64, mem[layout.sizes..][0..8], layout.segment, .little); std.mem.writeInt(u64, mem[layout.sizes + 8 ..][0..8], layout.segment, .little); diff --git a/zig/src/fabric/mcdma.zig b/zig/src/fabric/mcdma.zig index 290bcc1c61..d3292477a9 100644 --- a/zig/src/fabric/mcdma.zig +++ b/zig/src/fabric/mcdma.zig @@ -1,5 +1,6 @@ //! Registered windows over first-party MCDMA: one progress context per port, with staged sources held through flush. const std = @import("std"); +const platform = @import("sys"); const abi = @import("mcdma_abi.zig"); const rma = @import("rdma.zig"); const words = @import("words.zig"); @@ -367,7 +368,7 @@ test "an endpoint gives the library's placing thread user-interactive QoS unless for ([_]?[*:0]const u8{ null, "0", "1" }) |set| { _ = if (set) |v| setenv("MCDMA_FABRIC_QOS", v, 1) else unsetenv("MCDMA_FABRIC_QOS"); try std.testing.expectError(error.LibraryUnavailable, Endpoint.create(std.testing.allocator, "/nonexistent/libmcdma-fabric.dylib", config)); - try std.testing.expectEqualStrings(if (set) |v| std.mem.span(v) else "1", if (std.c.getenv("MCDMA_FABRIC_QOS")) |v| std.mem.span(v) else "(unset)"); + try std.testing.expectEqualStrings(if (set) |v| std.mem.span(v) else "1", if (platform.c.getenv("MCDMA_FABRIC_QOS")) |v| std.mem.span(v) else "(unset)"); } } diff --git a/zig/src/fabric/mcdma_abi.zig b/zig/src/fabric/mcdma_abi.zig index 0b4bcf062c..feeacbb75c 100644 --- a/zig/src/fabric/mcdma_abi.zig +++ b/zig/src/fabric/mcdma_abi.zig @@ -1,5 +1,6 @@ //! First-party libmcdma-fabric ABI 1, loaded only from an explicit absolute library path. const std = @import("std"); +const platform = @import("sys"); pub const Context = opaque {}; pub const Peer = opaque {}; @@ -27,7 +28,7 @@ pub const MaxLinks = *const fn () callconv(.c) c_uint; pub const ws_room: usize = 64; pub const Library = struct { - handle: std.DynLib, + handle: platform.DynLib, api: Api, write_signal: ?WriteSignal = null, connect_links: ?ConnectLinks = null, @@ -36,7 +37,7 @@ pub const Library = struct { /// No install-path or loader-search fallback; the caller supplies the inspected user-directory library. pub fn open(path: [:0]const u8) Error!Library { if (!std.fs.path.isAbsolute(path)) return error.InvalidPath; - var handle = std.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; + var handle = platform.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; errdefer handle.close(); var api: Api = undefined; inline for (@typeInfo(Api).@"struct".field_names) |name| { diff --git a/zig/src/fabric/registration.zig b/zig/src/fabric/registration.zig index 2bbaf084b4..11a18e92cc 100644 --- a/zig/src/fabric/registration.zig +++ b/zig/src/fabric/registration.zig @@ -1,13 +1,14 @@ //! A service's registration with its listen daemon ("MODE poll"), and the daemons' STATUS lines, as protocol 1 defines. const std = @import("std"); +const platform = @import("sys"); const layout = @import("layout.zig"); pub const Error = error{ InvalidName, NoDaemon, Busy, Refused, Io }; /// The registration channel: the daemon sends nothing after OK until the link drops (BYE or a closed socket). pub const Registration = struct { - in_fd: std.c.fd_t, - out_fd: std.c.fd_t, + in_fd: platform.c.fd_t, + out_fd: platform.c.fd_t, /// Register with the listen daemon at MCDMA_RPCD_SOCKET, else /tmp/mcdma-rpcd.NAME.sock. pub fn attach(name: []const u8) Error!Registration { @@ -35,21 +36,21 @@ pub const Registration = struct { /// Whether the daemon still holds the registration: anything readable means BYE or a closed socket. pub fn alive(r: *const Registration) bool { - var fds = [_]std.c.pollfd{.{ .fd = r.in_fd, .events = std.c.POLL.IN, .revents = 0 }}; - const n = std.c.poll(&fds, 1, 0); + var fds = [_]platform.c.pollfd{.{ .fd = r.in_fd, .events = platform.c.POLL.IN, .revents = 0 }}; + const n = platform.c.poll(&fds, 1, 0); return n == 0; } pub fn close(r: *Registration) void { - _ = std.c.close(r.in_fd); - if (r.out_fd != r.in_fd) _ = std.c.close(r.out_fd); + _ = platform.c.close(r.in_fd); + if (r.out_fd != r.in_fd) _ = platform.c.close(r.out_fd); r.* = undefined; } }; /// The listen daemon's socket for link `name`, honouring MCDMA_RPCD_SOCKET as mcdma-rpcd does. pub fn socketPath(buf: []u8, name: []const u8) ![:0]const u8 { - if (std.c.getenv("MCDMA_RPCD_SOCKET")) |env| { + if (platform.c.getenv("MCDMA_RPCD_SOCKET")) |env| { const s = std.mem.span(env); if (s.len > 0) return std.mem.printSentinel(buf, "{s}", .{s}, 0); } @@ -57,35 +58,35 @@ pub fn socketPath(buf: []u8, name: []const u8) ![:0]const u8 { } /// Connect to the daemon's local Unix socket; the only socket the fabric opens, and never TCP. -fn connectUnix(path: [:0]const u8) ?std.c.fd_t { - var addr: std.c.sockaddr.un = .{ .path = undefined }; +fn connectUnix(path: [:0]const u8) ?platform.c.fd_t { + var addr: platform.c.sockaddr.un = .{ .path = undefined }; if (path.len >= addr.path.len) return null; @memset(&addr.path, 0); @memcpy(addr.path[0..path.len], path); - const fd = std.c.socket(std.c.AF.UNIX, std.c.SOCK.STREAM, 0); + const fd = platform.c.socket(platform.c.AF.UNIX, platform.c.SOCK.STREAM, 0); if (fd < 0) return null; - if (std.c.connect(fd, @ptrCast(&addr), @sizeOf(std.c.sockaddr.un)) != 0) { - _ = std.c.close(fd); + if (platform.c.connect(fd, @ptrCast(&addr), @sizeOf(platform.c.sockaddr.un)) != 0) { + _ = platform.c.close(fd); return null; } return fd; } -fn writeAll(fd: std.c.fd_t, bytes: []const u8) Error!void { +fn writeAll(fd: platform.c.fd_t, bytes: []const u8) Error!void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes[done..].ptr, bytes.len - done); + const n = platform.c.write(fd, bytes[done..].ptr, bytes.len - done); if (n <= 0) return error.Io; done += @intCast(n); } } /// Read one line a byte at a time, as the daemons write them; the newline is dropped. -fn readLine(fd: std.c.fd_t, buf: []u8) Error![]const u8 { +fn readLine(fd: platform.c.fd_t, buf: []u8) Error![]const u8 { var n: usize = 0; while (n < buf.len) { var c: [1]u8 = undefined; - if (std.c.read(fd, &c, 1) != 1) return error.Io; + if (platform.c.read(fd, &c, 1) != 1) return error.Io; if (c[0] == '\n') return buf[0..n]; buf[n] = c[0]; n += 1; @@ -139,9 +140,9 @@ pub fn parsePeer(line: []const u8) ?Peer { return p; } -fn pipePair() ![2]std.c.fd_t { - var fds: [2]std.c.fd_t = undefined; - if (std.c.pipe(&fds) != 0) return error.Io; +fn pipePair() ![2]platform.c.fd_t { + var fds: [2]platform.c.fd_t = undefined; + if (platform.c.pipe(&fds) != 0) return error.Io; return fds; } @@ -149,14 +150,14 @@ test "registration answers OK, ERR busy or anything else, and BYE ends it" { for ([_][]const u8{ "OK\n", "ERR busy\n", "ERR unknown command\n" }, 0..) |answer, i| { const to_daemon = try pipePair(); const from_daemon = try pipePair(); - defer for ([_]std.c.fd_t{ to_daemon[0], to_daemon[1], from_daemon[0], from_daemon[1] }) |fd| { - _ = std.c.close(fd); + defer for ([_]platform.c.fd_t{ to_daemon[0], to_daemon[1], from_daemon[0], from_daemon[1] }) |fd| { + _ = platform.c.close(fd); }; try writeAll(from_daemon[1], answer); const r: Registration = .{ .in_fd = from_daemon[0], .out_fd = to_daemon[1] }; const got = r.handshake(); var sent: [16]u8 = undefined; - try std.testing.expectEqual(@as(isize, 10), std.c.read(to_daemon[0], &sent, sent.len)); + try std.testing.expectEqual(@as(isize, 10), platform.c.read(to_daemon[0], &sent, sent.len)); try std.testing.expectEqualStrings("MODE poll\n", sent[0..10]); switch (i) { 0 => { diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 1da50d386d..9b9353c554 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,5 +1,6 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const abi = @import("verbs_abi.zig"); @@ -52,13 +53,13 @@ const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; pub const Verbs = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, pub fn open(path: ?[:0]const u8) Error!Verbs { - var lib = std.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; + var lib = platform.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; @@ -195,7 +196,7 @@ test "a data-device GID scan skips zero entries and prefers link-local over an I } test "verbs smoke: open the device list and close it (TF_FABRIC_VERBS=1)" { - const flag = std.c.getenv("TF_FABRIC_VERBS") orelse return error.SkipZigTest; + const flag = platform.c.getenv("TF_FABRIC_VERBS") orelse return error.SkipZigTest; if (flag[0] != '1') return error.SkipZigTest; var v = Verbs.open(null) catch |err| { std.debug.print("verbs: {s} unavailable ({s})\n", .{ default_path, @errorName(err) }); diff --git a/zig/src/fabric/verbs_abi.zig b/zig/src/fabric/verbs_abi.zig index 1941fc3157..801fa9d80a 100644 --- a/zig/src/fabric/verbs_abi.zig +++ b/zig/src/fabric/verbs_abi.zig @@ -1,5 +1,6 @@ //! The verbs ABI (rdma-core's infiniband/verbs.h, which Apple's SDK ships for librdma): the structs and constants we use. const std = @import("std"); +const platform = @import("sys"); pub const Device = extern struct { _ops: [2]?*anyopaque, @@ -22,8 +23,8 @@ pub const Qp = extern struct { qp_num: u32, state: c_int, qp_type: c_int, - mutex: std.c.pthread_mutex_t, - cond: std.c.pthread_cond_t, + mutex: platform.c.pthread_mutex_t, + cond: platform.c.pthread_cond_t, events_completed: u32, }; @@ -33,8 +34,8 @@ pub const Cq = extern struct { cq_context: ?*anyopaque, handle: u32, cqe: c_int, - mutex: std.c.pthread_mutex_t, - cond: std.c.pthread_cond_t, + mutex: platform.c.pthread_mutex_t, + cond: platform.c.pthread_cond_t, comp_events_completed: u32, async_events_completed: u32, }; @@ -59,7 +60,7 @@ pub const Context = extern struct { cmd_fd: c_int, async_fd: c_int, num_comp_vectors: c_int, - mutex: std.c.pthread_mutex_t, + mutex: platform.c.pthread_mutex_t, abi_compat: ?*anyopaque, }; diff --git a/zig/src/fabric/words.zig b/zig/src/fabric/words.zig index 53d41a0f58..f9c749fc72 100644 --- a/zig/src/fabric/words.zig +++ b/zig/src/fabric/words.zig @@ -1,5 +1,6 @@ //! Ordered mailbox words: MCDMA's libmcdma-rpc (ABI 1, rpc/mcdma_rpc.h) by dlopen, or its same rules in Zig. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const layout = @import("layout.zig"); @@ -30,20 +31,20 @@ pub fn load(w: *const u64) u64 { /// CLOCK_MONOTONIC in nanoseconds, the clock libmcdma-rpc waits on. pub fn nowNs() u64 { - var ts: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &ts); + var ts: platform.c.timespec = undefined; + _ = platform.c.clock_gettime(.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); } /// Sleep about `ns` nanoseconds (the library's short step between polls). pub fn pause(ns: u64) void { - const ts: std.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } pub fn checkedNowNs() error{ClockUnavailable}!u64 { - var ts: std.c.timespec = undefined; - if (std.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0 or ts.nsec < 0 or ts.nsec >= std.time.ns_per_s) return error.ClockUnavailable; + var ts: platform.c.timespec = undefined; + if (platform.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0 or ts.nsec < 0 or ts.nsec >= std.time.ns_per_s) return error.ClockUnavailable; return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); } @@ -81,7 +82,7 @@ else /// libmcdma-rpc opened at run time; the three Metal functions are optional, as the header says. pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, abi_fn: AbiFn, wait_word: WaitFn, store_word: StoreFn, @@ -92,7 +93,7 @@ pub const Library = struct { /// Open `path`, else MCDMA_RPC_LIBRARY, else the install prefix; refuse any ABI but 1. pub fn open(path: ?[:0]const u8) Error!Library { if (path) |p| return openPath(p); - if (std.c.getenv("MCDMA_RPC_LIBRARY")) |env| { + if (platform.c.getenv("MCDMA_RPC_LIBRARY")) |env| { const p = std.mem.span(env); if (p.len > 0) return openPath(p); } @@ -104,7 +105,7 @@ pub const Library = struct { } pub fn openPath(path: [:0]const u8) Error!Library { - var lib = std.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; + var lib = platform.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; errdefer lib.close(); var l: Library = .{ .lib = lib, diff --git a/zig/src/families/flashnext/call_log.zig b/zig/src/families/flashnext/call_log.zig index 387ec89220..5afdca5da9 100644 --- a/zig/src/families/flashnext/call_log.zig +++ b/zig/src/families/flashnext/call_log.zig @@ -1,5 +1,6 @@ //! FZ_CALL_LOG: a prompt pass's calls (rows, parts, wall and GPU milliseconds), logged in one line after its first token. const std = @import("std"); +const platform = @import("sys"); const N = 16; @@ -17,7 +18,7 @@ pub const CallLog = struct { if (c.n == N) return; c.rows[c.n] = @intCast(rows); c.parts[c.n] = @intCast(parts); - c.wall[c.n] = @floatCast(@as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6); // mach ticks: 125/3 ns + c.wall[c.n] = @floatCast(@as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6); // mach ticks: 125/3 ns c.gpu[c.n] = @floatCast((g1 - g0) * 1e3); c.n += 1; } diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index c9f8481249..9ff39b2845 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -1,6 +1,7 @@ //! Flash Next served from the replay engine: one reply at a time, its prompt through prompt chunks (the MTP head's //! keys for every prompt row), then GPU-side rounds whose drafts the depth rule sets; tokens leave as rounds are read. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fz = @import("replay.zig"); const config = @import("config.zig"); @@ -151,8 +152,8 @@ pub const Engine = struct { e.copy_long = 6; e.segments = true; e.mark_taps = true; - e.pair_min = tp_settings.pairMin(if (std.c.getenv("FZ_PAIR_MIN")) |v| std.mem.span(v) else null, PAIR_MIN); // rank 1 takes rank 0's - e.call_log = std.c.getenv("FZ_CALL_LOG") != null; + e.pair_min = tp_settings.pairMin(if (platform.c.getenv("FZ_PAIR_MIN")) |v| std.mem.span(v) else null, PAIR_MIN); // rank 1 takes rank 0's + e.call_log = platform.c.getenv("FZ_CALL_LOG") != null; e.passed = null; e.snap_pool = .{}; e.handoff_ns = 0; @@ -177,8 +178,8 @@ pub const Engine = struct { r.event = try device.sharedEvent(); if (dump_dir) |d| try r.compile(d) else try r.compileChecked(); r.sel = try Select.init(r, MAXR); - r.lane_new = std.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (the recorded bits) - if (std.c.getenv("FZ_GDN")) |v| { + r.lane_new = platform.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (the recorded bits) + if (platform.c.getenv("FZ_GDN")) |v| { r.gdn_pipe = try fz.gdn_step.compile(r, false); if (v[0] == '2') r.gdn_kept = try fz.gdn_step.compile(r, true); // GPU-side rounds keep one state a layer } @@ -211,7 +212,7 @@ pub const Engine = struct { }; e.identity_hash = std.hash.Wyhash.hash(0x6964, identity); const owned_pack_dir = if (dump_dir == null) blk: { - const home = if (std.c.getenv("HOME")) |v| std.mem.span(v) else ""; + const home = if (platform.c.getenv("HOME")) |v| std.mem.span(v) else ""; const root = try pack_cache.defaultRoot(arena, home); break :blk try pack_cache.ensure(gpa, io, model_dir, root, identity, PackBuilder{}); } else null; @@ -434,7 +435,7 @@ pub const Engine = struct { e.tsel = try GSelect.init(r, 0); e.msel = try GSelect.init(r, 0); // speed-up mode: TP=2 with the peer the settings name (tp.zig); both ranks then run the same replies in lockstep - const tp_path: ?[]const u8 = speed_up orelse if (std.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; + const tp_path: ?[]const u8 = speed_up orelse if (platform.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; if (tp_path != null and !uniform) return error.SpeedUpNeeds6Bit32; // TP splits the 6/32 kernels only if (tp_path) |path| r.tp = try fz.Tp2.init(arena, r.device, path); if (r.tp == null) try fz.rows_load.attach( @@ -504,8 +505,8 @@ pub const Engine = struct { const pool = mtl.objc.Pool.push(); defer pool.pop(); if (r.tp) |tp| if (tp.rank == 0) { // served speed-up mode: rank 1 runs the same request (follow) - const t0 = std.c.mach_absolute_time(); - defer e.handoff_ns = @intCast((std.c.mach_absolute_time() - t0) * 125 / 3); // mach ticks are 125/3 ns on Apple silicon + const t0 = platform.c.mach_absolute_time(); + defer e.handoff_ns = @intCast((platform.c.mach_absolute_time() - t0) * 125 / 3); // mach ticks are 125/3 ns on Apple silicon try follow_mod.send(e, tp, prompt, from, marks, max_tokens, eos, depth, budget); if (!try tp.waitAck()) return error.PeerNotResumed; }; @@ -522,7 +523,7 @@ pub const Engine = struct { var calls: CallLog = .{}; while (at < prompt.len) { if (try e.agree(out.cancelled(out.ctx))) return .{ .reason = .cancelled }; - const t_call = std.c.mach_absolute_time(); + const t_call = platform.c.mach_absolute_time(); const g_call = m.gpu_seconds; if (r.tp) |tp| { // speed-up mode: the call's rows split across the two Macs, running on past up to MARKS marks const pair_taps = e.mark_taps and m.marks != null; diff --git a/zig/src/families/flashnext/follow.zig b/zig/src/families/flashnext/follow.zig index 4255faeb6e..482870c098 100644 --- a/zig/src/families/flashnext/follow.zig +++ b/zig/src/families/flashnext/follow.zig @@ -1,5 +1,6 @@ //! Speed-up mode's rank 1: rank 0's requests run here in step, and this Mac keeps its own states under rank 0's names. const std = @import("std"); +const platform = @import("sys"); const engine = @import("engine.zig"); const snapshot = @import("snapshot.zig"); const fz = @import("replay.zig"); @@ -9,7 +10,7 @@ const Engine = engine.Engine; /// Rank 0's requests as they come, until its empty request; between them a free buffer is readied inside the budget. pub fn run(e: *Engine) !void { while (try one(e, null)) |_| { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const got = e.snap_pool.ready(e.r.device, engine.NEXT_TURN, .{ .ctx = e.r.tp.?, .check = requestWaiting }, e.peer_budget -| e.peer_held -| e.snap_pool.spare()); if (got.cap > 0) std.log.info("speed-up rank 1: readied {d} of {d} MiB for the next save in {d:.1} ms", .{ got.touched >> 20, got.cap >> 20, msSince(t0) }); } @@ -21,7 +22,7 @@ fn requestWaiting(ctx: *anyopaque) bool { } fn msSince(t0: u64) f64 { - return @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6; // mach ticks are 125/3 ns on Apple silicon + return @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6; // mach ticks are 125/3 ns on Apple silicon } /// Rank 0's request for `one`: a 16-word head, then the prompt, marks, dropped states' keys and the budget. @@ -65,7 +66,7 @@ pub fn one(e: *Engine, out: ?engine.Out) !?engine.Result { const from = head[12]; if (head[15] == 0) return error.TpProtocol; // rank 0 sends its positive pair minimum with every request e.pair_min = head[15]; - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const ok = from == 0 or blk: { const st = e.peer_kept.get(engine.keyOf(prompt[0 .. from + 1])) orelse break :blk false; snapshot.restore(e, st) catch break :blk false; @@ -123,7 +124,7 @@ const Reply = struct { } fn marked(ctx: *anyopaque, at: usize) void { const r: *Reply = @ptrCast(@alignCast(ctx)); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); defer { r.kept += 1; r.keep_ns += msSince(t0) * 1e6; diff --git a/zig/src/families/flashnext/host_test.zig b/zig/src/families/flashnext/host_test.zig index bc198ce4ad..642812ce54 100644 --- a/zig/src/families/flashnext/host_test.zig +++ b/zig/src/families/flashnext/host_test.zig @@ -1,5 +1,6 @@ //! Host-only contracts exercise strict text configuration, affine six-bit packing and borrowed-view shapes. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; const family = @import("flashnext.zig"); const st = @import("../../core/safetensors.zig"); @@ -298,12 +299,12 @@ test "the prompt's checked-in text resolves and patches from another working dir const sources = @import("kernel_sources").flashnext_gen; var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); - const old_fd = std.c.open(".", .{}); // O_RDONLY via the flags struct's default - defer _ = std.c.close(old_fd); + const old_fd = platform.c.open(".", .{}); // O_RDONLY via the flags struct's default + defer _ = platform.c.close(old_fd); const tmpz = try std.fmt.allocPrintSentinel(std.testing.allocator, ".zig-cache/tmp/{s}", .{tmp.sub_path}, 0); defer std.testing.allocator.free(tmpz); - try std.testing.expect(std.c.chdir(tmpz) == 0); - defer _ = std.c.fchdir(old_fd); + try std.testing.expect(platform.c.chdir(tmpz) == 0); + defer _ = platform.c.fchdir(old_fd); for (table.entries) |e| { var text: []const u8 = ""; for (sources.sources) |s| { diff --git a/zig/src/families/flashnext/learned.zig b/zig/src/families/flashnext/learned.zig index 5d1485abdf..262f44f447 100644 --- a/zig/src/families/flashnext/learned.zig +++ b/zig/src/families/flashnext/learned.zig @@ -1,5 +1,6 @@ //! Flash Next's learned prompt states (--learn): a kept state's file on disk, and the probe keying this build's bits. const std = @import("std"); +const platform = @import("sys"); const ks = @import("kernel_sources"); const snapshot = @import("snapshot.zig"); const roles_gen = @import("roles_gen.zig"); @@ -21,24 +22,24 @@ pub fn path(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { pub fn writeFile(st: *const State, file: [:0]const u8) !void { var tmp_buf: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head: Head = .{ .magic = FILE_MAGIC, .layout = @backingInt(st.layout), .at = st.at, .bytes = st.bytes, .hist = st.hist }; try put(fd, std.mem.asBytes(&head)); try put(fd, st.buf.contents()[0..st.bytes]); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } /// `file`'s state after `at` tokens, in a buffer from the engine's pool; another position, layout or size is refused. pub fn readFile(e: *Engine, gpa: std.mem.Allocator, at: usize, file: [:0]const u8) !*State { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: Head = undefined; try get(fd, std.mem.asBytes(&head), 0); const layout: u32 = @backingInt(snapshot.layoutOf(e)); @@ -53,7 +54,7 @@ pub fn readFile(e: *Engine, gpa: std.mem.Allocator, at: usize, file: [:0]const u fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -62,7 +63,7 @@ fn put(fd: c_int, b: []const u8) !void { fn get(fd: c_int, b: []u8, at: u64) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.SnapshotRead; done += @intCast(n); } diff --git a/zig/src/families/flashnext/replay.zig b/zig/src/families/flashnext/replay.zig index 834c0a732a..7021f8c6b1 100644 --- a/zig/src/families/flashnext/replay.zig +++ b/zig/src/families/flashnext/replay.zig @@ -1,6 +1,7 @@ //! Flash Next on our Metal runtime: the Python engine's recorded kernels (a dump from tools/zig/flashnext_dump.py) //! replayed from Zig, with our own kernels for wide experts, prompt chunks, block selection and GPU-side rounds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const roles_gen = @import("roles_gen.zig"); @@ -43,7 +44,7 @@ pub const Buf = struct { return .{ .b = x.b, .off = x.off + bytes }; } }; -pub const Entry = struct { fd: std.c.fd_t, at: usize, len: usize }; +pub const Entry = struct { fd: platform.c.fd_t, at: usize, len: usize }; pub const Variant = struct { inputs: [][]const u8, outputs: [][]const u8, meta: [][]const u8, pipe: mtl.Pipeline, file: []const u8 = "", name: []const u8 = "", text: []const u8 = "" }; pub const Site = struct { v: *Variant, grid: mtl.Size, tg: mtl.Size }; @@ -814,10 +815,10 @@ const hc_up_source = \\ ; -pub fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { +pub fn readAll(fd: platform.c.fd_t, dest: []u8, at: usize) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } @@ -915,7 +916,7 @@ pub const Run = struct { /// Every tensor of a safetensors file in the index, at its absolute offset. pub fn indexFile(r: *Run, path: [:0]const u8) !void { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; var head: [8]u8 = undefined; try readAll(fd, &head, 0); diff --git a/zig/src/families/flashnext/rows_load.zig b/zig/src/families/flashnext/rows_load.zig index cf8e89cfbb..1932830cd1 100644 --- a/zig/src/families/flashnext/rows_load.zig +++ b/zig/src/families/flashnext/rows_load.zig @@ -1,5 +1,6 @@ //! Rows mode without tensor units (0.6.6's pre-M5 decode): every dense lane also as MLX-layout members for rows_w. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const replay = @import("replay.zig"); pub const rows_w = @import("rows_w.zig"); @@ -14,7 +15,7 @@ const KO = 6144; // the attention and DeltaNet output projections' inputs /// Whether decode reads its dense projections on the row kernels: no tensor units and FZ_M3_LANES unset. pub fn on(device: mtl.Device) bool { - return !device.tensorUnits() and std.c.getenv("FZ_M3_LANES") == null; + return !device.tensorUnits() and platform.c.getenv("FZ_M3_LANES") == null; } /// `d` on the run's encoder at `rows` rows, unless the role is knocked out. diff --git a/zig/src/families/flashnext/tp.zig b/zig/src/families/flashnext/tp.zig index e5718126d2..3fe8c3a514 100644 --- a/zig/src/families/flashnext/tp.zig +++ b/zig/src/families/flashnext/tp.zig @@ -1,5 +1,6 @@ //! Speed-up mode for Flash Next: two Macs with the whole model each, half the work each, over MCDMA. Decode: each rank computes every routed and the shared expert for its half of the window's rows (`ownRows`) and those rows' MoE branch with one Mac's sums; after each target layer's experts the GPU posts a sequence; the host sends the branch rows (bf16) to the peer, and the next layer's combine waits for the flag the peer's message sets, then updates every row's streams (one Mac's bits). DeltaNet layers split their heads: each rank's out-projection partial (fp32), exchanged and added in rank order. Prefill: a chunk's rows split across the two Macs (replay.zig `chunkPair`); the host sends each layer's handoff into the peer's slot for that layer and signals the layer's flag, which the peer's GPU waits for. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const host_wait = @import("tp_host_wait.zig"); @@ -341,9 +342,9 @@ pub const Tp2 = struct { .jobs = try gpa.alloc(Job, JOBS), }; t.one.slice(i32, 1)[0] = 1; - t.trace = std.c.getenv("TF_TP_TRACE") != null; - t.local = std.c.getenv("TF_TP_LOCAL") != null; - t.stats = std.c.getenv("TF_TP_STATS") != null; + t.trace = platform.c.getenv("TF_TP_TRACE") != null; + t.local = platform.c.getenv("TF_TP_LOCAL") != null; + t.stats = platform.c.getenv("TF_TP_STATS") != null; // both ranks up before the first round: a word each way try rd.signal(t.peer, HELLO, 1); while (@atomicLoad(u64, t.word64(HELLO), .acquire) < 1) std.atomic.spinLoopHint(); @@ -529,11 +530,11 @@ pub const Tp2 = struct { /// Served, rank 0: rank 1's answer for the request just sent: true when it restored the same resume state. pub fn waitAck(t: *Tp2) !bool { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (true) { const v = @atomicLoad(u64, t.word64(ACK_FLAG), .acquire); if (v >> 1 >= t.req) return v >> 1 == t.req and v & 1 != 0; - if (std.c.mach_absolute_time() - t0 > 7_200_000_000) return error.TpPeerSilent; // 300 s: rank 1 may still be loading + if (platform.c.mach_absolute_time() - t0 > 7_200_000_000) return error.TpPeerSilent; // 300 s: rank 1 may still be loading std.atomic.spinLoopHint(); } } @@ -551,8 +552,8 @@ pub const Tp2 = struct { if (t.quitting.load(.acquire)) return null; spins += 1; if (spins > 100_000) { // idle: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } return .{ .head = @ptrCast(t.win.ptr + REQ), .tokens = @ptrCast(@alignCast(t.win.ptr + REQ + 64)) }; @@ -565,11 +566,11 @@ pub const Tp2 = struct { try t.rd.signal(t.peer, CTRL_FLAG, (t.ctrl << 6) | (@as(u64, depth & 31) << 1) | @intFromBool(quit)); return .{ .quit = quit, .depth = depth }; } - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (true) { const v = @atomicLoad(u64, t.word64(CTRL_FLAG), .acquire); if (v >> 6 >= t.ctrl) return .{ .quit = quit or (v >> 6 == t.ctrl and v & 1 != 0), .depth = @intCast((v >> 1) & 31) }; - if (std.c.mach_absolute_time() - t0 > 240_000_000) return error.TpPeerSilent; // 10 s + if (platform.c.mach_absolute_time() - t0 > 240_000_000) return error.TpPeerSilent; // 10 s std.atomic.spinLoopHint(); } } @@ -604,7 +605,7 @@ pub const Tp2 = struct { /// for the peer's, by the flag its message sets) or a prefill send (its writes, then its flag). While idle it /// watches the last exchange: 10 s without the peer's answer fails the link. fn service(t: *Tp2, first: u32) void { - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send const posted = t.word32(SYNC + GPU * 4); const flag = t.word64(FLAG); var seq = first; @@ -614,7 +615,7 @@ pub const Tp2 = struct { var w = @atomicLoad(u32, posted, .acquire); while (!postReached(w, seq) or !reached(t.queued.load(.acquire), seq)) { if (t.stop.load(.acquire)) return; - if (want_at > 0 and !reached(@truncate(@atomicLoad(u64, flag, .acquire)), @truncate(want)) and !t.failed.load(.acquire) and std.c.mach_absolute_time() - want_at > 240_000_000) { // 10 s of 24 MHz ticks + if (want_at > 0 and !reached(@truncate(@atomicLoad(u64, flag, .acquire)), @truncate(want)) and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { // 10 s of 24 MHz ticks t.fail(seq, error.PeerSilent); t.land(want); } @@ -622,7 +623,7 @@ pub const Tp2 = struct { w = @atomicLoad(u32, posted, .acquire); } const job = t.jobs[seq % JOBS]; - const seen = if (t.stats) std.c.mach_absolute_time() else 0; + const seen = if (t.stats) platform.c.mach_absolute_time() else 0; if ((t.local and job.kind != .send) or t.failed.load(.acquire)) { // a failed link drains: the GPU never hangs if (job.kind != .send) t.land(job.value); seq +%= 1; @@ -638,11 +639,11 @@ pub const Tp2 = struct { } if (job.kind != .send) { want = x; - want_at = std.c.mach_absolute_time(); + want_at = platform.c.mach_absolute_time(); if (t.failed.load(.acquire)) t.land(x); } if (t.stats) { - const held = std.c.mach_absolute_time() - seen; // ticks: 41.67 ns each on Apple silicon + const held = platform.c.mach_absolute_time() - seen; // ticks: 41.67 ns each on Apple silicon t.held_ns += held; t.held_max = @max(t.held_max, held); t.held_kind[@backingInt(job.kind)][0] += held; @@ -705,13 +706,13 @@ test "the host service follows the GPU's posts across the 27-bit post wrap and t for (0..8) |_| { // local serving lands each exchange's flag once its post is seen x +%= 1; const seq = t.queue(.{ .kind = .exchange, .value = x }); - const ts: std.c.timespec = .{ .sec = 0, .nsec = 2_000_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 2_000_000 }; + _ = platform.c.nanosleep(&ts, null); try std.testing.expect(@as(u32, @truncate(@atomicLoad(u64, t.word64(FLAG), .acquire))) != x); // not before its post @atomicStore(u32, t.word32(SYNC + GPU * 4), (seq << 5) | 4, .release); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (@as(u32, @truncate(@atomicLoad(u64, t.word64(FLAG), .acquire))) != x) { - if (std.c.mach_absolute_time() - t0 > 24_000_000) return error.ServiceStuck; // 1 s + if (platform.c.mach_absolute_time() - t0 > 24_000_000) return error.ServiceStuck; // 1 s std.atomic.spinLoopHint(); } } @@ -729,8 +730,8 @@ test "partNext and hostWait across the 32-bit wrap: the next exchange's slot, an } const Peer = struct { // the peer's message landing a call's flag, a little later fn land(w: *u64, v: u64) void { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000_000 }; + _ = platform.c.nanosleep(&ts, null); @atomicStore(u64, w, v, .release); } }; diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 98404ca22a..f1a92ef8c3 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's dev-time checks on a loaded engine: captures, traces, teacher forcing and knock-out profiles. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -27,9 +28,9 @@ pub fn capture(e: *Engine, prompt: []const u32, path: []const u8) !void { fwd.head(&x, b.enc, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); try e.finish(b.cb, b.enc); fwd.flipKda(&x); - const file = std.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; std.debug.print("captured {d} rows ({d} bytes) in {s}\n", .{ n, bytes, path }); } @@ -47,8 +48,8 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, @memcpy(Engine.u32s(e.prompt_ids.at(@as(usize, P) * 4), total - P), reply[0 .. total - P]); const most = @as(usize, st.max_rows) * c.hidden * 2 * (2 + 4 * @as(usize, c.run)) + @as(usize, st.max_rows) * c.vocab * 2; const dump = try e.arena.buffer(most); - const file = std.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); var x = e.ctx(); var at: u32 = 0; var calls: usize = 0; @@ -62,7 +63,7 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, try e.finish(b.cb, b.enc); fwd.flipKda(&x); const bytes = x.dump_at + @as(usize, n) * c.vocab * 2; - if (std.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; + if (platform.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; at += n; } std.debug.print("traced {d} calls ({d} prompt rows, {d} steps) in {s}\n", .{ calls, P, total - P, path }); @@ -106,7 +107,7 @@ pub fn profile(e: *Engine, depth: u32, reps: usize, only: bool, parts: bool) !vo const c = &e.c; const D = c.hidden; const d: u32 = if (e.w.mtp == null) 0 else @min(depth, st.max_rows - 1); - const wide = if (std.c.getenv("GLM_PROFILE_ROWS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 1 else 1; // a copy window's rows, without the MTP head + const wide = if (platform.c.getenv("GLM_PROFILE_ROWS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 1 else 1; // a copy window's rows, without the MTP head const R = if (d == 0) std.math.clamp(wide, 1, st.max_rows) else d + 1; const pool = mtl.objc.Pool.push(); defer pool.pop(); diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index e1ffe5ac41..e030c7200a 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash on Metal: one greedy reply at a time, the prompt in 16-row windows or tensor-unit chunks, then drafted rounds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -95,7 +96,7 @@ pub const Engine = struct { /// The checkpoint in `dir`, caches for `cap` tokens; GLM_LAYERS=N: the first N layers only; GLM_EP=settings: half the experts. pub fn load(gpa: std.mem.Allocator, dir: []const u8, cap: u32) !*Engine { - return loadWith(gpa, dir, cap, if (std.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); + return loadWith(gpa, dir, cap, if (platform.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); } /// `load` with expert parallel over the link in `ep_path` (this Mac's settings), or on one Mac when null. @@ -104,23 +105,23 @@ pub const Engine = struct { errdefer gpa.destroy(e); const pool = mtl.objc.Pool.push(); defer pool.pop(); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); e.gpa = gpa; e.ev = 0; e.residency = null; e.gpu = .{ 0, 0 }; - e.fused_route = if (std.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; - e.draft_vocab = if (std.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; + e.fused_route = if (platform.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; + e.draft_vocab = if (platform.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; e.committed = 0; e.ep = null; e.pr = null; e.trace_last = null; e.margins = null; e.chunk_rows = prompt_mod.max_rows; - e.cut = if (std.c.getenv("GLM_CUTS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; - e.copy_min = if (std.c.getenv("GLM_COPY")) |v| std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0, 0, 8) else 0; - e.rank_log = if (std.c.getenv("GLM_RANKS")) |v| v[0] == '1' else false; - if (std.c.getenv("GLM_CHUNK")) |v| e.chunk_rows = std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch prompt_mod.max_rows, st.max_rows + 1, prompt_mod.max_rows); + e.cut = if (platform.c.getenv("GLM_CUTS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; + e.copy_min = if (platform.c.getenv("GLM_COPY")) |v| std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0, 0, 8) else 0; + e.rank_log = if (platform.c.getenv("GLM_RANKS")) |v| v[0] == '1' else false; + if (platform.c.getenv("GLM_CHUNK")) |v| e.chunk_rows = std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch prompt_mod.max_rows, st.max_rows + 1, prompt_mod.max_rows); e.ep_arena = .init(gpa); errdefer e.ep_arena.deinit(); e.device = try mtl.Device.init(); @@ -138,17 +139,17 @@ pub const Engine = struct { e.draft_vocab -= e.draft_vocab % 4; // the head's kernel takes four rows a simdgroup const model = try load_plan.modelHash(gpa, dir, f.bytes[0..f.size]); e.model_hash = model; - if (std.c.getenv("GLM_LAYERS")) |v| try cfg.subset(&e.c, std.fmt.parseInt(u32, std.mem.span(v), 10) catch return error.BadLayerCount); - if (std.c.getenv("GLM_MTP")) |v| if (v[0] == '0') { // no MTP head: no drafts, its ~4 GB left on disk + if (platform.c.getenv("GLM_LAYERS")) |v| try cfg.subset(&e.c, std.fmt.parseInt(u32, std.mem.span(v), 10) catch return error.BadLayerCount); + if (platform.c.getenv("GLM_MTP")) |v| if (v[0] == '0') { // no MTP head: no drafts, its ~4 GB left on disk e.c.mtp = 0; }; const link: ?ep_mod.Settings = if (ep_path) |sp| blk: { const sf = try mtl.MappedFile.open(try e.ep_arena.allocator().dupeSentinel(u8, sp, 0)); defer sf.deinit(); const s = try ep_mod.readSettings(e.ep_arena.allocator(), sf.bytes[0..sf.size]); - const by_rows = if (std.c.getenv("GLM_EP_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "experts") else true; + const by_rows = if (platform.c.getenv("GLM_EP_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "experts") else true; if (by_rows) try cfg.splitRows(&e.c, s.rank, 2) else try cfg.split(&e.c, s.rank, 2); - const tp = if (std.c.getenv("GLM_TP")) |v| v[0] != '0' else true; // TP2: each Mac half the KDA heads (GLM_TP=0: off) + const tp = if (platform.c.getenv("GLM_TP")) |v| v[0] != '0' else true; // TP2: each Mac half the KDA heads (GLM_TP=0: off) if (by_rows and tp) try cfg.splitHeads(&e.c, s.rank, 2); if (e.c.tp > 1) { // TP2's head halves: the whole vocabulary (its kernel's pitch); no host logits to rank e.draft_vocab = e.c.vocab; @@ -162,7 +163,7 @@ pub const Engine = struct { gpa.destroy(e.k); } try frags.check(e.device, e.queue, gpa); // the fragment layout nax.h gives GLM's prompt and attention kernels - const chunked = if (std.c.getenv("GLM_PROMPT")) |v| v[0] != '0' else true; + const chunked = if (platform.c.getenv("GLM_PROMPT")) |v| v[0] != '0' else true; e.limit = try load_plan.loadLimit(e.device); const limit = e.limit.bytes; var plan_bytes = try load_plan.planBytes(gpa, e.device, dir, &e.c); @@ -176,7 +177,7 @@ pub const Engine = struct { plan_bytes = lean_bytes; } } - if (std.c.getenv("GLM_DRY") != null) return error.DryRun; + if (platform.c.getenv("GLM_DRY") != null) return error.DryRun; e.arena = .{ .device = e.device, .gpa = gpa }; errdefer e.arena.deinit(); const both = try st.init(&e.arena, &e.c, cap); @@ -195,9 +196,9 @@ pub const Engine = struct { return error.OverMemoryLimit; } { // free memory short of the load: a cached copy of the checkpoint holds it, the paced reads would wait (#583) - const dropped_at = std.c.mach_absolute_time(); + const dropped_at = platform.c.mach_absolute_time(); if (page_cache.dropForLoad(gpa, dir, plan_bytes + e.arena.bytes)) |d| { - if (d.dropped > 0) std.log.info("glm: dropped {d:.1} GiB of the checkpoint from the file cache, the {d:.1} GiB free memory lacked for the load ({d:.2} s)", .{ @as(f64, @floatFromInt(d.dropped)) / (1 << 30), @as(f64, @floatFromInt(d.deficit)) / (1 << 30), @as(f64, @floatFromInt(std.c.mach_absolute_time() - dropped_at)) / 24e6 }); + if (d.dropped > 0) std.log.info("glm: dropped {d:.1} GiB of the checkpoint from the file cache, the {d:.1} GiB free memory lacked for the load ({d:.2} s)", .{ @as(f64, @floatFromInt(d.dropped)) / (1 << 30), @as(f64, @floatFromInt(d.deficit)) / (1 << 30), @as(f64, @floatFromInt(platform.c.mach_absolute_time() - dropped_at)) / 24e6 }); } else |err| std.log.warn("glm: could not drop the checkpoint's cached pages ({s}); the load reads beside them", .{@errorName(err)}); } e.w = try wts.load(gpa, e.device, dir, &e.c, 16, false); @@ -216,7 +217,7 @@ pub const Engine = struct { e.ep = try ep_mod.Ep.init(gpa, e.device, s, me); } // opt-in: wiring 181 GB leaves macOS nothing to reclaim if another model shares the Mac (Flash Next runs without) - if (std.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { + if (platform.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { for (e.w.buffers.items) |b| set.add(b); for (e.arena.buffers.items) |b| set.add(b); set.commit(); @@ -228,13 +229,13 @@ pub const Engine = struct { if (e.residency) |set| e.keepalive_sets[0] = set; e.keepalive_target = .{ .queue = e.queue, .sets = if (e.residency != null) &e.keepalive_sets else &.{} }; // a process's first prompt pays a one-time cost: every rank of a pair pays it here at the same step (GLM_NO_WARMUP: skip) - if (e.pr != null and cap >= 128 and std.c.getenv("GLM_NO_WARMUP") == null) { + if (e.pr != null and cap >= 128 and platform.c.getenv("GLM_NO_WARMUP") == null) { var ids: [64]u32 = undefined; for (&ids, 0..) |*t, i| t.* = @intCast(1000 + i); var quiet: Quiet = .{}; _ = try e.generate(&ids, 2, &.{}, 0, .{ .ctx = &quiet, .prefilled = Quiet.prefilled, .tokens = Quiet.tokens, .cancelled = Quiet.cancelled }); } - e.load_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) / 24e6; + e.load_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) / 24e6; return e; } @@ -324,7 +325,7 @@ pub const Engine = struct { e.ev += 1; cb.signal(e.event, e.ev); cb.commit(); - e.committed = std.c.mach_absolute_time(); + e.committed = platform.c.mach_absolute_time(); if (spin) { while (@backingInt(cb.status()) < @backingInt(mtl.types.CommandBufferStatus.completed)) std.atomic.spinLoopHint(); } else cb.wait(); @@ -408,7 +409,7 @@ pub const Engine = struct { e.s.reset(); @memcpy(u32s(e.prompt_ids, prompt.len), prompt); var x = e.ctx(); - const t_start = std.c.mach_absolute_time(); + const t_start = platform.c.mach_absolute_time(); // the prompt: windows of up to 16 rows, the MTP head taking each row with the token after it var at: u32 = 0; var last_n: u32 = 1; @@ -456,7 +457,7 @@ pub const Engine = struct { at += n; } e.s.pos = P; - const t_prompt = std.c.mach_absolute_time(); + const t_prompt = platform.c.mach_absolute_time(); out.prefilled(out.ctx); var res: Result = .{ .reason = .length, .prompt_seconds = @as(f64, @floatFromInt(t_prompt - t_start)) / 24e6 }; var tok = u32s(e.sc.picks, 1)[0]; @@ -469,7 +470,7 @@ pub const Engine = struct { if (max_tokens <= 1) return res; if (d == 0 and e.ep == null and e.copy_min == 0 and e.margins == null) { try plain_rounds.run(e, &x, tok, max_tokens, eos, out, &res); - res.decode_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t_prompt)) / 24e6; + res.decode_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t_prompt)) / 24e6; return res; } var emitted: usize = 1; @@ -490,7 +491,7 @@ pub const Engine = struct { u32s(e.sc.next, 1)[0] = tok; var last_end: f64 = 0; while (true) { - const t_enc = std.c.mach_absolute_time(); + const t_enc = platform.c.mach_absolute_time(); const b = e.begin(); if (window > 0) { // the last window's kept rows (the prompt's windows kept all theirs) fwd.keepKda(&x, b.enc, window, keep); @@ -584,7 +585,7 @@ pub const Engine = struct { break; } } - res.decode_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t_prompt)) / 24e6; + res.decode_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t_prompt)) / 24e6; return res; } }; diff --git a/zig/src/families/glm/ep.zig b/zig/src/families/glm/ep.zig index 74ea506c26..1fd8a94b9f 100644 --- a/zig/src/families/glm/ep.zig +++ b/zig/src/families/glm/ep.zig @@ -1,5 +1,6 @@ //! Expert parallel over two Macs (MCDMA): half the routed experts each, a MoE call's (row, slot) outputs swapped, all combined on both. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const settings = @import("../flashnext/tp_settings.zig"); @@ -120,8 +121,8 @@ pub const Ep = struct { t.ctl = try control.Control.init(gpa, rd, CONTROL, &t.failed); errdefer t.ctl.deinit(gpa); _ = try t.ctl.hello(me); // both Macs up, running the same thing, before the first exchange - t.trace = std.c.getenv("GLM_EP_TRACE") != null; - if (std.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); + t.trace = platform.c.getenv("GLM_EP_TRACE") != null; + if (platform.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); t.thread = try std.Thread.spawn(.{}, service, .{t}); std.log.info("expert parallel: rank {d} of 2 holds experts {d}-{d}, connected", .{ t.rank, me.own_lo, me.own_hi - 1 }); return t; @@ -307,7 +308,7 @@ pub const Ep = struct { /// Each posted exchange's entries to the peer, then its flag; 10 s without an answer fails the link and lands every flag. fn service(t: *Ep) void { - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send const posted = t.word32(POSTED); const flag = t.word64(FLAG); var x: u64 = 1; @@ -320,26 +321,26 @@ pub const Ep = struct { t.st.watch(@atomicLoad(u64, flag, .acquire), x - 1); if (t.stop.load(.acquire)) return; if (waiting and @atomicLoad(u64, flag, .acquire) >= want) waiting = false; - if (waiting and !t.failed.load(.acquire) and std.c.mach_absolute_time() - want_at > 240_000_000) { + if (waiting and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { t.fail(want, error.PeerSilent); t.land(want); waiting = false; } spins += 1; if (spins > 1_000_000) { // idle: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } - const seen = std.c.mach_absolute_time(); + const seen = platform.c.mach_absolute_time(); const n: usize = @atomicLoad(u32, t.word32(COUNT + 4 * @as(usize, @intCast(x % 2))), .acquire); const theirs: u32 = @bitCast(@as([*]const i32, @ptrCast(@alignCast(t.lists.contents() + COUNTS)))[1]); t.st.posted(x, seen, @intCast(n), theirs); t.st.watch(@atomicLoad(u64, flag, .acquire), x); if (t.trace) std.debug.print("EP rank{d} exchange {d}: {d} entries\n", .{ t.rank, x, n }); if (t.delay_ns > 0 and n * ENTRY > 1 << 20) { - const ts: std.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } if (t.failed.load(.acquire) or n * ENTRY > SLOT) { if (n * ENTRY > SLOT) t.fail(x, error.BadCount); @@ -349,7 +350,7 @@ pub const Ep = struct { t.land(x); }; want = x; - want_at = std.c.mach_absolute_time(); + want_at = platform.c.mach_absolute_time(); waiting = true; t.sent += 1; t.held_ticks += want_at - seen; @@ -394,7 +395,7 @@ pub const Ep = struct { /// The flag's new value seen now; every exchange up to `upto` whose post and landing are both known is counted. fn watch(s: *Stats, flag: u64, upto: u64) void { if (flag > s.flag_seen) { - const now = std.c.mach_absolute_time(); + const now = platform.c.mach_absolute_time(); var v = s.flag_seen + 1; while (v <= flag) : (v += 1) s.flag_t[v % 64] = now; s.flag_seen = flag; diff --git a/zig/src/families/glm/ep_control.zig b/zig/src/families/glm/ep_control.zig index 420d7411d4..4e69e3e39b 100644 --- a/zig/src/families/glm/ep_control.zig +++ b/zig/src/families/glm/ep_control.zig @@ -1,5 +1,6 @@ //! The expert-parallel pair's control protocol over a link's window: identities, requests and stop decisions, each acked. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const Rdma = fabric.rdma.Rdma; @@ -54,7 +55,7 @@ pub const Request = struct { max_tokens: usize, depth: usize, eos: []const u32, pub const Command = union(enum) { request: Request, lanes: struct { kind: u32, words: []const u32 } }; fn ticks() u64 { - return std.c.mach_absolute_time(); // 24 MHz on Apple silicon + return platform.c.mach_absolute_time(); // 24 MHz on Apple silicon } pub const Control = struct { @@ -96,8 +97,8 @@ pub const Control = struct { if (bounded and ticks() - t0 > c.limit) return error.EpPeerSilent; spins += 1; if (spins > 100_000) { // a long wait: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } } @@ -295,8 +296,8 @@ test "early cancels: rank 0 waits for rank 1's acks, nothing is overwritten or s }; var e0: ?anyerror = undefined; const th = try std.Thread.spawn(.{}, Rank0.run, .{ &p.c[0], &e0 }); - const ts: std.c.timespec = .{ .sec = 0, .nsec = 200_000_000 }; // rank 1 held back while rank 0 runs ahead - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 200_000_000 }; // rank 1 held back while rank 0 runs ahead + _ = platform.c.nanosleep(&ts, null); for (prompts, 0..) |pr, i| { const r = (try p.c[1].waitRequest()).?; try std.testing.expectEqualSlices(u32, &pr, r.prompt); diff --git a/zig/src/families/glm/learned_pair.zig b/zig/src/families/glm/learned_pair.zig index 4d15eb992f..b1c8c57321 100644 --- a/zig/src/families/glm/learned_pair.zig +++ b/zig/src/families/glm/learned_pair.zig @@ -1,5 +1,6 @@ //! A peer's disk reservation covers its snapshot half until rank 0 has committed or removed the shared index entry. const std = @import("std"); +const platform = @import("sys"); const disk = @import("lanes").learned_disk; pub const State = struct { admission: disk.Admission = .{}, @@ -98,12 +99,12 @@ test "each rank includes actual other-identity files in its independent cap befo var filename: [1400]u8 = undefined; const current = try std.fmt.bufPrintSentinel(&one, "{s}/0000000000000001", .{root}, 0); const other = try std.fmt.bufPrintSentinel(&two, "{s}/0000000000000002", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(current, 0o700)); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(other, 0o700)); - const fd = std.c.open(try std.fmt.bufPrintSentinel(&filename, "{s}/state.bin", .{other}, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(current, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(other, 0o700)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&filename, "{s}/state.bin", .{other}, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestCreate; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 96)); - _ = std.c.close(fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 96)); + _ = platform.c.close(fd); var rank0: State = .{ .admission = .{ .floor = 0, .free = Fake.free, .clock = Fake.clock }, .cap = 200 }; var rank1: State = .{ .admission = .{ .floor = 0, .free = Fake.free, .clock = Fake.clock }, .cap = 100 }; try std.testing.expectEqual(@as(u64, 0), rank0.need(current, 8, 2)); diff --git a/zig/src/families/glm/load_plan.zig b/zig/src/families/glm/load_plan.zig index 53a0dddbcd..94d2b98b68 100644 --- a/zig/src/families/glm/load_plan.zig +++ b/zig/src/families/glm/load_plan.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's load plan: the checkpoint's identity, the bytes its weights and caches take, the chunk height, the limit. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -137,8 +138,8 @@ pub fn limitFor(ram: u64, working_set: u64, text: ?[]const u8) error{BadLoadLimi pub fn loadLimit(device: mtl.Device) error{BadLoadLimit}!Limit { var mem: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; // RAM unknown: a zero limit refuses every load - const text: ?[]const u8 = if (std.c.getenv("GLM_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; + if (platform.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; // RAM unknown: a zero limit refuses every load + const text: ?[]const u8 = if (platform.c.getenv("GLM_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; const l = limitFor(mem, device.maxWorkingSet(), text) catch |err| { std.log.err("glm: GLM_LOAD_LIMIT_GB={s} is not a load limit: give a number of GB above 0, or unset it for 70% of RAM", .{text.?}); return err; diff --git a/zig/src/families/glm/mirror.zig b/zig/src/families/glm/mirror.zig index 6a3a7dbb70..4c9d862f9f 100644 --- a/zig/src/families/glm/mirror.zig +++ b/zig/src/families/glm/mirror.zig @@ -1,5 +1,6 @@ //! Speed-up mode's slot commands: rank 0 sends each slot operation before it runs it; rank 1 runs the same, in order. const std = @import("std"); +const platform = @import("sys"); const eng = @import("engine.zig"); const slots_mod = @import("slots.zig"); const st = @import("state.zig"); @@ -192,10 +193,10 @@ fn learned(sl: *Slots, k: Kind, w: []const u32) !void { } pub fn fileBytes(file: [:0]const u8) u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return 0; - defer _ = std.c.close(fd); - const n = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const n = platform.c.lseek(fd, 0, platform.c.SEEK.END); return if (n > 0) @intCast(n) else 0; } diff --git a/zig/src/families/glm/names.zig b/zig/src/families/glm/names.zig index abea94eabb..3058b447e8 100644 --- a/zig/src/families/glm/names.zig +++ b/zig/src/families/glm/names.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's checkpoint index in either layout (the original or mlx-lm's), each tensor by its original name. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const load_fix = @import("weights_fix.zig"); @@ -45,15 +46,15 @@ pub fn read(gpa: std.mem.Allocator, arena: std.mem.Allocator, dir: []const u8, m for (files.keys(), 0..) |name, si| { const full = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); try shards.append(gpa, full); - const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(full, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: [8]u8 = undefined; try load_fix.readAll(fd, &head, 0); const hlen = std.mem.readInt(u64, &head, .little); const header = try arena.alloc(u8, hlen); try load_fix.readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); const data = 8 + hlen; const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, arena, header, .{}); var eit = doc2.object.iterator(); diff --git a/zig/src/families/glm/page_cache.zig b/zig/src/families/glm/page_cache.zig index ab64f39d75..cdd9fb2ab9 100644 --- a/zig/src/families/glm/page_cache.zig +++ b/zig/src/families/glm/page_cache.zig @@ -1,5 +1,6 @@ //! A checkpoint's pages in the file cache, dropped before a load: F_NOCACHE reads keep out of it but evict nothing. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const paced_read = @import("../../core/paced_read.zig"); @@ -33,15 +34,15 @@ pub fn dropCached(gpa: std.mem.Allocator, dir: []const u8, want: u64, d: *Droppe const page: usize = std.heap.pageSize(); for (files.keys()) |name| { if (d.dropped >= want) break; - const fd = std.c.open(try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end <= 0) continue; const len: usize = @intCast(end); - const map = std.c.mmap(null, len, .{ .READ = true }, .{ .TYPE = .SHARED }, fd, 0); - if (map == std.c.MAP_FAILED) return error.MapFailed; - defer _ = std.c.munmap(@alignCast(map), len); + const map = platform.c.mmap(null, len, .{ .READ = true }, .{ .TYPE = .SHARED }, fd, 0); + if (map == platform.c.MAP_FAILED) return error.MapFailed; + defer _ = platform.c.munmap(@alignCast(map), len); const vec = try a.alloc(u8, (len + page - 1) / page); var cached: u64 = 0; if (mincore(map, len, vec.ptr) == 0) for (vec) |v| { @@ -57,31 +58,31 @@ pub fn dropCached(gpa: std.mem.Allocator, dir: []const u8, want: u64, d: *Droppe test "a checkpoint's cached pages are dropped only as far as wanted, and a shard with none is left alone" { const gpa = std.testing.allocator; var name_buf: [64]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-drop-cached-{d}", .{std.c.getpid()}, 0); - _ = std.c.mkdir(dir, 0o700); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-drop-cached-{d}", .{platform.c.getpid()}, 0); + _ = platform.c.mkdir(dir, 0o700); + defer _ = platform.c.rmdir(dir); var path_buf: [128]u8 = undefined; const index = try std.fmt.bufPrintSentinel(&path_buf, "{s}/model.safetensors.index.json", .{dir}, 0); const text = "{\"weight_map\": {\"a\": \"s.safetensors\", \"b\": \"s.safetensors\"}}"; var shard_buf: [128]u8 = undefined; const shard = try std.fmt.bufPrintSentinel(&shard_buf, "{s}/s.safetensors", .{dir}, 0); for ([_][:0]const u8{ index, shard }, [_][]const u8{ text, "" }) |file, body| { - const fd = std.c.open(file, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); if (body.len > 0) { - _ = std.c.write(fd, body.ptr, body.len); + _ = platform.c.write(fd, body.ptr, body.len); continue; } const data = try gpa.alloc(u8, 4 << 20); defer gpa.free(data); @memset(data, 7); - try std.testing.expectEqual(@as(isize, @intCast(data.len)), std.c.write(fd, data.ptr, data.len)); - _ = std.c.fsync(fd); - _ = std.c.pread(fd, data.ptr, data.len, 0); // read through the file cache + try std.testing.expectEqual(@as(isize, @intCast(data.len)), platform.c.write(fd, data.ptr, data.len)); + _ = platform.c.fsync(fd); + _ = platform.c.pread(fd, data.ptr, data.len, 0); // read through the file cache } - defer _ = std.c.unlink(index); - defer _ = std.c.unlink(shard); + defer _ = platform.c.unlink(index); + defer _ = platform.c.unlink(shard); var none: Dropped = .{ .deficit = 0 }; try dropCached(gpa, dir, 0, &none); // nothing wanted: no shard touched, the pages stay cached try std.testing.expectEqual(@as(u64, 0), none.dropped); diff --git a/zig/src/families/glm/plain_rounds.zig b/zig/src/families/glm/plain_rounds.zig index c5f2114916..3289fdc513 100644 --- a/zig/src/families/glm/plain_rounds.zig +++ b/zig/src/families/glm/plain_rounds.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's plain greedy rounds in generate: the next round committed before this one's pick is read. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fwd = @import("forward.zig"); const Fence = @import("../../core/fence.zig").Fence; @@ -12,7 +13,7 @@ const Pending = struct { cb: mtl.CommandBuffer, t_enc: u64, committed: u64 }; /// One plain round at `pos`, its token from `from` (null: the ids as set), its pick into `to`; after a round it waits on the fence, not the event. fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { - const t_enc = std.c.mach_absolute_time(); + const t_enc = platform.c.mach_absolute_time(); const fence = e.round_fence.?; const cb = e.queue.commandBuffer(); if (from == null and e.ev > 0) cb.waitFor(e.event, e.ev); @@ -31,7 +32,7 @@ fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { e.ev += 1; b.cb.signal(e.event, e.ev); b.cb.commit(); - e.committed = std.c.mach_absolute_time(); + e.committed = platform.c.mach_absolute_time(); return .{ .cb = b.cb, .t_enc = t_enc, .committed = e.committed }; } diff --git a/zig/src/families/glm/pool.zig b/zig/src/families/glm/pool.zig index fb14809842..56f36077f0 100644 --- a/zig/src/families/glm/pool.zig +++ b/zig/src/families/glm/pool.zig @@ -1,5 +1,6 @@ //! One Mac's streams share the engine's MLA cache, a contiguous region each; kept prefixes hold blocks until needed. const std = @import("std"); +const platform = @import("sys"); const st = @import("state.zig"); const snapshot = @import("snapshot.zig"); const Engine = @import("engine.zig").Engine; @@ -41,7 +42,7 @@ pub const Pool = struct { /// Whether the slots share one pool: on one Mac (a pair's ranks mirror slot commands), unless GLM_POOL=0. pub fn on(e: *const Engine) bool { if (e.ep != null or e.followsPeer()) return false; - const v = std.c.getenv("GLM_POOL") orelse return true; + const v = platform.c.getenv("GLM_POOL") orelse return true; return !std.mem.eql(u8, std.mem.span(v), "0"); } diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index 4f73d916b0..4be9a124cf 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -1,5 +1,6 @@ //! Prompt chunks: matmuls on the M5 tensor units at the row kernels' arithmetic (core/affine_mm.zig), the rest on the row kernels. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -83,7 +84,7 @@ pub fn init(gpa: std.mem.Allocator, arena: *st.Arena, device: mtl.Device, c: *co /// The prompt's sparse MLA on the tensor units (GLM_SPARSE_NAX=0: the row kernel, its output beside its queries). fn sparseNax() bool { - return if (std.c.getenv("GLM_SPARSE_NAX")) |v| v[0] != '0' else true; + return if (platform.c.getenv("GLM_SPARSE_NAX")) |v| v[0] != '0' else true; } /// The bytes `init` takes from the arena for `rows`-row chunks. diff --git a/zig/src/families/glm/slots.zig b/zig/src/families/glm/slots.zig index 90aa566962..fc97750612 100644 --- a/zig/src/families/glm/slots.zig +++ b/zig/src/families/glm/slots.zig @@ -1,5 +1,6 @@ //! Stream slots on one Mac: a stream's caches each; prompt chunks, shared windows, keeps and drafts run on them. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -76,7 +77,7 @@ pub const Slots = struct { .last = try e.arena.buffer(@as(usize, e.c.hidden) * 2), }; const plane = @as(usize, st.max_rows) * e.c.hidden * 2; - var sl: Slots = .{ .gpa = gpa, .e = e, .slots = slots, .log = std.c.getenv("GLM_WINDOWS") != null, .rows = try e.arena.buffer(plane), .ids = try e.arena.buffer(st.max_rows * 4), .lasts = try e.arena.buffer(plane), .picks = try e.arena.buffer(st.max_rows * 4), .ring = try e.arena.buffer(ring_len * 4) }; + var sl: Slots = .{ .gpa = gpa, .e = e, .slots = slots, .log = platform.c.getenv("GLM_WINDOWS") != null, .rows = try e.arena.buffer(plane), .ids = try e.arena.buffer(st.max_rows * 4), .lasts = try e.arena.buffer(plane), .picks = try e.arena.buffer(st.max_rows * 4), .ring = try e.arena.buffer(ring_len * 4) }; if (pool_mod.on(e)) pool_mod.setUp(&sl); return sl; } @@ -247,7 +248,7 @@ pub const Slots = struct { const e = sl.e; var segs: [st.max_rows]fwd.Seg = undefined; if (wins.len == 0 or wins.len > segs.len) return error.WindowOutOfStep; - const t_in = std.c.mach_absolute_time(); + const t_in = platform.c.mach_absolute_time(); sl.settleAll(); try sl.launch(); // the GPU runs the keeps and drafts while the window encodes const ids = Engine.u32s(e.sc.ids, st.max_rows); @@ -274,7 +275,7 @@ pub const Slots = struct { fwd.head(&x, enc, e.sc.hidden, e.sc.logits, e.sc.picks, total); sl.ahead = .{ 0, 0 }; try sl.flushBy(true); - const ret = std.c.mach_absolute_time(); + const ret = platform.c.mach_absolute_time(); if (sl.log) { const first = if (sl.ahead[1] > 0) sl.ahead[0] else e.gpu[0]; std.log.info("glm window: streams {d} rows {d} gpu {d:.2} ms gap {d:.2} ms ahead {d:.2} join {d:.3} host {d:.2} encode {d:.2} ms", .{ wins.len, total, (e.gpu[1] - e.gpu[0]) * 1e3, (first - sl.last_end) * 1e3, (sl.ahead[1] - sl.ahead[0]) * 1e3, if (sl.ahead[1] > 0) (e.gpu[0] - sl.ahead[1]) * 1e3 else 0, @as(f64, @floatFromInt(t_in -| sl.last_ret)) / 24e3, @as(f64, @floatFromInt(e.committed - t_in)) / 24e3 }); diff --git a/zig/src/families/glm/snapshot_file.zig b/zig/src/families/glm/snapshot_file.zig index 901b5d90b6..b8b058d7d3 100644 --- a/zig/src/families/glm/snapshot_file.zig +++ b/zig/src/families/glm/snapshot_file.zig @@ -1,45 +1,46 @@ //! GLM learned halves use a versioned checksum envelope; legacy, corrupt and partial files never enter a cache. const std = @import("std"); +const platform = @import("sys"); const checked = @import("../../core/snapshot_file.zig"); const magic: u32 = 0x474c4d54; pub const Head = extern struct { magic: u32 = magic, version: u32 = 1, at: u32, reserved: u32 = 0, bytes: u64 }; pub fn write(file: [:0]const u8, at: u32, payload: []const u8) !void { var tmp_buf: [1200]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head: Head = .{ .at = at, .bytes = payload.len }; try checked.write(fd, std.mem.asBytes(&head), payload); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } pub fn read(file: [:0]const u8, at: u32, payload: []u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const expected: Head = .{ .at = at, .bytes = payload.len }; var head: Head = undefined; - if (std.c.pread(fd, std.mem.asBytes(&head).ptr, @sizeOf(Head), 0) != @sizeOf(Head) or !std.mem.eql(u8, std.mem.asBytes(&head), std.mem.asBytes(&expected))) return error.SnapshotRead; + if (platform.c.pread(fd, std.mem.asBytes(&head).ptr, @sizeOf(Head), 0) != @sizeOf(Head) or !std.mem.eql(u8, std.mem.asBytes(&head), std.mem.asBytes(&expected))) return error.SnapshotRead; try checked.read(fd, std.mem.asBytes(&head), payload); } test "both GLM half files refuse a changed body, legacy header, truncation and another position" { var path: [160]u8 = undefined; - const file = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-glm-file-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(file); + const file = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-glm-file-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(file); try write(file, 128, "body"); var out: [4]u8 = undefined; try read(file, 128, &out); try std.testing.expectEqualStrings("body", &out); try std.testing.expectError(error.SnapshotRead, read(file, 129, &out)); - const fd = std.c.open(file, .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, @sizeOf(Head) + 1)); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, @sizeOf(Head) + 1)); try std.testing.expectError(error.SnapshotRead, read(file, 128, &out)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 16)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 16)); try std.testing.expectError(error.SnapshotRead, read(file, 128, &out)); } diff --git a/zig/src/families/glm/weights.zig b/zig/src/families/glm/weights.zig index 5beb758f14..d85bebec5b 100644 --- a/zig/src/families/glm/weights.zig +++ b/zig/src/families/glm/weights.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's weights in the Python family's layout (stacked projections and experts, repacked mixes), read by parallel preads. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const cfg = @import("config.zig"); diff --git a/zig/src/families/glm/weights_fix.zig b/zig/src/families/glm/weights_fix.zig index ec0861fe6e..3d12f8b75b 100644 --- a/zig/src/families/glm/weights_fix.zig +++ b/zig/src/families/glm/weights_fix.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash weights as they load: parallel reads into place, and rewrites (mixes, router, conv, norms, columns). const std = @import("std"); +const platform = @import("sys"); const st = @import("../../core/safetensors.zig"); const paced_read = @import("../../core/paced_read.zig"); @@ -12,16 +13,16 @@ pub const Transform = struct { src: Src, dst: [*]u8, fix: Fix, extra: [3]?Src = /// One tensor's bytes read straight into their place in a weights buffer. pub const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; -pub fn readAll(fd: std.c.fd_t, dest: []u8, at: u64) !void { +pub fn readAll(fd: platform.c.fd_t, dest: []u8, at: u64) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } } -pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void { +pub fn transform(gpa: std.mem.Allocator, fds: []platform.c.fd_t, t: Transform) !void { const raw = try gpa.alloc(u8, t.src.len); defer gpa.free(raw); try readAll(fds[t.src.shard], raw, t.src.off); @@ -110,7 +111,7 @@ test "fp32 values that are bf16 narrow in place; any other value is refused" { } const Pool = struct { - fds: []std.c.fd_t, + fds: []platform.c.fd_t, jobs: []const Copy, next: std.atomic.Value(usize) = .init(0), failed: std.atomic.Value(bool) = .init(false), @@ -132,15 +133,15 @@ const Pool = struct { /// Every copy on `threads` threads, uncached, paced by free memory (a large load otherwise fills the page cache too). pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: []const Copy, transforms: []const Transform, threads: usize) !void { - const fds = try gpa.alloc(std.c.fd_t, shards.len); + const fds = try gpa.alloc(platform.c.fd_t, shards.len); defer gpa.free(fds); for (shards, fds) |path, *fd| { - fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); + fd.* = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd.* < 0) return error.OpenFailed; - _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE + _ = platform.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE } defer for (fds) |fd| { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var pool: Pool = .{ .fds = fds, .jobs = copies }; const workers = try gpa.alloc(?std.Thread, threads); @@ -153,7 +154,7 @@ pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: [ const Tx = struct { gpa: std.mem.Allocator, jobs: []const Transform, - fds: []std.c.fd_t, + fds: []platform.c.fd_t, next: std.atomic.Value(usize) = .init(0), failed: std.atomic.Value(bool) = .init(false), fn run(x: *@This()) void { diff --git a/zig/src/families/glm/weights_test.zig b/zig/src/families/glm/weights_test.zig index 87e0c7fd01..92dc8ea962 100644 --- a/zig/src/families/glm/weights_test.zig +++ b/zig/src/families/glm/weights_test.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's loader on synthetic checkpoints: one sparse shard whose header names every tensor at its width. const std = @import("std"); +const platform = @import("sys"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); const quant = @import("quant.zig"); @@ -156,10 +157,10 @@ const Fixture = struct { std.mem.writeInt(u64, &len, head.items.len, .little); try tmp.dir.writeFile(io, .{ .sub_path = "model-00001.safetensors", .data = try std.mem.concat(f.a, u8, &.{ &len, head.items }) }); const path = try std.fmt.allocPrintSentinel(f.a, "{s}/model-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; } }; diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index f52e5a4a6c..ee3c374e4c 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/zig/src/families/nemotron/backend.zig @@ -1,5 +1,6 @@ //! The lane core's Metal backend for Nemotron: streams' rows in one forward a round, drawn tokens in GPU slots. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const Fence = @import("../../core/fence.zig").Fence; @@ -89,7 +90,7 @@ pub const Metal = struct { .tokens = try m.device.buffer(ring * 4, opts), .prompt = try m.device.buffer(@max(o.capacity, 16) * 4, opts), .fence = try Fence.init(m.device), - .trace = std.c.getenv("TF_SLIDE_TRACE") != null, + .trace = platform.c.getenv("TF_SLIDE_TRACE") != null, }; if (o.chunk > fused_rows) b.wide = try prefill.Prefill.init(m, @min(o.chunk, o.capacity)); if (o.drafts and m.weights.mtp != null) b.head = try mtp.Head.init(m.device, m.config, &m.kernels, &m.weights, &b.pool, max_lanes, o.capacity); diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..7483f02db6 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -1,6 +1,7 @@ //! The Nemotron CUDA engine (engine.py): weights, kernels and one sequence's buffers; windows and commits. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const Config = @import("config.zig").Config; const kern = @import("cuda_kernels.zig"); @@ -130,7 +131,7 @@ pub const Engine = struct { const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.* = .{ .gpa = gpa, .io = io, .ctx = ctx, .stream = undefined, .k = undefined, .w = undefined, .b = undefined, .c = undefined, .max_len = 0, .pinned = undefined, .history = undefined }; - e.trace = std.c.getenv("TF_SLIDE_TRACE") != null; + e.trace = platform.c.getenv("TF_SLIDE_TRACE") != null; e.c = try Config.read(gpa, io, model_dir); if (opts.segments < 1 or opts.segments > segs.MAX) return error.BadSegments; e.segments = opts.segments; diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 64594342c4..f09ad552e6 100644 --- a/zig/src/families/nemotron/cuda_native.zig +++ b/zig/src/families/nemotron/cuda_native.zig @@ -1,5 +1,6 @@ //! Nemotron's CUDA engine, MTP head and lane backend behind the native family interface. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); @@ -134,7 +135,7 @@ fn release(p: *anyopaque) void { /// TF_CUDA_CARVEOUT=1: the display carveout KV planes go to first (cuda/carveout.zig); null when it is off. fn openCarveout(gpa: std.mem.Allocator, ctx: *const cuda.Context) !?*cuda.Carveout { const bytes = (try cuda.carveout.requested(false, env("TF_CUDA_CARVEOUT"), env("TF_CUDA_CARVEOUT_MIB"))) orelse return null; - const card: [*:0]const u8 = std.c.getenv("TF_DRM_CARD") orelse cuda.carveout.default_card; + const card: [*:0]const u8 = platform.c.getenv("TF_DRM_CARD") orelse cuda.carveout.default_card; const c = try gpa.create(cuda.Carveout); errdefer gpa.destroy(c); c.* = cuda.Carveout.open(ctx.d, card, bytes) catch |err| { @@ -151,5 +152,5 @@ fn closeCarveout(gpa: std.mem.Allocator, c: *cuda.Carveout) void { } fn env(name: [:0]const u8) ?[]const u8 { - return std.mem.span(std.c.getenv(name) orelse return null); + return std.mem.span(platform.c.getenv(name) orelse return null); } diff --git a/zig/src/families/nemotron/cuda_seq_graphs.zig b/zig/src/families/nemotron/cuda_seq_graphs.zig index 9e51e1556b..09fb6965a4 100644 --- a/zig/src/families/nemotron/cuda_seq_graphs.zig +++ b/zig/src/families/nemotron/cuda_seq_graphs.zig @@ -1,6 +1,7 @@ //! Graphs a sequence other than the engine's own captures on its own buffers, once each, and replays after. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); /// One sequence's captured graphs, by the step they run (a window's or a head level's id). @@ -27,6 +28,6 @@ pub const Stats = struct { /// Whether TF_GRAPH_LOG=1 asks for a line a request: its captures, their time, its replays and eager windows. pub fn logFromEnv() bool { - const v = std.c.getenv("TF_GRAPH_LOG") orelse return false; + const v = platform.c.getenv("TF_GRAPH_LOG") orelse return false; return std.mem.eql(u8, std.mem.span(v), "1"); } diff --git a/zig/src/families/nemotron/forward.zig b/zig/src/families/nemotron/forward.zig index 7dc93e2de2..10751fbbf7 100644 --- a/zig/src/families/nemotron/forward.zig +++ b/zig/src/families/nemotron/forward.zig @@ -1,5 +1,6 @@ //! One Nemotron-H forward over streams' row segments, encoded as FusedDecode's kernels with their geometry. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -376,7 +377,7 @@ var split_mode: enum(u8) { unread, on, off } = .unread; /// TF_COOP_SPLIT=0 keeps one threadgroup a column tile, as the Python engine dispatches. fn splitOn() bool { if (split_mode == .unread) { - const v = std.c.getenv("TF_COOP_SPLIT"); + const v = platform.c.getenv("TF_COOP_SPLIT"); split_mode = if (v != null and v.?[0] == '0') .off else .on; } return split_mode == .on; diff --git a/zig/src/families/nemotron/gpu_round.zig b/zig/src/families/nemotron/gpu_round.zig index f8db78584b..6515ec7750 100644 --- a/zig/src/families/nemotron/gpu_round.zig +++ b/zig/src/families/nemotron/gpu_round.zig @@ -1,5 +1,6 @@ //! A lone greedy stream's rounds with no host wait: the GPU accepts, keeps rows and drafts; the host reads behind it. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const st = @import("state.zig"); @@ -349,7 +350,7 @@ pub fn run(gpa: std.mem.Allocator, b: *Metal, s: *lanes.Stream, o: Options, hook var mtp_known: usize = mtp0 + 1; var queued: u64 = 0; // rounds submitted var read: u64 = 0; // rounds whose records the host took - const log = std.c.getenv("TF_ROUND_LOG") != null; + const log = platform.c.getenv("TF_ROUND_LOG") != null; var depth = Depth{ .levels = o.depth, .most = o.depth, .fixed = o.fixed }; var windows: [ring]usize = undefined; // each queued round's verify rows var levels = depth.levels; // the levels the last queued head drafts (the next verify's rows - 1) @@ -502,7 +503,7 @@ pub fn run(gpa: std.mem.Allocator, b: *Metal, s: *lanes.Stream, o: Options, hook if (o.level_bench) { // every level live (zero bars), each chain alone in its command buffer r.o.stop = false; - if (std.c.getenv("TF_HEAD_GEOMETRY")) |g| b.head.?.geometry = std.fmt.parseInt(usize, std.mem.span(g), 10) catch 0; + if (platform.c.getenv("TF_HEAD_GEOMETRY")) |g| b.head.?.geometry = std.fmt.parseInt(usize, std.mem.span(g), 10) catch 0; var med: [2]f64 = undefined; for ([_]usize{ 1, @min(8, o.depth) }, 0..) |n, k| { var ms: [9]f64 = undefined; diff --git a/zig/src/families/nemotron/learned_prompts.zig b/zig/src/families/nemotron/learned_prompts.zig index 3f0cdfb1f6..7e7edc8818 100644 --- a/zig/src/families/nemotron/learned_prompts.zig +++ b/zig/src/families/nemotron/learned_prompts.zig @@ -1,5 +1,6 @@ //! Nemotron's learned prompt states (--learn): a kept state's file on disk, and the probe that keys this build's bits. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const snapshot = @import("snapshot.zig"); @@ -20,23 +21,23 @@ pub fn path(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { pub fn writeFile(snap: *const Snap, file: [:0]const u8) !void { var tmp_buf: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head = [HEAD]u32{ FILE_MAGIC, snap.at, @intFromBool(snap.head), @truncate(snap.bytes), @truncate(snap.bytes >> 32) }; try checked.write(fd, std.mem.sliceAsBytes(&head), snap.buf.contents()[0..snap.bytes]); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } /// `file`'s state into `snap`, whose buffer holds its bytes; a file of another position, layout or size is refused. pub fn readFile(snap: *Snap, file: [:0]const u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: [HEAD]u32 = undefined; try get(fd, std.mem.sliceAsBytes(&head), 0); const bytes = @as(u64, head[3]) | @as(u64, head[4]) << 32; @@ -47,7 +48,7 @@ pub fn readFile(snap: *Snap, file: [:0]const u8) !void { fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -56,7 +57,7 @@ fn put(fd: c_int, b: []const u8) !void { fn get(fd: c_int, b: []u8, at: u64) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.SnapshotRead; done += @intCast(n); } @@ -127,7 +128,7 @@ pub fn snapRead(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror! /// --learn: learned state `key`'s file removed. pub fn snapForget(_: *anyopaque, dir: [:0]const u8, key: u64) void { var buf: [1100]u8 = undefined; - _ = std.c.unlink(path(&buf, dir, key) catch return); + _ = platform.c.unlink(path(&buf, dir, key) catch return); } test "a learned state's file name is its key in hex" { diff --git a/zig/src/families/nemotron/learner.zig b/zig/src/families/nemotron/learner.zig index 0607eb90c1..69d70cb360 100644 --- a/zig/src/families/nemotron/learner.zig +++ b/zig/src/families/nemotron/learner.zig @@ -1,5 +1,6 @@ //! Sliding Weights' lessons, one learner for every backend over a host that gives the trainer, weights and files. const std = @import("std"); +const platform = @import("sys"); const adapters = @import("learn_shape.zig"); const choice = @import("choice.zig"); const router = @import("router.zig"); @@ -96,7 +97,7 @@ pub fn Learner(comptime H: type) type { pocketed: bool = false, // the last lesson put neurons in the pocket pub fn init(gpa: std.mem.Allocator, io: std.Io, b: H.Ctx) Self { - return .{ .gpa = gpa, .io = io, .b = b, .whole = std.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; + return .{ .gpa = gpa, .io = io, .b = b, .whole = platform.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; } pub fn deinit(l: *Self) void { diff --git a/zig/src/families/nemotron/living.zig b/zig/src/families/nemotron/living.zig index 1f0e8b073b..f479c17c2d 100644 --- a/zig/src/families/nemotron/living.zig +++ b/zig/src/families/nemotron/living.zig @@ -1,5 +1,6 @@ //! Saved lessons as shards beside the model's, a file a lesson (its blocks and router); both backends pass in `core`. const std = @import("std"); +const platform = @import("sys"); const Config = @import("config.zig").Config; const shape = @import("learn_shape.zig"); @@ -27,7 +28,7 @@ pub fn Io(comptime core: type) type { while (true) : (k += 1) { const path = try std.fs.path.joinZ(gpa, &.{ dir, try fileName(&buf, k) }); defer gpa.free(path); - if (std.c.unlink(path) != 0) break; + if (platform.c.unlink(path) != 0) break; } return bytes; } @@ -127,7 +128,7 @@ pub fn Io(comptime core: type) type { pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !?Saved { const old = try std.fs.path.joinZ(gpa, &.{ dir, legacy }); defer gpa.free(old); - if (std.c.access(old, std.c.F_OK) == 0) std.log.warn("slide: {s} holds lessons saved by an earlier version; they are not read", .{legacy}); + if (platform.c.access(old, platform.c.F_OK) == 0) std.log.warn("slide: {s} holds lessons saved by an earlier version; they are not read", .{legacy}); var lessons: std.ArrayList(Lesson) = .empty; errdefer { for (lessons.items) |*l| l.file.close(io); diff --git a/zig/src/families/nemotron/model.zig b/zig/src/families/nemotron/model.zig index 8ff4dae263..740846b96d 100644 --- a/zig/src/families/nemotron/model.zig +++ b/zig/src/families/nemotron/model.zig @@ -1,5 +1,6 @@ //! A loaded Nemotron-H model: config, checkpoint buffers, kernel-layout weights and compiled pipelines. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ckpt = @import("../../core/checkpoint_metal.zig"); const shards = @import("../../core/checkpoint.zig"); @@ -84,7 +85,7 @@ pub const Model = struct { if (mtp) { const path = try std.fmt.allocPrintSentinel(allocator, "{s}/mtp-4bit.safetensors", .{dir}, 0); defer allocator.free(path); - if (std.c.access(path, 0) == 0) try m.checkpoint.addFile(m.device, path, "mtp."); + if (platform.c.access(path, 0) == 0) try m.checkpoint.addFile(m.device, path, "mtp."); } const drafting = m.checkpoint.has("mtp.layers.0.eh_proj.weight"); m.draft_ids = if (drafting) try draft_ids.load(allocator, m.config.vocab) else &.{}; diff --git a/zig/src/families/nemotron/pocket.zig b/zig/src/families/nemotron/pocket.zig index ff24ac5ee9..90232ef763 100644 --- a/zig/src/families/nemotron/pocket.zig +++ b/zig/src/families/nemotron/pocket.zig @@ -1,5 +1,6 @@ //! A pocket of new neurons beside each MoE layer's shared expert, where facts are learned; nothing the model had moves. const std = @import("std"); +const platform = @import("sys"); const examples = @import("lesson_examples.zig"); const lesson_route = @import("lesson_route.zig"); const pocket_fit = @import("pocket_fit.zig"); @@ -52,7 +53,7 @@ const Refit = struct { slot: usize, neuron: usize }; /// Whether environment variable `name` is set to `value`. fn set(name: [*:0]const u8, value: []const u8) bool { - const v = std.c.getenv(name) orelse return false; + const v = platform.c.getenv(name) orelse return false; return std.mem.eql(u8, std.mem.span(v), value); } @@ -136,7 +137,7 @@ pub fn Pocket(comptime T: type) type { read: reading.Seen = .{}, // what the last read pass read, which a lesson begun again over the same keeps pub fn init(gpa: std.mem.Allocator) Self { - const asked: Mode = if (std.c.getenv("TF_SLIDE_POCKET_KEY")) |v| std.meta.stringToEnum(Mode, std.mem.span(v)) orelse .auto else .auto; + const asked: Mode = if (platform.c.getenv("TF_SLIDE_POCKET_KEY")) |v| std.meta.stringToEnum(Mode, std.mem.span(v)) orelse .auto else .auto; return .{ .gpa = gpa, .mode = asked, .strict = !set("TF_SLIDE_POCKET_LEADS", "off"), .keeping = !set("TF_SLIDE_POCKET_KEEPS", "off"), .valuing = set("TF_SLIDE_POCKET_VALUES", "on"), .silencing = set("TF_SLIDE_POCKET_SAME", "silent"), .worst = set("TF_SLIDE_POCKET_REPLAY", "worst"), .joining = set("TF_SLIDE_POCKET_GROUP", "joint") }; } @@ -360,7 +361,7 @@ pub fn Pocket(comptime T: type) type { const d = c.hidden; const n = p.slots(); const m = p.labels.items.len; - if (std.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); + if (platform.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); const kinds = try p.gpa.alloc(Kind, 3 * m); defer p.gpa.free(kinds); for (p.labels.items, 0..) |l, i| for (0..3) |k| { diff --git a/zig/src/families/nemotron/pocket_dump.zig b/zig/src/families/nemotron/pocket_dump.zig index 1653589631..b1f5d152a6 100644 --- a/zig/src/families/nemotron/pocket_dump.zig +++ b/zig/src/families/nemotron/pocket_dump.zig @@ -1,11 +1,12 @@ //! A lesson's read rows written out for offline study (TF_SLIDE_POCKET_DUMP=dir): inputs, labels and examples. const std = @import("std"); +const platform = @import("sys"); /// `bytes` written to `path`. fn put(path: [:0]const u8, bytes: []const u8) !void { - const f = std.c.fopen(path, "wb") orelse return error.DumpOpen; - defer _ = std.c.fclose(f); - if (std.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; + const f = platform.c.fopen(path, "wb") orelse return error.DumpOpen; + defer _ = platform.c.fclose(f); + if (platform.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; } /// Pocket `p`'s read rows under `dir` for `tag`: bf16 [read, slots, D], then their labels and examples as text. diff --git a/zig/src/families/nemotron/whole.zig b/zig/src/families/nemotron/whole.zig index 4fa4a9fc01..b8f6fc8a99 100644 --- a/zig/src/families/nemotron/whole.zig +++ b/zig/src/families/nemotron/whole.zig @@ -1,5 +1,6 @@ //! Whole-mode lessons: pocket neurons (default), a weight edit (TF_SLIDE_EDIT) or a probe (TF_SLIDE_PROBE). const std = @import("std"); +const platform = @import("sys"); const examples = @import("lesson_examples.zig"); const adapters = @import("learn_shape.zig"); @@ -10,8 +11,8 @@ pub const Method = enum { pocket, edit, probe }; /// The method the environment picks for whole lessons. pub fn method() Method { - if (std.c.getenv("TF_SLIDE_PROBE") != null) return .probe; - if (std.c.getenv("TF_SLIDE_EDIT") != null) return .edit; + if (platform.c.getenv("TF_SLIDE_PROBE") != null) return .probe; + if (platform.c.getenv("TF_SLIDE_EDIT") != null) return .edit; return .pocket; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 53771c8eda..428f764854 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -1,5 +1,6 @@ //! A native text model keeps immutable checkpoint bytes and prepared affine layouts under one owner. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const core = @import("core"); const cfg = @import("config.zig"); @@ -71,7 +72,7 @@ pub const Model = struct { m.frame = try Frame.init(m.device, c, rows); errdefer m.frame.deinit(); m.ane = null; - const columns = if (std.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; + const columns = if (platform.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available() and m.weights.uniformMlp(); const fits = if (share_fits) |f| !share or f(ane_mlp.Share.bytes(c, columns)) else true; if (!fits) std.log.info("qwen27: the Neural Engine's prompt share ({d} MiB of fp16 weights) would take this Mac past its load limit; prompts run on the GPU alone", .{ane_mlp.Share.bytes(c, columns) >> 20}); diff --git a/zig/src/hip/driver.zig b/zig/src/hip/driver.zig index 0982c96660..7fa1ef9cca 100644 --- a/zig/src/hip/driver.zig +++ b/zig/src/hip/driver.zig @@ -1,5 +1,6 @@ //! The first HIP admission boundary, with signatures from ROCm's hip_runtime_api.h; no CUDA symbol aliases or GPU kernels. const std = @import("std"); +const platform = @import("sys"); pub const Error = error{ DriverUnavailable, MissingSymbol, HipFailed, InvalidDeviceCount }; pub const Api = struct { @@ -8,11 +9,11 @@ pub const Api = struct { }; pub const Driver = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn openPath(path: []const u8) Error!Driver { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); const api = try @import("symbols.zig").resolve(Api, &lib); try initialize(api); diff --git a/zig/src/hip/runtime.zig b/zig/src/hip/runtime.zig index 8f5527bd84..2583cf6bbd 100644 --- a/zig/src/hip/runtime.zig +++ b/zig/src/hip/runtime.zig @@ -1,10 +1,11 @@ //! Complete model-free HIP runtime admission; partial symbol tables never escape. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Error = error{ DriverUnavailable, MissingSymbol, HipFailed, Invalid }; pub const Runtime = struct { - lib: std.DynLib, + lib: platform.DynLib, api: abi.Api, pub fn open() Error!Runtime { @@ -12,7 +13,7 @@ pub const Runtime = struct { } pub fn openPath(path: []const u8) Error!Runtime { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); const api = try @import("symbols.zig").resolve(abi.Api, &lib); try check(api.hipInit(0)); diff --git a/zig/src/hip/symbols.zig b/zig/src/hip/symbols.zig index e334f716cb..54c2c14bc6 100644 --- a/zig/src/hip/symbols.zig +++ b/zig/src/hip/symbols.zig @@ -1,7 +1,8 @@ //! One symbol-resolution path for admission fixtures and the full runtime. const std = @import("std"); +const platform = @import("sys"); -pub fn resolve(comptime Api: type, lib: *std.DynLib) error{MissingSymbol}!Api { +pub fn resolve(comptime Api: type, lib: *platform.DynLib) error{MissingSymbol}!Api { var api: Api = undefined; const info = @typeInfo(Api).@"struct"; inline for (info.field_names, info.field_types) |name, T| { diff --git a/zig/src/learned_warning_test.zig b/zig/src/learned_warning_test.zig index cb9a9346b1..382ac2659e 100644 --- a/zig/src/learned_warning_test.zig +++ b/zig/src/learned_warning_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const pc = @import("core/prompt_cache.zig"); const Imprint = @import("core/prompt_imprint.zig").Imprint; const Fake = @import("core/prompt_cache_fake.zig").Fake; @@ -56,7 +57,7 @@ fn expectedWarnings(n: usize) !void { pub fn main() !void { const a = std.heap.page_allocator; var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-warning-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-warning-{d}", .{platform.c.getpid()}); defer @import("core/prompt_imprint_test.zig").rmTree(root); { var im = try Imprint.open(a, root, 1, 1 << 20); diff --git a/zig/src/metal/clock.zig b/zig/src/metal/clock.zig index 6461c8e227..9e0572c597 100644 --- a/zig/src/metal/clock.zig +++ b/zig/src/metal/clock.zig @@ -1,13 +1,14 @@ //! Host time in seconds on the clock Metal's GPUStartTime and GPUEndTime use (mach_absolute_time). const std = @import("std"); +const platform = @import("sys"); var scale: f64 = 0; pub fn seconds() f64 { if (scale == 0) { - var info: std.c.mach_timebase_info_data = undefined; - _ = std.c.mach_timebase_info(&info); + var info: platform.c.mach_timebase_info_data = undefined; + _ = platform.c.mach_timebase_info(&info); scale = @as(f64, @floatFromInt(info.numer)) / @as(f64, @floatFromInt(info.denom)) / 1e9; } - return @as(f64, @floatFromInt(std.c.mach_absolute_time())) * scale; + return @as(f64, @floatFromInt(platform.c.mach_absolute_time())) * scale; } diff --git a/zig/src/metal/device.zig b/zig/src/metal/device.zig index 69d67e108c..353cd6d8ce 100644 --- a/zig/src/metal/device.zig +++ b/zig/src/metal/device.zig @@ -1,5 +1,6 @@ //! The GPU device: queues, buffers (zero-copy over mmap too), events and residency sets. const std = @import("std"); +const platform = @import("sys"); const objc = @import("objc.zig"); const types = @import("types.zig"); const command = @import("command.zig"); @@ -139,10 +140,10 @@ pub const MappedFile = struct { size: usize, pub fn open(path: [*:0]const u8) !MappedFile { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end <= 0) return error.EmptyFile; const size: usize = @intCast(end); const page = std.heap.pageSize(); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..50af1bd955 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,5 +1,6 @@ //! CUDA family registration, device setup, memory admission and the native lane host. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); @@ -44,7 +45,7 @@ fn chipClass(a: Allocator, capability: u32) ?[]const u8 { } fn getenv(name: [:0]const u8) ?[]const u8 { - return std.mem.span(std.c.getenv(name) orelse return null); + return std.mem.span(platform.c.getenv(name) orelse return null); } /// A variable's whole number: null when unset or empty, an error for anything else that is not one. diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index cdba9b907b..539964a682 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -1,6 +1,7 @@ //! Flash Next behind the native server: one reply at a time on the replay engine (prompt chunks, then GPU-side //! rounds), greedy only; the engine's thread owns the GPU, requests wait in arrival order (foreground first). const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -420,7 +421,7 @@ const Snaps = struct { /// --learn: learned state `key`'s file removed. fn forget(_: *anyopaque, dir: [:0]const u8, key: u64) void { var buf: [1100]u8 = undefined; - _ = std.c.unlink(learned.path(&buf, dir, key) catch return); + _ = platform.c.unlink(learned.path(&buf, dir, key) catch return); } }; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 39fe531408..8370f45ed7 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -1,5 +1,6 @@ //! The Metal engines: lane core (Nemotron, Qwen3.5-2B, GLM, Qwen3.8-27B), Flash Next replay. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -260,9 +261,9 @@ fn shardStamps(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8) ! /// A safetensors file's header (its tensors' names, shapes and offsets), in `a`. fn safetensorsHeader(a: Allocator, dir: []const u8, name: []const u8) ![]const u8 { const path = try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.FileNotFound; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var n: u64 = 0; if (!api.prompt_imprint.readAt(fd, std.mem.asBytes(&n), 0) or n > 1 << 24) return error.BadHeader; const head = try a.alloc(u8, @intCast(n)); @@ -305,7 +306,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ /// Flash Next uses the dump named by TF_FLASHNEXT_DUMP, or the checked-in kernels when that variable is unset. fn openFlashNext(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { - const dump: ?[]const u8 = if (std.c.getenv("TF_FLASHNEXT_DUMP")) |d| std.mem.span(d) else null; + const dump: ?[]const u8 = if (platform.c.getenv("TF_FLASHNEXT_DUMP")) |d| std.mem.span(d) else null; const native = modelContext(a, io, o.dir); const window: i64 = o.context orelse native; if (window < 0 or (native > 0 and window > native)) { @@ -373,7 +374,7 @@ fn openGlmDsa(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ } var ram: u64 = 0; var len: usize = 8; - if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; + if (platform.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; const pool = mtl.objc.Pool.push(); defer pool.pop(); const h = glm_dsa.open(gpa, io, o.dir, @intCast(window), settings, @max(1, @min(o.lanes, 16)), ram / 100 * 66, o.prompt_cache_gib) catch |e| { @@ -469,10 +470,10 @@ test "a shard rewritten in place, same size and bytes, gives the learned states try std.testing.expectEqual(first.final(), notes.final()); std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; const path = try std.fmt.allocPrintSentinel(a, "{s}/model-00001-of-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); try std.testing.expect(fd >= 0); - try std.testing.expectEqual(@as(isize, 7), std.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size - _ = std.c.close(fd); + try std.testing.expectEqual(@as(isize, 7), platform.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size + _ = platform.c.close(fd); var moved = std.hash.Wyhash.init(0); try shardStamps(&moved, a, io, dir); try std.testing.expect(first.final() != moved.final()); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index e09a508a8c..e7da14ef17 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -1,5 +1,6 @@ //! The 27B's lane core backend: a stream a runner slot, shared rounds across slots, DFlash2 trees for slot 0. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -125,7 +126,7 @@ pub const Metal = struct { const drafting = b.draft != null; const shared = b.runner.slots > 1; // a lone stream drafts DFlash2 trees, as in 1.0.5: copies land less often there (TF_QWEN27_COPIES=1 opts in) - const lone_copies = if (std.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const lone_copies = if (platform.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; return .{ .exact_width = if (drafting or shared) max_rows else 1, .mtp = drafting, .speculate = drafting, .speculate_early = false, .drafts = b.nodes, .draft_probabilities = true, .window_costs = b.costs[0..b.timed], .hidden_rows = shared, .streams_exact = true, .batch_rows = if (shared) batch_rows else max_rows, .max_streams = b.runner.slots, .lone_copies = lone_copies }; } diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 09da59e53f..514cecdc8d 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -1,5 +1,6 @@ //! Qwen3.8-27B on the lane host: a stream a runner slot in shared rounds, DFlash2 drafts for slot 0's stream. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const mtl = @import("metal"); const api = @import("engine_api"); @@ -289,7 +290,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off // a lone greedy request takes the serial loop on tensor units, where it is faster (TF_QWEN27_LANES overrides) const units = model.device.tensorUnits(); - const serial = if (std.c.getenv("TF_QWEN27_LANES")) |v| !std.mem.eql(u8, std.mem.span(v), "1") else units; + const serial = if (platform.c.getenv("TF_QWEN27_LANES")) |v| !std.mem.eql(u8, std.mem.span(v), "1") else units; if (h.draft != null and serial) h.host.lone = .{ .ctx = h, .run = Host.lone }; h.warm = .{ .queue = model.queue }; // the ready slots are wired before the cache's budget reads the footprint, and the later ones count as taken diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig new file mode 100644 index 0000000000..ee614d7507 --- /dev/null +++ b/zig/src/platform/sys.zig @@ -0,0 +1,47 @@ +//! The operating system as TensorFold calls it: `c` is libc's surface (std.c on POSIX, win_c.zig on Windows) and +//! `DynLib` opens shared libraries (std.DynLib on POSIX, LoadLibraryW on Windows). Every module imports this as `sys`. + +const std = @import("std"); +const builtin = @import("builtin"); + +pub const windows = builtin.os.tag == .windows; + +pub const c = if (windows) @import("win_c.zig") else std.c; + +/// std.posix, or its sockets, signals and limits over Winsock and console events. +pub const posix = if (windows) @import("win_posix.zig") else std.posix; + +pub const DynLib = if (windows) WinDynLib else std.DynLib; + +/// The CUDA driver's file name: nvcuda.dll comes with the display driver on Windows. +pub const cuda_driver = if (windows) "nvcuda.dll" else "libcuda.so.1"; + +extern "kernel32" fn LoadLibraryW(name: [*:0]const u16) callconv(.winapi) ?*anyopaque; +extern "kernel32" fn GetProcAddress(module: *anyopaque, name: [*:0]const u8) callconv(.winapi) ?*anyopaque; +extern "kernel32" fn FreeLibrary(module: *anyopaque) callconv(.winapi) c_int; + +const WinDynLib = struct { + module: *anyopaque, + + pub const Error = error{ FileNotFound, NameTooLong, InvalidWtf8 }; + + pub fn open(path: []const u8) Error!WinDynLib { + var buf: [1024:0]u16 = undefined; + const n = std.unicode.wtf8ToWtf16Le(buf[0 .. buf.len - 1], path) catch return error.InvalidWtf8; + buf[n] = 0; + return .{ .module = LoadLibraryW(buf[0..n :0].ptr) orelse return error.FileNotFound }; + } + + pub fn openZ(path: [*:0]const u8) Error!WinDynLib { + return open(std.mem.span(path)); + } + + pub fn close(self: *WinDynLib) void { + _ = FreeLibrary(self.module); + } + + pub fn lookup(self: *WinDynLib, comptime T: type, name: [:0]const u8) ?T { + const p = GetProcAddress(self.module, name.ptr) orelse return null; + return @ptrCast(@alignCast(p)); + } +}; diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig new file mode 100644 index 0000000000..82b4dc7b3f --- /dev/null +++ b/zig/src/platform/win_c.zig @@ -0,0 +1,577 @@ +//! The libc surface TensorFold calls through `sys.c`, on Windows: integer descriptors over the UCRT (files opened by +//! CreateFileW with every share mode, so unlink and rename behave as on POSIX), directory descriptors as remembered +//! paths, readdir over FindFirstFileW, and POSIX's names for the rest. Each call keeps libc's contract: -1 and errno. + +const std = @import("std"); + +pub const fd_t = c_int; +pub const mode_t = c_uint; +pub const off_t = i64; +pub const pid_t = c_int; +pub const time_t = i64; +pub const E = std.c.E; +pub const errno = std.c.errno; +pub const timespec = extern struct { sec: time_t, nsec: c_long }; + +pub const SEEK = struct { + pub const SET: c_int = 0; + pub const CUR: c_int = 1; + pub const END: c_int = 2; +}; +pub const F_OK: c_int = 0; +pub const R_OK: c_int = 4; +pub const W_OK: c_int = 2; +pub const X_OK: c_int = 0; // Windows has no execute bit; existence is the answer +pub const AT = struct { + pub const FDCWD: fd_t = -100; + pub const REMOVEDIR: u32 = 0x200; +}; +pub const DT = struct { + pub const UNKNOWN: u8 = 0; + pub const DIR: u8 = 4; + pub const REG: u8 = 8; + pub const LNK: u8 = 10; +}; +pub const S = struct { + pub const IFMT: mode_t = 0o170000; + pub const IFDIR: mode_t = 0o040000; + pub const IFREG: mode_t = 0o100000; + pub const IFLNK: mode_t = 0o120000; +}; +pub const RLIM = struct { + pub const INFINITY: u64 = std.math.maxInt(u64); +}; +pub const timeval = extern struct { sec: time_t, usec: c_long }; +pub const CLOCK = struct { + pub const REALTIME: c_int = 0; + pub const MONOTONIC: c_int = 1; +}; + +/// Linux's open flags by name; DIRECT asks for unbuffered reads (FILE_FLAG_NO_BUFFERING), the rest map or are moot. +pub const O = packed struct(u32) { + ACCMODE: enum(u2) { RDONLY = 0, WRONLY = 1, RDWR = 2 } = .RDONLY, + CREAT: bool = false, + EXCL: bool = false, + NOCTTY: bool = false, + TRUNC: bool = false, + APPEND: bool = false, + NONBLOCK: bool = false, + DSYNC: bool = false, + ASYNC: bool = false, + DIRECT: bool = false, + LARGEFILE: bool = false, + DIRECTORY: bool = false, + NOFOLLOW: bool = false, + NOATIME: bool = false, + CLOEXEC: bool = false, + SYNC: bool = false, + PATH: bool = false, + TMPFILE: bool = false, + _: u13 = 0, +}; + +// ---- Win32 and UCRT entry points ---- + +const HANDLE = *anyopaque; +const BOOL = c_int; +const invalid_handle: isize = -1; + +const OVERLAPPED = extern struct { internal: usize = 0, internal_high: usize = 0, offset: u32, offset_high: u32, event: ?HANDLE = null }; +const FILETIME = extern struct { low: u32, high: u32 }; +const FindData = extern struct { + attributes: u32, + created: FILETIME, + accessed: FILETIME, + written: FILETIME, + size_high: u32, + size_low: u32, + reserved0: u32, + reserved1: u32, + name: [260]u16, + alternate: [14]u16, +}; + +extern "kernel32" fn CreateFileW(name: [*:0]const u16, access: u32, share: u32, security: ?*anyopaque, disposition: u32, flags: u32, template: ?HANDLE) callconv(.winapi) isize; +extern "kernel32" fn ReadFile(h: HANDLE, buf: [*]u8, len: u32, got: ?*u32, ov: ?*OVERLAPPED) callconv(.winapi) BOOL; +extern "kernel32" fn WriteFile(h: HANDLE, buf: [*]const u8, len: u32, put: ?*u32, ov: ?*OVERLAPPED) callconv(.winapi) BOOL; +extern "kernel32" fn GetLastError() callconv(.winapi) u32; +extern "kernel32" fn GetFileAttributesW(name: [*:0]const u16) callconv(.winapi) u32; +extern "kernel32" fn MoveFileExW(from: [*:0]const u16, to: [*:0]const u16, flags: u32) callconv(.winapi) BOOL; +extern "kernel32" fn FindFirstFileW(pattern: [*:0]const u16, data: *FindData) callconv(.winapi) isize; +extern "kernel32" fn FindNextFileW(h: HANDLE, data: *FindData) callconv(.winapi) BOOL; +extern "kernel32" fn FindClose(h: HANDLE) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; +extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; +extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; +extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; +extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; +extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; +extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; +extern "kernel32" fn CloseHandle(h: HANDLE) callconv(.winapi) BOOL; + +extern "c" fn _errno() *c_int; +extern "c" fn _open_osfhandle(h: isize, flags: c_int) c_int; +extern "c" fn _get_osfhandle(fd: c_int) isize; +extern "c" fn _close(fd: c_int) c_int; +extern "c" fn _read(fd: c_int, buf: [*]u8, len: c_uint) c_int; +extern "c" fn _write(fd: c_int, buf: [*]const u8, len: c_uint) c_int; +extern "c" fn _lseeki64(fd: c_int, off: i64, whence: c_int) i64; +extern "c" fn _commit(fd: c_int) c_int; +extern "c" fn _chsize_s(fd: c_int, size: i64) c_int; +extern "c" fn _wunlink(path: [*:0]const u16) c_int; +extern "c" fn _wrmdir(path: [*:0]const u16) c_int; +extern "c" fn _wmkdir(path: [*:0]const u16) c_int; +extern "c" fn _waccess(path: [*:0]const u16, mode: c_int) c_int; +extern "c" fn _getpid() c_int; +extern "c" fn _isatty(fd: c_int) c_int; +extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; +extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; +pub extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +pub extern "c" fn fopen(path: [*:0]const u8, mode: [*:0]const u8) ?*std.c.FILE; +pub extern "c" fn fwrite(ptr: [*]const u8, size: usize, n: usize, f: *std.c.FILE) usize; +pub extern "c" fn fclose(f: *std.c.FILE) c_int; +pub extern "c" fn free(ptr: ?*anyopaque) void; +pub extern "c" fn malloc(n: usize) ?*anyopaque; + +/// The UCRT's struct _stat64. +const Stat64 = extern struct { dev: u32, ino: u16, mode: u16, nlink: i16, uid: i16, gid: i16, rdev: u32, size: i64, atime: i64, mtime: i64, ctime: i64 }; + +fn fail(e: E) c_int { + _errno().* = @intFromEnum(e); + return -1; +} + +/// The errno a failed Win32 call means. +fn lastError() E { + return switch (GetLastError()) { + 2, 3, 15, 123, 161 => .NOENT, // file, path, drive not found; bad name + 5, 32, 33 => .ACCES, // access denied, sharing and lock violations + 80, 183 => .EXIST, + 112 => .NOSPC, + 6 => .BADF, + 145 => .NOTEMPTY, + 267 => .NOTDIR, + 87 => .INVAL, + else => .IO, + }; +} + +// ---- paths ---- + +const Wide = [32767:0]u16; + +/// UTF-8 path to a NUL-terminated wide path in `buf`; null (errno set) when it does not fit or is not UTF-8. +/// POSIX's /tmp is the user's temporary folder and /proc/self/exe this executable, as Cygwin and MSYS map them. +fn wide(buf: *Wide, path: []const u8) ?[*:0]const u16 { + var at: usize = 0; + var rest = path; + if (std.mem.eql(u8, path, "/proc/self/exe")) { + const n = GetModuleFileNameW(null, buf, buf.len - 1); + if (n == 0 or n >= buf.len - 1) { + _ = fail(.NAMETOOLONG); + return null; + } + buf[n] = 0; + return buf[0..n :0].ptr; + } + if (std.mem.startsWith(u8, path, "/tmp") and (path.len == 4 or path[4] == '/')) { + at = GetTempPathW(buf.len - 1, buf); // ends in a separator + if (at == 0 or at >= buf.len - 1) { + _ = fail(.NAMETOOLONG); + return null; + } + if (path.len == 4) at -= 1; + rest = if (path.len > 5) path[5..] else ""; + } + const n = std.unicode.wtf8ToWtf16Le(buf[at .. buf.len - 1], rest) catch { + _ = fail(.NAMETOOLONG); + return null; + }; + buf[at + n] = 0; + return buf[0 .. at + n :0].ptr; +} + +/// Directory descriptors: open(O.DIRECTORY) remembers the path; openat and unlinkat join onto it. +const dir_base: fd_t = 0x4000_0000; +var dirs_lock: std.atomic.Mutex = .unlocked; +var dir_paths: [256]?[:0]u8 = @splat(null); + +fn lockDirs() void { + while (!dirs_lock.tryLock()) std.atomic.spinLoopHint(); +} + +fn dirPath(fd: fd_t) ?[:0]const u8 { + if (fd < dir_base or fd >= dir_base + dir_paths.len) return null; + lockDirs(); + defer dirs_lock.unlock(); + return dir_paths[@intCast(fd - dir_base)]; +} + +fn openDir(path: []const u8) c_int { + var buf: Wide = undefined; + const w = wide(&buf, path) orelse return -1; + const attrs = GetFileAttributesW(w); + if (attrs == 0xFFFF_FFFF) return fail(lastError()); + if (attrs & 0x10 == 0) return fail(.NOTDIR); + const copy = std.heap.c_allocator.dupeSentinel(u8, path, 0) catch return fail(.NOMEM); + lockDirs(); + defer dirs_lock.unlock(); + for (&dir_paths, 0..) |*slot, i| if (slot.* == null) { + slot.* = copy; + return dir_base + @as(fd_t, @intCast(i)); + }; + std.heap.c_allocator.free(copy); + return fail(.MFILE); +} + +/// `name` under the directory `fd` (or the working directory, or absolute) in `buf`. +fn joinAt(buf: []u8, fd: fd_t, name: [*:0]const u8) ?[:0]const u8 { + const leaf = std.mem.span(name); + if (fd == AT.FDCWD or std.fs.path.isAbsolute(leaf)) return std.fmt.bufPrintSentinel(buf, "{s}", .{leaf}, 0) catch null; + const base = dirPath(fd) orelse { + _ = fail(.BADF); + return null; + }; + return std.fmt.bufPrintSentinel(buf, "{s}\\{s}", .{ base, leaf }, 0) catch { + _ = fail(.NAMETOOLONG); + return null; + }; +} + +// ---- descriptors ---- + +/// Variadic as libc's is, so callers pass a mode or not; the mode is moot on NTFS without ACL work. +pub fn open(path: [*:0]const u8, flags: O, ...) callconv(.c) c_int { + return openFile(path, flags); +} + +fn openFile(path: [*:0]const u8, flags: O) c_int { + const p = std.mem.span(path); + if (flags.DIRECTORY) return openDir(p); + var buf: Wide = undefined; + const w = wide(&buf, p) orelse return -1; + const generic_read: u32 = 0x8000_0000; + const generic_write: u32 = 0x4000_0000; + const rights: u32 = switch (flags.ACCMODE) { + .RDONLY => generic_read, + .WRONLY => generic_write, + .RDWR => generic_read | generic_write, + }; + const disposition: u32 = if (flags.CREAT and flags.EXCL) 1 // CREATE_NEW + else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS + else if (flags.CREAT) 4 // OPEN_ALWAYS + else if (flags.TRUNC) 5 // TRUNCATE_EXISTING + else 3; // OPEN_EXISTING + var attrs: u32 = 0x80; // FILE_ATTRIBUTE_NORMAL + if (flags.DIRECT) attrs |= 0x2000_0000; // FILE_FLAG_NO_BUFFERING + const h = CreateFileW(w, rights, 0x7, null, disposition, attrs, null); // share read, write and delete + if (h == invalid_handle) { + // a directory opened without O.DIRECTORY, as POSIX allows for reading its entries + if (GetLastError() == 5 and flags.ACCMODE == .RDONLY and GetFileAttributesW(w) & 0x10 != 0 and GetFileAttributesW(w) != 0xFFFF_FFFF) return openDir(p); + return fail(lastError()); + } + var crt: c_int = 0x8000; // _O_BINARY + crt |= switch (flags.ACCMODE) { + .RDONLY => 0, + .WRONLY => 1, + .RDWR => 2, + }; + if (flags.APPEND) crt |= 0x8; + const fd = _open_osfhandle(h, crt); + if (fd < 0) { + _ = CloseHandle(@ptrFromInt(@as(usize, @bitCast(h)))); + return fail(.MFILE); + } + return fd; +} + +pub fn openat(dir: fd_t, path: [*:0]const u8, flags: O, ...) callconv(.c) c_int { + var buf: [4096]u8 = undefined; + const p = joinAt(&buf, dir, path) orelse return -1; + return openFile(p.ptr, flags); +} + +pub fn close(fd: fd_t) c_int { + if (fd >= dir_base and fd < dir_base + dir_paths.len) { + lockDirs(); + defer dirs_lock.unlock(); + const slot = &dir_paths[@intCast(fd - dir_base)]; + const p = slot.* orelse return fail(.BADF); + std.heap.c_allocator.free(p); + slot.* = null; + return 0; + } + return _close(fd); +} + +const chunk: usize = 1 << 30; + +pub fn read(fd: fd_t, buf: [*]u8, len: usize) isize { + return _read(fd, buf, @intCast(@min(len, chunk))); +} + +pub fn write(fd: fd_t, buf: [*]const u8, len: usize) isize { + return _write(fd, buf, @intCast(@min(len, chunk))); +} + +fn osHandle(fd: fd_t) ?HANDLE { + const h = _get_osfhandle(fd); + if (h == invalid_handle or h == -2) { + _ = fail(.BADF); + return null; + } + return @ptrFromInt(@as(usize, @bitCast(h))); +} + +/// Positional read by an OVERLAPPED offset; end of file reads 0 as pread does. +pub fn pread(fd: fd_t, buf: [*]u8, len: usize, offset: off_t) isize { + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = @truncate(@as(u64, @bitCast(offset))), .offset_high = @truncate(@as(u64, @bitCast(offset)) >> 32) }; + var got: u32 = 0; + if (ReadFile(h, buf, @intCast(@min(len, chunk)), &got, &ov) == 0) { + if (GetLastError() == 38) return 0; // ERROR_HANDLE_EOF + return fail(lastError()); + } + return got; +} + +pub fn pwrite(fd: fd_t, buf: [*]const u8, len: usize, offset: off_t) isize { + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = @truncate(@as(u64, @bitCast(offset))), .offset_high = @truncate(@as(u64, @bitCast(offset)) >> 32) }; + var put: u32 = 0; + if (WriteFile(h, buf, @intCast(@min(len, chunk)), &put, &ov) == 0) return fail(lastError()); + return put; +} + +pub fn lseek(fd: fd_t, offset: off_t, whence: c_int) off_t { + return _lseeki64(fd, offset, whence); +} + +pub fn fsync(fd: fd_t) c_int { + if (dirPath(fd) != null) return 0; // directory entries are durable once NTFS has them + return _commit(fd); +} + +pub fn ftruncate(fd: fd_t, size: off_t) c_int { + const rc = _chsize_s(fd, size); + if (rc != 0) return fail(@enumFromInt(rc)); + return 0; +} + +pub fn isatty(fd: fd_t) c_int { + return _isatty(fd); +} + +pub const Stat = struct { + dev: u64, + ino: u64, + mode: mode_t, + size: i64, + mtim: timespec, + + pub fn mtime(s: Stat) timespec { + return s.mtim; + } +}; + +pub fn fstat(fd: fd_t, out: *Stat) c_int { + if (dirPath(fd) != null) { + out.* = .{ .dev = 0, .ino = 0, .mode = S.IFDIR | 0o700, .size = 0, .mtim = .{ .sec = 0, .nsec = 0 } }; + return 0; + } + var st: Stat64 = undefined; + if (_fstat64(fd, &st) != 0) return -1; + out.* = .{ .dev = st.dev, .ino = st.ino, .mode = st.mode, .size = st.size, .mtim = .{ .sec = st.mtime, .nsec = 0 } }; + return 0; +} + +// ---- names ---- + +pub fn unlink(path: [*:0]const u8) c_int { + var buf: Wide = undefined; + return _wunlink(wide(&buf, std.mem.span(path)) orelse return -1); +} + +pub fn rmdir(path: [*:0]const u8) c_int { + var buf: Wide = undefined; + return _wrmdir(wide(&buf, std.mem.span(path)) orelse return -1); +} + +pub fn unlinkat(dir: fd_t, path: [*:0]const u8, flags: u32) c_int { + var buf: [4096]u8 = undefined; + const p = joinAt(&buf, dir, path) orelse return -1; + return if (flags & AT.REMOVEDIR != 0) rmdir(p.ptr) else unlink(p.ptr); +} + +pub fn mkdir(path: [*:0]const u8, mode: mode_t) c_int { + _ = mode; + var buf: Wide = undefined; + return _wmkdir(wide(&buf, std.mem.span(path)) orelse return -1); +} + +/// POSIX rename: replaces an existing target. +pub fn rename(from: [*:0]const u8, to: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wa = wide(&a, std.mem.span(from)) orelse return -1; + const wb = wide(&b, std.mem.span(to)) orelse return -1; + if (MoveFileExW(wa, wb, 0x1 | 0x8) == 0) return fail(lastError()); // REPLACE_EXISTING | WRITE_THROUGH + return 0; +} + +pub fn access(path: [*:0]const u8, mode: c_int) c_int { + var buf: Wide = undefined; + return _waccess(wide(&buf, std.mem.span(path)) orelse return -1, mode & 6); +} + +// ---- directories ---- + +pub const dirent = extern struct { + type: u8, + name: [1024]u8, +}; + +pub const DIR = struct { + find: ?HANDLE, + data: FindData, + first: bool, + fd: ?fd_t, + ent: dirent, +}; + +pub fn opendir(path: [*:0]const u8) ?*DIR { + var pattern: [4096]u8 = undefined; + const p = std.fmt.bufPrint(&pattern, "{s}\\*", .{std.mem.span(path)}) catch { + _ = fail(.NAMETOOLONG); + return null; + }; + var buf: Wide = undefined; + const w = wide(&buf, p) orelse return null; + const d = std.heap.c_allocator.create(DIR) catch { + _ = fail(.NOMEM); + return null; + }; + const h = FindFirstFileW(w, &d.data); + if (h == invalid_handle) { + const e = lastError(); + std.heap.c_allocator.destroy(d); + _ = fail(e); + return null; + } + d.* = .{ .find = @ptrFromInt(@as(usize, @bitCast(h))), .data = d.data, .first = true, .fd = null, .ent = undefined }; + return d; +} + +/// The stream owns `fd` from here: closedir releases it. +pub fn fdopendir(fd: fd_t) ?*DIR { + const path = dirPath(fd) orelse { + _ = fail(.BADF); + return null; + }; + const d = opendir(path.ptr) orelse return null; + d.fd = fd; + return d; +} + +pub fn readdir(d: *DIR) ?*dirent { + if (d.first) { + d.first = false; + } else if (FindNextFileW(d.find.?, &d.data) == 0) { + return null; + } + const name = std.mem.sliceTo(&d.data.name, 0); + const n = std.unicode.wtf16LeToWtf8(d.ent.name[0 .. d.ent.name.len - 1], name); + d.ent.name[n] = 0; + d.ent.type = if (d.data.attributes & 0x400 != 0) DT.LNK else if (d.data.attributes & 0x10 != 0) DT.DIR else DT.REG; + return &d.ent; +} + +pub fn closedir(d: *DIR) c_int { + if (d.find) |h| _ = FindClose(h); + if (d.fd) |fd| _ = close(fd); + std.heap.c_allocator.destroy(d); + return 0; +} + +// ---- process, environment, time ---- + +pub fn getpid() pid_t { + return _getpid(); +} + +pub fn getuid() u32 { + return 0; +} + +pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { + if (overwrite == 0 and getenv(name) != null) return 0; + if (_putenv_s(name, value) != 0) return fail(.INVAL); + return 0; +} + +pub fn unsetenv(name: [*:0]const u8) c_int { + if (_putenv_s(name, "") != 0) return fail(.INVAL); + return 0; +} + +threadlocal var sleep_timer: ?HANDLE = null; + +/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. +pub fn nanosleep(req: *const timespec, rem: ?*timespec) c_int { + if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; + const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; + if (ns <= 0) return 0; + if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS + if (sleep_timer) |t| { + const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units + if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { + _ = WaitForSingleObject(t, 0xFFFF_FFFF); + return 0; + } + } + Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); + return 0; +} + +pub fn clock_gettime(clock: c_int, out: *timespec) c_int { + if (clock == CLOCK.MONOTONIC) { + var count: i64 = 0; + var freq: i64 = 0; + _ = QueryPerformanceCounter(&count); + _ = QueryPerformanceFrequency(&freq); + out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; + return 0; + } + var ft: FILETIME = undefined; + GetSystemTimePreciseAsFileTime(&ft); + const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); + const unix = since_1601 - 116_444_736_000_000_000; + out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; + return 0; +} + +pub fn gettimeofday(tv: *timeval, tz: ?*anyopaque) c_int { + _ = tz; + var ts: timespec = undefined; + _ = clock_gettime(CLOCK.REALTIME, &ts); + tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; + return 0; +} + +/// Bytes free to this user under `dir` (statvfs's bavail * frsize). +pub fn diskAvailable(dir: [*:0]const u8) ?u64 { + var buf: Wide = undefined; + const w = wide(&buf, std.mem.span(dir)) orelse return null; + var avail: u64 = 0; + if (GetDiskFreeSpaceExW(w, &avail, null, null) == 0) return null; + return avail; +} + +/// NTFS keeps access in ACLs; mode bits are not stored, so a mode change succeeds and changes nothing. +pub fn fchmodat(dir: fd_t, path: [*:0]const u8, mode: mode_t, flags: u32) c_int { + _ = .{ dir, path, mode, flags }; + return 0; +} diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig new file mode 100644 index 0000000000..024e8bcc15 --- /dev/null +++ b/zig/src/platform/win_posix.zig @@ -0,0 +1,279 @@ +//! The std.posix surface the HTTP server calls through `sys.posix`, on Windows: Winsock sockets (SOCKET as an isize, +//! INVALID_SOCKET as -1), WSAPoll, console control events as SIGINT and SIGTERM, and no descriptor limits to raise. +//! Calls in `system` return -1 with errno set from WSAGetLastError, as libc's do. + +const std = @import("std"); +const c = @import("win_c.zig"); + +pub const E = c.E; +pub const errno = c.errno; +pub const fd_t = c.fd_t; +pub const AT = c.AT; +pub const S = c.S; + +// ---- sockets ---- + +pub const socket_t = isize; +pub const socklen_t = c_int; +pub const sockaddr = extern struct { family: u16, data: [14]u8 }; + +pub const AF = struct { + pub const UNIX: u32 = 1; + pub const INET: u32 = 2; + pub const INET6: u32 = 23; +}; +pub const SOCK = struct { + pub const STREAM: u32 = 1; +}; +pub const SOL = struct { + pub const SOCKET: i32 = 0xffff; +}; +pub const SO = struct { + /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, which the listener + /// sets REUSEADDR to prevent; exclusive use is the POSIX meaning. + pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); +}; +pub const SHUT = struct { + pub const RD: i32 = 0; + pub const WR: i32 = 1; + pub const RDWR: i32 = 2; +}; +pub const MSG = struct { + pub const PEEK: u32 = 0x2; + /// Winsock has no per-call nonblocking flag; every caller polls first, so the receive never waits. + pub const DONTWAIT: u32 = 0; +}; +pub const POLL = struct { + pub const IN: i16 = 0x0300; // POLLRDNORM | POLLRDBAND + /// WSAPoll refuses POLLPRI; out-of-band data is not a signal the server uses. + pub const PRI: i16 = 0; + pub const OUT: i16 = 0x0010; + pub const ERR: i16 = 0x0001; + pub const HUP: i16 = 0x0002; + pub const NVAL: i16 = 0x0004; +}; +pub const pollfd = extern struct { fd: socket_t, events: i16, revents: i16 }; + +const WSAData = extern struct { version: u16, high: u16, rest: [400]u8 }; +const ws = struct { + extern "ws2_32" fn WSAStartup(version: u16, data: *WSAData) callconv(.winapi) c_int; + extern "ws2_32" fn WSAGetLastError() callconv(.winapi) c_int; + extern "ws2_32" fn WSAPoll(fds: [*]pollfd, n: u32, timeout: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn socket(af: c_int, kind: c_int, proto: c_int) callconv(.winapi) isize; + extern "ws2_32" fn bind(s: isize, addr: *const anyopaque, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn listen(s: isize, backlog: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn connect(s: isize, addr: *const anyopaque, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn accept(s: isize, addr: ?*anyopaque, len: ?*c_int) callconv(.winapi) isize; + extern "ws2_32" fn getsockname(s: isize, addr: *anyopaque, len: *c_int) callconv(.winapi) c_int; + extern "ws2_32" fn setsockopt(s: isize, level: c_int, opt: c_int, val: [*]const u8, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn shutdown(s: isize, how: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn closesocket(s: isize) callconv(.winapi) c_int; + extern "ws2_32" fn send(s: isize, buf: [*]const u8, len: c_int, flags: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn recv(s: isize, buf: [*]u8, len: c_int, flags: c_int) callconv(.winapi) c_int; +}; +extern "c" fn _errno() *c_int; + +var started = std.atomic.Value(u8).init(0); + +/// WSAStartup once per process, before the first socket call. +fn start() void { + if (started.load(.acquire) == 2) return; + if (started.cmpxchgStrong(0, 1, .acq_rel, .acquire) == null) { + var data: WSAData = undefined; + _ = ws.WSAStartup(0x0202, &data); + started.store(2, .release); + return; + } + while (started.load(.acquire) != 2) std.atomic.spinLoopHint(); +} + +/// -1 with errno from the thread's Winsock error. +fn failed() c_int { + const e: E = switch (ws.WSAGetLastError()) { + 10004 => .INTR, + 10013 => .ACCES, + 10014 => .FAULT, + 10022 => .INVAL, + 10024 => .MFILE, + 10035 => .AGAIN, + 10038 => .NOTSOCK, + 10048 => .ADDRINUSE, + 10049 => .ADDRNOTAVAIL, + 10050 => .NETDOWN, + 10053 => .CONNABORTED, + 10054 => .CONNRESET, + 10055 => .NOBUFS, + 10057 => .NOTCONN, + 10058 => .PIPE, + 10060 => .TIMEDOUT, + 10061 => .CONNREFUSED, + else => .IO, + }; + _errno().* = @intFromEnum(e); + return -1; +} + +fn len32(n: usize) c_int { + return @intCast(@min(n, std.math.maxInt(c_int))); +} + +pub const system = struct { + pub fn socket(family: u32, kind: u32, proto: u32) isize { + start(); + const fd = ws.socket(@intCast(family), @intCast(kind), @intCast(proto)); + return if (fd == -1) failed() else fd; + } + pub fn bind(fd: socket_t, addr: *const anyopaque, len: socklen_t) c_int { + return if (ws.bind(fd, addr, len) != 0) failed() else 0; + } + pub fn listen(fd: socket_t, backlog: u31) c_int { + return if (ws.listen(fd, backlog) != 0) failed() else 0; + } + pub fn accept(fd: socket_t, addr: ?*anyopaque, len: ?*socklen_t) isize { + const s = ws.accept(fd, addr, len); + return if (s == -1) failed() else s; + } + pub fn getsockname(fd: socket_t, addr: *anyopaque, len: *socklen_t) c_int { + return if (ws.getsockname(fd, addr, len) != 0) failed() else 0; + } + pub fn shutdown(fd: socket_t, how: i32) c_int { + return if (ws.shutdown(fd, how) != 0) failed() else 0; + } + /// Closes a socket; files close through `sys.c.close`. + pub fn close(fd: socket_t) c_int { + return if (ws.closesocket(fd) != 0) failed() else 0; + } + /// Sends on a socket. + pub fn write(fd: socket_t, buf: [*]const u8, len: usize) isize { + const n = ws.send(fd, buf, len32(len), 0); + return if (n < 0) failed() else n; + } + pub fn recvfrom(fd: socket_t, buf: [*]u8, len: usize, flags: u32, addr: ?*anyopaque, alen: ?*socklen_t) isize { + _ = addr; + _ = alen; + const n = ws.recv(fd, buf, len32(len), @intCast(flags)); + return if (n < 0) failed() else n; + } +}; + +pub fn setsockopt(fd: socket_t, level: i32, opt: u32, value: []const u8) error{SetSockOptFailed}!void { + if (ws.setsockopt(fd, level, @bitCast(opt), value.ptr, @intCast(value.len)) != 0) return error.SetSockOptFailed; +} + +/// WSAPoll, with POSIX's rule that an entry whose fd is negative is skipped (WSAPoll would fail the call). +pub fn poll(fds: []pollfd, timeout_ms: i32) error{PollFailed}!usize { + start(); + var live: [64]pollfd = undefined; + var where: [64]usize = undefined; + var n: usize = 0; + for (fds, 0..) |*f, i| { + f.revents = 0; + if (f.fd < 0) continue; + if (n == live.len) return error.PollFailed; + live[n] = f.*; + where[n] = i; + n += 1; + } + if (n == 0) { + if (timeout_ms > 0) Sleep(@intCast(timeout_ms)); + return 0; + } + const rc = ws.WSAPoll(&live, @intCast(n), timeout_ms); + if (rc < 0) return error.PollFailed; + for (live[0..n], where[0..n]) |f, i| fds[i].revents = f.revents; + return @intCast(rc); +} + +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; + +/// Receives from a socket; 0 is the peer's orderly close. +pub fn read(fd: socket_t, buf: []u8) error{ConnectionResetByPeer}!usize { + const n = ws.recv(fd, buf.ptr, len32(buf.len), 0); + if (n < 0) return error.ConnectionResetByPeer; + return @intCast(n); +} + +// ---- signals: console control events ---- + +pub const SIG = enum(c_int) { + HUP = 1, + INT = 2, + PIPE = 13, + TERM = 15, + _, + + pub const IGN: ?Handler = @ptrFromInt(1); + pub const DFL: ?Handler = null; +}; +pub const Handler = *const fn (SIG) callconv(.c) void; +pub const sigset_t = u32; +pub const Sigaction = struct { + handler: extern union { handler: ?Handler, sigaction: ?*const anyopaque }, + mask: sigset_t, + flags: u32, +}; + +pub fn sigemptyset() sigset_t { + return 0; +} + +var on_int: ?Handler = null; +var on_term: ?Handler = null; +var console_hooked = std.atomic.Value(bool).init(false); + +extern "kernel32" fn SetConsoleCtrlHandler(routine: ?*const fn (u32) callconv(.winapi) c_int, add: c_int) callconv(.winapi) c_int; + +/// Ctrl+C and Ctrl+Break are SIGINT; closing the console, logoff and shutdown are SIGTERM. +fn consoleEvent(kind: u32) callconv(.winapi) c_int { + const h = switch (kind) { + 0, 1 => on_int, + else => on_term, + } orelse return 0; + if (h == SIG.IGN.?) return 1; + h(if (kind <= 1) .INT else .TERM); + return 1; +} + +/// SIGINT and SIGTERM through console events; Windows has no SIGHUP or SIGPIPE (a closed socket fails the send). +pub fn sigaction(sig: SIG, act: ?*const Sigaction, old: ?*Sigaction) void { + if (old) |o| o.* = .{ .handler = .{ .handler = SIG.DFL }, .mask = 0, .flags = 0 }; + const a = act orelse return; + switch (sig) { + .INT => on_int = a.handler.handler, + .TERM => on_term = a.handler.handler, + else => return, + } + if (!console_hooked.swap(true, .acq_rel)) _ = SetConsoleCtrlHandler(consoleEvent, 1); +} + +// ---- descriptor limits: Windows has none to raise ---- + +pub const rlim_t = u64; +pub const rlimit = struct { cur: rlim_t, max: rlim_t }; +pub const rlimit_resource = enum { NOFILE }; + +pub fn getrlimit(_: rlimit_resource) error{Unsupported}!rlimit { + return error.Unsupported; +} + +pub fn setrlimit(_: rlimit_resource, _: rlimit) error{Unsupported}!void { + return error.Unsupported; +} + +/// A connected stream pair over loopback TCP (Windows has no socketpair); for tests that drive a connection. +pub fn socketpair(pair: *[2]socket_t) error{SocketPair}!void { + start(); + const l = ws.socket(@intCast(AF.INET), @intCast(SOCK.STREAM), 0); + if (l == -1) return error.SocketPair; + defer _ = ws.closesocket(l); + var addr: extern struct { family: u16 = 2, port: u16 = 0, addr: [4]u8 = .{ 127, 0, 0, 1 }, zero: [8]u8 = @splat(0) } = .{}; + var len: c_int = @sizeOf(@TypeOf(addr)); + if (ws.bind(l, &addr, len) != 0 or ws.listen(l, 1) != 0 or ws.getsockname(l, &addr, &len) != 0) return error.SocketPair; + const a = ws.socket(@intCast(AF.INET), @intCast(SOCK.STREAM), 0); + if (a == -1) return error.SocketPair; + errdefer _ = ws.closesocket(a); + if (ws.connect(a, &addr, len) != 0) return error.SocketPair; + const b = ws.accept(l, null, null); + if (b == -1) return error.SocketPair; + pair.* = .{ a, b }; +} diff --git a/zig/src/server/auth.zig b/zig/src/server/auth.zig index 5457651dd9..b2274e2dc0 100644 --- a/zig/src/server/auth.zig +++ b/zig/src/server/auth.zig @@ -1,6 +1,7 @@ //! API keys held as SHA-256 digests with labels; a restricted key file reread on change or SIGHUP. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const builtin = @import("builtin"); const log = @import("log.zig"); const Allocator = std.mem.Allocator; @@ -84,10 +85,10 @@ pub const Store = struct { const path = s.path.?; const pathz = try s.gpa.dupeSentinel(u8, path, 0); defer s.gpa.free(pathz); - const fd = posix.openatZ(posix.AT.FDCWD, pathz, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true, .NONBLOCK = true }, 0) catch return error.Unsafe; - defer _ = posix.system.close(fd); + const fd = openKeys(pathz) catch return error.Unsafe; + defer _ = platform.c.close(fd); const st = meta(fd, null) orelse return error.Unsafe; - if ((st.mode & posix.S.IFMT) != posix.S.IFREG or st.mode & 0o044 != 0) { + if ((st.mode & platform.c.S.IFMT) != platform.c.S.IFREG or st.mode & 0o044 != 0) { problem.* = "API key file must be a regular file unreadable by other users; use chmod 600"; return error.Invalid; } @@ -101,7 +102,7 @@ pub const Store = struct { const text = try a.alloc(u8, (1 << 20) + 1); var got: usize = 0; while (got < text.len) { - const n = posix.read(fd, text[got..]) catch return error.Unsafe; + const n = readKeys(fd, text[got..]) catch return error.Unsafe; if (n == 0) break; got += n; } @@ -290,7 +291,7 @@ test "auth cleanup on every initialization allocation failure" { try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "keys", .data = "ops: sk-file\nsk-unlabelled\n" }); const path = try std.fmt.allocPrintSentinel(gpa, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); defer gpa.free(path); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); const Case = struct { fn run(a: Allocator, file: ?[]const u8) !void { var problem: []const u8 = ""; @@ -304,6 +305,19 @@ test "auth cleanup on every initialization allocation failure" { try std.testing.checkAllAllocationFailures(gpa, Case.run, .{@as(?[]const u8, path)}); } +/// The key file opened for reading, never through a link. +fn openKeys(pathz: [*:0]const u8) !platform.c.fd_t { + if (builtin.os.tag != .windows) return posix.openatZ(posix.AT.FDCWD, pathz, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true, .NONBLOCK = true }, 0); + const fd = platform.c.open(pathz, .{ .ACCMODE = .RDONLY }); + return if (fd < 0) error.OpenFailed else fd; +} + +fn readKeys(fd: platform.c.fd_t, buf: []u8) !usize { + if (builtin.os.tag != .windows) return posix.read(fd, buf); + const n = platform.c.read(fd, buf.ptr, buf.len); + return if (n < 0) error.ReadFailed else @intCast(n); +} + /// A file's type and permission bits, size, identity and mtime: statx on Linux (its libc has no stat), stat elsewhere. const Meta = struct { mode: u32, size: u64, dev: u64, ino: u64, mtime_ns: i128 }; @@ -313,11 +327,23 @@ fn meta(fd: posix.fd_t, path: ?[*:0]const u8) ?Meta { const linux = std.os.linux; var sx: linux.Statx = undefined; const want: linux.STATX = .{ .TYPE = true, .MODE = true, .INO = true, .SIZE = true, .MTIME = true }; - const rc = if (path) |p| std.c.statx(posix.AT.FDCWD, p, 0, want, &sx) else std.c.statx(fd, "", 0x1000, want, &sx); // AT_EMPTY_PATH + const rc = if (path) |p| platform.c.statx(posix.AT.FDCWD, p, 0, want, &sx) else platform.c.statx(fd, "", 0x1000, want, &sx); // AT_EMPTY_PATH if (rc != 0) return null; const dev = (@as(u64, sx.dev_major) << 32) | sx.dev_minor; return .{ .mode = sx.mode, .size = sx.size, .dev = dev, .ino = sx.ino, .mtime_ns = @as(i128, sx.mtime.sec) * std.time.ns_per_s + sx.mtime.nsec }; } + if (builtin.os.tag == .windows) { + // NTFS keeps access in ACLs, not mode bits: the type and size are checked, and the profile's ACL is the guard + var st: platform.c.Stat = undefined; + const f = if (path) |p| platform.c.open(p, .{ .ACCMODE = .RDONLY }) else fd; + if (f < 0) return null; + defer if (path != null) { + _ = platform.c.close(f); + }; + if (platform.c.fstat(f, &st) != 0) return null; + const t = st.mtime(); + return .{ .mode = (st.mode & platform.c.S.IFMT) | 0o600, .size = @intCast(st.size), .dev = st.dev, .ino = st.ino, .mtime_ns = @as(i128, t.nsec) + @as(i128, t.sec) * std.time.ns_per_s }; + } var st: posix.Stat = undefined; const rc = if (path) |p| posix.system.stat(p, &st) else posix.system.fstat(fd, &st); if (posix.errno(rc) != .SUCCESS) return null; @@ -326,6 +352,7 @@ fn meta(fd: posix.fd_t, path: ?[*:0]const u8) ?Meta { } test "the key file: refused while others can read it, read at 0600, reread after an edit" { + if (builtin.os.tag == .windows) return error.SkipZigTest; // refusal by mode bits; NTFS keeps access in ACLs var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); @@ -334,11 +361,11 @@ test "the key file: refused while others can read it, read at 0600, reread after defer tmp.cleanup(); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-one\n" }); const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); var problem: []const u8 = ""; try std.testing.expectError(error.KeyFile, Store.init(a, &.{}, "", path, false, &problem)); try std.testing.expect(std.mem.indexOf(u8, problem, "chmod 600") != null); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); var s = try Store.init(std.testing.allocator, &.{"sk-cli"}, "sk-env", path, false, &problem); defer s.deinit(); try std.testing.expectEqualStrings("cli-1", s.match(a, "Bearer sk-cli", null).?); @@ -354,6 +381,7 @@ test "the key file: refused while others can read it, read at 0600, reread after } test "a request's key label outlives a reload of the key file" { + if (builtin.os.tag == .windows) return error.SkipZigTest; // refusal by mode bits; NTFS keeps access in ACLs const gpa = std.testing.allocator; const io = std.testing.io; var arena: std.heap.ArenaAllocator = .init(gpa); @@ -363,7 +391,7 @@ test "a request's key label outlives a reload of the key file" { defer tmp.cleanup(); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-one\n" }); const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); var problem: []const u8 = ""; var s = try Store.init(gpa, &.{}, "", path, false, &problem); defer if (s.file_arena) |*keys| keys.deinit(); @@ -373,7 +401,7 @@ test "a request's key label outlives a reload of the key file" { try std.testing.expectEqualStrings("ci", s.match(a, null, "sk-two").?); // another request rereads the file try std.testing.expectEqualStrings("ops", label); const kept = s.match(a, null, "sk-two").?; - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); s.checked_ns = null; try std.testing.expect(s.match(a, "Bearer sk-one", null) == null); // a refused reload frees the file's keys too try std.testing.expectEqualStrings("ci", kept); diff --git a/zig/src/server/http_conn.zig b/zig/src/server/http_conn.zig index f29d9f8684..536e8649d7 100644 --- a/zig/src/server/http_conn.zig +++ b/zig/src/server/http_conn.zig @@ -1,6 +1,7 @@ //! One HTTP/1.1 connection as Python's BaseHTTPRequestHandler serves it: bounded heads, keep-alive, error pages. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const Allocator = std.mem.Allocator; const log = @import("log.zig"); diff --git a/zig/src/server/late_system_test.zig b/zig/src/server/late_system_test.zig index e661ae2fea..7f4622dcd1 100644 --- a/zig/src/server/late_system_test.zig +++ b/zig/src/server/late_system_test.zig @@ -1,5 +1,6 @@ //! A later system message is re-roled to user on Messages when the template cannot keep it, matching chat. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const server = @import("root.zig"); const json = server.json; @@ -132,10 +133,10 @@ const Capture = struct { const Reply = struct { status: u16, body: []u8 }; fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { - var pair: [2]std.c.fd_t = undefined; - if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; - defer _ = std.c.close(pair[0]); - defer _ = std.c.close(pair[1]); + var pair: [2]platform.posix.socket_t = undefined; + if (@import("builtin").os.tag == .windows) try platform.posix.socketpair(&pair) else if (platform.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer _ = platform.posix.system.close(pair[0]); + defer _ = platform.posix.system.close(pair[1]); const head = try std.fmt.allocPrint(std.testing.allocator, "POST {s} HTTP/1.1\r\nHost: t\r\nContent-Length: {d}\r\n\r\n", .{ path, body.len }); defer std.testing.allocator.free(head); try writeAll(pair[0], head); @@ -157,23 +158,23 @@ fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { return .{ .status = status, .body = payload }; } -fn writeAll(fd: std.c.fd_t, data: []const u8) !void { +fn writeAll(fd: platform.posix.socket_t, data: []const u8) !void { var sent: usize = 0; while (sent < data.len) { - const n = std.c.write(fd, data[sent..].ptr, data.len - sent); + const n = platform.posix.system.write(fd, data[sent..].ptr, data.len - sent); if (n <= 0) return error.Closed; sent += @intCast(n); } } -fn readReady(a: std.mem.Allocator, fd: std.c.fd_t) ![]u8 { +fn readReady(a: std.mem.Allocator, fd: platform.posix.socket_t) ![]u8 { var out: std.ArrayList(u8) = .empty; errdefer out.deinit(a); var tmp: [4096]u8 = undefined; while (true) { - var pollfd = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.IN, .revents = 0 }}; - if (try std.posix.poll(&pollfd, 200) == 0) break; - const n = std.posix.read(fd, &tmp) catch |e| return e; + var pollfd = [_]platform.posix.pollfd{.{ .fd = fd, .events = platform.posix.POLL.IN, .revents = 0 }}; + if (try platform.posix.poll(&pollfd, 200) == 0) break; + const n = platform.posix.read(fd, &tmp) catch |e| return e; if (n == 0) break; try out.appendSlice(a, tmp[0..n]); } diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index 8eb7a27df6..a6492e737d 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -1,6 +1,7 @@ //! The listening socket: SO_REUSEADDR only (never REUSEPORT, so a second server on the port fails), backlog 128. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const net = std.Io.net; const log = @import("log.zig"); const Threaded = std.Io.Threaded; diff --git a/zig/src/server/live.zig b/zig/src/server/live.zig index cb0cbaf640..3d6f8b482e 100644 --- a/zig/src/server/live.zig +++ b/zig/src/server/live.zig @@ -1,5 +1,6 @@ //! The live status line (connections, decode and prefill tok/s), redrawn in a terminal, and /health's copy of it. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const log = @import("log.zig"); @@ -94,7 +95,7 @@ fn columns() usize { if (os == .linux) { if (std.os.linux.ioctl(1, std.os.linux.T.IOCGWINSZ, @intFromPtr(&ws)) == 0 and ws.col > 0) return ws.col; } else if (os.isDarwin()) { - if (std.c.ioctl(1, @bitCast(@as(u32, 0x40087468)), &ws) == 0 and ws.col > 0) return ws.col; + if (platform.c.ioctl(1, @bitCast(@as(u32, 0x40087468)), &ws) == 0 and ws.col > 0) return ws.col; } return 100; } diff --git a/zig/src/server/log.zig b/zig/src/server/log.zig index 40430985a2..629bf90bee 100644 --- a/zig/src/server/log.zig +++ b/zig/src/server/log.zig @@ -1,5 +1,6 @@ //! Server output: ``[tensorfold]`` lines on stdout, and the live status line they clear first. const std = @import("std"); +const platform = @import("sys"); const Conn = @import("http_conn.zig").Conn; const afterChars = @import("reply_text.zig").afterChars; @@ -23,8 +24,8 @@ const clear = "\r\x1b[2K"; fn writeOut(bytes: []const u8) void { var off: usize = 0; while (off < bytes.len) { - const rc = std.posix.system.write(1, bytes[off..].ptr, bytes.len - off); - if (std.posix.errno(rc) != .SUCCESS or rc <= 0) return; + const rc = platform.c.write(1, bytes[off..].ptr, bytes.len - off); + if (platform.c.errno(rc) != .SUCCESS or rc <= 0) return; off += @intCast(rc); } } diff --git a/zig/src/server/request_log.zig b/zig/src/server/request_log.zig index ea4c8b7997..02c26a0c0d 100644 --- a/zig/src/server/request_log.zig +++ b/zig/src/server/request_log.zig @@ -1,5 +1,6 @@ //! TENSORFOLD_REQUEST_LOG: each chat and completion body appended as one JSON line, images redacted, for replays. const std = @import("std"); +const platform = @import("sys"); const json = @import("json.zig"); const Value = json.Value; const Allocator = std.mem.Allocator; @@ -33,8 +34,9 @@ pub fn append(a: Allocator, path: []const u8, body: Value) void { if (body.strField("priority")) |p| if (std.mem.eql(u8, p, "background")) return; const line = json.stringify(a, redact(a, body) catch return, .{}) catch return; const pathz = a.dupeSentinel(u8, path, 0) catch return; - const fd = std.posix.openatZ(std.posix.AT.FDCWD, pathz, .{ .ACCMODE = .WRONLY, .APPEND = true, .CREAT = true }, 0o644) catch return; - defer _ = std.posix.system.close(fd); + const fd = platform.c.open(pathz, .{ .ACCMODE = .WRONLY, .APPEND = true, .CREAT = true }, @as(platform.c.mode_t, 0o644)); + if (fd < 0) return; + defer _ = platform.c.close(fd); const framed = std.mem.concat(a, u8, &.{ line, "\n" }) catch return; - _ = std.posix.system.write(fd, framed.ptr, framed.len); + _ = platform.c.write(fd, framed.ptr, framed.len); } diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..46a8415f3b 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -1,5 +1,6 @@ //! A server's life: keys, the listening socket, the startup lines, the live line, and a clean stop on SIGTERM or SIGINT. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const auth = @import("auth.zig"); @@ -10,7 +11,7 @@ const listener_mod = @import("listener.zig"); const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); const Allocator = std.mem.Allocator; -const posix = std.posix; +const posix = platform.posix; var stop_requested: std.atomic.Value(bool) = .init(false); @@ -33,7 +34,7 @@ pub fn stopping() bool { } /// The hard limit macOS reads as unlimited. -const rlim_infinity: u64 = std.c.RLIM.INFINITY; +const rlim_infinity: u64 = platform.c.RLIM.INFINITY; /// The targets to try, in order: the hard limit when finite, else the fallbacks 65536 then 10240, first that works; never an unlimited soft limit, since a loop over every descriptor would do absurd work. fn fillTargets(max: u64, out: *[2]u64) []const u64 { @@ -211,6 +212,7 @@ test "the targets are the hard limit, or the fallbacks under an unlimited one" { } test "the raise lifts a low soft limit, and a second raise does nothing" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // no descriptor limits to raise const before = try posix.getrlimit(.NOFILE); defer posix.setrlimit(.NOFILE, before) catch {}; if (@as(u64, @intCast(before.max)) <= 256) return error.SkipZigTest; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..4d76d98105 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -1,5 +1,6 @@ //! The server: one engine and tokenizer behind the HTTP routes, one thread a connection. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const errors = @import("errors.zig"); @@ -280,7 +281,7 @@ pub const Server = struct { pub fn serve(srv: *Server, listener: listener_mod.Listener, stop: *const std.atomic.Value(bool)) void { while (listener.accept(stop)) |c| { const t = std.Thread.spawn(.{ .stack_size = 8 << 20 }, connection, .{ srv, c }) catch { - _ = std.posix.system.close(c.fd); + _ = platform.posix.system.close(c.fd); continue; }; t.detach(); @@ -290,7 +291,7 @@ pub const Server = struct { fn connection(srv: *Server, accepted: listener_mod.Listener.Accepted) void { _ = srv.open_connections.fetchAdd(1, .acq_rel); defer _ = srv.open_connections.fetchSub(1, .acq_rel); - defer _ = std.posix.system.close(accepted.fd); + defer _ = platform.posix.system.close(accepted.fd); var peer_buf: [64]u8 = undefined; const peer = listener_mod.peerText(&peer_buf, accepted.peer); var conn = http_conn.Conn.init(srv.gpa, accepted.fd, peer) catch return; diff --git a/zig/tests/cuda/pack_compare.zig b/zig/tests/cuda/pack_compare.zig index b1cc379a36..10f9041c65 100644 --- a/zig/tests/cuda/pack_compare.zig +++ b/zig/tests/cuda/pack_compare.zig @@ -1,5 +1,6 @@ //! tf-train-pack's checks on either backend: packed passes against one a pass, and passes after kept prefixes. const std = @import("std"); +const platform = @import("sys"); /// What one sequence's pass left: its loss, whether it came back, its pushes and its keys at the pocket's layers. const Seen = struct { loss: f32, recalled: bool, push: []f32, keys: []u16 }; @@ -169,7 +170,7 @@ pub fn For(comptime train: type) type { /// TF_PACK_DUMP: the packed target pass's losses and pushes written there, for another build's to compare. fn dump(gpa: std.mem.Allocator, io: std.Io, seen: []const Seen) !void { - const path = std.c.getenv("TF_PACK_DUMP") orelse return; + const path = platform.c.getenv("TF_PACK_DUMP") orelse return; var bytes: std.ArrayList(u8) = .empty; for (seen) |x| { try bytes.appendSlice(gpa, std.mem.asBytes(&x.loss)); diff --git a/zig/tests/cuda/train_pack.zig b/zig/tests/cuda/train_pack.zig index 2ab91949b9..d721c19887 100644 --- a/zig/tests/cuda/train_pack.zig +++ b/zig/tests/cuda/train_pack.zig @@ -1,5 +1,6 @@ //! Training passes holding several sequences against one a pass, and from kept prefixes against whole, on CUDA. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const nemotron = @import("nemotron"); const check_mod = @import("check.zig"); @@ -17,7 +18,7 @@ const open = 2; pub fn run(gpu: check_mod.Gpu, model: []const u8, path: []const u8, count: usize) !void { const gpa = gpu.gpa; const io = gpu.io; - const kernels = std.c.getenv("TENSORFOLD_CUDA_KERNELS"); + const kernels = platform.c.getenv("TENSORFOLD_CUDA_KERNELS"); const e = try nemotron.Engine.init(gpa, io, gpu.ctx, model, if (kernels) |k| std.mem.span(k) else null, .{ .context = 8192, .mtp = false, .graphs = false, .sampling = null, .segments = 1 }); defer e.deinit(); var living: ?nemotron.sites.Sites = null; diff --git a/zig/tests/fabric_library.zig b/zig/tests/fabric_library.zig index 04141e546a..2e09f6aefb 100644 --- a/zig/tests/fabric_library.zig +++ b/zig/tests/fabric_library.zig @@ -1,5 +1,6 @@ //! The fabric bound to MCDMA's own libmcdma-rpc (built from MCDMA's rpc/ by tools/zig/test_fabric.sh, MCDMA_RPC_LIBRARY). const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const fabric = @import("fabric"); const layout = fabric.layout; @@ -9,7 +10,7 @@ const mailbox = fabric.mailbox; const gpa = std.testing.allocator; fn open() !words.Library { - if (std.c.getenv("MCDMA_RPC_LIBRARY") == null) return error.SkipZigTest; + if (platform.c.getenv("MCDMA_RPC_LIBRARY") == null) return error.SkipZigTest; return words.Library.open(null); } diff --git a/zig/tests/flashnext_bench.zig b/zig/tests/flashnext_bench.zig index 7a0c73d528..d15901c68a 100644 --- a/zig/tests/flashnext_bench.zig +++ b/zig/tests/flashnext_bench.zig @@ -1,5 +1,6 @@ //! FZ_DBENCH: Flash Next dense and DeltaNet kernels at 1-16 rows against the recorded ones, bit for bit, then timed. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const fz = tf.flashnext_replay; @@ -238,7 +239,7 @@ pub fn run(r: *Run, m: *Model, prompt: []const u32, gpa: std.mem.Allocator, aren }; const pfs = [_]usize{ 1, 2 }; - const bits = std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_DBENCH").?), "bits"); + const bits = std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_DBENCH").?), "bits"); // bits: core lane against the recorded kernel, every layer, one Mac and speed-up layouts, every read-ahead depth const keep = try gpa.alloc(u8, fz.MAXR * fz.VOCAB * 2); defer gpa.free(keep); diff --git a/zig/tests/flashnext_reuse.zig b/zig/tests/flashnext_reuse.zig index eeeda53247..c9c66a3ca1 100644 --- a/zig/tests/flashnext_reuse.zig +++ b/zig/tests/flashnext_reuse.zig @@ -1,5 +1,6 @@ //! Flash Next prompt reuse against fresh passes: agent-like turns, every depth, copy drafts off and on, state bytes. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const api = @import("engine_api"); @@ -121,9 +122,9 @@ pub fn main(init: std.process.Init) !void { } const pool = mtl.objc.Pool.push(); defer pool.pop(); - const n_out: usize = if (std.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 64; + const n_out: usize = if (platform.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 64; var depths: std.ArrayList(?usize) = .empty; - var dit = std.mem.tokenizeScalar(u8, if (std.c.getenv("FZ_DEPTHS")) |v| std.mem.span(v) else "1,2,3,5,0", ','); + var dit = std.mem.tokenizeScalar(u8, if (platform.c.getenv("FZ_DEPTHS")) |v| std.mem.span(v) else "1,2,3,5,0", ','); while (dit.next()) |d| { const k = try std.fmt.parseInt(usize, d, 10); try depths.append(a, if (k == 0) null else k); diff --git a/zig/tests/flashnext_run.zig b/zig/tests/flashnext_run.zig index 87efe89c96..5c1e1c3234 100644 --- a/zig/tests/flashnext_run.zig +++ b/zig/tests/flashnext_run.zig @@ -2,6 +2,7 @@ //! recorded variant, every weight the pack's or the checkpoint's. Checks one-row greedy tokens and every row of the //! Python engine's drafted windows (with rollback), then times each window size. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const tf = @import("tensorfold"); @@ -46,13 +47,13 @@ const jsonInt = fz.jsonInt; fn armName(seg_ab: bool, arm: usize) []const u8 { if (seg_ab) return if (arm == 1) "segments on " else "segments off"; - if (std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "fz")) return if (arm == 1) "fz kernels" else "recorded "; + if (platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "fz")) return if (arm == 1) "fz kernels" else "recorded "; return if (arm == 1) "copy on " else "copy off"; } /// FZ_PROMPT_N: only the prompt's first n tokens (all when unset). fn promptCut() !usize { - const v = std.c.getenv("FZ_PROMPT_N") orelse return std.math.maxInt(usize); + const v = platform.c.getenv("FZ_PROMPT_N") orelse return std.math.maxInt(usize); return std.fmt.parseInt(usize, std.mem.span(v), 10); } @@ -124,7 +125,7 @@ pub fn main(init: std.process.Init) !void { defer arena_state.deinit(); const arena = arena_state.allocator(); const device = try mtl.Device.init(); - if (std.c.getenv("FZ_COMPILE_CHECK") != null) { // every prompt-kernel source and the 6-bit file on this Metal compiler + if (platform.c.getenv("FZ_COMPILE_CHECK") != null) { // every prompt-kernel source and the 6-bit file on this Metal compiler var failed: usize = 0; var files: std.ArrayList(ks.File) = .empty; try files.appendSlice(gpa, &ks.prefill); @@ -149,8 +150,8 @@ pub fn main(init: std.process.Init) !void { } std.process.exit(if (failed == 0) 0 else 1); } - if (std.c.getenv("FZ_OPORDER") != null) return @import("flashnext_bench.zig").opOrder(device); // no model loaded - if (std.c.getenv("FZ_MMA_PEAK") != null) { // the tensor units' rate on register fragments: 16x32x16 ops a second + if (platform.c.getenv("FZ_OPORDER") != null) return @import("flashnext_bench.zig").opOrder(device); // no model loaded + if (platform.c.getenv("FZ_MMA_PEAK") != null) { // the tensor units' rate on register fragments: 16x32x16 ops a second const alib = try mtl.Library.fromSource(device, try fz.frags.source(device, arena, ks.flashnext_attn), mtl.CompileOptions.mlx()); const pipe = try mtl.Pipeline.init(device, alib, "tf_mma_peak", false); const queue = try device.queue(); @@ -174,9 +175,9 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_ENGINE") != null) { // the served engine: its own load, warm-up and reply, on FZ_REF's prompt + if (platform.c.getenv("FZ_ENGINE") != null) { // the served engine: its own load, warm-up and reply, on FZ_REF's prompt const fx = @import("tensorfold").flashnext_engine; - const rp = std.mem.span(std.c.getenv("FZ_REF") orelse return error.NoRef); + const rp = std.mem.span(platform.c.getenv("FZ_REF") orelse return error.NoRef); const rf = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}", .{rp}, 0)); const rj = try std.json.parseFromSliceLeaky(std.json.Value, arena, rf.bytes[0..rf.size], .{}); const all_items = rj.object.get("prompt").?.array.items; @@ -184,9 +185,9 @@ pub fn main(init: std.process.Init) !void { const toks = try arena.alloc(u32, items.len); for (items, 0..) |x, i| toks[i] = @intCast(x.integer); const e = try fx.Engine.load(gpa, init.io, args[1], args[2]); - if (std.c.getenv("FZ_COPY_MIN")) |v| e.copy_min = try std.fmt.parseInt(u32, std.mem.span(v), 10); - if (std.c.getenv("FZ_COPY_LONG")) |v| e.copy_long = try std.fmt.parseInt(u32, std.mem.span(v), 10); - if (std.c.getenv("FZ_ENGINE_WARM") != null) try e.warm(); + if (platform.c.getenv("FZ_COPY_MIN")) |v| e.copy_min = try std.fmt.parseInt(u32, std.mem.span(v), 10); + if (platform.c.getenv("FZ_COPY_LONG")) |v| e.copy_long = try std.fmt.parseInt(u32, std.mem.span(v), 10); + if (platform.c.getenv("FZ_ENGINE_WARM") != null) try e.warm(); const Show = struct { got: std.ArrayList(u32) = .empty, a: std.mem.Allocator, @@ -200,11 +201,11 @@ pub fn main(init: std.process.Init) !void { return false; } }; - const n_out: usize = if (std.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 256; + const n_out: usize = if (platform.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 256; var replies: [2][]u32 = undefined; - const seg_ab = std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "seg"); + const seg_ab = platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "seg"); var depths: std.ArrayList(?usize) = .empty; // FZ_DEPTHS=3,5,7: each depth's arms in turn (0: the depth rule, p: plain) - if (std.c.getenv("FZ_DEPTHS")) |v| { + if (platform.c.getenv("FZ_DEPTHS")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |d| { if (std.mem.eql(u8, d, "p")) { @@ -214,9 +215,9 @@ pub fn main(init: std.process.Init) !void { const n = try std.fmt.parseInt(usize, d, 10); try depths.append(arena, if (n == 0) null else n); } - } else try depths.append(arena, if (std.c.getenv("FZ_DEPTH")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else null); + } else try depths.append(arena, if (platform.c.getenv("FZ_DEPTH")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else null); // FZ_AB=fz (with FZ_LANE=1 FZ_GDN=2): recorded dense and DeltaNet kernels, then core lane, fz_gdn, kept states - const fz_ab = std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "fz"); + const fz_ab = platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "fz"); const fz_on = .{ e.r.lane_new, e.r.gdn_pipe, e.r.gdn_kept }; for (depths.items) |depth| for (0..2) |arm| { // copy drafts off, then on (FZ_AB=seg: staggered prompt segments off, then on) const plain = depth != null and depth.? == 0; // one token a round: one arm (copies need drafts) @@ -256,7 +257,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("{s}: reply hash {x:0>16}\n", .{ name, std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(sh.got.items)) }); if (arm == 1) for (res.copy_by_len, 0..) |cl, n| if (cl[0] > 0) std.debug.print(" match {d}: {d} rounds, {d:.2} landed\n", .{ n, cl[0], @as(f64, @floatFromInt(cl[1])) / @as(f64, @floatFromInt(cl[0])) }); }; - if (std.c.getenv("FZ_NO_REF") != null) return; // the reply hashes are the check (two Macs against one) + if (platform.c.getenv("FZ_NO_REF") != null) return; // the reply hashes are the check (two Macs against one) // the reference: the same prompt chunks, then one row a step e.hostMode(); const mm = e.m; @@ -287,28 +288,28 @@ pub fn main(init: std.process.Init) !void { // the replay below reads every expert, embedding and table at 6/32: other widths are the engine's (Engine.load) if (!try fz.engine_widths.uniformDir(arena, init.io, args[1])) return error.HarnessNeeds6Bit32; var r = Run{ .arena = arena, .device = device, .queue = try device.queue() }; - r.fused_xsum = std.c.getenv("FZ_FUSED_XSUM") != null; - r.serial = std.c.getenv("FZ_SERIAL") != null; - r.xnew = std.c.getenv("FZ_XNEW") != null; - r.split = std.c.getenv("FZ_SPLIT") != null; - r.gdn_step = std.c.getenv("FZ_GDN_STEP") != null; - r.hc_mma = std.c.getenv("FZ_HC_MMA") != null; - r.hc_up = std.c.getenv("FZ_HCCHECK") != null or if (std.c.getenv("FZ_HC_UP")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; + r.fused_xsum = platform.c.getenv("FZ_FUSED_XSUM") != null; + r.serial = platform.c.getenv("FZ_SERIAL") != null; + r.xnew = platform.c.getenv("FZ_XNEW") != null; + r.split = platform.c.getenv("FZ_SPLIT") != null; + r.gdn_step = platform.c.getenv("FZ_GDN_STEP") != null; + r.hc_mma = platform.c.getenv("FZ_HC_MMA") != null; + r.hc_up = platform.c.getenv("FZ_HCCHECK") != null or if (platform.c.getenv("FZ_HC_UP")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; r.event = try device.sharedEvent(); - r.dense = r.xnew and std.c.getenv("FZ_DENSE") != null; - r.dense_target = r.dense and std.c.getenv("FZ_DENSE_TARGET") != null; - r.xpack = r.xnew and std.c.getenv("FZ_XPACK") != null; - r.grouped = r.xnew and !r.xpack and std.c.getenv("FZ_GROUPED") != null; - r.xfused = r.xnew and !r.xpack and !r.grouped and std.c.getenv("FZ_XFUSED") != null; - r.xsx = r.xnew and std.c.getenv("FZ_XSX") != null; - r.copy = std.c.getenv("FZ_COPY") != null; - r.prefetch = !r.serial and std.c.getenv("FZ_PREFETCH") != null; + r.dense = r.xnew and platform.c.getenv("FZ_DENSE") != null; + r.dense_target = r.dense and platform.c.getenv("FZ_DENSE_TARGET") != null; + r.xpack = r.xnew and platform.c.getenv("FZ_XPACK") != null; + r.grouped = r.xnew and !r.xpack and platform.c.getenv("FZ_GROUPED") != null; + r.xfused = r.xnew and !r.xpack and !r.grouped and platform.c.getenv("FZ_XFUSED") != null; + r.xsx = r.xnew and platform.c.getenv("FZ_XSX") != null; + r.copy = platform.c.getenv("FZ_COPY") != null; + r.prefetch = !r.serial and platform.c.getenv("FZ_PREFETCH") != null; const t0 = mtl.clock.seconds(); try r.compile(args[2]); r.sel = try Select.init(&r, MAXR); - if (std.c.getenv("FZ_CATCH_CHECK") != null) return catchCheck(&r, arena); - r.lane_new = r.xnew and std.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (same bits) - if (r.xnew and std.c.getenv("FZ_GDN") != null) r.gdn_pipe = try fz.gdn_step.compile(&r, false); + if (platform.c.getenv("FZ_CATCH_CHECK") != null) return catchCheck(&r, arena); + r.lane_new = r.xnew and platform.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (same bits) + if (r.xnew and platform.c.getenv("FZ_GDN") != null) r.gdn_pipe = try fz.gdn_step.compile(&r, false); const t1 = mtl.clock.seconds(); const index_file = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{args[1]}, 0)); @@ -320,8 +321,8 @@ pub fn main(init: std.process.Init) !void { while (fit.next()) |name| try r.indexFile(try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ args[1], name.* }, 0)); try r.indexFile(try std.fmt.allocPrintSentinel(arena, "{s}/pack.safetensors", .{args[2]}, 0)); - const own_prompt = std.c.getenv("FZ_REF") != null; // another prompt: passes against Python's records are skipped - const ref_path = if (std.c.getenv("FZ_REF")) |p| try std.fmt.allocPrintSentinel(arena, "{s}", .{std.mem.span(p)}, 0) else try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{args[2]}, 0); + const own_prompt = platform.c.getenv("FZ_REF") != null; // another prompt: passes against Python's records are skipped + const ref_path = if (platform.c.getenv("FZ_REF")) |p| try std.fmt.allocPrintSentinel(arena, "{s}", .{std.mem.span(p)}, 0) else try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{args[2]}, 0); const ref_file = try mtl.MappedFile.open(ref_path); const ref = try std.json.parseFromSliceLeaky(std.json.Value, arena, ref_file.bytes[0..ref_file.size], .{}); const ple_ref = ref.object.get("ple").?.object; @@ -502,15 +503,15 @@ pub fn main(init: std.process.Init) !void { const t2 = mtl.clock.seconds(); std.debug.print("compiled in {d:.2} s, loaded {d:.1} GB in {d:.1} s\n", .{ t1 - t0, @as(f64, @floatFromInt(r.loaded)) / 1e9, t2 - t1 }); - if (std.c.getenv("FZ_DBENCH") != null) { // dense classes and DeltaNet timed by rows (flashnext_bench.zig) + if (platform.c.getenv("FZ_DBENCH") != null) { // dense classes and DeltaNet timed by rows (flashnext_bench.zig) var toks: [MAXR]u32 = undefined; for (0..MAXR) |i| toks[i] = @intCast(ref.object.get("prompt").?.array.items[i].integer); return @import("flashnext_bench.zig").run(&r, m, &toks, gpa, arena); } - if (std.c.getenv("FZ_PROFILE") != null) { - if (std.c.getenv("TF_FLASHNEXT_TP")) |path| r.tp = try fz.Tp2.init(arena, r.device, std.mem.span(path)); + if (platform.c.getenv("FZ_PROFILE") != null) { + if (platform.c.getenv("TF_FLASHNEXT_TP")) |path| r.tp = try fz.Tp2.init(arena, r.device, std.mem.span(path)); var hcskips: std.ArrayList(u32) = .empty; // FZ_HCSKIP=1,2,4: the profile again under each hc knock-out - if (std.c.getenv("FZ_HCSKIP")) |v| { + if (platform.c.getenv("FZ_HCSKIP")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try hcskips.append(arena, try std.fmt.parseInt(u32, x, 10)); } else try hcskips.append(arena, 0); @@ -520,7 +521,7 @@ pub fn main(init: std.process.Init) !void { var pk: [MAXR]u32 = undefined; m.reset(); for (0..3) |_| try m.window(toks[0..1], &pk); - if (std.c.getenv("FZ_PROFILE_AB") != null) { // in-model A/B, interleaved: recorded, core lane, fz_gdn, both + if (platform.c.getenv("FZ_PROFILE_AB") != null) { // in-model A/B, interleaved: recorded, core lane, fz_gdn, both const gpipe = try fz.gdn_step.compile(&r, false); const arms = [_][2]bool{ .{ false, false }, .{ true, false }, .{ false, true }, .{ true, true } }; for ([_]usize{ 1, 2, 4, 6, 8, 16 }) |rows| { @@ -537,12 +538,12 @@ pub fn main(init: std.process.Init) !void { return; } var ups: std.ArrayList(bool) = .empty; // FZ_HC_UP=0,1: the profile with the recorded up projection, then fz_hc_up - if (std.c.getenv("FZ_HC_UP")) |v| { + if (platform.c.getenv("FZ_HC_UP")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try ups.append(arena, !std.mem.eql(u8, x, "0")); } else try ups.append(arena, false); var prow: std.ArrayList(usize) = .empty; // FZ_PROFILE_ROWS=2,3: the windows' widths (default 1, 4, 8) - if (std.c.getenv("FZ_PROFILE_ROWS")) |v| { + if (platform.c.getenv("FZ_PROFILE_ROWS")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try prow.append(arena, try std.fmt.parseInt(usize, x, 10)); } else try prow.appendSlice(arena, &.{ 1, 4, 8 }); @@ -551,7 +552,7 @@ pub fn main(init: std.process.Init) !void { r.hcskip = hs; var base: f64 = 0; for (names, 0..) |name, c| { - const only = std.c.getenv("FZ_PROFILE_ONLY") != null; // every class but this one knocked out + const only = platform.c.getenv("FZ_PROFILE_ONLY") != null; // every class but this one knocked out r.skip = if (c == 0) (if (only) 0x1ff else 0) else if (only) 0x1ff & ~(@as(u32, 1) << @intCast(c - 1)) else @as(u32, 1) << @intCast(c - 1); m.gpu_seconds = 0; for (0..20) |_| try m.window(toks[0..rows], &pk); @@ -573,7 +574,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_DUAL") != null) { // two lane groups on two queues: one window, two in a row, two at once + if (platform.c.getenv("FZ_DUAL") != null) { // two lane groups on two queues: one window, two in a row, two at once const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -648,7 +649,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_HCCHECK") != null) { // fz_hc_up against the recorded up projection: every row's logits and streams, bit for bit + if (platform.c.getenv("FZ_HCCHECK") != null) { // fz_hc_up against the recorded up projection: every row's logits and streams, bit for bit const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -687,7 +688,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("hc check: {s}\n", .{if (bad == 0) "bit-identical" else "DIFFERENT"}); return; } - if (std.c.getenv("FZ_GCHECK") != null) { // grouped or fused experts against fz_xgu/fz_xdown: every row's logits, bit for bit + if (platform.c.getenv("FZ_GCHECK") != null) { // grouped or fused experts against fz_xgu/fz_xdown: every row's logits, bit for bit const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -714,12 +715,12 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_XTIME") != null) { // the experts' share of a window: chain rows against one token repeated + if (platform.c.getenv("FZ_XTIME") != null) { // the experts' share of a window: chain rows against one token repeated const want0 = ref.object.get("tokens").?.array.items; var pk: [MAXR]u32 = undefined; m.reset(); for (0..3) |_| try m.window(&.{@intCast(want0[0].integer)}, &pk); - if (std.c.getenv("FZ_GPARTS") != null) for ([_]usize{ 2, 8 }) |rows| { + if (platform.c.getenv("FZ_GPARTS") != null) for ([_]usize{ 2, 8 }) |rows| { var toks: [MAXR]u32 = undefined; for (0..rows) |i| toks[i] = @intCast(want0[i].integer); for ([_]u32{ 0, 1, 4, 6 }) |gs| { @@ -747,21 +748,21 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_PREFILL") != null) { // the prompt path against the exact 8-row windows, then one-row decode + if (platform.c.getenv("FZ_PREFILL") != null) { // the prompt path against the exact 8-row windows, then one-row decode const all_items = ref.object.get("prompt").?.array.items; const pr_items = all_items[0..@min(all_items.len, try promptCut())]; // FZ_PROMPT_N: the prompt's first n tokens const toks = try gpa.alloc(u32, pr_items.len); defer gpa.free(toks); for (pr_items, 0..) |x, i| toks[i] = @intCast(x.integer); - if (std.c.getenv("FZ_SATTN_CHECK") != null) { // prompt attention: the decode's kernels vs the tensor units + if (platform.c.getenv("FZ_SATTN_CHECK") != null) { // prompt attention: the decode's kernels vs the tensor units var pr = try Prompt.init(&r, args[2], r.xnew_header); - if (std.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); - const ab0: []const u8 = if (std.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; + if (platform.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); + const ab0: []const u8 = if (platform.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; const seg2 = std.mem.eql(u8, ab0, "seg2"); // seg2: one chunk at a time against staggered segments (core/segments.zig): FZ_SEGS=2 (default) the // engine's rule (two segments of at least FZ_SEG_MIN rows), 1 one segment, 3-4 that many of up to a chunk - const nseg: usize = if (std.c.getenv("FZ_SEGS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 2; - const seg_min: usize = if (std.c.getenv("FZ_SEG_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 1; + const nseg: usize = if (platform.c.getenv("FZ_SEGS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 2; + const seg_min: usize = if (platform.c.getenv("FZ_SEG_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 1; if (seg2 and (nseg == 0 or nseg > tf.segments.MAX)) return error.Segments; var extra: [3]Prompt = undefined; var ps: [4]*Prompt = undefined; @@ -773,7 +774,7 @@ pub fn main(init: std.process.Init) !void { var outs: [2][48]u32 = undefined; const logits: [2][]u16 = .{ try gpa.alloc(u16, VOCAB), try gpa.alloc(u16, VOCAB) }; var best: [2]f64 = .{ 1e9, 1e9 }; - const ab: []const u8 = if (std.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; + const ab: []const u8 = if (platform.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; for (0..2) |arm| { // arm 0 the old kernels, arm 1 the new: FZ_AB attn (default), scan or tiles if (std.mem.eql(u8, ab, "scan")) pr.scan4 = arm == 1 else if (std.mem.eql(u8, ab, "tiles")) pr.tall_tiles = arm == 1 else if (std.mem.eql(u8, ab, "gu")) pr.fused_gu = arm == 1 else if (std.mem.eql(u8, ab, "bm32")) pr.expert_bm = if (arm == 1) 32 else 0 else if (std.mem.eql(u8, ab, "scan8")) pr.scan8 = arm == 1 else if (std.mem.eql(u8, ab, "seg2")) {} else pr.fast_attn = arm == 1; for (0..3) |run| { @@ -830,7 +831,7 @@ pub fn main(init: std.process.Init) !void { const a0 = mtl.clock.seconds(); var at: usize = 0; var last_n: usize = 1; - const ref_upto: usize = if (std.c.getenv("FZ_PF_REF")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else toks.len; + const ref_upto: usize = if (platform.c.getenv("FZ_PF_REF")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else toks.len; while (at < @min(ref_upto, toks.len)) { const n: usize = @min(MAXR, toks.len - at); try m.window(toks[at .. at + n], &pk); @@ -865,7 +866,7 @@ pub fn main(init: std.process.Init) !void { } // the prompt path var pr = try Prompt.init(&r, args[2], r.xnew_header); - if (std.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); + if (platform.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); m.reset(); const b0 = mtl.clock.seconds(); var first: u32 = 0; @@ -927,7 +928,7 @@ pub fn main(init: std.process.Init) !void { pr.skip = 0; return; } - if (std.c.getenv("FZ_QMM6") != null) { // 6-bit tensor-unit projections for prompt chunks: checked and timed + if (platform.c.getenv("FZ_QMM6") != null) { // 6-bit tensor-unit projections for prompt chunks: checked and timed const lib = try mtl.Library.fromSource(device, try fz.frags.source(device, arena, ks.flashnext_qmm6), mtl.CompileOptions.mlx()); const qmm_pipe = try mtl.Pipeline.init(device, lib, "tf_qmm6_t_nax", false); const off_pipe = try mtl.Pipeline.init(device, lib, "tf_expert_offsets6", false); @@ -1061,7 +1062,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_OVERLAP") != null) { // how many experts a window's rows share, layer by layer + if (platform.c.getenv("FZ_OVERLAP") != null) { // how many experts a window's rows share, layer by layer const pr = ref.object.get("prompt").?.array.items; const want0 = ref.object.get("tokens").?.array.items; r.ar = .{ .b = try r.buffer(64) }; @@ -1148,10 +1149,10 @@ pub fn main(init: std.process.Init) !void { var pick: [MAXR]u32 = undefined; // FZ_LONG: prompts of any length through the prompt chunks (reference and GPU-side rounds alike) - const long_mode = std.c.getenv("FZ_LONG") != null; - if (std.c.getenv("FZ_PRE1") != null) try m.window(&.{1000}, &pick); // a decode window before the prompt chunks + const long_mode = platform.c.getenv("FZ_LONG") != null; + if (platform.c.getenv("FZ_PRE1") != null) try m.window(&.{1000}, &pick); // a decode window before the prompt chunks var pr_long: ?Prompt = if (long_mode) try Prompt.init(&r, args[2], r.xnew_header) else null; - if (pr_long) |*pr| if (std.c.getenv("FZ_STEP")) |v| { + if (pr_long) |*pr| if (platform.c.getenv("FZ_STEP")) |v| { pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); }; const ptoks = try gpa.alloc(u32, prompt.len); @@ -1190,20 +1191,20 @@ pub fn main(init: std.process.Init) !void { const steps: f64 = @floatFromInt(want.len - 1); std.debug.print("one-row reference starts {any}\n", .{got.items[0..@min(24, got.items.len)]}); std.debug.print("one row: {d}/{d} tokens equal to Python's; {d:.1} tok/s ({d:.2} ms a step, GPU {d:.2} ms)\n", .{ same, want.len, steps / (t4 - t3), (t4 - t3) * 1e3 / steps, m.gpu_seconds * 1e3 / steps }); - if (std.c.getenv("FZ_ONLY1") != null) { + if (platform.c.getenv("FZ_ONLY1") != null) { if (same < want.len) std.debug.print("first difference at token {d}: got {d}, Python {d}\n", .{ same, got.items[same], want[same].integer }); return; } // 2. the Python engine's drafted windows, every row's pick, with rollback m.reset(); - if (!long_mode and std.c.getenv("FZ_ONLY17") == null) for (prompt) |tok| { + if (!long_mode and platform.c.getenv("FZ_ONLY17") == null) for (prompt) |tok| { try m.window(&.{@intCast(tok.integer)}, &pick); m.keepRows(&.{@intCast(tok.integer)}, 1); }; var bad: usize = 0; var total_rows: usize = 0; - const skip_mid = long_mode or std.c.getenv("FZ_ONLY17") != null; // the one-row reference, then GPU-side rounds only + const skip_mid = long_mode or platform.c.getenv("FZ_ONLY17") != null; // the one-row reference, then GPU-side rounds only const rounds = if (own_prompt or skip_mid) &[_]std.json.Value{} else ref.object.get("rounds").?.array.items; for (rounds, 0..) |round, ri| { const o = round.object; @@ -1310,7 +1311,7 @@ pub fn main(init: std.process.Init) !void { const adapt_story = [_][3]usize{ .{ 3, 3, 0 }, .{ 0, 7, 9 } }; const adapt_own = [_][3]usize{ .{ 3, 3, 0 }, .{ 6, 6, 0 }, .{ 0, 7, 9 } }; const adapt: []const [3]usize = if (own_prompt) &adapt_own else &adapt_story; - const copy_min: usize = if (std.c.getenv("FZ_COPY_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 3; + const copy_min: usize = if (platform.c.getenv("FZ_COPY_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 3; var hist: std.ArrayList(u32) = .empty; for (prompt) |x| try hist.append(gpa, @intCast(x.integer)); const n_prompt = hist.items.len; @@ -1441,7 +1442,7 @@ pub fn main(init: std.process.Init) !void { // N+1 while round N runs and reads the emitted tokens from a ring. Depth 0: the depth rule picks each round's // drafts from the rounds read back so far (the window two rounds behind the one being encoded). var depths7: std.ArrayList(usize) = .empty; - if (std.c.getenv("FZ_D7")) |v| { // e.g. FZ_D7=6,0 + if (platform.c.getenv("FZ_D7")) |v| { // e.g. FZ_D7=6,0 var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |d| try depths7.append(gpa, try std.fmt.parseInt(usize, d, 10)); } else try depths7.appendSlice(gpa, &.{ 3, 6, 0 }); diff --git a/zig/tests/flashnext_think.zig b/zig/tests/flashnext_think.zig index 283d217b36..8c7bbe0e0b 100644 --- a/zig/tests/flashnext_think.zig +++ b/zig/tests/flashnext_think.zig @@ -1,5 +1,6 @@ //! tf-flashnext-think [CASES] [SEED]: Flash Next's verdict kernels under thinking budgets against the lane core. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -13,8 +14,8 @@ pub fn main(init: std.process.Init) !void { defer dev.deinit(); var g = try tf.flashnext_think_check.Rig.init(dev); defer g.deinit(); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const sum = try tf.flashnext_think_check.run(init.gpa, &g, cases, seed); - const s = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e9; // mach ticks are 125/3 ns on Apple silicon + const s = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e9; // mach ticks are 125/3 ns on Apple silicon std.debug.print("{s}: {d} cases (seed {d}), plain and drafted == the lane core in each; closed at the budget {d}, think end first {d}; {d:.1} s\n", .{ std.mem.span(dev.name()), sum.cases, seed, sum.closed, sum.ended_first, s }); } diff --git a/zig/tests/glm_attn_probe.zig b/zig/tests/glm_attn_probe.zig index c6bc8ac6e6..e86a001d61 100644 --- a/zig/tests/glm_attn_probe.zig +++ b/zig/tests/glm_attn_probe.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's dense latent attention (forward.attendDense) on given inputs, for a bit check against a reference. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -57,8 +58,8 @@ pub fn main(init: std.process.Init) !void { try out.appendSlice(arena, bufs[3].contents()[0 .. 64 * n * 2]); try out.appendSlice(arena, bufs[4].contents()[0..qb]); } - const file = std.c.fopen(try std.fmt.allocPrintSentinel(arena, "{s}", .{args[2]}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(out.items.ptr, 1, out.items.len, file) != out.items.len) return error.WriteFailed; + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(arena, "{s}", .{args[2]}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(out.items.ptr, 1, out.items.len, file) != out.items.len) return error.WriteFailed; std.debug.print("{d} cases\n", .{cases}); } diff --git a/zig/tests/glm_dsa_prompt_check.zig b/zig/tests/glm_dsa_prompt_check.zig index e88d50ebac..9141bf7fcf 100644 --- a/zig/tests/glm_dsa_prompt_check.zig +++ b/zig/tests/glm_dsa_prompt_check.zig @@ -1,5 +1,6 @@ //! A prompt window's experts (every expert once, on the tensor units) against the decode path's 16-row windows, by row. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const g = tf.glm_dsa; @@ -53,8 +54,8 @@ pub fn main(init: std.process.Init) !void { for (q.x.slice(f32, @as(usize, rows) * hid)) |*v| v.* = prng.random().floatNorm(f32); const dec = try device.buffer(@as(usize, rows) * hid * 4, mtl.ResourceOptions.shared); const out = try device.buffer(@as(usize, rows) * hid * 4, mtl.ResourceOptions.shared); - if (std.c.getenv("GLM_ATTN") != null) return attention(&K, &L, &s, &q, queue, device, &c, rows, w.full); - if (std.c.getenv("GLM_NARROW") != null) return narrow(&K, &L, &q, queue, device, &c); + if (platform.c.getenv("GLM_ATTN") != null) return attention(&K, &L, &s, &q, queue, device, &c, rows, w.full); + if (platform.c.getenv("GLM_NARROW") != null) return narrow(&K, &L, &q, queue, device, &c); // the decode path: route and experts 16 rows at a time (each window's rows copied into the 16-row scratch) var dec_ms: f64 = 0; var k: u32 = 0; diff --git a/zig/tests/glm_dsa_run.zig b/zig/tests/glm_dsa_run.zig index 6c918d5676..94132983a1 100644 --- a/zig/tests/glm_dsa_run.zig +++ b/zig/tests/glm_dsa_run.zig @@ -1,5 +1,6 @@ //! GLM-5.3 on the pair without the server: rank 0 runs a prompt and its greedy reply, rank 1 follows. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const tok = @import("tokenizer"); const g = tf.glm_dsa; @@ -13,7 +14,7 @@ const usage = ; fn env(name: [:0]const u8, default: u64) u64 { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(u64, std.mem.span(v), 10) catch default; } @@ -30,9 +31,9 @@ pub fn main(init: std.process.Init) !void { } var ram: u64 = 0; var len: usize = 8; - if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; + if (platform.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; const limit = env("GLM_LIMIT_GB", ram * 7 / 10 / 1_000_000_000) * 1_000_000_000; - const layers: ?u32 = if (std.c.getenv("GLM_LAYERS") != null) @intCast(env("GLM_LAYERS", 78)) else null; + const layers: ?u32 = if (platform.c.getenv("GLM_LAYERS") != null) @intCast(env("GLM_LAYERS", 78)) else null; const depth: u32 = @intCast(env("GLM_DEPTH", 0)); const e = try g.engine.Engine.load(a, io, argv[1], argv[2], layers, depth > 0, @intCast(env("GLM_CONTEXT", 8192)), 1, limit); defer e.deinit(); @@ -53,7 +54,7 @@ pub fn main(init: std.process.Init) !void { } fn once(e: *g.engine.Engine, io: std.Io, t: *const tok.Tokenizer, prompt: []const u32, max: u64, depth: u32) !void { - const forced = std.c.getenv("GLM_FORCED") != null; + const forced = platform.c.getenv("GLM_FORCED") != null; var agree: usize = 0; const t0 = std.Io.Timestamp.now(io, .awake); try g.follow.begin(e, 0, prompt); diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index f2ceb725a0..e91e0657db 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -1,5 +1,6 @@ //! tf-glm-rows-speed: GLM-5.3-Flash's decode products at the checkpoints' widths and shapes, GPU time and GB/s. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const sources = @import("kernel_sources"); @@ -555,13 +556,13 @@ pub fn main() !void { 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK + sh8[0].bytes + sh8[1].bytes, ); - if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + if (platform.c.getenv("GLM_SPEED_CANDIDATES") != null) { const lib = try cand.candidateLib(device); defer lib.deinit(); try cand.moeCandidates(&b, lib, p, gu2, dn2, sh8); } } - if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + if (platform.c.getenv("GLM_SPEED_CANDIDATES") != null) { const lib = try cand.candidateLib(device); defer lib.deinit(); std.debug.print("candidates at 8 bits (bits checked == glm_rows):\n", .{}); @@ -571,7 +572,7 @@ pub fn main() !void { try cand.candidates(&b, lib, p, m, s.name); } } - if (std.c.getenv("GLM_SPEED_PROMPT") != null) try speed_prompt.prompt(&b, k, p); + if (platform.c.getenv("GLM_SPEED_PROMPT") != null) try speed_prompt.prompt(&b, k, p); std.debug.print("uniform 4/64:\n", .{}); for (uni.rows[0..uni.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); uni.total("uniform 4/64"); diff --git a/zig/tests/glm_run.zig b/zig/tests/glm_run.zig index 13194d6d0f..97ad4105a8 100644 --- a/zig/tests/glm_run.zig +++ b/zig/tests/glm_run.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash replies on the native engine at each draft depth: hashes, first differences from depth 0 and the reference, speeds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const glm = tf.glm; @@ -53,7 +54,7 @@ const Collect = struct { }; fn env(name: [:0]const u8, default: []const u8) []const u8 { - return if (std.c.getenv(name)) |v| std.mem.span(v) else default; + return if (platform.c.getenv(name)) |v| std.mem.span(v) else default; } fn hash(t: []const u32) u64 { @@ -124,13 +125,13 @@ pub fn main(init: std.process.Init) !void { const e = try glm.engine.Engine.load(gpa, args[1], cap); defer e.deinit(); std.debug.print("loaded in {d:.1} s: {d} of {d} layers, routed experts {d}-{d}, {d:.1} GB of weights, MTP head {s}, {d:.2} GB of caches and buffers, prompt chunks of {d} rows\n", .{ e.load_seconds, e.c.run, e.c.layers, e.c.own[0], e.c.own[1] - 1, @as(f64, @floatFromInt(e.w.bytes)) / 1e9, if (e.hasMtp()) "yes" else "no", @as(f64, @floatFromInt(e.arena.bytes)) / 1e9, if (e.pr) |p| p.rows else 0 }); - const strict = std.c.getenv("GLM_REF_STRICT") != null; + const strict = platform.c.getenv("GLM_REF_STRICT") != null; const eos = e.c.eos[0..e.c.eos_n]; - if (std.c.getenv("GLM_CAPTURE")) |path| { // the first prompt's first window, sublayer by sublayer (glm_ref.py --capture) + if (platform.c.getenv("GLM_CAPTURE")) |path| { // the first prompt's first window, sublayer by sublayer (glm_ref.py --capture) const first = doc.object.get("prompts").?.array.items[0]; try e.capture(try ints(arena, first.object.get("ids").?), std.mem.span(path)); } - if (std.c.getenv("GLM_TRACE")) |spec| { // NAME:STEPS:PATH, glm_ref.py --trace NAME STEPS PATH's twin + if (platform.c.getenv("GLM_TRACE")) |spec| { // NAME:STEPS:PATH, glm_ref.py --trace NAME STEPS PATH's twin var parts = std.mem.splitScalar(u8, std.mem.span(spec), ':'); const want = parts.next() orelse return error.BadTrace; const steps = try std.fmt.parseInt(usize, parts.next() orelse return error.BadTrace, 10); @@ -139,18 +140,18 @@ pub fn main(init: std.process.Init) !void { try e.trace(try ints(arena, p.object.get("ids").?), try ints(arena, p.object.get("expect").?), steps, out_path); }; } - if (std.c.getenv("GLM_FORCED") != null) for (doc.object.get("prompts").?.array.items) |p| { + if (platform.c.getenv("GLM_FORCED") != null) for (doc.object.get("prompts").?.array.items) |p| { const want = try ints(arena, p.object.get("expect") orelse continue); const r = try e.forced(try ints(arena, p.object.get("ids").?), want); std.debug.print("{s} forced: {d} of {d} picks equal the reference's", .{ p.object.get("name").?.string, r.same, want.len }); if (r.first) |i| std.debug.print(", first differs at {d}\n", .{i}) else std.debug.print("\n", .{}); }; var failures: usize = 0; - const vs: ?std.json.Value = if (std.c.getenv("GLM_VS")) |path| blk: { + const vs: ?std.json.Value = if (platform.c.getenv("GLM_VS")) |path| blk: { const f = try mtl.MappedFile.open(path); break :blk try std.json.parseFromSliceLeaky(std.json.Value, arena, f.bytes[0..f.size], .{}); } else null; - const trace_prefix: ?[*:0]const u8 = std.c.getenv("GLM_TRACE_LAST"); + const trace_prefix: ?[*:0]const u8 = platform.c.getenv("GLM_TRACE_LAST"); const trace_bytes = (2 + 4 * @as(usize, e.c.run)) * e.c.hidden * 2; if (trace_prefix != null) e.trace_last = try e.arena.buffer(trace_bytes); var saved: std.ArrayList(u8) = .empty; @@ -168,10 +169,10 @@ pub fn main(init: std.process.Init) !void { if (trace_prefix != null and run == 0 and d == depths.items[0]) col.trace_from = e.trace_last.?.addr()[0..trace_bytes]; var margins: std.ArrayList(f32) = .empty; defer margins.deinit(gpa); - const want_margins = run == 0 and d == depths.items[0] and (std.c.getenv("GLM_MARGINS") != null or std.c.getenv("GLM_MARGINS_VS") != null); + const want_margins = run == 0 and d == depths.items[0] and (platform.c.getenv("GLM_MARGINS") != null or platform.c.getenv("GLM_MARGINS_VS") != null); e.margins = if (want_margins) &margins else null; defer e.margins = null; - const want_logits = run == 0 and d == depths.items[0] and (std.c.getenv("GLM_LOGITS") != null or std.c.getenv("GLM_LOGITS_VS") != null); + const want_logits = run == 0 and d == depths.items[0] and (platform.c.getenv("GLM_LOGITS") != null or platform.c.getenv("GLM_LOGITS_VS") != null); if (want_logits) col.logits_from = @as([*]const u16, @ptrCast(@alignCast(e.sc.logits.addr())))[0..e.c.vocab]; const r = try e.generate(ids, max, eos, d, .{ .ctx = &col, .prefilled = Collect.prefilled, .tokens = Collect.tokens, .cancelled = Collect.cancelled }); const toks = col.toks.items; @@ -182,7 +183,7 @@ pub fn main(init: std.process.Init) !void { const rn: f64 = @floatFromInt(r.rounds); std.debug.print(" a round: {d:.2} ms wall, {d:.2} ms GPU, {d:.2} ms encoding, {d:.2} ms GPU idle between rounds; {d:.2} tokens\n", .{ r.decode_seconds * 1e3 / rn, r.gpu_seconds * 1e3 / rn, r.encode_seconds * 1e3 / rn, r.gap_seconds * 1e3 / @max(rn - 1, 1), @as(f64, @floatFromInt(toks.len -| 1)) / rn }); if (r.copy_rounds > 0) std.debug.print(" copy drafts: {d} rounds, {d} drafts kept ({d:.2} a copy round)\n", .{ r.copy_rounds, r.copy_accepted, @as(f64, @floatFromInt(r.copy_accepted)) / @as(f64, @floatFromInt(r.copy_rounds)) }); - if (std.c.getenv("GLM_RANKS") != null) { + if (platform.c.getenv("GLM_RANKS") != null) { std.debug.print(" draft ranks (kept: all / miss at 2nd / 3rd / later):", .{}); for (r.draft_ranks, 0..) |row, m| if (row[0] + row[1] + row[2] + row[3] > 0) std.debug.print(" m{d} {d}/{d}/{d}/{d}", .{ m, row[0], row[1], row[2], row[3] }); std.debug.print("\n", .{}); @@ -192,17 +193,17 @@ pub fn main(init: std.process.Init) !void { failures += 1; std.debug.print(" DIFFERS from depth {d} at token {d}\n", .{ depths.items[0], at }); } else std.debug.print(" equal to depth {d}\n", .{depths.items[0]}); - if (want_margins) if (std.c.getenv("GLM_MARGINS")) |prefix| { + if (want_margins) if (platform.c.getenv("GLM_MARGINS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.margins", .{ prefix, name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(std.mem.sliceAsBytes(margins.items).ptr, 4, margins.items.len, file) != margins.items.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(std.mem.sliceAsBytes(margins.items).ptr, 4, margins.items.len, file) != margins.items.len) return error.WriteFailed; }; if (vs) |other| if (other.object.get(name)) |theirs| { const want = try ints(arena, theirs); if (firstDiff(want, toks)) |at| { std.debug.print(" vs GLM_VS: first token {s}, first difference at token {d} of {d}\n", .{ if (at == 0) "DIFFERS" else "equal", at, @min(want.len, toks.len) }); - if (want_margins) if (std.c.getenv("GLM_MARGINS_VS")) |prefix| { + if (want_margins) if (platform.c.getenv("GLM_MARGINS_VS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.margins", .{ prefix, name }, 0); if (mtl.MappedFile.open(path)) |f| { defer f.deinit(); @@ -216,18 +217,18 @@ pub fn main(init: std.process.Init) !void { }; if (col.trace.len > 0) { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.trace", .{ std.mem.span(trace_prefix.?), name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(col.trace.ptr, 1, col.trace.len, file) != col.trace.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(col.trace.ptr, 1, col.trace.len, file) != col.trace.len) return error.WriteFailed; } if (want_logits and col.logits.len == 0) return error.NoLogits; - if (want_logits) if (std.c.getenv("GLM_LOGITS")) |prefix| { + if (want_logits) if (platform.c.getenv("GLM_LOGITS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(std.mem.sliceAsBytes(col.logits).ptr, 2, col.logits.len, file) != col.logits.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(std.mem.sliceAsBytes(col.logits).ptr, 2, col.logits.len, file) != col.logits.len) return error.WriteFailed; }; - if (want_logits) if (std.c.getenv("GLM_LOGITS_VS")) |prefix| try compareLogits(gpa, col.logits, try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0)); + if (want_logits) if (platform.c.getenv("GLM_LOGITS_VS")) |prefix| try compareLogits(gpa, col.logits, try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0)); if (run == 0 and d == depths.items[0]) { if (saved.items.len > 1) try saved.append(arena, ','); try saved.print(arena, "\"{s}\":[", .{name}); @@ -241,7 +242,7 @@ pub fn main(init: std.process.Init) !void { } else std.debug.print(" reference: {d} of {d} tokens equal\n", .{ @min(want.len, toks.len), want.len }); } }; - if (std.c.getenv("GLM_CANCEL_TEST") != null) { // a reply cancelled after its first prompt window, then a fresh one + if (platform.c.getenv("GLM_CANCEL_TEST") != null) { // a reply cancelled after its first prompt window, then a fresh one var cut: Collect = .{ .gpa = gpa, .cancel_at = 1 }; defer cut.toks.deinit(gpa); const rc = try e.generate(ids, max, eos, depths.items[0], .{ .ctx = &cut, .prefilled = Collect.prefilled, .tokens = Collect.tokens, .cancelled = Collect.cancelled }); @@ -253,23 +254,23 @@ pub fn main(init: std.process.Init) !void { std.debug.print(" cancel test: cancelled reply ended {t} after {d} tokens; the next fresh reply {s} the plain one\n", .{ rc.reason, cut.toks.items.len, if (same) "equals" else "DIFFERS from" }); } } - if (std.c.getenv("GLM_PROFILE")) |spec| { + if (platform.c.getenv("GLM_PROFILE")) |spec| { const reps = try std.fmt.parseInt(usize, env("GLM_PROFILE_REPS", "12"), 10); var pit = std.mem.tokenizeScalar(u8, std.mem.span(spec), ','); - const only = std.c.getenv("GLM_PROFILE_ONLY") != null; // each class alone instead of each left out - const parts = std.c.getenv("GLM_PROFILE_PARTS") != null; // by launch instead of by class + const only = platform.c.getenv("GLM_PROFILE_ONLY") != null; // each class alone instead of each left out + const parts = platform.c.getenv("GLM_PROFILE_PARTS") != null; // by launch instead of by class while (pit.next()) |dv| try e.profile(try std.fmt.parseInt(u32, dv, 10), reps, only, parts); } - if (std.c.getenv("GLM_MM_CHECK") != null) try e.checkMatmul(try ints(arena, doc.object.get("prompts").?.array.items[0].object.get("ids").?)); - if (std.c.getenv("GLM_PROMPT_PROFILE")) |v| { // each prompt's first chunk by class (GLM_PROFILE_ONLY: each alone) + if (platform.c.getenv("GLM_MM_CHECK") != null) try e.checkMatmul(try ints(arena, doc.object.get("prompts").?.array.items[0].object.get("ids").?)); + if (platform.c.getenv("GLM_PROMPT_PROFILE")) |v| { // each prompt's first chunk by class (GLM_PROFILE_ONLY: each alone) const reps = try std.fmt.parseInt(usize, std.mem.span(v), 10); - for (doc.object.get("prompts").?.array.items) |p| try e.profilePrompt(try ints(arena, p.object.get("ids").?), reps, std.c.getenv("GLM_PROFILE_ONLY") != null, std.c.getenv("GLM_PROFILE_PARTS") != null); + for (doc.object.get("prompts").?.array.items) |p| try e.profilePrompt(try ints(arena, p.object.get("ids").?), reps, platform.c.getenv("GLM_PROFILE_ONLY") != null, platform.c.getenv("GLM_PROFILE_PARTS") != null); } try saved.append(arena, '}'); - if (std.c.getenv("GLM_OUT")) |path| { - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(saved.items.ptr, 1, saved.items.len, file) != saved.items.len) return error.WriteFailed; + if (platform.c.getenv("GLM_OUT")) |path| { + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(saved.items.ptr, 1, saved.items.len, file) != saved.items.len) return error.WriteFailed; } if (failures > 0) { std.debug.print("{d} replies differ from the plain ones{s}\n", .{ failures, if (strict) " or the references" else "" }); diff --git a/zig/tests/linktest.zig b/zig/tests/linktest.zig index dfdad16efe..bf251fd1ec 100644 --- a/zig/tests/linktest.zig +++ b/zig/tests/linktest.zig @@ -1,5 +1,6 @@ //! Link qualification over the system's verbs; `run` swaps connection details on stdin and stdout and opens no socket. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const suite = @import("linktest_suite.zig"); const probe = @import("linktest_probe.zig"); @@ -30,7 +31,7 @@ fn readLine(buf: []u8) ![]const u8 { var n: usize = 0; while (n < buf.len) { var c: [1]u8 = undefined; - if (std.c.read(0, &c, 1) != 1) return error.EndOfInput; + if (platform.c.read(0, &c, 1) != 1) return error.EndOfInput; if (c[0] == '\n') return buf[0..n]; buf[n] = c[0]; n += 1; @@ -61,7 +62,7 @@ fn parseHello(line: []const u8) ![2]vqp.Info { /// Exit when stdin closes: the run ends with the session that started it. fn tether() void { var c: [1]u8 = undefined; - while (std.c.read(0, &c, 1) == 1) {} + while (platform.c.read(0, &c, 1) == 1) {} out("R stdin_closed", .{}); std.process.exit(4); } @@ -79,7 +80,7 @@ fn mapped(len: usize) ![]align(vl.page) u8 { /// Two queue pairs on one device: collectives on one, raw measurements on the other, so their receives never mix. fn run(gpa: std.mem.Allocator, v: *const Verbs, a: Args) !void { _ = try std.Thread.spawn(.{}, watchdog, .{a.deadline_s}); - out("PID {d}", .{std.c.getpid()}); + out("PID {d}", .{platform.c.getpid()}); var peer: vl.Peer = .{ .ep = try vqp.Endpoint.open(v, a.dev, a.kind, 4095, a.gid) }; defer peer.ep.close(); var raw = try peer.ep.sibling(a.kind, 4095); @@ -138,7 +139,7 @@ pub fn main(init: std.process.Init) !void { if (!std.mem.eql(u8, args[1], "run")) std.process.exit(2); const a = try parse(args); run(init.gpa, &v, a) catch |err| { - out("R error what={s} errno={d}", .{ @errorName(err), std.c._errno().* }); + out("R error what={s} errno={d}", .{ @errorName(err), platform.c._errno().* }); std.process.exit(1); }; } diff --git a/zig/tests/linktest_probe.zig b/zig/tests/linktest_probe.zig index ce3402075a..a5c6af2541 100644 --- a/zig/tests/linktest_probe.zig +++ b/zig/tests/linktest_probe.zig @@ -1,5 +1,6 @@ //! What a verbs device reports and what it actually allows: registration sizes and counts, queue pair types and counts. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const abi = fabric.verbs_abi; const Verbs = fabric.verbs.Verbs; @@ -8,7 +9,7 @@ const out = @import("linktest_suite.zig").out; const mib = 1 << 20; fn errno() c_int { - return std.c._errno().*; + return platform.c._errno().*; } pub fn device(v: *const Verbs, dev: []const u8) !void { diff --git a/zig/tests/linktest_suite.zig b/zig/tests/linktest_suite.zig index 37b998b9f7..f356e6e519 100644 --- a/zig/tests/linktest_suite.zig +++ b/zig/tests/linktest_suite.zig @@ -1,5 +1,6 @@ //! The link measurements over Thunderbolt's two-sided UC: opcodes, ping-pong, bandwidth, integrity, barrier, all-reduce. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const collective = fabric.collective; const vqp = fabric.verbs_qp; @@ -53,7 +54,7 @@ pub const Ctx = struct { pub fn out(comptime fmt: []const u8, args: anytype) void { var buf: [4096]u8 = undefined; const line = std.fmt.bufPrint(&buf, fmt ++ "\n", args) catch return; - _ = std.c.write(1, line.ptr, line.len); + _ = platform.c.write(1, line.ptr, line.len); } fn us(ns: u64) f64 { diff --git a/zig/tests/mm_bench.zig b/zig/tests/mm_bench.zig index 1ec779f916..76f6402c12 100644 --- a/zig/tests/mm_bench.zig +++ b/zig/tests/mm_bench.zig @@ -1,5 +1,6 @@ //! Prompt matmul kernels compiled from .metal files at run time: chained speed at one shape, bits against the first, error against fp64. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -17,7 +18,7 @@ const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, sums_row: i32 = const Kernel = struct { path: []const u8, bm: u32 = 64, bn: u32 = 64, threads: u32 = 128, entry: []const u8 = "tf_affine_mm", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, sums_pipe: mtl.Pipeline = undefined, sums: mtl.Buffer = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined, sums_ms: f64 = 0 }; fn envInt(name: [:0]const u8, default: usize) usize { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; } @@ -55,7 +56,7 @@ pub fn main(init: std.process.Init) !void { const trials = @min(envInt("MM_TRIALS", 5), 64); const copies = envInt("MM_COPIES", 4); const check_rows = @min(envInt("MM_CHECK_ROWS", 16), M); - const out_f32 = std.c.getenv("MM_F32") != null; + const out_f32 = platform.c.getenv("MM_F32") != null; const ob: usize = if (out_f32) 4 else 2; const kernels = try arena.alloc(Kernel, args.len - 4); for (kernels, args[4..]) |*k, s| k.* = try parseKernel(s); diff --git a/zig/tests/moe_bench.zig b/zig/tests/moe_bench.zig index 6e381ff8d5..f611ef2575 100644 --- a/zig/tests/moe_bench.zig +++ b/zig/tests/moe_bench.zig @@ -1,5 +1,6 @@ //! Routed-expert prompt kernels from .metal files compiled at run time: a chunk's skewed top-k routing, chained speed, bits against the first. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -16,12 +17,12 @@ const Args = extern struct { rows: i32, n: i32, k: i32, experts: i32 }; const Kernel = struct { path: []const u8, bm: u32 = 64, entry: []const u8 = "", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined }; fn envInt(name: [:0]const u8, default: usize) usize { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; } fn envFloat(name: [:0]const u8, default: f64) f64 { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseFloat(f64, std.mem.span(v)) catch default; } @@ -37,7 +38,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("usage: {s}\n", .{usage}); std.process.exit(2); } - const down = if (std.c.getenv("MOE_MODE")) |v| std.mem.eql(u8, std.mem.span(v), "down") else false; + const down = if (platform.c.getenv("MOE_MODE")) |v| std.mem.eql(u8, std.mem.span(v), "down") else false; const T = try std.fmt.parseInt(u32, args[1], 10); const E: u32 = @intCast(envInt("MOE_EXPERTS", 288)); const topk: u32 = @intCast(envInt("MOE_TOPK", 8)); @@ -65,7 +66,7 @@ pub fn main(init: std.process.Init) !void { total += w.*; } const pick = try arena.alloc(u32, n); - const exact = std.c.getenv("MOE_EXACT") != null; // pairs dealt round-robin: every expert the same row count + const exact = platform.c.getenv("MOE_EXACT") != null; // pairs dealt round-robin: every expert the same row count if (exact) for (pick, 0..) |*q, i| { q.* = @intCast((i % topk + (i / topk) * topk) % E); }; diff --git a/zig/tests/nemotron_prefill_check.zig b/zig/tests/nemotron_prefill_check.zig index 5bca853939..f1b91302aa 100644 --- a/zig/tests/nemotron_prefill_check.zig +++ b/zig/tests/nemotron_prefill_check.zig @@ -1,5 +1,6 @@ //! One prompt chunk on the prefill kernels against tools/zig/prefill_dump.py's rows: embedding, every layer, final norm. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -17,7 +18,7 @@ pub fn main(init: std.process.Init) !void { var it = std.mem.tokenizeScalar(u8, args[2], ','); while (it.next()) |t| try list.append(arena, try std.fmt.parseInt(u32, t, 10)); // TF_PREFILL_PREFIX=N: the first N ids fill the caches as a chunk first; the rest is the checked chunk - const before: usize = if (std.c.getenv("TF_PREFILL_PREFIX")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 0; + const before: usize = if (platform.c.getenv("TF_PREFILL_PREFIX")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 0; const all = list.items; const ids = all[before..]; const n = ids.len; diff --git a/zig/tests/pocket_fit_check.zig b/zig/tests/pocket_fit_check.zig index 7f7bfaf2f5..0fe514684a 100644 --- a/zig/tests/pocket_fit_check.zig +++ b/zig/tests/pocket_fit_check.zig @@ -1,5 +1,6 @@ //! A dumped pocket lesson's read rows fitted as the learner fits them, timed; each fit's numbers and weights written. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const pocket_fit = tf.nemotron.pocket_fit; @@ -57,7 +58,7 @@ pub fn main(init: std.process.Init) !void { const us = try gpa.alloc(f32, jobs.items.len * d); const fits = try gpa.alloc(Fit, jobs.items.len); // TF_FIT_THREADS: that many threads take the fits in turn (default: a thread a fit, as the learner runs them) - const want = if (std.c.getenv("TF_FIT_THREADS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else jobs.items.len; + const want = if (platform.c.getenv("TF_FIT_THREADS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else jobs.items.len; const threads = try gpa.alloc(std.Thread, @min(@max(want, 1), jobs.items.len)); var next = std.atomic.Value(usize).init(0); const work: Work = .{ .gpa = gpa, .rows = rows, .kinds = all_kinds, .m = m, .slots = slots, .d = d, .jobs = jobs.items, .us = us, .fits = fits, .next = &next }; diff --git a/zig/tests/qwen27/dflash_graph.zig b/zig/tests/qwen27/dflash_graph.zig index bbd9db4865..e382276faf 100644 --- a/zig/tests/qwen27/dflash_graph.zig +++ b/zig/tests/qwen27/dflash_graph.zig @@ -1,5 +1,6 @@ //! Generated draft tensors exercise the real callback backend without checkpoint files. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const q = tf.qwen27; @@ -47,19 +48,19 @@ pub fn createQ4(a: std.mem.Allocator, target: *q.model.Model) !*Model { } try header.writer.writeByte('}'); var path_buf: [96]u8 = undefined; - const path = try std.fmt.bufPrintSentinel(&path_buf, "/tmp/tf-q27-q4-draft-{d}.safetensors", .{std.c.getpid()}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const path = try std.fmt.bufPrintSentinel(&path_buf, "/tmp/tf-q27-q4-draft-{d}.safetensors", .{platform.c.getpid()}, 0); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.Create; - defer _ = std.c.unlink(path); + defer _ = platform.c.unlink(path); var head: [8]u8 = undefined; std.mem.writeInt(u64, &head, header.written().len, .little); for ([_][]const u8{ &head, header.written(), std.mem.sliceAsBytes(data.items) }) |part| { - if (std.c.write(fd, part.ptr, part.len) != @as(isize, @intCast(part.len))) { - _ = std.c.close(fd); + if (platform.c.write(fd, part.ptr, part.len) != @as(isize, @intCast(part.len))) { + _ = platform.c.close(fd); return error.ShortWrite; } } - _ = std.c.close(fd); + _ = platform.c.close(fd); return Model.open(a, target, config, &.{path}, .prepared_q4_reference); } diff --git a/zig/tests/qwen27/dflash_run.zig b/zig/tests/qwen27/dflash_run.zig index 06660eed7e..2d75b9d491 100644 --- a/zig/tests/qwen27/dflash_run.zig +++ b/zig/tests/qwen27/dflash_run.zig @@ -1,5 +1,6 @@ //! The draft prototype on real weights: target tokens and final state against ordinary native generation. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("qwen27"); const profile = @import("core").gpu_profile; @@ -71,7 +72,7 @@ pub fn main(init: std.process.Init) !void { try generation.prefill(prompt, chunk); const prompt_ns = start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); var trace = profile.Trace{ .queue = target.queue }; - const profiled = std.c.getenv("QWEN27_PROFILE") != null; + const profiled = platform.c.getenv("QWEN27_PROFILE") != null; if (profiled) { runner.profile(&trace); trace.enabled = true; diff --git a/zig/tests/qwen27/dflash_served.zig b/zig/tests/qwen27/dflash_served.zig index ca9b3c2bfb..6955e546f0 100644 --- a/zig/tests/qwen27/dflash_served.zig +++ b/zig/tests/qwen27/dflash_served.zig @@ -1,5 +1,6 @@ //! Generated native draft and plain requests retain CLI tokens and final state through the served host. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("tensorfold").qwen27; const api = @import("engine_api"); @@ -73,7 +74,7 @@ const Box = struct { }; /// Drafted requests on the lane core (its own round counts), by the host's rule: TF_QWEN27_LANES, else no tensor units. fn lanesMode(h: *host.Host) bool { - if (std.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + if (platform.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); return !h.model.device.tensorUnits(); } fn stop(_: *anyopaque, ids: []const u32) bool { diff --git a/zig/tests/qwen27/reuse.zig b/zig/tests/qwen27/reuse.zig index 46ce1ed90d..e41b5d8c04 100644 --- a/zig/tests/qwen27/reuse.zig +++ b/zig/tests/qwen27/reuse.zig @@ -1,5 +1,6 @@ //! Qwen27 prompt reuse through the served host against fresh prompt passes: each prompt's state and reply, bit for bit. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("tensorfold").qwen27; const api = @import("engine_api"); @@ -90,7 +91,7 @@ const Run = struct { /// Drafted requests on the lane core, by the served host's rule: TF_QWEN27_LANES, else a chip without tensor units. fn lanesMode(h: *host.Host) bool { - if (std.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + if (platform.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); return !h.model.device.tensorUnits(); } diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 184c401968..d7f916f4a6 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -1,5 +1,6 @@ //! Frozen server goldens for `zig build test-golden`, with the ten differences that are still known. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; const Io = std.Io; const wire = @import("wire.zig"); @@ -379,7 +380,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con try fw.interface.writeAll(content); try fw.interface.flush(); const kpz = try std.fmt.allocPrintSentinel(a, "{s}", .{kp}, 0); - _ = std.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); + _ = platform.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); } } const child = try a.create(std.process.Child); @@ -411,7 +412,7 @@ pub fn main(init: std.process.Init) !void { const cases_path = args[2]; const golden_dir = args[3]; const fixtures = args[4]; - const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{std.c.getpid()}); + const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{platform.c.getpid()}); Io.Dir.cwd().createDir(io, tmp_dir, .default_dir) catch {}; defer Io.Dir.cwd().deleteTree(io, tmp_dir) catch {}; diff --git a/zig/tests/train_check.zig b/zig/tests/train_check.zig index 08601733c3..77acc8e7c1 100644 --- a/zig/tests/train_check.zig +++ b/zig/tests/train_check.zig @@ -1,5 +1,6 @@ //! The learning step's gradients against the engine's own forward, moving each layer along its own, on the real model. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const nm = tf.nemotron; @@ -51,8 +52,8 @@ pub fn main(init: std.process.Init) !void { const with_grad = (try t.step(&ids, start, .grad, turn)).loss; const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0.toNanoseconds())) / 1e6; std.debug.print("loss {d:.5} (with gradients {d:.5}), a step with gradients {d:.1} ms, GPU {d:.1} ms\n", .{ plain, with_grad, ms, b.last_ms }); - if (std.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; - if (std.c.getenv("TF_TRAIN_PROFILE") != null) { + if (platform.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; + if (platform.c.getenv("TF_TRAIN_PROFILE") != null) { var prof: nm.encoder.Profiler = .{ .queue = m.queue, .kernels = &m.kernels }; t.profile = &prof; _ = try t.step(&ids, start, .grad, turn); diff --git a/zig/tests/weight_read_check.zig b/zig/tests/weight_read_check.zig index d5acf5e7a6..1089d93666 100644 --- a/zig/tests/weight_read_check.zig +++ b/zig/tests/weight_read_check.zig @@ -1,5 +1,6 @@ //! Check the actual native loaders with a synthetic file and bounded I/O faults. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -63,7 +64,7 @@ pub fn main(init: std.process.Init) !void { } else { try replay.indexFile(path); const entry = try replay.entry(if (grouped) "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight" else "probe"); - defer _ = std.c.close(entry.fd); + defer _ = platform.c.close(entry.fd); const result = if (grouped) replay.group(0, 2, "weight") else replay.load("probe"); const buffer = result catch |err| { if (!eof or err != error.ShortRead) return err; From e5bb3bca166111cc3cef9dd118f27ab28e22a841 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:14:21 +0000 Subject: [PATCH 02/46] Windows: CUDA kernels, the native server and the GPU suites on sm_89 - Kernel fatbins build with nvcc and cl.exe: MSVC's conforming preprocessor is passed on Windows hosts (CUDA 13's headers need it), and explicit kernel instantiations carry the __restrict__ qualifiers their templates have (MSVC counts them in the function type; the generated code is unchanged). - Nemotron's checkpoint source keeps io_uring to Linux; elsewhere its existing thread reads run. - pull_exlr8 raises no descriptor limit on Windows. - modelopt test: the lane's slice partials are allocated. GPUs without clusters (before sm_90) reduce split K through them; the test passed none, so it wrote through a null pointer on Ampere and Ada. zig build native and install build with -Dnvcc=... -Dsm=89, and tf-cuda-test's synthetic suites pass on an RTX 4090. --- zig/build/cuda.zig | 2 ++ zig/kernels/cuda/experts.cu | 4 ++-- zig/kernels/cuda/experts_pack.cu | 4 ++-- zig/kernels/cuda/experts_prefill.cu | 2 +- zig/kernels/cuda/gdn.cu | 10 +++++----- zig/kernels/cuda/lane_gemv.cu | 2 +- zig/kernels/cuda/prefill_attention.cu | 4 ++-- zig/kernels/cuda/qmm_group.cu | 2 +- zig/kernels/cuda/qmm_prefill.cu | 2 +- zig/kernels/cuda/scan_rows.cu | 8 ++++---- zig/src/cli/pull_exlr8.zig | 1 + zig/src/families/nemotron/cuda_source.zig | 12 ++++++++---- zig/src/platform/win_posix.zig | 5 +++-- zig/tests/cuda/modelopt_tests.zig | 17 +++++++++++++---- 14 files changed, 46 insertions(+), 29 deletions(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..4f54eaf76d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -226,6 +226,8 @@ fn fatbin(b: *std.Build, nvcc: []const u8, version: std.Build.LazyPath, k: Kerne const run = b.addSystemCommand(&.{ nvcc, "-fatbin" }); run.addFileInput(version); run.addArgs(&torch_flags); + // CUDA 13's headers need MSVC's conforming preprocessor when cl.exe is the host compiler + if (b.graph.host.result.os.tag == .windows) run.addArg("-Xcompiler=/Zc:preprocessor"); run.addArgs(k.flags); const a = if (k.arch_specific) "a" else ""; var it = std.mem.tokenizeScalar(u8, sms, ','); diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..d19f6ee90c 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -276,7 +276,7 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const } // namespace tf_experts // Decode form for groups of 64: up with relu^2 (epilogue 1), down to fp32 (epilogue 0). -#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16*, \ - int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); +#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16* __restrict__, \ + int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_EXPERT(1) TF_EXPERT(0) diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..b47ded555f 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -38,5 +38,5 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca } // namespace tf_experts_pack // Groups of 64 inputs. -template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*, - uint32_t*, int, int, int, int); +template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t* __restrict__, const uint16_t* __restrict__, const uint16_t* __restrict__, + uint32_t* __restrict__, int, int, int, int); diff --git a/zig/kernels/cuda/experts_prefill.cu b/zig/kernels/cuda/experts_prefill.cu index c1c7fbbac9..3a1df1dfee 100644 --- a/zig/kernels/cuda/experts_prefill.cu +++ b/zig/kernels/cuda/experts_prefill.cu @@ -152,6 +152,6 @@ __global__ void __launch_bounds__(WM * WN * 32) // Prefill form for groups of 64: up with relu^2 (epilogue 1), down to bf16 (epilogue 3). #define TF_PREFILL(EPI) template __global__ void tf_experts_prefill::prefill_kernel<64, 1, EPI, 2, 2, 4>( \ - const __nv_bfloat16*, int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); + const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_PREFILL(1) TF_PREFILL(3) diff --git a/zig/kernels/cuda/gdn.cu b/zig/kernels/cuda/gdn.cu index 574fb92de5..2c96e1a3ad 100644 --- a/zig/kernels/cuda/gdn.cu +++ b/zig/kernels/cuda/gdn.cu @@ -305,14 +305,14 @@ __global__ void __launch_bounds__(32 * WARPS) replay_kernel( } // namespace tf_gdn // The instantiations the Python wrappers launch (gdn_replay_cuda, dispatch_tree for bf16 keys). -#define TF_REPLAY(QK) template __global__ void tf_gdn::replay_kernel(const long long*, int, const int*, int, \ - const int*, int, float*, int, int, int); +#define TF_REPLAY(QK) template __global__ void tf_gdn::replay_kernel(const long long* __restrict__, int, const int* __restrict__, int, \ + const int* __restrict__, int, float* __restrict__, int, int, int); TF_REPLAY(__nv_bfloat16) TF_REPLAY(float) #define TF_TREE(S, R, W, C) template __global__ void tf_gdn::tree_kernel<__nv_bfloat16, S, R, W, C>( \ - const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, const float*, const float*, const float*, \ - const long long*, const int*, const int*, int, __nv_bfloat16*, int, int, int, tf_gdn::Pending<__nv_bfloat16>, \ - float*, const long long*, bool); + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const float* __restrict__, const float* __restrict__, const float* __restrict__, \ + const long long* __restrict__, const int* __restrict__, const int* __restrict__, int, __nv_bfloat16* __restrict__, int, int, int, tf_gdn::Pending<__nv_bfloat16>, \ + float*, const long long* __restrict__, bool); TF_TREE(0, 8, 4, true) TF_TREE(1, 8, 4, false) TF_TREE(2, 8, 2, false) diff --git a/zig/kernels/cuda/lane_gemv.cu b/zig/kernels/cuda/lane_gemv.cu index 94a78a5ae8..af6f0990ac 100644 --- a/zig/kernels/cuda/lane_gemv.cu +++ b/zig/kernels/cuda/lane_gemv.cu @@ -266,7 +266,7 @@ __device__ __forceinline__ void gemv_split(const __nv_bfloat16* __restrict__ x, } // namespace tf_lane_gemv -template __global__ void tf_lane_gemv::gemv_kernel<64, 64, 8>(const __nv_bfloat16*, const float*, +template __global__ void tf_lane_gemv::gemv_kernel<64, 64, 8>(const __nv_bfloat16* __restrict__, const float* __restrict__, const __grid_constant__ tf_lane_gemv::Part, int, int, int); extern "C" __global__ void __launch_bounds__(128) tf_lane_gemv_split(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 122f404495..b2521ae4a6 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -241,5 +241,5 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric } // namespace tf_prefill_attention // Head dim 128, eight warps, eight query heads a block, eight staging slots (Nemotron's 16 heads a KV head). -template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16*, - const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float); +template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float); diff --git a/zig/kernels/cuda/qmm_group.cu b/zig/kernels/cuda/qmm_group.cu index f3f65e184f..f9b2f9f2a7 100644 --- a/zig/kernels/cuda/qmm_group.cu +++ b/zig/kernels/cuda/qmm_group.cu @@ -338,4 +338,4 @@ __global__ void __launch_bounds__(WM * WN * 32) group_kernel( // The instantiations Nemotron's windows launch on sm_121 (tile 2: rows <= 16, bf16 out). template __global__ void tf_qmm_group::group_kernel<64, 16, 64, 1, 4, 8, false, false, false, false>( - const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int); + const __nv_bfloat16* __restrict__, const float* __restrict__, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int); diff --git a/zig/kernels/cuda/qmm_prefill.cu b/zig/kernels/cuda/qmm_prefill.cu index 1610b22704..3c61031a25 100644 --- a/zig/kernels/cuda/qmm_prefill.cu +++ b/zig/kernels/cuda/qmm_prefill.cu @@ -149,5 +149,5 @@ __global__ void __launch_bounds__(WM * WN * 32) prefill_kernel( // The instantiation prefill_matmul launches (tile 0: 128x128 on 2x4 warps, 3 stages), bf16 out. template __global__ void tf_qmm_prefill::prefill_kernel<64, 128, 128, 2, 4, 3, false>( - const __nv_bfloat16*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, + const __nv_bfloat16* __restrict__, const uint32_t* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, void* __restrict__, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/scan_rows.cu b/zig/kernels/cuda/scan_rows.cu index 569ffedce8..b3ee6784cb 100644 --- a/zig/kernels/cuda/scan_rows.cu +++ b/zig/kernels/cuda/scan_rows.cu @@ -117,7 +117,7 @@ __global__ void __launch_bounds__(TPR * ROWS) scan_kernel( } // namespace tf_scan_rows -template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16*, const __nv_bfloat16*, float*, const float*, - const float*, const float*, __nv_bfloat16*, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); -template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16*, const __nv_bfloat16*, float*, const float*, - const float*, const float*, __nv_bfloat16*, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); +template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, float* __restrict__, const float* __restrict__, + const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); +template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, float* __restrict__, const float* __restrict__, + const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); diff --git a/zig/src/cli/pull_exlr8.zig b/zig/src/cli/pull_exlr8.zig index 3e8c641f7e..c56d16ada2 100644 --- a/zig/src/cli/pull_exlr8.zig +++ b/zig/src/cli/pull_exlr8.zig @@ -88,6 +88,7 @@ pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, err_out: *std.Io.Write /// Index maps, connections and scans pass the default 256 descriptors: raise this process's own limit. fn moreFiles() void { + if (@import("builtin").os.tag == .windows) return; // no descriptor limit to raise var l = std.posix.getrlimit(.NOFILE) catch return; const want: std.posix.rlim_t = if (l.max == std.posix.RLIM.INFINITY) 4096 else @min(l.max, 4096); if (l.cur >= want) return; diff --git a/zig/src/families/nemotron/cuda_source.zig b/zig/src/families/nemotron/cuda_source.zig index ffc4ba2eec..f4a1f07fdd 100644 --- a/zig/src/families/nemotron/cuda_source.zig +++ b/zig/src/families/nemotron/cuda_source.zig @@ -2,6 +2,8 @@ const std = @import("std"); const linux = std.os.linux; +/// io_uring is Linux's; elsewhere every read runs on its own thread. +const has_ring = @import("builtin").os.tag == .linux; const cuda = @import("cuda"); const core = @import("core"); const kern = @import("cuda_kernels.zig"); @@ -42,7 +44,7 @@ pub const Source = struct { copied: [slot_count]?cuda.Event = @splat(null), // a page-locked slot's last copy: its next read waits for that, not the stream pending: [slot_count]?Pending = @splat(null), in_flight: usize = 0, - ring: ?linux.IoUring = null, // null where io_uring is unavailable: each read then runs on its own thread + ring: if (has_ring) ?linux.IoUring else ?void = null, // null where io_uring is unavailable: each read then runs on its own thread jobs: [slot_count]Job = undefined, asked: [slot_count]u64 = @splat(0), // the order reads were asked in: copies follow it without a ring next: u64 = 0, @@ -69,7 +71,7 @@ pub const Source = struct { s.pinned[made] = h; s.slots[made] = @alignCast(h.bytes); } - s.ring = linux.IoUring.init(slot_count, 0) catch |err| blk: { + if (has_ring) s.ring = linux.IoUring.init(slot_count, 0) catch |err| blk: { std.log.info("io_uring unavailable ({t}): checkpoint reads run on threads", .{err}); break :blk null; }; @@ -79,7 +81,7 @@ pub const Source = struct { /// Waits for reads still in flight, then frees the ring and slots and closes the files. pub fn deinit(s: *Source) void { while (s.in_flight > 0) s.reap(1) catch break; - if (s.ring) |*r| r.deinit() else s.joinAll(); + if (!has_ring) s.joinAll() else if (s.ring) |*r| r.deinit() else s.joinAll(); s.freeSlots(slot_count); for (s.files.items) |*m| m.file.close(); s.files.deinit(s.gpa); @@ -156,7 +158,8 @@ pub const Source = struct { const lo = std.mem.alignBackward(u64, p.offset, dio.alignment); var q = p; q.at = @intCast(p.offset - lo); - if (s.ring) |*r| { + if (has_ring and s.ring != null) { + const r = &s.ring.?; _ = try r.read(k, p.file.fd, .{ .buffer = s.slots[k][0..dio.span(p.offset, p.len)] }, lo); _ = try r.submit(); } else { @@ -192,6 +195,7 @@ pub const Source = struct { /// Finished reads, copied on; waits for at least `wait` of them. fn reap(s: *Source, wait: u32) !void { + if (!has_ring) return s.reapThreads(wait); const r = if (s.ring) |*ring| ring else return s.reapThreads(wait); var cqes: [slot_count]linux.io_uring_cqe = undefined; const n = try r.copy_cqes(&cqes, wait); diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index 024e8bcc15..4fa75167d9 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -123,8 +123,9 @@ pub const system = struct { const fd = ws.socket(@intCast(family), @intCast(kind), @intCast(proto)); return if (fd == -1) failed() else fd; } - pub fn bind(fd: socket_t, addr: *const anyopaque, len: socklen_t) c_int { - return if (ws.bind(fd, addr, len) != 0) failed() else 0; + /// `len` as std.Io.Threaded's addressToPosix gives it (its socklen_t is unsigned on Windows). + pub fn bind(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { + return if (ws.bind(fd, addr, @intCast(len)) != 0) failed() else 0; } pub fn listen(fd: socket_t, backlog: u31) c_int { return if (ws.listen(fd, backlog) != 0) failed() else 0; diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index faa25fdfdd..cf4cf9fc62 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -187,22 +187,31 @@ fn lanes(gpu: Gpu, rng: std.Random, p: *const modelopt.Packer, lin: cuda.qlinear defer y.free(); var y2 = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); defer y2.free(); - try lin.decode(s, q, .{ .x = dx.ptr }, y.ptr, 16); + // slice partials for the lane's split K: GPUs without clusters (before sm_90) reduce through them + const lw = switch (q) { + .fp8, .bf16 => |v| v, + else => unreachable, + }; + var part_bytes: usize = 4; + for ([_]usize{ 16, 128, m }) |r| part_bytes = @max(part_bytes, lin.lane.?.partBytes(r, lw)); + var part = try cuda.DeviceBuffer.alloc(gpu.d, part_bytes); + defer part.free(); + try lin.decode(s, q, .{ .x = dx.ptr, .part = part.ptr }, y.ptr, 16); try s.synchronize(); const wide = try check.download(gpu, y); defer a.free(wide); for (1..16) |rows| { try y2.fill8(0xff, s.handle); - try lin.decode(s, q, .{ .x = dx.ptr }, y2.ptr, rows); + try lin.decode(s, q, .{ .x = dx.ptr, .part = part.ptr }, y2.ptr, rows); try s.synchronize(); try same(gpu, "decode rows at a narrower width", y2, wide[0 .. rows * n * 2]); } var worst = [2]f64{ 0, 0 }; - try lin.prompt(s, q, .{ .x = dx.ptr }, y.ptr, m); + try lin.prompt(s, q, .{ .x = dx.ptr, .part = part.ptr }, y.ptr, m); try s.synchronize(); const prompt = try check.download(gpu, y); defer a.free(prompt); - try lin.prompt(s, q, .{ .x = dx.ptr }, y2.ptr, 128); + try lin.prompt(s, q, .{ .x = dx.ptr, .part = part.ptr }, y2.ptr, 128); try s.synchronize(); try same(gpu, "prompt rows in a chunk of 128", y2, prompt[0 .. 128 * n * 2]); for (0..m) |r| for (0..n) |c| { From 1f87b27911f942eaf4171c4544ae8b80cfc8fc83 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:27:13 +0000 Subject: [PATCH 03/46] Windows: pull into the hub cache, HOME, and the whole test step - pull: where Windows refuses the snapshot's symlink, the snapshot path becomes a hard link to its blob (CreateHardLinkW through `sys.c.link`): no extra space, and learning already refuses linked shards. Link failures are logged with their reason. - HOME falls back to USERPROFILE (Windows sets no HOME) for the hub cache and every libc getenv("HOME"). - Tests read optional fixture variables through `platform.getenv`; std.process.Environ.getPosix does not build on Windows. - The learned-state symlink tests skip on Windows (symlinks need privilege); `sys.c.symlink` asks for an unprivileged one. zig build test passes on x86_64-windows. --- zig/src/cli/hub.zig | 2 +- zig/src/cli/pull_file.zig | 18 ++++++++++++- zig/src/cluster/admission_test.zig | 3 ++- zig/src/cluster/cli_test.zig | 3 ++- zig/src/cluster/cost_test.zig | 3 ++- zig/src/cluster/estimate.zig | 3 ++- zig/src/cluster/model.zig | 3 ++- zig/src/cluster/plan_test.zig | 3 ++- zig/src/cluster/status.zig | 3 ++- zig/src/cluster/topology.zig | 3 ++- zig/src/core/exlr8_oracle_test.zig | 13 +++++----- zig/src/core/learned_part_test.zig | 3 +++ zig/src/platform/sys.zig | 5 ++++ zig/src/platform/win_c.zig | 41 +++++++++++++++++++++++++++++- zig/src/server/hub.zig | 2 +- 15 files changed, 90 insertions(+), 18 deletions(-) diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 64347ae50c..9b5edacb0e 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -38,7 +38,7 @@ pub fn cacheDir(a: Allocator, env: ?*const std.process.Environ.Map, override: ?[ } }.f; if (get(env, "HF_HUB_CACHE")) |c| return c; - const home = get(env, "HOME") orelse ""; + const home = get(env, "HOME") orelse get(env, "USERPROFILE") orelse ""; // Windows sets USERPROFILE, not HOME if (get(env, "HF_HOME")) |h| return std.fs.path.join(a, &.{ h, "hub" }); return std.fs.path.join(a, &.{ home, ".cache", "huggingface", "hub" }); } diff --git a/zig/src/cli/pull_file.zig b/zig/src/cli/pull_file.zig index c3de2654cf..1926f553ff 100644 --- a/zig/src/cli/pull_file.zig +++ b/zig/src/cli/pull_file.zig @@ -1,5 +1,6 @@ //! One file from the hub into the blob store: range resume, size and sha256 checks, retries, and its snapshot link. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; /// One file the hub's tree lists for the revision: LFS files carry their sha256, small files their git blob sha1. @@ -143,7 +144,22 @@ pub fn linkIntoSnapshot(a: Allocator, io: std.Io, snapshot_dir: []const u8, path if (std.fs.path.dirname(link_path)) |parent| try std.Io.Dir.cwd().createDirPath(io, parent); std.Io.Dir.cwd().deleteFile(io, link_path) catch {}; const target = try std.fs.path.join(a, &.{ "..", "..", "blobs", blob_name }); - std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch return error.SymLinkFailed; + std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch |err| { + // Windows grants symlinks only with Developer Mode or elevation (and refuses some): a hard link to the blob, + // which costs no space (huggingface_hub copies instead); learning refuses linked shards either way + if (@import("builtin").os.tag != .windows) { + std.log.err("cannot link {s} to its blob: {t}", .{ link_path, err }); + return error.SymLinkFailed; + } + const repo_dir = std.fs.path.dirname(std.fs.path.dirname(snapshot_dir) orelse return error.SymLinkFailed) orelse return error.SymLinkFailed; + const blob = try std.fs.path.joinZ(a, &.{ repo_dir, "blobs", blob_name }); + const linked = try a.dupeSentinel(u8, link_path, 0); + _ = platform.c.unlink(linked); // the refused symlink can leave an empty file behind + if (platform.c.link(blob, linked) != 0) { + std.log.err("cannot link {s} to its blob: {t}", .{ link_path, platform.c.errno(-1) }); + return error.SymLinkFailed; + } + }; } test "a hub tree path stays inside the snapshot: relative, without . or .. segments" { diff --git a/zig/src/cluster/admission_test.zig b/zig/src/cluster/admission_test.zig index 1b5b0f1b3b..200fe2b111 100644 --- a/zig/src/cluster/admission_test.zig +++ b/zig/src/cluster/admission_test.zig @@ -1,5 +1,6 @@ //! 32 to 128 streams: what fits at 8k and 32k context, prompt rows a round, and 32 toy streams joining while others decode. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const plan = @import("plan.zig"); const budget = @import("budget.zig"); @@ -115,7 +116,7 @@ test "32 toy streams join while others decode, prompt chunks in the same rounds: } test "the admission tables for CLUSTER.md (TF_CLUSTER_REPORT=1)" { - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; + if (platform.getenv("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/cli_test.zig b/zig/src/cluster/cli_test.zig index 267fca93d6..bc7b7f58bd 100644 --- a/zig/src/cluster/cli_test.zig +++ b/zig/src/cluster/cli_test.zig @@ -1,5 +1,6 @@ //! The cluster commands end to end on fake hosts: init, check, status, serve with in-process bring-up, and node. const std = @import("std"); +const platform = @import("sys"); const cli = @import("cli.zig"); const fixtures = @import("probe_fixtures.zig"); @@ -86,7 +87,7 @@ const studios = ; test "check plans K3 from its config and Hub file list and names the raise a low node needs; serve stops before launching" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/cost_test.zig b/zig/src/cluster/cost_test.zig index dec13181ca..47293a6d20 100644 --- a/zig/src/cluster/cost_test.zig +++ b/zig/src/cluster/cost_test.zig @@ -1,5 +1,6 @@ //! K3 and GLM-5.3 rounds on four 512 GiB Studios at 1 to 128 rows: dedup, caches, collectives, and the MLA split. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const model = @import("model.zig"); const estimate = @import("estimate.zig"); @@ -66,7 +67,7 @@ test "GLM-5.3 in FP8 reads far less a token than K3, and 128 streams at 32k cont } test "the cost tables for CLUSTER.md (TF_CLUSTER_REPORT=1)" { - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; + if (platform.getenv("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/estimate.zig b/zig/src/cluster/estimate.zig index 74917ba740..9d33e43313 100644 --- a/zig/src/cluster/estimate.zig +++ b/zig/src/cluster/estimate.zig @@ -1,5 +1,6 @@ //! A checkpoint's tensor table from its config and the Hub's file list, so a cluster can be planned before the weights land. const std = @import("std"); +const platform = @import("sys"); const checkpoint = @import("checkpoint.zig"); const model = @import("model.zig"); const roles = @import("roles.zig"); @@ -282,7 +283,7 @@ test "the K3 table matches every shard size to within its header" { } test "the real K3 file list (TF_K3_DIR): 96 files, 1.561 TB, every header under 0.01% of its file" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/model.zig b/zig/src/cluster/model.zig index 644b117264..91b37ea140 100644 --- a/zig/src/cluster/model.zig +++ b/zig/src/cluster/model.zig @@ -1,5 +1,6 @@ //! A model's shape from its config.json (text_config when nested): what the planner and the memory budgets need. const std = @import("std"); +const platform = @import("sys"); pub const Attention = enum(u8) { gqa, mla, linear }; @@ -257,7 +258,7 @@ test "GLM-5.3: every layer MLA with a 2,048-token indexer, 75 MoE layers, one MT } test "the real K3 config (TF_K3_DIR) parses to the literal" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/plan_test.zig b/zig/src/cluster/plan_test.zig index 666ecb3ba2..0514244c35 100644 --- a/zig/src/cluster/plan_test.zig +++ b/zig/src/cluster/plan_test.zig @@ -1,5 +1,6 @@ //! Kimi K3 on four 512 GiB M3 Ultras: the default layout, layouts from a config, determinism, and refusals. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const model = @import("model.zig"); const estimate = @import("estimate.zig"); @@ -34,7 +35,7 @@ pub fn k3(a: std.mem.Allocator) !checkpoint.Checkpoint { const size: u64 = if (i == 0) 2_341_216_112 else if (i < s.layers) (if (kda) 16_990_911_504 else 16_567_501_776) else if (i == s.layers) 4_697_664_072 else if (i == s.layers + 1) 92_289_328 else 802_448_352; f.* = .{ .name = try std.fmt.allocPrint(a, "model-{d:0>5}-of-{d:0>6}.safetensors", .{ i + 1, s.layers + 3 }), .size = size }; } - if (std.testing.environ.getPosix("TF_K3_DIR")) |dir| { + if (platform.getenv("TF_K3_DIR")) |dir| { const text = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, try std.fs.path.join(a, &.{ dir, "files.tsv" }), a, .limited(1 << 20)); files = try estimate.manifest(a, text); } diff --git a/zig/src/cluster/status.zig b/zig/src/cluster/status.zig index ecb6dd7457..49e0395dbf 100644 --- a/zig/src/cluster/status.zig +++ b/zig/src/cluster/status.zig @@ -1,5 +1,6 @@ //! What `tensorfold cluster status` and `check` print: nodes, links, memory, the placement per rank, costs and load progress. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const topology = @import("topology.zig"); const plan_mod = @import("plan.zig"); @@ -99,7 +100,7 @@ test "the K3 status on four Studios prints every rank, fits, and the round costs try placement(&out.writer, &p, needs, fs); try costs(&out.writer, &p, needs, &s, 8192); const text = try a.dupe(u8, out.written()); - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") != null) { + if (platform.getenv("TF_CLUSTER_REPORT") != null) { std.debug.print("\n{s}\n", .{text}); const others = [_]struct { []const u8, plan_mod.Options }{ .{ "experts split by tensor parallelism", .{ .layout = .{ .tensor = 4, .expert = 1 } } }, diff --git a/zig/src/cluster/topology.zig b/zig/src/cluster/topology.zig index 6e5292c8dd..6c4d2d7564 100644 --- a/zig/src/cluster/topology.zig +++ b/zig/src/cluster/topology.zig @@ -1,5 +1,6 @@ //! The cluster's wiring from every node's ports: which port reaches which node, at what rate, and the graph's shape. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); /// One cable: node `a`'s port to node `b`'s port (indices into the node list and their port lists). @@ -134,7 +135,7 @@ test "four Macs with six cables form a mesh; three cables a chain or a star; a c } test "the lent cluster's real probes (TF_CLUSTER_PROBES): four M3 Ultras in a full mesh at 80 Gb/s" { - const root = std.testing.environ.getPosix("TF_CLUSTER_PROBES") orelse return error.SkipZigTest; + const root = platform.getenv("TF_CLUSTER_PROBES") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/core/exlr8_oracle_test.zig b/zig/src/core/exlr8_oracle_test.zig index 13c854dee9..6bdb2addc2 100644 --- a/zig/src/core/exlr8_oracle_test.zig +++ b/zig/src/core/exlr8_oracle_test.zig @@ -1,6 +1,7 @@ //! EXLR8 decoding against fixtures the reference loader (MIT) wrote outside the tree: layers 0, 3, 40, 78, vocab. const std = @import("std"); +const platform = @import("sys"); const codec = @import("exlr8_codec.zig"); const ex = @import("exlr8_experts.zig"); const dense = @import("exlr8_dense.zig"); @@ -169,7 +170,7 @@ const Layers = struct { }; test "expert halves of layers 3, 40 and 78: stored checksums, the reference's tiles and W bits" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -193,7 +194,7 @@ test "expert halves of layers 3, 40 and 78: stored checksums, the reference's ti h.suh(suh[0..h.rows()]); h.svh(svh[0..h.cols()]); try checkSamples(codes, h.rows(), h.cols(), suh[0..h.rows()], svh[0..h.cols()], unit.w_samples, false, &stats); - if (std.testing.environ.getPosix("TF_EXLR8_ORACLE_FULL")) |full| try wholeHalf(full, unit, codes, h, suh[0..h.rows()], svh[0..h.cols()]); + if (platform.getenv("TF_EXLR8_ORACLE_FULL")) |full| try wholeHalf(full, unit, codes, h, suh[0..h.rows()], svh[0..h.cols()]); } stats.print("expert halves"); } @@ -233,7 +234,7 @@ fn wholeHalf(dir: []const u8, unit: Unit, codes: []const u8, h: ex.Half, suh: [] } test "every half of every stored width verifies, as the reference's --verify counts them" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); const fx = try load(arena.allocator()); @@ -260,7 +261,7 @@ test "every half of every stored width verifies, as the reference's --verify cou } test "dense and vocabulary files: records, checksums, every format bit for bit, plain tensors as shipped" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -331,7 +332,7 @@ fn denseBits(m: dense.Matrix, unit: DenseUnit) !void { } test "the eight manifests parse to the reference's widths, and every width they ask of a present layer is stored" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -371,7 +372,7 @@ test "the eight manifests parse to the reference's widths, and every width they } test "every present release file matches its published sha256 and size" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var checked: usize = 0; for ([_][]const u8{ "experts", "dense" }) |sub| { diff --git a/zig/src/core/learned_part_test.zig b/zig/src/core/learned_part_test.zig index 2f7ab53027..a2984e345d 100644 --- a/zig/src/core/learned_part_test.zig +++ b/zig/src/core/learned_part_test.zig @@ -85,6 +85,7 @@ test "orphan temporary pressure cannot evict an older valid learned state" { } test "owned-name links and directories stay untouched and foreign directories are not traversed" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "types"); @@ -123,6 +124,7 @@ test "owned-name links and directories stay untouched and foreign directories ar } test "an identity-directory symlink refuses open and preserves foreign payloads" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "identity-link"); @@ -145,6 +147,7 @@ test "an identity-directory symlink refuses open and preserves foreign payloads" } test "an explicitly configured root alias still cleans only its real identity children" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "root-target"); diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig index ee614d7507..35bb853608 100644 --- a/zig/src/platform/sys.zig +++ b/zig/src/platform/sys.zig @@ -13,6 +13,11 @@ pub const posix = if (windows) @import("win_posix.zig") else std.posix; pub const DynLib = if (windows) WinDynLib else std.DynLib; +/// An environment variable's value, through libc (std.process.Environ.getPosix does not build on Windows). +pub fn getenv(name: [*:0]const u8) ?[:0]const u8 { + return std.mem.span(c.getenv(name) orelse return null); +} + /// The CUDA driver's file name: nvcuda.dll comes with the display driver on Windows. pub const cuda_driver = if (windows) "nvcuda.dll" else "libcuda.so.1"; diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index 82b4dc7b3f..b21dce83d1 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -110,6 +110,8 @@ extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; +extern "kernel32" fn CreateHardLinkW(new: [*:0]const u16, existing: [*:0]const u16, security: ?*anyopaque) callconv(.winapi) BOOL; +extern "kernel32" fn CreateSymbolicLinkW(link: [*:0]const u16, target: [*:0]const u16, flags: u32) callconv(.winapi) u8; extern "kernel32" fn CloseHandle(h: HANDLE) callconv(.winapi) BOOL; extern "c" fn _errno() *c_int; @@ -129,7 +131,16 @@ extern "c" fn _getpid() c_int; extern "c" fn _isatty(fd: c_int) c_int; extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; -pub extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +const ucrt = struct { + extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +}; + +/// libc's getenv, except that HOME, which Windows does not set, falls back to USERPROFILE. +pub fn getenv(name: [*:0]const u8) ?[*:0]u8 { + if (ucrt.getenv(name)) |v| return v; + if (std.mem.eql(u8, std.mem.span(name), "HOME")) return ucrt.getenv("USERPROFILE"); + return null; +} pub extern "c" fn fopen(path: [*:0]const u8, mode: [*:0]const u8) ?*std.c.FILE; pub extern "c" fn fwrite(ptr: [*]const u8, size: usize, n: usize, f: *std.c.FILE) usize; pub extern "c" fn fclose(f: *std.c.FILE) c_int; @@ -422,6 +433,16 @@ pub fn rename(from: [*:0]const u8, to: [*:0]const u8) c_int { return 0; } +/// POSIX link: `new` names the same file as `old` (same volume). +pub fn link(old: [*:0]const u8, new: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wa = wide(&a, std.mem.span(old)) orelse return -1; + const wb = wide(&b, std.mem.span(new)) orelse return -1; + if (CreateHardLinkW(wb, wa, null) == 0) return fail(lastError()); + return 0; +} + pub fn access(path: [*:0]const u8, mode: c_int) c_int { var buf: Wide = undefined; return _waccess(wide(&buf, std.mem.span(path)) orelse return -1, mode & 6); @@ -575,3 +596,21 @@ pub fn fchmodat(dir: fd_t, path: [*:0]const u8, mode: mode_t, flags: u32) c_int _ = .{ dir, path, mode, flags }; return 0; } + +/// Symbolic links need Developer Mode or elevation on Windows; this asks for an unprivileged one. +pub fn symlink(target: [*:0]const u8, path: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wt = wide(&a, std.mem.span(target)) orelse return -1; + const wp = wide(&b, std.mem.span(path)) orelse return -1; + const attrs = GetFileAttributesW(wt); + const dir: u32 = if (attrs != 0xFFFF_FFFF and attrs & 0x10 != 0) 0x1 else 0; + if (CreateSymbolicLinkW(wp, wt, dir | 0x2) == 0) return fail(lastError()); // ALLOW_UNPRIVILEGED_CREATE + return 0; +} + +/// Reading a link's target needs its reparse data; nothing on Windows asks for it yet. +pub fn readlink(path: [*:0]const u8, buf: [*]u8, len: usize) isize { + _ = .{ path, buf, len }; + return fail(.NOSYS); +} diff --git a/zig/src/server/hub.zig b/zig/src/server/hub.zig index 5a155e6a1b..4e0efb1a5a 100644 --- a/zig/src/server/hub.zig +++ b/zig/src/server/hub.zig @@ -21,7 +21,7 @@ fn isDir(io: std.Io, path: []const u8) bool { /// The checkpoint directory for ``model``, or null with ``problem`` set (no downloads: the cache must hold it). pub fn resolve(a: Allocator, io: std.Io, environ: ?*const std.process.Environ.Map, model: []const u8, problem: *[]const u8) !?[]const u8 { - const home = if (environ) |e| e.get("HOME") orelse "" else ""; + const home = if (environ) |e| e.get("HOME") orelse e.get("USERPROFILE") orelse "" else ""; // Windows sets USERPROFILE const expanded = if (std.mem.startsWith(u8, model, "~/")) try std.fs.path.join(a, &.{ home, model[2..] }) else model; if (isDir(io, expanded)) return expanded; if (!isRepoId(model)) { From 996a50fc3202b0a42fc86e417b1674ec8dff366d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:36:05 +0000 Subject: [PATCH 04/46] Windows: build notes, strip before install, formatting - docs/WINDOWS.md: requirements, build, run, how the port works and how it differs from Linux; the README points to it. - -Dstrip reaches tensorfold-native before its install step is made, so the step no longer expects a .pdb the stripped build does not write. - zig fmt on the platform files. --- README.md | 2 + docs/WINDOWS.md | 53 +++++++++++++++++++ zig/build/cuda.zig | 8 +-- zig/src/platform/win_c.zig | 95 ++++++---------------------------- zig/src/platform/win_clock.zig | 82 +++++++++++++++++++++++++++++ zig/src/platform/win_posix.zig | 2 +- 6 files changed, 159 insertions(+), 83 deletions(-) create mode 100644 docs/WINDOWS.md create mode 100644 zig/src/platform/win_clock.zig diff --git a/README.md b/README.md index d77b9f9f3e..31155dae9c 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,8 @@ tensorfold serve "$HOME/models/nemotron-lightning" \ Use a complete local model directory or an existing Hugging Face cache. The binary in a release archive is `bin/tensorfold-native`; the [runbook](RUNBOOK.md) covers archive installation and CUDA runtime files. + +On Windows with an NVIDIA GPU, build `tensorfold-native.exe` from source as [docs/WINDOWS.md](docs/WINDOWS.md) describes. Use `tensorfold-native --help` or `-h`, and `tensorfold-native capabilities --json`, to inspect the installed binary. Model weights have separate downloads and licenses. diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md new file mode 100644 index 0000000000..7869f11b8d --- /dev/null +++ b/docs/WINDOWS.md @@ -0,0 +1,53 @@ +# TensorFold on Windows (native) + +An unofficial port: `tensorfold-native.exe` runs on Windows x86_64 with an NVIDIA GPU through the same CUDA engine +the Linux build uses. Only the CUDA backend is ported; Metal is macOS's and HIP is not built. + +## What you need + +- An NVIDIA GPU and a current display driver (it ships `nvcuda.dll`; no CUDA toolkit is needed to run). +- To build: [Zig 0.17.0](https://ziglang.org/download/) (the version in `.zig-version`), the CUDA Toolkit + (13.x tested) for `nvcc`, and Visual Studio Build Tools with the C++ workload (`cl.exe`, nvcc's host compiler). + +## Build + +From a "x64 Native Tools" prompt (or after running `vcvars64.bat`), so `cl.exe` is on `PATH`: + +```bat +zig build native install -Dsm=89 -Dstrip=true ^ + -Dnvcc="%CUDA_PATH%\bin\nvcc.exe" +``` + +`-Dsm` takes the GPU's compute capability without the dot, comma separated for several: 86 (RTX 30), 89 (RTX 40), +120 (RTX 50). The server is `zig-out\native\bin\tensorfold-native.exe`; `zig-out\bin\tf-cuda-test.exe` runs the GPU +suites (`tf-cuda-test info`, `smoke`, `glue`, `sample`, `quant-ref`, `modelopt`, `exlr8`). + +`zig build test` runs the host tests without a GPU. + +## Run + +```bat +tensorfold-native pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit +tensorfold-native models +tensorfold-native serve --name local-model --parallel 1 --context 8192 --temperature 0 --no-thinking +``` + +The server listens on `http://127.0.0.1:8080/v1`. Ctrl+C (or closing the console) stops it cleanly. + +## How the port works + +- `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating + system as the engine calls it. On Linux and macOS it is `std.c`, `std.posix` and `std.DynLib` unchanged; on Windows + `win_c.zig` and `win_posix.zig` give the same names over the UCRT, Win32 and Winsock. +- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. `O_DIRECT` maps to + unbuffered reads. `/tmp` is the user's temporary folder and `HOME` falls back to `USERPROFILE`. +- The Hugging Face cache links snapshot files to blobs. Windows allows symlinks only with Developer Mode or elevation, + so `pull` falls back to hard links (no extra space). +- Checkpoint reads use threads (io_uring is Linux's). +- The API key file is checked for type and size; Windows keeps access in ACLs, not mode bits, so the "unreadable by + other users" check is the folder's ACL rather than `chmod 600`. + +## Differences from Linux + +- No io_uring, `ulimit` or GB10 carveout, and no `SIGHUP` (the key file is still reread when it changes). +- `TF_HEAT_*` thermal pacing reads Linux's `/sys/class/thermal` and does nothing on Windows. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4f54eaf76d..a291e9bbec 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -139,7 +139,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; + _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true, strip); } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. @@ -161,7 +161,7 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { const m = engines(b, target, optimize, cuda, lanes, nemotron); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); @@ -177,6 +177,8 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b exe.root_module.strip = true; return exe; } + // before the install step, which decides then whether a .pdb comes with the executable + exe.root_module.strip = strip; const install = b.addInstallArtifact(exe, .{ .dest_dir = .{ .override = .{ .custom = "native/bin" } } }); b.step("native", "tensorfold-native with the CUDA engines into zig-out/native/bin").dependOn(&install.step); return exe; @@ -248,7 +250,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false, null); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index b21dce83d1..eb4d0c746a 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -100,13 +100,6 @@ extern "kernel32" fn MoveFileExW(from: [*:0]const u16, to: [*:0]const u16, flags extern "kernel32" fn FindFirstFileW(pattern: [*:0]const u16, data: *FindData) callconv(.winapi) isize; extern "kernel32" fn FindNextFileW(h: HANDLE, data: *FindData) callconv(.winapi) BOOL; extern "kernel32" fn FindClose(h: HANDLE) callconv(.winapi) BOOL; -extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; -extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; -extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; -extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; -extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; -extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; -extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; @@ -127,10 +120,8 @@ extern "c" fn _wunlink(path: [*:0]const u16) c_int; extern "c" fn _wrmdir(path: [*:0]const u16) c_int; extern "c" fn _wmkdir(path: [*:0]const u16) c_int; extern "c" fn _waccess(path: [*:0]const u16, mode: c_int) c_int; -extern "c" fn _getpid() c_int; extern "c" fn _isatty(fd: c_int) c_int; extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; -extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; const ucrt = struct { extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; }; @@ -150,8 +141,8 @@ pub extern "c" fn malloc(n: usize) ?*anyopaque; /// The UCRT's struct _stat64. const Stat64 = extern struct { dev: u32, ino: u16, mode: u16, nlink: i16, uid: i16, gid: i16, rdev: u32, size: i64, atime: i64, mtime: i64, ctime: i64 }; -fn fail(e: E) c_int { - _errno().* = @intFromEnum(e); +pub fn fail(e: E) c_int { + _errno().* = @backingInt(e); return -1; } @@ -272,10 +263,10 @@ fn openFile(path: [*:0]const u8, flags: O) c_int { .RDWR => generic_read | generic_write, }; const disposition: u32 = if (flags.CREAT and flags.EXCL) 1 // CREATE_NEW - else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS - else if (flags.CREAT) 4 // OPEN_ALWAYS - else if (flags.TRUNC) 5 // TRUNCATE_EXISTING - else 3; // OPEN_EXISTING + else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS + else if (flags.CREAT) 4 // OPEN_ALWAYS + else if (flags.TRUNC) 5 // TRUNCATE_EXISTING + else 3; // OPEN_EXISTING var attrs: u32 = 0x80; // FILE_ATTRIBUTE_NORMAL if (flags.DIRECT) attrs |= 0x2000_0000; // FILE_FLAG_NO_BUFFERING const h = CreateFileW(w, rights, 0x7, null, disposition, attrs, null); // share read, write and delete @@ -368,7 +359,7 @@ pub fn fsync(fd: fd_t) c_int { pub fn ftruncate(fd: fd_t, size: off_t) c_int { const rc = _chsize_s(fd, size); - if (rc != 0) return fail(@enumFromInt(rc)); + if (rc != 0) return fail(@fromBackingInt(@intCast(rc))); return 0; } @@ -517,70 +508,16 @@ pub fn closedir(d: *DIR) c_int { return 0; } -// ---- process, environment, time ---- +// ---- process, environment, time: win_clock.zig ---- -pub fn getpid() pid_t { - return _getpid(); -} - -pub fn getuid() u32 { - return 0; -} - -pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { - if (overwrite == 0 and getenv(name) != null) return 0; - if (_putenv_s(name, value) != 0) return fail(.INVAL); - return 0; -} - -pub fn unsetenv(name: [*:0]const u8) c_int { - if (_putenv_s(name, "") != 0) return fail(.INVAL); - return 0; -} - -threadlocal var sleep_timer: ?HANDLE = null; - -/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. -pub fn nanosleep(req: *const timespec, rem: ?*timespec) c_int { - if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; - const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; - if (ns <= 0) return 0; - if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS - if (sleep_timer) |t| { - const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units - if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { - _ = WaitForSingleObject(t, 0xFFFF_FFFF); - return 0; - } - } - Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); - return 0; -} - -pub fn clock_gettime(clock: c_int, out: *timespec) c_int { - if (clock == CLOCK.MONOTONIC) { - var count: i64 = 0; - var freq: i64 = 0; - _ = QueryPerformanceCounter(&count); - _ = QueryPerformanceFrequency(&freq); - out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; - return 0; - } - var ft: FILETIME = undefined; - GetSystemTimePreciseAsFileTime(&ft); - const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); - const unix = since_1601 - 116_444_736_000_000_000; - out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; - return 0; -} - -pub fn gettimeofday(tv: *timeval, tz: ?*anyopaque) c_int { - _ = tz; - var ts: timespec = undefined; - _ = clock_gettime(CLOCK.REALTIME, &ts); - tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; - return 0; -} +const clock = @import("win_clock.zig"); +pub const getpid = clock.getpid; +pub const getuid = clock.getuid; +pub const setenv = clock.setenv; +pub const unsetenv = clock.unsetenv; +pub const nanosleep = clock.nanosleep; +pub const clock_gettime = clock.clock_gettime; +pub const gettimeofday = clock.gettimeofday; /// Bytes free to this user under `dir` (statvfs's bavail * frsize). pub fn diskAvailable(dir: [*:0]const u8) ?u64 { diff --git a/zig/src/platform/win_clock.zig b/zig/src/platform/win_clock.zig new file mode 100644 index 0000000000..ec41325cfc --- /dev/null +++ b/zig/src/platform/win_clock.zig @@ -0,0 +1,82 @@ +//! The process, environment and clock part of the libc surface on Windows (win_c.zig re-exports these): the process +//! id, environment writes over the UCRT, and nanosleep and clock_gettime over the kernel's timers and counters. + +const std = @import("std"); +const c = @import("win_c.zig"); + +const HANDLE = *anyopaque; +const BOOL = c_int; +const FILETIME = extern struct { low: u32, high: u32 }; + +extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; +extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; +extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; +extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; +extern "c" fn _getpid() c_int; +extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; + +pub fn getpid() c.pid_t { + return _getpid(); +} + +pub fn getuid() u32 { + return 0; +} + +pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { + if (overwrite == 0 and c.getenv(name) != null) return 0; + if (_putenv_s(name, value) != 0) return c.fail(.INVAL); + return 0; +} + +pub fn unsetenv(name: [*:0]const u8) c_int { + if (_putenv_s(name, "") != 0) return c.fail(.INVAL); + return 0; +} + +threadlocal var sleep_timer: ?HANDLE = null; + +/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. +pub fn nanosleep(req: *const c.timespec, rem: ?*c.timespec) c_int { + if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; + const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; + if (ns <= 0) return 0; + if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS + if (sleep_timer) |t| { + const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units + if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { + _ = WaitForSingleObject(t, 0xFFFF_FFFF); + return 0; + } + } + Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); + return 0; +} + +pub fn clock_gettime(clock: c_int, out: *c.timespec) c_int { + if (clock == c.CLOCK.MONOTONIC) { + var count: i64 = 0; + var freq: i64 = 0; + _ = QueryPerformanceCounter(&count); + _ = QueryPerformanceFrequency(&freq); + out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; + return 0; + } + var ft: FILETIME = undefined; + GetSystemTimePreciseAsFileTime(&ft); + const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); + const unix = since_1601 - 116_444_736_000_000_000; + out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; + return 0; +} + +pub fn gettimeofday(tv: *c.timeval, tz: ?*anyopaque) c_int { + _ = tz; + var ts: c.timespec = undefined; + _ = clock_gettime(c.CLOCK.REALTIME, &ts); + tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; + return 0; +} diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index 4fa75167d9..04c82d4dba 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -109,7 +109,7 @@ fn failed() c_int { 10061 => .CONNREFUSED, else => .IO, }; - _errno().* = @intFromEnum(e); + _errno().* = @backingInt(e); return -1; } From aac93d22e73c9df9ae6ac11b54533e4bad57d158 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:45:14 +0000 Subject: [PATCH 05/46] Default port 8888 8080 and 8000 are taken by too many other servers; serve listens on 8888 unless --port says otherwise, and the docs and bench scripts follow. --- CONTRIBUTING.md | 4 ++-- README.md | 10 +++++----- RUNBOOK.md | 8 ++++---- docs/WINDOWS.md | 2 +- docs/api.md | 4 ++-- docs/cache-benchmarks.md | 4 ++-- docs/control.md | 2 +- docs/cortheon.md | 4 ++-- docs/living-weights.md | 14 +++++++------- docs/recipes/README.md | 2 +- docs/recipes/qwen3.8-27b.md | 4 ++-- docs/recipes/qwen3.8-flash-next.md | 4 ++-- docs/speed-up-mode.md | 2 +- docs/vision.md | 2 +- tools/bench_concurrent.py | 2 +- tools/bench_openai.py | 2 +- zig/src/server/cli.zig | 2 +- 17 files changed, 36 insertions(+), 36 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5a9c4d1014..0630be4816 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -54,9 +54,9 @@ Record pass, fail and skip counts and explain any unrun platform or feature. For served decode, the existing standard-library client needs Python 3.11 or newer: ```sh -python3 tools/bench_openai.py http://127.0.0.1:8080 local-model \ +python3 tools/bench_openai.py http://127.0.0.1:8888 local-model \ --tokens 256 --reps 3 --temperatures 0 --label candidate --output decode.json -python3 tools/bench_concurrent.py http://127.0.0.1:8080 local-model \ +python3 tools/bench_concurrent.py http://127.0.0.1:8888 local-model \ --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial --output concurrent.json ``` diff --git a/README.md b/README.md index 31155dae9c..a6a018278f 100644 --- a/README.md +++ b/README.md @@ -28,12 +28,12 @@ On Windows with an NVIDIA GPU, build `tensorfold-native.exe` from source as [doc Use `tensorfold-native --help` or `-h`, and `tensorfold-native capabilities --json`, to inspect the installed binary. Model weights have separate downloads and licenses. -The API listens at `http://127.0.0.1:8080/v1` by default. +The API listens at `http://127.0.0.1:8888/v1` by default. It serves OpenAI chat completions, completions and Responses, plus Anthropic Messages. Use the model ID returned by `/v1/models`, or set one with `--name`. ```sh -curl -fsS http://127.0.0.1:8080/v1/chat/completions \ +curl -fsS http://127.0.0.1:8888/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` @@ -49,11 +49,11 @@ tensorfold pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit src=$(ls -d ~/.cache/huggingface/hub/models--TensorFold--NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit/snapshots/* | head -1) mkdir -p ~/models && cp -cRL "$src" ~/models/nemotron-living # always teach a copy tensorfold serve ~/models/nemotron-living --slide -curl -N http://127.0.0.1:8080/v1/slide/learn -H 'Content-Type: application/json' -d '{"text": "I like blue."}' +curl -N http://127.0.0.1:8888/v1/slide/learn -H 'Content-Type: application/json' -d '{"text": "I like blue."}' ``` Send it a document instead and it finds the facts in it and learns them one by one. Open -`http://127.0.0.1:8080/slide` to watch in 3D: the facts, the neurons each one took in the model's layers, the shard +`http://127.0.0.1:8888/slide` to watch in 3D: the facts, the neurons each one took in the model's layers, the shard files that hold them, and the lesson running now. Learning runs on Apple-silicon Macs and on an NVIDIA GB10 (DGX Spark), and a folder that learned on one serves on the @@ -152,7 +152,7 @@ The binary's `capabilities --json` response lists its supported flags and platfo | Flag | Meaning | | --- | --- | -| `--host HOST`, `--port PORT` | Listen address and port, default `127.0.0.1:8080`. | +| `--host HOST`, `--port PORT` | Listen address and port, default `127.0.0.1:8888`. | | `--name NAME` | Model ID advertised to clients. | | `--alias NAME` | Another accepted model ID; repeat for multiple aliases. | | `--api-key KEY` | Require a bearer key; repeat for multiple keys. | diff --git a/RUNBOOK.md b/RUNBOOK.md index 408f2d91e8..0d0282fb49 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -104,14 +104,14 @@ bin/tensorfold-native serve "$HOME/models/nemotron-lightning" \ Leave the server running and check it from another terminal: ```sh -curl -fsS http://127.0.0.1:8080/health -curl -fsS http://127.0.0.1:8080/v1/models -curl -fsS http://127.0.0.1:8080/v1/chat/completions \ +curl -fsS http://127.0.0.1:8888/health +curl -fsS http://127.0.0.1:8888/v1/models +curl -fsS http://127.0.0.1:8888/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` -Use `http://127.0.0.1:8080/v1` as the OpenAI client base URL. +Use `http://127.0.0.1:8888/v1` as the OpenAI client base URL. For streaming, add `"stream": true` to the request. The native server also has `/v1/responses` and Anthropic `/v1/messages` routes. Serve flags and request overrides are described in [README.md](README.md#serve-flags). diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index 7869f11b8d..fdb9b7f88f 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -32,7 +32,7 @@ tensorfold-native models tensorfold-native serve --name local-model --parallel 1 --context 8192 --temperature 0 --no-thinking ``` -The server listens on `http://127.0.0.1:8080/v1`. Ctrl+C (or closing the console) stops it cleanly. +The server listens on `http://127.0.0.1:8888/v1`. Ctrl+C (or closing the console) stops it cleanly. ## How the port works diff --git a/docs/api.md b/docs/api.md index a0563b0048..657eca442a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1,6 +1,6 @@ # Compatible APIs -The base URL is `http://127.0.0.1:8080/v1` with the default server settings. +The base URL is `http://127.0.0.1:8888/v1` with the default server settings. ## Native chat template override @@ -480,5 +480,5 @@ redacted thinking and other context edits return HTTP 400. Connect Claude Code using the ID from `/v1/models`: ```bash -ANTHROPIC_BASE_URL=http://127.0.0.1:8080 ANTHROPIC_API_KEY=local claude --model local-model +ANTHROPIC_BASE_URL=http://127.0.0.1:8888 ANTHROPIC_API_KEY=local claude --model local-model ``` diff --git a/docs/cache-benchmarks.md b/docs/cache-benchmarks.md index 2c9b9b6f95..e3977293f4 100644 --- a/docs/cache-benchmarks.md +++ b/docs/cache-benchmarks.md @@ -23,7 +23,7 @@ The standard-library-only cache client runs one first request, identical repeate concurrent group using the same prompt, greedy temperature, seed and reply limit. There is no warm-up: ```sh -python3 tools/bench_cache.py http://127.0.0.1:8080 MODEL \ +python3 tools/bench_cache.py http://127.0.0.1:8888 MODEL \ --tokens 64 --repeats 2 --streams 4 --require-cold --output cache.json ``` @@ -36,7 +36,7 @@ prefix retention (`--prompt-cache-gib 0`, where supported) for independent cold- For exact 2K/8K/32K/64K raw prefill, supply a JSON array of IDs from the checkpoint's own tokenizer: ```sh -python3 tools/bench_cache.py http://127.0.0.1:8080 MODEL \ +python3 tools/bench_cache.py http://127.0.0.1:8888 MODEL \ --prompt-file public-8192-tokens.json --expected-prompt-tokens 8192 \ --tokens 64 --require-cold --output cache-8192.json ``` diff --git a/docs/control.md b/docs/control.md index 4ed791242f..eb5e849364 100644 --- a/docs/control.md +++ b/docs/control.md @@ -118,7 +118,7 @@ The error from `tensorfold tui` prints that command with the venv's Python. `ric ```bash tensorfold tui tensorfold tui --profile code -tensorfold tui --url http://127.0.0.1:8080/v1 +tensorfold tui --url http://127.0.0.1:8888/v1 tensorfold tui --url https://spark.example.invalid --token-env TENSORFOLD_API_KEY tensorfold tui --demo tensorfold tui --demo --snapshot preview.svg diff --git a/docs/cortheon.md b/docs/cortheon.md index d8bee40e3e..bf6b242be1 100644 --- a/docs/cortheon.md +++ b/docs/cortheon.md @@ -100,8 +100,8 @@ tensorfold serve TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit --cor ``` ```bash -curl -s http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \ +curl -s http://127.0.0.1:8888/v1/chat/completions -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What changed in the latest Zig release?"}]}' ``` -Point any OpenAI-compatible client at `http://127.0.0.1:8080/v1` as before. +Point any OpenAI-compatible client at `http://127.0.0.1:8888/v1` as before. diff --git a/docs/living-weights.md b/docs/living-weights.md index 9fd99bdb83..b7004b1e6c 100644 --- a/docs/living-weights.md +++ b/docs/living-weights.md @@ -61,12 +61,12 @@ tensorfold serve ~/models/nemotron-living --slide --slide-graph ~/models/nemotro `--slide` turns learning on. `--slide-graph` keeps the list of learned facts in a file, so it survives a restart; without it the list lives in memory, though the weights keep what they learned either way. The server listens on -`http://127.0.0.1:8080` unless you pass `--host` or `--port`. +`http://127.0.0.1:8888` unless you pass `--host` or `--port`. ### 4. Teach it a fact ```sh -curl -N http://127.0.0.1:8080/v1/slide/learn \ +curl -N http://127.0.0.1:8888/v1/slide/learn \ -H 'Content-Type: application/json' \ -d '{"text": "I like blue."}' ``` @@ -82,7 +82,7 @@ answer. A fact takes 1.5 to 3 minutes. Use the normal chat API, from any client: ```sh -curl http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \ +curl http://127.0.0.1:8888/v1/chat/completions -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What colour do I like?"}]}' ``` @@ -112,7 +112,7 @@ nothing behind. An earlier fact's neuron that the new lesson would disturb is re - Several facts at once: put them in one text, as in `{"text": "I like blue. My sister is called Ana."}`. Each is learned in turn. - A file: send its contents as the text, for example - `jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8080/v1/slide/learn -H 'Content-Type: application/json' -d @-`. + `jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8888/v1/slide/learn -H 'Content-Type: application/json' -d @-`. - A web page: fetch its text first, or use omp's `/learn ` below, which does it for you. The server keeps answering chat requests while it learns: learning runs when no reply is being generated. Short, @@ -136,7 +136,7 @@ A panel shows each fact as queued, learning, learned or missed. ## See what it learned -Open `http://127.0.0.1:8080/slide` in a browser while the server runs. It draws the model in 3D and updates live: +Open `http://127.0.0.1:8888/slide` in a browser while the server runs. It draws the model in 3D and updates live: - The model as a tower of its 52 layers. Each MoE layer is a ring carrying its 64 pocket neurons, and a neuron glows in the colour of the fact that owns it. @@ -150,9 +150,9 @@ Drag to turn, scroll or pinch to zoom, and click a fact for its neurons, the que rounds. Phones get a bottom sheet. Without WebGL2, or with `/slide?flat`, the page draws a flat graph from the same data. -- `curl http://127.0.0.1:8080/v1/slide/graph` returns the same data as JSON: each fact's state, neurons (layer, index +- `curl http://127.0.0.1:8888/v1/slide/graph` returns the same data as JSON: each fact's state, neurons (layer, index and key), questions and rounds, and the model's layers, pocket and shards. -- `curl -X DELETE http://127.0.0.1:8080/v1/slide/graph` clears the graph. The weights keep what they learned. +- `curl -X DELETE http://127.0.0.1:8888/v1/slide/graph` clears the graph. The weights keep what they learned. ## Keep, move or reset what it learned diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 1e38e07605..28698ef465 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -61,7 +61,7 @@ how matrix multiplication uses a GPU. The chat fixture disables thinking. After `--name bench`, run from the repository root: ```bash -python3 tools/bench_openai.py http://127.0.0.1:8080 bench \ +python3 tools/bench_openai.py http://127.0.0.1:8888 bench \ --tokens 64 --reps 5 --temperatures 1.0,0 --output bench.json ``` diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index d4084dcfc1..6027d7d01a 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -90,7 +90,7 @@ input is qualified on this checkpoint. ```bash tensorfold pull nvidia/Qwen3.8-27B-NVFP4 z-lab/Qwen3.8-27B-DFlash2 -tensorfold serve nvidia/Qwen3.8-27B-NVFP4 --host 0.0.0.0 --port 8080 +tensorfold serve nvidia/Qwen3.8-27B-NVFP4 --host 0.0.0.0 --port 8888 ``` Prompts take bf16 rows on a prompt GEMM that reads the stored bytes: each FP8 byte and each NVFP4 code times its @@ -122,7 +122,7 @@ serve the folder: ```bash python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-27B-exl3', revision='3.00bpw', local_dir='qwen27b-exl3-3.00bpw')" tensorfold pull z-lab/Qwen3.8-27B-DFlash2 -tensorfold serve qwen27b-exl3-3.00bpw --host 0.0.0.0 --port 8080 +tensorfold serve qwen27b-exl3-3.00bpw --host 0.0.0.0 --port 8888 ``` Verify windows use the row-invariant EXL3 linear, so drafted replies equal `"draft": false` ones; prompts use the diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index af65c4fccd..bed3056e40 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -98,7 +98,7 @@ local-inference-lab's. ship, MTP head included, on one GPU: ```bash -tensorfold serve local-inference-lab/Qwen3.8-Flash-Next-NVFP4 --host 0.0.0.0 --port 8080 +tensorfold serve local-inference-lab/Qwen3.8-Flash-Next-NVFP4 --host 0.0.0.0 --port 8888 ``` | Checkpoint (revision) | Routed experts | DeltaNet, attention, shared expert | n-gram table | @@ -169,7 +169,7 @@ Download that branch, then serve the folder: ```bash python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-Flash-Next-exl3', revision='3.05bpw_h5_ng5', local_dir='flashnext-exl3-3.05bpw')" python -m tensorfold.cuda.exl3.inspect flashnext-exl3-3.05bpw # bits per tensor, from the headers -tensorfold serve flashnext-exl3-3.05bpw --host 0.0.0.0 --port 8080 +tensorfold serve flashnext-exl3-3.05bpw --host 0.0.0.0 --port 8888 ``` A pack maps its n-gram table from its own file and runs on one GPU: `--tp 2` and `--ple-on-ssd` are for the MLX diff --git a/docs/speed-up-mode.md b/docs/speed-up-mode.md index 69803a8df7..f9ad574d48 100644 --- a/docs/speed-up-mode.md +++ b/docs/speed-up-mode.md @@ -119,7 +119,7 @@ Flags that change decoding, such as `--no-drafts` or `--no-thinking`, go on both ### 6. Connect a client -Point any OpenAI-compatible client at `http://RANK0:8080/v1` with the model name `flash-next`, where RANK0 is rank 0's +Point any OpenAI-compatible client at `http://RANK0:8888/v1` with the model name `flash-next`, where RANK0 is rank 0's address. Send requests only to rank 0, since rank 1 refuses its own. ## Measured speeds diff --git a/docs/vision.md b/docs/vision.md index 9536f7e7fc..486c283663 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -64,7 +64,7 @@ body = { 'max_tokens': 128, 'stream': False, } -request = Request('http://127.0.0.1:8080/v1/chat/completions', +request = Request('http://127.0.0.1:8888/v1/chat/completions', data=json.dumps(body).encode(), headers={'Content-Type': 'application/json'}) with urlopen(request) as response: diff --git a/tools/bench_concurrent.py b/tools/bench_concurrent.py index 446592072e..c1c19eb591 100644 --- a/tools/bench_concurrent.py +++ b/tools/bench_concurrent.py @@ -16,7 +16,7 @@ server with a small listen backlog resets simultaneous connections; those resets are reported, not hidden). Standard library only. - python3 tools/bench_concurrent.py http://127.0.0.1:8080 MODEL --levels 1,2,4,8 --alone --output out.json + python3 tools/bench_concurrent.py http://127.0.0.1:8888 MODEL --levels 1,2,4,8 --alone --output out.json """ import argparse diff --git a/tools/bench_openai.py b/tools/bench_openai.py index fca9bf1067..d18685df6b 100644 --- a/tools/bench_openai.py +++ b/tools/bench_openai.py @@ -4,7 +4,7 @@ CUDA engine's bench times (the first sampled token excluded). Standard library only, so it runs on a bare host. - python3 tools/bench_openai.py http://127.0.0.1:8080 MODEL --tokens 64 --reps 5 --output out.json + python3 tools/bench_openai.py http://127.0.0.1:8888 MODEL --tokens 64 --reps 5 --output out.json """ import argparse diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..3ad95bf2fc 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -122,7 +122,7 @@ pub const env = [_][]const u8{ "TENSORFOLD_API_KEY", "TENSORFOLD_NO_LIVE", "TENS pub const Args = struct { model: []const u8 = "", host: []const u8 = "127.0.0.1", - port: u16 = 8080, + port: u16 = 8888, name: []const u8 = "", chat_template: ?[]const u8 = null, alias: []const []const u8 = &.{}, From 81f5576ca473652f712de2247f305f02985fce01 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:34:37 +0000 Subject: [PATCH 06/46] qwen3_5: serve Qwen3.5-2B on CUDA The MLX affine 4-bit checkpoint on the native CUDA server, beside Nemotron: qwen35.cu holds the Metal engine's glue at its bf16 rounding points (gated delta rule, causal convolution, gated attention), cuda.qlinear's affine-4 kernels do the projections, sample.cu draws on the device. Windows of up to 16 rows record each row's recurrence and convolution tail, so the lane core shares rounds across streams and keeps partial copies; a window's rows equal one-row steps bit for bit. draw.zig is the host draw both backends now share. On an RTX 4090: 210 tok/s for one stream, 1250 tok/s across eight; logits agree with an fp32 reference at every position of a 36-token check. --- zig/build/cuda.zig | 32 +- zig/kernels/cuda/qwen35.cu | 458 +++++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/backend.zig | 107 +---- zig/src/families/qwen3_5/cuda.zig | 9 + zig/src/families/qwen3_5/cuda_engine.zig | 484 +++++++++++++++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 211 ++++++++++ zig/src/families/qwen3_5/draw.zig | 108 +++++ zig/src/native/cuda.zig | 18 +- zig/tests/qwen35_cuda_main.zig | 58 +++ 10 files changed, 1367 insertions(+), 119 deletions(-) create mode 100644 zig/kernels/cuda/qwen35.cu create mode 100644 zig/src/families/qwen3_5/cuda.zig create mode 100644 zig/src/families/qwen3_5/cuda_engine.zig create mode 100644 zig/src/families/qwen3_5/cuda_native.zig create mode 100644 zig/src/families/qwen3_5/draw.zig create mode 100644 zig/tests/qwen35_cuda_main.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index a291e9bbec..8632ea6c7f 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -32,6 +32,7 @@ const kernels = [_]Kernel{ .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) + .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: Qwen3.5-2B's glue, the Metal kernels' rounding points .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, @@ -78,7 +79,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); @@ -90,7 +91,11 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("lanes", lanes); nemotron.addImport("heat", heat); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat }; + const qwen35 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + qwen35.addImport("cuda", cuda); + qwen35.addImport("core", core); + qwen35.addImport("lanes", lanes); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .qwen35 = qwen35, .tokenizer = tokenizer, .heat = heat }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -122,7 +127,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu cli.addImport("core", mods.core); cli.addImport("lanes", mods.lanes); cli.addImport("nemotron", mods.nemotron); - const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron); + const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.qwen35); cli.addOptions("build_options", build_options); cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines, mods.core)); b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); @@ -139,18 +144,23 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true, strip); + _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.tokenizer, build_options, true, strip); + // tf-qwen35-cuda: Qwen3.5-2B's CUDA forward, its logits written for a reference to compare + const q35 = b.createModule(.{ .root_source_file = b.path("zig/tests/qwen35_cuda_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); + q35.addImport("cuda", cuda); + q35.addImport("qwen35", mods.qwen35); + b.installArtifact(b.addExecutable(.{ .name = "tf-qwen35-cuda", .root_module = q35 })); } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. -fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { +fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } }, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron }, .{ .name = "qwen35", .module = qwen35 } }, }); return .{ .api = api, .engines = mod }; } @@ -161,8 +171,8 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { - const m = engines(b, target, optimize, cuda, lanes, nemotron); +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { + const m = engines(b, target, optimize, cuda, lanes, nemotron, qwen35); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const exe = b.addExecutable(.{ .name = "tensorfold-native", .root_module = b.createModule(.{ @@ -192,7 +202,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); - const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron); + const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.qwen35); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = host, .optimize = .debug, .link_libc = true }); @@ -212,7 +222,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); @@ -250,7 +260,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false, null); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.tokenizer, build_options, false, null); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu new file mode 100644 index 0000000000..a68fa30990 --- /dev/null +++ b/zig/kernels/cuda/qwen35.cu @@ -0,0 +1,458 @@ +// Qwen3.5-2B's glue on CUDA: the Metal engine's kernels (zig/kernels/metal/qwen3_5) with the same bf16 rounding points. +// Rows of one forward belong to segments (one stream's window each); a segment's state lives in its stream's blocks. + +#include +#include + +namespace tf_q35 { + +constexpr int HID = 2048; // hidden +constexpr int NH = 16; // linear-attention heads (keys and values) +constexpr int DK = 128; // their head width +constexpr int CONV = 6144; // q, k and v channels the convolution sees +constexpr int TAPS = 4; +constexpr int QH = 8; // attention query heads +constexpr int KVH = 2; // attention key/value heads +constexpr int HD = 256; // attention head width +constexpr size_t REC = (size_t)NH * DK * DK; // floats of one layer's recurrence +constexpr size_t CONV_ROWS = (size_t)(TAPS - 1) * CONV; + +// One stream's window in this forward: its state blocks, cache length before the window, and rows. +struct Seg { + unsigned long long rec; // float recurrences, one REC block a linear layer + unsigned long long conv; // bf16 convolution tails, CONV_ROWS a linear layer + unsigned long long kv; // bf16 keys then values, KVH * cap * HD each, an attention layer at a time + int cap, len, row0, rows; +}; + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +__device__ __forceinline__ float sigmoid(float x) { return 1.0f / (1.0f + __expf(-x)); } + +} // namespace tf_q35 + +using namespace tf_q35; + +// MLX affine 4-bit token rows: out = bf16(s * q + b). +extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, + const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, + __nv_bfloat16* __restrict__ out) { + const size_t id = ids[blockIdx.x]; + for (int b = threadIdx.x; b < HID / 2; b += blockDim.x) { + const uint8_t v = W[id * (HID / 2) + b]; + const size_t g = id * (HID / 64) + (2 * b) / 64; + const float s = bf(S[g]), z = bf(B[g]); + out[(size_t)blockIdx.x * HID + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); + out[(size_t)blockIdx.x * HID + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); + } +} + +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), xs = x's 64-group sums. One 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, + float* __restrict__ XS, float eps) { + constexpr int E = HID / 256; + __shared__ float red[8]; + __shared__ float xb[HID]; + const size_t row = blockIdx.x; + const int t = threadIdx.x; + float hv[E]; + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < E; i++) { + const size_t at = row * HID + t * E + i; + float h = bf(H[at]); + if (R != nullptr) { + const __nv_bfloat16 hb = __float2bfloat16_rn(h + bf(R[at])); + H[at] = hb; + h = bf(hb); + } + hv[i] = h; + ss = fmaf(h, h, ss); + } + ss = warp_sum(ss); + if ((t & 31) == 0) red[t >> 5] = ss; + __syncthreads(); + float total = 0.0f; + for (int i = 0; i < 8; i++) total += red[i]; + const float inv = rsqrtf(total / (float)HID + eps); +#pragma unroll + for (int i = 0; i < E; i++) { + const int d = t * E + i; + const __nv_bfloat16 x = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); + X[row * HID + d] = x; + xb[d] = bf(x); + } + __syncthreads(); + if (XS != nullptr && t < HID / 64) { + float s = 0.0f; + for (int j = 0; j < 64; j++) s += xb[t * 64 + j]; + XS[row * (HID / 64) + t] = s; + } +} + +// The 64-group sums of bf16 rows D wide (the decode projections' input sums). Grid (rows), a thread a group. +extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D) { + const size_t row = blockIdx.x; + for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { + const __nv_bfloat16* x = X + row * D + g * 64; + float s = 0.0f; + for (int j = 0; j < 64; j++) s += bf(x[j]); + XS[row * (D / 64) + g] = s; + } +} + +// The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. +// Grid (1, 3 * NH, rows), a warp a (row, head): heads [0, NH) are q, [NH, 2NH) k, [2NH, 3NH) v. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer, const __nv_bfloat16* __restrict__ CW, + const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, + const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, + __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA) { + const int lane = threadIdx.x; + const int head = blockIdx.y; + const int w = blockIdx.z; + const Seg sg = segs[rowseg[w]]; + const int r = w - sg.row0; + const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * CONV_ROWS; + const bool isq = head < NH, isk = !isq && head < 2 * NH; + const int c0 = head * DK; // q, k and v channels are consecutive heads of DK + constexpr int PER = DK / 32; + float vals[PER]; +#pragma unroll + for (int j = 0; j < PER; j++) { + const int c = c0 + lane * PER + j; + float acc = 0.0f; + for (int tap = 0; tap < TAPS; tap++) { + const int at = r + tap; // in [state rows (TAPS - 1)] ++ [this segment's rows] + const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * CONV + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]); + acc += bf(CW[c * TAPS + tap]) * xv; + } + const float conv = rbf(acc); + const float sig = rbf(1.0f / (1.0f + expf(-conv))); + vals[j] = rbf(conv * sig); + } + if (isq || isk) { + float ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) ss += vals[j] * vals[j]; + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + (1e-6f / (float)DK)); + const float scale = isq ? rbf(1.0f / (float)DK) : rbf(rsqrtf((float)DK)); + __nv_bfloat16* out = isq ? Q + ((size_t)w * NH + head) * DK : K + ((size_t)w * NH + head - NH) * DK; +#pragma unroll + for (int j = 0; j < PER; j++) out[lane * PER + j] = __float2bfloat16_rn(scale * rbf(vals[j] * inv)); + } else { + const int hv = head - 2 * NH; +#pragma unroll + for (int j = 0; j < PER; j++) V[((size_t)w * NH + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); + if (lane == 0) { + const float s = rbf(bf(Ain[w * NH + hv]) + bf(DT[hv])); + const float sp = rbf(fmaxf(s, 0.0f) + logf(1.0f + expf(-fabsf(s)))); + G[w * NH + hv] = expf(-expf(ALOG[hv]) * sp); + BETA[w * NH + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NH + hv])))); + } + } +} + +// Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). +// SNAP, if given, also takes the tail after each row (forward row m at m * CONV_ROWS), read before the update. +extern "C" __global__ void tf_q35_conv_commit(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, + __nv_bfloat16* __restrict__ SNAP) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= CONV) return; + const Seg sg = segs[blockIdx.y]; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * CONV_ROWS; + if (SNAP != nullptr) + for (int j = 0; j < sg.rows; j++) + for (int r = 0; r < TAPS - 1; r++) { + const int at = r + j + 1; + SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * CONV + c] = + at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + } + // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row + for (int r = 0; r < TAPS - 1; r++) { + const int at = r + sg.rows; + CS[(size_t)r * CONV + c] = at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + } +} + +// The gated delta rule over each segment's rows, its recurrence updated in place; SNAP, if given, takes the state after +// each row (forward row m at m * REC). Grid (1, DK / 4, NH * segments), block (32, 4): a warp a (head, value row), +// its 32 lanes over the key width. +extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, + const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, + __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP) { + const int hv = blockIdx.z % NH; + const Seg sg = segs[blockIdx.z / NH]; + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * REC + ((size_t)hv * DK + dv) * DK + PT * dk; + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int r = 0; r < sg.rows; r++) { + const size_t row = sg.row0 + r; + float kc[PT], qc[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) { + kc[i] = bf(k[(row * NH + hv) * DK + PT * dk + i]); + qc[i] = bf(q[(row * NH + hv) * DK + PT * dk + i]); + } + const float vc = bf(v[(row * NH + hv) * DK + dv]); + const float gc = g[row * NH + hv]; + const float bc = bf(beta[row * NH + hv]); + float kv_mem = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * gc; + kv_mem += state[i] * kc[i]; + } + kv_mem = warp_sum(kv_mem); + const float delta = (vc - kv_mem) * bc; + float out = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] + kc[i] * delta; + out += state[i] * qc[i]; + } + out = warp_sum(out); + if (dk == 0) y[(row * NH + hv) * DK + dv] = __float2bfloat16_rn(out); + if (SNAP != nullptr) { + float* sn = SNAP + row * REC + ((size_t)hv * DK + dv) * DK + PT * dk; +#pragma unroll + for (int i = 0; i < PT; i++) sn[i] = state[i]; + } + } +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear head. Grid (1, NH, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, + const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT) { + const int lane = threadIdx.x, hv = blockIdx.y; + const size_t m = blockIdx.z; + constexpr int PER = DK / 32; + float yv[PER], ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) { + yv[j] = bf(Y[(m * NH + hv) * DK + lane * PER + j]); + ss += yv[j] * yv[j]; + } + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + eps); +#pragma unroll + for (int j = 0; j < PER; j++) { + const int d = lane * PER + j; + const float x = rbf(bf(NW[d]) * (yv[j] * inv)); + const float zf = bf(Z[m * HID + hv * DK + d]); + OUT[m * HID + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); + } +} + +// Per-head RMSNorm of attention rows: X's head h of row m at m * stride + h * hstride, out packed (rows, heads, HD). +// Grid (rows * heads), 64 threads of 4 values. +extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps) { + __shared__ float sums[2]; + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * hstride; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + if ((t & 31) == 0) sums[t >> 5] = acc; + __syncthreads(); + const float inv = rsqrtf((sums[0] + sums[1]) / (float)HD + eps); +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + OUT[(m * heads + head) * HD + d] = __float2bfloat16_rn(bf(W[d]) * (v[i] * inv)); + } +} + +// Text RoPE on the first 64 of HD: pairs (i, i + 32) turned by position * theta^(-i / 32). +__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, int position) { + if (d >= 64) return x[d]; + const int i = d % 32; + const float angle = (float)position * powf(10000000.0f, -(float)i / 32.0f); + float sn, cs; + sincosf(angle, &sn, &cs); + const float a = bf(x[i]), b = bf(x[i + 32]); + return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); +} + +// Queries in place at their positions; keys rotated and values copied into their segment's cache. +// Grid (rows), HD threads. +extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(__nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, + const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer) { + const int d = threadIdx.x; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + const int pos = sg.len + (int)m - sg.row0; + __shared__ __nv_bfloat16 qs[QH * HD]; + for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; + __syncthreads(); + for (int h = 0; h < QH; h++) Qn[(m * QH + h) * HD + d] = rotate(qs + h * HD, d, pos); + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + for (int h = 0; h < KVH; h++) { + const size_t dst = ((size_t)h * sg.cap + pos) * HD + d; + keys[dst] = rotate(Kn + (m * KVH + h) * HD, d, pos); + values[dst] = Vr[(m * KVH + h) * HD + d]; + } +} + +// Causal attention over a segment's cache, ck keys a block: grid (chunks, KVH, rows), a warp a query head of the +// kv head's group of QH / KVH. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +extern "C" __global__ void __launch_bounds__(128) tf_q35_attn_partial(const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, + float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; + const int c = blockIdx.x, h = blockIdx.y; + const size_t m = blockIdx.z; + const int qh = h * (QH / KVH) + g; + const Seg sg = segs[rowseg[m]]; + const int last = sg.len + (int)m - sg.row0; // the row's own position, the last key it sees + const int k0 = c * ck, k1 = min(k0 + ck, last + 1); + const size_t slot = (m * QH + qh) * nch + c; + if (k0 >= k1) { + if (lane == 0) { + PM[slot] = -INFINITY; + PL[slot] = 0.0f; + } + return; + } + const __nv_bfloat16* K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + const __nv_bfloat16* V = K + (size_t)KVH * sg.cap * HD; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[(m * QH + qh) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (int pos = k0; pos < k1; pos++) { + const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* kb = reinterpret_cast(&kw); + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const uint4 vw = *reinterpret_cast(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* vb = reinterpret_cast(&vw); +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); + mx = mn; + } + if (lane == 0) { + PM[slot] = mx; + PL[slot] = l; + } +#pragma unroll + for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; +} + +// The chunks' partials into each (row, query head)'s output. Grid (QH, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(const float* __restrict__ PM, const float* __restrict__ PL, + const float* __restrict__ PO, int nch, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x, qh = blockIdx.x; + const size_t m = blockIdx.y; + const size_t base = (m * QH + qh) * nch; + float mx = -INFINITY; + for (int c = 0; c < nch; c++) + if (PL[base + c] > 0.0f) mx = fmaxf(mx, PM[base + c]); + float lsum = 0.0f, acc[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = 0.0f; + for (int c = 0; c < nch; c++) { + if (!(PL[base + c] > 0.0f)) continue; + const float e = __expf(PM[base + c] - mx); + lsum = fmaf(PL[base + c], e, lsum); +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = fmaf(PO[(base + c) * HD + lane * DPL + i], e, acc[i]); + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[(m * QH + qh) * HD + lane * DPL + i] = __float2bfloat16_rn(acc[i] / lsum); +} + +// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows. +extern "C" __global__ void tf_q35_attn_gate(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, + __nv_bfloat16* __restrict__ OUT, int rows) { + const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= (size_t)rows * HID) return; + const size_t r = i / HID; + const int d = (int)(i % HID), head = d / HD; + const float gate = bf(QRAW[r * 2 * HID + head * 2 * HD + HD + d % HD]); + const float s = rbf(1.0f / (1.0f + expf(-gate))); + OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); +} + +// out = bf16(silu(gate) * up), n wide rows. +extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, + __nv_bfloat16* __restrict__ OUT, long long total) { + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + const float gf = bf(GATE[i]); + OUT[i] = __float2bfloat16_rn(gf / (1.0f + expf(-gf)) * bf(UP[i])); +} + +// Each row's greedy token: the largest logit, a tie to the lower id (the host draw's rule). Grid (rows), 1024 threads. +extern "C" __global__ void __launch_bounds__(1024) tf_q35_argmax(const __nv_bfloat16* __restrict__ L, int V, unsigned* __restrict__ TOK) { + __shared__ float bv[32]; + __shared__ int bi[32]; + const __nv_bfloat16* x = L + (size_t)blockIdx.x * V; + float best = -INFINITY; + int id = 0x7fffffff; + for (int i = threadIdx.x; i < V; i += blockDim.x) { + const float v = bf(x[i]); + if (v > best) { + best = v; + id = i; + } + } +#pragma unroll + for (int m = 16; m; m >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, best, m); + const int oi = __shfl_xor_sync(0xffffffffu, id, m); + if (ov > best || (ov == best && oi < id)) { + best = ov; + id = oi; + } + } + if ((threadIdx.x & 31) == 0) { + bv[threadIdx.x >> 5] = best; + bi[threadIdx.x >> 5] = id; + } + __syncthreads(); + if (threadIdx.x == 0) { + for (int w = 1; w < 32; w++) + if (bv[w] > best || (bv[w] == best && bi[w] < id)) { + best = bv[w]; + id = bi[w]; + } + TOK[blockIdx.x] = (unsigned)id; + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..c70e8b5ccf 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -30,6 +30,7 @@ pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").b pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; +pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/backend.zig b/zig/src/families/qwen3_5/backend.zig index 55fc09eecb..ef369077a8 100644 --- a/zig/src/families/qwen3_5/backend.zig +++ b/zig/src/families/qwen3_5/backend.zig @@ -163,106 +163,11 @@ pub const Metal = struct { } }; -/// A row's allowed-token bits (structured output): one bit a logit. -const mask_words = (c.vocab + 31) / 32; +const draws = @import("draw.zig"); +const mask_words = draws.mask_words; +pub const value = draws.value; +pub const draw = draws.draw; -pub fn value(word: u16) f64 { - return @as(f32, @bitCast(@as(u32, word) << 16)); -} - -/// The unchanged top-k candidates for the shared fp64 keyed sampler; a `mask` keeps its allowed ones (-inf's draw). -pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { - var best: ?usize = null; - var allowed: usize = 0; - for (logits, 0..) |word, id| { - const x = value(word); - if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; - if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; - allowed += 1; - if (best == null or x > value(logits[best.?])) best = id; - } - if (best == null) return error.NoAllowedQwenToken; - const sampling = settings orelse return @intCast(best.?); - if (sampling.temperature <= 0) return @intCast(best.?); - const k = if (sampling.top_k == 0) allowed else @min(allowed, sampling.top_k); - const values = try gpa.alloc(f64, k); - defer gpa.free(values); - const ids = try gpa.alloc(u64, k); - defer gpa.free(ids); - var filled: usize = 0; - for (logits, 0..) |word, id| { - if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; - const x = value(word); - if (k == allowed) { - values[filled] = x; - ids[filled] = id; - filled += 1; - continue; - } - if (filled == k and x <= values[k - 1]) continue; - var at = @min(filled, k - 1); - while (at > 0 and x > values[at - 1]) : (at -= 1) { - values[at] = values[at - 1]; - ids[at] = ids[at - 1]; - } - values[at] = x; - ids[at] = id; - filled = @min(filled + 1, k); - } - return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); -} - -test "candidate selection preserves full-vocabulary keyed draws and tie order" { - const gpa = std.testing.allocator; - const words = try gpa.alloc(u16, c.vocab); - defer gpa.free(words); - const values = try gpa.alloc(f64, c.vocab); - defer gpa.free(values); - const ids = try gpa.alloc(u64, c.vocab); - defer gpa.free(ids); - for (words, values, ids, 0..) |*w, *v, *id, i| { - const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; - w.* = @intCast(@as(u32, @bitCast(x)) >> 16); - v.* = value(w.*); - id.* = i; - } - for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { - for ([_]u64{ 0, 128, 65536 }) |position| { - const s: lanes.Sampling = .{ .seed = 1234, .top_k = k, .temperature = 0.7, .top_p = 0.95, .min_p = 0.1 }; - const want = try lanes.sampling.choose(gpa, values, ids, position, s); - try std.testing.expectEqual(want, try draw(gpa, words, s, position, &.{})); - } - } - try std.testing.expectEqual(try draw(gpa, words, null, 0, &.{}), try draw(gpa, words, .{ .seed = 7, .temperature = 0 }, 99, &.{})); - words[0] = 0x7fc0; - try std.testing.expectError(error.NonfiniteQwenLogits, draw(gpa, words, null, 0, &.{})); -} - -test "a masked draw is the full-vocabulary draw with the other tokens at -inf" { - const gpa = std.testing.allocator; - const words = try gpa.alloc(u16, c.vocab); - defer gpa.free(words); - const values = try gpa.alloc(f64, c.vocab); - defer gpa.free(values); - const ids = try gpa.alloc(u64, c.vocab); - defer gpa.free(ids); - var mask: [mask_words]u32 = @splat(0); - for (words, values, ids, 0..) |*w, *v, *id, i| { - const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; - w.* = @intCast(@as(u32, @bitCast(x)) >> 16); - const on = i % 7 == 3 or i % 1000 == 1; - if (on) mask[i / 32] |= @as(u32, 1) << @intCast(i % 32); - v.* = if (on) value(w.*) else -std.math.inf(f64); - id.* = i; - } - for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { - for ([_]u64{ 0, 128, 65536 }) |position| { - const s: lanes.Sampling = .{ .seed = 99, .top_k = k, .temperature = 0.8, .top_p = 0.9, .min_p = 0.05 }; - const want = try lanes.sampling.choose(gpa, values, ids, position, s); - try std.testing.expectEqual(want, try draw(gpa, words, s, position, &mask)); - } - } - try std.testing.expectEqual(try lanes.sampling.choose(gpa, values, ids, 0, .{ .seed = 1, .temperature = 0 }), try draw(gpa, words, null, 0, &mask)); - @memset(&mask, 0); - try std.testing.expectError(error.NoAllowedQwenToken, draw(gpa, words, null, 0, &mask)); +test { + _ = draws; } diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig new file mode 100644 index 0000000000..fdecefed7f --- /dev/null +++ b/zig/src/families/qwen3_5/cuda.zig @@ -0,0 +1,9 @@ +//! Qwen3.5-2B on CUDA: the engine and its native family entry (the Metal engine is qwen3_5.zig). +pub const config = @import("config.zig"); +pub const engine = @import("cuda_engine.zig"); +pub const native = @import("cuda_native.zig"); + +test { + _ = config; + _ = @import("draw.zig"); +} diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig new file mode 100644 index 0000000000..96190737d1 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -0,0 +1,484 @@ +//! Qwen3.5-2B on CUDA: the affine-4 projections through cuda.qlinear, the glue in qwen35.cu, one forward for any +//! mix of streams' windows. Each stream owns its recurrences, convolution tails and KV cache (Cache). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const c = @import("config.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows of one forward at most: a prompt chunk, or a shared round's windows. +pub const max_rows = 128; +/// Streams' windows in one forward at most. +pub const max_segs = 32; +/// Rows of a verify forward at most: each row's recurrences are recorded for a partial acceptance. +pub const max_verify = 32; +/// A window's rows at most (exact: one-row steps' bits). +pub const window_rows = 16; +/// Partial attention chunks a row at most: the chunk width grows with the window instead. +const max_chunks = 64; + +const deltas = c.layers - c.layers / 4; +const attentions = c.layers / 4; +pub const rec_bytes = c.linear_heads * c.linear_dim * c.linear_dim * 4; +pub const conv_bytes = (c.conv_taps - 1) * c.conv_dim * 2; +const state_bytes = deltas * (rec_bytes + conv_bytes); + +/// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. +pub fn streamBytes(cap: usize) usize { + return state_bytes + kvBytes(cap); +} + +fn kvBytes(cap: usize) usize { + return attentions * 2 * c.kv_heads * cap * c.head_dim * 2; +} + +/// qwen35.cu's Seg. +const Seg = extern struct { rec: u64, conv: u64, kv: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int }; + +/// One stream's state on the device and the logits of its newest row on the host. +pub const Cache = struct { + state: DeviceBuffer, + kv: DeviceBuffer, + cap: usize, + len: usize = 0, + logits: []u16, + /// The newest verify window: where it started, its first row in that forward, its rows. + last: ?struct { start: usize, row0: usize, rows: usize } = null, + + pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + var state = try DeviceBuffer.alloc(e.d, state_bytes); + errdefer state.free(); + try state.fill8(0, null); + var kv = try DeviceBuffer.alloc(e.d, kvBytes(cap)); + errdefer kv.free(); + return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, c.vocab) }; + } + + pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { + self.state.free(); + self.kv.free(); + gpa.free(self.logits); + } + + fn seg(self: *const Cache, row0: usize, rows: usize) Seg { + return .{ .rec = self.state.ptr, .conv = self.state.ptr + deltas * rec_bytes, .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; + } +}; + +const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; +const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; +const Layer = struct { + input_norm: u64, + post_norm: u64, + gate: Affine4, + up: Affine4, + down: Affine4, + mixer: union(enum) { delta: Delta, attention: Attention }, +}; + +const Fns = struct { + embed: cuda.Function, + norm: cuda.Function, + sums: cuda.Function, + gdn_pre: cuda.Function, + conv_commit: cuda.Function, + gdn_chain: cuda.Function, + gdn_post: cuda.Function, + head_norm: cuda.Function, + rope: cuda.Function, + attn_partial: cuda.Function, + attn_merge: cuda.Function, + attn_gate: cuda.Function, + mlp_act: cuda.Function, + argmax: cuda.Function, +}; + +/// A window of one stream in a forward. +pub const Part = struct { cache: *Cache, rows: usize }; + +/// What a forward keeps. A prompt chunk (`last`, one window) puts its last row's logits in its cache's `logits`; +/// verify windows keep every row, drawn on the device into `tokens` (`draws`, one a row) or their logits for +/// `rowLogits` (`logits`), and record each row's state for `rewind`. +pub const Keep = union(enum) { last, draws: []const Draw, logits }; + +/// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. +pub const Draw = struct { rule: ?cuda.sample.Rule, position: u64 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + stream: cuda.Stream, + mods: [6]cuda.Module, + draws: cuda.sample.Draws, + affine: cuda.qlinear.Affine4Kernels, + fns: Fns, + held: std.ArrayList(DeviceBuffer) = .empty, // weights and scratch, freed together + // weights + embed_w: u64, + embed_s: u64, + embed_b: u64, + head: Affine4, + final_norm: u64, + layers: [c.layers]Layer, + // scratch, max_rows rows each + h: u64, + x: u64, + r: u64, + xs: u64, + qkv: u64, + z: u64, + a: u64, + b: u64, + q: u64, + k: u64, + v: u64, + g: u64, + beta: u64, + y: u64, + mix: u64, + qraw: u64, + kraw: u64, + vraw: u64, + qn: u64, + kn: u64, + gate: u64, + up: u64, + act: u64, + pm: u64, + pl: u64, + po: u64, + logits: u64, + // the per-forward tables: ids, segments, each row's segment and the draws' rules, staged in pinned memory + tables: DeviceBuffer, + staged: cuda.HostBuffer, + host_logits: cuda.HostBuffer, + /// The kept rows' device draws, when forward was given them. + tokens: [max_verify]u32 = undefined, + rec_snap: u64, // recorded recurrences: max_verify rows a linear layer + conv_snap: u64, // recorded convolution tails, the same + + const ids_at = 0; + const segs_at = max_rows * 4; + const rowseg_at = segs_at + max_segs * @sizeOf(Seg); + const rules_at = rowseg_at + max_rows * 4; + const meta_at = rules_at + max_verify * @sizeOf(cuda.sample.Rule); + const tokens_at = meta_at + max_verify * 4; + const tables_len = tokens_at + max_verify * 4; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const config = try c.Config.read(gpa, io, dir); + const d = ctx.d; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = undefined; + e.gpa = gpa; + e.d = d; + e.held = .empty; + e.stream = try cuda.Stream.init(d, true); + errdefer e.stream.deinit(); + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample }; + var loaded: usize = 0; + errdefer for (e.mods[0..loaded]) |*m| m.unload(); + for (images, 0..) |img, i| { + e.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + const gb10 = try ctx.attribute(.compute_capability_major) == 12 and try ctx.attribute(.compute_capability_minor) == 1; + e.affine = try cuda.qlinear.Affine4Kernels.resolve(d, e.mods[0], e.mods[1], e.mods[2], e.mods[3], sms, gb10); + errdefer e.affine.deinit(); + e.draws = try cuda.sample.Draws.resolve(e.mods[5]); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); + errdefer e.freeHeld(); + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + try e.loadWeights(&ck, config.tied); + try e.allocScratch(); + e.tables = try DeviceBuffer.alloc(d, tables_len); + errdefer e.tables.free(); + e.staged = try cuda.HostBuffer.alloc(d, tables_len); + errdefer e.staged.free(); + e.host_logits = try cuda.HostBuffer.alloc(d, max_verify * c.vocab * 2); + return e; + } + + pub fn deinit(e: *Engine) void { + e.host_logits.free(); + e.staged.free(); + e.tables.free(); + e.freeHeld(); + e.affine.deinit(); + for (&e.mods) |*m| m.unload(); + e.stream.deinit(); + e.gpa.destroy(e); + } + + fn freeHeld(e: *Engine) void { + for (e.held.items) |*b| b.free(); + e.held.deinit(e.gpa); + } + + fn own(e: *Engine, b: DeviceBuffer) !u64 { + try e.held.append(e.gpa, b); + return b.ptr; + } + + fn alloc(e: *Engine, bytes: usize) !u64 { + var b = try DeviceBuffer.alloc(e.d, bytes); + errdefer b.free(); + return e.own(b); + } + + fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { + var b = try DeviceBuffer.fromHost(e.d, t.bytes); + errdefer b.free(); + return e.own(b); + } + + fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [160]u8 = undefined; + return e.put(try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); + } + + /// An MLX affine-4 projection packed into qlinear's tiles; the checkpoint's own layout leaves the device after. + fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var raw: [3]DeviceBuffer = undefined; + var got: usize = 0; + defer for (raw[0..got]) |*b| b.free(); + const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / c.group }, .{ "biases", .bf16, k / c.group } }; + for (parts) |p| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + raw[got] = try DeviceBuffer.fromHost(e.d, t.bytes); + got += 1; + } + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0].ptr, raw[1].ptr, raw[2].ptr, q); + try e.stream.synchronize(); + return q; + } + + fn loadWeights(e: *Engine, ck: *core.Checkpoint, tied: bool) !void { + const m = "language_model.model."; + e.embed_w = try e.tensor(ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ c.vocab, c.hidden / 8 }); + e.embed_s = try e.tensor(ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); + e.embed_b = try e.tensor(ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); + e.head = if (tied) try e.linear(ck, m, "embed_tokens", c.vocab, c.hidden) else try e.linear(ck, "language_model.", "lm_head", c.vocab, c.hidden); + e.final_norm = try e.tensor(ck, "{s}norm.weight", .{m}, .bf16, &.{c.hidden}); + for (&e.layers, 0..) |*l, i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + l.* = .{ + .input_norm = try e.tensor(ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{c.hidden}), + .post_norm = try e.tensor(ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{c.hidden}), + .gate = try e.linear(ck, p, "mlp.gate_proj", c.intermediate, c.hidden), + .up = try e.linear(ck, p, "mlp.up_proj", c.intermediate, c.hidden), + .down = try e.linear(ck, p, "mlp.down_proj", c.hidden, c.intermediate), + .mixer = undefined, + }; + if (c.linear(i)) { + l.mixer = .{ .delta = .{ + .qkv = try e.linear(ck, p, "linear_attn.in_proj_qkv", c.conv_dim, c.hidden), + .z = try e.linear(ck, p, "linear_attn.in_proj_z", c.hidden, c.hidden), + .a = try e.linear(ck, p, "linear_attn.in_proj_a", c.linear_heads, c.hidden), + .b = try e.linear(ck, p, "linear_attn.in_proj_b", c.linear_heads, c.hidden), + .out = try e.linear(ck, p, "linear_attn.out_proj", c.hidden, c.hidden), + .conv = try e.tensor(ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ c.conv_dim, c.conv_taps, 1 }), + .norm = try e.tensor(ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{c.linear_dim}), + .a_log = try e.tensor(ck, "{s}linear_attn.A_log", .{p}, .f32, &.{c.linear_heads}), + .dt_bias = try e.tensor(ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{c.linear_heads}), + } }; + } else { + l.mixer = .{ .attention = .{ + .q = try e.linear(ck, p, "self_attn.q_proj", 2 * c.hidden, c.hidden), + .k = try e.linear(ck, p, "self_attn.k_proj", c.kv_heads * c.head_dim, c.hidden), + .v = try e.linear(ck, p, "self_attn.v_proj", c.kv_heads * c.head_dim, c.hidden), + .o = try e.linear(ck, p, "self_attn.o_proj", c.hidden, c.hidden), + .q_norm = try e.tensor(ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{c.head_dim}), + .k_norm = try e.tensor(ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{c.head_dim}), + } }; + } + } + } + + fn allocScratch(e: *Engine) !void { + const R = max_rows; + inline for (.{ "h", "x", "r", "z", "q", "k", "v", "y", "mix", "qn" }) |name| @field(e, name) = try e.alloc(R * c.hidden * 2); + inline for (.{ "qkv", "gate", "up", "act" }) |name| @field(e, name) = try e.alloc(R * c.intermediate * 2); + e.xs = try e.alloc(R * (c.intermediate / 64) * 4); + inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try e.alloc(R * c.linear_heads * 2); + e.g = try e.alloc(R * c.linear_heads * 4); + e.qraw = try e.alloc(R * 2 * c.hidden * 2); + inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try e.alloc(R * c.kv_heads * c.head_dim * 2); + e.pm = try e.alloc(R * c.query_heads * max_chunks * 4); + e.pl = try e.alloc(R * c.query_heads * max_chunks * 4); + e.po = try e.alloc(R * c.query_heads * max_chunks * c.head_dim * 4); + e.logits = try e.alloc(max_verify * c.vocab * 2); + e.rec_snap = try e.alloc(deltas * max_verify * rec_bytes); + e.conv_snap = try e.alloc(deltas * max_verify * conv_bytes); + } + + fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = @intCast(block[0]), .y = @intCast(block[1]), .z = @intCast(block[2]) } }; + try cuda.launch.launch(f, cfg, e.stream, &a); + } + + /// out (rows, n) = x (rows, k) through `l`: decode tiles 16 rows at a time (each row's bits whatever the batch), + /// or the prompt GEMM for a prompt chunk wider than one tile. + fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { + if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)) }); + var at: usize = 0; + while (at < rows) : (at += 16) { + const n = @min(16, rows - at); + try e.affine.decode(e.stream, x + at * l.k * 2, e.xs + at * (l.k / 64) * 4, l, out + at * l.n * 2, n); + } + } + + fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, @as(u64, 0), c.eps }); + } + + /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. + pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { + var total: usize = 0; + var widest: usize = 0; + for (parts) |p| { + if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; + total += p.rows; + widest = @max(widest, p.cache.len + p.rows); + } + if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; + const record = keep != .last; + if (record and total > max_verify) return error.InvalidQwenSegments; + if (!record and parts.len != 1) return error.InvalidQwenSegments; + if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; + for (ids) |id| if (id >= c.vocab) return error.InvalidQwenToken; + const heads: usize = if (record) total else 1; + // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps + const prompt = !record; + // the tables + const st = e.staged.bytes; + @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids); + const segs = std.mem.bytesAsSlice(Seg, st[segs_at..][0 .. parts.len * @sizeOf(Seg)]); + const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); + var row0: usize = 0; + for (parts, 0..) |p, i| { + segs[i] = p.cache.seg(row0, p.rows); + for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); + p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; + row0 += p.rows; + } + if (keep == .draws) { + const rules = std.mem.bytesAsSlice(cuda.sample.Rule, st[rules_at..][0 .. total * @sizeOf(cuda.sample.Rule)]); + const meta = std.mem.bytesAsSlice(c_int, st[meta_at..][0 .. total * 4]); + for (keep.draws, rules, meta) |p, *r, *m| { + if (p.rule) |x| r.* = x; + m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 + } + } + try e.tables.uploadAsync(0, st, e.stream.handle); + const t_ids = e.tables.ptr + ids_at; + const t_segs = e.tables.ptr + segs_at; + const t_rows = e.tables.ptr + rowseg_at; + const nseg = parts.len; + // attention chunks: at most max_chunks a row, each a multiple of 128 keys + const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); + const nch = (widest + ck - 1) / ck; + + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h }); + var di: usize = 0; + var ai: c_int = 0; + for (e.layers, 0..) |l, i| { + try e.norm(i != 0, l.input_norm, total); + switch (l.mixer) { + .delta => |m| { + const layer: c_int = @intCast(di); + const rec_snap: u64 = if (record) e.rec_snap + di * max_verify * rec_bytes else 0; + const conv_snap: u64 = if (record) e.conv_snap + di * max_verify * conv_bytes else 0; + try e.proj(m.qkv, e.x, e.qkv, total, prompt); + try e.proj(m.z, e.x, e.z, total, prompt); + try e.proj(m.a, e.x, e.a, total, prompt); + try e.proj(m.b, e.x, e.b, total, prompt); + try e.run(e.fns.gdn_pre, .{ 1, 3 * c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g, e.beta }); + try e.run(e.fns.conv_commit, .{ (c.conv_dim + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ e.qkv, t_segs, layer, conv_snap }); + try e.run(e.fns.gdn_chain, .{ 1, c.linear_dim / 4, c.linear_heads * nseg }, .{ 32, 4, 1 }, .{ e.q, e.k, e.v, e.g, e.beta, t_segs, layer, e.y, rec_snap }); + try e.run(e.fns.gdn_post, .{ 1, c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.y, e.z, m.norm, c.eps, e.mix }); + try e.proj(m.out, e.mix, e.r, total, prompt); + di += 1; + }, + .attention => |m| { + try e.proj(m.q, e.x, e.qraw, total, prompt); + try e.proj(m.k, e.x, e.kraw, total, prompt); + try e.proj(m.v, e.x, e.vraw, total, prompt); + try e.run(e.fns.head_norm, .{ total * c.query_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, c.query_heads), @as(c_int, 2 * c.hidden), @as(c_int, 2 * c.head_dim), c.eps }); + try e.run(e.fns.head_norm, .{ total * c.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, c.kv_heads), @as(c_int, c.kv_heads * c.head_dim), @as(c_int, c.head_dim), c.eps }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ c.head_dim, 1, 1 }, .{ e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); + try e.run(e.fns.attn_partial, .{ nch, c.kv_heads, total }, .{ 128, 1, 1 }, .{ e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ c.query_heads, total, 1 }, .{ 32, 1, 1 }, .{ e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); + try e.run(e.fns.attn_gate, .{ (total * c.hidden + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); + try e.proj(m.o, e.mix, e.r, total, prompt); + ai += 1; + }, + } + try e.norm(true, l.post_norm, total); + try e.proj(l.gate, e.x, e.gate, total, prompt); + try e.proj(l.up, e.x, e.up, total, prompt); + const n: c_longlong = @intCast(total * c.intermediate); + try e.run(e.fns.mlp_act, .{ (total * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); + try e.proj(l.down, e.act, e.r, total, prompt); + } + try e.norm(true, e.final_norm, total); + // the head over the kept rows + try e.proj(e.head, e.x + (total - heads) * c.hidden * 2, e.logits, heads, false); + for (parts) |p| p.cache.len += p.rows; + switch (keep) { + .draws => |pk| { + const toks = e.tables.ptr + tokens_at; + for (pk, 0..) |p, i| { + const row = e.logits + i * c.vocab * 2; + if (p.rule != null) { + try e.draws.draw(e.stream, row, c.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, c.vocab), toks + i * 4 }); + } + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); + }, + .logits, .last => { + const out = e.host_logits.bytes[0 .. heads * c.vocab * 2]; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); + }, + } + } + + /// Row `i`'s logits from the newest forward that kept logits. + pub fn rowLogits(e: *const Engine, i: usize) []const u16 { + return e.host_logits.slice(u16)[i * c.vocab ..][0..c.vocab]; + } + + /// Keeps the first `rows` of the cache's newest window: its recurrences and convolution tails as they were after + /// that row (the forward recorded them), its length at that row. The window's later keys stay as unread slots. + pub fn rewind(e: *Engine, cache: *Cache, rows: usize) !void { + const last = cache.last orelse return error.NothingToKeep; + if (rows == 0 or rows > last.rows) return error.NothingToKeep; + cache.last = null; + cache.len = last.start + rows; + if (rows == last.rows) return; + const row = last.row0 + rows - 1; + for (0..deltas) |di| { + try cache.state.copyFrom(di * rec_bytes, e.rec_snap + (di * max_verify + row) * rec_bytes, rec_bytes, e.stream.handle); + try cache.state.copyFrom(deltas * rec_bytes + di * conv_bytes, e.conv_snap + (di * max_verify + row) * conv_bytes, conv_bytes, e.stream.handle); + } + try e.stream.synchronize(); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig new file mode 100644 index 0000000000..4995a82b45 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -0,0 +1,211 @@ +//! Qwen3.5-2B's CUDA lane backend behind the native family interface (native/cuda.zig's registry): windows of up +//! to 16 rows (the lane core's copies), every stream's window in one forward, prompts in chunks, draws on the device. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const c = @import("config.zig"); +const E = @import("cuda_engine.zig"); +const draws = @import("draw.zig"); +const be = lanes.backend; + +pub const model_type = "qwen3_5"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = E.max_rows; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + +/// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. +pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + /// Device bytes each admitted stream allocates (recurrences, convolution tails, keys and values). + stream_bytes: usize, + free_streams: u32 = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, +}; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.QwenHasNoCapturedKernels; + const e = try E.Engine.init(gpa, io, ctx, dir); + errdefer e.deinit(); + const l = try gpa.create(Lanes); + l.* = .{ .gpa = gpa, .e = e, .capacity = o.context }; + return .{ + .backend = l.backend(), + .facts = .{ .exact_width = E.window_rows, .hidden_rows = true, .batch_rows = E.max_verify, .max_streams = E.max_segs, .streams_exact = true }, + .rows = E.window_rows, + .stream_bytes = E.streamBytes(o.context), + .ctx = l, + .deinit = release, + }; +} + +/// A request this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.InvalidQwenPrompt => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const l: *Lanes = @ptrCast(@alignCast(p)); + l.deinit(); +} + +/// The device draw of a stream's rule at `position`: sample.cu's keyed rule, greedy at temperature 0 or without one. +fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { + const s = sampling orelse return .{ .rule = null, .position = position }; + if (s.temperature <= 0) return .{ .rule = null, .position = position }; + return .{ .rule = cuda.sample.rule(s.seed, s.temperature, s.top_p, s.minLog(), s.top_k), .position = position }; +} + +const Lanes = struct { + gpa: std.mem.Allocator, + e: *E.Engine, + capacity: usize, + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + + fn deinit(self: *Lanes) void { + var it = self.caches.valueIterator(); + while (it.next()) |cache| { + cache.*.deinit(self.gpa); + self.gpa.destroy(cache.*); + } + self.caches.deinit(self.gpa); + self.e.deinit(); + self.gpa.destroy(self); + } + + fn backend(self: *Lanes) be.Backend { + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Lanes { + return @ptrCast(@alignCast(ptr)); + } + + fn cacheOf(self: *Lanes, stream: *lanes.Stream) !*E.Cache { + return self.caches.get(stream) orelse error.UnknownQwenStream; + } + + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + const self = cast(ptr); + const ids = stream.prompt(); + if (ids.len == 0 or ids.len + stream.max_new > self.capacity or self.caches.contains(stream)) return error.InvalidQwenPrompt; + const cache = try self.gpa.create(E.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try E.Cache.init(self.gpa, self.e, @min(self.capacity, ids.len + stream.max_new + 1)); + errdefer cache.deinit(self.gpa); + var at: usize = 0; + while (at < ids.len) { + var end: usize = @min(ids.len, at + E.max_rows); + for (stream.chunks) |boundary| if (boundary > at) { + end = @min(end, boundary); + break; + }; + try self.e.forward(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], .last); + at = end; + } + try self.caches.put(self.gpa, stream, cache); + } + + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + return firstMaskedFn(ptr, stream, position, &.{}); + } + + fn firstMaskedFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidQwenPosition; + return try draws.draw(self.gpa, cache.logits, stream.sampling, position, mask); + } + + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len + 1) return error.InvalidQwenPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| @intCast(h), + }; + try self.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{token}, .{ .draws = &.{pick(stream.sampling, position)} }); + cache.last = null; + return self.e.tokens[0]; + } + + fn readFn(_: *anyopaque, handle: u64) !u32 { + if (handle >= c.vocab) return error.InvalidQwenToken; + return @intCast(handle); + } + + /// Each window (a stream's pending token and copies) checked, then as many as fit in each forward. + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + if (windows.len != out.len) return error.InvalidQwenWindows; + for (windows, out) |w, o| { + const rows = w.rows(); + if (w.parents != null or w.held != 0 or rows > E.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; + if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; + const cache = try self.cacheOf(w.stream); + for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidQwenPosition; + } + var at: usize = 0; + while (at < windows.len) { + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_verify]u32 = undefined; + var picks: [E.max_verify]E.Draw = undefined; + var n: usize = 0; + var total: usize = 0; + var masked = false; + while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= E.max_verify) : (n += 1) { + const w = windows[at + n]; + parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows() }; + ids[total] = w.pending; + @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + for (w.positions, 0..) |p, r| picks[total + r] = pick(w.stream.sampling, p); + masked = masked or w.masks.len != 0; + total += w.rows(); + } + // structured output draws on the host, over each row's logits with its mask applied + try self.e.forward(parts[0..n], ids[0..total], if (masked) .logits else .{ .draws = picks[0..total] }); + var row: usize = 0; + for (windows[at..][0..n], out[at..][0..n]) |w, o| { + for (o.sampled, w.positions, 0..) |*t, p, r| { + const mask = if (w.masks.len > 0) w.masks[r * draws.mask_words ..][0..draws.mask_words] else &.{}; + t.* = if (masked) try draws.draw(self.gpa, self.e.rowLogits(row + r), w.stream.sampling, p, mask) else self.e.tokens[row + r]; + } + @memcpy(o.drafts, w.tokens); + row += w.rows(); + } + at += n; + } + } + + /// The accepted prefix of each window: its state as it was after that row. + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + for (windows, paths) |w, path| { + for (path, 0..) |r, i| if (r != i) return error.UnsupportedQwenTree; + try self.e.rewind(try self.cacheOf(w.stream), path.len); + } + } + + fn draftFn(_: *anyopaque, requests: []const be.DraftRequest) !void { + for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + if (self.caches.fetchRemove(stream)) |entry| { + entry.value.deinit(self.gpa); + self.gpa.destroy(entry.value); + } + } +}; diff --git a/zig/src/families/qwen3_5/draw.zig b/zig/src/families/qwen3_5/draw.zig new file mode 100644 index 0000000000..c9acd623bf --- /dev/null +++ b/zig/src/families/qwen3_5/draw.zig @@ -0,0 +1,108 @@ +//! The host draw over a row of bf16 logits, shared by the Metal and CUDA backends: greedy, keyed sampling or a mask. +const std = @import("std"); +const lanes = @import("lanes"); +const c = @import("config.zig"); + +/// A row's allowed-token bits (structured output): one bit a logit. +pub const mask_words = (c.vocab + 31) / 32; + +pub fn value(word: u16) f64 { + return @as(f32, @bitCast(@as(u32, word) << 16)); +} + +/// The unchanged top-k candidates for the shared fp64 keyed sampler; a `mask` keeps its allowed ones (-inf's draw). +pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { + var best: ?usize = null; + var allowed: usize = 0; + for (logits, 0..) |word, id| { + const x = value(word); + if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + allowed += 1; + if (best == null or x > value(logits[best.?])) best = id; + } + if (best == null) return error.NoAllowedQwenToken; + const sampling = settings orelse return @intCast(best.?); + if (sampling.temperature <= 0) return @intCast(best.?); + const k = if (sampling.top_k == 0) allowed else @min(allowed, sampling.top_k); + const values = try gpa.alloc(f64, k); + defer gpa.free(values); + const ids = try gpa.alloc(u64, k); + defer gpa.free(ids); + var filled: usize = 0; + for (logits, 0..) |word, id| { + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + const x = value(word); + if (k == allowed) { + values[filled] = x; + ids[filled] = id; + filled += 1; + continue; + } + if (filled == k and x <= values[k - 1]) continue; + var at = @min(filled, k - 1); + while (at > 0 and x > values[at - 1]) : (at -= 1) { + values[at] = values[at - 1]; + ids[at] = ids[at - 1]; + } + values[at] = x; + ids[at] = id; + filled = @min(filled + 1, k); + } + return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); +} + +test "candidate selection preserves full-vocabulary keyed draws and tie order" { + const gpa = std.testing.allocator; + const words = try gpa.alloc(u16, c.vocab); + defer gpa.free(words); + const values = try gpa.alloc(f64, c.vocab); + defer gpa.free(values); + const ids = try gpa.alloc(u64, c.vocab); + defer gpa.free(ids); + for (words, values, ids, 0..) |*w, *v, *id, i| { + const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; + w.* = @intCast(@as(u32, @bitCast(x)) >> 16); + v.* = value(w.*); + id.* = i; + } + for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { + for ([_]u64{ 0, 128, 65536 }) |position| { + const s: lanes.Sampling = .{ .seed = 1234, .top_k = k, .temperature = 0.7, .top_p = 0.95, .min_p = 0.1 }; + const want = try lanes.sampling.choose(gpa, values, ids, position, s); + try std.testing.expectEqual(want, try draw(gpa, words, s, position, &.{})); + } + } + try std.testing.expectEqual(try draw(gpa, words, null, 0, &.{}), try draw(gpa, words, .{ .seed = 7, .temperature = 0 }, 99, &.{})); + words[0] = 0x7fc0; + try std.testing.expectError(error.NonfiniteQwenLogits, draw(gpa, words, null, 0, &.{})); +} + +test "a masked draw is the full-vocabulary draw with the other tokens at -inf" { + const gpa = std.testing.allocator; + const words = try gpa.alloc(u16, c.vocab); + defer gpa.free(words); + const values = try gpa.alloc(f64, c.vocab); + defer gpa.free(values); + const ids = try gpa.alloc(u64, c.vocab); + defer gpa.free(ids); + var mask: [mask_words]u32 = @splat(0); + for (words, values, ids, 0..) |*w, *v, *id, i| { + const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; + w.* = @intCast(@as(u32, @bitCast(x)) >> 16); + const on = i % 7 == 3 or i % 1000 == 1; + if (on) mask[i / 32] |= @as(u32, 1) << @intCast(i % 32); + v.* = if (on) value(w.*) else -std.math.inf(f64); + id.* = i; + } + for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { + for ([_]u64{ 0, 128, 65536 }) |position| { + const s: lanes.Sampling = .{ .seed = 99, .top_k = k, .temperature = 0.8, .top_p = 0.9, .min_p = 0.05 }; + const want = try lanes.sampling.choose(gpa, values, ids, position, s); + try std.testing.expectEqual(want, try draw(gpa, words, s, position, &mask)); + } + } + try std.testing.expectEqual(try lanes.sampling.choose(gpa, values, ids, 0, .{ .seed = 1, .temperature = 0 }), try draw(gpa, words, null, 0, &mask)); + @memset(&mask, 0); + try std.testing.expectError(error.NoAllowedQwenToken, draw(gpa, words, null, 0, &mask)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 50af1bd955..b90c801e0d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -5,13 +5,14 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const qwen35 = @import("qwen35"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{nemotron.native}; +const registry = .{ nemotron.native, qwen35.native }; /// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ @@ -336,7 +337,10 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + // what a family's Loaded may leave out: no learner, no lone driver, no prompt-cache copies + const L = @TypeOf(loaded); + const learn_bytes: usize = if (@hasField(L, "learn_bytes")) loaded.learn_bytes else 0; + const room = after.room(model) -| learn_bytes; const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { @@ -347,14 +351,14 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const h = try gpa.create(Host); errdefer gpa.destroy(h); - h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; + const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasField(L, "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet @@ -365,7 +369,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (@hasField(L, "learns") and loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig new file mode 100644 index 0000000000..5f0dc8f906 --- /dev/null +++ b/zig/tests/qwen35_cuda_main.zig @@ -0,0 +1,58 @@ +//! `tf-qwen35-cuda `: Qwen3.5-2B's CUDA forward over comma-separated token ids, +//! one token a forward (step: every position's logits) or the ids as one prompt (prompt: the last position's), as +//! little-endian bf16 rows into for a reference to compare. Also prints the step mode's decode speed. +const std = @import("std"); +const cuda = @import("cuda"); +const q = @import("qwen35"); +const E = q.engine; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len != 5) { + std.debug.print("usage: tf-qwen35-cuda \n", .{}); + return 2; + } + var ids: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeAny(u8, args[3], ", "); + while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); + defer e.deinit(); + var cache = try E.Cache.init(init.gpa, e, 4096); + defer cache.deinit(init.gpa); + var out: std.ArrayList(u8) = .empty; + if (std.mem.eql(u8, args[2], "step")) { + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + for (ids.items) |id| { + try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, &.{id}, .logits); + @memcpy(cache.logits, e.rowLogits(0)); + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + const ns: u64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); + std.debug.print("{d} steps in {d:.1} ms: {d:.1} tok/s\n", .{ ids.items.len, @as(f64, @floatFromInt(ns)) / 1e6, @as(f64, @floatFromInt(ids.items.len)) * 1e9 / @as(f64, @floatFromInt(ns)) }); + } else if (std.mem.eql(u8, args[2], "window")) { + // windows of 1, 2, 3, ... rows: every row's logits, which must equal step's bit for bit + var at: usize = 0; + var w: usize = 1; + while (at < ids.items.len) : (w = w % E.window_rows + 1) { + const n = @min(w, ids.items.len - at); + try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .logits); + for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + at += n; + } + } else { + var at: usize = 0; + while (at < ids.items.len) { + const n = @min(E.max_rows, ids.items.len - at); + try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); + at += n; + } + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = out.items }); + return 0; +} From 3ace97ed415aa4fb47b7ff3cb18853b848a1d78e Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:10:43 +0000 Subject: [PATCH 07/46] qwen3_5: serve Qwen3.8-27B on CUDA, one arena for a model's device memory The 2B's CUDA engine now reads its widths from config.json (cuda_geometry.zig): linear value heads over fewer key heads, more query heads a key/value head, an untied head, an A_log stored as bf16. The 27B's logits agree with an fp32 reference at every position of a 36-token check, and its windows equal one-row steps bit for bit; the 2B's output is unchanged. Two Windows driver costs shaped the rest. Each kernel launch grew with the process's allocation count (1.2 ms a launch over 1,900 weight buffers), so a model's weights and scratch now live in one arena (cuda_weights.zig packs each projection through a reused staging buffer: the 27B loads in 25 s, not 200). And a process's first forwards through a buffer stall for seconds, so the engine warms every request path at load and the backend keeps a pool of stream caches that `reserve` allocates and runs a row through once the server knows its stream count. On an RTX 4090: the 27B decodes at 37 tok/s in the runner, 31 to 35 tok/s served, 210 tok/s across eight streams; the 2B stays at 210 and 1,250. --- zig/kernels/cuda/qwen35.cu | 195 +++++++------ zig/src/families/qwen3_5/cuda.zig | 5 +- zig/src/families/qwen3_5/cuda_engine.zig | 325 ++++++++++----------- zig/src/families/qwen3_5/cuda_geometry.zig | 177 +++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 80 +++-- zig/src/families/qwen3_5/cuda_weights.zig | 143 +++++++++ zig/src/native/cuda.zig | 5 + zig/tests/qwen35_cuda_main.zig | 21 +- 8 files changed, 656 insertions(+), 295 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_geometry.zig create mode 100644 zig/src/families/qwen3_5/cuda_weights.zig diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index a68fa30990..715f83f64e 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -1,27 +1,30 @@ -// Qwen3.5-2B's glue on CUDA: the Metal engine's kernels (zig/kernels/metal/qwen3_5) with the same bf16 rounding points. -// Rows of one forward belong to segments (one stream's window each); a segment's state lives in its stream's blocks. +// Qwen3.5's hybrid text models on CUDA (the 2B and Qwen3.8-27B): the Metal engine's kernels (zig/kernels/metal/qwen3_5) +// with the same bf16 rounding points, their widths from the model's Geo. Rows of one forward belong to segments (one +// stream's window each); a segment's state lives in its stream's blocks. #include #include namespace tf_q35 { -constexpr int HID = 2048; // hidden -constexpr int NH = 16; // linear-attention heads (keys and values) -constexpr int DK = 128; // their head width -constexpr int CONV = 6144; // q, k and v channels the convolution sees -constexpr int TAPS = 4; -constexpr int QH = 8; // attention query heads -constexpr int KVH = 2; // attention key/value heads -constexpr int HD = 256; // attention head width -constexpr size_t REC = (size_t)NH * DK * DK; // floats of one layer's recurrence -constexpr size_t CONV_ROWS = (size_t)(TAPS - 1) * CONV; +constexpr int DK = 128; // linear-attention key and value head width +constexpr int TAPS = 4; // causal convolution taps +constexpr int HD = 256; // attention head width +constexpr int MAXQ = 32; // attention query heads at most + +// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads. +struct Geo { + int hid, nk, nv, qh, kvh; + __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees + __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence + __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } +}; // One stream's window in this forward: its state blocks, cache length before the window, and rows. struct Seg { - unsigned long long rec; // float recurrences, one REC block a linear layer - unsigned long long conv; // bf16 convolution tails, CONV_ROWS a linear layer - unsigned long long kv; // bf16 keys then values, KVH * cap * HD each, an attention layer at a time + unsigned long long rec; // float recurrences, one rec() block a linear layer + unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer + unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time int cap, len, row0, rows; }; @@ -34,40 +37,39 @@ __device__ __forceinline__ float warp_sum(float x) { return x; } -__device__ __forceinline__ float sigmoid(float x) { return 1.0f / (1.0f + __expf(-x)); } - } // namespace tf_q35 using namespace tf_q35; -// MLX affine 4-bit token rows: out = bf16(s * q + b). +// MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, - __nv_bfloat16* __restrict__ out) { + __nv_bfloat16* __restrict__ out, int D) { const size_t id = ids[blockIdx.x]; - for (int b = threadIdx.x; b < HID / 2; b += blockDim.x) { - const uint8_t v = W[id * (HID / 2) + b]; - const size_t g = id * (HID / 64) + (2 * b) / 64; + for (int b = threadIdx.x; b < D / 2; b += blockDim.x) { + const uint8_t v = W[id * (D / 2) + b]; + const size_t g = id * (D / 64) + (2 * b) / 64; const float s = bf(S[g]), z = bf(B[g]); - out[(size_t)blockIdx.x * HID + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); - out[(size_t)blockIdx.x * HID + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); + out[(size_t)blockIdx.x * D + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); + out[(size_t)blockIdx.x * D + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); } } -// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), xs = x's 64-group sums. One 256-thread block a row. +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), D wide (a multiple of 256, at most 8192). One 256-thread +// block a row; each thread's elements in index order, warps summed in order. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, - const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, - float* __restrict__ XS, float eps) { - constexpr int E = HID / 256; + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D) { + constexpr int MAXE = 32; __shared__ float red[8]; - __shared__ float xb[HID]; const size_t row = blockIdx.x; const int t = threadIdx.x; - float hv[E]; + const int E = D / 256; + float hv[MAXE]; float ss = 0.0f; #pragma unroll - for (int i = 0; i < E; i++) { - const size_t at = row * HID + t * E + i; + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; + const size_t at = row * D + t * E + i; float h = bf(H[at]); if (R != nullptr) { const __nv_bfloat16 hb = __float2bfloat16_rn(h + bf(R[at])); @@ -82,19 +84,12 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __r __syncthreads(); float total = 0.0f; for (int i = 0; i < 8; i++) total += red[i]; - const float inv = rsqrtf(total / (float)HID + eps); + const float inv = rsqrtf(total / (float)D + eps); #pragma unroll - for (int i = 0; i < E; i++) { + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; const int d = t * E + i; - const __nv_bfloat16 x = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); - X[row * HID + d] = x; - xb[d] = bf(x); - } - __syncthreads(); - if (XS != nullptr && t < HID / 64) { - float s = 0.0f; - for (int j = 0; j < 64; j++) s += xb[t * 64 + j]; - XS[row * (HID / 64) + t] = s; + X[row * D + d] = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); } } @@ -110,8 +105,8 @@ extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, floa } // The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. -// Grid (1, 3 * NH, rows), a warp a (row, head): heads [0, NH) are q, [NH, 2NH) k, [2NH, 3NH) v. -extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, +// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, const int* __restrict__ rowseg, int layer, const __nv_bfloat16* __restrict__ CW, const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, @@ -120,10 +115,11 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloa const int lane = threadIdx.x; const int head = blockIdx.y; const int w = blockIdx.z; + const int NK = geo.nk, NV = geo.nv, C = geo.conv(); const Seg sg = segs[rowseg[w]]; const int r = w - sg.row0; - const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * CONV_ROWS; - const bool isq = head < NH, isk = !isq && head < 2 * NH; + const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * geo.conv_rows(); + const bool isq = head < NK, isk = !isq && head < 2 * NK; const int c0 = head * DK; // q, k and v channels are consecutive heads of DK constexpr int PER = DK / 32; float vals[PER]; @@ -133,7 +129,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloa float acc = 0.0f; for (int tap = 0; tap < TAPS; tap++) { const int at = r + tap; // in [state rows (TAPS - 1)] ++ [this segment's rows] - const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * CONV + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]); + const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * C + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]); acc += bf(CW[c * TAPS + tap]) * xv; } const float conv = rbf(acc); @@ -147,57 +143,61 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloa ss = warp_sum(ss); const float inv = rsqrtf(ss / (float)DK + (1e-6f / (float)DK)); const float scale = isq ? rbf(1.0f / (float)DK) : rbf(rsqrtf((float)DK)); - __nv_bfloat16* out = isq ? Q + ((size_t)w * NH + head) * DK : K + ((size_t)w * NH + head - NH) * DK; + __nv_bfloat16* out = isq ? Q + ((size_t)w * NK + head) * DK : K + ((size_t)w * NK + head - NK) * DK; #pragma unroll for (int j = 0; j < PER; j++) out[lane * PER + j] = __float2bfloat16_rn(scale * rbf(vals[j] * inv)); } else { - const int hv = head - 2 * NH; + const int hv = head - 2 * NK; #pragma unroll - for (int j = 0; j < PER; j++) V[((size_t)w * NH + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); + for (int j = 0; j < PER; j++) V[((size_t)w * NV + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); if (lane == 0) { - const float s = rbf(bf(Ain[w * NH + hv]) + bf(DT[hv])); + const float s = rbf(bf(Ain[w * NV + hv]) + bf(DT[hv])); const float sp = rbf(fmaxf(s, 0.0f) + logf(1.0f + expf(-fabsf(s)))); - G[w * NH + hv] = expf(-expf(ALOG[hv]) * sp); - BETA[w * NH + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NH + hv])))); + G[w * NV + hv] = expf(-expf(ALOG[hv]) * sp); + BETA[w * NV + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NV + hv])))); } } } // Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). -// SNAP, if given, also takes the tail after each row (forward row m at m * CONV_ROWS), read before the update. -extern "C" __global__ void tf_q35_conv_commit(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, +// SNAP, if given, also takes the tail after each row (forward row m at m * conv_rows()), read before the update. +extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, __nv_bfloat16* __restrict__ SNAP) { + const int C = geo.conv(); const int c = blockIdx.x * blockDim.x + threadIdx.x; - if (c >= CONV) return; + if (c >= C) return; const Seg sg = segs[blockIdx.y]; - __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * CONV_ROWS; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); if (SNAP != nullptr) for (int j = 0; j < sg.rows; j++) for (int r = 0; r < TAPS - 1; r++) { const int at = r + j + 1; - SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * CONV + c] = - at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * C + c] = + at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; } // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row for (int r = 0; r < TAPS - 1; r++) { const int at = r + sg.rows; - CS[(size_t)r * CONV + c] = at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + CS[(size_t)r * C + c] = at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; } } // The gated delta rule over each segment's rows, its recurrence updated in place; SNAP, if given, takes the state after -// each row (forward row m at m * REC). Grid (1, DK / 4, NH * segments), block (32, 4): a warp a (head, value row), -// its 32 lanes over the key width. -extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, +// each row (forward row m at m * rec()). Value head hv reads key head hv / (nv / nk). Grid (1, DK / 4, nv * segments), +// block (32, 4): a warp a (value head, value row), its 32 lanes over the key width. +extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP) { - const int hv = blockIdx.z % NH; - const Seg sg = segs[blockIdx.z / NH]; + const int NK = geo.nk, NV = geo.nv; + const int hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const Seg sg = segs[blockIdx.z / NV]; const int dk = threadIdx.x; const int dv = blockIdx.y * 4 + threadIdx.y; constexpr int PT = DK / 32; - float* st = reinterpret_cast(sg.rec) + (size_t)layer * REC + ((size_t)hv * DK + dv) * DK + PT * dk; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; float state[PT]; #pragma unroll for (int i = 0; i < PT; i++) state[i] = st[i]; @@ -206,12 +206,12 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bf float kc[PT], qc[PT]; #pragma unroll for (int i = 0; i < PT; i++) { - kc[i] = bf(k[(row * NH + hv) * DK + PT * dk + i]); - qc[i] = bf(q[(row * NH + hv) * DK + PT * dk + i]); + kc[i] = bf(k[(row * NK + hk) * DK + PT * dk + i]); + qc[i] = bf(q[(row * NK + hk) * DK + PT * dk + i]); } - const float vc = bf(v[(row * NH + hv) * DK + dv]); - const float gc = g[row * NH + hv]; - const float bc = bf(beta[row * NH + hv]); + const float vc = bf(v[(row * NV + hv) * DK + dv]); + const float gc = g[row * NV + hv]; + const float bc = bf(beta[row * NV + hv]); float kv_mem = 0.0f; #pragma unroll for (int i = 0; i < PT; i++) { @@ -227,9 +227,9 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bf out += state[i] * qc[i]; } out = warp_sum(out); - if (dk == 0) y[(row * NH + hv) * DK + dv] = __float2bfloat16_rn(out); + if (dk == 0) y[(row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); if (SNAP != nullptr) { - float* sn = SNAP + row * REC + ((size_t)hv * DK + dv) * DK + PT * dk; + float* sn = SNAP + row * geo.rec() + at; #pragma unroll for (int i = 0; i < PT; i++) sn[i] = state[i]; } @@ -238,16 +238,17 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bf for (int i = 0; i < PT; i++) st[i] = state[i]; } -// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear head. Grid (1, NH, rows), a warp each. -extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, +// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear value head; rows nv * DK wide. Grid (1, nv, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT) { const int lane = threadIdx.x, hv = blockIdx.y; const size_t m = blockIdx.z; + const size_t W = (size_t)geo.nv * DK; constexpr int PER = DK / 32; float yv[PER], ss = 0.0f; #pragma unroll for (int j = 0; j < PER; j++) { - yv[j] = bf(Y[(m * NH + hv) * DK + lane * PER + j]); + yv[j] = bf(Y[m * W + hv * DK + lane * PER + j]); ss += yv[j] * yv[j]; } ss = warp_sum(ss); @@ -256,8 +257,8 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(const __nv_bflo for (int j = 0; j < PER; j++) { const int d = lane * PER + j; const float x = rbf(bf(NW[d]) * (yv[j] * inv)); - const float zf = bf(Z[m * HID + hv * DK + d]); - OUT[m * HID + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); + const float zf = bf(Z[m * W + hv * DK + d]); + OUT[m * W + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); } } @@ -300,14 +301,15 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i // Queries in place at their positions; keys rotated and values copied into their segment's cache. // Grid (rows), HD threads. -extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(__nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, +extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, const int* __restrict__ rowseg, int layer) { + const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; const int pos = sg.len + (int)m - sg.row0; - __shared__ __nv_bfloat16 qs[QH * HD]; + __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; __syncthreads(); for (int h = 0; h < QH; h++) Qn[(m * QH + h) * HD + d] = rotate(qs + h * HD, d, pos); @@ -320,12 +322,13 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(__nv_bfloat16* __r } } -// Causal attention over a segment's cache, ck keys a block: grid (chunks, KVH, rows), a warp a query head of the -// kv head's group of QH / KVH. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. -extern "C" __global__ void __launch_bounds__(128) tf_q35_attn_partial(const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, - float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { +// Causal attention over a segment's cache, ck keys a block: grid (chunks, kvh, rows), block qh / kvh warps, a warp a +// query head of the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, + float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; + const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; const size_t m = blockIdx.z; @@ -374,10 +377,11 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_attn_partial(const __nv for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; } -// The chunks' partials into each (row, query head)'s output. Grid (QH, rows), a warp each. -extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(const float* __restrict__ PM, const float* __restrict__ PL, +// The chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, const float* __restrict__ PO, int nch, __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; + const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; const size_t m = blockIdx.y; const size_t base = (m * QH + qh) * nch; @@ -398,19 +402,20 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(const float* for (int i = 0; i < DPL; i++) OUT[(m * QH + qh) * HD + lane * DPL + i] = __float2bfloat16_rn(acc[i] / lsum); } -// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows. -extern "C" __global__ void tf_q35_attn_gate(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, +// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows; rows qh * HD wide. +extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, __nv_bfloat16* __restrict__ OUT, int rows) { + const size_t W = (size_t)geo.qh * HD; const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; - if (i >= (size_t)rows * HID) return; - const size_t r = i / HID; - const int d = (int)(i % HID), head = d / HD; - const float gate = bf(QRAW[r * 2 * HID + head * 2 * HD + HD + d % HD]); + if (i >= (size_t)rows * W) return; + const size_t r = i / W; + const int d = (int)(i % W), head = d / HD; + const float gate = bf(QRAW[r * 2 * W + head * 2 * HD + HD + d % HD]); const float s = rbf(1.0f / (1.0f + expf(-gate))); OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); } -// out = bf16(silu(gate) * up), n wide rows. +// out = bf16(silu(gate) * up) over `total` values. extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, __nv_bfloat16* __restrict__ OUT, long long total) { const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index fdecefed7f..ea5eaa465e 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -1,9 +1,12 @@ -//! Qwen3.5-2B on CUDA: the engine and its native family entry (the Metal engine is qwen3_5.zig). +//! The Qwen3.5 hybrids on CUDA (the 2B, Qwen3.8-27B): the engine and its native family entry (the Metal 2B engine is +//! qwen3_5.zig). pub const config = @import("config.zig"); +pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); test { _ = config; + _ = geometry; _ = @import("draw.zig"); } diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 96190737d1..316814eaf6 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,43 +1,40 @@ -//! Qwen3.5-2B on CUDA: the affine-4 projections through cuda.qlinear, the glue in qwen35.cu, one forward for any -//! mix of streams' windows. Each stream owns its recurrences, convolution tails and KV cache (Cache). +//! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the +//! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails +//! and KV cache (Cache); the widths come from config.json (cuda_geometry.zig), the weights from cuda_weights.zig. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); -const c = @import("config.zig"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_weights.zig"); +pub const Geometry = geometry.Geometry; const Affine4 = cuda.qlinear.Affine4; const DeviceBuffer = cuda.DeviceBuffer; +const head_dim = geometry.head_dim; /// Rows of one forward at most: a prompt chunk, or a shared round's windows. pub const max_rows = 128; /// Streams' windows in one forward at most. pub const max_segs = 32; /// Rows of a verify forward at most: each row's recurrences are recorded for a partial acceptance. -pub const max_verify = 32; +const most_verify = 32; /// A window's rows at most (exact: one-row steps' bits). -pub const window_rows = 16; +const most_window = 16; +/// Device bytes the recorded rows may take: a wide model records fewer rows a forward. +const snapshot_budget = 1280 << 20; /// Partial attention chunks a row at most: the chunk width grows with the window instead. const max_chunks = 64; -const deltas = c.layers - c.layers / 4; -const attentions = c.layers / 4; -pub const rec_bytes = c.linear_heads * c.linear_dim * c.linear_dim * 4; -pub const conv_bytes = (c.conv_taps - 1) * c.conv_dim * 2; -const state_bytes = deltas * (rec_bytes + conv_bytes); - -/// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. -pub fn streamBytes(cap: usize) usize { - return state_bytes + kvBytes(cap); -} - -fn kvBytes(cap: usize) usize { - return attentions * 2 * c.kv_heads * cap * c.head_dim * 2; +/// Device allocations are 256-byte aligned, as the driver's are. +pub fn aligned(n: usize) usize { + return std.mem.alignForward(usize, n, 256); } -/// qwen35.cu's Seg. +/// qwen35.cu's Geo and Seg. +const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int }; const Seg = extern struct { rec: u64, conv: u64, kv: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int }; -/// One stream's state on the device and the logits of its newest row on the host. +/// One stream's state on the device and the logits of its newest prompt chunk on the host. pub const Cache = struct { state: DeviceBuffer, kv: DeviceBuffer, @@ -48,12 +45,20 @@ pub const Cache = struct { last: ?struct { start: usize, row0: usize, rows: usize } = null, pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { - var state = try DeviceBuffer.alloc(e.d, state_bytes); + var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); errdefer state.free(); try state.fill8(0, null); - var kv = try DeviceBuffer.alloc(e.d, kvBytes(cap)); + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); errdefer kv.free(); - return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, c.vocab) }; + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). + pub fn reset(self: *Cache, e: *const Engine) !void { + self.len = 0; + self.last = null; + try self.state.fill8(0, e.stream.handle); } pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { @@ -62,14 +67,14 @@ pub const Cache = struct { gpa.free(self.logits); } - fn seg(self: *const Cache, row0: usize, rows: usize) Seg { - return .{ .rec = self.state.ptr, .conv = self.state.ptr + deltas * rec_bytes, .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; + fn seg(self: *const Cache, g: Geometry, row0: usize, rows: usize) Seg { + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; } }; -const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; -const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; -const Layer = struct { +pub const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; +pub const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; +pub const Layer = struct { input_norm: u64, post_norm: u64, gate: Affine4, @@ -109,19 +114,28 @@ pub const Draw = struct { rule: ?cuda.sample.Rule, position: u64 }; pub const Engine = struct { gpa: std.mem.Allocator, d: *const cuda.Driver, + g: Geometry, + geo: Geo, stream: cuda.Stream, mods: [6]cuda.Module, draws: cuda.sample.Draws, affine: cuda.qlinear.Affine4Kernels, fns: Fns, - held: std.ArrayList(DeviceBuffer) = .empty, // weights and scratch, freed together + /// Every weight and scratch buffer in one allocation: the WDDM driver's launch cost grows with a process's + /// allocation count, and a 64-layer model would have thousands. + arena: DeviceBuffer, + arena_at: usize = 0, + stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there + /// Rows a verify forward records (snapshot_budget over a row's states), and a window's rows at most. + verify_rows: usize, + window_rows: usize, // weights embed_w: u64, embed_s: u64, embed_b: u64, head: Affine4, final_norm: u64, - layers: [c.layers]Layer, + layers: []Layer, // scratch, max_rows rows each h: u64, x: u64, @@ -134,7 +148,7 @@ pub const Engine = struct { q: u64, k: u64, v: u64, - g: u64, + g_: u64, beta: u64, y: u64, mix: u64, @@ -150,35 +164,44 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, + rec_snap: u64, // recorded recurrences: verify_rows rows a linear layer + conv_snap: u64, // recorded convolution tails, the same // the per-forward tables: ids, segments, each row's segment and the draws' rules, staged in pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, host_logits: cuda.HostBuffer, /// The kept rows' device draws, when forward was given them. - tokens: [max_verify]u32 = undefined, - rec_snap: u64, // recorded recurrences: max_verify rows a linear layer - conv_snap: u64, // recorded convolution tails, the same + tokens: [most_verify]u32 = undefined, const ids_at = 0; const segs_at = max_rows * 4; const rowseg_at = segs_at + max_segs * @sizeOf(Seg); const rules_at = rowseg_at + max_rows * 4; - const meta_at = rules_at + max_verify * @sizeOf(cuda.sample.Rule); - const tokens_at = meta_at + max_verify * 4; - const tables_len = tokens_at + max_verify * 4; + const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); + const tokens_at = meta_at + most_verify * 4; + const tables_len = tokens_at + most_verify * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; - const config = try c.Config.read(gpa, io, dir); + const g = try Geometry.read(gpa, io, dir); const d = ctx.d; const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.* = undefined; e.gpa = gpa; e.d = d; - e.held = .empty; + e.g = g; + e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads) }; + e.verify_rows = std.math.clamp(snapshot_budget / (g.deltas() * (g.recBytes() + g.convBytes())), 1, most_verify); + e.window_rows = @min(most_window, e.verify_rows); + e.arena_at = 0; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); + e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); + errdefer e.arena.free(); + try e.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); + errdefer e.stage.free(); const kk = cuda.kernels; const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample }; var loaded: usize = 0; @@ -193,16 +216,26 @@ pub const Engine = struct { errdefer e.affine.deinit(); e.draws = try cuda.sample.Draws.resolve(e.mods[5]); inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); - errdefer e.freeHeld(); + e.layers = try gpa.alloc(Layer, g.layers); + errdefer gpa.free(e.layers); var ck = try core.Checkpoint.openModel(gpa, io, dir); defer ck.close(); - try e.loadWeights(&ck, config.tied); - try e.allocScratch(); + try weights.load(e, &ck); + _ = try e.scratch(.allocate); e.tables = try DeviceBuffer.alloc(d, tables_len); errdefer e.tables.free(); e.staged = try cuda.HostBuffer.alloc(d, tables_len); errdefer e.staged.free(); - e.host_logits = try cuda.HostBuffer.alloc(d, max_verify * c.vocab * 2); + e.host_logits = try cuda.HostBuffer.alloc(d, most_verify * g.vocab * 2); + errdefer e.host_logits.free(); + // every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds + // (a 16 GiB model: a 20 s launch, then a 3 s one), which no request should pay + var warm = try Cache.init(gpa, e, 64); + defer warm.deinit(gpa); + const ids: [32]u32 = @splat(0); + try e.forward(&.{.{ .cache = &warm, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs + try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = null, .position = 33 }} }); + try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = cuda.sample.rule(1, 0.7, 0.95, -std.math.inf(f64), 20), .position = 34 }} }); return e; } @@ -210,117 +243,62 @@ pub const Engine = struct { e.host_logits.free(); e.staged.free(); e.tables.free(); - e.freeHeld(); + e.stage.free(); + e.arena.free(); + e.gpa.free(e.layers); e.affine.deinit(); for (&e.mods) |*m| m.unload(); e.stream.deinit(); e.gpa.destroy(e); } - fn freeHeld(e: *Engine) void { - for (e.held.items) |*b| b.free(); - e.held.deinit(e.gpa); - } - - fn own(e: *Engine, b: DeviceBuffer) !u64 { - try e.held.append(e.gpa, b); - return b.ptr; - } - - fn alloc(e: *Engine, bytes: usize) !u64 { - var b = try DeviceBuffer.alloc(e.d, bytes); - errdefer b.free(); - return e.own(b); + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. + pub fn streamBytes(e: *const Engine, cap: usize) usize { + return e.g.stateBytes() + e.g.kvBytes(cap); } - fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { - var b = try DeviceBuffer.fromHost(e.d, t.bytes); - errdefer b.free(); - return e.own(b); + /// The next `bytes` of the arena, 256-byte aligned. + pub fn alloc(e: *Engine, bytes: usize) !u64 { + const at = aligned(e.arena_at); + if (at + bytes > e.arena.len) return error.ArenaExhausted; + e.arena_at = at + bytes; + return e.arena.ptr + at; } - fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { - var buf: [160]u8 = undefined; - return e.put(try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); - } - - /// An MLX affine-4 projection packed into qlinear's tiles; the checkpoint's own layout leaves the device after. - fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { - var buf: [160]u8 = undefined; - var raw: [3]DeviceBuffer = undefined; - var got: usize = 0; - defer for (raw[0..got]) |*b| b.free(); - const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / c.group }, .{ "biases", .bf16, k / c.group } }; - for (parts) |p| { - const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); - raw[got] = try DeviceBuffer.fromHost(e.d, t.bytes); - got += 1; - } - const l = Affine4.layout(n, k); - const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; - try e.affine.pack(e.stream, raw[0].ptr, raw[1].ptr, raw[2].ptr, q); - try e.stream.synchronize(); - return q; - } - - fn loadWeights(e: *Engine, ck: *core.Checkpoint, tied: bool) !void { - const m = "language_model.model."; - e.embed_w = try e.tensor(ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ c.vocab, c.hidden / 8 }); - e.embed_s = try e.tensor(ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); - e.embed_b = try e.tensor(ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); - e.head = if (tied) try e.linear(ck, m, "embed_tokens", c.vocab, c.hidden) else try e.linear(ck, "language_model.", "lm_head", c.vocab, c.hidden); - e.final_norm = try e.tensor(ck, "{s}norm.weight", .{m}, .bf16, &.{c.hidden}); - for (&e.layers, 0..) |*l, i| { - var pbuf: [96]u8 = undefined; - const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); - l.* = .{ - .input_norm = try e.tensor(ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{c.hidden}), - .post_norm = try e.tensor(ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{c.hidden}), - .gate = try e.linear(ck, p, "mlp.gate_proj", c.intermediate, c.hidden), - .up = try e.linear(ck, p, "mlp.up_proj", c.intermediate, c.hidden), - .down = try e.linear(ck, p, "mlp.down_proj", c.hidden, c.intermediate), - .mixer = undefined, - }; - if (c.linear(i)) { - l.mixer = .{ .delta = .{ - .qkv = try e.linear(ck, p, "linear_attn.in_proj_qkv", c.conv_dim, c.hidden), - .z = try e.linear(ck, p, "linear_attn.in_proj_z", c.hidden, c.hidden), - .a = try e.linear(ck, p, "linear_attn.in_proj_a", c.linear_heads, c.hidden), - .b = try e.linear(ck, p, "linear_attn.in_proj_b", c.linear_heads, c.hidden), - .out = try e.linear(ck, p, "linear_attn.out_proj", c.hidden, c.hidden), - .conv = try e.tensor(ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ c.conv_dim, c.conv_taps, 1 }), - .norm = try e.tensor(ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{c.linear_dim}), - .a_log = try e.tensor(ck, "{s}linear_attn.A_log", .{p}, .f32, &.{c.linear_heads}), - .dt_bias = try e.tensor(ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{c.linear_heads}), - } }; - } else { - l.mixer = .{ .attention = .{ - .q = try e.linear(ck, p, "self_attn.q_proj", 2 * c.hidden, c.hidden), - .k = try e.linear(ck, p, "self_attn.k_proj", c.kv_heads * c.head_dim, c.hidden), - .v = try e.linear(ck, p, "self_attn.v_proj", c.kv_heads * c.head_dim, c.hidden), - .o = try e.linear(ck, p, "self_attn.o_proj", c.hidden, c.hidden), - .q_norm = try e.tensor(ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{c.head_dim}), - .k_norm = try e.tensor(ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{c.head_dim}), - } }; + /// The scratch rows' bytes (`measure`: summed, nothing taken), or the rows themselves from the arena. + fn scratch(e: *Engine, mode: enum { measure, allocate }) !usize { + var total: usize = 0; + const Take = struct { + e: *Engine, + mode: @TypeOf(mode), + total: *usize, + fn alloc(t: @This(), bytes: usize) !u64 { + t.total.* += aligned(bytes); + return if (t.mode == .measure) 0 else try t.e.alloc(bytes); } - } - } - - fn allocScratch(e: *Engine) !void { + }; + const t = Take{ .e = e, .mode = mode, .total = &total }; + const g = e.g; const R = max_rows; - inline for (.{ "h", "x", "r", "z", "q", "k", "v", "y", "mix", "qn" }) |name| @field(e, name) = try e.alloc(R * c.hidden * 2); - inline for (.{ "qkv", "gate", "up", "act" }) |name| @field(e, name) = try e.alloc(R * c.intermediate * 2); - e.xs = try e.alloc(R * (c.intermediate / 64) * 4); - inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try e.alloc(R * c.linear_heads * 2); - e.g = try e.alloc(R * c.linear_heads * 4); - e.qraw = try e.alloc(R * 2 * c.hidden * 2); - inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try e.alloc(R * c.kv_heads * c.head_dim * 2); - e.pm = try e.alloc(R * c.query_heads * max_chunks * 4); - e.pl = try e.alloc(R * c.query_heads * max_chunks * 4); - e.po = try e.alloc(R * c.query_heads * max_chunks * c.head_dim * 4); - e.logits = try e.alloc(max_verify * c.vocab * 2); - e.rec_snap = try e.alloc(deltas * max_verify * rec_bytes); - e.conv_snap = try e.alloc(deltas * max_verify * conv_bytes); + const wide = @max(g.valueWidth(), g.attnWidth(), g.hidden); // y and mix: either mixer's rows + inline for (.{ "h", "x", "r" }) |name| @field(e, name) = try t.alloc(R * g.hidden * 2); + inline for (.{ "y", "mix", "z", "v" }) |name| @field(e, name) = try t.alloc(R * wide * 2); + inline for (.{ "q", "k" }) |name| @field(e, name) = try t.alloc(R * g.k_heads * geometry.dk * 2); + e.qkv = try t.alloc(R * g.convWidth() * 2); + inline for (.{ "gate", "up", "act" }) |name| @field(e, name) = try t.alloc(R * g.intermediate * 2); + e.xs = try t.alloc(R * (@max(g.intermediate, wide) / 64) * 4); + inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try t.alloc(R * g.v_heads * 2); + e.g_ = try t.alloc(R * g.v_heads * 4); + e.qraw = try t.alloc(R * 2 * g.attnWidth() * 2); + e.qn = try t.alloc(R * g.attnWidth() * 2); + inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try t.alloc(R * g.kv_heads * head_dim * 2); + e.pm = try t.alloc(R * g.q_heads * max_chunks * 4); + e.pl = try t.alloc(R * g.q_heads * max_chunks * 4); + e.po = try t.alloc(R * g.q_heads * max_chunks * head_dim * 4); + e.logits = try t.alloc(most_verify * g.vocab * 2); + e.rec_snap = try t.alloc(g.deltas() * e.verify_rows * g.recBytes()); + e.conv_snap = try t.alloc(g.deltas() * e.verify_rows * g.convBytes()); + return total; } fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { @@ -343,11 +321,12 @@ pub const Engine = struct { } fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { - try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, @as(u64, 0), c.eps }); + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)) }); } /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { + const g = e.g; var total: usize = 0; var widest: usize = 0; for (parts) |p| { @@ -357,10 +336,10 @@ pub const Engine = struct { } if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; const record = keep != .last; - if (record and total > max_verify) return error.InvalidQwenSegments; + if (record and total > e.verify_rows) return error.InvalidQwenSegments; if (!record and parts.len != 1) return error.InvalidQwenSegments; if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; - for (ids) |id| if (id >= c.vocab) return error.InvalidQwenToken; + for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; const heads: usize = if (record) total else 1; // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps const prompt = !record; @@ -371,7 +350,7 @@ pub const Engine = struct { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); var row0: usize = 0; for (parts, 0..) |p, i| { - segs[i] = p.cache.seg(row0, p.rows); + segs[i] = p.cache.seg(g, row0, p.rows); for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; row0 += p.rows; @@ -389,11 +368,13 @@ pub const Engine = struct { const t_segs = e.tables.ptr + segs_at; const t_rows = e.tables.ptr + rowseg_at; const nseg = parts.len; + const geo = e.geo; // attention chunks: at most max_chunks a row, each a multiple of 128 keys const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); const nch = (widest + ck - 1) / ck; + const group = g.q_heads / g.kv_heads; - try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h }); + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)) }); var di: usize = 0; var ai: c_int = 0; for (e.layers, 0..) |l, i| { @@ -401,29 +382,30 @@ pub const Engine = struct { switch (l.mixer) { .delta => |m| { const layer: c_int = @intCast(di); - const rec_snap: u64 = if (record) e.rec_snap + di * max_verify * rec_bytes else 0; - const conv_snap: u64 = if (record) e.conv_snap + di * max_verify * conv_bytes else 0; + const rec_snap: u64 = if (record) e.rec_snap + di * e.verify_rows * g.recBytes() else 0; + const conv_snap: u64 = if (record) e.conv_snap + di * e.verify_rows * g.convBytes() else 0; try e.proj(m.qkv, e.x, e.qkv, total, prompt); try e.proj(m.z, e.x, e.z, total, prompt); try e.proj(m.a, e.x, e.a, total, prompt); try e.proj(m.b, e.x, e.b, total, prompt); - try e.run(e.fns.gdn_pre, .{ 1, 3 * c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g, e.beta }); - try e.run(e.fns.conv_commit, .{ (c.conv_dim + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ e.qkv, t_segs, layer, conv_snap }); - try e.run(e.fns.gdn_chain, .{ 1, c.linear_dim / 4, c.linear_heads * nseg }, .{ 32, 4, 1 }, .{ e.q, e.k, e.v, e.g, e.beta, t_segs, layer, e.y, rec_snap }); - try e.run(e.fns.gdn_post, .{ 1, c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.y, e.z, m.norm, c.eps, e.mix }); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap }); + try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix }); try e.proj(m.out, e.mix, e.r, total, prompt); di += 1; }, .attention => |m| { + const aw = g.attnWidth(); try e.proj(m.q, e.x, e.qraw, total, prompt); try e.proj(m.k, e.x, e.kraw, total, prompt); try e.proj(m.v, e.x, e.vraw, total, prompt); - try e.run(e.fns.head_norm, .{ total * c.query_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, c.query_heads), @as(c_int, 2 * c.hidden), @as(c_int, 2 * c.head_dim), c.eps }); - try e.run(e.fns.head_norm, .{ total * c.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, c.kv_heads), @as(c_int, c.kv_heads * c.head_dim), @as(c_int, c.head_dim), c.eps }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ c.head_dim, 1, 1 }, .{ e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); - try e.run(e.fns.attn_partial, .{ nch, c.kv_heads, total }, .{ 128, 1, 1 }, .{ e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); - try e.run(e.fns.attn_merge, .{ c.query_heads, total, 1 }, .{ 32, 1, 1 }, .{ e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); - try e.run(e.fns.attn_gate, .{ (total * c.hidden + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); + try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps }); + try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); + try e.run(e.fns.attn_partial, .{ nch, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); + try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; }, @@ -431,29 +413,29 @@ pub const Engine = struct { try e.norm(true, l.post_norm, total); try e.proj(l.gate, e.x, e.gate, total, prompt); try e.proj(l.up, e.x, e.up, total, prompt); - const n: c_longlong = @intCast(total * c.intermediate); - try e.run(e.fns.mlp_act, .{ (total * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); + const n: c_longlong = @intCast(total * g.intermediate); + try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); try e.proj(l.down, e.act, e.r, total, prompt); } try e.norm(true, e.final_norm, total); // the head over the kept rows - try e.proj(e.head, e.x + (total - heads) * c.hidden * 2, e.logits, heads, false); + try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); for (parts) |p| p.cache.len += p.rows; switch (keep) { .draws => |pk| { const toks = e.tables.ptr + tokens_at; for (pk, 0..) |p, i| { - const row = e.logits + i * c.vocab * 2; + const row = e.logits + i * g.vocab * 2; if (p.rule != null) { - try e.draws.draw(e.stream, row, c.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); - } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, c.vocab), toks + i * 4 }); + try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); } try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); }, .logits, .last => { - const out = e.host_logits.bytes[0 .. heads * c.vocab * 2]; + const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); @@ -463,21 +445,24 @@ pub const Engine = struct { /// Row `i`'s logits from the newest forward that kept logits. pub fn rowLogits(e: *const Engine, i: usize) []const u16 { - return e.host_logits.slice(u16)[i * c.vocab ..][0..c.vocab]; + return e.host_logits.slice(u16)[i * e.g.vocab ..][0..e.g.vocab]; } /// Keeps the first `rows` of the cache's newest window: its recurrences and convolution tails as they were after /// that row (the forward recorded them), its length at that row. The window's later keys stay as unread slots. pub fn rewind(e: *Engine, cache: *Cache, rows: usize) !void { + const g = e.g; const last = cache.last orelse return error.NothingToKeep; if (rows == 0 or rows > last.rows) return error.NothingToKeep; cache.last = null; cache.len = last.start + rows; if (rows == last.rows) return; const row = last.row0 + rows - 1; - for (0..deltas) |di| { - try cache.state.copyFrom(di * rec_bytes, e.rec_snap + (di * max_verify + row) * rec_bytes, rec_bytes, e.stream.handle); - try cache.state.copyFrom(deltas * rec_bytes + di * conv_bytes, e.conv_snap + (di * max_verify + row) * conv_bytes, conv_bytes, e.stream.handle); + const rb = g.recBytes(); + const cb = g.convBytes(); + for (0..g.deltas()) |di| { + try cache.state.copyFrom(di * rb, e.rec_snap + (di * e.verify_rows + row) * rb, rb, e.stream.handle); + try cache.state.copyFrom(g.deltas() * rb + di * cb, e.conv_snap + (di * e.verify_rows + row) * cb, cb, e.stream.handle); } try e.stream.synchronize(); } diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig new file mode 100644 index 0000000000..34207bacb5 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -0,0 +1,177 @@ +//! The widths of a Qwen3.5 hybrid text model (the 2B, Qwen3.8-27B) as the CUDA engine reads them from config.json: +//! what varies is read, what qwen35.cu fixes (head widths, taps, rotary layout, 4-bit groups of 64) is checked. +const std = @import("std"); + +pub const max_layers = 128; +/// qwen35.cu's fixed widths. +pub const dk = 128; +pub const taps = 4; +pub const head_dim = 256; +pub const max_query_heads = 32; + +pub const Geometry = struct { + hidden: usize, + intermediate: usize, + layers: usize, + vocab: usize, + k_heads: usize, // linear-attention key heads + v_heads: usize, // linear-attention value heads + q_heads: usize, // attention query heads + kv_heads: usize, // attention key/value heads + eps: f32, + tied: bool, + context: usize, + linear: [max_layers]bool, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Geometry { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + /// Value-head width of the gated delta rule's rows (and of z and the out projection's input). + pub fn valueWidth(g: Geometry) usize { + return g.v_heads * dk; + } + + /// Channels the causal convolution sees: q, k and v. + pub fn convWidth(g: Geometry) usize { + return (2 * g.k_heads + g.v_heads) * dk; + } + + /// Attention output width (q_heads of head_dim), the o projection's input. + pub fn attnWidth(g: Geometry) usize { + return g.q_heads * head_dim; + } + + pub fn deltas(g: Geometry) usize { + var n: usize = 0; + for (g.linear[0..g.layers]) |l| n += @intFromBool(l); + return n; + } + + pub fn attentions(g: Geometry) usize { + return g.layers - g.deltas(); + } + + /// Bytes of one linear layer's recurrence (fp32) and convolution tail (bf16). + pub fn recBytes(g: Geometry) usize { + return g.v_heads * dk * dk * 4; + } + + pub fn convBytes(g: Geometry) usize { + return (taps - 1) * g.convWidth() * 2; + } + + /// A stream's recurrences and tails. + pub fn stateBytes(g: Geometry) usize { + return g.deltas() * (g.recBytes() + g.convBytes()); + } + + /// A stream's keys and values at a `cap`-token window. + pub fn kvBytes(g: Geometry, cap: usize) usize { + return g.attentions() * 2 * g.kv_heads * cap * head_dim * 2; + } +}; + +fn object(v: std.json.Value) !std.json.ObjectMap { + return if (v == .object) v.object else error.BadQwenConfig; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) !usize { + const v = o.get(key) orelse return error.BadQwenConfig; + if (v != .integer or v.integer < 1) return error.BadQwenConfig; + return @intCast(v.integer); +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + return switch (o.get(key) orelse return error.BadQwenConfig) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadQwenConfig, + }; +} + +fn is(o: std.json.ObjectMap, key: []const u8, comptime T: type, want: T) !void { + const v = o.get(key) orelse return error.UnsupportedQwenConfig; + const ok = switch (T) { + bool => v == .bool and v.bool == want, + []const u8 => v == .string and std.mem.eql(u8, v.string, want), + usize => v == .integer and v.integer == want, + else => @compileError("type"), + }; + if (!ok) return error.UnsupportedQwenConfig; +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer doc.deinit(); + const root = try object(doc.value); + try is(root, "model_type", []const u8, "qwen3_5"); + const t = try object(root.get("text_config") orelse return error.BadQwenConfig); + try is(t, "model_type", []const u8, "qwen3_5_text"); + try is(t, "attention_bias", bool, false); + try is(t, "attn_output_gate", bool, true); + try is(t, "hidden_act", []const u8, "silu"); + try is(t, "head_dim", usize, head_dim); + try is(t, "linear_key_head_dim", usize, dk); + try is(t, "linear_value_head_dim", usize, dk); + try is(t, "linear_conv_kernel_dim", usize, taps); + const rope = try object(t.get("rope_parameters") orelse return error.UnsupportedQwenRotary); + try is(rope, "rope_type", []const u8, "default"); + if (try number(rope, "rope_theta") != 10000000 or try number(rope, "partial_rotary_factor") != 0.25) return error.UnsupportedQwenRotary; + const q = try object(root.get("quantization") orelse return error.UnsupportedQwenQuantization); + if (try integer(q, "bits") != 4 or try integer(q, "group_size") != 64) return error.UnsupportedQwenQuantization; + if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return error.UnsupportedQwenQuantization; + var g: Geometry = .{ + .hidden = try integer(t, "hidden_size"), + .intermediate = try integer(t, "intermediate_size"), + .layers = try integer(t, "num_hidden_layers"), + .vocab = try integer(t, "vocab_size"), + .k_heads = try integer(t, "linear_num_key_heads"), + .v_heads = try integer(t, "linear_num_value_heads"), + .q_heads = try integer(t, "num_attention_heads"), + .kv_heads = try integer(t, "num_key_value_heads"), + .eps = @floatCast(try number(t, "rms_norm_eps")), + .tied = blk: { + const v = t.get("tie_word_embeddings") orelse root.get("tie_word_embeddings") orelse break :blk false; + if (v != .bool) return error.BadQwenConfig; + break :blk v.bool; + }, + .context = try integer(t, "max_position_embeddings"), + .linear = @splat(false), + }; + if (g.layers > max_layers or g.hidden % 256 != 0 or g.hidden > 8192 or g.intermediate % 64 != 0) return error.UnsupportedQwenGeometry; + if (g.v_heads % g.k_heads != 0 or g.q_heads % g.kv_heads != 0 or g.q_heads > max_query_heads) return error.UnsupportedQwenGeometry; + const types = t.get("layer_types") orelse return error.UnsupportedQwenGeometry; + if (types != .array or types.array.items.len != g.layers) return error.UnsupportedQwenGeometry; + for (types.array.items, 0..) |kind, i| { + if (kind != .string) return error.UnsupportedQwenGeometry; + if (std.mem.eql(u8, kind.string, "linear_attention")) g.linear[i] = true else if (!std.mem.eql(u8, kind.string, "full_attention")) return error.UnsupportedQwenGeometry; + } + return g; +} + +test "the 27B's widths: 48 value heads over 16 key heads, 24 query heads over 4" { + const text = + \\{"model_type":"qwen3_5","quantization":{"group_size":64,"bits":4,"mode":"affine"},"tie_word_embeddings":false, + \\ "text_config":{"model_type":"qwen3_5_text","attention_bias":false,"attn_output_gate":true,"hidden_act":"silu", + \\ "head_dim":256,"linear_key_head_dim":128,"linear_value_head_dim":128,"linear_conv_kernel_dim":4, + \\ "hidden_size":5120,"intermediate_size":17408,"num_hidden_layers":4,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":48,"num_attention_heads":24,"num_key_value_heads":4,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144,"tie_word_embeddings":false, + \\ "layer_types":["linear_attention","linear_attention","linear_attention","full_attention"], + \\ "rope_parameters":{"rope_type":"default","rope_theta":10000000,"partial_rotary_factor":0.25}}} + ; + const g = try parse(std.testing.allocator, text); + try std.testing.expectEqual(@as(usize, 3), g.deltas()); + try std.testing.expectEqual(@as(usize, 10240), g.convWidth()); + try std.testing.expectEqual(@as(usize, 6144), g.attnWidth()); + try std.testing.expectEqual(@as(usize, 48 * 128 * 128 * 4), g.recBytes()); + try std.testing.expect(!g.tied); + const wide = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"head_dim\":256", "\"head_dim\":128"); + defer std.testing.allocator.free(wide); + try std.testing.expectError(error.UnsupportedQwenConfig, parse(std.testing.allocator, wide)); +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 4995a82b45..4b46a5a896 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -1,5 +1,6 @@ -//! Qwen3.5-2B's CUDA lane backend behind the native family interface (native/cuda.zig's registry): windows of up -//! to 16 rows (the lane core's copies), every stream's window in one forward, prompts in chunks, draws on the device. +//! The Qwen3.5 hybrids' CUDA lane backend (the 2B, Qwen3.8-27B) behind the native family interface (native/cuda.zig's +//! registry): windows of up to 16 rows (the lane core's copies), every stream's window in one forward, prompts in +//! chunks, draws on the device. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -26,24 +27,44 @@ pub const Loaded = struct { free_streams: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, + /// Called with the admitted stream count: that many caches allocated and touched now, since the driver commits + /// memory on first touch (seconds for a wide model's keys), and a request should not pay it. + reserve: *const fn (*anyopaque, u32) anyerror!void, }; pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; const e = try E.Engine.init(gpa, io, ctx, dir); errdefer e.deinit(); + if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary const l = try gpa.create(Lanes); l.* = .{ .gpa = gpa, .e = e, .capacity = o.context }; return .{ .backend = l.backend(), - .facts = .{ .exact_width = E.window_rows, .hidden_rows = true, .batch_rows = E.max_verify, .max_streams = E.max_segs, .streams_exact = true }, - .rows = E.window_rows, - .stream_bytes = E.streamBytes(o.context), + .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, + .rows = @intCast(e.window_rows), + .stream_bytes = e.streamBytes(o.context), .ctx = l, .deinit = release, + .reserve = reserve, }; } +fn reserve(p: *anyopaque, streams: u32) anyerror!void { + const l: *Lanes = @ptrCast(@alignCast(p)); + try l.pool.ensureTotalCapacity(l.gpa, streams); + while (l.pool.items.len < streams) { + const cache = try l.gpa.create(E.Cache); + errdefer l.gpa.destroy(cache); + cache.* = try E.Cache.init(l.gpa, l.e, l.capacity); + errdefer cache.deinit(l.gpa); + // a row through it: kernels have to have referenced a buffer once before the driver's stall is behind it + try l.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{0}, .last); + try cache.reset(l.e); + l.pool.appendAssumeCapacity(cache); + } +} + /// A request this engine refuses, in words; null: none of its own. pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { @@ -70,18 +91,40 @@ const Lanes = struct { e: *E.Engine, capacity: usize, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. + pool: std.ArrayList(*E.Cache) = .empty, fn deinit(self: *Lanes) void { var it = self.caches.valueIterator(); - while (it.next()) |cache| { - cache.*.deinit(self.gpa); - self.gpa.destroy(cache.*); - } + while (it.next()) |cache| self.drop(cache.*); self.caches.deinit(self.gpa); + for (self.pool.items) |cache| self.drop(cache); + self.pool.deinit(self.gpa); self.e.deinit(); self.gpa.destroy(self); } + fn drop(self: *Lanes, cache: *E.Cache) void { + cache.deinit(self.gpa); + self.gpa.destroy(cache); + } + + /// A cache for a new stream: a pooled one reset, else one more (a window beyond the reserve). + fn take(self: *Lanes) !*E.Cache { + if (self.pool.pop()) |cache| { + try cache.reset(self.e); + return cache; + } + const cache = try self.gpa.create(E.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try E.Cache.init(self.gpa, self.e, self.capacity); + return cache; + } + + fn give(self: *Lanes, cache: *E.Cache) void { + self.pool.append(self.gpa, cache) catch self.drop(cache); + } + fn backend(self: *Lanes) be.Backend { return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; } @@ -98,10 +141,8 @@ const Lanes = struct { const self = cast(ptr); const ids = stream.prompt(); if (ids.len == 0 or ids.len + stream.max_new > self.capacity or self.caches.contains(stream)) return error.InvalidQwenPrompt; - const cache = try self.gpa.create(E.Cache); - errdefer self.gpa.destroy(cache); - cache.* = try E.Cache.init(self.gpa, self.e, @min(self.capacity, ids.len + stream.max_new + 1)); - errdefer cache.deinit(self.gpa); + const cache = try self.take(); + errdefer self.give(cache); var at: usize = 0; while (at < ids.len) { var end: usize = @min(ids.len, at + E.max_rows); @@ -150,7 +191,7 @@ const Lanes = struct { if (windows.len != out.len) return error.InvalidQwenWindows; for (windows, out) |w, o| { const rows = w.rows(); - if (w.parents != null or w.held != 0 or rows > E.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.parents != null or w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; const cache = try self.cacheOf(w.stream); @@ -159,12 +200,12 @@ const Lanes = struct { var at: usize = 0; while (at < windows.len) { var parts: [E.max_segs]E.Part = undefined; - var ids: [E.max_verify]u32 = undefined; - var picks: [E.max_verify]E.Draw = undefined; + var ids: [E.max_rows]u32 = undefined; + var picks: [E.max_rows]E.Draw = undefined; var n: usize = 0; var total: usize = 0; var masked = false; - while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= E.max_verify) : (n += 1) { + while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= self.e.verify_rows) : (n += 1) { const w = windows[at + n]; parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows() }; ids[total] = w.pending; @@ -203,9 +244,6 @@ const Lanes = struct { fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { const self = cast(ptr); - if (self.caches.fetchRemove(stream)) |entry| { - entry.value.deinit(self.gpa); - self.gpa.destroy(entry.value); - } + if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); } }; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig new file mode 100644 index 0000000000..a105b876f4 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -0,0 +1,143 @@ +//! The checkpoint's weights onto the device: each affine-4 projection staged as stored and packed into qlinear's +//! tiles, the vectors as they are, all from the engine's one arena (cuda_engine.zig). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Geometry = geometry.Geometry; +const Engine = engine.Engine; +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const aligned = engine.aligned; +const head_dim = geometry.head_dim; + +/// The packed tiles' bytes of an n-by-k affine-4 projection (qlinear's words, scales and biases). +fn packedBytes(n: usize, k: usize) usize { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// The checkpoint's own bytes of an n-by-k affine-4 projection: words, then its scales and biases. +fn rawBytes(n: usize, k: usize) usize { + return n * k / 8 * 4 + 2 * (n * k / 64 * 2); +} + +/// What `load` takes from the arena: every packed projection and vector, each allocation aligned. +pub fn total(g: Geometry) usize { + var t: usize = 0; + t += aligned(g.vocab * g.hidden / 8 * 4) + 2 * aligned(g.vocab * g.hidden / 64 * 2); // the embedding as stored + t += packedBytes(g.vocab, g.hidden); // the head, tied or not + t += aligned(g.hidden * 2); + for (g.linear[0..g.layers]) |lin| { + t += 2 * aligned(g.hidden * 2) + 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); + if (lin) { + t += packedBytes(g.convWidth(), g.hidden) + packedBytes(g.valueWidth(), g.hidden) + 2 * packedBytes(g.v_heads, g.hidden) + packedBytes(g.hidden, g.valueWidth()); + t += aligned(g.convWidth() * geometry.taps * 2) + aligned(geometry.dk * 2) + aligned(g.v_heads * 4) + aligned(g.v_heads * 2); + } else { + t += packedBytes(2 * g.attnWidth(), g.hidden) + 2 * packedBytes(g.kv_heads * head_dim, g.hidden) + packedBytes(g.hidden, g.attnWidth()); + t += 2 * aligned(head_dim * 2); + } + } + return t; +} + +/// The largest projection's bytes as stored: the staging buffer the load packs each one from. +pub fn stageBytes(g: Geometry) usize { + var most: usize = 0; + for ([_][2]usize{ .{ g.vocab, g.hidden }, .{ g.intermediate, g.hidden }, .{ g.hidden, g.intermediate }, .{ g.convWidth(), g.hidden }, .{ g.valueWidth(), g.hidden }, .{ 2 * g.attnWidth(), g.hidden } }) |nk| most = @max(most, rawBytes(nk[0], nk[1])); + return most; +} + +fn upload(e: *Engine, bytes: []const u8) !u64 { + const p = try e.alloc(bytes.len); + try e.arena.upload(p - e.arena.ptr, bytes); + return p; +} + +fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { + return upload(e, t.bytes); +} + +/// A vector the kernels read as f32, stored as f32 (the 2B) or bf16 (the 27B's A_log). +fn f32vec(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, n: usize) !u64 { + var buf: [160]u8 = undefined; + const t = try ck.get(try std.fmt.bufPrint(&buf, fmt, args)); + if (t.is(.f32, &.{n})) return put(e, t); + if (!t.is(.bf16, &.{n})) return error.UnexpectedTensor; + const words = try e.gpa.alloc(u32, n); + defer e.gpa.free(words); + for (words, std.mem.bytesAsSlice(u16, t.bytes)) |*w, h| w.* = @as(u32, h) << 16; + return upload(e, std.mem.sliceAsBytes(words)); +} + +fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [160]u8 = undefined; + return put(e, try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); +} + +/// An MLX affine-4 projection packed into qlinear's tiles from the staging buffer, which holds it as stored. +fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var raw: [3]u64 = undefined; + var at: usize = 0; + const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (parts, &raw) |p, *r| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + try e.stage.upload(at, t.bytes); + r.* = e.stage.ptr + at; + at += t.bytes.len; + } + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], q); + try e.stream.synchronize(); // the next projection's staging overwrites this one's + return q; +} + +pub fn load(e: *Engine, ck: *core.Checkpoint) !void { + const g = e.g; + const m = "language_model.model."; + e.embed_w = try tensor(e, ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ g.vocab, g.hidden / 8 }); + e.embed_s = try tensor(e, ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); + e.embed_b = try tensor(e, ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); + e.head = if (g.tied) try linear(e, ck, m, "embed_tokens", g.vocab, g.hidden) else try linear(e, ck, "language_model.", "lm_head", g.vocab, g.hidden); + e.final_norm = try tensor(e, ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); + const conv = g.convWidth(); + const vw = g.valueWidth(); + for (e.layers, 0..) |*l, i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + l.* = .{ + .input_norm = try tensor(e, ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}), + .post_norm = try tensor(e, ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}), + .gate = try linear(e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), + .up = try linear(e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), + .down = try linear(e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + .mixer = undefined, + }; + if (g.linear[i]) { + l.mixer = .{ .delta = .{ + .qkv = try linear(e, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), + .z = try linear(e, ck, p, "linear_attn.in_proj_z", vw, g.hidden), + .a = try linear(e, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), + .b = try linear(e, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), + .out = try linear(e, ck, p, "linear_attn.out_proj", g.hidden, vw), + .conv = try tensor(e, ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ conv, geometry.taps, 1 }), + .norm = try tensor(e, ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{geometry.dk}), + .a_log = try f32vec(e, ck, "{s}linear_attn.A_log", .{p}, g.v_heads), + .dt_bias = try tensor(e, ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{g.v_heads}), + } }; + } else { + l.mixer = .{ .attention = .{ + .q = try linear(e, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), + .k = try linear(e, ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden), + .v = try linear(e, ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden), + .o = try linear(e, ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()), + .q_norm = try tensor(e, ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{head_dim}), + .k_norm = try tensor(e, ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{head_dim}), + } }; + } + } +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index b90c801e0d..cf64c4f8b9 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -349,6 +349,11 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; return null; }; + // a family that pre-allocates its streams' buffers does so now that their count is known + if (@hasField(L, "reserve")) loaded.reserve(loaded.ctx, streams) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); + return null; + }; const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 5f0dc8f906..78414fdb32 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,6 +1,7 @@ -//! `tf-qwen35-cuda `: Qwen3.5-2B's CUDA forward over comma-separated token ids, -//! one token a forward (step: every position's logits) or the ids as one prompt (prompt: the last position's), as -//! little-endian bf16 rows into for a reference to compare. Also prints the step mode's decode speed. +//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over comma-separated +//! token ids: one token a forward (step: every position's logits), windows of growing widths (window: the same rows, +//! which must equal step's) or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into +//! for a reference to compare. Step mode also prints the decode speed. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -26,19 +27,23 @@ pub fn main(init: std.process.Init) !u8 { defer cache.deinit(init.gpa); var out: std.ArrayList(u8) = .empty; if (std.mem.eql(u8, args[2], "step")) { - const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); - for (ids.items) |id| { + var first: i64 = 0; + var rest: i64 = 0; + for (ids.items, 0..) |id, i| { + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, &.{id}, .logits); + const dt: i64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); + if (i == 0) first = dt else rest += dt; @memcpy(cache.logits, e.rowLogits(0)); try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); } - const ns: u64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); - std.debug.print("{d} steps in {d:.1} ms: {d:.1} tok/s\n", .{ ids.items.len, @as(f64, @floatFromInt(ns)) / 1e6, @as(f64, @floatFromInt(ids.items.len)) * 1e9 / @as(f64, @floatFromInt(ns)) }); + const n: f64 = @floatFromInt(ids.items.len - 1); + std.debug.print("{d} steps: the first {d:.1} ms, then {d:.1} ms a step ({d:.1} tok/s)\n", .{ ids.items.len, @as(f64, @floatFromInt(first)) / 1e6, @as(f64, @floatFromInt(rest)) / 1e6 / n, n * 1e9 / @as(f64, @floatFromInt(rest)) }); } else if (std.mem.eql(u8, args[2], "window")) { // windows of 1, 2, 3, ... rows: every row's logits, which must equal step's bit for bit var at: usize = 0; var w: usize = 1; - while (at < ids.items.len) : (w = w % E.window_rows + 1) { + while (at < ids.items.len) : (w = w % e.window_rows + 1) { const n = @min(w, ids.items.len - at); try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .logits); for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); From 6b85f22799d796fc19982657c1e68834f63ebff5 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:36:52 +0000 Subject: [PATCH 08/46] qwen3_5: replay each decode shape as a CUDA graph A forward's launches (the embedding to the kept rows' logits) are captured the first time a shape (rows, windows, recording) is seen and replayed after; the attention chunk count moves into the per-forward table so a graph stays valid as caches grow. Prompt chunks stay uncaptured: their GEMMs dwarf the launches, and each length would be a graph of its own. A new graph's first run pays the Windows driver's stall (about a second on a wide model), so `reserve` also captures the shared rounds' shapes, a row of each of 1 to n streams, at load. Bit for bit the same logits; on an RTX 4090 the 2B steps in 2.8 ms instead of 5.3 (served: 324 tok/s, 1,560 across eight streams) and the 27B in 23.8 ms instead of 26.9 (served: 40 tok/s, 243 across eight). --- zig/kernels/cuda/qwen35.cu | 15 ++-- zig/src/families/qwen3_5/cuda_engine.zig | 110 +++++++++++++++++------ zig/src/families/qwen3_5/cuda_native.zig | 11 +++ 3 files changed, 103 insertions(+), 33 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 715f83f64e..15d77f706d 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -322,15 +322,18 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo } } -// Causal attention over a segment's cache, ck keys a block: grid (chunks, kvh, rows), block qh / kvh warps, a warp a -// query head of the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +// Causal attention over a segment's cache, P[1] keys a block over P[0] chunks (read from the device, so a captured +// forward replays as caches grow): grid (chunks at most, kvh, rows), block qh / kvh warps, a warp a query head of +// the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, + const int* __restrict__ rowseg, int layer, const int* __restrict__ P, float scale, float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; + const int nch = P[0], ck = P[1]; const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; + if (c >= nch) return; const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; @@ -377,10 +380,12 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; } -// The chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +// The P[0] chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, - const float* __restrict__ PO, int nch, __nv_bfloat16* __restrict__ OUT) { + const float* __restrict__ PO, const int* __restrict__ P, + __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; + const int nch = P[0]; const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; const size_t m = blockIdx.y; diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 316814eaf6..05e69afa11 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,6 +1,7 @@ //! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the //! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails //! and KV cache (Cache); the widths come from config.json (cuda_geometry.zig), the weights from cuda_weights.zig. +//! A forward's launches are captured into a CUDA graph the first time its shape is seen and replayed after. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -172,6 +173,8 @@ pub const Engine = struct { host_logits: cuda.HostBuffer, /// The kept rows' device draws, when forward was given them. tokens: [most_verify]u32 = undefined, + /// The captured forwards by shape (`Shape.key`): the launches from the embedding to the head. + graphs: std.AutoHashMapUnmanaged(u32, cuda.graph.Exec) = .empty, const ids_at = 0; const segs_at = max_rows * 4; @@ -179,7 +182,8 @@ pub const Engine = struct { const rules_at = rowseg_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const tables_len = tokens_at + most_verify * 4; + const params_at = tokens_at + most_verify * 4; // the attention chunks and their width + const tables_len = params_at + 2 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -195,6 +199,7 @@ pub const Engine = struct { e.verify_rows = std.math.clamp(snapshot_budget / (g.deltas() * (g.recBytes() + g.convBytes())), 1, most_verify); e.window_rows = @min(most_window, e.verify_rows); e.arena_at = 0; + e.graphs = .empty; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); @@ -240,6 +245,9 @@ pub const Engine = struct { } pub fn deinit(e: *Engine) void { + var it = e.graphs.valueIterator(); + while (it.next()) |x| x.deinit(); + e.graphs.deinit(e.gpa); e.host_logits.free(); e.staged.free(); e.tables.free(); @@ -363,15 +371,82 @@ pub const Engine = struct { m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 } } + // attention chunks: at most max_chunks a row, each a multiple of 128 keys + const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..8]); + params[0] = @intCast((widest + ck - 1) / ck); + params[1] = @intCast(ck); try e.tables.uploadAsync(0, st, e.stream.handle); + const shape: Shape = .{ .total = @intCast(total), .nseg = @intCast(parts.len), .record = record, .prompt = prompt }; + if (prompt) { + // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length would be a graph of its own, + // paying the driver's stall on a new graph's first run (about a second on a wide model) + try e.body(shape); + } else { + const gop = try e.graphs.getOrPut(e.gpa, shape.key()); + if (!gop.found_existing) { + errdefer _ = e.graphs.remove(shape.key()); + // captured, nothing runs: the replay below is this shape's first run + try cuda.graph.beginCapture(e.stream, .thread_local); + e.body(shape) catch |err| { + var broken = cuda.graph.endCapture(e.stream) catch return err; + broken.deinit(); + return err; + }; + var captured = try cuda.graph.endCapture(e.stream); + defer captured.deinit(); + gop.value_ptr.* = try captured.instantiate(); + } + try gop.value_ptr.launchOn(e.stream); + } + for (parts) |p| p.cache.len += p.rows; + switch (keep) { + .draws => |pk| { + const toks = e.tables.ptr + tokens_at; + for (pk, 0..) |p, i| { + const row = e.logits + i * g.vocab * 2; + if (p.rule != null) { + try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); + } + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); + }, + .logits, .last => { + const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); + }, + } + } + + /// A forward's shape, which fixes its launches: its tables and the caches' lengths are read on the device. + const Shape = struct { + total: u8, + nseg: u8, + record: bool, + prompt: bool, + + fn key(s: Shape) u32 { + return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17; + } + }; + + /// The launches of a forward of `shape`, from the embedding to the kept rows' logits (captured once a shape). + fn body(e: *Engine, shape: Shape) !void { + const g = e.g; + const total: usize = shape.total; + const nseg: usize = shape.nseg; + const record = shape.record; + const prompt = shape.prompt; + const heads: usize = if (record) total else 1; const t_ids = e.tables.ptr + ids_at; const t_segs = e.tables.ptr + segs_at; const t_rows = e.tables.ptr + rowseg_at; - const nseg = parts.len; + const t_params = e.tables.ptr + params_at; const geo = e.geo; - // attention chunks: at most max_chunks a row, each a multiple of 128 keys - const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); - const nch = (widest + ck - 1) / ck; const group = g.q_heads / g.kv_heads; try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)) }); @@ -403,8 +478,8 @@ pub const Engine = struct { try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps }); try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps }); try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); - try e.run(e.fns.attn_partial, .{ nch, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); - try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); + try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; @@ -420,27 +495,6 @@ pub const Engine = struct { try e.norm(true, e.final_norm, total); // the head over the kept rows try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); - for (parts) |p| p.cache.len += p.rows; - switch (keep) { - .draws => |pk| { - const toks = e.tables.ptr + tokens_at; - for (pk, 0..) |p, i| { - const row = e.logits + i * g.vocab * 2; - if (p.rule != null) { - try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); - } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); - } - try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); - try e.stream.synchronize(); - @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); - }, - .logits, .last => { - const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; - try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); - try e.stream.synchronize(); - if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); - }, - } } /// Row `i`'s logits from the newest forward that kept logits. diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 4b46a5a896..7fdea31ae7 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -63,6 +63,17 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try cache.reset(l.e); l.pool.appendAssumeCapacity(cache); } + // the shared rounds' shapes, a row of each of 1 to `streams` streams, captured now: a new graph's first run + // pays the driver's stall too (about a second on a wide model) + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_segs]u32 = @splat(0); + var picks: [E.max_segs]E.Draw = undefined; + for (0..@min(streams, E.max_segs)) |n| { + parts[n] = .{ .cache = l.pool.items[n], .rows = 1 }; + picks[n] = .{ .rule = null, .position = 1 }; + try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); + for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); + } } /// A request this engine refuses, in words; null: none of its own. From b19e0d73a27c25d3579b12373ca74a94cdacc592 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 02:00:27 +0000 Subject: [PATCH 09/46] qwen3_5: DFlash2 drafts for Qwen3.8-27B on CUDA --drafter loads z-lab's DFlash2 draft model behind the lane core's Drafted wrapper: the target writes each row's residual after the drafter's five tapped layers into a ring a stream, the drafter folds those rows into a context of keys and values, and after every verify it drafts a block of seven tokens past the pending one (the target's own embedding and head at its ends, a codebook selector tracing one path through each row's best sixteen). Its projections are quantized to affine-4 on the host at load and run through the target's kernels; its glue is qwen35_draft.cu. The Windows driver stalls a buffer's first kernel and a new graph's first run for seconds, so the drafter pools its contexts and a captured forward now serves every verify shape: its launches cover a bucket of rows and every window, and the kernels read the true counts from the device. On an RTX 4090 the 27B serves 52 to 58 tok/s greedy with drafts against 40 plain, about 100 tok/s on long replies, the same tokens either way; plain decoding keeps its 24 ms step. The 2B is unchanged. --- zig/build/cuda.zig | 3 +- zig/kernels/cuda/qwen35.cu | 63 ++- zig/kernels/cuda/qwen35_draft.cu | 156 ++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda.zig | 2 + zig/src/families/qwen3_5/cuda_draft.zig | 482 ++++++++++++++++++ .../families/qwen3_5/cuda_draft_weights.zig | 224 ++++++++ zig/src/families/qwen3_5/cuda_engine.zig | 94 +++- zig/src/families/qwen3_5/cuda_native.zig | 54 +- zig/src/native/cuda.zig | 27 +- zig/src/native/cuda_drafted.zig | 57 +++ 11 files changed, 1112 insertions(+), 51 deletions(-) create mode 100644 zig/kernels/cuda/qwen35_draft.cu create mode 100644 zig/src/families/qwen3_5/cuda_draft.zig create mode 100644 zig/src/families/qwen3_5/cuda_draft_weights.zig create mode 100644 zig/src/native/cuda_drafted.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 8632ea6c7f..78fdc3e649 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -32,7 +32,8 @@ const kernels = [_]Kernel{ .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) - .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: Qwen3.5-2B's glue, the Metal kernels' rounding points + .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points + .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 15d77f706d..de830da26f 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -25,7 +25,8 @@ struct Seg { unsigned long long rec; // float recurrences, one rec() block a linear layer unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time - int cap, len, row0, rows; + unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each + int cap, len, row0, rows, tap_rows, taps; }; __device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } @@ -41,10 +42,16 @@ __device__ __forceinline__ float warp_sum(float x) { using namespace tf_q35; +// The per-forward counts a captured forward reads from the device, so one graph serves every shape: P[0] attention +// chunks, P[1] their width, P[2] rows, P[3] windows. A null P means the launch's grid is exact. +__device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } +__device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } + // MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, - __nv_bfloat16* __restrict__ out, int D) { + __nv_bfloat16* __restrict__ out, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; const size_t id = ids[blockIdx.x]; for (int b = threadIdx.x; b < D / 2; b += blockDim.x) { const uint8_t v = W[id * (D / 2) + b]; @@ -58,9 +65,11 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __rest // h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), D wide (a multiple of 256, at most 8192). One 256-thread // block a row; each thread's elements in index order, warps summed in order. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, - const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D) { + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, + const int* __restrict__ P) { constexpr int MAXE = 32; __shared__ float red[8]; + if (past_rows(P, blockIdx.x)) return; const size_t row = blockIdx.x; const int t = threadIdx.x; const int E = D / 256; @@ -94,7 +103,8 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __r } // The 64-group sums of bf16 rows D wide (the decode projections' input sums). Grid (rows), a thread a group. -extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D) { +extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; const size_t row = blockIdx.x; for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { const __nv_bfloat16* x = X + row * D + g * 64; @@ -111,10 +121,12 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, - __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA) { + __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA, + const int* __restrict__ P) { const int lane = threadIdx.x; const int head = blockIdx.y; const int w = blockIdx.z; + if (past_rows(P, w)) return; const int NK = geo.nk, NV = geo.nv, C = geo.conv(); const Seg sg = segs[rowseg[w]]; const int r = w - sg.row0; @@ -162,10 +174,10 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ // Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). // SNAP, if given, also takes the tail after each row (forward row m at m * conv_rows()), read before the update. extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ SNAP) { + __nv_bfloat16* __restrict__ SNAP, const int* __restrict__ P) { const int C = geo.conv(); const int c = blockIdx.x * blockDim.x + threadIdx.x; - if (c >= C) return; + if (c >= C || past_windows(P, blockIdx.y)) return; const Seg sg = segs[blockIdx.y]; __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); if (SNAP != nullptr) @@ -188,8 +200,9 @@ extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __re extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP) { + __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP, const int* __restrict__ P) { const int NK = geo.nk, NV = geo.nv; + if (past_windows(P, blockIdx.z / NV)) return; const int hv = blockIdx.z % NV; const int hk = hv / (NV / NK); const Seg sg = segs[blockIdx.z / NV]; @@ -240,8 +253,10 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons // out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear value head; rows nv * DK wide. Grid (1, nv, rows), a warp each. extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, - const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT) { + const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT, + const int* __restrict__ P) { const int lane = threadIdx.x, hv = blockIdx.y; + if (past_rows(P, blockIdx.z)) return; const size_t m = blockIdx.z; const size_t W = (size_t)geo.nv * DK; constexpr int PER = DK / 32; @@ -265,9 +280,11 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const // Per-head RMSNorm of attention rows: X's head h of row m at m * stride + h * hstride, out packed (rows, heads, HD). // Grid (rows * heads), 64 threads of 4 values. extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, - __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps) { + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps, + const int* __restrict__ P) { __shared__ float sums[2]; const int head = blockIdx.x % heads; + if (past_rows(P, blockIdx.x / heads)) return; const size_t m = blockIdx.x / heads; const int t = threadIdx.x; const __nv_bfloat16* x = X + m * stride + head * hstride; @@ -303,9 +320,10 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i // Grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer) { + const int* __restrict__ rowseg, int layer, const int* __restrict__ P) { const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; + if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; const int pos = sg.len + (int)m - sg.row0; @@ -333,7 +351,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; - if (c >= nch) return; + if (c >= nch || past_rows(P, blockIdx.z)) return; const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; @@ -388,6 +406,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, cons const int nch = P[0]; const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; + if (past_rows(P, blockIdx.y)) return; const size_t m = blockIdx.y; const size_t base = (m * QH + qh) * nch; float mx = -INFINITY; @@ -409,8 +428,9 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, cons // out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows; rows qh * HD wide. extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, - __nv_bfloat16* __restrict__ OUT, int rows) { + __nv_bfloat16* __restrict__ OUT, int rows, const int* __restrict__ P) { const size_t W = (size_t)geo.qh * HD; + if (P != nullptr) rows = P[2]; const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; if (i >= (size_t)rows * W) return; const size_t r = i / W; @@ -420,9 +440,24 @@ extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __rest OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); } +// A tapped layer's residual stream, bf16(h + r) as the next norm rounds it, into each row's segment's tap ring at +// the row's position: plane `tap` of the row's taps. Grid (rows), 256 threads. +extern "C" __global__ void __launch_bounds__(256) tf_q35_tap(const __nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const Seg* __restrict__ segs, const int* __restrict__ rowseg, int tap, int D, + const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + if (sg.tap == 0) return; + const long long pos = sg.len + (long long)m - sg.row0; + __nv_bfloat16* out = reinterpret_cast<__nv_bfloat16*>(sg.tap) + ((size_t)(pos % sg.tap_rows) * sg.taps + tap) * D; + for (int i = threadIdx.x; i < D; i += blockDim.x) out[i] = __float2bfloat16_rn(bf(H[m * D + i]) + bf(R[m * D + i])); +} + // out = bf16(silu(gate) * up) over `total` values. extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, - __nv_bfloat16* __restrict__ OUT, long long total) { + __nv_bfloat16* __restrict__ OUT, long long total, const int* __restrict__ P, int D) { + if (P != nullptr) total = (long long)P[2] * D; const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; if (i >= total) return; const float gf = bf(GATE[i]); diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu new file mode 100644 index 0000000000..5d6286ea2b --- /dev/null +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -0,0 +1,156 @@ +// The DFlash2 drafter's glue on CUDA (zig/src/families/qwen3_5/cuda_draft.zig): a block of rows attends over the +// target's tapped context, with two-tap dynamic convolutions around attention and MLP. Rounding points follow the +// Metal drafter's draft_ops.metal; the projections run through the target engine's affine-4 kernels. + +#include +#include + +namespace tf_q35d { + +constexpr int HD = 128; // the drafter's head width (q, k, v) + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35d + +using namespace tf_q35d; + +// Tap rows (start + i) of a stream's ring of R rows, `width` bf16 each, into OUT rows 0..n. Grid (n), 256 threads. +extern "C" __global__ void tf_q35d_gather_taps(const __nv_bfloat16* __restrict__ RING, int R, long long start, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + blockIdx.x; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + +// y[t, c] = sum over lags l of (BASE[branch, l, c] + DYN[t, (branch * 2 + l) * groups + c / group]) * x[t - l, c], +// each product and sum rounded to bf16 in turn; x[-1] is zero (the block's first row has no past); RES, if given, +// is added last. D wide rows, `group` channels a dynamic coefficient. Grid (rows), 256 threads. +extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ DYN, + const __nv_bfloat16* __restrict__ BASE, const __nv_bfloat16* __restrict__ RES, + __nv_bfloat16* __restrict__ OUT, int D, int group, int branch) { + const int t = blockIdx.x; + const int groups = D / group; + for (int c = threadIdx.x; c < D; c += blockDim.x) { + const int g = c / group; + float total = 0.0f; +#pragma unroll + for (int l = 0; l < 2; l++) { + if (t - l < 0) continue; + const float x = bf(X[(size_t)(t - l) * D + c]); + total = rbf(total + rbf(bf(BASE[(branch * 2 + l) * D + c]) * x)); + total = rbf(total + rbf(bf(DYN[(size_t)t * (4 * groups) + (branch * 2 + l) * groups + g]) * x)); + } + if (RES != nullptr) total = rbf(total + bf(RES[(size_t)t * D + c])); + OUT[(size_t)t * D + c] = __float2bfloat16_rn(total); + } +} + +// Per-head RMSNorm over HD then rotate-half RoPE over all HD at position pos0 + row: X's head h of row m at +// m * stride + h * HD, OUT packed (rows, heads, HD). Grid (rows * heads), 32 threads of 4 values. +extern "C" __global__ void __launch_bounds__(32) tf_q35d_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, long long pos0, + float eps, float theta) { + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * HD; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + const float inv = rsqrtf(acc / (float)HD + eps); + float n[4]; +#pragma unroll + for (int i = 0; i < 4; i++) n[i] = rbf(rbf(v[i] * inv) * bf(W[4 * t + i])); + // the pair of lane element d is d + HD / 2: this lane's partner values come from the lane HD / 8 = 16 away + const float pos = (float)(pos0 + (long long)m); + __nv_bfloat16* out = OUT + (m * heads + head) * HD; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + const float partner = __shfl_xor_sync(0xffffffffu, n[i], 16); + const int j = d < HD / 2 ? d : d - HD / 2; + const float angle = pos * powf(theta, -(float)j / (float)(HD / 2)); + float sn, cs; + sincosf(angle, &sn, &cs); + out[d] = __float2bfloat16_rn(d < HD / 2 ? n[i] * cs - partner * sn : n[i] * cs + partner * sn); + } +} + +// Rows of packed (kv_heads, HD) keys and values at positions pos0 + i into a layer's rings (slot pos % W). +extern "C" __global__ void tf_q35d_ring_put(const __nv_bfloat16* __restrict__ K, const __nv_bfloat16* __restrict__ V, + __nv_bfloat16* __restrict__ RK, __nv_bfloat16* __restrict__ RV, int width, long long pos0, int W) { + const int i = blockIdx.x; + const size_t slot = (size_t)((pos0 + i) % W); + for (int c = threadIdx.x; c < width; c += blockDim.x) { + RK[slot * width + c] = K[(size_t)i * width + c]; + RV[slot * width + c] = V[(size_t)i * width + c]; + } +} + +// The block's attention: query row m, head h sees the context ring's positions [max(begin, p - window + 1), pos0) +// where p = pos0 + m, and every block row (the drafter is not causal within its block). GQA over kv_heads; scale +// 1 / sqrt(HD). Grid (rows, heads), a warp each; online softmax in fp32, bf16 out (rows, heads * HD). +extern "C" __global__ void __launch_bounds__(32) tf_q35d_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KB, + const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, + const __nv_bfloat16* __restrict__ RV, int rows, int heads, int kv_heads, int W, + long long begin, long long pos0, int window, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = rsqrtf((float)HD); + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + const long long p = pos0 + m; + long long first = p - window + 1; + if (first < begin) first = begin; + for (long long pos = first; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + for (int r = 0; r < rows; r++) { + const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(o[i] / l); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index c70e8b5ccf..ebc3722cf6 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -31,6 +31,7 @@ pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; +pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index ea5eaa465e..f21778c7bb 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -4,9 +4,11 @@ pub const config = @import("config.zig"); pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); +pub const draft = @import("cuda_draft.zig"); test { _ = config; _ = geometry; _ = @import("draw.zig"); + _ = @import("cuda_draft_weights.zig"); } diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig new file mode 100644 index 0000000000..90cce5c4cd --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -0,0 +1,482 @@ +//! The DFlash2 drafter on CUDA behind the lane core's Drafter contract (core/lanes/drafter.zig): it absorbs the +//! target's tapped rows into a context ring a stream, then drafts a block of 7 tokens after the pending one, the +//! target's own embedding and head at its ends and a codebook selector tracing one path through the top 16 a row. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const weights = @import("cuda_draft_weights.zig"); + +const E = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; +const dr = lanes.drafter; +const Stream = lanes.Stream; + +/// Rows an absorb runs at once (the target's prompt chunks are at most this). +const max_rows = engine.max_rows; +/// The logprob row's words (sample.cu's LP_WORDS) and its best-token capacity. +const lp_words = 3 + 2 * 20; +const lp_top = 20; + +pub const Config = struct { + hidden: usize, + intermediate: usize, + layers: usize, + vocab: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + block: usize, + mask: u32, + window: usize, + group: usize, + rank: usize, + topk: usize, + eps: f32, + theta: f32, + taps: [8]u32, + n_taps: usize, + + fn qWidth(c: Config) usize { + return c.heads * c.head_dim; + } + fn kvWidth(c: Config) usize { + return c.kv_heads * c.head_dim; + } + fn dynamic(c: Config) usize { + return 4 * (c.hidden / c.group); + } + /// The ring's rows: the window and a block more. + fn ringRows(c: Config) usize { + return c.window + 16; + } + fn ringBytes(c: Config) usize { + return c.layers * 2 * c.ringRows() * c.kvWidth() * 2; + } + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, target: engine.Geometry) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadDraftConfig; + const o = doc.value.object; + const d = (o.get("dflash_config") orelse return error.BadDraftConfig); + if (d != .object) return error.BadDraftConfig; + var c: Config = .{ + .hidden = try int(o, "hidden_size"), + .intermediate = try int(o, "intermediate_size"), + .layers = try int(o, "num_hidden_layers"), + .vocab = try int(o, "vocab_size"), + .heads = try int(o, "num_attention_heads"), + .kv_heads = try int(o, "num_key_value_heads"), + .head_dim = try int(o, "head_dim"), + .block = try int(d.object, "block_size"), + .mask = @intCast(try int(d.object, "mask_token_id")), + .window = try int(o, "sliding_window"), + .group = try int(d.object, "conv_group_size"), + .rank = try int(d.object, "selector_rank"), + .topk = try int(d.object, "selector_top_k"), + .eps = @floatCast(try num(o, "rms_norm_eps")), + .theta = @floatCast(try num((o.get("rope_parameters") orelse return error.BadDraftConfig).object, "rope_theta")), + .taps = undefined, + .n_taps = 0, + }; + const taps = d.object.get("target_layer_ids") orelse return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > c.taps.len) return error.BadDraftConfig; + for (taps.array.items, 0..) |t, i| { + if (t != .integer or t.integer < 0 or t.integer >= target.layers) return error.BadDraftConfig; + c.taps[i] = @intCast(t.integer); + } + c.n_taps = taps.array.items.len; + // what the kernels fix: the head width, two-tap convolutions, a block of 8, the selector's top-k within a row + if (c.head_dim != 128 or try int(d.object, "conv_kernel_size") != 2 or c.block != 8 or c.topk > lp_top or c.mask >= c.vocab) return error.UnsupportedDraftConfig; + if (c.hidden != target.hidden or c.vocab != target.vocab or c.hidden % 256 != 0 or c.hidden % c.group != 0 or c.heads % c.kv_heads != 0 or c.layers > 8) return error.UnsupportedDraftConfig; + if (c.window < c.block or c.window > engine.tap_rows - max_rows) return error.UnsupportedDraftConfig; + return c; + } +}; + +fn int(o: std.json.ObjectMap, key: []const u8) !usize { + const v = o.get(key) orelse return error.BadDraftConfig; + if (v != .integer or v.integer < 0) return error.BadDraftConfig; + return @intCast(v.integer); +} + +fn num(o: std.json.ObjectMap, key: []const u8) !f64 { + return switch (o.get(key) orelse return error.BadDraftConfig) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +const Fns = struct { gather_taps: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; + +/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds. +const Ctx = struct { + ring: DeviceBuffer, + begin: u64 = 0, + end: u64 = 0, + held: [8]u32 = undefined, + held_n: usize = 0, +}; + +pub const Draft = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E, + c: Config, + w: weights.Weights, + mod: cuda.Module, + fns: Fns, + scratch: DeviceBuffer, + at: usize = 0, + // absorb rows (max_rows) and block rows (block), all bf16 unless said + taps_in: u64, + th: u64, + thn: u64, + kabs: u64, + kn: u64, + vabs: u64, + xs: u64, // 64-group sums for the projections (f32) + h: u64, + h2: u64, + x: u64, + xin: u64, + dyn: u64, + qraw: u64, + q: u64, + kraw: u64, + kb: u64, + vb: u64, + att: u64, + o: u64, + gate: u64, + up: u64, + act: u64, + logits: u64, + words: u64, // u32 logprob words a row + picks: u64, // u32 zeros: the logprob rows' pick (unused) + hp: u64, + ids: u64, // u32 the block's ids + staged: cuda.HostBuffer, // pinned: the ids up, the words and projected rows down + streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, + /// Contexts between streams: allocated and run through once (`reserve`), since the driver stalls a buffer's + /// first kernel for seconds. + pool: std.ArrayList(*Ctx) = .empty, + step_ms: f64 = 4.0, + + const words_at = 64; + const hp_at = words_at + 16 * lp_words * 4; + const staged_len = hp_at + 16 * 256 * 2; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8) !*Draft { + const c = try Config.read(gpa, io, dir, e.g); + const d = try gpa.create(Draft); + errdefer gpa.destroy(d); + d.* = .{ .gpa = gpa, .io = io, .e = e, .c = c, .w = undefined, .mod = undefined, .fns = undefined, .scratch = undefined, .taps_in = 0, .th = 0, .thn = 0, .kabs = 0, .kn = 0, .vabs = 0, .xs = 0, .h = 0, .h2 = 0, .x = 0, .xin = 0, .dyn = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .o = 0, .gate = 0, .up = 0, .act = 0, .logits = 0, .words = 0, .picks = 0, .hp = 0, .ids = 0, .staged = undefined }; + d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen35_draft); + errdefer d.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(d.fns, name) = try d.mod.function("tf_q35d_" ++ name); + d.w = try weights.load(gpa, io, e, dir, .{ .hidden = c.hidden, .intermediate = c.intermediate, .layers = c.layers, .vocab = c.vocab, .q_width = c.qWidth(), .kv_width = c.kvWidth(), .head_dim = c.head_dim, .dynamic = c.dynamic(), .rank = c.rank, .taps = c.n_taps }); + errdefer d.w.deinit(gpa); + d.scratch = try DeviceBuffer.alloc(e.d, d.scratchBytes()); + errdefer d.scratch.free(); + try d.scratch.fill8(0, null); + d.at = 0; + _ = try d.layout(true); + d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); + errdefer d.staged.free(); + try d.picksZero(); + return d; + } + + pub fn deinit(d: *Draft) void { + var it = d.streams.valueIterator(); + while (it.next()) |ctx| d.drop(ctx.*); + d.streams.deinit(d.gpa); + for (d.pool.items) |ctx| d.drop(ctx); + d.pool.deinit(d.gpa); + d.staged.free(); + d.scratch.free(); + d.w.deinit(d.gpa); + d.mod.unload(); + d.gpa.destroy(d); + } + + fn scratchBytes(d: *Draft) usize { + d.at = 0; + return d.layout(false) catch unreachable; + } + + /// The scratch rows, measured (nothing assigned) or assigned from the scratch buffer. + fn layout(d: *Draft, assign: bool) !usize { + const c = d.c; + const B = c.block; + const R = max_rows; + const sizes = .{ + .{ "taps_in", R * c.n_taps * c.hidden * 2 }, .{ "th", R * c.hidden * 2 }, .{ "thn", R * c.hidden * 2 }, .{ "kabs", R * c.kvWidth() * 2 }, + .{ "kn", R * c.kvWidth() * 2 }, .{ "vabs", R * c.kvWidth() * 2 }, .{ "xs", R * (c.n_taps * c.hidden / 64) * 4 }, .{ "h", B * c.hidden * 2 }, + .{ "h2", B * c.hidden * 2 }, .{ "x", B * c.hidden * 2 }, .{ "xin", B * c.hidden * 2 }, .{ "dyn", B * c.dynamic() * 2 }, + .{ "qraw", B * c.qWidth() * 2 }, .{ "q", B * c.qWidth() * 2 }, .{ "kraw", B * c.kvWidth() * 2 }, .{ "kb", B * c.kvWidth() * 2 }, + .{ "vb", B * c.kvWidth() * 2 }, .{ "att", B * c.qWidth() * 2 }, .{ "o", B * c.hidden * 2 }, .{ "gate", B * c.intermediate * 2 }, + .{ "up", B * c.intermediate * 2 }, .{ "act", B * c.intermediate * 2 }, .{ "logits", B * c.vocab * 2 }, .{ "words", B * lp_words * 4 }, + .{ "picks", B * 4 }, .{ "hp", B * c.rank * 2 }, .{ "ids", B * 4 }, + }; + var total: usize = 0; + inline for (sizes) |s| { + const at = engine.aligned(total); + total = at + s[1]; + if (assign) @field(d, s[0]) = d.scratch.ptr + at; + } + return total; + } + + fn picksZero(d: *Draft) !void { + const zero: [64]u8 = @splat(0); + try d.e.d.check(d.e.d.api.cuMemcpyHtoD_v2(d.picks, &zero, 64), "cuMemcpyHtoD"); + } + + fn drop(d: *Draft, ctx: *Ctx) void { + ctx.ring.free(); + d.gpa.destroy(ctx); + } + + /// A new context, its ring touched by the kernels once: a row absorbed from the first tap ring at hand, then + /// a block drafted over it, so no request pays the driver's stall. + fn create(d: *Draft) !*Ctx { + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.c.ringBytes()) }; + errdefer ctx.ring.free(); + try ctx.ring.fill8(0, null); + try d.absorb(ctx, d.taps_in, 0, 1); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.c.block - 1); + ctx.* = .{ .ring = ctx.ring }; + return ctx; + } + + /// `n` contexts ready in the pool, for the streams the server admits. + pub fn reserve(d: *Draft, n: usize) !void { + try d.pool.ensureTotalCapacity(d.gpa, n); + while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + } + + fn ringK(d: *Draft, ctx: *Ctx, layer: usize) u64 { + return ctx.ring.ptr + layer * 2 * d.c.ringRows() * d.c.kvWidth() * 2; + } + + fn ringV(d: *Draft, ctx: *Ctx, layer: usize) u64 { + return d.ringK(ctx, layer) + d.c.ringRows() * d.c.kvWidth() * 2; + } + + /// out (rows, n) = x (rows, k) through `l`, as the target's decode tiles or prompt GEMM, with the drafter's sums. + fn proj(d: *Draft, l: cuda.qlinear.Affine4, x: u64, out: u64, rows: usize) !void { + const e = d.e; + if (rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, d.xs, @as(c_int, @intCast(l.k)), @as(u64, 0) }); + try e.affine.decode(e.stream, x, d.xs, l, out, rows); + } + + fn norm(d: *Draft, h: u64, w: u64, out: u64, rows: usize) !void { + const e = d.e; + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ h, @as(u64, 0), w, out, d.c.eps, @as(c_int, @intCast(d.c.hidden)), @as(u64, 0) }); + } + + fn dynConv(d: *Draft, x: u64, dyn: u64, base: u64, res: u64, out: u64, rows: usize, branch: usize) !void { + try d.e.run(d.fns.dyn_conv, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ x, dyn, base, res, out, @as(c_int, @intCast(d.c.hidden)), @as(c_int, @intCast(d.c.group)), @as(c_int, @intCast(branch)) }); + } + + fn normRope(d: *Draft, x: u64, w: u64, out: u64, rows: usize, heads: usize, pos0: u64) !void { + try d.e.run(d.fns.norm_rope, .{ rows * heads, 1, 1 }, .{ 32, 1, 1 }, .{ x, w, out, @as(c_int, @intCast(heads)), @as(c_int, @intCast(heads * d.c.head_dim)), @as(c_longlong, @intCast(pos0)), d.c.eps, d.c.theta }); + } + + /// Target rows at positions [start, start + n) from the stream's tap ring into its context (replacing what was + /// held at or past `start`); rows older than the window are skipped. + fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize) !void { + const c = d.c; + var start = start_; + var n = n_; + if (start < ctx.end) ctx.end = start; + if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here + if (n > c.window) { + start += n - c.window; + n = c.window; + ctx.begin = start; + } + ctx.end = start; + const e = d.e; + while (n > 0) { + const rows = @min(n, max_rows); + try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + try d.proj(d.w.fc, d.taps_in, d.th, rows); + try d.norm(d.th, d.w.hidden_norm, d.thn, rows); + for (d.w.layers, 0..) |l, i| { + try d.proj(l.k, d.thn, d.kabs, rows); + try d.normRope(d.kabs, l.k_norm, d.kn, rows, c.kv_heads, start); + try d.proj(l.v, d.thn, d.vabs, rows); + try e.run(d.fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kn, d.vabs, d.ringK(ctx, i), d.ringV(ctx, i), @as(c_int, @intCast(c.kvWidth())), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.ringRows())) }); + } + start += rows; + n -= rows; + ctx.end = start; + } + if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; + } + + /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held. + fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize) !void { + const c = d.c; + const e = d.e; + const B = c.block; + if (ctx.end != pos0) return error.DraftContextBehind; + const ids = std.mem.bytesAsSlice(u32, d.staged.bytes[0 .. B * 4]); + ids[0] = pending; + @memset(ids[1..B], c.mask); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.ids, ids.ptr, B * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + try e.run(e.fns.embed, .{ B, 1, 1 }, .{ 256, 1, 1 }, .{ d.ids, e.embed_w, e.embed_s, e.embed_b, d.h, @as(c_int, @intCast(c.hidden)), @as(u64, 0) }); + for (d.w.layers, 0..) |l, i| { + try d.norm(d.h, l.input_norm, d.x, B); + try d.proj(l.attn_conv, d.x, d.dyn, B); + try d.dynConv(d.x, d.dyn, l.attn_base, 0, d.xin, B, 0); + try d.proj(l.q, d.xin, d.qraw, B); + try d.proj(l.k, d.xin, d.kraw, B); + try d.proj(l.v, d.xin, d.vb, B); + try d.normRope(d.qraw, l.q_norm, d.q, B, c.heads, pos0); + try d.normRope(d.kraw, l.k_norm, d.kb, B, c.kv_heads, pos0); + try e.run(d.fns.attn, .{ B, c.heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.kb, d.vb, d.ringK(ctx, i), d.ringV(ctx, i), @as(c_int, @intCast(B)), @as(c_int, @intCast(c.heads)), @as(c_int, @intCast(c.kv_heads)), @as(c_int, @intCast(c.ringRows())), @as(c_longlong, @intCast(ctx.begin)), @as(c_longlong, @intCast(pos0)), @as(c_int, @intCast(c.window)), d.att }); + try d.proj(l.o, d.att, d.o, B); + try d.dynConv(d.o, d.dyn, l.attn_base, d.h, d.h2, B, 1); + try d.norm(d.h2, l.post_norm, d.x, B); + try d.proj(l.mlp_conv, d.x, d.dyn, B); + try d.dynConv(d.x, d.dyn, l.mlp_base, 0, d.xin, B, 0); + try d.proj(l.gate, d.xin, d.gate, B); + try d.proj(l.up, d.xin, d.up, B); + try e.run(e.fns.mlp_act, .{ (B * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ d.gate, d.up, d.act, @as(c_longlong, @intCast(B * c.intermediate)), @as(u64, 0), @as(c_int, @intCast(c.intermediate)) }); + try d.proj(l.down, d.act, d.o, B); + try d.dynConv(d.o, d.dyn, l.mlp_base, d.h2, d.h, B, 1); + } + // the mask rows' logits through the target's head, their best tokens and their selector projections + const rows = B - 1; + try d.norm(d.h + c.hidden * 2, d.w.final_norm, d.x, rows); + try d.proj(e.head, d.x, d.logits, rows); + try e.draws.logprobRows(e.stream, d.logits, c.vocab, d.picks, @intCast(c.topk), false, d.words, rows); + try d.proj(d.w.selector, d.x, d.hp, rows); + const st = d.staged.bytes; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, rows * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + hp_at, d.hp, rows * c.rank * 2, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + d.select(ctx, pending, rows); + ctx.held_n = @min(depth, rows); + } + + /// The path through the rows' candidates: each row's best of unary logit + predecessor . projected . successor. + fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { + const c = d.c; + const words = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][0 .. rows * lp_words * 4]); + const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][0 .. rows * c.rank * 2]); + var pred = pending; + for (0..rows) |r| { + const w = words[r * lp_words ..][0..lp_words]; + var best: f32 = -std.math.inf(f32); + var pick: u32 = w[2]; + for (0..c.topk) |j| { + const cand = w[2 + j]; + const unary: f32 = @bitCast(w[2 + lp_top + j]); + var edge: f32 = 0; + for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[r * c.rank + k]) * bf16(d.w.successor[cand * c.rank + k]); + const score = unary + edge; + if (score > best) { + best = score; + pick = cand; + } + } + ctx.held[r] = pick; + pred = pick; + } + } + + // the Drafter contract + + pub fn drafter(d: *Draft) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(p: *anyopaque) *Draft { + return @ptrCast(@alignCast(p)); + } + + fn tapsFn(p: *anyopaque) []const u32 { + const d = cast(p); + return d.c.taps[0..d.c.n_taps]; + } + + fn factsFn(p: *anyopaque) dr.Facts { + const d = cast(p); + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes() }; + } + + fn openFn(p: *anyopaque, s: *Stream) anyerror!void { + const d = cast(p); + if (d.streams.contains(s)) return error.StreamOpen; + const ctx = d.pool.pop() orelse try d.create(); + errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); + ctx.* = .{ .ring = ctx.ring }; + try d.streams.put(d.gpa, s, ctx); + } + + fn absorbFn(p: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("absorb", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) + try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows); + } + } + + fn holdFn(p: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("hold", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + try d.hold(ctx, it.pending, it.position - 1, it.depth); + } + } + + fn heldFn(p: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(p); + for (streams, out) |s, o| { + const ctx = d.streams.get(s) orelse return error.UnknownStream; + if (o.len > ctx.held_n) return error.TooFewDrafts; + @memcpy(o, ctx.held[0..o.len]); + } + } + + fn releaseFn(p: *anyopaque, s: *Stream) void { + const d = cast(p); + if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } +}; + +/// Acceptance by depth until a stream has its own (the reference's chat text: measured, see the commit). +const prior = [_]f64{ 0.75, 0.65, 0.58, 0.52, 0.47, 0.43, 0.4 }; + +fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// TF_TIMING=1: each call's wall time on stderr. +pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { + if (std.c.getenv("TF_TIMING") == null) return; + const now = std.Io.Clock.awake.now(io); + const ns: f64 = @floatFromInt(t0.durationTo(now).toNanoseconds()); + const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); + std.debug.print("[timing] {s} {d:.2} ms at {d:.1}\n", .{ what, ns / 1e6, at / 1e6 }); +} diff --git a/zig/src/families/qwen3_5/cuda_draft_weights.zig b/zig/src/families/qwen3_5/cuda_draft_weights.zig new file mode 100644 index 0000000000..2b0cf9f350 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_draft_weights.zig @@ -0,0 +1,224 @@ +//! The DFlash2 drafter's checkpoint onto the device: its bf16 projections quantized to affine-4 (group 64) on the +//! host and packed into qlinear's tiles through the target engine's staging buffer, its norms and base kernels as +//! they are, its codebooks kept on the host for the selector. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const engine = @import("cuda_engine.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const Engine = engine.Engine; +const aligned = engine.aligned; + +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + q: Affine4, + k: Affine4, + v: Affine4, + o: Affine4, + q_norm: u64, + k_norm: u64, + gate: Affine4, + up: Affine4, + down: Affine4, + attn_conv: Affine4, // kernel_projection: hidden -> 4 * groups dynamic coefficients + attn_base: u64, // base_kernel (2 branches, 2 lags, hidden) bf16 + mlp_conv: Affine4, + mlp_base: u64, +}; + +pub const Weights = struct { + fc: Affine4, // the taps' concatenation -> hidden + hidden_norm: u64, + final_norm: u64, + selector: Affine4, // hidden_projection: hidden -> rank + predecessor: []u16, // codebooks (vocab, rank) bf16, on the host + successor: []u16, + layers: []Layer, + arena: DeviceBuffer, + arena_at: usize = 0, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + w.arena.free(); + gpa.free(w.predecessor); + gpa.free(w.successor); + gpa.free(w.layers); + } + + fn alloc(w: *Weights, bytes: usize) !u64 { + const at = aligned(w.arena_at); + if (at + bytes > w.arena.len) return error.ArenaExhausted; + w.arena_at = at + bytes; + return w.arena.ptr + at; + } +}; + +/// The drafter's widths the loader needs (cuda_draft.zig's Config). +pub const Shape = struct { hidden: usize, intermediate: usize, layers: usize, vocab: usize, q_width: usize, kv_width: usize, head_dim: usize, dynamic: usize, rank: usize, taps: usize }; + +fn packedBytes(n: usize, k: usize) usize { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// Device bytes of everything `load` packs. +pub fn total(s: Shape) usize { + var t = packedBytes(s.hidden, s.taps * s.hidden) + packedBytes(s.rank, s.hidden) + 2 * aligned(s.hidden * 2); + t += s.layers * (2 * aligned(s.hidden * 2) + packedBytes(s.q_width, s.hidden) + 2 * packedBytes(s.kv_width, s.hidden) + packedBytes(s.hidden, s.q_width) + + 2 * aligned(s.head_dim * 2) + 2 * packedBytes(s.intermediate, s.hidden) + packedBytes(s.hidden, s.intermediate) + 2 * packedBytes(s.dynamic, s.hidden) + 2 * aligned(4 * s.hidden * 2)); + return t; +} + +fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +fn toBf16(x: f32) u16 { + // round to nearest even, as the device does + const u: u32 = @bitCast(x); + const rounding: u32 = 0x7fff + ((u >> 16) & 1); + return @intCast((u +% rounding) >> 16); +} + +/// MLX affine-4, group 64, over bf16 rows: a group's scale is bf16((max - min) / 15) and bias bf16(min); each weight's +/// nibble is round((w - bias) / scale) clamped to [0, 15], eight to a word, the first in the low bits. +fn quantize(gpa: std.mem.Allocator, w: []align(1) const u16, n: usize, k: usize) !struct { words: []u32, scales: []u16, biases: []u16 } { + const groups = k / 64; + const words = try gpa.alloc(u32, n * k / 8); + errdefer gpa.free(words); + const scales = try gpa.alloc(u16, n * groups); + errdefer gpa.free(scales); + const biases = try gpa.alloc(u16, n * groups); + errdefer gpa.free(biases); + for (0..n) |r| { + for (0..groups) |g| { + const row = w[r * k + g * 64 ..][0..64]; + var lo: f32 = std.math.inf(f32); + var hi: f32 = -std.math.inf(f32); + for (row) |h| { + const x = bf16(h); + lo = @min(lo, x); + hi = @max(hi, x); + } + const scale = toBf16((hi - lo) / 15.0); + const bias = toBf16(lo); + scales[r * groups + g] = scale; + biases[r * groups + g] = bias; + const s = bf16(scale); + const b = bf16(bias); + const inv = if (s == 0) 0.0 else 1.0 / s; + for (0..8) |j| { + var word: u32 = 0; + for (0..8) |i| { + const q = std.math.clamp(@round((bf16(row[j * 8 + i]) - b) * inv), 0.0, 15.0); + word |= @as(u32, @intFromFloat(q)) << @intCast(4 * i); + } + words[(r * k + g * 64) / 8 + j] = word; + } + } + } + return .{ .words = words, .scales = scales, .biases = biases }; +} + +/// A bf16 (n, k) tensor quantized and packed through the target's staging buffer into the drafter's arena. +fn linear(gpa: std.mem.Allocator, e: *Engine, w: *Weights, ck: *core.Checkpoint, name: []const u8, n: usize, k: usize) !Affine4 { + const t = try ck.expect(name, .bf16, &.{ n, k }); + const q = try quantize(gpa, std.mem.bytesAsSlice(u16, t.bytes), n, k); + defer { + gpa.free(q.words); + gpa.free(q.scales); + gpa.free(q.biases); + } + const parts = [_][]const u8{ std.mem.sliceAsBytes(q.words), std.mem.sliceAsBytes(q.scales), std.mem.sliceAsBytes(q.biases) }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (parts, &raw) |p, *r| { + if (at + p.len > e.stage.len) return error.StageTooSmall; + try e.stage.upload(at, p); + r.* = e.stage.ptr + at; + at += p.len; + } + const l = Affine4.layout(n, k); + const out: Affine4 = .{ .w = try w.alloc(l.words), .s = try w.alloc(l.scales), .b = try w.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], out); + try e.stream.synchronize(); + return out; +} + +fn vector(w: *Weights, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + const t = try ck.expect(name, .bf16, shape); + const p = try w.alloc(t.bytes.len); + try w.arena.upload(p - w.arena.ptr, t.bytes); + return p; +} + +fn hostBf16(gpa: std.mem.Allocator, ck: *core.Checkpoint, name: []const u8, shape: []const usize) ![]u16 { + const t = try ck.expect(name, .bf16, shape); + const out = try gpa.alloc(u16, t.bytes.len / 2); + @memcpy(out, std.mem.bytesAsSlice(u16, t.bytes)); + return out; +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *Engine, dir: []const u8, s: Shape) !Weights { + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var w: Weights = .{ .fc = undefined, .hidden_norm = 0, .final_norm = 0, .selector = undefined, .predecessor = &.{}, .successor = &.{}, .layers = &.{}, .arena = try DeviceBuffer.alloc(e.d, total(s)) }; + errdefer w.arena.free(); + try w.arena.fill8(0, null); + w.layers = try gpa.alloc(Layer, s.layers); + errdefer gpa.free(w.layers); + w.fc = try linear(gpa, e, &w, &ck, "fc.weight", s.hidden, s.taps * s.hidden); + w.hidden_norm = try vector(&w, &ck, "hidden_norm.weight", &.{s.hidden}); + w.final_norm = try vector(&w, &ck, "norm.weight", &.{s.hidden}); + w.selector = try linear(gpa, e, &w, &ck, "candidate_selector.hidden_projection.weight", s.rank, s.hidden); + w.predecessor = try hostBf16(gpa, &ck, "candidate_selector.predecessor_codebook", &.{ s.vocab, s.rank }); + errdefer gpa.free(w.predecessor); + w.successor = try hostBf16(gpa, &ck, "candidate_selector.successor_codebook", &.{ s.vocab, s.rank }); + errdefer gpa.free(w.successor); + for (w.layers, 0..) |*l, i| { + var buf: [96]u8 = undefined; + const Name = struct { + buf: *[96]u8, + i: usize, + fn of(self: @This(), suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(self.buf, "layers.{d}.{s}", .{ self.i, suffix }); + } + }; + const nm: Name = .{ .buf = &buf, .i = i }; + l.* = .{ + .input_norm = try vector(&w, &ck, try nm.of("input_layernorm.weight"), &.{s.hidden}), + .post_norm = try vector(&w, &ck, try nm.of("post_attention_layernorm.weight"), &.{s.hidden}), + .q = try linear(gpa, e, &w, &ck, try nm.of("self_attn.q_proj.weight"), s.q_width, s.hidden), + .k = try linear(gpa, e, &w, &ck, try nm.of("self_attn.k_proj.weight"), s.kv_width, s.hidden), + .v = try linear(gpa, e, &w, &ck, try nm.of("self_attn.v_proj.weight"), s.kv_width, s.hidden), + .o = try linear(gpa, e, &w, &ck, try nm.of("self_attn.o_proj.weight"), s.hidden, s.q_width), + .q_norm = try vector(&w, &ck, try nm.of("self_attn.q_norm.weight"), &.{s.head_dim}), + .k_norm = try vector(&w, &ck, try nm.of("self_attn.k_norm.weight"), &.{s.head_dim}), + .gate = try linear(gpa, e, &w, &ck, try nm.of("mlp.gate_proj.weight"), s.intermediate, s.hidden), + .up = try linear(gpa, e, &w, &ck, try nm.of("mlp.up_proj.weight"), s.intermediate, s.hidden), + .down = try linear(gpa, e, &w, &ck, try nm.of("mlp.down_proj.weight"), s.hidden, s.intermediate), + .attn_conv = try linear(gpa, e, &w, &ck, try nm.of("attention_conv.kernel_projection.weight"), s.dynamic, s.hidden), + .attn_base = try vector(&w, &ck, try nm.of("attention_conv.base_kernel"), &.{ 2, 2, s.hidden }), + .mlp_conv = try linear(gpa, e, &w, &ck, try nm.of("mlp_conv.kernel_projection.weight"), s.dynamic, s.hidden), + .mlp_base = try vector(&w, &ck, try nm.of("mlp_conv.base_kernel"), &.{ 2, 2, s.hidden }), + }; + } + return w; +} + +test "bf16 rounding of a quantized group's scale and bias" { + try std.testing.expectEqual(@as(u16, 0x3f80), toBf16(1.0)); + try std.testing.expectEqual(@as(f32, 1.0), bf16(0x3f80)); + const w: [64]u16 = @splat(0x3f80); // a flat group: scale 0, bias 1, every nibble 0 + const q = try quantize(std.testing.allocator, &w, 1, 64); + defer { + std.testing.allocator.free(q.words); + std.testing.allocator.free(q.scales); + std.testing.allocator.free(q.biases); + } + try std.testing.expectEqual(@as(u16, 0), q.scales[0]); + try std.testing.expectEqual(@as(u16, 0x3f80), q.biases[0]); + try std.testing.expectEqual(@as(u32, 0), q.words[0]); +} diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 05e69afa11..4ce0cfceb5 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -33,12 +33,17 @@ pub fn aligned(n: usize) usize { /// qwen35.cu's Geo and Seg. const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int }; -const Seg = extern struct { rec: u64, conv: u64, kv: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int }; +const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; + +/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. +pub const tap_rows = 2048 + max_rows; /// One stream's state on the device and the logits of its newest prompt chunk on the host. pub const Cache = struct { state: DeviceBuffer, kv: DeviceBuffer, + /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. + tap: ?DeviceBuffer = null, cap: usize, len: usize = 0, logits: []u16, @@ -52,7 +57,13 @@ pub const Cache = struct { var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); errdefer kv.free(); try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) - return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + var tap: ?DeviceBuffer = null; + if (e.n_taps > 0) { + tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); + try tap.?.fill8(0, null); + } + errdefer if (tap) |*t| t.free(); + return .{ .state = state, .kv = kv, .tap = tap, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; } /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). @@ -65,11 +76,13 @@ pub const Cache = struct { pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { self.state.free(); self.kv.free(); + if (self.tap) |*t| t.free(); gpa.free(self.logits); } - fn seg(self: *const Cache, g: Geometry, row0: usize, rows: usize) Seg { - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; + fn seg(self: *const Cache, e: *const Engine, row0: usize, rows: usize) Seg { + const g = e.g; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; } }; @@ -99,6 +112,7 @@ const Fns = struct { attn_gate: cuda.Function, mlp_act: cuda.Function, argmax: cuda.Function, + tap: cuda.Function, }; /// A window of one stream in a forward. @@ -175,6 +189,9 @@ pub const Engine = struct { tokens: [most_verify]u32 = undefined, /// The captured forwards by shape (`Shape.key`): the launches from the embedding to the head. graphs: std.AutoHashMapUnmanaged(u32, cuda.graph.Exec) = .empty, + /// A drafter's tapped layers (`setTaps`): each row's residual after them lands in its stream's tap ring. + taps: [8]u32 = undefined, + n_taps: usize = 0, const ids_at = 0; const segs_at = max_rows * 4; @@ -182,8 +199,8 @@ pub const Engine = struct { const rules_at = rowseg_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const params_at = tokens_at + most_verify * 4; // the attention chunks and their width - const tables_len = params_at + 2 * 4; + const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows + const tables_len = params_at + 4 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -200,6 +217,7 @@ pub const Engine = struct { e.window_rows = @min(most_window, e.verify_rows); e.arena_at = 0; e.graphs = .empty; + e.n_taps = 0; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); @@ -260,9 +278,22 @@ pub const Engine = struct { e.gpa.destroy(e); } - /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps. pub fn streamBytes(e: *const Engine, cap: usize) usize { - return e.g.stateBytes() + e.g.kvBytes(cap); + return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes(); + } + + /// A stream's tap ring: tap_rows rows of the tapped layers' residuals. + pub fn tapBytes(e: *const Engine) usize { + return tap_rows * e.n_taps * e.g.hidden * 2; + } + + /// The layers a drafter taps, told before the streams' caches exist; each forward then fills their rings. + pub fn setTaps(e: *Engine, taps: []const u32) !void { + if (taps.len > e.taps.len) return error.TooManyTaps; + for (taps) |t| if (t >= e.g.layers) return error.BadTap; + @memcpy(e.taps[0..taps.len], taps); + e.n_taps = taps.len; } /// The next `bytes` of the arena, 256-byte aligned. @@ -309,7 +340,7 @@ pub const Engine = struct { return total; } - fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { + pub fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { var a: cuda.Args = .{}; inline for (args) |v| a.add(v); const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = @intCast(block[0]), .y = @intCast(block[1]), .z = @intCast(block[2]) } }; @@ -318,9 +349,9 @@ pub const Engine = struct { /// out (rows, n) = x (rows, k) through `l`: decode tiles 16 rows at a time (each row's bits whatever the batch), /// or the prompt GEMM for a prompt chunk wider than one tile. - fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { + pub fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); - try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)) }); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)), e.tables.ptr + params_at }); var at: usize = 0; while (at < rows) : (at += 16) { const n = @min(16, rows - at); @@ -329,7 +360,7 @@ pub const Engine = struct { } fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { - try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)) }); + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. @@ -358,7 +389,7 @@ pub const Engine = struct { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); var row0: usize = 0; for (parts, 0..) |p, i| { - segs[i] = p.cache.seg(g, row0, p.rows); + segs[i] = p.cache.seg(e, row0, p.rows); for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; row0 += p.rows; @@ -373,11 +404,18 @@ pub const Engine = struct { } // attention chunks: at most max_chunks a row, each a multiple of 128 keys const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); - const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..8]); + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..16]); params[0] = @intCast((widest + ck - 1) / ck); params[1] = @intCast(ck); + params[2] = @intCast(total); + params[3] = @intCast(parts.len); try e.tables.uploadAsync(0, st, e.stream.handle); - const shape: Shape = .{ .total = @intCast(total), .nseg = @intCast(parts.len), .record = record, .prompt = prompt }; + // a captured forward's launches cover a bucket of rows and every window; its kernels read the true counts + // (a lone stream's one-row step, the common case without a drafter, keeps its own exact graph) + const lone = total == 1 and parts.len == 1; + const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; + const windows: usize = if (lone) 1 else max_segs; + const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; if (prompt) { // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length would be a graph of its own, // paying the driver's stall on a new graph's first run (about a second on a wide model) @@ -428,9 +466,10 @@ pub const Engine = struct { nseg: u8, record: bool, prompt: bool, + taps: bool, fn key(s: Shape) u32 { - return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17; + return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17 | @as(u32, @intFromBool(s.taps)) << 18; } }; @@ -449,7 +488,7 @@ pub const Engine = struct { const geo = e.geo; const group = g.q_heads / g.kv_heads; - try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)) }); + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)), t_params }); var di: usize = 0; var ai: c_int = 0; for (e.layers, 0..) |l, i| { @@ -463,10 +502,10 @@ pub const Engine = struct { try e.proj(m.z, e.x, e.z, total, prompt); try e.proj(m.a, e.x, e.a, total, prompt); try e.proj(m.b, e.x, e.b, total, prompt); - try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta }); - try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap }); - try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap }); - try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix }); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap, t_params }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap, t_params }); + try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix, t_params }); try e.proj(m.out, e.mix, e.r, total, prompt); di += 1; }, @@ -475,12 +514,12 @@ pub const Engine = struct { try e.proj(m.q, e.x, e.qraw, total, prompt); try e.proj(m.k, e.x, e.kraw, total, prompt); try e.proj(m.v, e.x, e.vraw, total, prompt); - try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps }); - try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); + try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); + try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai, t_params }); try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); - try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); + try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; }, @@ -489,8 +528,11 @@ pub const Engine = struct { try e.proj(l.gate, e.x, e.gate, total, prompt); try e.proj(l.up, e.x, e.up, total, prompt); const n: c_longlong = @intCast(total * g.intermediate); - try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); + try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); try e.proj(l.down, e.act, e.r, total, prompt); + if (shape.taps) for (e.taps[0..e.n_taps], 0..) |t, ti| if (t == i) { + try e.run(e.fns.tap, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.r, t_segs, t_rows, @as(c_int, @intCast(ti)), @as(c_int, @intCast(g.hidden)), t_params }); + }; } try e.norm(true, e.final_norm, total); // the head over the kept rows diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 7fdea31ae7..66b25c02c8 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -7,6 +7,7 @@ const lanes = @import("lanes"); const c = @import("config.zig"); const E = @import("cuda_engine.zig"); const draws = @import("draw.zig"); +const Draft = @import("cuda_draft.zig").Draft; const be = lanes.backend; pub const model_type = "qwen3_5"; @@ -14,6 +15,8 @@ pub const formats: []const []const u8 = &.{"mlx-q4g64"}; pub const default_context: i64 = 32768; pub const max_segments: u32 = 1; pub const prompt_rows: u32 = E.max_rows; +/// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. +pub const takes_drafter = true; pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; @@ -38,7 +41,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer e.deinit(); if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary const l = try gpa.create(Lanes); - l.* = .{ .gpa = gpa, .e = e, .capacity = o.context }; + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context }; return .{ .backend = l.backend(), .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, @@ -74,6 +77,30 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); } + // with a drafter: a lone stream's windows of the pending token and 1 to 7 drafts + if (l.e.n_taps > 0) for (2..@min(8, l.e.verify_rows) + 1) |rows| { + for (picks[0..rows]) |*pk| pk.* = .{ .rule = null, .position = 1 }; + try l.e.forward(&.{.{ .cache = l.pool.items[0], .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + try l.pool.items[0].reset(l.e); + }; +} + +/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). +pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { + if (bits != 4) return error.DrafterBitsUnsupported; + const l: *Lanes = @ptrCast(@alignCast(p)); + const d = try Draft.init(gpa, io, l.e, dir); + return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; +} + +fn reserveDrafter(p: *anyopaque, streams: u32) anyerror!void { + const d: *Draft = @ptrCast(@alignCast(p)); + try d.reserve(streams); +} + +fn closeDrafter(p: *anyopaque) void { + const d: *Draft = @ptrCast(@alignCast(p)); + d.deinit(); } /// A request this engine refuses, in words; null: none of its own. @@ -99,6 +126,7 @@ fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { const Lanes = struct { gpa: std.mem.Allocator, + io: std.Io, e: *E.Engine, capacity: usize, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, @@ -137,7 +165,7 @@ const Lanes = struct { } fn backend(self: *Lanes) be.Backend { - return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; } fn cast(ptr: *anyopaque) *Lanes { @@ -199,6 +227,8 @@ const Lanes = struct { /// Each window (a stream's pending token and copies) checked, then as many as fit in each forward. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("verify", self.io, t0); if (windows.len != out.len) return error.InvalidQwenWindows; for (windows, out) |w, o| { const rows = w.rows(); @@ -243,6 +273,8 @@ const Lanes = struct { /// The accepted prefix of each window: its state as it was after that row. fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("keep", self.io, t0); for (windows, paths) |w, path| { for (path, 0..) |r, i| if (r != i) return error.UnsupportedQwenTree; try self.e.rewind(try self.cacheOf(w.stream), path.len); @@ -253,6 +285,24 @@ const Lanes = struct { for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; } + /// A drafter's tapped layers, before any stream's caches exist: each forward fills their rings after. + fn prepareFeaturesFn(ptr: *anyopaque, taps: []const u32) !void { + const self = cast(ptr); + if (self.caches.count() != 0 or self.pool.items.len != 0) return error.TapsAfterStreams; + try self.e.setTaps(taps); + } + + /// Rows [start, start + count) of a stream's tap ring: `offset` carries the first row's position, the drafter + /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. + fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { + const self = cast(ptr); + if (taps.len != self.e.n_taps) return error.TapsMismatch; + const cache = try self.cacheOf(stream); + const ring = cache.tap orelse return error.NoTaps; + if (start + count > cache.len or count > E.tap_rows) return error.RowsNotHeld; + return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; + } + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { const self = cast(ptr); if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index cf64c4f8b9..27ca52012f 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -10,6 +10,7 @@ const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); +const drafted = @import("cuda_drafted.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{ nemotron.native, qwen35.native }; @@ -154,6 +155,7 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + served: drafted.Served = undefined, // the family's backend, wrapped with --drafter's draft model when given fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -167,6 +169,7 @@ const Host = struct { } if (h.slide) |s| h.slide_free.?(h.gpa, s); h.gpa.free(h.dir); + h.served.close(h.gpa); h.release(h.family); if (h.gpu) |g| { g.ctx.deinit(); @@ -335,6 +338,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; }; defer if (!opened) loaded.deinit(loaded.ctx); + var served = try drafted.wrap(F, a, gpa, io, loaded, o, problem) orelse return null; + defer if (!opened) served.close(gpa); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; // what a family's Loaded may leave out: no learner, no lone driver, no prompt-cache copies @@ -342,24 +347,30 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const learn_bytes: usize = if (@hasField(L, "learn_bytes")) loaded.learn_bytes else 0; const room = after.room(model) -| learn_bytes; const free = loaded.free_streams; - const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { + const streams = budget.admit(room, served.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { - error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(loaded.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), - error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(loaded.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / loaded.stream_bytes }), + error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(served.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), + error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(served.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / served.stream_bytes }), }; return null; }; // a family that pre-allocates its streams' buffers does so now that their count is known - if (@hasField(L, "reserve")) loaded.reserve(loaded.ctx, streams) catch |e| { + if (@hasField(L, "reserve")) { + loaded.reserve(loaded.ctx, streams) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); + return null; + }; + } + served.reserve(streams) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); return null; }; const h = try gpa.create(Host); errdefer gpa.destroy(h); - h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; - const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = served.backend, .served = served, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; + const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, served.stream_bytes * (streams -| free)); h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(served.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); @@ -394,7 +405,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, served.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = served.backend.vtable.first_row != null, .prompt_logprobs = served.backend.vtable.prompt_rows, .structures = served.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } diff --git a/zig/src/native/cuda_drafted.zig b/zig/src/native/cuda_drafted.zig new file mode 100644 index 0000000000..d259af0993 --- /dev/null +++ b/zig/src/native/cuda_drafted.zig @@ -0,0 +1,57 @@ +//! --drafter on the CUDA host: a family's draft model behind the lane core's Drafted wrapper over its backend. +const std = @import("std"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const Allocator = std.mem.Allocator; + +/// What the host serves: the family's backend as is, or wrapped with a drafter in front. +pub const Served = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + stream_bytes: usize, + held: ?*Held = null, + + /// The drafter's own per-stream buffers for `streams`, once the server knows the count. + pub fn reserve(s: *Served, streams: u32) !void { + if (s.held) |h| try h.reserve(h.ctx, streams); + } + + pub fn close(s: *Served, gpa: Allocator) void { + if (s.held) |h| { + h.drafted.deinit(); + h.close(h.ctx); + gpa.destroy(h); + } + s.held = null; + } +}; + +const Held = struct { drafted: lanes.drafted.Drafted, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void }; + +/// `loaded`'s backend, wrapped with the drafter `o.drafter` names when it does; null with `problem` set on refusal. +pub fn wrap(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, loaded: anytype, o: api.Open, problem: *[]const u8) !?Served { + var s: Served = .{ .backend = loaded.backend, .facts = loaded.facts, .stream_bytes = loaded.stream_bytes }; + const dir = o.drafter orelse return s; + if (!o.drafts) return s; + if (!@hasDecl(F, "openDrafter")) { + problem.* = try std.fmt.allocPrint(a, "--drafter: {s} checkpoints load no draft model on CUDA", .{F.model_type}); + return null; + } + const opened = F.openDrafter(gpa, io, loaded.ctx, dir, o.drafter_bits) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load the drafter {s} ({s})", .{ dir, @errorName(e) }); + return null; + }; + const h = try gpa.create(Held); + errdefer gpa.destroy(h); + h.* = .{ .drafted = undefined, .ctx = opened.ctx, .close = opened.close, .reserve = opened.reserve }; + h.drafted = lanes.drafted.Drafted.init(gpa, loaded.backend, opened.drafter) catch |e| { + opened.close(opened.ctx); + problem.* = try std.fmt.allocPrint(a, "the drafter {s} cannot attach to this model ({s})", .{ dir, @errorName(e) }); + return null; + }; + s.backend = h.drafted.backend(); + s.facts = h.drafted.facts(loaded.facts); + s.stream_bytes = h.drafted.streamBytes(loaded.stream_bytes); + s.held = h; + return s; +} From 288ba03d6a64292cf2e44bd26197c9051cebe70d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:05:12 +0000 Subject: [PATCH 10/46] qwen3_5: tree drafts on CUDA, the kept path replayed from a log A verify window's rows may form a tree (each row's parent given): the convolution taps, the delta rule and attention follow each row's own branch, and the DFlash2 drafter fills the lane core's planned lanes (each lane its parent's rank-r choice by the codebook selector), so a lone stream verifies up to 12 lanes a round and the planner learns which shapes land. A verify no longer snapshots every row's state (1.3 GB shared, which capped the 27B at 8 rows): each stream logs its rows' decay, write and key, and keep replays the kept path onto the recurrence in one pass (and moves a tree's kept keys down), so any window width fits. One-row steps and prompts update in place as before. The lane host times the target's windows at load, which the depth rule had lacked (it had been capping chains at 3 drafts). The runner's new tree mode verifies tree windows with junk branches against step logits (bit-exact on both models). --- zig/build/cuda.zig | 1 + zig/kernels/cuda/qwen35.cu | 164 +++++++++------- zig/kernels/cuda/qwen35.cuh | 81 ++++++++ zig/kernels/cuda/qwen35_draft.cu | 15 ++ zig/kernels/cuda/qwen35_keep.cu | 63 ++++++ zig/src/core/lanes/drafted.zig | 16 +- zig/src/core/lanes/drafter.zig | 11 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda_cache.zig | 104 ++++++++++ zig/src/families/qwen3_5/cuda_draft.zig | 149 ++++++++++---- zig/src/families/qwen3_5/cuda_engine.zig | 216 +++++++++------------ zig/src/families/qwen3_5/cuda_geometry.zig | 11 ++ zig/src/families/qwen3_5/cuda_native.zig | 61 ++++-- zig/src/native/cuda.zig | 4 + zig/tests/qwen35_cuda_main.zig | 47 ++++- 15 files changed, 684 insertions(+), 260 deletions(-) create mode 100644 zig/kernels/cuda/qwen35.cuh create mode 100644 zig/kernels/cuda/qwen35_keep.cu create mode 100644 zig/src/families/qwen3_5/cuda_cache.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 78fdc3e649..1f3e047eb4 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -34,6 +34,7 @@ const kernels = [_]Kernel{ .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) + .{ .name = "qwen35_keep", .flags = &.{"-O3"} }, // ours: a verified window's kept path applied to its stream's state .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index de830da26f..2186037011 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -2,50 +2,18 @@ // with the same bf16 rounding points, their widths from the model's Geo. Rows of one forward belong to segments (one // stream's window each); a segment's state lives in its stream's blocks. -#include -#include - -namespace tf_q35 { - -constexpr int DK = 128; // linear-attention key and value head width -constexpr int TAPS = 4; // causal convolution taps -constexpr int HD = 256; // attention head width -constexpr int MAXQ = 32; // attention query heads at most - -// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads. -struct Geo { - int hid, nk, nv, qh, kvh; - __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees - __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence - __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } -}; - -// One stream's window in this forward: its state blocks, cache length before the window, and rows. -struct Seg { - unsigned long long rec; // float recurrences, one rec() block a linear layer - unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer - unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time - unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each - int cap, len, row0, rows, tap_rows, taps; -}; - -__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } -__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } - -__device__ __forceinline__ float warp_sum(float x) { -#pragma unroll - for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); - return x; -} - -} // namespace tf_q35 +#include "qwen35.cuh" using namespace tf_q35; // The per-forward counts a captured forward reads from the device, so one graph serves every shape: P[0] attention -// chunks, P[1] their width, P[2] rows, P[3] windows. A null P means the launch's grid is exact. +// chunks, P[1] their width, P[2] rows, P[3] windows, P[4] whether the rows are recorded for keep (a verify window: +// its state stays as it was until the kept path is applied). A null P means the launch's grid is exact. +// A window's rows may form a tree: rowpar names each forward row's parent forward row (-1: the window's first row, +// after the stream's cache), rowanc each row's ancestors and itself as bits over its window's rows. __device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } __device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } +__device__ __forceinline__ bool recorded(const int* P) { return P != nullptr && P[4] != 0; } // MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, @@ -115,9 +83,11 @@ extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, floa } // The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. -// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. +// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. A row's taps +// are its parents (then the stream's tail), so a tree's rows each see their own branch. extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, const __nv_bfloat16* __restrict__ CW, + const int* __restrict__ rowseg, const int* __restrict__ rowpar, int layer, + const __nv_bfloat16* __restrict__ CW, const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, @@ -129,19 +99,22 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ if (past_rows(P, w)) return; const int NK = geo.nk, NV = geo.nv, C = geo.conv(); const Seg sg = segs[rowseg[w]]; - const int r = w - sg.row0; const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * geo.conv_rows(); const bool isq = head < NK, isk = !isq && head < 2 * NK; const int c0 = head * DK; // q, k and v channels are consecutive heads of DK + int src[TAPS]; // tap t reads the row TAPS - 1 - t tokens back: a forward row, or -1 - (tail row) +#pragma unroll + for (int tap = 0; tap < TAPS; tap++) src[tap] = back(rowpar, w, TAPS - 1 - tap); constexpr int PER = DK / 32; float vals[PER]; #pragma unroll for (int j = 0; j < PER; j++) { const int c = c0 + lane * PER + j; float acc = 0.0f; +#pragma unroll for (int tap = 0; tap < TAPS; tap++) { - const int at = r + tap; // in [state rows (TAPS - 1)] ++ [this segment's rows] - const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * C + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]); + const int at = src[tap]; + const float xv = at < 0 ? bf(CS[(size_t)(-1 - at) * C + c]) : bf(QKV[(size_t)at * C + c]); acc += bf(CW[c * TAPS + tap]) * xv; } const float conv = rbf(acc); @@ -171,22 +144,26 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ } } -// Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). -// SNAP, if given, also takes the tail after each row (forward row m at m * conv_rows()), read before the update. -extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ SNAP, const int* __restrict__ P) { +// Each segment's convolution tail: a recorded window logs the tail after each of its rows (the last TAPS - 1 tokens of +// the row's branch) for keep, and leaves the stream's own; a prompt chunk or one-row window updates the stream's in +// place. A thread a (segment, channel). +extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, const int* __restrict__ P) { const int C = geo.conv(); const int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= C || past_windows(P, blockIdx.y)) return; const Seg sg = segs[blockIdx.y]; __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); - if (SNAP != nullptr) - for (int j = 0; j < sg.rows; j++) + if (recorded(P)) { + for (int j = 0; j < sg.rows; j++) { + __nv_bfloat16* out = log_row(geo, sg, layer, j).conv; for (int r = 0; r < TAPS - 1; r++) { - const int at = r + j + 1; - SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * C + c] = - at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; + const int at = back(rowpar, sg.row0 + j, TAPS - 2 - r); + out[(size_t)r * C + c] = at < 0 ? CS[(size_t)(-1 - at) * C + c] : QKV[(size_t)at * C + c]; } + } + if (sg.rows != 1) return; + } // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row for (int r = 0; r < TAPS - 1; r++) { const int at = r + sg.rows; @@ -194,13 +171,17 @@ extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __re } } -// The gated delta rule over each segment's rows, its recurrence updated in place; SNAP, if given, takes the state after -// each row (forward row m at m * rec()). Value head hv reads key head hv / (nv / nk). Grid (1, DK / 4, nv * segments), -// block (32, 4): a warp a (value head, value row), its 32 lanes over the key width. +// The gated delta rule over each segment's rows, each row's state its parent's updated: a chain carries the state in +// registers, a branch rebuilds its parent's from the stream's recurrence and the logged updates of the ancestors. +// A recorded window logs each row's (decay, write, key) and leaves the recurrence as it was (keep applies the kept +// path, qwen35_keep.cu); a prompt chunk or one-row window updates it in place. Value head hv reads key head +// hv / (nv / nk). Grid (1, DK / 4, nv * segments), block (32, 4): a warp a (value head, value row), its 32 lanes +// over the key width. extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, - const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP, const int* __restrict__ P) { + const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, __nv_bfloat16* __restrict__ y, + const int* __restrict__ P) { const int NK = geo.nk, NV = geo.nv; if (past_windows(P, blockIdx.z / NV)) return; const int hv = blockIdx.z % NV; @@ -211,18 +192,40 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons constexpr int PT = DK / 32; const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + const bool record = recorded(P); float state[PT]; #pragma unroll for (int i = 0; i < PT; i++) state[i] = st[i]; + int prev = -1; // the row the registers' state follows for (int r = 0; r < sg.rows; r++) { - const size_t row = sg.row0 + r; + const int row = sg.row0 + r; + const int parent = rowpar[row]; + if (parent != prev) { + // a branch: its parent's state from the stream's and the ancestors' updates, oldest first (their keys + // and decays from this forward's rows, their writes from the log, which this thread wrote itself) + int anc[MAXW]; + int n = 0; + for (int a = parent; a >= 0; a = rowpar[a]) anc[n++] = a; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = n - 1; j >= 0; j--) { + const int a = anc[j]; + const float ga = g[a * NV + hv], da = log_row(geo, sg, layer, a - sg.row0).delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * ga; + state[i] = fmaf(bf(k[((size_t)a * NK + hk) * DK + PT * dk + i]), da, state[i]); + } + } + } + prev = row; float kc[PT], qc[PT]; #pragma unroll for (int i = 0; i < PT; i++) { - kc[i] = bf(k[(row * NK + hk) * DK + PT * dk + i]); - qc[i] = bf(q[(row * NK + hk) * DK + PT * dk + i]); + kc[i] = bf(k[((size_t)row * NK + hk) * DK + PT * dk + i]); + qc[i] = bf(q[((size_t)row * NK + hk) * DK + PT * dk + i]); } - const float vc = bf(v[(row * NV + hv) * DK + dv]); + const float vc = bf(v[((size_t)row * NV + hv) * DK + dv]); const float gc = g[row * NV + hv]; const float bc = bf(beta[row * NV + hv]); float kv_mem = 0.0f; @@ -236,17 +239,23 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons float out = 0.0f; #pragma unroll for (int i = 0; i < PT; i++) { - state[i] = state[i] + kc[i] * delta; + state[i] = fmaf(kc[i], delta, state[i]); out += state[i] * qc[i]; } out = warp_sum(out); - if (dk == 0) y[(row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); - if (SNAP != nullptr) { - float* sn = SNAP + row * geo.rec() + at; + if (dk == 0) y[((size_t)row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); + if (record) { + const LogRow l = log_row(geo, sg, layer, r); + if (dk == 0) { + l.delta[hv * DK + dv] = delta; + if (dv == 0) l.g[hv] = gc; + } + if (dv == 0 && hv % (NV / NK) == 0) #pragma unroll - for (int i = 0; i < PT; i++) sn[i] = state[i]; + for (int i = 0; i < PT; i++) l.k[hk * DK + PT * dk + i] = k[((size_t)row * NK + hk) * DK + PT * dk + i]; } } + if (record && sg.rows != 1) return; #pragma unroll for (int i = 0; i < PT; i++) st[i] = state[i]; } @@ -316,17 +325,19 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); } -// Queries in place at their positions; keys rotated and values copied into their segment's cache. -// Grid (rows), HD threads. +// Queries in place at their positions (the cache length plus the row's depth in its window); keys rotated and values +// copied into their segment's cache at the row's own slot (the cache length plus its row). Grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, const int* __restrict__ P) { + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P) { const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; - const int pos = sg.len + (int)m - sg.row0; + const int pos = sg.len + __popc(rowanc[m]) - 1; + const int slot = sg.len + (int)m - sg.row0; __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; __syncthreads(); @@ -334,7 +345,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; for (int h = 0; h < KVH; h++) { - const size_t dst = ((size_t)h * sg.cap + pos) * HD + d; + const size_t dst = ((size_t)h * sg.cap + slot) * HD + d; keys[dst] = rotate(Kn + (m * KVH + h) * HD, d, pos); values[dst] = Vr[(m * KVH + h) * HD + d]; } @@ -342,9 +353,11 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo // Causal attention over a segment's cache, P[1] keys a block over P[0] chunks (read from the device, so a captured // forward replays as caches grow): grid (chunks at most, kvh, rows), block qh / kvh warps, a warp a query head of -// the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +// the kv head's group. A row sees the cache before its window and its window's ancestors and itself (rowanc). +// Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, const int* __restrict__ P, float scale, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P, float scale, float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; const int nch = P[0], ck = P[1]; @@ -355,8 +368,8 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; - const int last = sg.len + (int)m - sg.row0; // the row's own position, the last key it sees - const int k0 = c * ck, k1 = min(k0 + ck, last + 1); + const unsigned anc = rowanc[m]; + const int k0 = c * ck, k1 = min(k0 + ck, sg.len + sg.rows); const size_t slot = (m * QH + qh) * nch + c; if (k0 >= k1) { if (lane == 0) { @@ -375,6 +388,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r } float mx = -INFINITY, l = 0.0f; for (int pos = k0; pos < k1; pos++) { + if (pos >= sg.len && ((anc >> (pos - sg.len)) & 1u) == 0) continue; const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); const __nv_bfloat16* kb = reinterpret_cast(&kw); float dot = 0.0f; diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh new file mode 100644 index 0000000000..f3baa2be6c --- /dev/null +++ b/zig/kernels/cuda/qwen35.cuh @@ -0,0 +1,81 @@ +// What qwen35.cu and qwen35_keep.cu share: the model's widths (Geo), a stream's window in a forward (Seg), the per-row +// record a verify window leaves in its stream's log, and the walk from a row back through its parents. + +#pragma once +#include +#include + +namespace tf_q35 { + +constexpr int DK = 128; // linear-attention key and value head width +constexpr int TAPS = 4; // causal convolution taps +constexpr int HD = 256; // attention head width +constexpr int MAXQ = 32; // attention query heads at most +constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) + +// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. +struct Geo { + int hid, nk, nv, qh, kvh, nl; + __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees + __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence + __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } + // one row's record of one linear layer: decay and write strength a value head, its write a value (head, dv), its + // key a key head, and the convolution tail after it + __device__ size_t log_bytes() const { return (size_t)nv * 4 + (size_t)nv * DK * 4 + (size_t)nk * DK * 2 + conv_rows() * 2; } +}; + +// One stream's window in this forward: its state blocks, cache length before the window, and rows. +struct Seg { + unsigned long long rec; // float recurrences, one rec() block a linear layer + unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer + unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time + unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each + unsigned long long log; // the window's rows' records (MAXW rows of nl layers of log_bytes), for keep + int cap, len, row0, rows, tap_rows, taps; +}; + +// A window row's record of one linear layer (Seg.log). +struct LogRow { + float* g; // [nv] decay + float* delta; // [nv][DK] the write's strength along dv + __nv_bfloat16* k; // [nk][DK] the key + __nv_bfloat16* conv; // [TAPS - 1][conv()] the tail after the row +}; + +__device__ __forceinline__ LogRow log_row(const Geo& geo, const Seg& sg, int layer, int row) { + char* base = reinterpret_cast(sg.log) + ((size_t)row * geo.nl + layer) * geo.log_bytes(); + LogRow l; + l.g = reinterpret_cast(base); + l.delta = l.g + geo.nv; + l.k = reinterpret_cast<__nv_bfloat16*>(l.delta + (size_t)geo.nv * DK); + l.conv = l.k + (size_t)geo.nk * DK; + return l; +} + +// The rows of a window to keep, in order (row 0 first), passed by value. +struct Path { + int n; + int rows[MAXW]; +}; + +// The row `steps` tokens before forward row `w` through its parents (rowpar: each forward row's parent forward row, -1 +// for a window's first row): a forward row, or -1 - t for the stream's convolution tail row t (TAPS - 2 the newest). +__device__ __forceinline__ int back(const int* __restrict__ rowpar, int w, int steps) { + int a = w; + for (int s = 1; s <= steps; s++) { + a = rowpar[a]; + if (a < 0) return -1 - (TAPS - 2 - (steps - s)); + } + return a; +} + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35 diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu index 5d6286ea2b..a4726d0711 100644 --- a/zig/kernels/cuda/qwen35_draft.cu +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -31,6 +31,21 @@ extern "C" __global__ void tf_q35d_gather_taps(const __nv_bfloat16* __restrict__ for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; } +// Rows of a window to gather, by value (a tree's kept path). +struct TapRows { + int n; + int rows[16]; +}; + +// Tap rows (start + rows.rows[i]) of a stream's ring into OUT rows 0..n: a kept path's rows in order. Grid (n). +extern "C" __global__ void tf_q35d_gather_rows(const __nv_bfloat16* __restrict__ RING, int R, long long start, TapRows rows, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + rows.rows[blockIdx.x]; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + // y[t, c] = sum over lags l of (BASE[branch, l, c] + DYN[t, (branch * 2 + l) * groups + c / group]) * x[t - l, c], // each product and sum rounded to bf16 in turn; x[-1] is zero (the block's first row has no past); RES, if given, // is added last. D wide rows, `group` channels a dynamic coefficient. Grid (rows), 256 threads. diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu new file mode 100644 index 0000000000..b2dae0b2b5 --- /dev/null +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -0,0 +1,63 @@ +// Keeping a verified window's path on CUDA (zig/src/families/qwen3_5/cuda_cache.zig): the stream's recurrences take +// the path's logged updates in order, its convolution tails the path's last row's, and a tree's kept keys move down +// to the slots after the cache. The rows of a window were recorded by qwen35.cu's kernels (Seg.log). + +#include "qwen35.cuh" + +using namespace tf_q35; + +// Each linear layer's recurrence through the path's rows' (decay, key, write), as the forward's chain would have +// left it. Grid (1, DK / 4, nv * nl), block (32, 4): a warp a (layer, value head, value row), its lanes over the key. +extern "C" __global__ void __launch_bounds__(128) tf_q35k_apply(Geo geo, Seg sg, Path path) { + const int NK = geo.nk, NV = geo.nv; + const int layer = blockIdx.z / NV, hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = 0; j < path.n; j++) { + const LogRow l = log_row(geo, sg, layer, path.rows[j]); + const float g = l.g[hv], d = l.delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * g; + state[i] = fmaf(bf(l.k[hk * DK + PT * dk + i]), d, state[i]); + } + } +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// Each linear layer's convolution tail as logged after window row `last`. Grid (nl, tail values / 256), 256 threads. +extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { + const int layer = blockIdx.x; + const size_t n = geo.conv_rows(); + const size_t i = (size_t)blockIdx.y * blockDim.x + threadIdx.x; + if (i >= n) return; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * n; + CS[i] = log_row(geo, sg, layer, last).conv[i]; +} + +// The path's rows' keys and values (slot cache length + row) down to the slots cache length + 0, 1, ... in path +// order: row j of the path is at least j, so a move never lands on a slot a later move still reads. One block a +// (attention layer, kv head): grid (attention layers, kvh), HD threads, a barrier between moves. +extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { + const int KVH = geo.kvh; + const int layer = blockIdx.x, h = blockIdx.y; + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + ((size_t)layer * 2 * KVH + h) * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + const int d = threadIdx.x; + for (int j = 0; j < path.n; j++) { + const int from = sg.len + path.rows[j], to = sg.len + j; + if (from != to) { + keys[(size_t)to * HD + d] = keys[(size_t)from * HD + d]; + values[(size_t)to * HD + d] = values[(size_t)from * HD + d]; + } + __syncthreads(); + } +} diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 3bdffe7c7a..e1a79fd2de 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -60,7 +60,9 @@ pub const Drafted = struct { m.plain_guard = d.plain_guard; m.draft_streams = d.batched; m.draft_probabilities = false; - m.head_trees = false; + m.head_trees = d.trees; + if (d.window_costs.len > 0) m.window_costs = d.window_costs; + if (d.lane_costs.len > 0) m.lane_costs = d.lane_costs; return m; } @@ -157,15 +159,17 @@ pub const Drafted = struct { x.holds.clearRetainingCapacity(); try x.firsts.resize(x.gpa, requests.len); for (requests, x.firsts.items) |r, *tok| { - if (r.lanes != null) return error.TreesNotBuilt; + if (r.lanes != null and !x.drafter.facts().trees) return error.TreesNotBuilt; const lane = x.lanes.getPtr(r.stream) orelse return error.NotDrafting; lane.filler = 0; // a new request supersedes fillers a copy round left unused var pending: u32 = undefined; if (r.rows) |rows| { - for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; - const n: u32 = @intCast(rows.len); + // a tree's kept path names the window's rows out of order: the features of all up to its last + var prefix = true; + for (rows, 0..) |row, i| prefix = prefix and row == i; + const n: u32 = @intCast(if (prefix) rows.len else rows[rows.len - 1] + 1); const f = try x.target.features(r.stream, x.drafter.taps(), r.start, n); - try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len] }); + try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len], .rows = if (prefix) null else rows }); pending = r.follow[rows.len - 1]; lane.rows = true; } else { @@ -187,7 +191,7 @@ pub const Drafted = struct { lane.filler = r.depth; continue; } - try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth }); + try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth, .lanes = r.lanes }); } if (x.absorbs.items.len > 0) try x.drafter.absorb(x.absorbs.items); if (x.holds.items.len > 0) try x.drafter.hold(x.holds.items); diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 0f8eab242a..038ee35782 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -1,6 +1,8 @@ //! A drafter outside the target family reading its tapped states; each call takes a round's streams together. const Stream = @import("stream.zig").Stream; const Features = @import("backend.zig").Features; +const Cost = @import("config.zig").Cost; +const Shape = @import("shape.zig").Shape; /// What the round loop's depth rule needs from the drafter itself (not from the target, whose own head may not exist). pub const Facts = struct { @@ -10,13 +12,18 @@ pub const Facts = struct { plain_guard: bool = true, // plain rounds compete with drafted depths (the rule may choose no drafts) batched: bool = false, // hold() runs a shared round's streams as one batch stream_bytes: usize = 0, // device bytes each opened stream holds (its caches), for memory admission + trees: bool = false, // hold() takes a tree of lanes (Hold.lanes), each lane its parent's rank-r token + window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule + lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) }; /// Target rows [start, start + features.rows) and the token after each; rows at or past `start` replace what was held. -pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32 }; +/// `rows`, when given, picks the features' rows to absorb (a tree's kept path), in order; null: all of them. +pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32, rows: ?[]const u32 = null }; /// Draft `depth` tokens after `pending`, the first at `position`; asked only once row `position - 2` was absorbed. -pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32 }; +/// With `lanes`, a tree instead: lane i's token is the rank-r choice after its parent (depth unused). +pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32, lanes: ?*const Shape = null }; pub const Drafter = struct { ptr: *anyopaque, diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index ebc3722cf6..756e535a76 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -32,6 +32,7 @@ pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; +pub const qwen35_keep: []const u8 = if (available) &Blob("fatbin_qwen35_keep").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig new file mode 100644 index 0000000000..1d2a31d531 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -0,0 +1,104 @@ +//! One stream's state on the CUDA engine: its recurrences and convolution tails, keys and values, a drafter's tap +//! ring, and the log of its newest verify window's rows, from which `keep` lands any path of that window (a prefix +//! of a chain, or a branch of a tree) without the forward having touched the stream's state. +const std = @import("std"); +const cuda = @import("cuda"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Engine = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. +pub const tap_rows = 2048 + engine.max_rows; + +/// qwen35.cuh's Seg and Path. +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; +pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +pub const Cache = struct { + state: DeviceBuffer, + kv: DeviceBuffer, + /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. + tap: ?DeviceBuffer = null, + /// The newest verify window's rows' records (qwen35.cuh's LogRow), most_window rows a linear layer. + log: DeviceBuffer, + cap: usize, + len: usize = 0, + /// Positions the newest forward wrote up to: a tap ring holds them until the next forward, whatever `keep` lands. + written: usize = 0, + logits: []u16, + /// The newest verify window of more than one row, until its path is kept: where it started, its rows. + last: ?struct { start: usize, rows: usize } = null, + + pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); + errdefer state.free(); + try state.fill8(0, null); + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); + errdefer kv.free(); + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + var log = try DeviceBuffer.alloc(e.d, e.g.logBytes(engine.most_window)); + errdefer log.free(); + try log.fill8(0, null); + var tap: ?DeviceBuffer = null; + if (e.n_taps > 0) { + tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); + try tap.?.fill8(0, null); + } + errdefer if (tap) |*t| t.free(); + return .{ .state = state, .kv = kv, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). + pub fn reset(self: *Cache, e: *const Engine) !void { + self.len = 0; + self.written = 0; + self.last = null; + try self.state.fill8(0, e.stream.handle); + } + + pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { + self.state.free(); + self.kv.free(); + self.log.free(); + if (self.tap) |*t| t.free(); + gpa.free(self.logits); + } + + /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. + pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { + const g = e.g; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; + } + + /// Keeps `path` (the window's rows, the first 0, ascending) of the newest window: the recurrences and tails as a + /// chain of those rows would have left them, a tree's kept keys moved down to the slots after the cache, the + /// length at the path's end. Queued on the engine's stream; nothing waits. + pub fn keep(self: *Cache, e: *Engine, path: []const u32) !void { + const last = self.last orelse return error.NothingToKeep; + if (path.len == 0 or path.len > last.rows or path[0] != 0) return error.NothingToKeep; + var p: Path = .{ .n = @intCast(path.len), .rows = @splat(0) }; + var prefix = true; + for (path, 0..) |r, i| { + if (r >= last.rows or (i > 0 and r <= path[i - 1])) return error.NothingToKeep; + p.rows[i] = @intCast(r); + prefix = prefix and r == i; + } + self.last = null; + self.len = last.start + path.len; + const g = e.g; + const sg = self.seg(e, last.start, 0, last.rows); + try e.run(e.keeps.apply, .{ 1, geometry.dk / 4, g.v_heads * g.deltas() }, .{ 32, 4, 1 }, .{ e.geo, sg, p }); + try e.run(e.keeps.conv, .{ g.deltas(), (g.convBytes() / 2 + 255) / 256, 1 }, .{ 256, 1, 1 }, .{ e.geo, sg, @as(c_int, @intCast(path[path.len - 1])) }); + if (!prefix and g.attentions() > 0) try e.run(e.keeps.compact, .{ g.attentions(), g.kv_heads, 1 }, .{ geometry.head_dim, 1, 1 }, .{ e.geo, sg, p }); + } + + /// The newest window kept whole (the round loop keeps nothing when every row landed), before the next forward. + pub fn settle(self: *Cache, e: *Engine) !void { + const last = self.last orelse return; + var all: [engine.most_window]u32 = undefined; + for (all[0..last.rows], 0..) |*r, i| r.* = @intCast(i); + try self.keep(e, all[0..last.rows]); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 90cce5c4cd..9b5a9c1c2c 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -1,6 +1,7 @@ //! The DFlash2 drafter on CUDA behind the lane core's Drafter contract (core/lanes/drafter.zig): it absorbs the //! target's tapped rows into a context ring a stream, then drafts a block of 7 tokens after the pending one, the -//! target's own embedding and head at its ends and a codebook selector tracing one path through the top 16 a row. +//! target's own embedding and head at its ends and a codebook selector scoring the top 16 a row under each parent: +//! one path (a chain), or the lanes of the round loop's tree (each lane its parent's rank-r choice). const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -113,14 +114,18 @@ fn num(o: std.json.ObjectMap, key: []const u8) !f64 { }; } -const Fns = struct { gather_taps: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; +const Fns = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; -/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds. +/// qwen35_draft.cu's TapRows. +const TapRows = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds (a +/// chain, or a tree's lanes in the shape's order). const Ctx = struct { ring: DeviceBuffer, begin: u64 = 0, end: u64 = 0, - held: [8]u32 = undefined, + held: [engine.most_window]u32 = undefined, held_n: usize = 0, }; @@ -167,7 +172,11 @@ pub const Draft = struct { /// Contexts between streams: allocated and run through once (`reserve`), since the driver stalls a buffer's /// first kernel for seconds. pool: std.ArrayList(*Ctx) = .empty, - step_ms: f64 = 4.0, + /// One level of a hold (a hold drafts every level at once; timed at load, over the block's levels). + step_ms: f64 = 1.0, + /// The target's window ms by rows (cuda_native.zig times them), for the round loop's prices. + costs: [engine.most_window]lanes.config.Cost = undefined, + n_costs: usize = 0, const words_at = 64; const hp_at = words_at + 16 * lp_words * 4; @@ -191,9 +200,37 @@ pub const Draft = struct { d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); errdefer d.staged.free(); try d.picksZero(); + // the first context, and a hold timed on it (the quickest of several: the first ones pay the driver's stalls) + const ctx = try d.create(); + errdefer d.drop(ctx); + try d.absorb(ctx, d.taps_in, 0, 1, null); + var quickest: f64 = std.math.inf(f64); + for (0..8) |_| { + ctx.end = 1; + const t0 = std.Io.Clock.awake.now(io); + try d.hold(ctx, 0, 1, c.block - 1, null); + quickest = @min(quickest, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6); + } + d.step_ms = quickest / @as(f64, @floatFromInt(c.block - 1)); + ctx.* = .{ .ring = ctx.ring }; + try d.pool.append(gpa, ctx); return d; } + /// The drafter's taps (the target's layers it reads). + pub fn taps(d: *const Draft) []const u32 { + return d.c.taps[0..d.c.n_taps]; + } + + /// The target's window of `rows` rows takes `ms` (told by cuda_native.zig before the facts are read). + pub fn setCost(d: *Draft, rows: u32, ms: f64) void { + if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); + if (d.n_costs < d.costs.len) { + d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; + d.n_costs += 1; + } + } + pub fn deinit(d: *Draft) void { var it = d.streams.valueIterator(); while (it.next()) |ctx| d.drop(ctx.*); @@ -253,8 +290,8 @@ pub const Draft = struct { ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.c.ringBytes()) }; errdefer ctx.ring.free(); try ctx.ring.fill8(0, null); - try d.absorb(ctx, d.taps_in, 0, 1); // any tap rows do: the drafter's own zeroed staging - try d.hold(ctx, 0, 1, d.c.block - 1); + try d.absorb(ctx, d.taps_in, 0, 1, null); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.c.block - 1, null); ctx.* = .{ .ring = ctx.ring }; return ctx; } @@ -295,11 +332,13 @@ pub const Draft = struct { } /// Target rows at positions [start, start + n) from the stream's tap ring into its context (replacing what was - /// held at or past `start`); rows older than the window are skipped. - fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize) !void { + /// held at or past `start`); rows older than the window are skipped. With `rows`, the ring's rows start + rows[i] + /// (a tree's kept path, which the target keeps at the positions start + i) instead, n of them. + fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32) !void { const c = d.c; var start = start_; var n = n_; + if (rows_) |rs| if (rs.len != n or n > engine.most_window) return error.BadDraftRows; if (start < ctx.end) ctx.end = start; if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here if (n > c.window) { @@ -311,7 +350,11 @@ pub const Draft = struct { const e = d.e; while (n > 0) { const rows = @min(n, max_rows); - try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + if (rows_) |rs| { + var tr: TapRows = .{ .n = @intCast(rows), .rows = @splat(0) }; + for (rs, 0..) |r, i| tr.rows[i] = @intCast(r); + try e.run(d.fns.gather_rows, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), tr, @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + } else try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); try d.proj(d.w.fc, d.taps_in, d.th, rows); try d.norm(d.th, d.w.hidden_norm, d.thn, rows); for (d.w.layers, 0..) |l, i| { @@ -327,8 +370,9 @@ pub const Draft = struct { if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; } - /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held. - fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize) !void { + /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held as a chain, + /// or the `lanes` tree's tokens. + fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { const c = d.c; const e = d.e; const B = c.block; @@ -369,34 +413,62 @@ pub const Draft = struct { try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, rows * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + hp_at, d.hp, rows * c.rank * 2, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); - d.select(ctx, pending, rows); - ctx.held_n = @min(depth, rows); + if (lanes_) |l| { + try d.selectLanes(ctx, pending, rows, l); + if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); + } else { + d.select(ctx, pending, rows); + ctx.held_n = @min(depth, rows); + if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); + } } - /// The path through the rows' candidates: each row's best of unary logit + predecessor . projected . successor. - fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { + /// Row `r`'s candidate of `rank` after the token `pred`: by unary logit + predecessor . projected . successor, + /// ties to the earlier candidate (the row's top-k order). + fn choice(d: *const Draft, r: usize, pred: u32, rank: usize) u32 { const c = d.c; - const words = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][0 .. rows * lp_words * 4]); - const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][0 .. rows * c.rank * 2]); + const w = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][r * lp_words * 4 ..][0 .. lp_words * 4]); + const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][r * c.rank * 2 ..][0 .. c.rank * 2]); + var scores: [lp_top]f32 = undefined; + for (0..c.topk) |j| { + const cand = w[2 + j]; + const unary: f32 = @bitCast(w[2 + lp_top + j]); + var edge: f32 = 0; + for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[k]) * bf16(d.w.successor[cand * c.rank + k]); + scores[j] = unary + edge; + } + var taken: [lp_top]bool = @splat(false); + var pick: usize = 0; + for (0..rank + 1) |_| { + var best: f32 = -std.math.inf(f32); + pick = 0; + for (0..c.topk) |j| if (!taken[j] and scores[j] > best) { + best = scores[j]; + pick = j; + }; + taken[pick] = true; + } + return w[2 + pick]; + } + + /// The path through the rows' candidates: each row's best after the one before. + fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { var pred = pending; for (0..rows) |r| { - const w = words[r * lp_words ..][0..lp_words]; - var best: f32 = -std.math.inf(f32); - var pick: u32 = w[2]; - for (0..c.topk) |j| { - const cand = w[2 + j]; - const unary: f32 = @bitCast(w[2 + lp_top + j]); - var edge: f32 = 0; - for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[r * c.rank + k]) * bf16(d.w.successor[cand * c.rank + k]); - const score = unary + edge; - if (score > best) { - best = score; - pick = cand; - } - } - ctx.held[r] = pick; - pred = pick; + ctx.held[r] = d.choice(r, pred, 0); + pred = ctx.held[r]; + } + } + + /// The tree's lanes in order: lane i the rank-r candidate at its depth's row after its parent lane's token. + fn selectLanes(d: *Draft, ctx: *Ctx, pending: u32, rows: usize, l: *const lanes.shape.Shape) !void { + if (l.parents.len > ctx.held.len) return error.TooManyLanes; + for (l.parents, l.depths, l.ranks, 0..) |par, depth, rank, i| { + if (depth >= rows or rank >= d.c.topk or par >= @as(i32, @intCast(i))) return error.BadLanes; + const pred = if (par < 0) pending else ctx.held[@intCast(par)]; + ctx.held[i] = d.choice(depth, pred, rank); } + ctx.held_n = l.parents.len; } // the Drafter contract @@ -416,7 +488,8 @@ pub const Draft = struct { fn factsFn(p: *anyopaque) dr.Facts { const d = cast(p); - return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes() }; + // TF_NO_TREES=1: chains only (for comparing) + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = std.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; } fn openFn(p: *anyopaque, s: *Stream) anyerror!void { @@ -436,7 +509,9 @@ pub const Draft = struct { const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) - try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows); + if (it.rows) |rs| { + try d.absorb(ctx, it.features.buffer, it.features.offset, rs.len, rs); + } else try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows, null); } } @@ -446,7 +521,7 @@ pub const Draft = struct { defer timing("hold", d.io, t0); for (items) |it| { const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; - try d.hold(ctx, it.pending, it.position - 1, it.depth); + try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); } } diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 4ce0cfceb5..be59726fc6 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,14 +1,20 @@ //! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the //! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails -//! and KV cache (Cache); the widths come from config.json (cuda_geometry.zig), the weights from cuda_weights.zig. -//! A forward's launches are captured into a CUDA graph the first time its shape is seen and replayed after. +//! and KV cache (cuda_cache.zig); the widths come from config.json (cuda_geometry.zig), the weights from +//! cuda_weights.zig. A forward's launches are captured into a CUDA graph the first time its shape is seen and +//! replayed after. A verify window's rows may form a tree (each row's parent given): the forward computes every +//! branch, logs each row, and the stream's state stays as it was until a path is kept. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); const geometry = @import("cuda_geometry.zig"); const weights = @import("cuda_weights.zig"); +const caches = @import("cuda_cache.zig"); pub const Geometry = geometry.Geometry; +pub const Cache = caches.Cache; +pub const tap_rows = caches.tap_rows; +const Seg = caches.Seg; const Affine4 = cuda.qlinear.Affine4; const DeviceBuffer = cuda.DeviceBuffer; const head_dim = geometry.head_dim; @@ -17,12 +23,10 @@ const head_dim = geometry.head_dim; pub const max_rows = 128; /// Streams' windows in one forward at most. pub const max_segs = 32; -/// Rows of a verify forward at most: each row's recurrences are recorded for a partial acceptance. -const most_verify = 32; -/// A window's rows at most (exact: one-row steps' bits). -const most_window = 16; -/// Device bytes the recorded rows may take: a wide model records fewer rows a forward. -const snapshot_budget = 1280 << 20; +/// Rows of a verify forward at most: each row's logits are kept and drawn. +pub const most_verify = 32; +/// A window's rows at most (exact: one-row steps' bits; qwen35.cuh's MAXW, the rows a stream's log holds). +pub const most_window = 16; /// Partial attention chunks a row at most: the chunk width grows with the window instead. const max_chunks = 64; @@ -31,60 +35,8 @@ pub fn aligned(n: usize) usize { return std.mem.alignForward(usize, n, 256); } -/// qwen35.cu's Geo and Seg. -const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int }; -const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; - -/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. -pub const tap_rows = 2048 + max_rows; - -/// One stream's state on the device and the logits of its newest prompt chunk on the host. -pub const Cache = struct { - state: DeviceBuffer, - kv: DeviceBuffer, - /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. - tap: ?DeviceBuffer = null, - cap: usize, - len: usize = 0, - logits: []u16, - /// The newest verify window: where it started, its first row in that forward, its rows. - last: ?struct { start: usize, row0: usize, rows: usize } = null, - - pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { - var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); - errdefer state.free(); - try state.fill8(0, null); - var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); - errdefer kv.free(); - try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) - var tap: ?DeviceBuffer = null; - if (e.n_taps > 0) { - tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); - try tap.?.fill8(0, null); - } - errdefer if (tap) |*t| t.free(); - return .{ .state = state, .kv = kv, .tap = tap, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; - } - - /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). - pub fn reset(self: *Cache, e: *const Engine) !void { - self.len = 0; - self.last = null; - try self.state.fill8(0, e.stream.handle); - } - - pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { - self.state.free(); - self.kv.free(); - if (self.tap) |*t| t.free(); - gpa.free(self.logits); - } - - fn seg(self: *const Cache, e: *const Engine, row0: usize, rows: usize) Seg { - const g = e.g; - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; - } -}; +/// qwen35.cuh's Geo. +pub const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int, nl: c_int }; pub const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; pub const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; @@ -115,12 +67,15 @@ const Fns = struct { tap: cuda.Function, }; -/// A window of one stream in a forward. -pub const Part = struct { cache: *Cache, rows: usize }; +/// qwen35_keep.cu's kernels (Cache.keep). +pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: cuda.Function }; + +/// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). +pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; /// What a forward keeps. A prompt chunk (`last`, one window) puts its last row's logits in its cache's `logits`; /// verify windows keep every row, drawn on the device into `tokens` (`draws`, one a row) or their logits for -/// `rowLogits` (`logits`), and record each row's state for `rewind`. +/// `rowLogits` (`logits`), and log each row for `Cache.keep`. pub const Keep = union(enum) { last, draws: []const Draw, logits }; /// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. @@ -132,16 +87,17 @@ pub const Engine = struct { g: Geometry, geo: Geo, stream: cuda.Stream, - mods: [6]cuda.Module, + mods: [7]cuda.Module, draws: cuda.sample.Draws, affine: cuda.qlinear.Affine4Kernels, fns: Fns, + keeps: Keeps, /// Every weight and scratch buffer in one allocation: the WDDM driver's launch cost grows with a process's /// allocation count, and a 64-layer model would have thousands. arena: DeviceBuffer, arena_at: usize = 0, stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there - /// Rows a verify forward records (snapshot_budget over a row's states), and a window's rows at most. + /// Rows a verify forward takes at most, and a window's rows at most. verify_rows: usize, window_rows: usize, // weights @@ -179,9 +135,8 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, - rec_snap: u64, // recorded recurrences: verify_rows rows a linear layer - conv_snap: u64, // recorded convolution tails, the same - // the per-forward tables: ids, segments, each row's segment and the draws' rules, staged in pinned memory + // the per-forward tables: ids, segments, each row's segment, parent and ancestors, the draws' rules, staged in + // pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, host_logits: cuda.HostBuffer, @@ -196,11 +151,13 @@ pub const Engine = struct { const ids_at = 0; const segs_at = max_rows * 4; const rowseg_at = segs_at + max_segs * @sizeOf(Seg); - const rules_at = rowseg_at + max_rows * 4; + const rowpar_at = rowseg_at + max_rows * 4; + const rowanc_at = rowpar_at + max_rows * 4; + const rules_at = rowanc_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows - const tables_len = params_at + 4 * 4; + const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows, record + const tables_len = params_at + 8 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -212,9 +169,9 @@ pub const Engine = struct { e.gpa = gpa; e.d = d; e.g = g; - e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads) }; - e.verify_rows = std.math.clamp(snapshot_budget / (g.deltas() * (g.recBytes() + g.convBytes())), 1, most_verify); - e.window_rows = @min(most_window, e.verify_rows); + e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads), .nl = @intCast(g.deltas()) }; + e.verify_rows = most_verify; + e.window_rows = most_window; e.arena_at = 0; e.graphs = .empty; e.n_taps = 0; @@ -226,7 +183,7 @@ pub const Engine = struct { e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); errdefer e.stage.free(); const kk = cuda.kernels; - const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample }; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample, kk.qwen35_keep }; var loaded: usize = 0; errdefer for (e.mods[0..loaded]) |*m| m.unload(); for (images, 0..) |img, i| { @@ -239,6 +196,7 @@ pub const Engine = struct { errdefer e.affine.deinit(); e.draws = try cuda.sample.Draws.resolve(e.mods[5]); inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); + inline for (@typeInfo(Keeps).@"struct".field_names) |name| @field(e.keeps, name) = try e.mods[6].function("tf_q35k_" ++ name); e.layers = try gpa.alloc(Layer, g.layers); errdefer gpa.free(e.layers); var ck = try core.Checkpoint.openModel(gpa, io, dir); @@ -251,17 +209,26 @@ pub const Engine = struct { errdefer e.staged.free(); e.host_logits = try cuda.HostBuffer.alloc(d, most_verify * g.vocab * 2); errdefer e.host_logits.free(); - // every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds - // (a 16 GiB model: a 20 s launch, then a 3 s one), which no request should pay - var warm = try Cache.init(gpa, e, 64); - defer warm.deinit(gpa); - const ids: [32]u32 = @splat(0); - try e.forward(&.{.{ .cache = &warm, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs - try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = null, .position = 33 }} }); - try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = cuda.sample.rule(1, 0.7, 0.95, -std.math.inf(f64), 20), .position = 34 }} }); + try e.warm(gpa); return e; } + /// Every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds (a 16 GiB + /// model: a 20 s launch, then a 3 s one), which no request should pay. + fn warm(e: *Engine, gpa: std.mem.Allocator) !void { + var warm_cache = try Cache.init(gpa, e, 64); + defer warm_cache.deinit(gpa); + const ids: [32]u32 = @splat(0); + const greedy: Draw = .{ .rule = null, .position = 33 }; + try e.forward(&.{.{ .cache = &warm_cache, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs + try e.forward(&.{.{ .cache = &warm_cache, .rows = 1 }}, ids[0..1], .{ .draws = &.{greedy} }); + try e.forward(&.{.{ .cache = &warm_cache, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = .{ .seed = 1, .inv_t = 1.0 / 0.7, .top_p = 0.95, .near = 20.0, .min_log = -std.math.inf(f32), .top_k = 20 }, .position = 34 }} }); + // a tree window and a branch of it kept: the keep kernels, and the forward's branch path + try e.forward(&.{.{ .cache = &warm_cache, .rows = 3, .parents = &.{ -1, 0, 0 } }}, ids[0..3], .{ .draws = &.{ greedy, greedy, greedy } }); + try warm_cache.keep(e, &.{ 0, 2 }); + try e.stream.synchronize(); + } + pub fn deinit(e: *Engine) void { var it = e.graphs.valueIterator(); while (it.next()) |x| x.deinit(); @@ -278,9 +245,9 @@ pub const Engine = struct { e.gpa.destroy(e); } - /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps. + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps, log. pub fn streamBytes(e: *const Engine, cap: usize) usize { - return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes(); + return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes() + e.g.logBytes(most_window); } /// A stream's tap ring: tap_rows rows of the tapped layers' residuals. @@ -335,8 +302,6 @@ pub const Engine = struct { e.pl = try t.alloc(R * g.q_heads * max_chunks * 4); e.po = try t.alloc(R * g.q_heads * max_chunks * head_dim * 4); e.logits = try t.alloc(most_verify * g.vocab * 2); - e.rec_snap = try t.alloc(g.deltas() * e.verify_rows * g.recBytes()); - e.conv_snap = try t.alloc(g.deltas() * e.verify_rows * g.convBytes()); return total; } @@ -363,12 +328,34 @@ pub const Engine = struct { try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } - /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. + /// The row tables of `parts`: each forward row's segment, parent forward row (-1: its window's first) and the + /// ancestors-and-self bits over its window's rows (a chain without `parents`). + fn rowTables(st: []u8, parts: []const Part) !void { + const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); + const rowpar = std.mem.bytesAsSlice(c_int, st[rowpar_at..][0 .. max_rows * 4]); + const rowanc = std.mem.bytesAsSlice(u32, st[rowanc_at..][0 .. max_rows * 4]); + var row0: usize = 0; + for (parts, 0..) |p, i| { + if (p.parents) |ps| if (ps.len != p.rows or p.rows > most_window or ps[0] != -1) return error.InvalidQwenParents; + for (0..p.rows) |r| { + const parent: i32 = if (p.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (parent >= @as(i32, @intCast(r)) or parent < -1) return error.InvalidQwenParents; + rowseg[row0 + r] = @intCast(i); + rowpar[row0 + r] = if (parent < 0) -1 else @intCast(row0 + @as(usize, @intCast(parent))); + rowanc[row0 + r] = (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); + } + row0 += p.rows; + } + } + + /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. A + /// window not yet kept from the stream's last verify is kept whole first. pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { const g = e.g; var total: usize = 0; var widest: usize = 0; for (parts) |p| { + try p.cache.settle(e); if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; total += p.rows; widest = @max(widest, p.cache.len + p.rows); @@ -379,6 +366,7 @@ pub const Engine = struct { if (!record and parts.len != 1) return error.InvalidQwenSegments; if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; + for (parts) |p| if (p.parents != null and !record) return error.InvalidQwenParents; const heads: usize = if (record) total else 1; // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps const prompt = !record; @@ -386,12 +374,10 @@ pub const Engine = struct { const st = e.staged.bytes; @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids); const segs = std.mem.bytesAsSlice(Seg, st[segs_at..][0 .. parts.len * @sizeOf(Seg)]); - const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); + try rowTables(st, parts); var row0: usize = 0; for (parts, 0..) |p, i| { - segs[i] = p.cache.seg(e, row0, p.rows); - for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); - p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; + segs[i] = p.cache.seg(e, p.cache.len, row0, p.rows); row0 += p.rows; } if (keep == .draws) { @@ -404,11 +390,12 @@ pub const Engine = struct { } // attention chunks: at most max_chunks a row, each a multiple of 128 keys const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); - const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..16]); + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..32]); params[0] = @intCast((widest + ck - 1) / ck); params[1] = @intCast(ck); params[2] = @intCast(total); params[3] = @intCast(parts.len); + params[4] = @intFromBool(record); try e.tables.uploadAsync(0, st, e.stream.handle); // a captured forward's launches cover a bucket of rows and every window; its kernels read the true counts // (a lone stream's one-row step, the common case without a drafter, keeps its own exact graph) @@ -437,7 +424,11 @@ pub const Engine = struct { } try gop.value_ptr.launchOn(e.stream); } - for (parts) |p| p.cache.len += p.rows; + for (parts) |p| { + p.cache.last = if (record and p.rows > 1) .{ .start = p.cache.len, .rows = p.rows } else null; + p.cache.len += p.rows; + p.cache.written = p.cache.len; + } switch (keep) { .draws => |pk| { const toks = e.tables.ptr + tokens_at; @@ -484,6 +475,8 @@ pub const Engine = struct { const t_ids = e.tables.ptr + ids_at; const t_segs = e.tables.ptr + segs_at; const t_rows = e.tables.ptr + rowseg_at; + const t_par = e.tables.ptr + rowpar_at; + const t_anc = e.tables.ptr + rowanc_at; const t_params = e.tables.ptr + params_at; const geo = e.geo; const group = g.q_heads / g.kv_heads; @@ -496,15 +489,13 @@ pub const Engine = struct { switch (l.mixer) { .delta => |m| { const layer: c_int = @intCast(di); - const rec_snap: u64 = if (record) e.rec_snap + di * e.verify_rows * g.recBytes() else 0; - const conv_snap: u64 = if (record) e.conv_snap + di * e.verify_rows * g.convBytes() else 0; try e.proj(m.qkv, e.x, e.qkv, total, prompt); try e.proj(m.z, e.x, e.z, total, prompt); try e.proj(m.a, e.x, e.a, total, prompt); try e.proj(m.b, e.x, e.b, total, prompt); - try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); - try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap, t_params }); - try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap, t_params }); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, t_par, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, t_par, layer, t_params }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, t_par, layer, e.y, t_params }); try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix, t_params }); try e.proj(m.out, e.mix, e.r, total, prompt); di += 1; @@ -516,8 +507,8 @@ pub const Engine = struct { try e.proj(m.v, e.x, e.vraw, total, prompt); try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai, t_params }); - try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); + try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); try e.proj(m.o, e.mix, e.r, total, prompt); @@ -543,23 +534,4 @@ pub const Engine = struct { pub fn rowLogits(e: *const Engine, i: usize) []const u16 { return e.host_logits.slice(u16)[i * e.g.vocab ..][0..e.g.vocab]; } - - /// Keeps the first `rows` of the cache's newest window: its recurrences and convolution tails as they were after - /// that row (the forward recorded them), its length at that row. The window's later keys stay as unread slots. - pub fn rewind(e: *Engine, cache: *Cache, rows: usize) !void { - const g = e.g; - const last = cache.last orelse return error.NothingToKeep; - if (rows == 0 or rows > last.rows) return error.NothingToKeep; - cache.last = null; - cache.len = last.start + rows; - if (rows == last.rows) return; - const row = last.row0 + rows - 1; - const rb = g.recBytes(); - const cb = g.convBytes(); - for (0..g.deltas()) |di| { - try cache.state.copyFrom(di * rb, e.rec_snap + (di * e.verify_rows + row) * rb, rb, e.stream.handle); - try cache.state.copyFrom(g.deltas() * rb + di * cb, e.conv_snap + (di * e.verify_rows + row) * cb, cb, e.stream.handle); - } - try e.stream.synchronize(); - } }; diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index 34207bacb5..7b327fd1a9 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -70,6 +70,17 @@ pub const Geometry = struct { return g.deltas() * (g.recBytes() + g.convBytes()); } + /// One verified row's record of one linear layer (qwen35.cuh's log_bytes): decay and write by value head, + /// key by key head, the tail after it. + pub fn logRowBytes(g: Geometry) usize { + return g.v_heads * 4 + g.v_heads * dk * 4 + g.k_heads * dk * 2 + g.convBytes(); + } + + /// A stream's log of a window's `rows` rows. + pub fn logBytes(g: Geometry, rows: usize) usize { + return rows * g.deltas() * g.logRowBytes(); + } + /// A stream's keys and values at a `cap`-token window. pub fn kvBytes(g: Geometry, cap: usize) usize { return g.attentions() * 2 * g.kv_heads * cap * head_dim * 2; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 66b25c02c8..d4fa634cb1 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -77,19 +77,24 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); } - // with a drafter: a lone stream's windows of the pending token and 1 to 7 drafts - if (l.e.n_taps > 0) for (2..@min(8, l.e.verify_rows) + 1) |rows| { + // with a drafter: a lone stream's windows of the pending token and its drafts, chains and trees + if (l.e.n_taps > 0) for (2..l.e.window_rows + 1) |rows| { for (picks[0..rows]) |*pk| pk.* = .{ .rule = null, .position = 1 }; try l.e.forward(&.{.{ .cache = l.pool.items[0], .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); try l.pool.items[0].reset(l.e); }; } -/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). +/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). The +/// target's windows of 1 to window_rows rows are timed here (with the drafter's taps on) for the round loop's prices. pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; const d = try Draft.init(gpa, io, l.e, dir); + errdefer d.deinit(); + try l.e.setTaps(d.taps()); + try l.timeWindows(d); return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; } @@ -121,7 +126,7 @@ fn release(p: *anyopaque) void { fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { const s = sampling orelse return .{ .rule = null, .position = position }; if (s.temperature <= 0) return .{ .rule = null, .position = position }; - return .{ .rule = cuda.sample.rule(s.seed, s.temperature, s.top_p, s.minLog(), s.top_k), .position = position }; + return .{ .rule = cuda.sample.Rule.from(lanes.gpu_rule.rule(s)), .position = position }; } const Lanes = struct { @@ -164,6 +169,26 @@ const Lanes = struct { self.pool.append(self.gpa, cache) catch self.drop(cache); } + /// Each window width's forward, timed on a cache of its own (its third run: the shape's graph captured and the + /// driver's stall behind it, the previous window's keep in front as in a round), into the drafter's facts. + fn timeWindows(self: *Lanes, d: *Draft) !void { + var cache = try E.Cache.init(self.gpa, self.e, 64); + defer cache.deinit(self.gpa); + const ids: [E.most_window]u32 = @splat(0); + var picks: [E.most_window]E.Draw = undefined; + for (&picks) |*pk| pk.* = .{ .rule = null, .position = 1 }; + for (1..self.e.window_rows + 1) |rows| { + var ms: f64 = 0; + for (0..3) |run| { + try cache.reset(self.e); + const t0 = std.Io.Clock.awake.now(self.io); + try self.e.forward(&.{.{ .cache = &cache, .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + if (run == 2) ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(self.io)).toNanoseconds())) / 1e6; + } + d.setCost(@intCast(rows), ms); + } + } + fn backend(self: *Lanes) be.Backend { return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; } @@ -224,7 +249,8 @@ const Lanes = struct { return @intCast(handle); } - /// Each window (a stream's pending token and copies) checked, then as many as fit in each forward. + /// Each window (a stream's pending token and drafts, a chain or a tree) checked, then as many as fit in each + /// forward. A tree's rows are keyed at the cache length plus their depth, as the round loop positions them. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { const self = cast(ptr); const t0 = std.Io.Clock.awake.now(self.io); @@ -232,11 +258,18 @@ const Lanes = struct { if (windows.len != out.len) return error.InvalidQwenWindows; for (windows, out) |w, o| { const rows = w.rows(); - if (w.parents != null or w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; const cache = try self.cacheOf(w.stream); - for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidQwenPosition; + var depth: [E.most_window]u64 = undefined; + for (w.positions, 0..) |p, r| { + const parent: i32 = if (w.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (r == 0 and parent != -1) return error.UnsupportedQwenWindow; + if (r > 0 and (parent < 0 or parent >= @as(i32, @intCast(r)))) return error.UnsupportedQwenWindow; + depth[r] = if (parent < 0) 0 else depth[@intCast(parent)] + 1; + if (p != cache.len + depth[r] + 1) return error.InvalidQwenPosition; + } } var at: usize = 0; while (at < windows.len) { @@ -248,7 +281,7 @@ const Lanes = struct { var masked = false; while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= self.e.verify_rows) : (n += 1) { const w = windows[at + n]; - parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows() }; + parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows(), .parents = w.parents }; ids[total] = w.pending; @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); for (w.positions, 0..) |p, r| picks[total + r] = pick(w.stream.sampling, p); @@ -270,14 +303,15 @@ const Lanes = struct { } } - /// The accepted prefix of each window: its state as it was after that row. + /// The accepted path of each window (a chain's prefix or a tree's branch): the state as that chain leaves it. fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { const self = cast(ptr); const t0 = std.Io.Clock.awake.now(self.io); defer @import("cuda_draft.zig").timing("keep", self.io, t0); for (windows, paths) |w, path| { - for (path, 0..) |r, i| if (r != i) return error.UnsupportedQwenTree; - try self.e.rewind(try self.cacheOf(w.stream), path.len); + const cache = try self.cacheOf(w.stream); + if (cache.last == null) continue; // a one-row window (or one kept already): nothing to roll back + try cache.keep(self.e, path); } } @@ -293,13 +327,14 @@ const Lanes = struct { } /// Rows [start, start + count) of a stream's tap ring: `offset` carries the first row's position, the drafter - /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. + /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. The newest window's + /// rows stay readable after a keep (a tree's kept path reads them by row) until the next forward. fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { const self = cast(ptr); if (taps.len != self.e.n_taps) return error.TapsMismatch; const cache = try self.cacheOf(stream); const ring = cache.tap orelse return error.NoTaps; - if (start + count > cache.len or count > E.tap_rows) return error.RowsNotHeld; + if (start + count > cache.written or count > E.tap_rows) return error.RowsNotHeld; return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; } diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 27ca52012f..ae7e82d1e5 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -193,6 +193,10 @@ const Host = struct { fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain) !void { h.cfg = try lanes.Config.init(h.gpa, facts, rows, rows - 1); errdefer h.cfg.deinit(h.gpa); + if (facts.head_trees) { // a drafter's trees: the planner picks a lone stream's lanes, as deep as it drafts + h.cfg.head_lanes = rows - 1; + h.cfg.head_depth = @max(1, facts.drafts); + } h.clock = .{ .io = io }; h.core = lanes.Engine.init(h.gpa, &h.cfg, h.backend(), h.clock.clock()); errdefer h.core.deinit(); diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 78414fdb32..858dc34c07 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,7 +1,8 @@ -//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over comma-separated -//! token ids: one token a forward (step: every position's logits), windows of growing widths (window: the same rows, -//! which must equal step's) or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into -//! for a reference to compare. Step mode also prints the decode speed. +//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over +//! comma-separated token ids: one token a forward (step: every position's logits), windows of growing widths +//! (window: the same rows, which must equal step's), tree windows with junk branches kept to the true one (tree: +//! the same), or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into for a +//! reference to compare. Step mode also prints the decode speed. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -11,7 +12,7 @@ pub fn main(init: std.process.Init) !u8 { const a = init.arena.allocator(); const args = try init.minimal.args.toSlice(a); if (args.len != 5) { - std.debug.print("usage: tf-qwen35-cuda \n", .{}); + std.debug.print("usage: tf-qwen35-cuda \n", .{}); return 2; } var ids: std.ArrayList(u32) = .empty; @@ -49,6 +50,42 @@ pub fn main(init: std.process.Init) !u8 { for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); at += n; } + } else if (std.mem.eql(u8, args[2], "tree")) { + // tree windows of growing widths: the true continuation down one branch, a junk sibling before each of its + // rows and a junk grandchild, the true branch kept; its rows' logits must equal step's bit for bit + var at: usize = 0; + var w: usize = 2; + while (at < ids.items.len) : (w = w % e.window_rows + 1) { + var toks: [E.most_window]u32 = undefined; + var parents: [E.most_window]i32 = undefined; + var path: [E.most_window]u32 = undefined; + toks[0] = ids.items[at]; + parents[0] = -1; + path[0] = 0; + var n: usize = 1; + var kept: usize = 1; + var last: i32 = 0; + var depth: usize = 1; + while (n + 2 <= w and at + depth < ids.items.len) : (depth += 1) { + toks[n] = (ids.items[at + depth] + 7) % @as(u32, @intCast(e.g.vocab)); + parents[n] = last; + toks[n + 1] = ids.items[at + depth]; + parents[n + 1] = last; + path[kept] = @intCast(n + 1); + last = @intCast(n + 1); + n += 2; + kept += 1; + } + if (n < w and n >= 3) { + toks[n] = 1; + parents[n] = 1; + n += 1; + } + try e.forward(&.{.{ .cache = &cache, .rows = n, .parents = parents[0..n] }}, toks[0..n], .logits); + for (path[0..kept]) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + if (n > 1) try cache.keep(e, path[0..kept]); + at += kept; + } } else { var at: usize = 0; while (at < ids.items.len) { From 61fc2bf9412eca769848af16eef7eed322c569dd Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:06:09 +0000 Subject: [PATCH 11/46] docs: the Qwen hybrids on CUDA in the Windows page and the README What the port serves beyond Nemotron: Qwen3.5-2B and Qwen3.8-27B with its DFlash2 drafter, the speeds measured on an RTX 4090, the load time and why, the timing and chain-only switches, and what is not there. --- README.md | 5 +++-- docs/WINDOWS.md | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a6a018278f..c838faa144 100644 --- a/README.md +++ b/README.md @@ -73,8 +73,8 @@ The platform column names the hardware tested for each model. | Qwen3.8 Flash Next | MLX affine 6-bit, group 32 | Metal on M5 Ultra | | GLM-5.3-Flash | MLX affine 4-bit, group 64 | Metal on two M5 Ultras | | GLM-5.3 | EXLR8 `home` profile (trellis K2/K3/K4 experts) | Metal on two M5 Ultras | -| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max | -| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra | +| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (unofficial Windows port) | +| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (unofficial Windows port) | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. NVIDIA's NVFP4 checkpoint, `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`, serves on CUDA as stored; its bf16 MTP head is quantized to 4-bit at load for drafts. @@ -90,6 +90,7 @@ Its prompt processing trails mlx_lm for now, and we are fixing it. An M3 Ultra r Bonsai, Gemma 4, Qwen3.6 and DeepSeek-V4 are still under qualification for 1.0.x. The Python 0.6.6 engine remains available at the `v0.6.6` tag for those models and other backends. On CUDA, 1.0.11 serves Nemotron on a GB10, greedy and sampled, with concurrent requests sharing each round. +The unofficial Windows port ([docs/WINDOWS.md](docs/WINDOWS.md)) also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. The Linux x86_64 archive adds NVIDIA Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), tested on an RTX 3090, and Blackwell cards with compute capability 12.0 (the RTX 50 series and RTX PRO 6000 Blackwell), tested on an RTX PRO 6000; the A100 (8.0) is not supported yet. ## Exact decoding diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index fdb9b7f88f..b11d72f4fd 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -34,6 +34,32 @@ tensorfold-native serve --name local-model --parallel 1 --context 8 The server listens on `http://127.0.0.1:8888/v1`. Ctrl+C (or closing the console) stops it cleanly. +## Models on CUDA + +Nemotron 3.5 Lightning serves as on Linux. The Qwen3.5 hybrids have a CUDA engine of their own in this port +(`zig/src/families/qwen3_5/cuda*.zig`, kernels `zig/kernels/cuda/qwen35*.cu`), the Metal kernels' arithmetic with +the same rounding points, so a window of up to 16 rows reproduces one-token steps bit for bit: + +```bat +tensorfold-native serve --name qwen2b --context 32768 +tensorfold-native serve --name qwen27b --context 8192 --parallel 4 ^ + --drafter +``` + +- Qwen3.5-2B: about 340 tokens/s for one request on an RTX 4090, 1,500 across eight. +- Qwen3.8-27B (16 GiB): 40 tokens/s plain, at the card's memory bandwidth; 240 across eight requests. +- `--drafter` (4-bit only) loads z-lab's DFlash2 draft model: the target writes five layers' residuals into a ring + for it, and it drafts a tree of up to 12 lanes a round that the round loop's planner shapes from what the target + took. A lone greedy request decodes at 80 to 160 tokens/s by how predictable the text is (prose at the low end, + code at the high), sampled ones at 50 to 70, and four at once reach about 160 aggregate; replies stay identical to + plain decoding. Loading takes about 80 s: the drafter's weights are quantized to 4 bits on the host and every + shape is run once, as the driver stalls a process's first use of each kernel and buffer. +- A verify window's rows are not snapshotted: each stream logs its rows' recurrence updates and a keep replays the + accepted path, so the 27B records 32 rows a forward where 8 fit before. +- `TF_TIMING=1` prints each verify, keep, absorb and hold's milliseconds and the drafted shape (chain or lanes); + `TF_NO_TREES=1` drafts chains only. +- No prompt cache (`--prompt-cache-gib`), `--learn` or `--slide` for these two models. + ## How the port works - `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating From 6aaf241affa27864e08898a4648bc4e32a491015 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:37 +0000 Subject: [PATCH 12/46] qwen3_5, platform: comments on one line The lean rule (tools/zig/lean_check.py) counts a comment that runs on to a second line; the Windows port's new files had a hundred such runs. Each is one line of at most 120 columns now, and the checker reports the same 539 problems as main. --- zig/build/sys.zig | 3 +- zig/kernels/cuda/qwen35.cu | 45 +++++-------------- zig/kernels/cuda/qwen35.cuh | 9 ++-- zig/kernels/cuda/qwen35_draft.cu | 15 ++----- zig/kernels/cuda/qwen35_keep.cu | 11 ++--- zig/src/cli/pull_file.zig | 3 +- zig/src/core/lanes/drafter.zig | 6 +-- zig/src/families/qwen3_5/cuda.zig | 3 +- zig/src/families/qwen3_5/cuda_cache.zig | 8 +--- zig/src/families/qwen3_5/cuda_draft.zig | 24 +++------- .../families/qwen3_5/cuda_draft_weights.zig | 7 +-- zig/src/families/qwen3_5/cuda_engine.zig | 35 +++++---------- zig/src/families/qwen3_5/cuda_geometry.zig | 6 +-- zig/src/families/qwen3_5/cuda_native.zig | 23 +++------- zig/src/families/qwen3_5/cuda_weights.zig | 3 +- zig/src/platform/sys.zig | 3 +- zig/src/platform/win_c.zig | 7 +-- zig/src/platform/win_clock.zig | 3 +- zig/src/platform/win_posix.zig | 7 +-- zig/tests/qwen35_cuda_main.zig | 9 +--- 20 files changed, 65 insertions(+), 165 deletions(-) diff --git a/zig/build/sys.zig b/zig/build/sys.zig index b34be717ed..1c1309af78 100644 --- a/zig/build/sys.zig +++ b/zig/build/sys.zig @@ -1,5 +1,4 @@ -//! `sys` (zig/src/platform/sys.zig) for every module the build compiles: one walk of the step graph after it is -//! declared, so no module's declaration needs to name it. +//! `sys` (zig/src/platform/sys.zig) for every module the build compiles, by one walk of the step graph once declared. const std = @import("std"); diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 2186037011..cef7c7b5a2 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -1,16 +1,10 @@ -// Qwen3.5's hybrid text models on CUDA (the 2B and Qwen3.8-27B): the Metal engine's kernels (zig/kernels/metal/qwen3_5) -// with the same bf16 rounding points, their widths from the model's Geo. Rows of one forward belong to segments (one -// stream's window each); a segment's state lives in its stream's blocks. +// Qwen3.5's hybrids (2B, 27B) on CUDA: the Metal kernels' arithmetic and rounding; a segment is one stream's window. #include "qwen35.cuh" using namespace tf_q35; -// The per-forward counts a captured forward reads from the device, so one graph serves every shape: P[0] attention -// chunks, P[1] their width, P[2] rows, P[3] windows, P[4] whether the rows are recorded for keep (a verify window: -// its state stays as it was until the kept path is applied). A null P means the launch's grid is exact. -// A window's rows may form a tree: rowpar names each forward row's parent forward row (-1: the window's first row, -// after the stream's cache), rowanc each row's ancestors and itself as bits over its window's rows. +// P (device counts, one graph for every shape): chunks, width, rows, windows, recorded; rowpar/rowanc give the tree. __device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } __device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } __device__ __forceinline__ bool recorded(const int* P) { return P != nullptr && P[4] != 0; } @@ -30,8 +24,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __rest } } -// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), D wide (a multiple of 256, at most 8192). One 256-thread -// block a row; each thread's elements in index order, warps summed in order. +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)); D a multiple of 256 up to 8192; a 256-thread block a row. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, const int* __restrict__ P) { @@ -82,9 +75,7 @@ extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, floa } } -// The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. -// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. A row's taps -// are its parents (then the stream's tail), so a tree's rows each see their own branch. +// Conv, silu, q/k norms, the linear heads' decay and write: a warp a (row, head) over q, k, v; taps from the parents. extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const int* __restrict__ rowpar, int layer, const __nv_bfloat16* __restrict__ CW, @@ -144,9 +135,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ } } -// Each segment's convolution tail: a recorded window logs the tail after each of its rows (the last TAPS - 1 tokens of -// the row's branch) for keep, and leaves the stream's own; a prompt chunk or one-row window updates the stream's in -// place. A thread a (segment, channel). +// A segment's conv tail: a recorded window logs each row's for keep and leaves the stream's; others update it in place. extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, const int* __restrict__ rowpar, int layer, const int* __restrict__ P) { const int C = geo.conv(); @@ -171,12 +160,7 @@ extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __re } } -// The gated delta rule over each segment's rows, each row's state its parent's updated: a chain carries the state in -// registers, a branch rebuilds its parent's from the stream's recurrence and the logged updates of the ancestors. -// A recorded window logs each row's (decay, write, key) and leaves the recurrence as it was (keep applies the kept -// path, qwen35_keep.cu); a prompt chunk or one-row window updates it in place. Value head hv reads key head -// hv / (nv / nk). Grid (1, DK / 4, nv * segments), block (32, 4): a warp a (value head, value row), its 32 lanes -// over the key width. +// Gated delta rule, each row's state its parent's: a chain keeps it in registers, a branch rebuilds it from the log. extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, @@ -201,8 +185,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons const int row = sg.row0 + r; const int parent = rowpar[row]; if (parent != prev) { - // a branch: its parent's state from the stream's and the ancestors' updates, oldest first (their keys - // and decays from this forward's rows, their writes from the log, which this thread wrote itself) + // a branch: its parent's state from the stream's and the ancestors' logged updates, oldest first int anc[MAXW]; int n = 0; for (int a = parent; a >= 0; a = rowpar[a]) anc[n++] = a; @@ -286,8 +269,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const } } -// Per-head RMSNorm of attention rows: X's head h of row m at m * stride + h * hstride, out packed (rows, heads, HD). -// Grid (rows * heads), 64 threads of 4 values. +// Per-head RMSNorm of attention rows, out packed (rows, heads, HD); grid (rows * heads), 64 threads of 4 values. extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps, const int* __restrict__ P) { @@ -325,8 +307,7 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); } -// Queries in place at their positions (the cache length plus the row's depth in its window); keys rotated and values -// copied into their segment's cache at the row's own slot (the cache length plus its row). Grid (rows), HD threads. +// Queries rotated in place, keys rotated and values copied into the row's cache slot; grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, @@ -351,10 +332,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo } } -// Causal attention over a segment's cache, P[1] keys a block over P[0] chunks (read from the device, so a captured -// forward replays as caches grow): grid (chunks at most, kvh, rows), block qh / kvh warps, a warp a query head of -// the kv head's group. A row sees the cache before its window and its window's ancestors and itself (rowanc). -// Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +// Causal attention over the cache in P[0] chunks of P[1] keys, a warp a query head; partials by (row, head, chunk). extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, const int* __restrict__ P, float scale, @@ -454,8 +432,7 @@ extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __rest OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); } -// A tapped layer's residual stream, bf16(h + r) as the next norm rounds it, into each row's segment's tap ring at -// the row's position: plane `tap` of the row's taps. Grid (rows), 256 threads. +// A tapped layer's residual bf16(h + r) into the segment's tap ring at the row's position; a 256-thread block a row. extern "C" __global__ void __launch_bounds__(256) tf_q35_tap(const __nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, const Seg* __restrict__ segs, const int* __restrict__ rowseg, int tap, int D, const int* __restrict__ P) { diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index f3baa2be6c..b4f789814f 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -1,5 +1,4 @@ -// What qwen35.cu and qwen35_keep.cu share: the model's widths (Geo), a stream's window in a forward (Seg), the per-row -// record a verify window leaves in its stream's log, and the walk from a row back through its parents. +// Shared by qwen35.cu and qwen35_keep.cu: the widths (Geo), a stream's window (Seg), a row's log record, parent walks. #pragma once #include @@ -19,8 +18,7 @@ struct Geo { __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } - // one row's record of one linear layer: decay and write strength a value head, its write a value (head, dv), its - // key a key head, and the convolution tail after it + // a row's record of a linear layer: decay, strength and write a value head, key a key head, the tail after it __device__ size_t log_bytes() const { return (size_t)nv * 4 + (size_t)nv * DK * 4 + (size_t)nk * DK * 2 + conv_rows() * 2; } }; @@ -58,8 +56,7 @@ struct Path { int rows[MAXW]; }; -// The row `steps` tokens before forward row `w` through its parents (rowpar: each forward row's parent forward row, -1 -// for a window's first row): a forward row, or -1 - t for the stream's convolution tail row t (TAPS - 2 the newest). +// The row `steps` tokens before forward row `w` through rowpar: a forward row, or -1 - t for the stream's tail row t. __device__ __forceinline__ int back(const int* __restrict__ rowpar, int w, int steps) { int a = w; for (int s = 1; s <= steps; s++) { diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu index a4726d0711..af2533bc51 100644 --- a/zig/kernels/cuda/qwen35_draft.cu +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -1,6 +1,4 @@ -// The DFlash2 drafter's glue on CUDA (zig/src/families/qwen3_5/cuda_draft.zig): a block of rows attends over the -// target's tapped context, with two-tap dynamic convolutions around attention and MLP. Rounding points follow the -// Metal drafter's draft_ops.metal; the projections run through the target engine's affine-4 kernels. +// DFlash2 drafter glue on CUDA (cuda_draft.zig): a block of rows over the target's taps, draft_ops.metal's rounding. #include #include @@ -46,9 +44,7 @@ extern "C" __global__ void tf_q35d_gather_rows(const __nv_bfloat16* __restrict__ for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; } -// y[t, c] = sum over lags l of (BASE[branch, l, c] + DYN[t, (branch * 2 + l) * groups + c / group]) * x[t - l, c], -// each product and sum rounded to bf16 in turn; x[-1] is zero (the block's first row has no past); RES, if given, -// is added last. D wide rows, `group` channels a dynamic coefficient. Grid (rows), 256 threads. +// y[t, c] = sum over lags l of (BASE[l, c] + DYN[t, l, c / group]) * x[t - l, c] in bf16 steps; x[-1] = 0; RES last. extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ DYN, const __nv_bfloat16* __restrict__ BASE, const __nv_bfloat16* __restrict__ RES, __nv_bfloat16* __restrict__ OUT, int D, int group, int branch) { @@ -69,8 +65,7 @@ extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, } } -// Per-head RMSNorm over HD then rotate-half RoPE over all HD at position pos0 + row: X's head h of row m at -// m * stride + h * HD, OUT packed (rows, heads, HD). Grid (rows * heads), 32 threads of 4 values. +// Per-head RMSNorm then rotate-half RoPE at position pos0 + row; OUT packed (rows, heads, HD), grid (rows * heads). extern "C" __global__ void __launch_bounds__(32) tf_q35d_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ OUT, int heads, int stride, long long pos0, float eps, float theta) { @@ -115,9 +110,7 @@ extern "C" __global__ void tf_q35d_ring_put(const __nv_bfloat16* __restrict__ K, } } -// The block's attention: query row m, head h sees the context ring's positions [max(begin, p - window + 1), pos0) -// where p = pos0 + m, and every block row (the drafter is not causal within its block). GQA over kv_heads; scale -// 1 / sqrt(HD). Grid (rows, heads), a warp each; online softmax in fp32, bf16 out (rows, heads * HD). +// Block attention: a row sees ring positions [max(begin, p - window + 1), pos0) and every block row; a warp a row. extern "C" __global__ void __launch_bounds__(32) tf_q35d_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int rows, int heads, int kv_heads, int W, diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu index b2dae0b2b5..20d00fc212 100644 --- a/zig/kernels/cuda/qwen35_keep.cu +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -1,13 +1,10 @@ -// Keeping a verified window's path on CUDA (zig/src/families/qwen3_5/cuda_cache.zig): the stream's recurrences take -// the path's logged updates in order, its convolution tails the path's last row's, and a tree's kept keys move down -// to the slots after the cache. The rows of a window were recorded by qwen35.cu's kernels (Seg.log). +// Keeping a verified window's path (cuda_cache.zig): the recurrences take its logged updates and its keys move down. #include "qwen35.cuh" using namespace tf_q35; -// Each linear layer's recurrence through the path's rows' (decay, key, write), as the forward's chain would have -// left it. Grid (1, DK / 4, nv * nl), block (32, 4): a warp a (layer, value head, value row), its lanes over the key. +// Each linear layer's recurrence through the path's (decay, key, write), as the forward's chain would have left it. extern "C" __global__ void __launch_bounds__(128) tf_q35k_apply(Geo geo, Seg sg, Path path) { const int NK = geo.nk, NV = geo.nv; const int layer = blockIdx.z / NV, hv = blockIdx.z % NV; @@ -43,9 +40,7 @@ extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { CS[i] = log_row(geo, sg, layer, last).conv[i]; } -// The path's rows' keys and values (slot cache length + row) down to the slots cache length + 0, 1, ... in path -// order: row j of the path is at least j, so a move never lands on a slot a later move still reads. One block a -// (attention layer, kv head): grid (attention layers, kvh), HD threads, a barrier between moves. +// The path's keys and values move down to the slots after the cache; row j lands at or below j, never on one unread. extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { const int KVH = geo.kvh; const int layer = blockIdx.x, h = blockIdx.y; diff --git a/zig/src/cli/pull_file.zig b/zig/src/cli/pull_file.zig index 1926f553ff..a5750e3e79 100644 --- a/zig/src/cli/pull_file.zig +++ b/zig/src/cli/pull_file.zig @@ -145,8 +145,7 @@ pub fn linkIntoSnapshot(a: Allocator, io: std.Io, snapshot_dir: []const u8, path std.Io.Dir.cwd().deleteFile(io, link_path) catch {}; const target = try std.fs.path.join(a, &.{ "..", "..", "blobs", blob_name }); std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch |err| { - // Windows grants symlinks only with Developer Mode or elevation (and refuses some): a hard link to the blob, - // which costs no space (huggingface_hub copies instead); learning refuses linked shards either way + // Windows grants symlinks only with Developer Mode or elevation: a hard link to the blob costs no space if (@import("builtin").os.tag != .windows) { std.log.err("cannot link {s} to its blob: {t}", .{ link_path, err }); return error.SymLinkFailed; diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 038ee35782..c956522f50 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -17,12 +17,10 @@ pub const Facts = struct { lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) }; -/// Target rows [start, start + features.rows) and the token after each; rows at or past `start` replace what was held. -/// `rows`, when given, picks the features' rows to absorb (a tree's kept path), in order; null: all of them. +/// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32, rows: ?[]const u32 = null }; -/// Draft `depth` tokens after `pending`, the first at `position`; asked only once row `position - 2` was absorbed. -/// With `lanes`, a tree instead: lane i's token is the rank-r choice after its parent (depth unused). +/// `depth` drafts after `pending` at `position`; with `lanes` a tree, lane i the rank-r choice after its parent. pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32, lanes: ?*const Shape = null }; pub const Drafter = struct { diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index f21778c7bb..435ebe162e 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -1,5 +1,4 @@ -//! The Qwen3.5 hybrids on CUDA (the 2B, Qwen3.8-27B): the engine and its native family entry (the Metal 2B engine is -//! qwen3_5.zig). +//! The Qwen3.5 hybrids on CUDA (2B, Qwen3.8-27B): the engine and its native family entry; Metal's 2B is qwen3_5.zig. pub const config = @import("config.zig"); pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 1d2a31d531..b4376c9490 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -1,6 +1,4 @@ -//! One stream's state on the CUDA engine: its recurrences and convolution tails, keys and values, a drafter's tap -//! ring, and the log of its newest verify window's rows, from which `keep` lands any path of that window (a prefix -//! of a chain, or a branch of a tree) without the forward having touched the stream's state. +//! A stream's state on CUDA: recurrences, conv tails, keys and values, the drafter's tap ring, the last window's log. const std = @import("std"); const cuda = @import("cuda"); const geometry = @import("cuda_geometry.zig"); @@ -72,9 +70,7 @@ pub const Cache = struct { return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; } - /// Keeps `path` (the window's rows, the first 0, ascending) of the newest window: the recurrences and tails as a - /// chain of those rows would have left them, a tree's kept keys moved down to the slots after the cache, the - /// length at the path's end. Queued on the engine's stream; nothing waits. + /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. pub fn keep(self: *Cache, e: *Engine, path: []const u32) !void { const last = self.last orelse return error.NothingToKeep; if (path.len == 0 or path.len > last.rows or path[0] != 0) return error.NothingToKeep; diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 9b5a9c1c2c..5df707de3c 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -1,7 +1,4 @@ -//! The DFlash2 drafter on CUDA behind the lane core's Drafter contract (core/lanes/drafter.zig): it absorbs the -//! target's tapped rows into a context ring a stream, then drafts a block of 7 tokens after the pending one, the -//! target's own embedding and head at its ends and a codebook selector scoring the top 16 a row under each parent: -//! one path (a chain), or the lanes of the round loop's tree (each lane its parent's rank-r choice). +//! The DFlash2 drafter on CUDA behind core/lanes/drafter.zig: a tap ring a stream, 7 drafts as a chain or a tree. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -119,8 +116,7 @@ const Fns = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, dyn /// qwen35_draft.cu's TapRows. const TapRows = extern struct { n: c_int, rows: [engine.most_window]c_int }; -/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds (a -/// chain, or a tree's lanes in the shape's order). +/// A stream's context: its layers' key and value rings over [begin, end), and the drafts it holds (a chain or tree). const Ctx = struct { ring: DeviceBuffer, begin: u64 = 0, @@ -169,8 +165,7 @@ pub const Draft = struct { ids: u64, // u32 the block's ids staged: cuda.HostBuffer, // pinned: the ids up, the words and projected rows down streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, - /// Contexts between streams: allocated and run through once (`reserve`), since the driver stalls a buffer's - /// first kernel for seconds. + /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. pool: std.ArrayList(*Ctx) = .empty, /// One level of a hold (a hold drafts every level at once; timed at load, over the block's levels). step_ms: f64 = 1.0, @@ -282,8 +277,7 @@ pub const Draft = struct { d.gpa.destroy(ctx); } - /// A new context, its ring touched by the kernels once: a row absorbed from the first tap ring at hand, then - /// a block drafted over it, so no request pays the driver's stall. + /// A new context, its ring touched once (an absorb, then a draft), so no request pays the driver's stall. fn create(d: *Draft) !*Ctx { const ctx = try d.gpa.create(Ctx); errdefer d.gpa.destroy(ctx); @@ -331,9 +325,7 @@ pub const Draft = struct { try d.e.run(d.fns.norm_rope, .{ rows * heads, 1, 1 }, .{ 32, 1, 1 }, .{ x, w, out, @as(c_int, @intCast(heads)), @as(c_int, @intCast(heads * d.c.head_dim)), @as(c_longlong, @intCast(pos0)), d.c.eps, d.c.theta }); } - /// Target rows at positions [start, start + n) from the stream's tap ring into its context (replacing what was - /// held at or past `start`); rows older than the window are skipped. With `rows`, the ring's rows start + rows[i] - /// (a tree's kept path, which the target keeps at the positions start + i) instead, n of them. + /// Target rows [start, start + n) from the tap ring into the context; with `rows`, ring rows start + rows[i]. fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32) !void { const c = d.c; var start = start_; @@ -370,8 +362,7 @@ pub const Draft = struct { if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; } - /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held as a chain, - /// or the `lanes` tree's tokens. + /// The block [pending, mask x 7] at pos0.. over the context: `depth` drafts as a chain, or the `lanes` tree. fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { const c = d.c; const e = d.e; @@ -423,8 +414,7 @@ pub const Draft = struct { } } - /// Row `r`'s candidate of `rank` after the token `pred`: by unary logit + predecessor . projected . successor, - /// ties to the earlier candidate (the row's top-k order). + /// Row `r`'s `rank`th candidate after `pred`: unary logit + pred . projected . successor, ties to the earlier one. fn choice(d: *const Draft, r: usize, pred: u32, rank: usize) u32 { const c = d.c; const w = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][r * lp_words * 4 ..][0 .. lp_words * 4]); diff --git a/zig/src/families/qwen3_5/cuda_draft_weights.zig b/zig/src/families/qwen3_5/cuda_draft_weights.zig index 2b0cf9f350..6d614f7982 100644 --- a/zig/src/families/qwen3_5/cuda_draft_weights.zig +++ b/zig/src/families/qwen3_5/cuda_draft_weights.zig @@ -1,6 +1,4 @@ -//! The DFlash2 drafter's checkpoint onto the device: its bf16 projections quantized to affine-4 (group 64) on the -//! host and packed into qlinear's tiles through the target engine's staging buffer, its norms and base kernels as -//! they are, its codebooks kept on the host for the selector. +//! The DFlash2 drafter's weights onto the device: projections quantized to affine-4 on the host, codebooks kept there. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -82,8 +80,7 @@ fn toBf16(x: f32) u16 { return @intCast((u +% rounding) >> 16); } -/// MLX affine-4, group 64, over bf16 rows: a group's scale is bf16((max - min) / 15) and bias bf16(min); each weight's -/// nibble is round((w - bias) / scale) clamped to [0, 15], eight to a word, the first in the low bits. +/// MLX affine-4, group 64: scale bf16((max - min) / 15), bias bf16(min), nibbles round((w - bias) / scale), low first. fn quantize(gpa: std.mem.Allocator, w: []align(1) const u16, n: usize, k: usize) !struct { words: []u32, scales: []u16, biases: []u16 } { const groups = k / 64; const words = try gpa.alloc(u32, n * k / 8); diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index be59726fc6..c9cf79d096 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,9 +1,4 @@ -//! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the -//! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails -//! and KV cache (cuda_cache.zig); the widths come from config.json (cuda_geometry.zig), the weights from -//! cuda_weights.zig. A forward's launches are captured into a CUDA graph the first time its shape is seen and -//! replayed after. A verify window's rows may form a tree (each row's parent given): the forward computes every -//! branch, logs each row, and the stream's state stays as it was until a path is kept. +//! Qwen3.5's hybrids on CUDA: qlinear projections, qwen35.cu glue, one captured forward over every stream's window. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -73,9 +68,7 @@ pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: c /// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; -/// What a forward keeps. A prompt chunk (`last`, one window) puts its last row's logits in its cache's `logits`; -/// verify windows keep every row, drawn on the device into `tokens` (`draws`, one a row) or their logits for -/// `rowLogits` (`logits`), and log each row for `Cache.keep`. +/// What a forward keeps: a prompt chunk its last row's logits, a verify window every row's draw or logits and log. pub const Keep = union(enum) { last, draws: []const Draw, logits }; /// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. @@ -92,8 +85,7 @@ pub const Engine = struct { affine: cuda.qlinear.Affine4Kernels, fns: Fns, keeps: Keeps, - /// Every weight and scratch buffer in one allocation: the WDDM driver's launch cost grows with a process's - /// allocation count, and a 64-layer model would have thousands. + /// Every weight and scratch buffer in one allocation: WDDM's launch cost grows with the allocation count. arena: DeviceBuffer, arena_at: usize = 0, stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there @@ -135,8 +127,7 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, - // the per-forward tables: ids, segments, each row's segment, parent and ancestors, the draws' rules, staged in - // pinned memory + // per-forward tables (ids, segments, rows' segment, parent and ancestors, the draws' rules) staged in pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, host_logits: cuda.HostBuffer, @@ -213,8 +204,7 @@ pub const Engine = struct { return e; } - /// Every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds (a 16 GiB - /// model: a 20 s launch, then a 3 s one), which no request should pay. + /// Every path a request takes, run once now: WDDM stalls a process's first forwards for seconds (a 20 s launch). fn warm(e: *Engine, gpa: std.mem.Allocator) !void { var warm_cache = try Cache.init(gpa, e, 64); defer warm_cache.deinit(gpa); @@ -312,8 +302,7 @@ pub const Engine = struct { try cuda.launch.launch(f, cfg, e.stream, &a); } - /// out (rows, n) = x (rows, k) through `l`: decode tiles 16 rows at a time (each row's bits whatever the batch), - /// or the prompt GEMM for a prompt chunk wider than one tile. + /// out (rows, n) = x (rows, k) through `l`: 16-row decode tiles (a row's bits batch-independent) or a prompt GEMM. pub fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)), e.tables.ptr + params_at }); @@ -328,8 +317,7 @@ pub const Engine = struct { try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } - /// The row tables of `parts`: each forward row's segment, parent forward row (-1: its window's first) and the - /// ancestors-and-self bits over its window's rows (a chain without `parents`). + /// The row tables of `parts`: each row's segment, parent row (-1: a window's first) and ancestors-and-self bits. fn rowTables(st: []u8, parts: []const Part) !void { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); const rowpar = std.mem.bytesAsSlice(c_int, st[rowpar_at..][0 .. max_rows * 4]); @@ -348,8 +336,7 @@ pub const Engine = struct { } } - /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. A - /// window not yet kept from the stream's last verify is kept whole first. + /// One forward over `parts`' windows, `ids` in order; `keep` names the kept rows; an unkept verify is kept first. pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { const g = e.g; var total: usize = 0; @@ -397,15 +384,13 @@ pub const Engine = struct { params[3] = @intCast(parts.len); params[4] = @intFromBool(record); try e.tables.uploadAsync(0, st, e.stream.handle); - // a captured forward's launches cover a bucket of rows and every window; its kernels read the true counts - // (a lone stream's one-row step, the common case without a drafter, keeps its own exact graph) + // a captured forward covers a row bucket and every window; kernels read the true counts; a lone step is exact const lone = total == 1 and parts.len == 1; const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; const windows: usize = if (lone) 1 else max_segs; const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; if (prompt) { - // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length would be a graph of its own, - // paying the driver's stall on a new graph's first run (about a second on a wide model) + // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length's graph would pay the stall try e.body(shape); } else { const gop = try e.graphs.getOrPut(e.gpa, shape.key()); diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index 7b327fd1a9..fdd6ecb0c2 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -1,5 +1,4 @@ -//! The widths of a Qwen3.5 hybrid text model (the 2B, Qwen3.8-27B) as the CUDA engine reads them from config.json: -//! what varies is read, what qwen35.cu fixes (head widths, taps, rotary layout, 4-bit groups of 64) is checked. +//! A Qwen3.5 hybrid's widths (2B, 27B) as the CUDA engine reads them from config.json; what qwen35.cu fixes is checked. const std = @import("std"); pub const max_layers = 128; @@ -70,8 +69,7 @@ pub const Geometry = struct { return g.deltas() * (g.recBytes() + g.convBytes()); } - /// One verified row's record of one linear layer (qwen35.cuh's log_bytes): decay and write by value head, - /// key by key head, the tail after it. + /// A row's record of a linear layer (qwen35.cuh log_bytes): decay and write a value head, key a key head, the tail. pub fn logRowBytes(g: Geometry) usize { return g.v_heads * 4 + g.v_heads * dk * 4 + g.k_heads * dk * 2 + g.convBytes(); } diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index d4fa634cb1..161b34a882 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -1,6 +1,4 @@ -//! The Qwen3.5 hybrids' CUDA lane backend (the 2B, Qwen3.8-27B) behind the native family interface (native/cuda.zig's -//! registry): windows of up to 16 rows (the lane core's copies), every stream's window in one forward, prompts in -//! chunks, draws on the device. +//! The Qwen3.5 hybrids' CUDA lane backend (native/cuda.zig's registry): 16-row windows, every stream in one forward. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -30,8 +28,7 @@ pub const Loaded = struct { free_streams: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, - /// Called with the admitted stream count: that many caches allocated and touched now, since the driver commits - /// memory on first touch (seconds for a wide model's keys), and a request should not pay it. + /// Called with the admitted stream count: that many caches allocated and touched now, as memory commits on touch. reserve: *const fn (*anyopaque, u32) anyerror!void, }; @@ -66,8 +63,7 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try cache.reset(l.e); l.pool.appendAssumeCapacity(cache); } - // the shared rounds' shapes, a row of each of 1 to `streams` streams, captured now: a new graph's first run - // pays the driver's stall too (about a second on a wide model) + // the shared rounds' shapes, a row of 1 to `streams` streams, captured now: a new graph's first run stalls too var parts: [E.max_segs]E.Part = undefined; var ids: [E.max_segs]u32 = @splat(0); var picks: [E.max_segs]E.Draw = undefined; @@ -85,8 +81,7 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { }; } -/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). The -/// target's windows of 1 to window_rows rows are timed here (with the drafter's taps on) for the round loop's prices. +/// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); @@ -169,8 +164,7 @@ const Lanes = struct { self.pool.append(self.gpa, cache) catch self.drop(cache); } - /// Each window width's forward, timed on a cache of its own (its third run: the shape's graph captured and the - /// driver's stall behind it, the previous window's keep in front as in a round), into the drafter's facts. + /// Each width's forward timed on its own cache, on its third run: graph captured, stall paid, a keep in front. fn timeWindows(self: *Lanes, d: *Draft) !void { var cache = try E.Cache.init(self.gpa, self.e, 64); defer cache.deinit(self.gpa); @@ -249,8 +243,7 @@ const Lanes = struct { return @intCast(handle); } - /// Each window (a stream's pending token and drafts, a chain or a tree) checked, then as many as fit in each - /// forward. A tree's rows are keyed at the cache length plus their depth, as the round loop positions them. + /// Each window (a pending token and drafts) checked, then as many as fit a forward; a tree's rows keyed by depth. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { const self = cast(ptr); const t0 = std.Io.Clock.awake.now(self.io); @@ -326,9 +319,7 @@ const Lanes = struct { try self.e.setTaps(taps); } - /// Rows [start, start + count) of a stream's tap ring: `offset` carries the first row's position, the drafter - /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. The newest window's - /// rows stay readable after a keep (a tree's kept path reads them by row) until the next forward. + /// Rows [start, start + count) of a stream's tap ring at slot position % tap_rows; a request may span its end. fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { const self = cast(ptr); if (taps.len != self.e.n_taps) return error.TapsMismatch; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig index a105b876f4..2a2b75c361 100644 --- a/zig/src/families/qwen3_5/cuda_weights.zig +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -1,5 +1,4 @@ -//! The checkpoint's weights onto the device: each affine-4 projection staged as stored and packed into qlinear's -//! tiles, the vectors as they are, all from the engine's one arena (cuda_engine.zig). +//! The checkpoint's weights onto the device: affine-4 projections packed as qlinear tiles, out of the engine's arena. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig index 35bb853608..204eeb3d4a 100644 --- a/zig/src/platform/sys.zig +++ b/zig/src/platform/sys.zig @@ -1,5 +1,4 @@ -//! The operating system as TensorFold calls it: `c` is libc's surface (std.c on POSIX, win_c.zig on Windows) and -//! `DynLib` opens shared libraries (std.DynLib on POSIX, LoadLibraryW on Windows). Every module imports this as `sys`. +//! The OS as TensorFold calls it, imported as `sys`: `c` is std.c or win_c.zig, `DynLib` is std.DynLib or LoadLibraryW. const std = @import("std"); const builtin = @import("builtin"); diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index eb4d0c746a..21e057d6f5 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -1,6 +1,4 @@ -//! The libc surface TensorFold calls through `sys.c`, on Windows: integer descriptors over the UCRT (files opened by -//! CreateFileW with every share mode, so unlink and rename behave as on POSIX), directory descriptors as remembered -//! paths, readdir over FindFirstFileW, and POSIX's names for the rest. Each call keeps libc's contract: -1 and errno. +//! libc's surface on Windows (`sys.c`): descriptors over the UCRT, CreateFileW with every share mode, -1 and errno. const std = @import("std"); @@ -165,8 +163,7 @@ fn lastError() E { const Wide = [32767:0]u16; -/// UTF-8 path to a NUL-terminated wide path in `buf`; null (errno set) when it does not fit or is not UTF-8. -/// POSIX's /tmp is the user's temporary folder and /proc/self/exe this executable, as Cygwin and MSYS map them. +/// UTF-8 path as a wide path in `buf`, null with errno if it does not fit; /tmp and /proc/self/exe map as in Cygwin. fn wide(buf: *Wide, path: []const u8) ?[*:0]const u16 { var at: usize = 0; var rest = path; diff --git a/zig/src/platform/win_clock.zig b/zig/src/platform/win_clock.zig index ec41325cfc..0c0c67ab4e 100644 --- a/zig/src/platform/win_clock.zig +++ b/zig/src/platform/win_clock.zig @@ -1,5 +1,4 @@ -//! The process, environment and clock part of the libc surface on Windows (win_c.zig re-exports these): the process -//! id, environment writes over the UCRT, and nanosleep and clock_gettime over the kernel's timers and counters. +//! Process, environment and clock calls of the Windows libc surface, re-exported by win_c.zig: UCRT and kernel timers. const std = @import("std"); const c = @import("win_c.zig"); diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index 04c82d4dba..c0ab151e25 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -1,6 +1,4 @@ -//! The std.posix surface the HTTP server calls through `sys.posix`, on Windows: Winsock sockets (SOCKET as an isize, -//! INVALID_SOCKET as -1), WSAPoll, console control events as SIGINT and SIGTERM, and no descriptor limits to raise. -//! Calls in `system` return -1 with errno set from WSAGetLastError, as libc's do. +//! std.posix as the server calls it, on Windows: Winsock sockets, WSAPoll, console events as SIGINT and SIGTERM. const std = @import("std"); const c = @import("win_c.zig"); @@ -29,8 +27,7 @@ pub const SOL = struct { pub const SOCKET: i32 = 0xffff; }; pub const SO = struct { - /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, which the listener - /// sets REUSEADDR to prevent; exclusive use is the POSIX meaning. + /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, the opposite of POSIX's. pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); }; pub const SHUT = struct { diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 858dc34c07..0d1f711252 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,8 +1,4 @@ -//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over -//! comma-separated token ids: one token a forward (step: every position's logits), windows of growing widths -//! (window: the same rows, which must equal step's), tree windows with junk branches kept to the true one (tree: -//! the same), or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into for a -//! reference to compare. Step mode also prints the decode speed. +//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -51,8 +47,7 @@ pub fn main(init: std.process.Init) !u8 { at += n; } } else if (std.mem.eql(u8, args[2], "tree")) { - // tree windows of growing widths: the true continuation down one branch, a junk sibling before each of its - // rows and a junk grandchild, the true branch kept; its rows' logits must equal step's bit for bit + // tree windows of growing widths, junk branches beside the kept true one, whose logits must equal step's bits var at: usize = 0; var w: usize = 2; while (at < ids.items.len) : (w = w % e.window_rows + 1) { From 8107ee2b90c92f3d038c1db8e22caa0f20f0ca23 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:37 +0000 Subject: [PATCH 13/46] Windows: checkpoint reads skip the page cache direct_io opened files with O_DIRECT on Linux alone. On Windows the same open now asks for FILE_FLAG_NO_BUFFERING (which win_c maps O_DIRECT to), and the 4 KiB-aligned spans it already reads satisfy the sector alignment that needs; a file system that refuses it still falls back to buffered reads. The direct-read test runs on Windows too, checking the executable's bytes against buffered reads of it. --- docs/WINDOWS.md | 5 +++-- zig/src/core/direct_io.zig | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index b11d72f4fd..db0a6040e1 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -65,8 +65,9 @@ tensorfold-native serve --name qwen27b --context 8192 --p - `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating system as the engine calls it. On Linux and macOS it is `std.c`, `std.posix` and `std.DynLib` unchanged; on Windows `win_c.zig` and `win_posix.zig` give the same names over the UCRT, Win32 and Winsock. -- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. `O_DIRECT` maps to - unbuffered reads. `/tmp` is the user's temporary folder and `HOME` falls back to `USERPROFILE`. +- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. Checkpoint reads + skip the page cache (`FILE_FLAG_NO_BUFFERING`, as `O_DIRECT` on Linux). `/tmp` is the user's temporary folder and + `HOME` falls back to `USERPROFILE`. - The Hugging Face cache links snapshot files to blobs. Windows allows symlinks only with Developer Mode or elevation, so `pull` falls back to hard links (no extra space). - Checkpoint reads use threads (io_uring is Linux's). diff --git a/zig/src/core/direct_io.zig b/zig/src/core/direct_io.zig index 2f4cc546f6..f5bba757e2 100644 --- a/zig/src/core/direct_io.zig +++ b/zig/src/core/direct_io.zig @@ -26,7 +26,7 @@ pub const File = struct { /// `path` read-only, direct when its file system allows it, else through the page cache. pub fn open(path: [:0]const u8) !File { - if (builtin.os.tag == .linux) { + if (builtin.os.tag == .linux or builtin.os.tag == .windows) { const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); if (fd >= 0) return .{ .fd = fd, .direct = true }; } @@ -85,7 +85,7 @@ test "span and fits" { } test "direct reads return the same bytes as buffered reads" { - if (builtin.os.tag != .linux) return error.SkipZigTest; + if (builtin.os.tag != .linux and builtin.os.tag != .windows) return error.SkipZigTest; const gpa = std.testing.allocator; var f = try File.open("/proc/self/exe"); defer f.close(); From 693533e9d4d80ac3642bc5b89bebcbda2d70c9f5 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:37 +0000 Subject: [PATCH 14/46] Windows: the signal handler type allows address 1, for aarch64 SIG.IGN is a function pointer made from 1, as std.posix has it. aarch64 requires function pointers to be aligned unless the type says align(1), so the cross-build for aarch64-windows refused it. With align(1), as std's own handler type has, tensorfold-native cross-compiles for aarch64-windows with the GB10 (sm_121) fatbins; no hardware has run it. --- zig/src/platform/win_posix.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index c0ab151e25..c154eb7153 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -203,7 +203,7 @@ pub const SIG = enum(c_int) { pub const IGN: ?Handler = @ptrFromInt(1); pub const DFL: ?Handler = null; }; -pub const Handler = *const fn (SIG) callconv(.c) void; +pub const Handler = *align(1) const fn (SIG) callconv(.c) void; // align(1): IGN is address 1, as std.posix has it pub const sigset_t = u32; pub const Sigaction = struct { handler: extern union { handler: ?Handler, sigaction: ?*const anyopaque }, From f53ae20aba85a7043e90dd5dce23511c48d69f90 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:38 +0000 Subject: [PATCH 15/46] docs: Windows as a target in the README and the Windows page The Windows page opens with what the build is, what it is tested on and that it cross-compiles for aarch64-windows; the README's model table and the CUDA paragraph say Windows where they said unofficial port. --- README.md | 6 +++--- docs/WINDOWS.md | 5 +++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index c838faa144..48c7b8d27f 100644 --- a/README.md +++ b/README.md @@ -73,8 +73,8 @@ The platform column names the hardware tested for each model. | Qwen3.8 Flash Next | MLX affine 6-bit, group 32 | Metal on M5 Ultra | | GLM-5.3-Flash | MLX affine 4-bit, group 64 | Metal on two M5 Ultras | | GLM-5.3 | EXLR8 `home` profile (trellis K2/K3/K4 experts) | Metal on two M5 Ultras | -| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (unofficial Windows port) | -| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (unofficial Windows port) | +| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (Windows) | +| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (Windows) | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. NVIDIA's NVFP4 checkpoint, `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`, serves on CUDA as stored; its bf16 MTP head is quantized to 4-bit at load for drafts. @@ -90,7 +90,7 @@ Its prompt processing trails mlx_lm for now, and we are fixing it. An M3 Ultra r Bonsai, Gemma 4, Qwen3.6 and DeepSeek-V4 are still under qualification for 1.0.x. The Python 0.6.6 engine remains available at the `v0.6.6` tag for those models and other backends. On CUDA, 1.0.11 serves Nemotron on a GB10, greedy and sampled, with concurrent requests sharing each round. -The unofficial Windows port ([docs/WINDOWS.md](docs/WINDOWS.md)) also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. +On Windows ([docs/WINDOWS.md](docs/WINDOWS.md)), the native build also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. The Linux x86_64 archive adds NVIDIA Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), tested on an RTX 3090, and Blackwell cards with compute capability 12.0 (the RTX 50 series and RTX PRO 6000 Blackwell), tested on an RTX PRO 6000; the A100 (8.0) is not supported yet. ## Exact decoding diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index db0a6040e1..52b17191c6 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -1,7 +1,8 @@ # TensorFold on Windows (native) -An unofficial port: `tensorfold-native.exe` runs on Windows x86_64 with an NVIDIA GPU through the same CUDA engine -the Linux build uses. Only the CUDA backend is ported; Metal is macOS's and HIP is not built. +`tensorfold-native.exe` runs on Windows with an NVIDIA GPU through the same CUDA engine the Linux build uses. Only the +CUDA backend is built on Windows; Metal is macOS's and HIP is not built. It is tested on an RTX 4090 (x86_64); the +host code also cross-compiles for `aarch64-windows` with the GB10 (sm_121) kernels, untested on hardware so far. ## What you need From e875ee2ac3e2fd3e742f67074cc8441de915afad Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:38 +0000 Subject: [PATCH 16/46] CI: the host unit tests on Windows The no-GPU workflow runs zig build test on windows-latest beside ubuntu-latest, so the platform layer is built and tested on every pull request. --- .github/workflows/no-gpu.yml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 6acdde321e..09dcde39d9 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -38,3 +38,13 @@ jobs: (cd ../lean-base && python3 tools/zig/lean_check.py > ../lean-base.txt || true) python3 tools/zig/lean_check.py > ../lean-head.txt || true python3 tools/zig/lean_new.py ../lean-base.txt ../lean-head.txt + zig-windows: + runs-on: windows-latest + steps: + - uses: actions/checkout@v4 + - uses: mlugg/setup-zig@v2 + with: + version: 0.17.0 + # The same host unit tests on Windows: the platform layer (zig/src/platform) over the UCRT, Win32 and Winsock. + - name: zig build test + run: zig build test From 38ecce53f7d3f678c29c313a8ee88100adf19708 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:21:56 +0000 Subject: [PATCH 17/46] qwen3_5: the drafter reads TF_TIMING through the platform layer Every other module calls libc through `sys`; the drafter's four getenv calls were the last direct std.c calls in a file the Windows build compiles. --- zig/src/families/qwen3_5/cuda_draft.zig | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 5df707de3c..cfb8670cf5 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -1,5 +1,6 @@ //! The DFlash2 drafter on CUDA behind core/lanes/drafter.zig: a tap ring a stream, 7 drafts as a chain or a tree. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); @@ -219,7 +220,7 @@ pub const Draft = struct { /// The target's window of `rows` rows takes `ms` (told by cuda_native.zig before the facts are read). pub fn setCost(d: *Draft, rows: u32, ms: f64) void { - if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); if (d.n_costs < d.costs.len) { d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; d.n_costs += 1; @@ -406,11 +407,11 @@ pub const Draft = struct { try e.stream.synchronize(); if (lanes_) |l| { try d.selectLanes(ctx, pending, rows, l); - if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); } else { d.select(ctx, pending, rows); ctx.held_n = @min(depth, rows); - if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); } } @@ -479,7 +480,7 @@ pub const Draft = struct { fn factsFn(p: *anyopaque) dr.Facts { const d = cast(p); // TF_NO_TREES=1: chains only (for comparing) - return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = std.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = platform.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; } fn openFn(p: *anyopaque, s: *Stream) anyerror!void { @@ -539,7 +540,7 @@ fn bf16(h: u16) f32 { /// TF_TIMING=1: each call's wall time on stderr. pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { - if (std.c.getenv("TF_TIMING") == null) return; + if (platform.c.getenv("TF_TIMING") == null) return; const now = std.Io.Clock.awake.now(io); const ns: f64 = @floatFromInt(t0.durationTo(now).toNanoseconds()); const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); From 181782342dbfbd4a508eec3ce77a39e2c36136e3 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:33:03 +0000 Subject: [PATCH 18/46] qwen3_5: the drafter takes a prompt longer than its tap ring After a prompt the lane core hands the drafter every prompt row's features; the CUDA backend refused any request wider than its tap ring (2,176 rows), so a 27B request with --drafter and a prompt past that ended with RowsNotHeld, where the same prompt without the drafter ran. The ring keeps the newest tap_rows rows, which hold the drafter's whole context window, and absorb already reads only that window of what it is handed; the backend now refuses only rows not yet written or all gone. On an RTX 4090 the 27B with its drafter takes a 6,041-token prompt in 7.4 s, and its reply equals the one with "draft": false. --- zig/src/families/qwen3_5/cuda_native.zig | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 161b34a882..8773dffc71 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -319,13 +319,15 @@ const Lanes = struct { try self.e.setTaps(taps); } - /// Rows [start, start + count) of a stream's tap ring at slot position % tap_rows; a request may span its end. + /// Rows [start, start + count) of a stream's tap ring (slot position % tap_rows), the newest tap_rows rows kept. fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { const self = cast(ptr); if (taps.len != self.e.n_taps) return error.TapsMismatch; const cache = try self.cacheOf(stream); const ring = cache.tap orelse return error.NoTaps; - if (start + count > cache.written or count > E.tap_rows) return error.RowsNotHeld; + // a long prompt asks for more rows than the ring keeps; the drafter reads only its window of the newest + const held_from = cache.written -| E.tap_rows; + if (start + count > cache.written or start + count <= held_from) return error.RowsNotHeld; return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; } From b811a3b3f316c223fae6288b47418962b4b09f42 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 16:16:38 +0000 Subject: [PATCH 19/46] qwen3_5 cuda: prompt chunks past 32 rows see the right keys, and tree attention cuts chunks at logical positions - rowanc held a window's rows as bits of a u32, but prompt chunks run up to 128 rows: rows from 32 on rotated at the wrong positions and attended the wrong keys. A chain row (no parents) now carries its index in its window under the top bit; tree windows (at most 16 rows) keep their bits. - tf_q35_attn_partial cut its 128-key chunks at a key's cache row, so a kept tree path whose window crossed a chunk boundary summed in other chunks than one-row steps and changed bits. Chunks are now cut at the key's logical position (the stream's length plus the ancestors before it). Checked on an RTX 5090 (Linux, sm_120), tf-qwen35-cuda over 300 ids on Qwen3.5-2B and Qwen3.8-27B: step bits unchanged, window and tree equal step byte for byte (tree differed from row 126 on the 2B and 254 on the 27B before). Prompt mode's last row against step's at 16 to 300 ids: KL at most 0.003 and the same argmax in 20 of 20 (before: 27B at 40 ids KL 2.3 with a different argmax, at 128 ids KL 1.35). --- zig/kernels/cuda/qwen35.cu | 19 +++++++++++++++---- zig/src/families/qwen3_5/cuda_engine.zig | 5 ++++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index cef7c7b5a2..558baa7908 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -307,6 +307,16 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); } +// A row's ancestors-and-self: its tree window's row bits, or with the top bit set a chain row's index in its window. +__device__ __forceinline__ int rows_seen(unsigned anc) { return (anc >> 31) ? (int)(anc & 0x7fffffffu) + 1 : __popc(anc); } + +// The window row holding a row's nth seen window key. +__device__ __forceinline__ int seen_row(unsigned anc, int n) { + if (anc >> 31) return n; + for (int j = 0; j < n; j++) anc &= anc - 1; + return __ffs(anc) - 1; +} + // Queries rotated in place, keys rotated and values copied into the row's cache slot; grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, @@ -317,7 +327,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; - const int pos = sg.len + __popc(rowanc[m]) - 1; + const int pos = sg.len + rows_seen(rowanc[m]) - 1; const int slot = sg.len + (int)m - sg.row0; __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; @@ -347,7 +357,8 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; const unsigned anc = rowanc[m]; - const int k0 = c * ck, k1 = min(k0 + ck, sg.len + sg.rows); + // chunks cut at logical positions (a window key's: the stream's length plus its ancestors), as a step's are + const int k0 = c * ck, k1 = min(k0 + ck, sg.len + rows_seen(anc)); const size_t slot = (m * QH + qh) * nch + c; if (k0 >= k1) { if (lane == 0) { @@ -365,8 +376,8 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (int pos = k0; pos < k1; pos++) { - if (pos >= sg.len && ((anc >> (pos - sg.len)) & 1u) == 0) continue; + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); const __nv_bfloat16* kb = reinterpret_cast(&kw); float dot = 0.0f; diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index c9cf79d096..389b32f766 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -317,6 +317,9 @@ pub const Engine = struct { try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } + /// A chain row's rowanc: its index in its window, as a prompt chunk's rows outnumber the bits. + const chain_row: u32 = 1 << 31; + /// The row tables of `parts`: each row's segment, parent row (-1: a window's first) and ancestors-and-self bits. fn rowTables(st: []u8, parts: []const Part) !void { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); @@ -330,7 +333,7 @@ pub const Engine = struct { if (parent >= @as(i32, @intCast(r)) or parent < -1) return error.InvalidQwenParents; rowseg[row0 + r] = @intCast(i); rowpar[row0 + r] = if (parent < 0) -1 else @intCast(row0 + @as(usize, @intCast(parent))); - rowanc[row0 + r] = (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); + rowanc[row0 + r] = if (p.parents == null) chain_row | @as(u32, @intCast(r)) else (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); } row0 += p.rows; } From 0a0856ba97fd704b55ea39975821d76c0edcf8b2 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:47:32 +0000 Subject: [PATCH 20/46] Windows: the server goldens and the Cortheon fuzz step on Windows hosts zig build test-golden and zig build fuzz-cortheon exist on every non-macOS host since the golden step moved into zig/build/golden.zig, so the golden check compiles against the platform layer on Windows: - the platform's posix gains sockaddr.in, timeval, connect and kill for a child's handle (TerminateProcess; HUP has no equivalent), setsockopt turns a timeval for SO_RCVTIMEO or SO_SNDTIMEO into Winsock's milliseconds, and a socket read that times out is WouldBlock, as std.posix reports it. - the check carries SystemRoot into the spawned server's environment, which Winsock needs to start, and drains the server's stdout after the port line on every host: the fake server logs each request there, and a full pipe (4 KiB on Windows) blocked it after sixteen cases. - on Windows the check cannot send SIGTERM or SIGHUP, so the sigterm-exit-0 outcome and the cases after a key rewrite are skipped and printed as such. On Windows the goldens give 306 of 315 equal with the same nine known differences as Linux, and the fuzz step passes its eight tests; the no-GPU workflow's Windows job runs both. --- .github/workflows/no-gpu.yml | 5 ++++ zig/src/platform/win_posix.zig | 39 +++++++++++++++++++++++++++---- zig/tests/server/golden_check.zig | 33 +++++++++++++++++++++++++- 3 files changed, 72 insertions(+), 5 deletions(-) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 09dcde39d9..e74bb9e833 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -48,3 +48,8 @@ jobs: # The same host unit tests on Windows: the platform layer (zig/src/platform) over the UCRT, Win32 and Winsock. - name: zig build test run: zig build test + # The server goldens and the Cortheon fuzz harness on Windows too; the goldens skip the signal-driven cases. + - name: zig build test-golden + run: zig build test-golden + - name: zig build fuzz-cortheon + run: zig build fuzz-cortheon diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index c154eb7153..4088f3214b 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -13,7 +13,13 @@ pub const S = c.S; pub const socket_t = isize; pub const socklen_t = c_int; -pub const sockaddr = extern struct { family: u16, data: [14]u8 }; +pub const sockaddr = extern struct { + family: u16, + data: [14]u8, + + pub const in = extern struct { family: u16, port: u16, addr: u32, zero: [8]u8 = @splat(0) }; +}; +pub const timeval = c.timeval; pub const AF = struct { pub const UNIX: u32 = 1; @@ -29,6 +35,8 @@ pub const SOL = struct { pub const SO = struct { /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, the opposite of POSIX's. pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); + pub const RCVTIMEO: u32 = 0x1006; + pub const SNDTIMEO: u32 = 0x1005; }; pub const SHUT = struct { pub const RD: i32 = 0; @@ -124,6 +132,9 @@ pub const system = struct { pub fn bind(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { return if (ws.bind(fd, addr, @intCast(len)) != 0) failed() else 0; } + pub fn connect(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { + return if (ws.connect(fd, addr, @intCast(len)) != 0) failed() else 0; + } pub fn listen(fd: socket_t, backlog: u31) c_int { return if (ws.listen(fd, backlog) != 0) failed() else 0; } @@ -155,6 +166,13 @@ pub const system = struct { }; pub fn setsockopt(fd: socket_t, level: i32, opt: u32, value: []const u8) error{SetSockOptFailed}!void { + if (opt == SO.RCVTIMEO or opt == SO.SNDTIMEO) { + // POSIX gives a timeval; Winsock takes the milliseconds as a DWORD + const tv: *align(1) const timeval = @ptrCast(value.ptr); + const ms: u32 = @intCast(@as(i64, tv.sec) * 1000 + @divTrunc(@as(i64, tv.usec), 1000)); + if (ws.setsockopt(fd, level, @bitCast(opt), @ptrCast(&ms), @sizeOf(u32)) != 0) return error.SetSockOptFailed; + return; + } if (ws.setsockopt(fd, level, @bitCast(opt), value.ptr, @intCast(value.len)) != 0) return error.SetSockOptFailed; } @@ -184,13 +202,26 @@ pub fn poll(fds: []pollfd, timeout_ms: i32) error{PollFailed}!usize { extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; -/// Receives from a socket; 0 is the peer's orderly close. -pub fn read(fd: socket_t, buf: []u8) error{ConnectionResetByPeer}!usize { +/// Receives from a socket; 0 is the peer's orderly close, and a receive timeout (SO.RCVTIMEO) is WouldBlock. +pub fn read(fd: socket_t, buf: []u8) error{ ConnectionResetByPeer, WouldBlock }!usize { const n = ws.recv(fd, buf.ptr, len32(buf.len), 0); - if (n < 0) return error.ConnectionResetByPeer; + if (n < 0) return switch (ws.WSAGetLastError()) { + 10035, 10060 => error.WouldBlock, + else => error.ConnectionResetByPeer, + }; return @intCast(n); } +// ---- processes ---- + +extern "kernel32" fn TerminateProcess(process: *anyopaque, code: c_uint) callconv(.winapi) c_int; + +/// kill for a spawned child's handle: TERM ends it with exit code 1; HUP and the rest have no Windows equivalent. +pub fn kill(process: std.os.windows.HANDLE, sig: SIG) error{ Unsupported, Unexpected }!void { + if (sig != .TERM) return error.Unsupported; + if (TerminateProcess(process, 1) == 0) return error.Unexpected; +} + // ---- signals: console control events ---- pub const SIG = enum(c_int) { diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index d7f916f4a6..4976658e6b 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -1,7 +1,8 @@ //! Frozen server goldens for `zig build test-golden`, with the ten differences that are still known. const std = @import("std"); +const builtin = @import("builtin"); const platform = @import("sys"); -const posix = std.posix; +const posix = platform.posix; const Io = std.Io; const wire = @import("wire.zig"); const Reply = wire.Reply; @@ -22,6 +23,18 @@ const known_diffs = [_]Known{ .{ .group = "anthropic", .name = "tool-use-stream" }, }; +/// Reads a spawned server's stdout to its end, so its log lines never fill the pipe (4 KiB on Windows). +fn drainStdout(io: Io, out: Io.File) void { + var buf: [4096]u8 = undefined; + var reader = out.reader(io, &buf); + while (true) { + _ = reader.interface.discardShort(buf.len) catch return; + } +} + +/// The parent's SystemRoot, which a Windows child needs for Winsock; null elsewhere. +var system_root: ?[]const u8 = null; + fn isKnown(group: []const u8, name: []const u8) bool { for (known_diffs) |k| { if (std.mem.eql(u8, k.group, group) and std.mem.eql(u8, k.name, name)) return true; @@ -355,6 +368,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con } else try argv.append(a, arg); } var map = std.process.Environ.Map.init(a); + if (system_root) |root| try map.put("SystemRoot", root); // Winsock refuses to start without it try map.put("TENSORFOLD_NO_LIVE", "1"); try map.put("TF_FAKE_CONTEXT", context_tokens); var it = env.iterator(); @@ -398,12 +412,16 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con const line = reader.interface.takeDelimiterExclusive('\n') catch return error.NoPortLine; if (std.mem.startsWith(u8, line, "PORT ")) { const port = try std.fmt.parseInt(u16, std.mem.trim(u8, line[5..], " \r"), 10); + // the server's log lines keep coming on stdout: drained, or a full pipe blocks the server + const drain = try std.Thread.spawn(.{}, drainStdout, .{ io, child.stdout.? }); + drain.detach(); return .{ .child = child, .port = port, .log_path = log_path, .key_path = key_path }; } } } pub fn main(init: std.process.Init) !void { + if (builtin.os.tag == .windows) system_root = init.environ_map.get("SystemRoot"); const gpa = init.gpa; const io = init.io; const args = try init.minimal.args.toSlice(init.arena.allocator()); @@ -436,15 +454,26 @@ pub fn main(init: std.process.Init) !void { if (obj.get("cases")) |cases| { for (cases.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, server.?.port, case); } + var after_signal = false; // Windows: a step's cases after a signal it cannot send are skipped, and said so for (obj.get("then").?.array.items) |then| { const t = then.object; const kind = t.get("kind").?.string; const s = &server.?; if (std.mem.eql(u8, kind, "cases")) { + if (after_signal) { + std.debug.print("SKIPPED {d} cases after a signal Windows cannot send\n", .{t.get("cases").?.array.items.len}); + continue; + } for (t.get("cases").?.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, s.port, case); } else if (std.mem.eql(u8, kind, "sleep")) { std.Io.sleep(io, .fromMilliseconds(@intFromFloat(t.get("seconds").?.float * 1000.0)), .awake) catch {}; } else if (std.mem.eql(u8, kind, "sigterm")) { + if (builtin.os.tag == .windows) { + std.debug.print("SKIPPED lifecycle/sigterm-exit-0: Windows has no SIGTERM to send, the server is ended\n", .{}); + posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; + _ = s.child.wait(io) catch {}; + continue; + } posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; const term = s.child.wait(io) catch return error.WaitFailed; try outcomes.append(init.arena.allocator(), .{ @@ -471,8 +500,10 @@ pub fn main(init: std.process.Init) !void { try fw.interface.writeAll(t.get("content").?.string); try fw.interface.flush(); posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; // the server rereads the key file, as the reference does + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "sighup")) { posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "terminate")) { posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; _ = s.child.wait(io) catch {}; From ccd06b94acfe63c82e3497a4d6e762b245b6a3a8 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:56:21 +0000 Subject: [PATCH 21/46] cuda: cuBLASLt opens by its Windows DLL name The loader tried only libcublasLt.so.13 and libcublasLt.so, so the cublaslt suite of tf-cuda-test reported LibraryUnavailable on Windows; there it opens cublasLt64_13.dll from the toolkit's bin folder. --- zig/src/cuda/cublaslt.zig | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index 9e1f154172..cf5b167fc3 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -58,9 +58,10 @@ pub const Library = struct { lib: platform.DynLib, api: Api, - /// The CUDA 13 soname first, then the unversioned link name. + /// The CUDA 13 soname first, then the unversioned link name; on Windows the toolkit's DLL. pub fn open() Error!Library { - for ([_][]const u8{ "libcublasLt.so.13", "libcublasLt.so" }) |path| { + const names: []const []const u8 = if (@import("builtin").os.tag == .windows) &.{"cublasLt64_13.dll"} else &.{ "libcublasLt.so.13", "libcublasLt.so" }; + for (names) |path| { return openPath(path) catch |e| switch (e) { error.LibraryUnavailable => continue, else => return e, From 045abbd023a31553b006c8fb12c090acfbc102a8 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:20:52 +0000 Subject: [PATCH 22/46] Windows: writes on an O_APPEND descriptor land at the end whole The platform handed O_APPEND to the C runtime as _O_APPEND, whose write seeks to the end and then writes, so two threads appending at once could overwrite each other's bytes: the request log lost lines under concurrent requests, which the server goldens' request-log case caught on a slow runner. A descriptor opened with O_APPEND now writes through WriteFile at the end-of-file offset, which the kernel places whole, as POSIX does; the file keeps its write access, so ftruncate on such a descriptor still works (prompt_imprint repairs a torn index that way). A request_log test appends 200 lines from each of 8 threads and expects every line whole; it lost lines on Windows before this change. --- zig/src/platform/win_c.zig | 13 +++++++++++ zig/src/server/request_log.zig | 42 ++++++++++++++++++++++++++++++++++ zig/src/server/root.zig | 1 + 3 files changed, 56 insertions(+) diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index 21e057d6f5..bfb41928c5 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -284,6 +284,7 @@ fn openFile(path: [*:0]const u8, flags: O) c_int { _ = CloseHandle(@ptrFromInt(@as(usize, @bitCast(h)))); return fail(.MFILE); } + if (fd < append_fds.len) @atomicStore(u8, &append_fds[@intCast(fd)], @intFromBool(flags.APPEND), .release); return fd; } @@ -293,7 +294,11 @@ pub fn openat(dir: fd_t, path: [*:0]const u8, flags: O, ...) callconv(.c) c_int return openFile(p.ptr, flags); } +/// Descriptors opened with O_APPEND, whose writes go to the end of the file whole, as POSIX appends do. +var append_fds: [8192]u8 = @splat(0); + pub fn close(fd: fd_t) c_int { + if (fd >= 0 and fd < append_fds.len) @atomicStore(u8, &append_fds[@intCast(fd)], 0, .release); if (fd >= dir_base and fd < dir_base + dir_paths.len) { lockDirs(); defer dirs_lock.unlock(); @@ -313,6 +318,14 @@ pub fn read(fd: fd_t, buf: [*]u8, len: usize) isize { } pub fn write(fd: fd_t, buf: [*]const u8, len: usize) isize { + if (fd >= 0 and fd < append_fds.len and @atomicLoad(u8, &append_fds[@intCast(fd)], .acquire) != 0) { + // WriteFile at the end-of-file offset lands whole; the CRT's seek-then-write lets writers interleave + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = 0xFFFF_FFFF, .offset_high = 0xFFFF_FFFF }; + var put: u32 = 0; + if (WriteFile(h, buf, @intCast(@min(len, chunk)), &put, &ov) == 0) return fail(lastError()); + return put; + } return _write(fd, buf, @intCast(@min(len, chunk))); } diff --git a/zig/src/server/request_log.zig b/zig/src/server/request_log.zig index 02c26a0c0d..6d6513113e 100644 --- a/zig/src/server/request_log.zig +++ b/zig/src/server/request_log.zig @@ -40,3 +40,45 @@ pub fn append(a: Allocator, path: []const u8, body: Value) void { const framed = std.mem.concat(a, u8, &.{ line, "\n" }) catch return; _ = platform.c.write(fd, framed.ptr, framed.len); } + +test "appends from several threads keep every line whole" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const cwd = try std.process.currentPathAlloc(std.testing.io, a); + defer a.free(cwd); + const path = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/log.jsonl", .{ cwd, tmp.sub_path }); + defer a.free(path); + const threads = 8; + const lines = 200; + const head = "{\"line\": \""; + const Worker = struct { + fn run(p: []const u8, k: u8) void { + var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); + defer arena.deinit(); + const text = arena.allocator().alloc(u8, 300) catch return; + @memset(text, 'a' + k); + for (0..lines) |_| { + const o = json.newObject(arena.allocator()) catch return; + o.put(arena.allocator(), "line", .{ .string = text }) catch return; + append(arena.allocator(), p, .{ .object = o }); + } + } + }; + var workers: [threads]std.Thread = undefined; + for (&workers, 0..) |*w, k| w.* = try std.Thread.spawn(.{}, Worker.run, .{ path, @as(u8, @intCast(k)) }); + for (workers) |w| w.join(); + const got = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, path, a, .limited(1 << 24)); + defer a.free(got); + var counts: [threads]usize = @splat(0); + var it = std.mem.splitScalar(u8, got, '\n'); + while (it.next()) |line| { + if (line.len == 0) continue; + try std.testing.expectEqual(head.len + 300 + 2, line.len); + const k = line[head.len] - 'a'; + try std.testing.expect(k < threads); + for (line[head.len..][0..300]) |c| try std.testing.expectEqual('a' + k, c); + counts[k] += 1; + } + for (counts) |n| try std.testing.expectEqual(@as(usize, lines), n); +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 7630f5217a..9cfa2dc1d4 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -44,6 +44,7 @@ test { _ = @import("logprobs_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); + _ = @import("request_log.zig"); _ = @import("completion_prompt_test.zig"); _ = @import("server_init_test.zig"); _ = @import("decisions.zig"); From bdfeb8464af099b0be544c4b6552034aa937a830 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:30:32 +0000 Subject: [PATCH 23/46] golden check: a request-log mismatch names the first differing line The lifecycle/request-log outcome said only that the lines differ; it now gives both logs' line counts and the first line that differs, cut to 100 bytes each, so a mismatch seen on a CI runner can be read from its log. --- zig/tests/server/golden_check.zig | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 4976658e6b..13dc2d2b13 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -23,6 +23,26 @@ const known_diffs = [_]Known{ .{ .group = "anthropic", .name = "tool-use-stream" }, }; +/// Why two request logs differ: their line counts and the first line that does, each cut to 100 bytes. +fn logDiff(a: std.mem.Allocator, want: []const u8, got: []const u8) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + var wl = std.mem.splitScalar(u8, want, '\n'); + var gl = std.mem.splitScalar(u8, got, '\n'); + var n: usize = 0; + while (true) : (n += 1) { + const w = wl.next(); + const g = gl.next(); + if (w == null and g == null) break; + if (w != null and g != null and std.mem.eql(u8, w.?, g.?)) continue; + const ws = w orelse ""; + const gs = g orelse ""; + try out.append(a, try std.fmt.allocPrint(a, "line {d}: golden [{s}] zig [{s}]", .{ n + 1, ws[0..@min(ws.len, 100)], gs[0..@min(gs.len, 100)] })); + break; + } + try out.append(a, try std.fmt.allocPrint(a, "TENSORFOLD_REQUEST_LOG: golden {d} lines, zig {d} lines", .{ std.mem.count(u8, want, "\n"), std.mem.count(u8, got, "\n") })); + return out.items; +} + /// Reads a spawned server's stdout to its end, so its log lines never fill the pipe (4 KiB on Windows). fn drainStdout(io: Io, out: Io.File) void { var buf: [4096]u8 = undefined; @@ -490,7 +510,7 @@ pub fn main(init: std.process.Init) !void { .group = "lifecycle", .name = "request-log", .verdict = if (std.mem.eql(u8, want, got)) .equal else .differ, - .why = if (std.mem.eql(u8, want, got)) &.{} else &.{"TENSORFOLD_REQUEST_LOG lines differ from the frozen ones"}, + .why = if (std.mem.eql(u8, want, got)) &.{} else try logDiff(init.arena.allocator(), want, got), }); } else if (std.mem.eql(u8, kind, "rewrite_key")) { var file = try Io.Dir.cwd().createFile(io, s.key_path.?, .{}); From 39de0393a2a3d52844d8d0d7714504d0261aedfb Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 18:36:00 +0000 Subject: [PATCH 24/46] lane core and native CUDA host for two-rank families with a vision tower: batched prompt passes (prefill_many, addStreams), a backend's refusals for a memory guard, follow() for a tensor-parallel rank 1, shared-pool admission fields, image and video inputs for the native server (vision_inputs, media_fetch, qwen_image, stb_image) --- build.zig | 5 +- zig/build/cuda.zig | 6 +- zig/build/vision.zig | 10 + zig/src/core/engine_api.zig | 21 +- zig/src/core/lane_host.zig | 91 +- zig/src/core/lanes/backend.zig | 6 + zig/src/core/lanes/engine.zig | 68 + zig/src/core/lanes/fake.zig | 8 + zig/src/core/lanes/lanes.zig | 1 + zig/src/core/lanes/sampling.zig | 8 +- zig/src/core/lanes/stream.zig | 7 + zig/src/native/cuda.zig | 86 +- zig/src/native/cuda_bound.zig | 12 + zig/src/native/no_family.zig | 2 + zig/src/native/no_family_glm.zig | 2 + zig/src/server/chat.zig | 35 +- zig/src/server/cli.zig | 31 +- zig/src/server/main.zig | 4 + zig/src/server/media_fetch.zig | 336 ++ zig/src/server/openai.zig | 7 +- zig/src/server/prompt.zig | 14 + zig/src/server/serve.zig | 4 + zig/src/server/server.zig | 5 + zig/src/server/vision_inputs.zig | 342 ++ zig/src/vision/glm_image.zig | 194 + zig/src/vision/qwen_image.zig | 432 ++ zig/src/vision/tf_video.c | 220 + zig/vendor/stb/README.md | 3 + zig/vendor/stb/stb_image.h | 7988 ++++++++++++++++++++++++++++++ zig/vendor/stb/stb_image_impl.c | 10 + 30 files changed, 9922 insertions(+), 36 deletions(-) create mode 100644 zig/build/vision.zig create mode 100644 zig/src/native/no_family.zig create mode 100644 zig/src/native/no_family_glm.zig create mode 100644 zig/src/server/media_fetch.zig create mode 100644 zig/src/server/vision_inputs.zig create mode 100644 zig/src/vision/glm_image.zig create mode 100644 zig/src/vision/qwen_image.zig create mode 100644 zig/src/vision/tf_video.c create mode 100644 zig/vendor/stb/README.md create mode 100644 zig/vendor/stb/stb_image.h create mode 100644 zig/vendor/stb/stb_image_impl.c diff --git a/build.zig b/build.zig index 08765bd1df..481ad34d26 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const vision_build = @import("zig/build/vision.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -267,7 +268,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b counters_build.targets(b, target, optimize, metal, mods.sources, test_step); // The engines registry: the families this platform's binaries serve; the CLI reads it for models/info/pull. - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = mods.lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = mods.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, @@ -572,7 +573,7 @@ pub fn distMetalServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_id m.metal.addSystemFrameworkPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "System/Library/Frameworks" }))); m.metal.addLibraryPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/lib" }))); m.metal.addSystemIncludePath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/include" }))); - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = m.lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = m.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, .ReleaseFast) } } }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = m.metal }, .{ .name = "engine_api", .module = api }, .{ .name = "tensorfold", .module = m.engine } } }); return nativeServer(b, target, .fast, m.metal, m.engine, api, engines, build_options, null); } diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..4b4fdbb2dc 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -3,6 +3,7 @@ const std = @import("std"); const exlr8_build = @import("exlr8.zig"); const golden_build = @import("golden.zig"); +const vision_build = @import("vision.zig"); /// Each .cu in zig/kernels/cuda (`src`, else `name`) with the flags its Python extension passes in `extra_cuda_cflags`. const Kernel = struct { name: []const u8, flags: []const []const u8, src: ?[]const u8 = null, arch_specific: bool = false }; @@ -129,6 +130,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); runner.addImport("cuda", cuda); runner.addImport("core", mods.core); + runner.addImport("qwen_image", vision_build.module(b, target, optimize)); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); @@ -144,7 +146,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, @@ -210,7 +212,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug), vision_build.module(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/build/vision.zig b/zig/build/vision.zig new file mode 100644 index 0000000000..11eb50c143 --- /dev/null +++ b/zig/build/vision.zig @@ -0,0 +1,10 @@ +//! The image-input module (zig/src/vision/qwen_image.zig with stb_image compiled in), shared by every build that +//! carries the engine API (the server reaches it as engine_api.qwen_image) and by the CUDA test runner. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addCSourceFile(.{ .file = b.path("zig/vendor/stb/stb_image_impl.c"), .flags = &.{ "-std=c11", "-O2" } }); + m.addIncludePath(b.path("zig/vendor/stb")); + return m; +} diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 903dc4d40b..3c633640ad 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -1,5 +1,8 @@ //! The engine a server drives (``Engine``), and ``LaneHost``: the lane core served on one thread, rounds left to it. const std = @import("std"); +/// Image inputs (decode, resize, patches) for families with a vision tower; the server reaches them here. +pub const Image = lanes.Image; +pub const qwen_image = @import("qwen_image"); const lanes = @import("lanes"); const Allocator = std.mem.Allocator; @@ -88,6 +91,8 @@ pub const Request = struct { logprobs: ?u8 = null, /// Rows for each prompt token after the first, with ranks and this many best tokens, when ``Info.prompt_logprobs``. prompt_logprobs: ?u8 = null, + /// An image prompt's inputs in prompt order (Info.vision engines only; the prompt holds their placeholder rows). + images: []const lanes.Image = &.{}, }; pub const Reason = enum { stop, length, cancelled, failed }; @@ -162,6 +167,10 @@ pub const Info = struct { prompt_cache_plan: ?PromptCachePlan = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, + /// The engine encodes image inputs (Request.images): it was opened with ``vision`` and its family has a tower. + vision: bool = false, + /// Its tower takes GLM-5.3-Flash's preparation and placeholders (qwen_image.glm) instead of Qwen's. + vision_glm: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. admit: ?Admit = null, }; @@ -225,10 +234,20 @@ pub const Open = struct { master_port: u16 = 29551, /// --policy: the engine's policy as `key=value,...`; empty: its defaults. policy: []const u8 = "", + /// --vision: load the family's vision tower and accept image inputs. + vision: bool = false, + /// --kv-dtype: bits of the attention cache's codes (8 or 4) for families that keep a quantized cache. + kv_bits: u8 = 8, }; /// 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, + /// Rank 1 of two: replay rank 0's engine calls until it stops (the server then serves no HTTP); null elsewhere. + follow: ?*const fn (ctx: *anyopaque) anyerror!void = null, +}; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0 }; diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4394fdac05..7a557c20c3 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -286,15 +286,16 @@ pub const LaneHost = struct { } /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - fn admitOne(h: *LaneHost) bool { + /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit yet). + fn takeJob(h: *LaneHost) ?*Job { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); - return false; + return null; } const at = admit.pick(h) orelse { // a shared cache pool's head waits for room h.unlock(); - return false; + return null; }; const job = h.queued.orderedRemove(at); job.seq = h.admissions; @@ -302,9 +303,14 @@ pub const LaneHost = struct { h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return null; }; h.unlock(); + return job; + } + + /// The admitted request's stream (its kept state, drafter and grammar); false when it was dropped instead. + fn setup(h: *LaneHost, job: *Job) bool { const r = job.request; if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran var reuse: lanes.stream.Reuse = .{}; @@ -315,7 +321,10 @@ pub const LaneHost = struct { 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.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch { + _ = h.drop(job, "the drafter could not start"); + return false; + }; job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", .prompt = r.prompt, @@ -339,18 +348,31 @@ pub const LaneHost = struct { .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)"); + { + _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return false; + } } else null, + .images = r.images, .decode_spans = h.spans(r), .decide = r.decide, .untaught = r.untaught, }) catch { job.proposer.deinit(); - return h.drop(job, "out of memory"); + { + _ = h.drop(job, "out of memory"); + return false; + } }; job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; + } + + /// An admitted, set-up request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + fn launch(h: *LaneHost, job: *Job) bool { + const r = job.request; + const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (lone.fits(h, job)) return lone.run(h, job, began, false); if (h.core.fills()) { @@ -366,6 +388,55 @@ pub const LaneHost = struct { if (h.deliver(job)) h.remove(job); return true; } + fn admitOne(h: *LaneHost) bool { + const job = h.takeJob() orelse return false; + if (!h.setup(job)) return true; + return h.launch(job); + } + + /// Every queued request a lane is free for, their prompts in one pass where the backend takes several + /// (lanes.Engine.addStreams); the engine's special paths (chunked prompts, warm passes, the lone driver) go + /// through admitOne as before. + fn admitMany(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return h.admitOne(); + h.lock(); + const waiting = h.queued.items.len; + const free = h.info_.lanes -| h.admitted.items.len; + h.unlock(); + if (waiting < 2 or free < 2) return h.admitOne(); + var jobs: [lanes.Engine.max_together]*Job = undefined; + var n: usize = 0; + while (n < jobs.len and n < free) { + const job = h.takeJob() orelse break; + if (!h.setup(job)) continue; + if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { + _ = h.launch(job); + continue; + } + jobs[n] = job; + n += 1; + } + if (n == 0) return true; + if (n == 1) return h.launch(jobs[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], errs[0..n]) |job, err| { + if (err) |e| { + _ = switch (e) { + error.Cancelled => h.cancel(job), + error.LaneMemory => h.requeue(job), + else => h.drop(job, h.words(e)), + }; + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + return true; + } + /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt. fn fillOne(h: *LaneHost) void { @@ -494,7 +565,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitOne()) {} + while (h.admitMany()) {} h.fillOne(); status.noteLive(h); h.lock(); diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index fe26d3f003..723dc5effc 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -68,6 +68,9 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// Prefill several streams' prompts in one pass (each then as after `prefill`); false, having done nothing, + /// when they do not go together (each is then prefilled alone). + prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it. @@ -102,6 +105,9 @@ pub const Backend = struct { prompt_rows: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, + /// After a failed round: the streams it failed for alone (their caches could not grow), which the caller + /// ends while the round's other streams go on; 0 when the failure is the round's. + refused: ?*const fn (ptr: *anyopaque, out: []*Stream) usize = null, }; pub fn prefill(b: Backend, s: *Stream) !void { diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index cf0a97ea4a..d5be72c0bb 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -83,6 +83,68 @@ pub const Engine = struct { pub const fills = open.fills; pub const fillStream = open.fillStream; + /// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. + pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { + _ = e.arena.reset(.retain_capacity); + s.context.shrinkRetainingCapacity(s.prompt_len); + s.pending = null; + s.cache_len = s.prompt_len; + if (e.cfg.family_mtp and s.drafts) { + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; + }; + } + if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row + _ = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.finished) return e.release(s); + try e.live.append(e.gpa, s); + } + + /// Streams a backend prefills in one pass at most. + pub const max_together = 16; + + /// Whether the backend prefills several prompts in one pass (addStreams batches them). + pub fn batchesPrompts(e: *const Engine) bool { + return e.backend.vtable.prefill_many != null; + } + + /// open.add for several streams: their prompts in one pass when the backend takes them together (else each + /// alone), then each opened as after its own prefill. ``errs[i]``: stream i's error (a cancelled one is released + /// already; the caller discards a failed one). + pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return e.addEach(ss, errs); + if (ss.len < 2 or ss.len > max_together) return e.addEach(ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch |err| { + for (errs) |*x| x.* = err; + return; + }; + if (!together) return e.addEach(ss, errs); + for (ss, errs) |s, *x| { + if (s.isCancelled()) { // cancelled in its pass: no first token + e.backend.release(s); + x.* = error.Cancelled; + continue; + } + open.opened(e, s) catch |err| { + x.* = err; + }; + } + } + + fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| open.add(e, s) catch |err| { + x.* = err; + }; + } + /// Take over a stream another driver decoded so far: its cache settled, `pending` and `cache_len` set, no drafts held. pub fn adopt(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); @@ -100,6 +162,12 @@ pub const Engine = struct { return false; } + /// After ``step`` failed: the streams the backend refused alone (the caller ends them; the rest go on). + pub fn refused(e: *Engine, out: []*Stream) usize { + const ask = e.backend.vtable.refused orelse return 0; + return ask(e.backend.ptr, out); + } + /// Release a cancelled stream between rounds: no pending draws or drafts. pub fn discard(e: *Engine, s: *Stream) void { s.finished = true; diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index 9434736c46..fc90ee7483 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -243,6 +243,14 @@ pub const Fake = struct { return self(ptr).drawn.items[handle]; } + fn refusedFn(ptr: *anyopaque, out: []*Stream) usize { + const x = self(ptr); + const n = @min(out.len, x.refused.items.len); + @memcpy(out[0..n], x.refused.items[0..n]); + x.refused.clearRetainingCapacity(); + return n; + } + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const x = self(ptr); x.rounds += 1; diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 5d50a16ef3..57d8ae9eb7 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -29,6 +29,7 @@ pub const grammar = @import("grammar.zig"); pub const Engine = engine.Engine; pub const Config = config.Config; +pub const Image = @import("stream.zig").Image; pub const Model = config.Model; pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; diff --git a/zig/src/core/lanes/sampling.zig b/zig/src/core/lanes/sampling.zig index e9d469756c..43bddfa3a1 100644 --- a/zig/src/core/lanes/sampling.zig +++ b/zig/src/core/lanes/sampling.zig @@ -75,7 +75,7 @@ pub fn pairwiseSum(a: []const f64) f64 { return pairwiseSum(a[0..n2]) + pairwiseSum(a[n2..]); } -const Candidate = struct { value: f64, id: u64 }; +pub const Candidate = struct { value: f64, id: u64 }; fn byValueThenId(_: void, a: Candidate, b: Candidate) bool { return a.value > b.value or (a.value == b.value and a.id < b.id); @@ -88,6 +88,12 @@ pub fn choose(gpa: Allocator, values: []const f64, ids: []const u64, position: u defer gpa.free(order); for (order, values, ids) |*o, v, id| o.* = .{ .value = v, .id = id }; std.mem.sort(Candidate, order, {}, byValueThenId); + return chooseOrdered(gpa, order, position, s); +} + +/// ``choose`` over candidates already in its order (by value, then the lower id at a tie). +pub fn chooseOrdered(gpa: Allocator, order: []const Candidate, position: u64, s: Sampling) !u64 { + const width = order.len; if (s.temperature <= 0) return order[0].id; // temperature 0 is greedy: the first maximum, ties to the lower id const k = @max(1, @min(if (s.top_k != 0) @as(usize, s.top_k) else width, width)); const scaled = try gpa.alloc(f64, k); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index f028d10cc4..aade9ad50c 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -97,6 +97,10 @@ pub const Held = struct { parents: ?[]i32 = null, }; +/// An image input prepared for a vision tower: its patches (fp32 [grid h * w, 1536], merge-block order) and grid +/// [1, h, w]; the prompt holds its placeholder rows. +pub const Image = struct { patches: []const f32, grid: [3]i64 }; + pub const Spec = struct { id: []const u8, prompt: []const u32, @@ -123,6 +127,7 @@ pub const Spec = struct { prompt_logprobs: ?u8 = null, // and for each prompt token after the first, with ranks (the prompt pass starts at 0) grammar: ?*grammar.Constraint = null, // structured output: the reply stays in it (the creator frees it) decode_spans: []const [2]u32 = &.{}, // prompt rows an earlier reply decoded (Request.decode_spans), for backends that keep their bits + images: []const Image = &.{}, // an image prompt's inputs, in prompt order (the backend encodes them at prefill) }; pub const Stream = struct { @@ -140,6 +145,7 @@ pub const Stream = struct { think_close: []const u32, think_end: i64, think_open: bool, + images: []const Image = &.{}, loop_guard: bool, loop_period: ?u32 = null, chunks: []const u32, @@ -202,6 +208,7 @@ pub const Stream = struct { .cancel_check = spec.cancel_check, .yield_check = spec.yield_check, .think_budget = spec.think_budget, + .images = spec.images, .think_close = spec.think_close, .think_end = spec.think_end, .think_open = spec.think_open orelse (spec.think_budget > 0 or (spec.loop_guard and spec.think_end >= 0)), diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..a34186a4a3 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,4 +1,7 @@ //! CUDA family registration, device setup, memory admission and the native lane host. +/// The serve flags this build's engine takes past the common ones (server/cli.zig served()): two ranks. +pub const serves: []const []const u8 = &.{"--tp"}; + const std = @import("std"); const cuda = @import("cuda"); const api = @import("engine_api"); @@ -152,6 +155,14 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + follow_fn: ?*const fn (*anyopaque) anyerror!void = null, // rank 1 of two: the family's replay loop + + /// Rank 1: the family replays rank 0's calls on this thread (the context was made current here at open). + fn followFamily(p: *anyopaque) anyerror!void { + const h: *Host = @ptrCast(@alignCast(p)); + if (h.gpu) |g| try g.ctx.makeCurrent(); + return h.follow_fn.?(h.family); + } fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -241,6 +252,22 @@ const Host = struct { const refusals = @import("cuda_refusal.zig"); +/// A tensor-parallel rank above 0 (server/engines.zig follow): the family opened as that rank, then its replay loop +/// until rank 0 stops; false with ``problem`` set when it cannot. +pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { + const opened = (try open(a, gpa, io, o, problem)) orelse return false; + defer opened.close(opened.ctx); + const f = opened.follow orelse { + problem.* = "this family has no tensor-parallel ranks"; + return false; + }; + f(opened.ctx) catch |e| { + problem.* = try std.fmt.allocPrint(a, "rank {d} stopped: {s}", .{ o.rank, @errorName(e) }); + return false; + }; + return true; +} + /// The engine for `o.dir`, or null with `problem` set when `refusal` names a flag or no CUDA family reads it. pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { if (refusals.backend(o)) |why| { @@ -322,20 +349,42 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + const two = @hasDecl(F, "two_ranks") and F.two_ranks; + if (two and o.tp != 2) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on two GPUs here: serve it with --tp 2 --rank 0 on one host and --tp 2 --rank 1 --master on the other", .{F.model_type}); + return null; + } + if (!two and o.tp != 1) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on one GPU: drop --tp", .{F.model_type}); + return null; + } + const weights = if (@hasDecl(F, "deviceWeightBytes")) F.deviceWeightBytes(io, o.dir) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; } - const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }) catch |e| { + var fo: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; + if (@hasField(F.Options, "vision")) fo.vision = o.vision; + if (@hasField(F.Options, "kv_bits")) fo.kv_bits = o.kv_bits else if (o.kv_bits != 8) { + problem.* = try std.fmt.allocPrint(a, "{s} keeps an int8 attention cache only: drop --kv-dtype", .{F.model_type}); + return null; + } + if (@hasField(F.Options, "slide")) fo.slide = o.slide; + if (@hasField(F.Options, "rank")) { + fo.tp = @intCast(o.tp); + fo.rank = @intCast(o.rank); + fo.master = if (o.master.len > 0) o.master else null; + fo.master_port = o.master_port; + } + const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, fo) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; }; defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + const room = after.room(model) -| learnBytes(loaded); const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { @@ -344,16 +393,22 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; return null; }; + if (@hasDecl(F, "setStreams")) F.setStreams(loaded.ctx, streams); + if (@hasDecl(F, "setRoom")) F.setRoom(loaded.ctx, room) catch |e| { // a family that sizes each stream's caches out of one pool + problem.* = try std.fmt.allocPrint(a, "the ranks could not agree on the stream pool ({s})", .{@errorName(e)}); + return null; + }; const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); - h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + // the lane core's kept states (Nemotron's snapshots); a family that keeps its own prompt states has no `target` + const cache = if (@hasField(@TypeOf(loaded), "target")) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; + h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in {d}-row chunks{s}{s}; {s}", .{ + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learnBytes(loaded) > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learnBytes(loaded))}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasField(@TypeOf(loaded), "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet @@ -364,7 +419,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (comptime @hasField(@TypeOf(loaded), "learns")) { // a family with a Living Weights learner + if (loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); @@ -378,17 +434,29 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O } }.f; h.learn = s.hook(); + } } errdefer if (h.slide) |s| { h.slide_free.?(gpa, s); gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, + .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; + const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; + if (follow_loop) |f| { + h.follow_fn = f; + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h, .follow = Host.followFamily }; + } return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// A family's bytes kept for --slide (0 for families without a learner). +fn learnBytes(loaded: anytype) u64 { + return if (@hasField(@TypeOf(loaded), "learn_bytes")) loaded.learn_bytes else 0; +} + fn toGib(bytes: u64) f64 { return @as(f64, @floatFromInt(bytes)) / gib; } diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 40977dfa50..2144981686 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -18,6 +18,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.prefill_step.?(x.inner.ptr, s); } }.f else null, + .prefill_many = if (v.prefill_many != null) struct { + fn f(p: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { + const x = H.bind(p); + return x.inner.vtable.prefill_many.?(x.inner.ptr, ss); + } + }.f else null, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = H.bind(p); @@ -108,6 +114,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.held.?(x.inner.ptr, s, out); } }.f else null, + .refused = if (v.refused != null) struct { + fn f(p: *anyopaque, out: []*lanes.Stream) usize { + const x = H.bind(p); + return x.inner.vtable.refused.?(x.inner.ptr, out); + } + }.f else null, .prompt_rows = v.prompt_rows, .release = struct { fn f(p: *anyopaque, s: *lanes.Stream) void { diff --git a/zig/src/native/no_family.zig b/zig/src/native/no_family.zig new file mode 100644 index 0000000000..d81d7567ac --- /dev/null +++ b/zig/src/native/no_family.zig @@ -0,0 +1,2 @@ +//! Stands in for a CUDA family the build leaves out (-Dflashnext=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/native/no_family_glm.zig b/zig/src/native/no_family_glm.zig new file mode 100644 index 0000000000..2dcb07a292 --- /dev/null +++ b/zig/src/native/no_family_glm.zig @@ -0,0 +1,2 @@ +//! Stands in for the GLM CUDA family when the build leaves it out (-Dglm=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 7c702ef432..82c72d52ad 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -35,6 +35,8 @@ pub const Input = struct { logprobs: ?u8 = null, /// ``/v1/completions``' ``prompt_logprobs``: rows for the prompt with this many alternatives; null: not asked. prompt_logprobs: ?u8 = null, + /// Image inputs (vision_inputs.extract), in prompt order; the rendered prompt's image pads expand to their tokens. + images: []const api.Image = &.{}, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -90,9 +92,31 @@ pub const Prepared = struct { preparing: bool, }; +/// --tool-system: the server's instruction as the first message of a chat request that offers tools and sends no +/// system message (the client's own instruction wins); other requests pass unchanged. +fn withToolSystem(srv: *Server, a: Allocator, input: Input) Allocator.Error!Input { + const text = srv.config.tool_system orelse return input; + if (input.tools.len == 0 or input.prompt != null or input.messages != .array) return input; + for (input.messages.array) |m| { + if (m != .object) continue; + const role = m.object.get("role") orelse continue; + if (role == .string and std.mem.eql(u8, role.string, "system")) return input; + } + const sys = try json.newObject(a); + try sys.put(a, "role", .{ .string = "system" }); + try sys.put(a, "content", .{ .string = text }); + const list = try a.alloc(Value, input.messages.array.len + 1); + list[0] = .{ .object = sys }; + @memcpy(list[1..], input.messages.array); + var out = input; + out.messages = .{ .array = list }; + return out; +} + /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. -pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepared { +pub fn prepare(srv: *Server, cx: *Cx, input_in: Input, gone: anytype) Failure!Prepared { const a = cx.a; + const input = try withToolSystem(srv, a, input_in); const io = srv.io; const received = nowNs(io); const f = input.fields; @@ -108,7 +132,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); - const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + var rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + if (input.images.len > 0) rendered.history_len = 0; // image prompts keep no prefix for a later turn (as Python) if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -119,7 +144,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (input.max_tokens != null and limit > room) return cx.fail(.context_length, "{s} {d} tokens, but the rendered prompt has {d} tokens and requests {d} reply tokens, which exceeds the context window. Reduce the prompt to at most {d} prompt tokens or request at most {d} reply tokens, including chat template and thinking tokens.", .{ errors.context_limit, window, n, limit, @max(0, window - limit), room }); limit = @min(limit, room); } - const system_len: usize = if (input.prompt != null) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); + const system_len: usize = if (input.prompt != null or input.images.len > 0) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); var shared: std.ArrayList(u32) = .empty; if (system_len > 0) for ([_]i64{ @as(i64, @intCast(system_len)) - 2048, @as(i64, @intCast(system_len)) - 512, @intCast(system_len) }) |cut| { if (cut >= 512) try shared.append(a, @intCast(cut)); @@ -137,7 +162,9 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .history_len = @intCast(rendered.history_len), .rewind_len = @intCast(rendered.rewind_len), .shared_prefixes = shared.items, - .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), // a cut before the conversation's own text: fresh sessions resume their whole harness + .images = input.images, + // a cut just before the conversation's own text: fresh sessions resume their whole harness + .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), .decode_spans = try prompt_mod.replySpans(srv, cx, rendered.ids), // from the tokens alone, for every request: cached states are keyed by tokens .tools_json = if (input.tools.len > 0) try json.stringify(a, .{ .array = @constCast(input.tools) }, .{ .ascii = false }) else "", .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..852b3730b7 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -49,15 +49,16 @@ pub const flags = [_]Flag{ .{ .name = "--api-key-file", .native = true }, .{ .name = "--metrics-open", .kind = .store_true, .native = true }, .{ .name = "--dashboard", .kind = .store_true, .native = true }, - .{ .name = "--vision", .kind = .store_true }, - .{ .name = "--vision-urls", .kind = .store_true }, + .{ .name = "--vision", .kind = .store_true, .native = true }, + .{ .name = "--vision-urls", .kind = .store_true, .native = true }, .{ .name = "--vision-offload", .kind = .store_true }, - .{ .name = "--vision-max-images" }, - .{ .name = "--vision-image-tokens" }, + .{ .name = "--vision-max-images", .native = true }, + .{ .name = "--vision-image-tokens", .native = true }, .{ .name = "--context", .native = true }, .{ .name = "--speed-up", .native = !cuda_build }, .{ .name = "--max-tokens", .native = true }, .{ .name = "--temperature", .native = true }, + .{ .name = "--tool-system", .native = true }, .{ .name = "--top-p", .native = true }, .{ .name = "--top-k", .native = true }, .{ .name = "--min-p", .native = true }, @@ -99,7 +100,7 @@ pub const flags = [_]Flag{ .{ .name = "--master-port", .native = served("--tp") }, // the GPU engine's policy switches as key=value pairs, repeatable .{ .name = "--policy", .kind = .append, .native = served("--policy") }, - .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" } }, + .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" }, .native = served("--kv-dtype"), .native_values = &.{ "int8", "int4" } }, .{ .name = "--prefill-fp8", .kind = .store_true }, .{ .name = "--no-prefill-fp8", .kind = .store_true }, .{ .name = "--precision", .choices = &.{ "checkpoint", "full" } }, @@ -140,6 +141,11 @@ pub const Args = struct { learn_gib: f64 = 32, // disk for learned states, every model and build together max_tokens: i64 = 4096, temperature: ?f64 = null, + vision: bool = false, // accept image inputs (an engine whose family has a vision tower) + vision_urls: bool = false, // also public HTTPS image and video URLs + vision_max_images: i64 = 4, + vision_image_tokens: i64 = 4096, // visual tokens a request's images share + tool_system: ?[]const u8 = null, // an instruction added as the system message of tool requests that send none top_p: ?f64 = null, top_k: ?i64 = null, min_p: ?f64 = null, @@ -170,6 +176,7 @@ pub const Args = struct { master_port: u16 = 29551, policy: []const u8 = "", // every --policy joined by commas cortheon: bool = false, // Cortheon's program, research before building, and cortheon_search/cortheon_fetch run by the server + kv_dtype: []const u8 = "int8", // --kv-dtype (CUDA): the attention cache's codes, int8 or int4 }; /// A usage error's message (argparse's ``error:`` line); the caller exits 2. @@ -265,7 +272,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa } else if (is(name, "--learn-gib")) { out.learn = true; out.learn_gib = try gib(u, a, name, v); - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--tool-system")) out.tool_system = v else if (is(name, "--vision")) out.vision = true else if (is(name, "--vision-urls")) out.vision_urls = true else if (is(name, "--vision-max-images")) out.vision_max_images = try int(u, a, name, v) else if (is(name, "--vision-image-tokens")) out.vision_image_tokens = try int(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -282,9 +289,19 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa fn cudaFlag(a: Allocator, out: *Args, name: []const u8, v: []const u8, u: *Usage) error{ Usage, OutOfMemory }!bool { if (std.mem.eql(u8, name, "--device")) { out.device = std.math.cast(u32, try int(u, a, name, v)) orelse return fail(u, a, "argument --device: a GPU ordinal from 0: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--kv-dtype")) { + out.kv_dtype = v; } else if (std.mem.eql(u8, name, "--segments")) { const n = try int(u, a, name, v); out.segments = if (n >= 1) @intCast(@min(n, std.math.maxInt(u32))) else return fail(u, a, "argument --segments: a count from 1: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--tp")) { + out.tp = if (std.mem.eql(u8, v, "2")) 2 else 1; + } else if (std.mem.eql(u8, name, "--rank")) { + out.rank = if (std.mem.eql(u8, v, "1")) 1 else 0; + } else if (std.mem.eql(u8, name, "--master")) { + out.master = v; + } else if (std.mem.eql(u8, name, "--master-port")) { + out.master_port = std.math.cast(u16, try int(u, a, name, v)) orelse return fail(u, a, "argument --master-port: a TCP port: '{s}'", .{v}); } else return false; return true; } @@ -362,6 +379,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .mtp_drafts = args.mtp_drafts, .mtp_confidence = args.mtp_confidence, .keep = args.checkpoint_slots, + .vision = args.vision, + .kv_bits = if (std.mem.eql(u8, args.kv_dtype, "int4")) 4 else 8, .tp = args.tp, .rank = args.rank, .master = args.master, diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 4b85fca8cd..0cd1b03c6f 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -78,6 +78,10 @@ pub fn main(init: std.process.Init) !u8 { } orelse return fail(problem); defer up.text.deinit(); defer up.engine.close(up.engine.ctx); + if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls + follow(up.engine.ctx) catch |e| return fail(@errorName(e)); + return 0; + } return serve.run(gpa, io, args, .{ .engine = up.engine.engine, .text = up.text.text(), diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig new file mode 100644 index 0000000000..e175033bf6 --- /dev/null +++ b/zig/src/server/media_fetch.zig @@ -0,0 +1,336 @@ +//! Public HTTPS media for --vision-urls, as the Python server fetches it (vision/images_http.py): an https URL on +//! port 443 without credentials or a fragment, its host resolved and every address checked public, a TLS connection +//! verified for the host on an address that was checked, GET with identity encoding, up to 3 redirects (each checked +//! the same way), a declared media type the caller accepts, the bytes within the caller's limit, all inside a deadline. + +const std = @import("std"); +const Io = std.Io; +const net = std.Io.net; + +pub const Error = error{ OutOfMemory, Media }; + +/// The failure's words for the client (the request is refused with them). +pub const Failure = struct { text: []const u8 = "" }; + +pub const max_url_chars = 4096; +pub const max_redirects = 3; +pub const timeout_s = 10; // one download +pub const total_s = 30; // a request's downloads together + +fn fail(f: *Failure, text: []const u8) Error { + f.text = text; + return error.Media; +} + +/// Python ipaddress: is_global and not multicast or reserved, nor 192.0.0.0/24 or 168.63.129.16; no IPv6 address +/// that embeds an IPv4 one (mapped, 6to4, Teredo). +pub fn publicIp(a: net.IpAddress) bool { + switch (a) { + .ip4 => |v| { + const b = v.bytes; + const in = struct { + fn net4(x: [4]u8, base: [4]u8, bits: u5) bool { + const xi = std.mem.readInt(u32, &x, .big); + const bi = std.mem.readInt(u32, &base, .big); + const mask: u32 = if (bits == 0) 0 else ~@as(u32, 0) << @intCast(32 - @as(u6, bits)); + return xi & mask == bi & mask; + } + }.net4; + const blocked = [_]struct { [4]u8, u5 }{ + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, + }; + for (blocked) |blk| if (in(b, blk[0], blk[1])) return false; + if (std.mem.eql(u8, &b, &.{ 168, 63, 129, 16 })) return false; + return true; + }, + .ip6 => |v| { + const b = v.bytes; + // global unicast 2000::/3 only, less documentation (2001:db8::/32), the 2001::/23 IETF block (Teredo + // 2001::/32 among it), 6to4 (2002::/16) + if (b[0] & 0xe0 != 0x20) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; + if (b[0] == 0x20 and b[1] == 0x02) return false; + return true; + }, + } +} + +const Url = struct { host: []const u8, target: []const u8 }; + +/// grammar of images_http._url: https, a host, port 443, no credentials, fragment, backslash, whitespace or controls. +fn parseUrl(a: std.mem.Allocator, value: []const u8, f: *Failure) Error!Url { + if (value.len > max_url_chars) return fail(f, "image URL is too long or contains whitespace/control characters"); + for (value) |c| if (c <= 32 or c == 127) return fail(f, "image URL is too long or contains whitespace/control characters"); + const bad = "image URL must be HTTPS on port 443, without credentials or a fragment"; + if (!std.ascii.startsWithIgnoreCase(value, "https://")) return fail(f, bad); + if (std.mem.indexOfScalar(u8, value, '\\') != null or std.mem.indexOfScalar(u8, value, '#') != null) return fail(f, bad); + const rest = value["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const authority = rest[0..end]; + if (std.mem.indexOfScalar(u8, authority, '@') != null) return fail(f, bad); + var host = authority; + if (std.mem.lastIndexOfScalar(u8, authority, ':')) |c| { + if (authority.len > 0 and authority[0] == '[') return fail(f, bad); // a literal IPv6 host: not a name + if (!std.mem.eql(u8, authority[c + 1 ..], "443") and authority[c + 1 ..].len > 0) return fail(f, bad); + host = authority[0..c]; + } + if (host.len == 0 or std.mem.indexOfScalar(u8, host, '%') != null) return fail(f, bad); + for (host) |c| if (c >= 128) return fail(f, bad); // IDNA names: their ASCII (xn--) form + const lower = try std.ascii.allocLowerString(a, std.mem.trimEnd(u8, host, ".")); + for ([_][]const u8{ "localhost", "metadata.google.internal", "instance-data" }) |n| if (std.mem.eql(u8, lower, n)) return fail(f, "image URLs must use public internet hosts"); + // the path and query, percent-quoted as urllib.parse.quote(safe=...) leaves them + var t: std.ArrayList(u8) = .empty; + var tail = rest[end..]; + if (tail.len == 0 or tail[0] == '?') try t.append(a, '/'); + const safe_path = "/%:@!$&'()*+,;=-._~"; + var in_query = false; + for (tail) |c| { + if (c == '?' and !in_query) { + in_query = true; + try t.append(a, '?'); + continue; + } + const safe = std.ascii.isAlphanumeric(c) or std.mem.indexOfScalar(u8, safe_path, c) != null or (in_query and c == '?'); + if (safe) try t.append(a, c) else try t.print(a, "%{X:0>2}", .{c}); + } + tail = t.items; + return .{ .host = lower, .target = tail }; +} + +/// The location a redirect names, resolved against ``base`` (urljoin for absolute and absolute-path forms). +fn join(a: std.mem.Allocator, base: []const u8, location: []const u8) ![]const u8 { + if (std.ascii.startsWithIgnoreCase(location, "https://") or std.ascii.startsWithIgnoreCase(location, "http://")) return a.dupe(u8, location); + if (std.mem.startsWith(u8, location, "//")) return std.fmt.allocPrint(a, "https:{s}", .{location}); + const rest = base["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const origin = base[0 .. "https://".len + end]; + if (std.mem.startsWith(u8, location, "/")) return std.fmt.allocPrint(a, "{s}{s}", .{ origin, location }); + const path_end = std.mem.indexOfScalar(u8, rest[end..], '?') orelse rest.len - end; + const path = rest[end..][0..path_end]; + const dir = if (std.mem.lastIndexOfScalar(u8, path, '/')) |s| path[0 .. s + 1] else "/"; + return std.fmt.allocPrint(a, "{s}{s}{s}", .{ origin, dir, location }); +} + +fn remaining(io: Io, deadline: Io.Timestamp, f: *Failure) Error!i64 { + const left = Io.Clock.awake.now(io).durationTo(deadline).toMilliseconds(); + if (left <= 0) return fail(f, "image download timed out"); + return left; +} + +fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { + const tv: std.posix.timeval = .{ .sec = @intCast(@divTrunc(ms, 1000)), .usec = @intCast(@mod(ms, 1000) * 1000) }; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.RCVTIMEO, std.mem.asBytes(&tv)) catch {}; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, std.mem.asBytes(&tv)) catch {}; +} + +/// A TCP connection to the checked address, its connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds a +/// Linux connect too; std's connect has no timeout yet). +fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + const posix = std.posix; + const linux = std.os.linux; + const family: u32 = switch (addr) { + .ip4 => posix.AF.INET, + .ip6 => posix.AF.INET6, + }; + const rc = linux.socket(family, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (linux.errno(rc) != .SUCCESS) return fail(f, "image download failed or timed out"); + const fd: posix.fd_t = @intCast(rc); + errdefer _ = linux.close(fd); + setTimeouts(fd, ms); + var sa4: posix.sockaddr.in = undefined; + var sa6: posix.sockaddr.in6 = undefined; + const r = switch (addr) { + .ip4 => |v| blk: { + sa4 = .{ .port = std.mem.nativeToBig(u16, 443), .addr = @bitCast(v.bytes) }; + break :blk linux.connect(fd, @ptrCast(&sa4), @sizeOf(posix.sockaddr.in)); + }, + .ip6 => |v| blk: { + sa6 = .{ .port = std.mem.nativeToBig(u16, 443), .flowinfo = 0, .addr = v.bytes, .scope_id = 0 }; + break :blk linux.connect(fd, @ptrCast(&sa6), @sizeOf(posix.sockaddr.in6)); + }, + }; + if (linux.errno(r) != .SUCCESS) return fail(f, "image download failed or timed out"); + return .{ .socket = .{ .handle = fd, .address = addr } }; +} + +/// The bytes of ``url`` (and its media type, one of ``media_types``), at most ``max_bytes``, before ``deadline``. +pub fn fetch(a: std.mem.Allocator, io: Io, url_in: []const u8, max_bytes: usize, deadline_in: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!struct { data: []u8, media: []const u8 } { + var url = url_in; + const deadline = blk: { + const one = Io.Clock.awake.now(io).addDuration(.fromSeconds(timeout_s)); + break :blk if (one.nanoseconds < deadline_in.nanoseconds) one else deadline_in; + }; + var hop: usize = 0; + while (true) : (hop += 1) { + const u = try parseUrl(a, url, f); + const got = try request(a, io, u, max_bytes, deadline, media_types, f); + switch (got) { + .body => |b| return .{ .data = b.data, .media = b.media }, + .redirect => |loc| { + if (hop == max_redirects) return fail(f, "image download has too many redirects"); + url = join(a, url, loc) catch return error.OutOfMemory; + }, + } + } +} + +const Got = union(enum) { body: struct { data: []u8, media: []const u8 }, redirect: []const u8 }; + +var bundle_lock: Io.RwLock = .init; +var bundle: std.crypto.Certificate.Bundle = .empty; +var bundle_loaded = false; + +fn request(a: std.mem.Allocator, io: Io, u: Url, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + // every address the name resolves to must be public; the connection goes to the first of them + const name = net.HostName.init(u.host) catch return fail(f, "image host could not be resolved"); + var results: [32]net.HostName.LookupResult = undefined; + var q: Io.Queue(net.HostName.LookupResult) = .init(&results); + name.lookup(io, &q, .{ .port = 443 }) catch return fail(f, "image host could not be resolved"); + var first: ?net.IpAddress = null; + var count: usize = 0; + while (q.getOneUncancelable(io)) |r| switch (r) { + .address => |addr| { + if (!publicIp(addr)) return fail(f, "image URLs must resolve only to public internet addresses"); + if (first == null) first = addr; + count += 1; + }, + .canonical_name => {}, + } else |_| {} + const addr = first orelse return fail(f, "image host could not be resolved"); + _ = try remaining(io, deadline, f); + // the CA bundle, read once + { + bundle_lock.lockUncancelable(io); + defer bundle_lock.unlock(io); + if (!bundle_loaded) { + bundle.rescan(std.heap.page_allocator, io, Io.Clock.real.now(io)) catch return fail(f, "image download failed: no CA certificates on the server"); + bundle_loaded = true; + } + } + const stream = try connect(addr, try remaining(io, deadline, f), f); + defer stream.close(io); + const tls_len = std.crypto.tls.Client.min_buffer_len; + const bufs = try a.alloc(u8, 4 * tls_len); + var sr = stream.reader(io, bufs[0..tls_len]); + var sw = stream.writer(io, bufs[tls_len .. 2 * tls_len]); + var entropy: [std.crypto.tls.Client.Options.entropy_len]u8 = undefined; + io.random(&entropy); + var tls = std.crypto.tls.Client.init(&sr.interface, &sw.interface, .{ + .host = .{ .explicit = u.host }, + .ca = .{ .bundle = .{ .gpa = std.heap.page_allocator, .io = io, .lock = &bundle_lock, .bundle = &bundle } }, + .read_buffer = bufs[2 * tls_len .. 3 * tls_len], + .write_buffer = bufs[3 * tls_len ..], + .entropy = &entropy, + .realtime_now = Io.Clock.real.now(io), + .allow_truncation_attacks = true, // HTTP's own lengths and framing end the body + }) catch return fail(f, "image download failed: TLS could not be verified for the host"); + // the request + const accept = std.mem.join(a, ", ", media_types) catch return error.OutOfMemory; + tls.writer.print("GET {s} HTTP/1.1\r\nHost: {s}\r\nAccept: {s}\r\nAccept-Encoding: identity\r\nUser-Agent: TensorFold-native\r\nConnection: close\r\n\r\n", .{ u.target, u.host, accept }) catch return fail(f, "image download failed or timed out"); + tls.writer.flush() catch return fail(f, "image download failed or timed out"); + sw.interface.flush() catch return fail(f, "image download failed or timed out"); + const r = &tls.reader; + // status and headers + const status_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + if (status_line.len < 12 or !std.mem.startsWith(u8, status_line, "HTTP/1.")) return fail(f, "image download failed"); + const status = std.fmt.parseInt(u16, status_line[9..12], 10) catch return fail(f, "image download failed"); + var location: ?[]const u8 = null; + var media: []const u8 = ""; + var length: ?usize = null; + var chunked = false; + var encoded = false; + var header_bytes: usize = 0; + while (true) { + const line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + header_bytes += line.len; + if (header_bytes > 64 * 1024) return fail(f, "image download failed"); + const l = std.mem.trimEnd(u8, line, "\r\n"); + if (l.len == 0) break; + const colon = std.mem.indexOfScalar(u8, l, ':') orelse continue; + const key = std.mem.trim(u8, l[0..colon], " \t"); + const val = std.mem.trim(u8, l[colon + 1 ..], " \t"); + if (std.ascii.eqlIgnoreCase(key, "location")) location = try a.dupe(u8, val); + if (std.ascii.eqlIgnoreCase(key, "content-type")) { + const semi = std.mem.indexOfScalar(u8, val, ';') orelse val.len; + media = try std.ascii.allocLowerString(a, std.mem.trim(u8, val[0..semi], " \t")); + } + if (std.ascii.eqlIgnoreCase(key, "content-length")) { + for (val) |c| if (!std.ascii.isDigit(c)) return fail(f, "image response exceeds the encoded byte limit"); + length = std.fmt.parseInt(usize, val, 10) catch return fail(f, "image response exceeds the encoded byte limit"); + } + if (std.ascii.eqlIgnoreCase(key, "transfer-encoding")) chunked = std.mem.indexOf(u8, try std.ascii.allocLowerString(a, val), "chunked") != null; + if (std.ascii.eqlIgnoreCase(key, "content-encoding") and !std.ascii.eqlIgnoreCase(val, "identity")) encoded = true; + } + switch (status) { + 301, 302, 303, 307, 308 => return .{ .redirect = location orelse return fail(f, "image redirect has no destination") }, + 200 => {}, + else => return fail(f, try std.fmt.allocPrint(a, "image download returned HTTP {d}", .{status})), + } + if (encoded) return fail(f, "compressed HTTP image responses are unsupported"); + for (media_types) |m| { + if (std.mem.eql(u8, m, media)) break; + } else return fail(f, if (std.mem.startsWith(u8, media_types[0], "image/")) "image URL content type must be JPEG, PNG or WebP" else try std.fmt.allocPrint(a, "media URL content type must be one of {s}", .{accept})); + if (length) |n| if (n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + var data: std.ArrayList(u8) = .empty; + if (chunked) { + while (true) { + _ = try remaining(io, deadline, f); + const size_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + const semi = std.mem.indexOfScalar(u8, size_line, ';') orelse size_line.len; + const hex = std.mem.trim(u8, size_line[0..semi], " \t\r\n"); + const n = std.fmt.parseInt(usize, hex, 16) catch return fail(f, "image download failed"); + if (n == 0) break; + if (data.items.len + n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + const piece = r.take(n) catch return fail(f, "image download failed or timed out"); + try data.appendSlice(a, piece); + _ = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed"); + } + } else { + while (length == null or data.items.len < length.?) { + _ = try remaining(io, deadline, f); + const want = if (length) |n| @min(n - data.items.len, 64 * 1024) else 64 * 1024; + const piece = r.peekGreedy(1) catch |e| switch (e) { + error.EndOfStream => break, + else => return fail(f, "image download failed or timed out"), + }; + const take = @min(piece.len, want); + if (data.items.len + take > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + try data.appendSlice(a, piece[0..take]); + r.toss(take); + } + if (length) |n| if (data.items.len != n) return fail(f, "image download failed or timed out"); + } + _ = try remaining(io, deadline, f); + return .{ .body = .{ .data = data.items, .media = media } }; +} + +test "public addresses as Python's ipaddress checks them" { + const ip4 = struct { + fn of(b: [4]u8) net.IpAddress { + return .{ .ip4 = .{ .bytes = b, .port = 443 } }; + } + }.of; + try std.testing.expect(publicIp(ip4(.{ 8, 8, 8, 8 }))); + try std.testing.expect(publicIp(ip4(.{ 151, 101, 1, 69 }))); + for ([_][4]u8{ .{ 10, 1, 2, 3 }, .{ 127, 0, 0, 1 }, .{ 169, 254, 169, 254 }, .{ 172, 20, 0, 1 }, .{ 192, 168, 1, 1 }, .{ 100, 100, 0, 1 }, .{ 0, 0, 0, 0 }, .{ 224, 0, 0, 1 }, .{ 255, 255, 255, 255 }, .{ 168, 63, 129, 16 }, .{ 192, 0, 0, 9 } }) |b| + try std.testing.expect(!publicIp(ip4(b))); +} + +test "URLs as images_http._url takes them" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Failure = .{}; + try std.testing.expectError(error.Media, parseUrl(a, "https://example.com/a b/c.png", &f)); + const ok = try parseUrl(a, "https://example.com:443/img/c.png?x=%20&y=ä", &f); + try std.testing.expectEqualStrings("example.com", ok.host); + try std.testing.expectEqualStrings("/img/c.png?x=%20&y=%C3%A4", ok.target); + for ([_][]const u8{ "http://example.com/x.png", "https://user@example.com/x", "https://example.com:8443/x", "https://localhost/x", "https://example.com/x#frag", "https://[::1]/x" }) |bad| + try std.testing.expectError(error.Media, parseUrl(a, bad, &f)); + try std.testing.expectEqualStrings("https://e.com/b/c.png", try join(a, "https://e.com/b/a.png", "c.png")); + try std.testing.expectEqualStrings("https://e.com/z", try join(a, "https://e.com/b/a.png", "/z")); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..aaf23610ae 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -127,7 +127,12 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (srv.config.request_log) |path| request_log.append(cx.a, path, body); var input: chat.Input = .{ .fields = undefined }; if (is_chat) { - input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + var msgs = body.get("messages"); + if (try @import("vision_inputs.zig").extract(srv, cx, msgs)) |x| { // images: decoded, prepared, their parts as text + msgs = x.messages; + input.images = x.images; + } + input.messages = try messages.normalize(cx, msgs, "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..6b7faaf8e2 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -39,6 +39,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t error.OutOfMemory => return error.OutOfMemory, error.Template => return cx.fail(.server, "{s}", .{problem}), }; + if (!thinking) ids = try withoutEffortLine(srv, cx, ids); if (!thinking and generation and ids.len > 0) { const last = try srv.text.decode(cx.a, ids[ids.len - 1 ..]); if (std.mem.eql(u8, @import("reply_text.zig").pyStrip(last), "")) { @@ -52,6 +53,19 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } +/// GLM-5.3's template writes a reasoning-effort system line whatever the request asks; with thinking off the Python +/// frontend renders it without (glm5_next/prompts.py thinking_off: no effort line, an empty think block). The line is +/// its own segment between special tokens, so its ids alone are the ids it renders to; the first is dropped. +fn withoutEffortLine(srv: *Server, cx: *Cx, ids: []u32) errors.Refused![]u32 { + const line = srv.text.encode(cx.a, "<|system|>Reasoning Effort: Max", false) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return ids, + }; + if (line.len < 2) return ids; + const at = std.mem.indexOf(u32, ids, line) orelse return ids; + return try std.mem.concat(cx.a, u32, &.{ ids[0..at], ids[at + line.len ..] }); +} + /// The template's generation prompt (thinking on or off): what a probe conversation renders past its history. fn generationPrompt(srv: *Server, cx: *Cx, thinking: bool) errors.Refused![]const u32 { const user = try json.newObject(cx.a); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..1bebd524e3 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -109,6 +109,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .model_dir = s.model_dir, .model_ids = ids.items, .default_max_tokens = args.max_tokens, + .vision_max_images = @intCast(@max(args.vision_max_images, 1)), + .vision_urls = args.vision_urls, + .vision_image_tokens = @intCast(std.math.clamp(args.vision_image_tokens, 1, 65536)), + .tool_system = if (args.tool_system) |t| (if (std.mem.trim(u8, t, " \t\r\n").len > 0) std.mem.trim(u8, t, " \t\r\n") else null) else null, .enable_thinking = args.thinking, .reasoning_effort = args.reasoning_effort, .thinking_budget = args.thinking_budget, diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..0407030322 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -36,6 +36,10 @@ pub const Config = struct { /// The served name first, then aliases, without repeats. model_ids: []const []const u8, default_max_tokens: i64 = 4096, + tool_system: ?[]const u8 = null, // --tool-system (trimmed; null when empty) + vision_max_images: usize = 4, // --vision-max-images + vision_urls: bool = false, // --vision-urls: public HTTPS media URLs are fetched (media_fetch.zig) + vision_image_tokens: u64 = 4096, // --vision-image-tokens: the visual tokens a request's images share enable_thinking: bool = true, reasoning_effort: ?[]const u8 = null, thinking_budget: i64 = 0, @@ -135,6 +139,7 @@ pub const Server = struct { if (srv.keepalive) |k| k.stop(); // before the engine's queue goes away if (srv.grammars) |g| g.close(); srv.metrics.deinit(); + if (srv.grammars) |g| g.close(); srv.arena.deinit(); srv.store.deinit(); srv.slide.deinit(); diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig new file mode 100644 index 0000000000..fe457dfb34 --- /dev/null +++ b/zig/src/server/vision_inputs.zig @@ -0,0 +1,342 @@ +//! Image and video inputs of a chat request, as the Python server's frontend takes them (tensorfold/vision/images.py, +//! videos.py, qwen_processing.py): each ``image_url`` part (a data: URL; JPEG or PNG) decoded and prepared, each +//! ``video_url`` part (a data: URL; MP4, WebM...) decoded by libtfvideo and cut into frame groups, and each replaced in +//! its message by the text Python's frontend expands it to: ``<|vision_start|>`` + one ``<|image_pad|>`` a visual +//! token + ``<|vision_end|>`` for an image, and ``<|vision_start|>`` + ``<|video_pad|>``s + ``<|vision_end|>`` +//! for each frame group of a video. The template copies the text, and one tokenization gives Python's ids. + +const std = @import("std"); +const media_fetch = @import("media_fetch.zig"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const Server = @import("server.zig").Server; +const Value = json.Value; +const Cx = errors.Cx; +const qi = api.qwen_image; + +const vision_start = "<|vision_start|>"; +const vision_end = "<|vision_end|>"; +const max_video_bytes: usize = 16 * 1024 * 1024; +const max_videos = 2; +const max_bytes_one: usize = 10 * 1024 * 1024; +const max_bytes_all: usize = 20 * 1024 * 1024; +const max_pixels_all: u64 = 32 * 1024 * 1024; +const max_tokens_one: u64 = 4096; + +pub const Extracted = struct { messages: Value, images: []api.Image }; + +/// The image URL of a content part (an OpenAI ``image_url`` part, or ``image`` / ``input_image`` spellings) and its +/// ``detail``; null for a part that is not an image. +const Kind = enum { image, video }; + +fn mediaPart(part: Value) ?struct { kind: Kind, url: ?[]const u8, detail: []const u8 } { + if (part != .object) return null; + const ty = part.get("type"); + const is_video = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "video_url") or std.mem.eql(u8, ty.?.string, "video"))) or + part.get("video_url") != null; + if (is_video) { + var vurl: ?[]const u8 = null; + if (part.get("video_url")) |vu| switch (vu) { + .string => |x| vurl = x, + .object => if (vu.get("url")) |u| if (u == .string) { + vurl = u.string; + }, + else => {}, + }; + if (vurl == null) if (part.get("video")) |v| if (v == .string) { + vurl = v.string; + }; + return .{ .kind = .video, .url = vurl, .detail = "auto" }; + } + const is_image = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "image_url") or std.mem.eql(u8, ty.?.string, "image") or + std.mem.eql(u8, ty.?.string, "input_image"))) or part.get("image_url") != null or part.get("image") != null; + if (!is_image) return null; + var url: ?[]const u8 = null; + var detail: []const u8 = "auto"; + if (part.get("image_url")) |iu| switch (iu) { + .string => |s| url = s, + .object => { + if (iu.get("url")) |u| if (u == .string) { + url = u.string; + }; + if (iu.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + }, + else => {}, + }; + if (url == null) if (part.get("image")) |im| if (im == .string) { + url = im.string; + }; + if (part.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + return .{ .kind = .image, .url = url, .detail = detail }; +} + +fn hasImages(messages: Value) bool { + if (messages != .array) return false; + for (messages.array) |m| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + for (c.array) |p| if (mediaPart(p) != null) return true; + } + return false; +} + +const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; +const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; + +/// The bytes of a ``data:;base64,`` URL, or (--vision-urls) of a public HTTPS URL, fetched before +/// ``deadline`` (images.py _check_source and load_images). +fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, deadline: std.Io.Timestamp) errors.Refused![]u8 { + if (!std.mem.startsWith(u8, url, "data:")) { + if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse(if (kind == .image) "images require data URLs or public HTTPS URLs" else "videos require data URLs or public HTTPS URLs"); + if (!srv.config.vision_urls) return cx.refuse("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls"); + if (url.len > media_fetch.max_url_chars) return cx.refuse("image URL is too long"); + var f: media_fetch.Failure = .{}; + const got = media_fetch.fetch(cx.a, srv.io, url, limit, deadline, if (kind == .image) &image_media else &video_media, &f) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Media => return cx.fail(.request, "{s}", .{f.text}), + }; + return got.data; + } + const comma = std.mem.indexOfScalar(u8, url, ',') orelse return cx.refuse("an image data: URL needs a comma before its bytes"); + if (!std.mem.endsWith(u8, url[0..comma], ";base64")) return cx.refuse("an image data: URL must be base64-encoded"); + const b64 = url[comma + 1 ..]; + const dec = std.base64.standard.Decoder; + const n = dec.calcSizeForSlice(b64) catch return cx.refuse("an image's base64 bytes are invalid"); + if (n > limit) return cx.refuse(if (kind == .image) "an image exceeds the 10 MiB limit (20 MiB for a request's media together)" else "a video exceeds the 16 MiB limit (20 MiB for a request's media together)"); + const out = try cx.a.alloc(u8, n); + dec.decode(out, b64) catch return cx.refuse("an image's base64 bytes are invalid"); + return out; +} + +/// A request's images extracted and prepared, its messages with their image parts as the template's image text; +/// null when it has none. A server without --vision refuses images. +pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { + const list = messages orelse return null; + if (!hasImages(list)) return null; + if (!srv.info.vision) return cx.refuse("image and video inputs need a server started with --vision"); + // every media part, in prompt order, with where its text goes + const Slot = struct { kind: Kind, url: ?[]const u8, detail: []const u8, msg: usize, part: usize }; + var slots: std.ArrayList(Slot) = .empty; + for (list.array, 0..) |m, i| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + const role = m.get("role"); + for (c.array, 0..) |p, j| { + const im = mediaPart(p) orelse continue; + if (role != null and role.? == .string and (std.mem.eql(u8, role.?.string, "system") or std.mem.eql(u8, role.?.string, "developer"))) + return cx.refuse("a system message cannot contain images or videos"); + try slots.append(cx.a, .{ .kind = im.kind, .url = im.url, .detail = im.detail, .msg = i, .part = j }); + } + } + var n_images: usize = 0; + var n_videos: usize = 0; + for (slots.items) |sl| switch (sl.kind) { + .image => n_images += 1, + .video => n_videos += 1, + }; + if (n_images > srv.config.vision_max_images) return cx.fail(.request, "a request may carry at most {d} images; this one has {d}", .{ srv.config.vision_max_images, n_images }); + if (n_videos > max_videos) return cx.fail(.request, "a request may carry at most {d} videos; this one has {d}", .{ max_videos, n_videos }); + const budget: u64 = srv.config.vision_image_tokens; + if (n_images > 0 and budget < n_images) return cx.refuse("the image count exceeds the visual-token budget"); + const per: u64 = if (n_images > 0) @min(budget / n_images, max_tokens_one) else 0; + var images: std.ArrayList(api.Image) = .empty; + const texts = try cx.a.alloc([]const u8, slots.items.len); + var total_bytes: usize = 0; + var total_pixels: u64 = 0; + const deadline = std.Io.Clock.awake.now(srv.io).addDuration(.fromSeconds(media_fetch.total_s)); // every URL's download + for (slots.items, 0..) |sl, k| { + const url = sl.url orelse return cx.refuse(if (sl.kind == .image) "an image part needs an image_url with a url" else "a video part needs a video_url with a url"); + const bytes = try dataBytes(srv, cx, url, @min(if (sl.kind == .image) max_bytes_one else max_video_bytes, max_bytes_all - total_bytes), sl.kind, deadline); + total_bytes += bytes.len; + if (total_bytes > max_bytes_all) return cx.refuse("a request's images and videos exceed the 20 MiB limit"); + var text: std.ArrayList(u8) = .empty; + switch (sl.kind) { + .image => { + var img = qi.decode(cx.a, bytes, .{}) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.ImageTooLarge => return cx.refuse("image dimensions exceed the decoded pixel limit"), + else => return cx.refuse("image bytes are invalid or unsupported; use JPEG or PNG"), + }; + total_pixels += @as(u64, img.w) * img.h; + if (total_pixels > max_pixels_all) return cx.refuse("a request's images exceed the decoded pixel limit"); + const cap = if (std.mem.eql(u8, sl.detail, "low")) @min(per, 256) else per; + if (srv.info.vision_glm) { // GLM: its canvas within [min(16, cap), cap] tokens, <|image|> placeholders + const p = qi.glm.prepare(cx.a, img, @min(16, cap), cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the image does not fit the visual-token budget"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + texts[k] = text.items; + continue; + } + const p = qi.prepare(cx.a, img, cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("absolute aspect ratio must be smaller than 200"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, vision_start); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image_pad|>"); + try text.appendSlice(cx.a, vision_end); + }, + .video => if (srv.info.vision_glm) { + // GLM: <|begin_of_video|>, each frame group as <|begin_of_image|> <|image|>s <|end_of_image|> and its + // first frame's "t.t seconds", <|end_of_video|> (mlx-vlm's glm5_next processor) + var v = try videoGlm(srv, cx, bytes); + defer v.deinit(cx.a); + try text.appendSlice(cx.a, "<|begin_of_video|>"); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds"); + } + try text.appendSlice(cx.a, "<|end_of_video|>"); + } else { + var v = try video(srv, cx, bytes); + defer v.deinit(cx.a); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, "<"); + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds>"); + try text.appendSlice(cx.a, vision_start); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|video_pad|>"); + try text.appendSlice(cx.a, vision_end); + } + }, + } + texts[k] = text.items; + } + // the messages with each media part as its text + const out = try cx.a.alloc(Value, list.array.len); + @memcpy(out, list.array); + var k: usize = 0; + while (k < slots.items.len) { + const i = slots.items[k].msg; + const c = list.array[i].get("content").?; + const parts = try cx.a.dupe(Value, c.array); + while (k < slots.items.len and slots.items[k].msg == i) : (k += 1) { + const t = try json.newObject(cx.a); + try t.put(cx.a, "type", .{ .string = "text" }); + try t.put(cx.a, "text", .{ .string = texts[k] }); + parts[slots.items[k].part] = .{ .object = t }; + } + const copy = try json.copyObject(cx.a, list.array[i].object); + try copy.put(cx.a, "content", .{ .array = parts }); + out[i] = .{ .object = copy }; + } + return .{ .messages = .{ .array = out }, .images = images.items }; +} + +// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives -------------------------------------------------- + +const TfVideo = extern struct { rgb: ?[*]u8, frames: c_int, height: c_int, width: c_int, indices: ?[*]c_int, fps: f64 }; +const SizeFn = *const fn (ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int; +const DecodeFn = *const fn (data: [*]const u8, len: usize, rate: f64, min_frames: c_int, max_frames: c_int, max_seconds: f64, max_dim: c_int, size: SizeFn, ctx: ?*anyopaque, out: *TfVideo) callconv(.c) c_int; +const FreeFn = *const fn (v: *TfVideo) callconv(.c) void; + +var lib_mutex: std.Io.Mutex = .init; +var lib: ?struct { decode: DecodeFn, free: FreeFn } = null; + +fn library(io: std.Io) ?@TypeOf(lib.?) { + lib_mutex.lockUncancelable(io); + defer lib_mutex.unlock(io); + if (lib == null) { + var dl = std.DynLib.open("libtfvideo.so") catch return null; + const d = dl.lookup(DecodeFn, "tf_video_decode") orelse return null; + const f = dl.lookup(FreeFn, "tf_video_free") orelse return null; + lib = .{ .decode = d, .free = f }; // kept open for the server's life + } + return lib; +} + +fn sizeFor(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const budget: *const u64 = @ptrCast(@alignCast(ctx.?)); + const s = qi.videoSize(@intCast(frames), @intCast(height), @intCast(width), budget.*) catch return -1; + oh.* = @intCast(s[0]); + ow.* = @intCast(s[1]); + return 0; +} + +/// GLM's video canvas: the frames' content size for libtfvideo (geometry within the budget), the canvas kept. +const GlmSize = struct { budget: u64, th: u32 = 0, tw: u32 = 0 }; + +fn sizeForGlm(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const gs: *GlmSize = @ptrCast(@alignCast(ctx.?)); + const g = qi.glm.geometry(@intCast(frames), @intCast(height), @intCast(width), 16, gs.budget) catch return -1; + gs.th = g[0]; + gs.tw = g[1]; + oh.* = @intCast(g[2]); + ow.* = @intCast(g[3]); + return 0; +} + +fn videoGlm(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.glm.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var gs: GlmSize = .{ .budget = 16384 }; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| gs.budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch gs.budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 2, 256, 3600.0, 8192, sizeForGlm, @ptrCast(&gs), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.glm.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, gs.th, gs.tw, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} + +fn video(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var budget: u64 = 16384; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 4, 256, 3600.0, 8192, sizeFor, @ptrCast(&budget), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} diff --git a/zig/src/vision/glm_image.zig b/zig/src/vision/glm_image.zig new file mode 100644 index 0000000000..b46e75f076 --- /dev/null +++ b/zig/src/vision/glm_image.zig @@ -0,0 +1,194 @@ +//! GLM-5.3-Flash's image and video preparation (mlx-vlm's glm5_next/processing.py, as TensorFold's GLM frontend runs +//! it): the canvas (smart_resize to multiples of 28 within a visual-token budget), the content resized into it (PIL's +//! antialiased bicubic) and zero-padded at the bottom and right, CLIP's normalization (x * f32(1/255) - mean) / std, +//! and the patches in merge-block order ([c, t, 14, 14] each; an image's one frame twice as the temporal pair). + +const std = @import("std"); +const qi = @import("qwen_image.zig"); + +pub const PATCH = 14; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; +const MEAN = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +const STD = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; +const PER = 3 * TEMPORAL * PATCH * PATCH; + +pub const Error = qi.Error; + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const f = @floor(x); + const d = x - f; + if (d > 0.5) return f + 1; + if (d < 0.5) return f; + return if (@mod(f, 2) == 0) f else f + 1; +} + +fn alignUp(v: f64, f: f64) f64 { + return @ceil(v / f) * f; +} + +/// smart_resize: the canvas (height, width) of ``frames`` frames within [min_tokens, max_tokens] merged tokens. +pub fn smartResize(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![2]u32 { + if (frames == 0 or height == 0 or width == 0 or min_tokens == 0 or min_tokens > max_tokens) return error.UnsupportedImage; + const f: f64 = FACTOR; + const ppt: f64 = TEMPORAL * f * f; + const min_px = @as(f64, @floatFromInt(min_tokens)) * ppt; + const max_px = @as(f64, @floatFromInt(max_tokens)) * ppt; + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const n: f64 = @floatFromInt(frames); + const af = @max(@as(f64, TEMPORAL), pyRound(n / TEMPORAL) * TEMPORAL); + var ah = alignUp(h, f); + var aw = alignUp(w, f); + var px = af * ah * aw; + if (px < min_px) { + const scale = @sqrt(min_px / (n * h * w)); + ah = alignUp(@max(1, @ceil(h * scale)), f); + aw = alignUp(@max(1, @ceil(w * scale)), f); + px = af * ah * aw; + } + if (px > max_px) { // the tallest content within the budget, by bisection on its height + if (max_px < af * f * f) return error.UnsupportedImage; + var low: f64 = 1; + var high: f64 = h; + var best: [2]f64 = .{ f, f }; + while (low <= high) { + const ch = @floor((low + high) / 2); + const cw = @max(1, @floor(w * ch / h)); + const cand = [2]f64{ alignUp(ch, f), alignUp(cw, f) }; + if (af * cand[0] * cand[1] <= max_px) { + best = cand; + low = ch + 1; + } else high = ch - 1; + } + ah = best[0]; + aw = best[1]; + } + return .{ @intFromFloat(ah), @intFromFloat(aw) }; +} + +/// _resize_geometry: the canvas (th, tw) and the content (ch, cw) resized into its top-left corner. +pub fn geometry(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![4]u32 { + const t = try smartResize(frames, height, width, min_tokens, max_tokens); + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + var scale = @min(@as(f64, @floatFromInt(t[0])) / h, @as(f64, @floatFromInt(t[1])) / w); + const ppt: f64 = TEMPORAL * FACTOR * FACTOR; + if (@as(f64, @floatFromInt(frames)) * h * w >= @as(f64, @floatFromInt(min_tokens)) * ppt) scale = @min(1.0, scale); + const ch: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[0])), @floor(h * scale)))); + const cw: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[1])), @floor(w * scale)))); + return .{ t[0], t[1], ch, cw }; +} + +pub const Prepared = struct { + patches: []f32, // [gh * gw, 1176] + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], MERGE * MERGE)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +fn norm(v: u8, c: usize) f32 { + const x = @as(f32, @floatFromInt(v)) * @as(f32, @floatCast(@as(f64, 1.0) / 255.0)); + return (x - MEAN[c]) / STD[c]; +} + +/// A frame's normalized value at canvas (y, x) (zero padding past the content's ch x cw). +fn at(frame: []const u8, cw: usize, ch: usize, y: usize, x: usize, c: usize) f32 { + return norm(if (y < ch and x < cw) frame[(y * cw + x) * 3 + c] else 0, c); +} + +/// ``im`` on its canvas within [min_tokens, max_tokens] tokens, normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: qi.Image, min_tokens: u64, max_tokens: u64) Error!Prepared { + const g = try geometry(TEMPORAL, im.h, im.w, min_tokens, max_tokens); + var r = if (g[2] == im.h and g[3] == im.w) null else try qi.resize(gpa, im, g[3], g[2]); + defer if (r) |*x| x.deinit(gpa); + const src = if (r) |x| x.rgb else im.rgb; + const gh: usize = g[0] / PATCH; + const gw: usize = g[1] / PATCH; + const out = try gpa.alloc(f32, gh * gw * PER); + var n: usize = 0; + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..PATCH) |y| for (0..PATCH) |x| { + const v = at(src, g[3], g[2], py + y, px + x, c); + for (0..TEMPORAL) |t| dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = v; + }; + n += 1; + }; + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +/// A video's frame groups: each group's patches (two consecutive frames as the temporal pair; an odd last frame +/// repeats) and its time (its first frame's seconds, as the processor's timestamps[::2]). +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1176 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * PER; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames at the content size (ch x cw: geometry's), placed on the th x tw canvas. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, ch: usize, cw: usize, th: usize, tw: usize, indices: []const i32, fps: f64) Error!Video { + if (th % FACTOR != 0 or tw % FACTOR != 0 or ch > th or cw > tw or frames < 1) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = th / PATCH; + const gw = tw / PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * PER); + errdefer gpa.free(out); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * ch * cw * 3 ..][0 .. ch * cw * 3]; + for (0..PATCH) |y| for (0..PATCH) |x| { + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = at(frame, cw, ch, py + y, px + x, c); + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| times[g] = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +test "smart_resize and geometry match the processor's" { + // 640 x 480 within 8000 tokens: aligned to 504 x 644, content unscaled + try std.testing.expectEqual([4]u32{ 504, 644, 480, 640 }, try geometry(TEMPORAL, 480, 640, 16, 8000)); + // 40 x 30 below 16 tokens: scaled up, content scaled + const g = try geometry(TEMPORAL, 30, 40, 16, 8000); + try std.testing.expectEqual(@as(u32, 112), g[0]); + try std.testing.expectEqual(@as(u32, 140), g[1]); + // 300 x 900 (tall) + const t = try geometry(TEMPORAL, 900, 300, 16, 8000); + try std.testing.expectEqual([2]u32{ 924, 308 }, .{ t[0], t[1] }); +} diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig new file mode 100644 index 0000000000..f918651811 --- /dev/null +++ b/zig/src/vision/qwen_image.zig @@ -0,0 +1,432 @@ +//! Image inputs for Qwen-style vision towers, as the Python engine prepares them (tensorfold/vision/images.py and +//! qwen_processing.py with transformers' Qwen2VL processor): decode (JPEG, PNG; EXIF orientation applied, alpha over +//! white), smart_resize to multiples of patch x merge within a pixel budget, an antialiased bicubic resize (PIL's +//! filter and bounds with torch's uint8 fixed point: int16 weights at the most precision the axis allows), +//! (x - 127.5) / 127.5, and the patches in merge-block order, +//! each [channel, frame (the image twice), 16, 16] -- 1,536 values. + +const std = @import("std"); +/// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. +pub const glm = @import("glm_image.zig"); + +extern fn stbi_info_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int) c_int; +extern fn stbi_load_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int, req: c_int) ?[*]u8; +extern fn stbi_image_free(p: ?*anyopaque) void; + +pub const Limits = struct { + max_dimension: u32 = 8192, + max_pixels: u64 = 16 * 1024 * 1024, +}; + +/// An RGB image, rows top to bottom, three bytes a pixel. +pub const Image = struct { + w: u32, + h: u32, + rgb: []u8, + + pub fn deinit(im: *Image, gpa: std.mem.Allocator) void { + gpa.free(im.rgb); + } +}; + +pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio }; + +/// JPEG or PNG bytes to RGB: EXIF orientation applied (JPEG), any alpha composited over white. +pub fn decode(gpa: std.mem.Allocator, bytes: []const u8, limits: Limits) Error!Image { + if (bytes.len == 0 or bytes.len > std.math.maxInt(c_int)) return error.UnsupportedImage; + var w: c_int = 0; + var h: c_int = 0; + var comp: c_int = 0; + if (stbi_info_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp) == 0) return error.UnsupportedImage; + if (w <= 0 or h <= 0 or w > limits.max_dimension or h > limits.max_dimension or + @as(u64, @intCast(w)) * @as(u64, @intCast(h)) > limits.max_pixels) return error.ImageTooLarge; + const alpha = comp == 2 or comp == 4; + const want: c_int = if (alpha) 4 else 3; + var got: c_int = 0; + const px = stbi_load_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp, want) orelse return error.UnsupportedImage; + defer stbi_image_free(px); + const n: usize = @intCast(w * h); + _ = &got; + var rgb = try gpa.alloc(u8, n * 3); + errdefer gpa.free(rgb); + if (alpha) { + // PIL's paste of RGBA onto white with the alpha as mask: out = white + (src - white) * a / 255, rounded + for (0..n) |i| { + const a: i32 = px[i * 4 + 3]; + for (0..3) |c| { + const s: i32 = px[i * 4 + c]; + const t: i32 = (s - 255) * a + 128; + rgb[i * 3 + c] = @intCast(255 + ((t + (t >> 8)) >> 8)); + } + } + } else @memcpy(rgb, px[0 .. n * 3]); + var im: Image = .{ .w = @intCast(w), .h = @intCast(h), .rgb = rgb }; + const o = orientation(bytes); + if (o > 1) { + const t = try transpose(gpa, im, o); + im.deinit(gpa); + im = t; + } + return im; +} + +/// A JPEG's EXIF orientation (1-8; 1 when absent or not a JPEG). +fn orientation(b: []const u8) u16 { + if (b.len < 4 or b[0] != 0xFF or b[1] != 0xD8) return 1; + var i: usize = 2; + while (i + 4 <= b.len and b[i] == 0xFF) { + const marker = b[i + 1]; + const len = (@as(usize, b[i + 2]) << 8) | b[i + 3]; + if (marker == 0xDA or len < 2 or i + 2 + len > b.len) break; + if (marker == 0xE1 and len >= 16 and std.mem.eql(u8, b[i + 4 .. i + 10], "Exif\x00\x00")) { + const t = b[i + 10 .. i + 2 + len]; + const le = t[0] == 'I'; + const rd16 = struct { + fn f(x: []const u8, at: usize, little: bool) u16 { + return if (little) std.mem.readInt(u16, x[at..][0..2], .little) else std.mem.readInt(u16, x[at..][0..2], .big); + } + }.f; + const rd32 = struct { + fn f(x: []const u8, at: usize, little: bool) u32 { + return if (little) std.mem.readInt(u32, x[at..][0..4], .little) else std.mem.readInt(u32, x[at..][0..4], .big); + } + }.f; + if (t.len < 8) return 1; + const ifd = rd32(t, 4, le); + if (ifd + 2 > t.len) return 1; + const count = rd16(t, ifd, le); + for (0..count) |e| { + const at = ifd + 2 + e * 12; + if (at + 12 > t.len) return 1; + if (rd16(t, at, le) == 0x0112) { + const v = rd16(t, at + 8, le); + return if (v >= 1 and v <= 8) v else 1; + } + } + return 1; + } + i += 2 + len; + } + return 1; +} + +/// PIL's exif_transpose for orientations 2-8. +fn transpose(gpa: std.mem.Allocator, im: Image, o: u16) Error!Image { + const swap = o >= 5; + const W: u32 = if (swap) im.h else im.w; + const H: u32 = if (swap) im.w else im.h; + const out = try gpa.alloc(u8, im.rgb.len); + for (0..H) |y| for (0..W) |x| { + // the source pixel of output (x, y) + const sx: u32, const sy: u32 = switch (o) { + 2 => .{ im.w - 1 - @as(u32, @intCast(x)), @intCast(y) }, + 3 => .{ im.w - 1 - @as(u32, @intCast(x)), im.h - 1 - @as(u32, @intCast(y)) }, + 4 => .{ @intCast(x), im.h - 1 - @as(u32, @intCast(y)) }, + 5 => .{ @intCast(y), @intCast(x) }, + 6 => .{ @intCast(y), im.h - 1 - @as(u32, @intCast(x)) }, + 7 => .{ im.w - 1 - @as(u32, @intCast(y)), im.h - 1 - @as(u32, @intCast(x)) }, + 8 => .{ im.w - 1 - @as(u32, @intCast(y)), @intCast(x) }, + else => .{ @intCast(x), @intCast(y) }, + }; + @memcpy(out[(y * W + x) * 3 ..][0..3], im.rgb[(sy * im.w + sx) * 3 ..][0..3]); + }; + return .{ .w = W, .h = H, .rgb = out }; +} + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const r = @round(x); + if (@abs(x - @trunc(x)) == 0.5) return 2.0 * @round(x / 2.0); + return r; +} + +/// transformers' smart_resize: both sides multiples of ``factor``, the area within [min_pixels, max_pixels]. +pub fn smartResize(height: u32, width: u32, factor: u32, min_pixels: u64, max_pixels: u64) Error![2]u32 { + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const f: f64 = @floatFromInt(factor); + if (@max(h, w) / @min(h, w) > 200) return error.BadAspectRatio; + var hb = pyRound(h / f) * f; + var wb = pyRound(w / f) * f; + if (hb * wb > @as(f64, @floatFromInt(max_pixels))) { + const beta = @sqrt(h * w / @as(f64, @floatFromInt(max_pixels))); + hb = @max(f, @floor(h / beta / f) * f); + wb = @max(f, @floor(w / beta / f) * f); + } else if (hb * wb < @as(f64, @floatFromInt(min_pixels))) { + const beta = @sqrt(@as(f64, @floatFromInt(min_pixels)) / (h * w)); + hb = @ceil(h * beta / f) * f; + wb = @ceil(w * beta / f) * f; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + + +fn bicubic(x0: f64) f64 { + const a = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1; + if (x < 2.0) return (((x - 5) * x + 8) * x - 4) * a; + return 0.0; +} + +const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; + +/// PIL's precompute_coeffs for an axis of ``in_size`` resampled to ``out_size``, as fixed point the way torch's uint8 +/// antialias kernel takes them: int16 weights at the largest precision (< 22 bits) that keeps the largest in range. +fn coeffs(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { + const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); + const filterscale = @max(scale, 1.0); + const support = 2.0 * filterscale; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const bounds = try gpa.alloc(usize, out_size * 2); + const kk = try gpa.alloc(i32, out_size * ksize); + const pre_all = try gpa.alloc(f64, out_size * ksize); + defer gpa.free(pre_all); + for (0..out_size) |xx| { + const pre = pre_all[xx * ksize ..][0..ksize]; + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + var ww: f64 = 0; + const ss = 1.0 / filterscale; + var xmin_i: i64 = @intFromFloat(center - support + 0.5); + if (xmin_i < 0) xmin_i = 0; + var xmax_i: i64 = @intFromFloat(center + support + 0.5); + if (xmax_i > @as(i64, @intCast(in_size))) xmax_i = @intCast(in_size); + const xmin: usize = @intCast(xmin_i); + const xmax: usize = @intCast(xmax_i - xmin_i); + @memset(pre, 0); + for (0..xmax) |x| { + const wgt = bicubic((@as(f64, @floatFromInt(x + xmin)) - center + 0.5) * ss); + pre[x] = wgt; + ww += wgt; + } + for (0..xmax) |x| { + if (ww != 0.0) pre[x] /= ww; + } + bounds[xx * 2] = xmin; + bounds[xx * 2 + 1] = xmax; + } + var w_max: f64 = 0; + for (pre_all) |v| w_max = @max(w_max, @abs(v)); + var precision: u5 = 0; + while (precision < 22) : (precision += 1) { + const next: i64 = @intFromFloat(0.5 + w_max * @as(f64, @floatFromInt(@as(i64, 1) << (precision + 1)))); + if (next >= (1 << 15)) break; + } + const one: f64 = @floatFromInt(@as(i64, 1) << precision); + for (pre_all, 0..) |v, i| kk[i] = @intFromFloat(if (v < 0) -0.5 + v * one else 0.5 + v * one); + return .{ .ksize = ksize, .bounds = bounds, .kk = kk, .precision = precision }; +} + +fn clip8(v: i64, precision: u5) u8 { + const r = v >> precision; + return @intCast(std.math.clamp(r, 0, 255)); +} + +/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits: the horizontal pass, then the vertical, each +/// rounding to bytes). +pub fn resize(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error!Image { + var cur = im; + var owned = false; + defer if (owned) cur.deinit(gpa); + if (ow != im.w) { + const c = try coeffs(gpa, im.w, ow); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, ow) * im.h * 3); + for (0..im.h) |y| for (0..ow) |x| { + const xmin = c.bounds[x * 2]; + const xmax = c.bounds[x * 2 + 1]; + const k = c.kk[x * c.ksize ..]; + for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..xmax) |j| ss += @as(i64, cur.rgb[(y * cur.w + xmin + j) * 3 + ch]) * k[j]; + out[(y * ow + x) * 3 + ch] = clip8(ss, c.precision); + } + }; + cur = .{ .w = ow, .h = im.h, .rgb = out }; + owned = true; + } + if (oh != cur.h) { + const c = try coeffs(gpa, cur.h, oh); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, cur.w) * oh * 3); + for (0..oh) |y| { + const ymin = c.bounds[y * 2]; + const ymax = c.bounds[y * 2 + 1]; + const k = c.kk[y * c.ksize ..]; + for (0..cur.w) |x| for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..ymax) |j| ss += @as(i64, cur.rgb[((ymin + j) * cur.w + x) * 3 + ch]) * k[j]; + out[(y * cur.w + x) * 3 + ch] = clip8(ss, c.precision); + }; + } + if (owned) cur.deinit(gpa); + cur = .{ .w = cur.w, .h = oh, .rgb = out }; + owned = true; + } + if (!owned) return .{ .w = im.w, .h = im.h, .rgb = try gpa.dupe(u8, im.rgb) }; + owned = false; + return cur; +} + +/// One image prepared for the tower: its patches (fp32 [gh * gw, 1536], merge-block order) and grid [1, gh, gw]. +pub const Prepared = struct { + patches: []f32, + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], 4)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +pub const PATCH = 16; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; + +/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels as the Python frontend passes it: one merged patch), +/// normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepared { + const cap = max_tokens * FACTOR * FACTOR; + const size = try smartResize(im.h, im.w, FACTOR, @min(FACTOR * FACTOR, cap), cap); + var r = try resize(gpa, im, size[1], size[0]); + defer r.deinit(gpa); + const gh: usize = size[0] / PATCH; + const gw: usize = size[1] / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, gh * gw * per); + var n: usize = 0; + var bh: usize = 0; + while (bh < gh / MERGE) : (bh += 1) { + var bw: usize = 0; + while (bw < gw / MERGE) : (bw += 1) { + for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @floatFromInt(r.rgb[((py + y) * r.w + px + x) * 3 + c]); + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 127.5) / 127.5; + }; + n += 1; + }; + } + } + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +// -- videos --------------------------------------------------------------------------------------------------------- + +/// A video's frame size for the tower (the Python frontend's video_size): Qwen3-VL's smart_resize with its per-frame +/// cap (at most 768 tokens a frame group, at least ~134), the whole video within ``budget`` tokens. +pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u32 { + const factor: f64 = FACTOR; + const shortest: f64 = 128 * FACTOR * FACTOR; + const longest: f64 = @as(f64, @floatFromInt(budget)) * TEMPORAL * FACTOR * FACTOR; + const per_frame = @max(@min(768 * FACTOR * FACTOR, @floor(longest / @as(f64, @floatFromInt(@max(1, frames))))), @floor(shortest * 1.05)); + const max_pixels = per_frame * @as(f64, @floatFromInt(frames)); + if (frames < TEMPORAL) return error.UnsupportedImage; + var height: f64 = @floatFromInt(height0); + var width: f64 = @floatFromInt(width0); + if (height < factor or width < factor) { + const scale = @max(factor / height, factor / width); + height = @trunc(height * scale); + width = @trunc(width * scale); + } + if (@max(height, width) / @min(height, width) > 200) return error.BadAspectRatio; + var hb = pyRound(height / factor) * factor; + var wb = pyRound(width / factor) * factor; + const tb = pyRound(@as(f64, @floatFromInt(frames)) / TEMPORAL) * TEMPORAL; + const n: f64 = @floatFromInt(frames); + if (tb * hb * wb > max_pixels) { + const beta = @sqrt(n * height * width / max_pixels); + hb = @max(factor, @floor(height / beta / factor) * factor); + wb = @max(factor, @floor(width / beta / factor) * factor); + } else if (tb * hb * wb < shortest) { + const beta = @sqrt(shortest / (n * height * width)); + hb = @ceil(height * beta / factor) * factor; + wb = @ceil(width * beta / factor) * factor; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + +/// A decoded video's frame groups for the tower: each group's patches (fp32 [gh * gw, 1536], two consecutive frames +/// as the temporal pair; an odd last frame repeats), normalized as numpy does ((x * f32(1/255) - 0.5) / 0.5), and +/// each group's time (the mean of its two frames' seconds). +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1536 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * 3 * TEMPORAL * PATCH * PATCH; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames of ``h`` x ``w`` (already at the tower's size), their source ``indices`` and ``fps``. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, h: usize, w: usize, indices: []const i32, fps: f64) Error!Video { + if (h % FACTOR != 0 or w % FACTOR != 0 or frames < 2) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = h / PATCH; + const gw = w / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * per); + errdefer gpa.free(out); + const scale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * h * w * 3 ..]; + for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @as(f32, @floatFromInt(frame[((py + y) * w + px + x) * 3 + c])) * scale; + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 0.5) / 0.5; + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| { + const a: f64 = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + const b: f64 = @as(f64, @floatFromInt(indices[@min(g * TEMPORAL + TEMPORAL - 1, frames - 1)])) / fps; + times[g] = (a + b) / 2; + } + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +/// Python's f"{t:.1f}": one decimal, an exact tie rounded to even. +pub fn formatSeconds(buf: []u8, t: f64) ![]const u8 { + const x: f128 = @as(f128, t) * 10; // exact: a double's 53 bits times ten fit f128's 113 + const fl = @floor(x); + const frac = x - fl; + var r = if (frac > 0.5) fl + 1 else fl; + if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; + const v: i64 = @intFromFloat(r); + return std.fmt.bufPrint(buf, "{d}.{d}", .{ @divTrunc(v, 10), @as(u64, @intCast(@mod(v, 10))) }); +} + +test { + _ = glm; +} diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c new file mode 100644 index 0000000000..1049416565 --- /dev/null +++ b/zig/src/vision/tf_video.c @@ -0,0 +1,220 @@ +// libtfvideo: a video's sampled frames for the native server's vision inputs, as the Python engine decodes them +// (tensorfold/vision/videos.py with PyAV): the first video stream; its rate (average, else guessed, else 24) and +// frame count (the stream's, else its duration times the rate); frames sampled at ``rate`` a second over the whole +// video (Qwen3-VL's rule: at least min_frames, at most max_frames, numpy's linspace rounded half to even), decoded +// in order and scaled straight to the size ``size`` picks as RGB the way PyAV's VideoFrame.to_ndarray(rgb24, width, +// height, interpolation="BICUBIC") does it (VideoReformatter: the frame's colour space and range kept, transfer and +// primaries unspecified, sws_scale_frame with SWS_BICUBIC alone). Built against FFmpeg 9 (the FFmpeg PyAV 19.0.1 +// ships, so the frames equal Python's); the server loads it when a request carries a video. +#include +#include +#include +#include +#include +#include +#include +#include + +typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); + +typedef struct { + uint8_t* rgb; // frames x height x width x 3 + int frames, height, width; + int* indices; // each frame's index in the source + double fps; // the source's rate +} tf_video; + +enum { + TF_OK = 0, + TF_INVALID = -1, // not a decodable container or stream + TF_NO_STREAM = -2, // no video stream + TF_DIMENSIONS = -3, // missing or over max_dim + TF_NO_COUNT = -4, // the frame count or rate is missing + TF_TOO_LONG = -5, // over max_seconds + TF_FEW_FRAMES = -6, // fewer than two decodable frames + TF_MEMORY = -7, + TF_SIZE = -8, // the size callback refused +}; + +typedef struct { const uint8_t* data; size_t len, pos; } mem_t; + +static int mem_read(void* o, uint8_t* buf, int n) { + mem_t* m = (mem_t*)o; + if (m->pos >= m->len) return AVERROR_EOF; + size_t k = m->len - m->pos < (size_t)n ? m->len - m->pos : (size_t)n; + memcpy(buf, m->data + m->pos, k); + m->pos += k; + return (int)k; +} + +static int64_t mem_seek(void* o, int64_t off, int whence) { + mem_t* m = (mem_t*)o; + if (whence == AVSEEK_SIZE) return (int64_t)m->len; + int64_t at = whence == SEEK_SET ? off : whence == SEEK_CUR ? (int64_t)m->pos + off : (int64_t)m->len + off; + if (at < 0 || at > (int64_t)m->len) return -1; + m->pos = (size_t)at; + return at; +} + +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): an SWS_CS_* value names the frame's colour space +// (the frame's own AVColorSpace number passes through it as PyAV passes it) +static void set_colorspace(AVFrame* f, int cs) { + if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; + else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; + else if (cs == SWS_CS_ITU601) f->colorspace = AVCOL_SPC_SMPTE170M; + else if (cs == SWS_CS_SMPTE240M) f->colorspace = AVCOL_SPC_SMPTE240M; + else if (cs == SWS_CS_BT2020) f->colorspace = AVCOL_SPC_BT2020_NCL; +} + +static double round_half_even(double x) { + double r = round(x); + if (fabs(x - trunc(x)) == 0.5) r = 2.0 * round(x / 2.0); + return r; +} + +void tf_video_free(tf_video* v) { + if (!v) return; + free(v->rgb); + free(v->indices); + v->rgb = NULL; + v->indices = NULL; +} + +int tf_video_decode(const uint8_t* data, size_t len, double rate, int min_frames, int max_frames, double max_seconds, + int max_dim, tf_size_fn size, void* ctx, tf_video* out) { + memset(out, 0, sizeof(*out)); + int rc = TF_INVALID; + mem_t mem = {data, len, 0}; + AVFormatContext* fmt = NULL; + AVIOContext* io = NULL; + AVCodecContext* dec = NULL; + AVPacket* pkt = NULL; + AVFrame* frame = NULL; + struct SwsContext* sws = NULL; + int* wanted = NULL; + unsigned char* iobuf = av_malloc(1 << 16); + if (!iobuf) return TF_MEMORY; + io = avio_alloc_context(iobuf, 1 << 16, 0, &mem, mem_read, NULL, mem_seek); + if (!io) { av_free(iobuf); return TF_MEMORY; } + fmt = avformat_alloc_context(); + if (!fmt) { rc = TF_MEMORY; goto done; } + fmt->pb = io; + if (avformat_open_input(&fmt, NULL, NULL, NULL) < 0) { fmt = NULL; goto done; } + if (avformat_find_stream_info(fmt, NULL) < 0) goto done; + int si = av_find_best_stream(fmt, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0); + if (si < 0) { rc = TF_NO_STREAM; goto done; } + // PyAV's streams.video[0] is the first video stream + for (unsigned i = 0; i < fmt->nb_streams; ++i) + if (fmt->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { si = (int)i; break; } + AVStream* st = fmt->streams[si]; + int width = st->codecpar->width, height = st->codecpar->height; + if (width <= 0 || height <= 0 || width > max_dim || height > max_dim) { rc = TF_DIMENSIONS; goto done; } + AVRational r = st->avg_frame_rate; + if (r.num <= 0 || r.den <= 0) r = av_guess_frame_rate(fmt, st, NULL); + double fps = (r.num > 0 && r.den > 0) ? av_q2d(r) : 24.0; + double seconds = 0.0; + if (st->duration > 0 && st->duration != AV_NOPTS_VALUE) seconds = (double)st->duration * av_q2d(st->time_base); + else if (fmt->duration > 0 && fmt->duration != AV_NOPTS_VALUE) seconds = (double)fmt->duration / 1e6; + int64_t total = st->nb_frames > 0 ? st->nb_frames : (int64_t)round_half_even(seconds * fps); + if (!(fps > 0 && fps <= 1000) || total <= 0) { rc = TF_NO_COUNT; goto done; } + if ((double)total / fps > max_seconds) { rc = TF_TOO_LONG; goto done; } + int64_t count = (int64_t)((double)total / fps * rate); + if (count < min_frames) count = min_frames; + if (count > max_frames) count = max_frames; + if (count > total) count = total; + wanted = malloc(sizeof(int) * (size_t)count); + if (!wanted) { rc = TF_MEMORY; goto done; } + // numpy.linspace(0, total - 1, count): step = (stop - start) / (count - 1), the last exactly stop + double stop = (double)(total - 1); + for (int64_t i = 0; i < count; ++i) { + double v = count > 1 ? (double)i * (stop / (double)(count - 1)) : 0.0; + if (count > 1 && i == count - 1) v = stop; + wanted[i] = (int)round_half_even(v); + } + int oh = 0, ow = 0; + if (size(ctx, (int)count, height, width, &oh, &ow) != 0 || oh <= 0 || ow <= 0) { rc = TF_SIZE; goto done; } + const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id); + if (!codec) goto done; + dec = avcodec_alloc_context3(codec); + if (!dec || avcodec_parameters_to_context(dec, st->codecpar) < 0) goto done; + dec->thread_count = 0; // PyAV's thread_type "AUTO" + dec->thread_type = FF_THREAD_FRAME | FF_THREAD_SLICE; + if (avcodec_open2(dec, codec, NULL) < 0) goto done; + out->rgb = malloc((size_t)count * (size_t)oh * (size_t)ow * 3); + out->indices = malloc(sizeof(int) * (size_t)count); + pkt = av_packet_alloc(); + frame = av_frame_alloc(); + if (!out->rgb || !out->indices || !pkt || !frame) { rc = TF_MEMORY; goto done; } + AVFrame* rgb = av_frame_alloc(); + if (!rgb) { rc = TF_MEMORY; goto done; } + sws = sws_alloc_context(); + if (!sws) { av_frame_free(&rgb); rc = TF_MEMORY; goto done; } + sws->threads = 0; + sws->flags = SWS_BICUBIC; + int n = 0, next = 0, got = 0, last = wanted[count - 1]; + int flushing = 0, finished = 0; + while (!finished) { + if (!flushing) { + int e = av_read_frame(fmt, pkt); + if (e < 0) { flushing = 1; avcodec_send_packet(dec, NULL); } + else { + if (pkt->stream_index == si) avcodec_send_packet(dec, pkt); + av_packet_unref(pkt); + } + } + for (;;) { + int e = avcodec_receive_frame(dec, frame); + if (e == AVERROR(EAGAIN)) { + if (flushing) finished = 1; // a flushed decoder only ends + break; + } + if (e < 0) { finished = 1; break; } + if (n > last) { finished = 1; av_frame_unref(frame); break; } + if (next < count && n == wanted[next]) { + // PyAV's _reformat: the destination copies the frame's properties; its colour space and range pass + // through _set_frame_colorspace with the frame's own values; transfer and primaries unspecified + av_frame_unref(rgb); + av_frame_copy_props(rgb, frame); + set_colorspace(frame, (int)frame->colorspace); + rgb->colorspace = frame->colorspace; + rgb->color_range = frame->color_range; + set_colorspace(rgb, (int)frame->colorspace); + const enum AVColorTransferCharacteristic trc = frame->color_trc; + const enum AVColorPrimaries pri = frame->color_primaries; + frame->color_trc = rgb->color_trc = AVCOL_TRC_UNSPECIFIED; + frame->color_primaries = rgb->color_primaries = AVCOL_PRI_UNSPECIFIED; + rgb->format = AV_PIX_FMT_RGB24; + rgb->width = ow; + rgb->height = oh; + if (av_frame_get_buffer(rgb, 0) < 0) { rc = TF_MEMORY; av_frame_unref(frame); av_frame_free(&rgb); goto done; } + int e2 = sws_scale_frame(sws, rgb, frame); + frame->color_trc = trc; + frame->color_primaries = pri; + if (e2 < 0) { av_frame_unref(frame); av_frame_free(&rgb); goto done; } + uint8_t* to = out->rgb + (size_t)got * oh * ow * 3; + for (int y = 0; y < oh; ++y) memcpy(to + (size_t)y * ow * 3, rgb->data[0] + (size_t)y * rgb->linesize[0], (size_t)ow * 3); + out->indices[got++] = n; + while (next < count && wanted[next] == n) ++next; + } + ++n; + av_frame_unref(frame); + } + } + av_frame_free(&rgb); + if (got < 2) { rc = TF_FEW_FRAMES; goto done; } + out->frames = got; + out->height = oh; + out->width = ow; + out->fps = fps; + rc = TF_OK; +done: + if (rc != TF_OK) tf_video_free(out); + free(wanted); + sws_freeContext(sws); + av_frame_free(&frame); + av_packet_free(&pkt); + avcodec_free_context(&dec); + if (fmt) avformat_close_input(&fmt); + if (io) { av_freep(&io->buffer); avio_context_free(&io); } + return rc; +} diff --git a/zig/vendor/stb/README.md b/zig/vendor/stb/README.md new file mode 100644 index 0000000000..11466f4316 --- /dev/null +++ b/zig/vendor/stb/README.md @@ -0,0 +1,3 @@ +stb_image.h v2.30 from https://github.com/nothings/stb at commit 2c980bb59875b0d32144a71867fbdebb2f77cd20 +(sha256 594c2fe35d49488b4382dbfaec8f98366defca819d916ac95becf3e75f4200b3), unmodified; dual public domain / MIT +(see the license at the end of the file). stb_image_impl.c compiles it for JPEG and PNG only. diff --git a/zig/vendor/stb/stb_image.h b/zig/vendor/stb/stb_image.h new file mode 100644 index 0000000000..9eedabedc4 --- /dev/null +++ b/zig/vendor/stb/stb_image.h @@ -0,0 +1,7988 @@ +/* stb_image - v2.30 - public domain image loader - http://nothings.org/stb + no warranty implied; use at your own risk + + Do this: + #define STB_IMAGE_IMPLEMENTATION + before you include this file in *one* C or C++ file to create the implementation. + + // i.e. it should look like this: + #include ... + #include ... + #include ... + #define STB_IMAGE_IMPLEMENTATION + #include "stb_image.h" + + You can #define STBI_ASSERT(x) before the #include to avoid using assert.h. + And #define STBI_MALLOC, STBI_REALLOC, and STBI_FREE to avoid using malloc,realloc,free + + + QUICK NOTES: + Primarily of interest to game developers and other people who can + avoid problematic images and only need the trivial interface + + JPEG baseline & progressive (12 bpc/arithmetic not supported, same as stock IJG lib) + PNG 1/2/4/8/16-bit-per-channel + + TGA (not sure what subset, if a subset) + BMP non-1bpp, non-RLE + PSD (composited view only, no extra channels, 8/16 bit-per-channel) + + GIF (*comp always reports as 4-channel) + HDR (radiance rgbE format) + PIC (Softimage PIC) + PNM (PPM and PGM binary only) + + Animated GIF still needs a proper API, but here's one way to do it: + http://gist.github.com/urraka/685d9a6340b26b830d49 + + - decode from memory or through FILE (define STBI_NO_STDIO to remove code) + - decode from arbitrary I/O callbacks + - SIMD acceleration on x86/x64 (SSE2) and ARM (NEON) + + Full documentation under "DOCUMENTATION" below. + + +LICENSE + + See end of file for license information. + +RECENT REVISION HISTORY: + + 2.30 (2024-05-31) avoid erroneous gcc warning + 2.29 (2023-05-xx) optimizations + 2.28 (2023-01-29) many error fixes, security errors, just tons of stuff + 2.27 (2021-07-11) document stbi_info better, 16-bit PNM support, bug fixes + 2.26 (2020-07-13) many minor fixes + 2.25 (2020-02-02) fix warnings + 2.24 (2020-02-02) fix warnings; thread-local failure_reason and flip_vertically + 2.23 (2019-08-11) fix clang static analysis warning + 2.22 (2019-03-04) gif fixes, fix warnings + 2.21 (2019-02-25) fix typo in comment + 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs + 2.19 (2018-02-11) fix warning + 2.18 (2018-01-30) fix warnings + 2.17 (2018-01-29) bugfix, 1-bit BMP, 16-bitness query, fix warnings + 2.16 (2017-07-23) all functions have 16-bit variants; optimizations; bugfixes + 2.15 (2017-03-18) fix png-1,2,4; all Imagenet JPGs; no runtime SSE detection on GCC + 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs + 2.13 (2016-12-04) experimental 16-bit API, only for PNG so far; fixes + 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes + 2.11 (2016-04-02) 16-bit PNGS; enable SSE2 in non-gcc x64 + RGB-format JPEG; remove white matting in PSD; + allocate large structures on the stack; + correct channel count for PNG & BMP + 2.10 (2016-01-22) avoid warning introduced in 2.09 + 2.09 (2016-01-16) 16-bit TGA; comments in PNM files; STBI_REALLOC_SIZED + + See end of file for full revision history. + + + ============================ Contributors ========================= + + Image formats Extensions, features + Sean Barrett (jpeg, png, bmp) Jetro Lauha (stbi_info) + Nicolas Schulz (hdr, psd) Martin "SpartanJ" Golini (stbi_info) + Jonathan Dummer (tga) James "moose2000" Brown (iPhone PNG) + Jean-Marc Lienher (gif) Ben "Disch" Wenger (io callbacks) + Tom Seddon (pic) Omar Cornut (1/2/4-bit PNG) + Thatcher Ulrich (psd) Nicolas Guillemot (vertical flip) + Ken Miller (pgm, ppm) Richard Mitton (16-bit PSD) + github:urraka (animated gif) Junggon Kim (PNM comments) + Christopher Forseth (animated gif) Daniel Gibson (16-bit TGA) + socks-the-fox (16-bit PNG) + Jeremy Sawicki (handle all ImageNet JPGs) + Optimizations & bugfixes Mikhail Morozov (1-bit BMP) + Fabian "ryg" Giesen Anael Seghezzi (is-16-bit query) + Arseny Kapoulkine Simon Breuss (16-bit PNM) + John-Mark Allen + Carmelo J Fdez-Aguera + + Bug & warning fixes + Marc LeBlanc David Woo Guillaume George Martins Mozeiko + Christpher Lloyd Jerry Jansson Joseph Thomson Blazej Dariusz Roszkowski + Phil Jordan Dave Moore Roy Eltham + Hayaki Saito Nathan Reed Won Chun + Luke Graham Johan Duparc Nick Verigakis the Horde3D community + Thomas Ruf Ronny Chevalier github:rlyeh + Janez Zemva John Bartholomew Michal Cichon github:romigrou + Jonathan Blow Ken Hamada Tero Hanninen github:svdijk + Eugene Golushkov Laurent Gomila Cort Stratton github:snagar + Aruelien Pocheville Sergio Gonzalez Thibault Reuille github:Zelex + Cass Everitt Ryamond Barbiero github:grim210 + Paul Du Bois Engin Manap Aldo Culquicondor github:sammyhw + Philipp Wiesemann Dale Weiler Oriol Ferrer Mesia github:phprus + Josh Tobin Neil Bickford Matthew Gregan github:poppolopoppo + Julian Raschke Gregory Mullen Christian Floisand github:darealshinji + Baldur Karlsson Kevin Schmidt JR Smith github:Michaelangel007 + Brad Weinberger Matvey Cherevko github:mosra + Luca Sas Alexander Veselov Zack Middleton [reserved] + Ryan C. Gordon [reserved] [reserved] + DO NOT ADD YOUR NAME HERE + + Jacko Dirks + + To add your name to the credits, pick a random blank space in the middle and fill it. + 80% of merge conflicts on stb PRs are due to people adding their name at the end + of the credits. +*/ + +#ifndef STBI_INCLUDE_STB_IMAGE_H +#define STBI_INCLUDE_STB_IMAGE_H + +// DOCUMENTATION +// +// Limitations: +// - no 12-bit-per-channel JPEG +// - no JPEGs with arithmetic coding +// - GIF always returns *comp=4 +// +// Basic usage (see HDR discussion below for HDR usage): +// int x,y,n; +// unsigned char *data = stbi_load(filename, &x, &y, &n, 0); +// // ... process data if not NULL ... +// // ... x = width, y = height, n = # 8-bit components per pixel ... +// // ... replace '0' with '1'..'4' to force that many components per pixel +// // ... but 'n' will always be the number that it would have been if you said 0 +// stbi_image_free(data); +// +// Standard parameters: +// int *x -- outputs image width in pixels +// int *y -- outputs image height in pixels +// int *channels_in_file -- outputs # of image components in image file +// int desired_channels -- if non-zero, # of image components requested in result +// +// The return value from an image loader is an 'unsigned char *' which points +// to the pixel data, or NULL on an allocation failure or if the image is +// corrupt or invalid. The pixel data consists of *y scanlines of *x pixels, +// with each pixel consisting of N interleaved 8-bit components; the first +// pixel pointed to is top-left-most in the image. There is no padding between +// image scanlines or between pixels, regardless of format. The number of +// components N is 'desired_channels' if desired_channels is non-zero, or +// *channels_in_file otherwise. If desired_channels is non-zero, +// *channels_in_file has the number of components that _would_ have been +// output otherwise. E.g. if you set desired_channels to 4, you will always +// get RGBA output, but you can check *channels_in_file to see if it's trivially +// opaque because e.g. there were only 3 channels in the source image. +// +// An output image with N components has the following components interleaved +// in this order in each pixel: +// +// N=#comp components +// 1 grey +// 2 grey, alpha +// 3 red, green, blue +// 4 red, green, blue, alpha +// +// If image loading fails for any reason, the return value will be NULL, +// and *x, *y, *channels_in_file will be unchanged. The function +// stbi_failure_reason() can be queried for an extremely brief, end-user +// unfriendly explanation of why the load failed. Define STBI_NO_FAILURE_STRINGS +// to avoid compiling these strings at all, and STBI_FAILURE_USERMSG to get slightly +// more user-friendly ones. +// +// Paletted PNG, BMP, GIF, and PIC images are automatically depalettized. +// +// To query the width, height and component count of an image without having to +// decode the full file, you can use the stbi_info family of functions: +// +// int x,y,n,ok; +// ok = stbi_info(filename, &x, &y, &n); +// // returns ok=1 and sets x, y, n if image is a supported format, +// // 0 otherwise. +// +// Note that stb_image pervasively uses ints in its public API for sizes, +// including sizes of memory buffers. This is now part of the API and thus +// hard to change without causing breakage. As a result, the various image +// loaders all have certain limits on image size; these differ somewhat +// by format but generally boil down to either just under 2GB or just under +// 1GB. When the decoded image would be larger than this, stb_image decoding +// will fail. +// +// Additionally, stb_image will reject image files that have any of their +// dimensions set to a larger value than the configurable STBI_MAX_DIMENSIONS, +// which defaults to 2**24 = 16777216 pixels. Due to the above memory limit, +// the only way to have an image with such dimensions load correctly +// is for it to have a rather extreme aspect ratio. Either way, the +// assumption here is that such larger images are likely to be malformed +// or malicious. If you do need to load an image with individual dimensions +// larger than that, and it still fits in the overall size limit, you can +// #define STBI_MAX_DIMENSIONS on your own to be something larger. +// +// =========================================================================== +// +// UNICODE: +// +// If compiling for Windows and you wish to use Unicode filenames, compile +// with +// #define STBI_WINDOWS_UTF8 +// and pass utf8-encoded filenames. Call stbi_convert_wchar_to_utf8 to convert +// Windows wchar_t filenames to utf8. +// +// =========================================================================== +// +// Philosophy +// +// stb libraries are designed with the following priorities: +// +// 1. easy to use +// 2. easy to maintain +// 3. good performance +// +// Sometimes I let "good performance" creep up in priority over "easy to maintain", +// and for best performance I may provide less-easy-to-use APIs that give higher +// performance, in addition to the easy-to-use ones. Nevertheless, it's important +// to keep in mind that from the standpoint of you, a client of this library, +// all you care about is #1 and #3, and stb libraries DO NOT emphasize #3 above all. +// +// Some secondary priorities arise directly from the first two, some of which +// provide more explicit reasons why performance can't be emphasized. +// +// - Portable ("ease of use") +// - Small source code footprint ("easy to maintain") +// - No dependencies ("ease of use") +// +// =========================================================================== +// +// I/O callbacks +// +// I/O callbacks allow you to read from arbitrary sources, like packaged +// files or some other source. Data read from callbacks are processed +// through a small internal buffer (currently 128 bytes) to try to reduce +// overhead. +// +// The three functions you must define are "read" (reads some bytes of data), +// "skip" (skips some bytes of data), "eof" (reports if the stream is at the end). +// +// =========================================================================== +// +// SIMD support +// +// The JPEG decoder will try to automatically use SIMD kernels on x86 when +// supported by the compiler. For ARM Neon support, you must explicitly +// request it. +// +// (The old do-it-yourself SIMD API is no longer supported in the current +// code.) +// +// On x86, SSE2 will automatically be used when available based on a run-time +// test; if not, the generic C versions are used as a fall-back. On ARM targets, +// the typical path is to have separate builds for NEON and non-NEON devices +// (at least this is true for iOS and Android). Therefore, the NEON support is +// toggled by a build flag: define STBI_NEON to get NEON loops. +// +// If for some reason you do not want to use any of SIMD code, or if +// you have issues compiling it, you can disable it entirely by +// defining STBI_NO_SIMD. +// +// =========================================================================== +// +// HDR image support (disable by defining STBI_NO_HDR) +// +// stb_image supports loading HDR images in general, and currently the Radiance +// .HDR file format specifically. You can still load any file through the existing +// interface; if you attempt to load an HDR file, it will be automatically remapped +// to LDR, assuming gamma 2.2 and an arbitrary scale factor defaulting to 1; +// both of these constants can be reconfigured through this interface: +// +// stbi_hdr_to_ldr_gamma(2.2f); +// stbi_hdr_to_ldr_scale(1.0f); +// +// (note, do not use _inverse_ constants; stbi_image will invert them +// appropriately). +// +// Additionally, there is a new, parallel interface for loading files as +// (linear) floats to preserve the full dynamic range: +// +// float *data = stbi_loadf(filename, &x, &y, &n, 0); +// +// If you load LDR images through this interface, those images will +// be promoted to floating point values, run through the inverse of +// constants corresponding to the above: +// +// stbi_ldr_to_hdr_scale(1.0f); +// stbi_ldr_to_hdr_gamma(2.2f); +// +// Finally, given a filename (or an open file or memory block--see header +// file for details) containing image data, you can query for the "most +// appropriate" interface to use (that is, whether the image is HDR or +// not), using: +// +// stbi_is_hdr(char *filename); +// +// =========================================================================== +// +// iPhone PNG support: +// +// We optionally support converting iPhone-formatted PNGs (which store +// premultiplied BGRA) back to RGB, even though they're internally encoded +// differently. To enable this conversion, call +// stbi_convert_iphone_png_to_rgb(1). +// +// Call stbi_set_unpremultiply_on_load(1) as well to force a divide per +// pixel to remove any premultiplied alpha *only* if the image file explicitly +// says there's premultiplied data (currently only happens in iPhone images, +// and only if iPhone convert-to-rgb processing is on). +// +// =========================================================================== +// +// ADDITIONAL CONFIGURATION +// +// - You can suppress implementation of any of the decoders to reduce +// your code footprint by #defining one or more of the following +// symbols before creating the implementation. +// +// STBI_NO_JPEG +// STBI_NO_PNG +// STBI_NO_BMP +// STBI_NO_PSD +// STBI_NO_TGA +// STBI_NO_GIF +// STBI_NO_HDR +// STBI_NO_PIC +// STBI_NO_PNM (.ppm and .pgm) +// +// - You can request *only* certain decoders and suppress all other ones +// (this will be more forward-compatible, as addition of new decoders +// doesn't require you to disable them explicitly): +// +// STBI_ONLY_JPEG +// STBI_ONLY_PNG +// STBI_ONLY_BMP +// STBI_ONLY_PSD +// STBI_ONLY_TGA +// STBI_ONLY_GIF +// STBI_ONLY_HDR +// STBI_ONLY_PIC +// STBI_ONLY_PNM (.ppm and .pgm) +// +// - If you use STBI_NO_PNG (or _ONLY_ without PNG), and you still +// want the zlib decoder to be available, #define STBI_SUPPORT_ZLIB +// +// - If you define STBI_MAX_DIMENSIONS, stb_image will reject images greater +// than that size (in either width or height) without further processing. +// This is to let programs in the wild set an upper bound to prevent +// denial-of-service attacks on untrusted data, as one could generate a +// valid image of gigantic dimensions and force stb_image to allocate a +// huge block of memory and spend disproportionate time decoding it. By +// default this is set to (1 << 24), which is 16777216, but that's still +// very big. + +#ifndef STBI_NO_STDIO +#include +#endif // STBI_NO_STDIO + +#define STBI_VERSION 1 + +enum +{ + STBI_default = 0, // only used for desired_channels + + STBI_grey = 1, + STBI_grey_alpha = 2, + STBI_rgb = 3, + STBI_rgb_alpha = 4 +}; + +#include +typedef unsigned char stbi_uc; +typedef unsigned short stbi_us; + +#ifdef __cplusplus +extern "C" { +#endif + +#ifndef STBIDEF +#ifdef STB_IMAGE_STATIC +#define STBIDEF static +#else +#define STBIDEF extern +#endif +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// PRIMARY API - works on images of any type +// + +// +// load image by filename, open file, or memory buffer +// + +typedef struct +{ + int (*read) (void *user,char *data,int size); // fill 'data' with 'size' bytes. return number of bytes actually read + void (*skip) (void *user,int n); // skip the next 'n' bytes, or 'unget' the last -n bytes if negative + int (*eof) (void *user); // returns nonzero if we are at end of file/data +} stbi_io_callbacks; + +//////////////////////////////////// +// +// 8-bits-per-channel interface +// + +STBIDEF stbi_uc *stbi_load_from_memory (stbi_uc const *buffer, int len , int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk , void *user, int *x, int *y, int *channels_in_file, int desired_channels); + +#ifndef STBI_NO_STDIO +STBIDEF stbi_uc *stbi_load (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_uc *stbi_load_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); +// for stbi_load_from_file, file pointer is left pointing immediately after image +#endif + +#ifndef STBI_NO_GIF +STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp); +#endif + +#ifdef STBI_WINDOWS_UTF8 +STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input); +#endif + +//////////////////////////////////// +// +// 16-bits-per-channel interface +// + +STBIDEF stbi_us *stbi_load_16_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); + +#ifndef STBI_NO_STDIO +STBIDEF stbi_us *stbi_load_16 (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_us *stbi_load_from_file_16(FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); +#endif + +//////////////////////////////////// +// +// float-per-channel interface +// +#ifndef STBI_NO_LINEAR + STBIDEF float *stbi_loadf_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); + STBIDEF float *stbi_loadf_from_callbacks (stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); + + #ifndef STBI_NO_STDIO + STBIDEF float *stbi_loadf (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); + STBIDEF float *stbi_loadf_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); + #endif +#endif + +#ifndef STBI_NO_HDR + STBIDEF void stbi_hdr_to_ldr_gamma(float gamma); + STBIDEF void stbi_hdr_to_ldr_scale(float scale); +#endif // STBI_NO_HDR + +#ifndef STBI_NO_LINEAR + STBIDEF void stbi_ldr_to_hdr_gamma(float gamma); + STBIDEF void stbi_ldr_to_hdr_scale(float scale); +#endif // STBI_NO_LINEAR + +// stbi_is_hdr is always defined, but always returns false if STBI_NO_HDR +STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user); +STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len); +#ifndef STBI_NO_STDIO +STBIDEF int stbi_is_hdr (char const *filename); +STBIDEF int stbi_is_hdr_from_file(FILE *f); +#endif // STBI_NO_STDIO + + +// get a VERY brief reason for failure +// on most compilers (and ALL modern mainstream compilers) this is threadsafe +STBIDEF const char *stbi_failure_reason (void); + +// free the loaded image -- this is just free() +STBIDEF void stbi_image_free (void *retval_from_stbi_load); + +// get image dimensions & components without fully decoding +STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp); +STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp); +STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len); +STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *clbk, void *user); + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_info (char const *filename, int *x, int *y, int *comp); +STBIDEF int stbi_info_from_file (FILE *f, int *x, int *y, int *comp); +STBIDEF int stbi_is_16_bit (char const *filename); +STBIDEF int stbi_is_16_bit_from_file(FILE *f); +#endif + + + +// for image formats that explicitly notate that they have premultiplied alpha, +// we just return the colors as stored in the file. set this flag to force +// unpremultiplication. results are undefined if the unpremultiply overflow. +STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply); + +// indicate whether we should process iphone images back to canonical format, +// or just pass them through "as-is" +STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert); + +// flip the image vertically, so the first pixel in the output array is the bottom left +STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip); + +// as above, but only applies to images loaded on the thread that calls the function +// this function is only available if your compiler supports thread-local variables; +// calling it will fail to link if your compiler doesn't +STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply); +STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert); +STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip); + +// ZLIB client - used by PNG, available for other purposes + +STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen); +STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header); +STBIDEF char *stbi_zlib_decode_malloc(const char *buffer, int len, int *outlen); +STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); + +STBIDEF char *stbi_zlib_decode_noheader_malloc(const char *buffer, int len, int *outlen); +STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); + + +#ifdef __cplusplus +} +#endif + +// +// +//// end header file ///////////////////////////////////////////////////// +#endif // STBI_INCLUDE_STB_IMAGE_H + +#ifdef STB_IMAGE_IMPLEMENTATION + +#if defined(STBI_ONLY_JPEG) || defined(STBI_ONLY_PNG) || defined(STBI_ONLY_BMP) \ + || defined(STBI_ONLY_TGA) || defined(STBI_ONLY_GIF) || defined(STBI_ONLY_PSD) \ + || defined(STBI_ONLY_HDR) || defined(STBI_ONLY_PIC) || defined(STBI_ONLY_PNM) \ + || defined(STBI_ONLY_ZLIB) + #ifndef STBI_ONLY_JPEG + #define STBI_NO_JPEG + #endif + #ifndef STBI_ONLY_PNG + #define STBI_NO_PNG + #endif + #ifndef STBI_ONLY_BMP + #define STBI_NO_BMP + #endif + #ifndef STBI_ONLY_PSD + #define STBI_NO_PSD + #endif + #ifndef STBI_ONLY_TGA + #define STBI_NO_TGA + #endif + #ifndef STBI_ONLY_GIF + #define STBI_NO_GIF + #endif + #ifndef STBI_ONLY_HDR + #define STBI_NO_HDR + #endif + #ifndef STBI_ONLY_PIC + #define STBI_NO_PIC + #endif + #ifndef STBI_ONLY_PNM + #define STBI_NO_PNM + #endif +#endif + +#if defined(STBI_NO_PNG) && !defined(STBI_SUPPORT_ZLIB) && !defined(STBI_NO_ZLIB) +#define STBI_NO_ZLIB +#endif + + +#include +#include // ptrdiff_t on osx +#include +#include +#include + +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) +#include // ldexp, pow +#endif + +#ifndef STBI_NO_STDIO +#include +#endif + +#ifndef STBI_ASSERT +#include +#define STBI_ASSERT(x) assert(x) +#endif + +#ifdef __cplusplus +#define STBI_EXTERN extern "C" +#else +#define STBI_EXTERN extern +#endif + + +#ifndef _MSC_VER + #ifdef __cplusplus + #define stbi_inline inline + #else + #define stbi_inline + #endif +#else + #define stbi_inline __forceinline +#endif + +#ifndef STBI_NO_THREAD_LOCALS + #if defined(__cplusplus) && __cplusplus >= 201103L + #define STBI_THREAD_LOCAL thread_local + #elif defined(__GNUC__) && __GNUC__ < 5 + #define STBI_THREAD_LOCAL __thread + #elif defined(_MSC_VER) + #define STBI_THREAD_LOCAL __declspec(thread) + #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_THREADS__) + #define STBI_THREAD_LOCAL _Thread_local + #endif + + #ifndef STBI_THREAD_LOCAL + #if defined(__GNUC__) + #define STBI_THREAD_LOCAL __thread + #endif + #endif +#endif + +#if defined(_MSC_VER) || defined(__SYMBIAN32__) +typedef unsigned short stbi__uint16; +typedef signed short stbi__int16; +typedef unsigned int stbi__uint32; +typedef signed int stbi__int32; +#else +#include +typedef uint16_t stbi__uint16; +typedef int16_t stbi__int16; +typedef uint32_t stbi__uint32; +typedef int32_t stbi__int32; +#endif + +// should produce compiler error if size is wrong +typedef unsigned char validate_uint32[sizeof(stbi__uint32)==4 ? 1 : -1]; + +#ifdef _MSC_VER +#define STBI_NOTUSED(v) (void)(v) +#else +#define STBI_NOTUSED(v) (void)sizeof(v) +#endif + +#ifdef _MSC_VER +#define STBI_HAS_LROTL +#endif + +#ifdef STBI_HAS_LROTL + #define stbi_lrot(x,y) _lrotl(x,y) +#else + #define stbi_lrot(x,y) (((x) << (y)) | ((x) >> (-(y) & 31))) +#endif + +#if defined(STBI_MALLOC) && defined(STBI_FREE) && (defined(STBI_REALLOC) || defined(STBI_REALLOC_SIZED)) +// ok +#elif !defined(STBI_MALLOC) && !defined(STBI_FREE) && !defined(STBI_REALLOC) && !defined(STBI_REALLOC_SIZED) +// ok +#else +#error "Must define all or none of STBI_MALLOC, STBI_FREE, and STBI_REALLOC (or STBI_REALLOC_SIZED)." +#endif + +#ifndef STBI_MALLOC +#define STBI_MALLOC(sz) malloc(sz) +#define STBI_REALLOC(p,newsz) realloc(p,newsz) +#define STBI_FREE(p) free(p) +#endif + +#ifndef STBI_REALLOC_SIZED +#define STBI_REALLOC_SIZED(p,oldsz,newsz) STBI_REALLOC(p,newsz) +#endif + +// x86/x64 detection +#if defined(__x86_64__) || defined(_M_X64) +#define STBI__X64_TARGET +#elif defined(__i386) || defined(_M_IX86) +#define STBI__X86_TARGET +#endif + +#if defined(__GNUC__) && defined(STBI__X86_TARGET) && !defined(__SSE2__) && !defined(STBI_NO_SIMD) +// gcc doesn't support sse2 intrinsics unless you compile with -msse2, +// which in turn means it gets to use SSE2 everywhere. This is unfortunate, +// but previous attempts to provide the SSE2 functions with runtime +// detection caused numerous issues. The way architecture extensions are +// exposed in GCC/Clang is, sadly, not really suited for one-file libs. +// New behavior: if compiled with -msse2, we use SSE2 without any +// detection; if not, we don't use it at all. +#define STBI_NO_SIMD +#endif + +#if defined(__MINGW32__) && defined(STBI__X86_TARGET) && !defined(STBI_MINGW_ENABLE_SSE2) && !defined(STBI_NO_SIMD) +// Note that __MINGW32__ doesn't actually mean 32-bit, so we have to avoid STBI__X64_TARGET +// +// 32-bit MinGW wants ESP to be 16-byte aligned, but this is not in the +// Windows ABI and VC++ as well as Windows DLLs don't maintain that invariant. +// As a result, enabling SSE2 on 32-bit MinGW is dangerous when not +// simultaneously enabling "-mstackrealign". +// +// See https://github.com/nothings/stb/issues/81 for more information. +// +// So default to no SSE2 on 32-bit MinGW. If you've read this far and added +// -mstackrealign to your build settings, feel free to #define STBI_MINGW_ENABLE_SSE2. +#define STBI_NO_SIMD +#endif + +#if !defined(STBI_NO_SIMD) && (defined(STBI__X86_TARGET) || defined(STBI__X64_TARGET)) +#define STBI_SSE2 +#include + +#ifdef _MSC_VER + +#if _MSC_VER >= 1400 // not VC6 +#include // __cpuid +static int stbi__cpuid3(void) +{ + int info[4]; + __cpuid(info,1); + return info[3]; +} +#else +static int stbi__cpuid3(void) +{ + int res; + __asm { + mov eax,1 + cpuid + mov res,edx + } + return res; +} +#endif + +#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name + +#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) +static int stbi__sse2_available(void) +{ + int info3 = stbi__cpuid3(); + return ((info3 >> 26) & 1) != 0; +} +#endif + +#else // assume GCC-style if not VC++ +#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) + +#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) +static int stbi__sse2_available(void) +{ + // If we're even attempting to compile this on GCC/Clang, that means + // -msse2 is on, which means the compiler is allowed to use SSE2 + // instructions at will, and so are we. + return 1; +} +#endif + +#endif +#endif + +// ARM NEON +#if defined(STBI_NO_SIMD) && defined(STBI_NEON) +#undef STBI_NEON +#endif + +#ifdef STBI_NEON +#include +#ifdef _MSC_VER +#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name +#else +#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) +#endif +#endif + +#ifndef STBI_SIMD_ALIGN +#define STBI_SIMD_ALIGN(type, name) type name +#endif + +#ifndef STBI_MAX_DIMENSIONS +#define STBI_MAX_DIMENSIONS (1 << 24) +#endif + +/////////////////////////////////////////////// +// +// stbi__context struct and start_xxx functions + +// stbi__context structure is our basic context used by all images, so it +// contains all the IO context, plus some basic image information +typedef struct +{ + stbi__uint32 img_x, img_y; + int img_n, img_out_n; + + stbi_io_callbacks io; + void *io_user_data; + + int read_from_callbacks; + int buflen; + stbi_uc buffer_start[128]; + int callback_already_read; + + stbi_uc *img_buffer, *img_buffer_end; + stbi_uc *img_buffer_original, *img_buffer_original_end; +} stbi__context; + + +static void stbi__refill_buffer(stbi__context *s); + +// initialize a memory-decode context +static void stbi__start_mem(stbi__context *s, stbi_uc const *buffer, int len) +{ + s->io.read = NULL; + s->read_from_callbacks = 0; + s->callback_already_read = 0; + s->img_buffer = s->img_buffer_original = (stbi_uc *) buffer; + s->img_buffer_end = s->img_buffer_original_end = (stbi_uc *) buffer+len; +} + +// initialize a callback-based context +static void stbi__start_callbacks(stbi__context *s, stbi_io_callbacks *c, void *user) +{ + s->io = *c; + s->io_user_data = user; + s->buflen = sizeof(s->buffer_start); + s->read_from_callbacks = 1; + s->callback_already_read = 0; + s->img_buffer = s->img_buffer_original = s->buffer_start; + stbi__refill_buffer(s); + s->img_buffer_original_end = s->img_buffer_end; +} + +#ifndef STBI_NO_STDIO + +static int stbi__stdio_read(void *user, char *data, int size) +{ + return (int) fread(data,1,size,(FILE*) user); +} + +static void stbi__stdio_skip(void *user, int n) +{ + int ch; + fseek((FILE*) user, n, SEEK_CUR); + ch = fgetc((FILE*) user); /* have to read a byte to reset feof()'s flag */ + if (ch != EOF) { + ungetc(ch, (FILE *) user); /* push byte back onto stream if valid. */ + } +} + +static int stbi__stdio_eof(void *user) +{ + return feof((FILE*) user) || ferror((FILE *) user); +} + +static stbi_io_callbacks stbi__stdio_callbacks = +{ + stbi__stdio_read, + stbi__stdio_skip, + stbi__stdio_eof, +}; + +static void stbi__start_file(stbi__context *s, FILE *f) +{ + stbi__start_callbacks(s, &stbi__stdio_callbacks, (void *) f); +} + +//static void stop_file(stbi__context *s) { } + +#endif // !STBI_NO_STDIO + +static void stbi__rewind(stbi__context *s) +{ + // conceptually rewind SHOULD rewind to the beginning of the stream, + // but we just rewind to the beginning of the initial buffer, because + // we only use it after doing 'test', which only ever looks at at most 92 bytes + s->img_buffer = s->img_buffer_original; + s->img_buffer_end = s->img_buffer_original_end; +} + +enum +{ + STBI_ORDER_RGB, + STBI_ORDER_BGR +}; + +typedef struct +{ + int bits_per_channel; + int num_channels; + int channel_order; +} stbi__result_info; + +#ifndef STBI_NO_JPEG +static int stbi__jpeg_test(stbi__context *s); +static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PNG +static int stbi__png_test(stbi__context *s); +static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__png_is16(stbi__context *s); +#endif + +#ifndef STBI_NO_BMP +static int stbi__bmp_test(stbi__context *s); +static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_TGA +static int stbi__tga_test(stbi__context *s); +static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PSD +static int stbi__psd_test(stbi__context *s); +static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc); +static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__psd_is16(stbi__context *s); +#endif + +#ifndef STBI_NO_HDR +static int stbi__hdr_test(stbi__context *s); +static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PIC +static int stbi__pic_test(stbi__context *s); +static void *stbi__pic_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_GIF +static int stbi__gif_test(stbi__context *s); +static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp); +static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PNM +static int stbi__pnm_test(stbi__context *s); +static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__pnm_is16(stbi__context *s); +#endif + +static +#ifdef STBI_THREAD_LOCAL +STBI_THREAD_LOCAL +#endif +const char *stbi__g_failure_reason; + +STBIDEF const char *stbi_failure_reason(void) +{ + return stbi__g_failure_reason; +} + +#ifndef STBI_NO_FAILURE_STRINGS +static int stbi__err(const char *str) +{ + stbi__g_failure_reason = str; + return 0; +} +#endif + +static void *stbi__malloc(size_t size) +{ + return STBI_MALLOC(size); +} + +// stb_image uses ints pervasively, including for offset calculations. +// therefore the largest decoded image size we can support with the +// current code, even on 64-bit targets, is INT_MAX. this is not a +// significant limitation for the intended use case. +// +// we do, however, need to make sure our size calculations don't +// overflow. hence a few helper functions for size calculations that +// multiply integers together, making sure that they're non-negative +// and no overflow occurs. + +// return 1 if the sum is valid, 0 on overflow. +// negative terms are considered invalid. +static int stbi__addsizes_valid(int a, int b) +{ + if (b < 0) return 0; + // now 0 <= b <= INT_MAX, hence also + // 0 <= INT_MAX - b <= INTMAX. + // And "a + b <= INT_MAX" (which might overflow) is the + // same as a <= INT_MAX - b (no overflow) + return a <= INT_MAX - b; +} + +// returns 1 if the product is valid, 0 on overflow. +// negative factors are considered invalid. +static int stbi__mul2sizes_valid(int a, int b) +{ + if (a < 0 || b < 0) return 0; + if (b == 0) return 1; // mul-by-0 is always safe + // portable way to check for no overflows in a*b + return a <= INT_MAX/b; +} + +#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) +// returns 1 if "a*b + add" has no negative terms/factors and doesn't overflow +static int stbi__mad2sizes_valid(int a, int b, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__addsizes_valid(a*b, add); +} +#endif + +// returns 1 if "a*b*c + add" has no negative terms/factors and doesn't overflow +static int stbi__mad3sizes_valid(int a, int b, int c, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && + stbi__addsizes_valid(a*b*c, add); +} + +// returns 1 if "a*b*c*d + add" has no negative terms/factors and doesn't overflow +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) +static int stbi__mad4sizes_valid(int a, int b, int c, int d, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && + stbi__mul2sizes_valid(a*b*c, d) && stbi__addsizes_valid(a*b*c*d, add); +} +#endif + +#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) +// mallocs with size overflow checking +static void *stbi__malloc_mad2(int a, int b, int add) +{ + if (!stbi__mad2sizes_valid(a, b, add)) return NULL; + return stbi__malloc(a*b + add); +} +#endif + +static void *stbi__malloc_mad3(int a, int b, int c, int add) +{ + if (!stbi__mad3sizes_valid(a, b, c, add)) return NULL; + return stbi__malloc(a*b*c + add); +} + +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) +static void *stbi__malloc_mad4(int a, int b, int c, int d, int add) +{ + if (!stbi__mad4sizes_valid(a, b, c, d, add)) return NULL; + return stbi__malloc(a*b*c*d + add); +} +#endif + +// returns 1 if the sum of two signed ints is valid (between -2^31 and 2^31-1 inclusive), 0 on overflow. +static int stbi__addints_valid(int a, int b) +{ + if ((a >= 0) != (b >= 0)) return 1; // a and b have different signs, so no overflow + if (a < 0 && b < 0) return a >= INT_MIN - b; // same as a + b >= INT_MIN; INT_MIN - b cannot overflow since b < 0. + return a <= INT_MAX - b; +} + +// returns 1 if the product of two ints fits in a signed short, 0 on overflow. +static int stbi__mul2shorts_valid(int a, int b) +{ + if (b == 0 || b == -1) return 1; // multiplication by 0 is always 0; check for -1 so SHRT_MIN/b doesn't overflow + if ((a >= 0) == (b >= 0)) return a <= SHRT_MAX/b; // product is positive, so similar to mul2sizes_valid + if (b < 0) return a <= SHRT_MIN / b; // same as a * b >= SHRT_MIN + return a >= SHRT_MIN / b; +} + +// stbi__err - error +// stbi__errpf - error returning pointer to float +// stbi__errpuc - error returning pointer to unsigned char + +#ifdef STBI_NO_FAILURE_STRINGS + #define stbi__err(x,y) 0 +#elif defined(STBI_FAILURE_USERMSG) + #define stbi__err(x,y) stbi__err(y) +#else + #define stbi__err(x,y) stbi__err(x) +#endif + +#define stbi__errpf(x,y) ((float *)(size_t) (stbi__err(x,y)?NULL:NULL)) +#define stbi__errpuc(x,y) ((unsigned char *)(size_t) (stbi__err(x,y)?NULL:NULL)) + +STBIDEF void stbi_image_free(void *retval_from_stbi_load) +{ + STBI_FREE(retval_from_stbi_load); +} + +#ifndef STBI_NO_LINEAR +static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp); +#endif + +#ifndef STBI_NO_HDR +static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp); +#endif + +static int stbi__vertically_flip_on_load_global = 0; + +STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip) +{ + stbi__vertically_flip_on_load_global = flag_true_if_should_flip; +} + +#ifndef STBI_THREAD_LOCAL +#define stbi__vertically_flip_on_load stbi__vertically_flip_on_load_global +#else +static STBI_THREAD_LOCAL int stbi__vertically_flip_on_load_local, stbi__vertically_flip_on_load_set; + +STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip) +{ + stbi__vertically_flip_on_load_local = flag_true_if_should_flip; + stbi__vertically_flip_on_load_set = 1; +} + +#define stbi__vertically_flip_on_load (stbi__vertically_flip_on_load_set \ + ? stbi__vertically_flip_on_load_local \ + : stbi__vertically_flip_on_load_global) +#endif // STBI_THREAD_LOCAL + +static void *stbi__load_main(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) +{ + memset(ri, 0, sizeof(*ri)); // make sure it's initialized if we add new fields + ri->bits_per_channel = 8; // default is 8 so most paths don't have to be changed + ri->channel_order = STBI_ORDER_RGB; // all current input & output are this, but this is here so we can add BGR order + ri->num_channels = 0; + + // test the formats with a very explicit header first (at least a FOURCC + // or distinctive magic number first) + #ifndef STBI_NO_PNG + if (stbi__png_test(s)) return stbi__png_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_BMP + if (stbi__bmp_test(s)) return stbi__bmp_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_GIF + if (stbi__gif_test(s)) return stbi__gif_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_PSD + if (stbi__psd_test(s)) return stbi__psd_load(s,x,y,comp,req_comp, ri, bpc); + #else + STBI_NOTUSED(bpc); + #endif + #ifndef STBI_NO_PIC + if (stbi__pic_test(s)) return stbi__pic_load(s,x,y,comp,req_comp, ri); + #endif + + // then the formats that can end up attempting to load with just 1 or 2 + // bytes matching expectations; these are prone to false positives, so + // try them later + #ifndef STBI_NO_JPEG + if (stbi__jpeg_test(s)) return stbi__jpeg_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_PNM + if (stbi__pnm_test(s)) return stbi__pnm_load(s,x,y,comp,req_comp, ri); + #endif + + #ifndef STBI_NO_HDR + if (stbi__hdr_test(s)) { + float *hdr = stbi__hdr_load(s, x,y,comp,req_comp, ri); + return stbi__hdr_to_ldr(hdr, *x, *y, req_comp ? req_comp : *comp); + } + #endif + + #ifndef STBI_NO_TGA + // test tga last because it's a crappy test! + if (stbi__tga_test(s)) + return stbi__tga_load(s,x,y,comp,req_comp, ri); + #endif + + return stbi__errpuc("unknown image type", "Image not of any known type, or corrupt"); +} + +static stbi_uc *stbi__convert_16_to_8(stbi__uint16 *orig, int w, int h, int channels) +{ + int i; + int img_len = w * h * channels; + stbi_uc *reduced; + + reduced = (stbi_uc *) stbi__malloc(img_len); + if (reduced == NULL) return stbi__errpuc("outofmem", "Out of memory"); + + for (i = 0; i < img_len; ++i) + reduced[i] = (stbi_uc)((orig[i] >> 8) & 0xFF); // top half of each byte is sufficient approx of 16->8 bit scaling + + STBI_FREE(orig); + return reduced; +} + +static stbi__uint16 *stbi__convert_8_to_16(stbi_uc *orig, int w, int h, int channels) +{ + int i; + int img_len = w * h * channels; + stbi__uint16 *enlarged; + + enlarged = (stbi__uint16 *) stbi__malloc(img_len*2); + if (enlarged == NULL) return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); + + for (i = 0; i < img_len; ++i) + enlarged[i] = (stbi__uint16)((orig[i] << 8) + orig[i]); // replicate to high and low byte, maps 0->0, 255->0xffff + + STBI_FREE(orig); + return enlarged; +} + +static void stbi__vertical_flip(void *image, int w, int h, int bytes_per_pixel) +{ + int row; + size_t bytes_per_row = (size_t)w * bytes_per_pixel; + stbi_uc temp[2048]; + stbi_uc *bytes = (stbi_uc *)image; + + for (row = 0; row < (h>>1); row++) { + stbi_uc *row0 = bytes + row*bytes_per_row; + stbi_uc *row1 = bytes + (h - row - 1)*bytes_per_row; + // swap row0 with row1 + size_t bytes_left = bytes_per_row; + while (bytes_left) { + size_t bytes_copy = (bytes_left < sizeof(temp)) ? bytes_left : sizeof(temp); + memcpy(temp, row0, bytes_copy); + memcpy(row0, row1, bytes_copy); + memcpy(row1, temp, bytes_copy); + row0 += bytes_copy; + row1 += bytes_copy; + bytes_left -= bytes_copy; + } + } +} + +#ifndef STBI_NO_GIF +static void stbi__vertical_flip_slices(void *image, int w, int h, int z, int bytes_per_pixel) +{ + int slice; + int slice_size = w * h * bytes_per_pixel; + + stbi_uc *bytes = (stbi_uc *)image; + for (slice = 0; slice < z; ++slice) { + stbi__vertical_flip(bytes, w, h, bytes_per_pixel); + bytes += slice_size; + } +} +#endif + +static unsigned char *stbi__load_and_postprocess_8bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + stbi__result_info ri; + void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 8); + + if (result == NULL) + return NULL; + + // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. + STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); + + if (ri.bits_per_channel != 8) { + result = stbi__convert_16_to_8((stbi__uint16 *) result, *x, *y, req_comp == 0 ? *comp : req_comp); + ri.bits_per_channel = 8; + } + + // @TODO: move stbi__convert_format to here + + if (stbi__vertically_flip_on_load) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi_uc)); + } + + return (unsigned char *) result; +} + +static stbi__uint16 *stbi__load_and_postprocess_16bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + stbi__result_info ri; + void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 16); + + if (result == NULL) + return NULL; + + // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. + STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); + + if (ri.bits_per_channel != 16) { + result = stbi__convert_8_to_16((stbi_uc *) result, *x, *y, req_comp == 0 ? *comp : req_comp); + ri.bits_per_channel = 16; + } + + // @TODO: move stbi__convert_format16 to here + // @TODO: special case RGB-to-Y (and RGBA-to-YA) for 8-bit-to-16-bit case to keep more precision + + if (stbi__vertically_flip_on_load) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi__uint16)); + } + + return (stbi__uint16 *) result; +} + +#if !defined(STBI_NO_HDR) && !defined(STBI_NO_LINEAR) +static void stbi__float_postprocess(float *result, int *x, int *y, int *comp, int req_comp) +{ + if (stbi__vertically_flip_on_load && result != NULL) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(float)); + } +} +#endif + +#ifndef STBI_NO_STDIO + +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) +STBI_EXTERN __declspec(dllimport) int __stdcall MultiByteToWideChar(unsigned int cp, unsigned long flags, const char *str, int cbmb, wchar_t *widestr, int cchwide); +STBI_EXTERN __declspec(dllimport) int __stdcall WideCharToMultiByte(unsigned int cp, unsigned long flags, const wchar_t *widestr, int cchwide, char *str, int cbmb, const char *defchar, int *used_default); +#endif + +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) +STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input) +{ + return WideCharToMultiByte(65001 /* UTF8 */, 0, input, -1, buffer, (int) bufferlen, NULL, NULL); +} +#endif + +static FILE *stbi__fopen(char const *filename, char const *mode) +{ + FILE *f; +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) + wchar_t wMode[64]; + wchar_t wFilename[1024]; + if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, filename, -1, wFilename, sizeof(wFilename)/sizeof(*wFilename))) + return 0; + + if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, mode, -1, wMode, sizeof(wMode)/sizeof(*wMode))) + return 0; + +#if defined(_MSC_VER) && _MSC_VER >= 1400 + if (0 != _wfopen_s(&f, wFilename, wMode)) + f = 0; +#else + f = _wfopen(wFilename, wMode); +#endif + +#elif defined(_MSC_VER) && _MSC_VER >= 1400 + if (0 != fopen_s(&f, filename, mode)) + f=0; +#else + f = fopen(filename, mode); +#endif + return f; +} + + +STBIDEF stbi_uc *stbi_load(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + unsigned char *result; + if (!f) return stbi__errpuc("can't fopen", "Unable to open file"); + result = stbi_load_from_file(f,x,y,comp,req_comp); + fclose(f); + return result; +} + +STBIDEF stbi_uc *stbi_load_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + unsigned char *result; + stbi__context s; + stbi__start_file(&s,f); + result = stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); + if (result) { + // need to 'unget' all the characters in the IO buffer + fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); + } + return result; +} + +STBIDEF stbi__uint16 *stbi_load_from_file_16(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + stbi__uint16 *result; + stbi__context s; + stbi__start_file(&s,f); + result = stbi__load_and_postprocess_16bit(&s,x,y,comp,req_comp); + if (result) { + // need to 'unget' all the characters in the IO buffer + fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); + } + return result; +} + +STBIDEF stbi_us *stbi_load_16(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + stbi__uint16 *result; + if (!f) return (stbi_us *) stbi__errpuc("can't fopen", "Unable to open file"); + result = stbi_load_from_file_16(f,x,y,comp,req_comp); + fclose(f); + return result; +} + + +#endif //!STBI_NO_STDIO + +STBIDEF stbi_us *stbi_load_16_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); +} + +STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *)clbk, user); + return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); +} + +STBIDEF stbi_uc *stbi_load_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); +} + +STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); +} + +#ifndef STBI_NO_GIF +STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp) +{ + unsigned char *result; + stbi__context s; + stbi__start_mem(&s,buffer,len); + + result = (unsigned char*) stbi__load_gif_main(&s, delays, x, y, z, comp, req_comp); + if (stbi__vertically_flip_on_load) { + stbi__vertical_flip_slices( result, *x, *y, *z, *comp ); + } + + return result; +} +#endif + +#ifndef STBI_NO_LINEAR +static float *stbi__loadf_main(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + unsigned char *data; + #ifndef STBI_NO_HDR + if (stbi__hdr_test(s)) { + stbi__result_info ri; + float *hdr_data = stbi__hdr_load(s,x,y,comp,req_comp, &ri); + if (hdr_data) + stbi__float_postprocess(hdr_data,x,y,comp,req_comp); + return hdr_data; + } + #endif + data = stbi__load_and_postprocess_8bit(s, x, y, comp, req_comp); + if (data) + return stbi__ldr_to_hdr(data, *x, *y, req_comp ? req_comp : *comp); + return stbi__errpf("unknown image type", "Image not of any known type, or corrupt"); +} + +STBIDEF float *stbi_loadf_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} + +STBIDEF float *stbi_loadf_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} + +#ifndef STBI_NO_STDIO +STBIDEF float *stbi_loadf(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + float *result; + FILE *f = stbi__fopen(filename, "rb"); + if (!f) return stbi__errpf("can't fopen", "Unable to open file"); + result = stbi_loadf_from_file(f,x,y,comp,req_comp); + fclose(f); + return result; +} + +STBIDEF float *stbi_loadf_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_file(&s,f); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} +#endif // !STBI_NO_STDIO + +#endif // !STBI_NO_LINEAR + +// these is-hdr-or-not is defined independent of whether STBI_NO_LINEAR is +// defined, for API simplicity; if STBI_NO_LINEAR is defined, it always +// reports false! + +STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len) +{ + #ifndef STBI_NO_HDR + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__hdr_test(&s); + #else + STBI_NOTUSED(buffer); + STBI_NOTUSED(len); + return 0; + #endif +} + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_is_hdr (char const *filename) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result=0; + if (f) { + result = stbi_is_hdr_from_file(f); + fclose(f); + } + return result; +} + +STBIDEF int stbi_is_hdr_from_file(FILE *f) +{ + #ifndef STBI_NO_HDR + long pos = ftell(f); + int res; + stbi__context s; + stbi__start_file(&s,f); + res = stbi__hdr_test(&s); + fseek(f, pos, SEEK_SET); + return res; + #else + STBI_NOTUSED(f); + return 0; + #endif +} +#endif // !STBI_NO_STDIO + +STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user) +{ + #ifndef STBI_NO_HDR + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__hdr_test(&s); + #else + STBI_NOTUSED(clbk); + STBI_NOTUSED(user); + return 0; + #endif +} + +#ifndef STBI_NO_LINEAR +static float stbi__l2h_gamma=2.2f, stbi__l2h_scale=1.0f; + +STBIDEF void stbi_ldr_to_hdr_gamma(float gamma) { stbi__l2h_gamma = gamma; } +STBIDEF void stbi_ldr_to_hdr_scale(float scale) { stbi__l2h_scale = scale; } +#endif + +static float stbi__h2l_gamma_i=1.0f/2.2f, stbi__h2l_scale_i=1.0f; + +STBIDEF void stbi_hdr_to_ldr_gamma(float gamma) { stbi__h2l_gamma_i = 1/gamma; } +STBIDEF void stbi_hdr_to_ldr_scale(float scale) { stbi__h2l_scale_i = 1/scale; } + + +////////////////////////////////////////////////////////////////////////////// +// +// Common code used by all image loaders +// + +enum +{ + STBI__SCAN_load=0, + STBI__SCAN_type, + STBI__SCAN_header +}; + +static void stbi__refill_buffer(stbi__context *s) +{ + int n = (s->io.read)(s->io_user_data,(char*)s->buffer_start,s->buflen); + s->callback_already_read += (int) (s->img_buffer - s->img_buffer_original); + if (n == 0) { + // at end of file, treat same as if from memory, but need to handle case + // where s->img_buffer isn't pointing to safe memory, e.g. 0-byte file + s->read_from_callbacks = 0; + s->img_buffer = s->buffer_start; + s->img_buffer_end = s->buffer_start+1; + *s->img_buffer = 0; + } else { + s->img_buffer = s->buffer_start; + s->img_buffer_end = s->buffer_start + n; + } +} + +stbi_inline static stbi_uc stbi__get8(stbi__context *s) +{ + if (s->img_buffer < s->img_buffer_end) + return *s->img_buffer++; + if (s->read_from_callbacks) { + stbi__refill_buffer(s); + return *s->img_buffer++; + } + return 0; +} + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_HDR) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +stbi_inline static int stbi__at_eof(stbi__context *s) +{ + if (s->io.read) { + if (!(s->io.eof)(s->io_user_data)) return 0; + // if feof() is true, check if buffer = end + // special case: we've only got the special 0 character at the end + if (s->read_from_callbacks == 0) return 1; + } + + return s->img_buffer >= s->img_buffer_end; +} +#endif + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) +// nothing +#else +static void stbi__skip(stbi__context *s, int n) +{ + if (n == 0) return; // already there! + if (n < 0) { + s->img_buffer = s->img_buffer_end; + return; + } + if (s->io.read) { + int blen = (int) (s->img_buffer_end - s->img_buffer); + if (blen < n) { + s->img_buffer = s->img_buffer_end; + (s->io.skip)(s->io_user_data, n - blen); + return; + } + } + s->img_buffer += n; +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_TGA) && defined(STBI_NO_HDR) && defined(STBI_NO_PNM) +// nothing +#else +static int stbi__getn(stbi__context *s, stbi_uc *buffer, int n) +{ + if (s->io.read) { + int blen = (int) (s->img_buffer_end - s->img_buffer); + if (blen < n) { + int res, count; + + memcpy(buffer, s->img_buffer, blen); + + count = (s->io.read)(s->io_user_data, (char*) buffer + blen, n - blen); + res = (count == (n-blen)); + s->img_buffer = s->img_buffer_end; + return res; + } + } + + if (s->img_buffer+n <= s->img_buffer_end) { + memcpy(buffer, s->img_buffer, n); + s->img_buffer += n; + return 1; + } else + return 0; +} +#endif + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) +// nothing +#else +static int stbi__get16be(stbi__context *s) +{ + int z = stbi__get8(s); + return (z << 8) + stbi__get8(s); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) +// nothing +#else +static stbi__uint32 stbi__get32be(stbi__context *s) +{ + stbi__uint32 z = stbi__get16be(s); + return (z << 16) + stbi__get16be(s); +} +#endif + +#if defined(STBI_NO_BMP) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) +// nothing +#else +static int stbi__get16le(stbi__context *s) +{ + int z = stbi__get8(s); + return z + (stbi__get8(s) << 8); +} +#endif + +#ifndef STBI_NO_BMP +static stbi__uint32 stbi__get32le(stbi__context *s) +{ + stbi__uint32 z = stbi__get16le(s); + z += (stbi__uint32)stbi__get16le(s) << 16; + return z; +} +#endif + +#define STBI__BYTECAST(x) ((stbi_uc) ((x) & 255)) // truncate int to byte without warnings + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +////////////////////////////////////////////////////////////////////////////// +// +// generic converter from built-in img_n to req_comp +// individual types do this automatically as much as possible (e.g. jpeg +// does all cases internally since it needs to colorspace convert anyway, +// and it never has alpha, so very few cases ). png can automatically +// interleave an alpha=255 channel, but falls back to this for other cases +// +// assume data buffer is malloced, so malloc a new one and free that one +// only failure mode is malloc failing + +static stbi_uc stbi__compute_y(int r, int g, int b) +{ + return (stbi_uc) (((r*77) + (g*150) + (29*b)) >> 8); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +static unsigned char *stbi__convert_format(unsigned char *data, int img_n, int req_comp, unsigned int x, unsigned int y) +{ + int i,j; + unsigned char *good; + + if (req_comp == img_n) return data; + STBI_ASSERT(req_comp >= 1 && req_comp <= 4); + + good = (unsigned char *) stbi__malloc_mad3(req_comp, x, y, 0); + if (good == NULL) { + STBI_FREE(data); + return stbi__errpuc("outofmem", "Out of memory"); + } + + for (j=0; j < (int) y; ++j) { + unsigned char *src = data + j * x * img_n ; + unsigned char *dest = good + j * x * req_comp; + + #define STBI__COMBO(a,b) ((a)*8+(b)) + #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) + // convert source image with img_n components to one with req_comp components; + // avoid switch per pixel, so use switch per scanline and massive macros + switch (STBI__COMBO(img_n, req_comp)) { + STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=255; } break; + STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=255; } break; + STBI__CASE(2,1) { dest[0]=src[0]; } break; + STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; + STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=255; } break; + STBI__CASE(3,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; + STBI__CASE(3,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = 255; } break; + STBI__CASE(4,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; + STBI__CASE(4,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = src[3]; } break; + STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; + default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return stbi__errpuc("unsupported", "Unsupported format conversion"); + } + #undef STBI__CASE + } + + STBI_FREE(data); + return good; +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) +// nothing +#else +static stbi__uint16 stbi__compute_y_16(int r, int g, int b) +{ + return (stbi__uint16) (((r*77) + (g*150) + (29*b)) >> 8); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) +// nothing +#else +static stbi__uint16 *stbi__convert_format16(stbi__uint16 *data, int img_n, int req_comp, unsigned int x, unsigned int y) +{ + int i,j; + stbi__uint16 *good; + + if (req_comp == img_n) return data; + STBI_ASSERT(req_comp >= 1 && req_comp <= 4); + + good = (stbi__uint16 *) stbi__malloc(req_comp * x * y * 2); + if (good == NULL) { + STBI_FREE(data); + return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); + } + + for (j=0; j < (int) y; ++j) { + stbi__uint16 *src = data + j * x * img_n ; + stbi__uint16 *dest = good + j * x * req_comp; + + #define STBI__COMBO(a,b) ((a)*8+(b)) + #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) + // convert source image with img_n components to one with req_comp components; + // avoid switch per pixel, so use switch per scanline and massive macros + switch (STBI__COMBO(img_n, req_comp)) { + STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=0xffff; } break; + STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=0xffff; } break; + STBI__CASE(2,1) { dest[0]=src[0]; } break; + STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; + STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=0xffff; } break; + STBI__CASE(3,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; + STBI__CASE(3,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = 0xffff; } break; + STBI__CASE(4,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; + STBI__CASE(4,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = src[3]; } break; + STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; + default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return (stbi__uint16*) stbi__errpuc("unsupported", "Unsupported format conversion"); + } + #undef STBI__CASE + } + + STBI_FREE(data); + return good; +} +#endif + +#ifndef STBI_NO_LINEAR +static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp) +{ + int i,k,n; + float *output; + if (!data) return NULL; + output = (float *) stbi__malloc_mad4(x, y, comp, sizeof(float), 0); + if (output == NULL) { STBI_FREE(data); return stbi__errpf("outofmem", "Out of memory"); } + // compute number of non-alpha components + if (comp & 1) n = comp; else n = comp-1; + for (i=0; i < x*y; ++i) { + for (k=0; k < n; ++k) { + output[i*comp + k] = (float) (pow(data[i*comp+k]/255.0f, stbi__l2h_gamma) * stbi__l2h_scale); + } + } + if (n < comp) { + for (i=0; i < x*y; ++i) { + output[i*comp + n] = data[i*comp + n]/255.0f; + } + } + STBI_FREE(data); + return output; +} +#endif + +#ifndef STBI_NO_HDR +#define stbi__float2int(x) ((int) (x)) +static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp) +{ + int i,k,n; + stbi_uc *output; + if (!data) return NULL; + output = (stbi_uc *) stbi__malloc_mad3(x, y, comp, 0); + if (output == NULL) { STBI_FREE(data); return stbi__errpuc("outofmem", "Out of memory"); } + // compute number of non-alpha components + if (comp & 1) n = comp; else n = comp-1; + for (i=0; i < x*y; ++i) { + for (k=0; k < n; ++k) { + float z = (float) pow(data[i*comp+k]*stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f; + if (z < 0) z = 0; + if (z > 255) z = 255; + output[i*comp + k] = (stbi_uc) stbi__float2int(z); + } + if (k < comp) { + float z = data[i*comp+k] * 255 + 0.5f; + if (z < 0) z = 0; + if (z > 255) z = 255; + output[i*comp + k] = (stbi_uc) stbi__float2int(z); + } + } + STBI_FREE(data); + return output; +} +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// "baseline" JPEG/JFIF decoder +// +// simple implementation +// - doesn't support delayed output of y-dimension +// - simple interface (only one output format: 8-bit interleaved RGB) +// - doesn't try to recover corrupt jpegs +// - doesn't allow partial loading, loading multiple at once +// - still fast on x86 (copying globals into locals doesn't help x86) +// - allocates lots of intermediate memory (full size of all components) +// - non-interleaved case requires this anyway +// - allows good upsampling (see next) +// high-quality +// - upsampled channels are bilinearly interpolated, even across blocks +// - quality integer IDCT derived from IJG's 'slow' +// performance +// - fast huffman; reasonable integer IDCT +// - some SIMD kernels for common paths on targets with SSE2/NEON +// - uses a lot of intermediate memory, could cache poorly + +#ifndef STBI_NO_JPEG + +// huffman decoding acceleration +#define FAST_BITS 9 // larger handles more cases; smaller stomps less cache + +typedef struct +{ + stbi_uc fast[1 << FAST_BITS]; + // weirdly, repacking this into AoS is a 10% speed loss, instead of a win + stbi__uint16 code[256]; + stbi_uc values[256]; + stbi_uc size[257]; + unsigned int maxcode[18]; + int delta[17]; // old 'firstsymbol' - old 'firstcode' +} stbi__huffman; + +typedef struct +{ + stbi__context *s; + stbi__huffman huff_dc[4]; + stbi__huffman huff_ac[4]; + stbi__uint16 dequant[4][64]; + stbi__int16 fast_ac[4][1 << FAST_BITS]; + +// sizes for components, interleaved MCUs + int img_h_max, img_v_max; + int img_mcu_x, img_mcu_y; + int img_mcu_w, img_mcu_h; + +// definition of jpeg image component + struct + { + int id; + int h,v; + int tq; + int hd,ha; + int dc_pred; + + int x,y,w2,h2; + stbi_uc *data; + void *raw_data, *raw_coeff; + stbi_uc *linebuf; + short *coeff; // progressive only + int coeff_w, coeff_h; // number of 8x8 coefficient blocks + } img_comp[4]; + + stbi__uint32 code_buffer; // jpeg entropy-coded buffer + int code_bits; // number of valid bits + unsigned char marker; // marker seen while filling entropy buffer + int nomore; // flag if we saw a marker so must stop + + int progressive; + int spec_start; + int spec_end; + int succ_high; + int succ_low; + int eob_run; + int jfif; + int app14_color_transform; // Adobe APP14 tag + int rgb; + + int scan_n, order[4]; + int restart_interval, todo; + +// kernels + void (*idct_block_kernel)(stbi_uc *out, int out_stride, short data[64]); + void (*YCbCr_to_RGB_kernel)(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step); + stbi_uc *(*resample_row_hv_2_kernel)(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs); +} stbi__jpeg; + +static int stbi__build_huffman(stbi__huffman *h, int *count) +{ + int i,j,k=0; + unsigned int code; + // build size list for each symbol (from JPEG spec) + for (i=0; i < 16; ++i) { + for (j=0; j < count[i]; ++j) { + h->size[k++] = (stbi_uc) (i+1); + if(k >= 257) return stbi__err("bad size list","Corrupt JPEG"); + } + } + h->size[k] = 0; + + // compute actual symbols (from jpeg spec) + code = 0; + k = 0; + for(j=1; j <= 16; ++j) { + // compute delta to add to code to compute symbol id + h->delta[j] = k - code; + if (h->size[k] == j) { + while (h->size[k] == j) + h->code[k++] = (stbi__uint16) (code++); + if (code-1 >= (1u << j)) return stbi__err("bad code lengths","Corrupt JPEG"); + } + // compute largest code + 1 for this size, preshifted as needed later + h->maxcode[j] = code << (16-j); + code <<= 1; + } + h->maxcode[j] = 0xffffffff; + + // build non-spec acceleration table; 255 is flag for not-accelerated + memset(h->fast, 255, 1 << FAST_BITS); + for (i=0; i < k; ++i) { + int s = h->size[i]; + if (s <= FAST_BITS) { + int c = h->code[i] << (FAST_BITS-s); + int m = 1 << (FAST_BITS-s); + for (j=0; j < m; ++j) { + h->fast[c+j] = (stbi_uc) i; + } + } + } + return 1; +} + +// build a table that decodes both magnitude and value of small ACs in +// one go. +static void stbi__build_fast_ac(stbi__int16 *fast_ac, stbi__huffman *h) +{ + int i; + for (i=0; i < (1 << FAST_BITS); ++i) { + stbi_uc fast = h->fast[i]; + fast_ac[i] = 0; + if (fast < 255) { + int rs = h->values[fast]; + int run = (rs >> 4) & 15; + int magbits = rs & 15; + int len = h->size[fast]; + + if (magbits && len + magbits <= FAST_BITS) { + // magnitude code followed by receive_extend code + int k = ((i << len) & ((1 << FAST_BITS) - 1)) >> (FAST_BITS - magbits); + int m = 1 << (magbits - 1); + if (k < m) k += (~0U << magbits) + 1; + // if the result is small enough, we can fit it in fast_ac table + if (k >= -128 && k <= 127) + fast_ac[i] = (stbi__int16) ((k * 256) + (run * 16) + (len + magbits)); + } + } + } +} + +static void stbi__grow_buffer_unsafe(stbi__jpeg *j) +{ + do { + unsigned int b = j->nomore ? 0 : stbi__get8(j->s); + if (b == 0xff) { + int c = stbi__get8(j->s); + while (c == 0xff) c = stbi__get8(j->s); // consume fill bytes + if (c != 0) { + j->marker = (unsigned char) c; + j->nomore = 1; + return; + } + } + j->code_buffer |= b << (24 - j->code_bits); + j->code_bits += 8; + } while (j->code_bits <= 24); +} + +// (1 << n) - 1 +static const stbi__uint32 stbi__bmask[17]={0,1,3,7,15,31,63,127,255,511,1023,2047,4095,8191,16383,32767,65535}; + +// decode a jpeg huffman value from the bitstream +stbi_inline static int stbi__jpeg_huff_decode(stbi__jpeg *j, stbi__huffman *h) +{ + unsigned int temp; + int c,k; + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + + // look at the top FAST_BITS and determine what symbol ID it is, + // if the code is <= FAST_BITS + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + k = h->fast[c]; + if (k < 255) { + int s = h->size[k]; + if (s > j->code_bits) + return -1; + j->code_buffer <<= s; + j->code_bits -= s; + return h->values[k]; + } + + // naive test is to shift the code_buffer down so k bits are + // valid, then test against maxcode. To speed this up, we've + // preshifted maxcode left so that it has (16-k) 0s at the + // end; in other words, regardless of the number of bits, it + // wants to be compared against something shifted to have 16; + // that way we don't need to shift inside the loop. + temp = j->code_buffer >> 16; + for (k=FAST_BITS+1 ; ; ++k) + if (temp < h->maxcode[k]) + break; + if (k == 17) { + // error! code not found + j->code_bits -= 16; + return -1; + } + + if (k > j->code_bits) + return -1; + + // convert the huffman code to the symbol id + c = ((j->code_buffer >> (32 - k)) & stbi__bmask[k]) + h->delta[k]; + if(c < 0 || c >= 256) // symbol id out of bounds! + return -1; + STBI_ASSERT((((j->code_buffer) >> (32 - h->size[c])) & stbi__bmask[h->size[c]]) == h->code[c]); + + // convert the id to a symbol + j->code_bits -= k; + j->code_buffer <<= k; + return h->values[c]; +} + +// bias[n] = (-1<code_bits < n) stbi__grow_buffer_unsafe(j); + if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing + + sgn = j->code_buffer >> 31; // sign bit always in MSB; 0 if MSB clear (positive), 1 if MSB set (negative) + k = stbi_lrot(j->code_buffer, n); + j->code_buffer = k & ~stbi__bmask[n]; + k &= stbi__bmask[n]; + j->code_bits -= n; + return k + (stbi__jbias[n] & (sgn - 1)); +} + +// get some unsigned bits +stbi_inline static int stbi__jpeg_get_bits(stbi__jpeg *j, int n) +{ + unsigned int k; + if (j->code_bits < n) stbi__grow_buffer_unsafe(j); + if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing + k = stbi_lrot(j->code_buffer, n); + j->code_buffer = k & ~stbi__bmask[n]; + k &= stbi__bmask[n]; + j->code_bits -= n; + return k; +} + +stbi_inline static int stbi__jpeg_get_bit(stbi__jpeg *j) +{ + unsigned int k; + if (j->code_bits < 1) stbi__grow_buffer_unsafe(j); + if (j->code_bits < 1) return 0; // ran out of bits from stream, return 0s intead of continuing + k = j->code_buffer; + j->code_buffer <<= 1; + --j->code_bits; + return k & 0x80000000; +} + +// given a value that's at position X in the zigzag stream, +// where does it appear in the 8x8 matrix coded as row-major? +static const stbi_uc stbi__jpeg_dezigzag[64+15] = +{ + 0, 1, 8, 16, 9, 2, 3, 10, + 17, 24, 32, 25, 18, 11, 4, 5, + 12, 19, 26, 33, 40, 48, 41, 34, + 27, 20, 13, 6, 7, 14, 21, 28, + 35, 42, 49, 56, 57, 50, 43, 36, + 29, 22, 15, 23, 30, 37, 44, 51, + 58, 59, 52, 45, 38, 31, 39, 46, + 53, 60, 61, 54, 47, 55, 62, 63, + // let corrupt input sample past end + 63, 63, 63, 63, 63, 63, 63, 63, + 63, 63, 63, 63, 63, 63, 63 +}; + +// decode one 64-entry block-- +static int stbi__jpeg_decode_block(stbi__jpeg *j, short data[64], stbi__huffman *hdc, stbi__huffman *hac, stbi__int16 *fac, int b, stbi__uint16 *dequant) +{ + int diff,dc,k; + int t; + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + t = stbi__jpeg_huff_decode(j, hdc); + if (t < 0 || t > 15) return stbi__err("bad huffman code","Corrupt JPEG"); + + // 0 all the ac values now so we can do it 32-bits at a time + memset(data,0,64*sizeof(data[0])); + + diff = t ? stbi__extend_receive(j, t) : 0; + if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta","Corrupt JPEG"); + dc = j->img_comp[b].dc_pred + diff; + j->img_comp[b].dc_pred = dc; + if (!stbi__mul2shorts_valid(dc, dequant[0])) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + data[0] = (short) (dc * dequant[0]); + + // decode AC components, see JPEG spec + k = 1; + do { + unsigned int zig; + int c,r,s; + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + r = fac[c]; + if (r) { // fast-AC path + k += (r >> 4) & 15; // run + s = r & 15; // combined length + if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); + j->code_buffer <<= s; + j->code_bits -= s; + // decode into unzigzag'd location + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) ((r >> 8) * dequant[zig]); + } else { + int rs = stbi__jpeg_huff_decode(j, hac); + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (rs != 0xf0) break; // end block + k += 16; + } else { + k += r; + // decode into unzigzag'd location + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) (stbi__extend_receive(j,s) * dequant[zig]); + } + } + } while (k < 64); + return 1; +} + +static int stbi__jpeg_decode_block_prog_dc(stbi__jpeg *j, short data[64], stbi__huffman *hdc, int b) +{ + int diff,dc; + int t; + if (j->spec_end != 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + + if (j->succ_high == 0) { + // first scan for DC coefficient, must be first + memset(data,0,64*sizeof(data[0])); // 0 all the ac values now + t = stbi__jpeg_huff_decode(j, hdc); + if (t < 0 || t > 15) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + diff = t ? stbi__extend_receive(j, t) : 0; + + if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta", "Corrupt JPEG"); + dc = j->img_comp[b].dc_pred + diff; + j->img_comp[b].dc_pred = dc; + if (!stbi__mul2shorts_valid(dc, 1 << j->succ_low)) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + data[0] = (short) (dc * (1 << j->succ_low)); + } else { + // refinement scan for DC coefficient + if (stbi__jpeg_get_bit(j)) + data[0] += (short) (1 << j->succ_low); + } + return 1; +} + +// @OPTIMIZE: store non-zigzagged during the decode passes, +// and only de-zigzag when dequantizing +static int stbi__jpeg_decode_block_prog_ac(stbi__jpeg *j, short data[64], stbi__huffman *hac, stbi__int16 *fac) +{ + int k; + if (j->spec_start == 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + + if (j->succ_high == 0) { + int shift = j->succ_low; + + if (j->eob_run) { + --j->eob_run; + return 1; + } + + k = j->spec_start; + do { + unsigned int zig; + int c,r,s; + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + r = fac[c]; + if (r) { // fast-AC path + k += (r >> 4) & 15; // run + s = r & 15; // combined length + if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); + j->code_buffer <<= s; + j->code_bits -= s; + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) ((r >> 8) * (1 << shift)); + } else { + int rs = stbi__jpeg_huff_decode(j, hac); + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (r < 15) { + j->eob_run = (1 << r); + if (r) + j->eob_run += stbi__jpeg_get_bits(j, r); + --j->eob_run; + break; + } + k += 16; + } else { + k += r; + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) (stbi__extend_receive(j,s) * (1 << shift)); + } + } + } while (k <= j->spec_end); + } else { + // refinement scan for these AC coefficients + + short bit = (short) (1 << j->succ_low); + + if (j->eob_run) { + --j->eob_run; + for (k = j->spec_start; k <= j->spec_end; ++k) { + short *p = &data[stbi__jpeg_dezigzag[k]]; + if (*p != 0) + if (stbi__jpeg_get_bit(j)) + if ((*p & bit)==0) { + if (*p > 0) + *p += bit; + else + *p -= bit; + } + } + } else { + k = j->spec_start; + do { + int r,s; + int rs = stbi__jpeg_huff_decode(j, hac); // @OPTIMIZE see if we can use the fast path here, advance-by-r is so slow, eh + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (r < 15) { + j->eob_run = (1 << r) - 1; + if (r) + j->eob_run += stbi__jpeg_get_bits(j, r); + r = 64; // force end of block + } else { + // r=15 s=0 should write 16 0s, so we just do + // a run of 15 0s and then write s (which is 0), + // so we don't have to do anything special here + } + } else { + if (s != 1) return stbi__err("bad huffman code", "Corrupt JPEG"); + // sign bit + if (stbi__jpeg_get_bit(j)) + s = bit; + else + s = -bit; + } + + // advance by r + while (k <= j->spec_end) { + short *p = &data[stbi__jpeg_dezigzag[k++]]; + if (*p != 0) { + if (stbi__jpeg_get_bit(j)) + if ((*p & bit)==0) { + if (*p > 0) + *p += bit; + else + *p -= bit; + } + } else { + if (r == 0) { + *p = (short) s; + break; + } + --r; + } + } + } while (k <= j->spec_end); + } + } + return 1; +} + +// take a -128..127 value and stbi__clamp it and convert to 0..255 +stbi_inline static stbi_uc stbi__clamp(int x) +{ + // trick to use a single test to catch both cases + if ((unsigned int) x > 255) { + if (x < 0) return 0; + if (x > 255) return 255; + } + return (stbi_uc) x; +} + +#define stbi__f2f(x) ((int) (((x) * 4096 + 0.5))) +#define stbi__fsh(x) ((x) * 4096) + +// derived from jidctint -- DCT_ISLOW +#define STBI__IDCT_1D(s0,s1,s2,s3,s4,s5,s6,s7) \ + int t0,t1,t2,t3,p1,p2,p3,p4,p5,x0,x1,x2,x3; \ + p2 = s2; \ + p3 = s6; \ + p1 = (p2+p3) * stbi__f2f(0.5411961f); \ + t2 = p1 + p3*stbi__f2f(-1.847759065f); \ + t3 = p1 + p2*stbi__f2f( 0.765366865f); \ + p2 = s0; \ + p3 = s4; \ + t0 = stbi__fsh(p2+p3); \ + t1 = stbi__fsh(p2-p3); \ + x0 = t0+t3; \ + x3 = t0-t3; \ + x1 = t1+t2; \ + x2 = t1-t2; \ + t0 = s7; \ + t1 = s5; \ + t2 = s3; \ + t3 = s1; \ + p3 = t0+t2; \ + p4 = t1+t3; \ + p1 = t0+t3; \ + p2 = t1+t2; \ + p5 = (p3+p4)*stbi__f2f( 1.175875602f); \ + t0 = t0*stbi__f2f( 0.298631336f); \ + t1 = t1*stbi__f2f( 2.053119869f); \ + t2 = t2*stbi__f2f( 3.072711026f); \ + t3 = t3*stbi__f2f( 1.501321110f); \ + p1 = p5 + p1*stbi__f2f(-0.899976223f); \ + p2 = p5 + p2*stbi__f2f(-2.562915447f); \ + p3 = p3*stbi__f2f(-1.961570560f); \ + p4 = p4*stbi__f2f(-0.390180644f); \ + t3 += p1+p4; \ + t2 += p2+p3; \ + t1 += p2+p4; \ + t0 += p1+p3; + +static void stbi__idct_block(stbi_uc *out, int out_stride, short data[64]) +{ + int i,val[64],*v=val; + stbi_uc *o; + short *d = data; + + // columns + for (i=0; i < 8; ++i,++d, ++v) { + // if all zeroes, shortcut -- this avoids dequantizing 0s and IDCTing + if (d[ 8]==0 && d[16]==0 && d[24]==0 && d[32]==0 + && d[40]==0 && d[48]==0 && d[56]==0) { + // no shortcut 0 seconds + // (1|2|3|4|5|6|7)==0 0 seconds + // all separate -0.047 seconds + // 1 && 2|3 && 4|5 && 6|7: -0.047 seconds + int dcterm = d[0]*4; + v[0] = v[8] = v[16] = v[24] = v[32] = v[40] = v[48] = v[56] = dcterm; + } else { + STBI__IDCT_1D(d[ 0],d[ 8],d[16],d[24],d[32],d[40],d[48],d[56]) + // constants scaled things up by 1<<12; let's bring them back + // down, but keep 2 extra bits of precision + x0 += 512; x1 += 512; x2 += 512; x3 += 512; + v[ 0] = (x0+t3) >> 10; + v[56] = (x0-t3) >> 10; + v[ 8] = (x1+t2) >> 10; + v[48] = (x1-t2) >> 10; + v[16] = (x2+t1) >> 10; + v[40] = (x2-t1) >> 10; + v[24] = (x3+t0) >> 10; + v[32] = (x3-t0) >> 10; + } + } + + for (i=0, v=val, o=out; i < 8; ++i,v+=8,o+=out_stride) { + // no fast case since the first 1D IDCT spread components out + STBI__IDCT_1D(v[0],v[1],v[2],v[3],v[4],v[5],v[6],v[7]) + // constants scaled things up by 1<<12, plus we had 1<<2 from first + // loop, plus horizontal and vertical each scale by sqrt(8) so together + // we've got an extra 1<<3, so 1<<17 total we need to remove. + // so we want to round that, which means adding 0.5 * 1<<17, + // aka 65536. Also, we'll end up with -128 to 127 that we want + // to encode as 0..255 by adding 128, so we'll add that before the shift + x0 += 65536 + (128<<17); + x1 += 65536 + (128<<17); + x2 += 65536 + (128<<17); + x3 += 65536 + (128<<17); + // tried computing the shifts into temps, or'ing the temps to see + // if any were out of range, but that was slower + o[0] = stbi__clamp((x0+t3) >> 17); + o[7] = stbi__clamp((x0-t3) >> 17); + o[1] = stbi__clamp((x1+t2) >> 17); + o[6] = stbi__clamp((x1-t2) >> 17); + o[2] = stbi__clamp((x2+t1) >> 17); + o[5] = stbi__clamp((x2-t1) >> 17); + o[3] = stbi__clamp((x3+t0) >> 17); + o[4] = stbi__clamp((x3-t0) >> 17); + } +} + +#ifdef STBI_SSE2 +// sse2 integer IDCT. not the fastest possible implementation but it +// produces bit-identical results to the generic C version so it's +// fully "transparent". +static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) +{ + // This is constructed to match our regular (generic) integer IDCT exactly. + __m128i row0, row1, row2, row3, row4, row5, row6, row7; + __m128i tmp; + + // dot product constant: even elems=x, odd elems=y + #define dct_const(x,y) _mm_setr_epi16((x),(y),(x),(y),(x),(y),(x),(y)) + + // out(0) = c0[even]*x + c0[odd]*y (c0, x, y 16-bit, out 32-bit) + // out(1) = c1[even]*x + c1[odd]*y + #define dct_rot(out0,out1, x,y,c0,c1) \ + __m128i c0##lo = _mm_unpacklo_epi16((x),(y)); \ + __m128i c0##hi = _mm_unpackhi_epi16((x),(y)); \ + __m128i out0##_l = _mm_madd_epi16(c0##lo, c0); \ + __m128i out0##_h = _mm_madd_epi16(c0##hi, c0); \ + __m128i out1##_l = _mm_madd_epi16(c0##lo, c1); \ + __m128i out1##_h = _mm_madd_epi16(c0##hi, c1) + + // out = in << 12 (in 16-bit, out 32-bit) + #define dct_widen(out, in) \ + __m128i out##_l = _mm_srai_epi32(_mm_unpacklo_epi16(_mm_setzero_si128(), (in)), 4); \ + __m128i out##_h = _mm_srai_epi32(_mm_unpackhi_epi16(_mm_setzero_si128(), (in)), 4) + + // wide add + #define dct_wadd(out, a, b) \ + __m128i out##_l = _mm_add_epi32(a##_l, b##_l); \ + __m128i out##_h = _mm_add_epi32(a##_h, b##_h) + + // wide sub + #define dct_wsub(out, a, b) \ + __m128i out##_l = _mm_sub_epi32(a##_l, b##_l); \ + __m128i out##_h = _mm_sub_epi32(a##_h, b##_h) + + // butterfly a/b, add bias, then shift by "s" and pack + #define dct_bfly32o(out0, out1, a,b,bias,s) \ + { \ + __m128i abiased_l = _mm_add_epi32(a##_l, bias); \ + __m128i abiased_h = _mm_add_epi32(a##_h, bias); \ + dct_wadd(sum, abiased, b); \ + dct_wsub(dif, abiased, b); \ + out0 = _mm_packs_epi32(_mm_srai_epi32(sum_l, s), _mm_srai_epi32(sum_h, s)); \ + out1 = _mm_packs_epi32(_mm_srai_epi32(dif_l, s), _mm_srai_epi32(dif_h, s)); \ + } + + // 8-bit interleave step (for transposes) + #define dct_interleave8(a, b) \ + tmp = a; \ + a = _mm_unpacklo_epi8(a, b); \ + b = _mm_unpackhi_epi8(tmp, b) + + // 16-bit interleave step (for transposes) + #define dct_interleave16(a, b) \ + tmp = a; \ + a = _mm_unpacklo_epi16(a, b); \ + b = _mm_unpackhi_epi16(tmp, b) + + #define dct_pass(bias,shift) \ + { \ + /* even part */ \ + dct_rot(t2e,t3e, row2,row6, rot0_0,rot0_1); \ + __m128i sum04 = _mm_add_epi16(row0, row4); \ + __m128i dif04 = _mm_sub_epi16(row0, row4); \ + dct_widen(t0e, sum04); \ + dct_widen(t1e, dif04); \ + dct_wadd(x0, t0e, t3e); \ + dct_wsub(x3, t0e, t3e); \ + dct_wadd(x1, t1e, t2e); \ + dct_wsub(x2, t1e, t2e); \ + /* odd part */ \ + dct_rot(y0o,y2o, row7,row3, rot2_0,rot2_1); \ + dct_rot(y1o,y3o, row5,row1, rot3_0,rot3_1); \ + __m128i sum17 = _mm_add_epi16(row1, row7); \ + __m128i sum35 = _mm_add_epi16(row3, row5); \ + dct_rot(y4o,y5o, sum17,sum35, rot1_0,rot1_1); \ + dct_wadd(x4, y0o, y4o); \ + dct_wadd(x5, y1o, y5o); \ + dct_wadd(x6, y2o, y5o); \ + dct_wadd(x7, y3o, y4o); \ + dct_bfly32o(row0,row7, x0,x7,bias,shift); \ + dct_bfly32o(row1,row6, x1,x6,bias,shift); \ + dct_bfly32o(row2,row5, x2,x5,bias,shift); \ + dct_bfly32o(row3,row4, x3,x4,bias,shift); \ + } + + __m128i rot0_0 = dct_const(stbi__f2f(0.5411961f), stbi__f2f(0.5411961f) + stbi__f2f(-1.847759065f)); + __m128i rot0_1 = dct_const(stbi__f2f(0.5411961f) + stbi__f2f( 0.765366865f), stbi__f2f(0.5411961f)); + __m128i rot1_0 = dct_const(stbi__f2f(1.175875602f) + stbi__f2f(-0.899976223f), stbi__f2f(1.175875602f)); + __m128i rot1_1 = dct_const(stbi__f2f(1.175875602f), stbi__f2f(1.175875602f) + stbi__f2f(-2.562915447f)); + __m128i rot2_0 = dct_const(stbi__f2f(-1.961570560f) + stbi__f2f( 0.298631336f), stbi__f2f(-1.961570560f)); + __m128i rot2_1 = dct_const(stbi__f2f(-1.961570560f), stbi__f2f(-1.961570560f) + stbi__f2f( 3.072711026f)); + __m128i rot3_0 = dct_const(stbi__f2f(-0.390180644f) + stbi__f2f( 2.053119869f), stbi__f2f(-0.390180644f)); + __m128i rot3_1 = dct_const(stbi__f2f(-0.390180644f), stbi__f2f(-0.390180644f) + stbi__f2f( 1.501321110f)); + + // rounding biases in column/row passes, see stbi__idct_block for explanation. + __m128i bias_0 = _mm_set1_epi32(512); + __m128i bias_1 = _mm_set1_epi32(65536 + (128<<17)); + + // load + row0 = _mm_load_si128((const __m128i *) (data + 0*8)); + row1 = _mm_load_si128((const __m128i *) (data + 1*8)); + row2 = _mm_load_si128((const __m128i *) (data + 2*8)); + row3 = _mm_load_si128((const __m128i *) (data + 3*8)); + row4 = _mm_load_si128((const __m128i *) (data + 4*8)); + row5 = _mm_load_si128((const __m128i *) (data + 5*8)); + row6 = _mm_load_si128((const __m128i *) (data + 6*8)); + row7 = _mm_load_si128((const __m128i *) (data + 7*8)); + + // column pass + dct_pass(bias_0, 10); + + { + // 16bit 8x8 transpose pass 1 + dct_interleave16(row0, row4); + dct_interleave16(row1, row5); + dct_interleave16(row2, row6); + dct_interleave16(row3, row7); + + // transpose pass 2 + dct_interleave16(row0, row2); + dct_interleave16(row1, row3); + dct_interleave16(row4, row6); + dct_interleave16(row5, row7); + + // transpose pass 3 + dct_interleave16(row0, row1); + dct_interleave16(row2, row3); + dct_interleave16(row4, row5); + dct_interleave16(row6, row7); + } + + // row pass + dct_pass(bias_1, 17); + + { + // pack + __m128i p0 = _mm_packus_epi16(row0, row1); // a0a1a2a3...a7b0b1b2b3...b7 + __m128i p1 = _mm_packus_epi16(row2, row3); + __m128i p2 = _mm_packus_epi16(row4, row5); + __m128i p3 = _mm_packus_epi16(row6, row7); + + // 8bit 8x8 transpose pass 1 + dct_interleave8(p0, p2); // a0e0a1e1... + dct_interleave8(p1, p3); // c0g0c1g1... + + // transpose pass 2 + dct_interleave8(p0, p1); // a0c0e0g0... + dct_interleave8(p2, p3); // b0d0f0h0... + + // transpose pass 3 + dct_interleave8(p0, p2); // a0b0c0d0... + dct_interleave8(p1, p3); // a4b4c4d4... + + // store + _mm_storel_epi64((__m128i *) out, p0); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p0, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p2); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p2, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p1); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p1, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p3); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p3, 0x4e)); + } + +#undef dct_const +#undef dct_rot +#undef dct_widen +#undef dct_wadd +#undef dct_wsub +#undef dct_bfly32o +#undef dct_interleave8 +#undef dct_interleave16 +#undef dct_pass +} + +#endif // STBI_SSE2 + +#ifdef STBI_NEON + +// NEON integer IDCT. should produce bit-identical +// results to the generic C version. +static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) +{ + int16x8_t row0, row1, row2, row3, row4, row5, row6, row7; + + int16x4_t rot0_0 = vdup_n_s16(stbi__f2f(0.5411961f)); + int16x4_t rot0_1 = vdup_n_s16(stbi__f2f(-1.847759065f)); + int16x4_t rot0_2 = vdup_n_s16(stbi__f2f( 0.765366865f)); + int16x4_t rot1_0 = vdup_n_s16(stbi__f2f( 1.175875602f)); + int16x4_t rot1_1 = vdup_n_s16(stbi__f2f(-0.899976223f)); + int16x4_t rot1_2 = vdup_n_s16(stbi__f2f(-2.562915447f)); + int16x4_t rot2_0 = vdup_n_s16(stbi__f2f(-1.961570560f)); + int16x4_t rot2_1 = vdup_n_s16(stbi__f2f(-0.390180644f)); + int16x4_t rot3_0 = vdup_n_s16(stbi__f2f( 0.298631336f)); + int16x4_t rot3_1 = vdup_n_s16(stbi__f2f( 2.053119869f)); + int16x4_t rot3_2 = vdup_n_s16(stbi__f2f( 3.072711026f)); + int16x4_t rot3_3 = vdup_n_s16(stbi__f2f( 1.501321110f)); + +#define dct_long_mul(out, inq, coeff) \ + int32x4_t out##_l = vmull_s16(vget_low_s16(inq), coeff); \ + int32x4_t out##_h = vmull_s16(vget_high_s16(inq), coeff) + +#define dct_long_mac(out, acc, inq, coeff) \ + int32x4_t out##_l = vmlal_s16(acc##_l, vget_low_s16(inq), coeff); \ + int32x4_t out##_h = vmlal_s16(acc##_h, vget_high_s16(inq), coeff) + +#define dct_widen(out, inq) \ + int32x4_t out##_l = vshll_n_s16(vget_low_s16(inq), 12); \ + int32x4_t out##_h = vshll_n_s16(vget_high_s16(inq), 12) + +// wide add +#define dct_wadd(out, a, b) \ + int32x4_t out##_l = vaddq_s32(a##_l, b##_l); \ + int32x4_t out##_h = vaddq_s32(a##_h, b##_h) + +// wide sub +#define dct_wsub(out, a, b) \ + int32x4_t out##_l = vsubq_s32(a##_l, b##_l); \ + int32x4_t out##_h = vsubq_s32(a##_h, b##_h) + +// butterfly a/b, then shift using "shiftop" by "s" and pack +#define dct_bfly32o(out0,out1, a,b,shiftop,s) \ + { \ + dct_wadd(sum, a, b); \ + dct_wsub(dif, a, b); \ + out0 = vcombine_s16(shiftop(sum_l, s), shiftop(sum_h, s)); \ + out1 = vcombine_s16(shiftop(dif_l, s), shiftop(dif_h, s)); \ + } + +#define dct_pass(shiftop, shift) \ + { \ + /* even part */ \ + int16x8_t sum26 = vaddq_s16(row2, row6); \ + dct_long_mul(p1e, sum26, rot0_0); \ + dct_long_mac(t2e, p1e, row6, rot0_1); \ + dct_long_mac(t3e, p1e, row2, rot0_2); \ + int16x8_t sum04 = vaddq_s16(row0, row4); \ + int16x8_t dif04 = vsubq_s16(row0, row4); \ + dct_widen(t0e, sum04); \ + dct_widen(t1e, dif04); \ + dct_wadd(x0, t0e, t3e); \ + dct_wsub(x3, t0e, t3e); \ + dct_wadd(x1, t1e, t2e); \ + dct_wsub(x2, t1e, t2e); \ + /* odd part */ \ + int16x8_t sum15 = vaddq_s16(row1, row5); \ + int16x8_t sum17 = vaddq_s16(row1, row7); \ + int16x8_t sum35 = vaddq_s16(row3, row5); \ + int16x8_t sum37 = vaddq_s16(row3, row7); \ + int16x8_t sumodd = vaddq_s16(sum17, sum35); \ + dct_long_mul(p5o, sumodd, rot1_0); \ + dct_long_mac(p1o, p5o, sum17, rot1_1); \ + dct_long_mac(p2o, p5o, sum35, rot1_2); \ + dct_long_mul(p3o, sum37, rot2_0); \ + dct_long_mul(p4o, sum15, rot2_1); \ + dct_wadd(sump13o, p1o, p3o); \ + dct_wadd(sump24o, p2o, p4o); \ + dct_wadd(sump23o, p2o, p3o); \ + dct_wadd(sump14o, p1o, p4o); \ + dct_long_mac(x4, sump13o, row7, rot3_0); \ + dct_long_mac(x5, sump24o, row5, rot3_1); \ + dct_long_mac(x6, sump23o, row3, rot3_2); \ + dct_long_mac(x7, sump14o, row1, rot3_3); \ + dct_bfly32o(row0,row7, x0,x7,shiftop,shift); \ + dct_bfly32o(row1,row6, x1,x6,shiftop,shift); \ + dct_bfly32o(row2,row5, x2,x5,shiftop,shift); \ + dct_bfly32o(row3,row4, x3,x4,shiftop,shift); \ + } + + // load + row0 = vld1q_s16(data + 0*8); + row1 = vld1q_s16(data + 1*8); + row2 = vld1q_s16(data + 2*8); + row3 = vld1q_s16(data + 3*8); + row4 = vld1q_s16(data + 4*8); + row5 = vld1q_s16(data + 5*8); + row6 = vld1q_s16(data + 6*8); + row7 = vld1q_s16(data + 7*8); + + // add DC bias + row0 = vaddq_s16(row0, vsetq_lane_s16(1024, vdupq_n_s16(0), 0)); + + // column pass + dct_pass(vrshrn_n_s32, 10); + + // 16bit 8x8 transpose + { +// these three map to a single VTRN.16, VTRN.32, and VSWP, respectively. +// whether compilers actually get this is another story, sadly. +#define dct_trn16(x, y) { int16x8x2_t t = vtrnq_s16(x, y); x = t.val[0]; y = t.val[1]; } +#define dct_trn32(x, y) { int32x4x2_t t = vtrnq_s32(vreinterpretq_s32_s16(x), vreinterpretq_s32_s16(y)); x = vreinterpretq_s16_s32(t.val[0]); y = vreinterpretq_s16_s32(t.val[1]); } +#define dct_trn64(x, y) { int16x8_t x0 = x; int16x8_t y0 = y; x = vcombine_s16(vget_low_s16(x0), vget_low_s16(y0)); y = vcombine_s16(vget_high_s16(x0), vget_high_s16(y0)); } + + // pass 1 + dct_trn16(row0, row1); // a0b0a2b2a4b4a6b6 + dct_trn16(row2, row3); + dct_trn16(row4, row5); + dct_trn16(row6, row7); + + // pass 2 + dct_trn32(row0, row2); // a0b0c0d0a4b4c4d4 + dct_trn32(row1, row3); + dct_trn32(row4, row6); + dct_trn32(row5, row7); + + // pass 3 + dct_trn64(row0, row4); // a0b0c0d0e0f0g0h0 + dct_trn64(row1, row5); + dct_trn64(row2, row6); + dct_trn64(row3, row7); + +#undef dct_trn16 +#undef dct_trn32 +#undef dct_trn64 + } + + // row pass + // vrshrn_n_s32 only supports shifts up to 16, we need + // 17. so do a non-rounding shift of 16 first then follow + // up with a rounding shift by 1. + dct_pass(vshrn_n_s32, 16); + + { + // pack and round + uint8x8_t p0 = vqrshrun_n_s16(row0, 1); + uint8x8_t p1 = vqrshrun_n_s16(row1, 1); + uint8x8_t p2 = vqrshrun_n_s16(row2, 1); + uint8x8_t p3 = vqrshrun_n_s16(row3, 1); + uint8x8_t p4 = vqrshrun_n_s16(row4, 1); + uint8x8_t p5 = vqrshrun_n_s16(row5, 1); + uint8x8_t p6 = vqrshrun_n_s16(row6, 1); + uint8x8_t p7 = vqrshrun_n_s16(row7, 1); + + // again, these can translate into one instruction, but often don't. +#define dct_trn8_8(x, y) { uint8x8x2_t t = vtrn_u8(x, y); x = t.val[0]; y = t.val[1]; } +#define dct_trn8_16(x, y) { uint16x4x2_t t = vtrn_u16(vreinterpret_u16_u8(x), vreinterpret_u16_u8(y)); x = vreinterpret_u8_u16(t.val[0]); y = vreinterpret_u8_u16(t.val[1]); } +#define dct_trn8_32(x, y) { uint32x2x2_t t = vtrn_u32(vreinterpret_u32_u8(x), vreinterpret_u32_u8(y)); x = vreinterpret_u8_u32(t.val[0]); y = vreinterpret_u8_u32(t.val[1]); } + + // sadly can't use interleaved stores here since we only write + // 8 bytes to each scan line! + + // 8x8 8-bit transpose pass 1 + dct_trn8_8(p0, p1); + dct_trn8_8(p2, p3); + dct_trn8_8(p4, p5); + dct_trn8_8(p6, p7); + + // pass 2 + dct_trn8_16(p0, p2); + dct_trn8_16(p1, p3); + dct_trn8_16(p4, p6); + dct_trn8_16(p5, p7); + + // pass 3 + dct_trn8_32(p0, p4); + dct_trn8_32(p1, p5); + dct_trn8_32(p2, p6); + dct_trn8_32(p3, p7); + + // store + vst1_u8(out, p0); out += out_stride; + vst1_u8(out, p1); out += out_stride; + vst1_u8(out, p2); out += out_stride; + vst1_u8(out, p3); out += out_stride; + vst1_u8(out, p4); out += out_stride; + vst1_u8(out, p5); out += out_stride; + vst1_u8(out, p6); out += out_stride; + vst1_u8(out, p7); + +#undef dct_trn8_8 +#undef dct_trn8_16 +#undef dct_trn8_32 + } + +#undef dct_long_mul +#undef dct_long_mac +#undef dct_widen +#undef dct_wadd +#undef dct_wsub +#undef dct_bfly32o +#undef dct_pass +} + +#endif // STBI_NEON + +#define STBI__MARKER_none 0xff +// if there's a pending marker from the entropy stream, return that +// otherwise, fetch from the stream and get a marker. if there's no +// marker, return 0xff, which is never a valid marker value +static stbi_uc stbi__get_marker(stbi__jpeg *j) +{ + stbi_uc x; + if (j->marker != STBI__MARKER_none) { x = j->marker; j->marker = STBI__MARKER_none; return x; } + x = stbi__get8(j->s); + if (x != 0xff) return STBI__MARKER_none; + while (x == 0xff) + x = stbi__get8(j->s); // consume repeated 0xff fill bytes + return x; +} + +// in each scan, we'll have scan_n components, and the order +// of the components is specified by order[] +#define STBI__RESTART(x) ((x) >= 0xd0 && (x) <= 0xd7) + +// after a restart interval, stbi__jpeg_reset the entropy decoder and +// the dc prediction +static void stbi__jpeg_reset(stbi__jpeg *j) +{ + j->code_bits = 0; + j->code_buffer = 0; + j->nomore = 0; + j->img_comp[0].dc_pred = j->img_comp[1].dc_pred = j->img_comp[2].dc_pred = j->img_comp[3].dc_pred = 0; + j->marker = STBI__MARKER_none; + j->todo = j->restart_interval ? j->restart_interval : 0x7fffffff; + j->eob_run = 0; + // no more than 1<<31 MCUs if no restart_interal? that's plenty safe, + // since we don't even allow 1<<30 pixels +} + +static int stbi__parse_entropy_coded_data(stbi__jpeg *z) +{ + stbi__jpeg_reset(z); + if (!z->progressive) { + if (z->scan_n == 1) { + int i,j; + STBI_SIMD_ALIGN(short, data[64]); + int n = z->order[0]; + // non-interleaved data, we just need to process one block at a time, + // in trivial scanline order + // number of blocks to do just depends on how many actual "pixels" this + // component has, independent of interleaved MCU blocking and such + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); + // every data block is an MCU, so countdown the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + // if it's NOT a restart, then just bail, so we get corrupt data + // rather than no data + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } else { // interleaved + int i,j,k,x,y; + STBI_SIMD_ALIGN(short, data[64]); + for (j=0; j < z->img_mcu_y; ++j) { + for (i=0; i < z->img_mcu_x; ++i) { + // scan an interleaved mcu... process scan_n components in order + for (k=0; k < z->scan_n; ++k) { + int n = z->order[k]; + // scan out an mcu's worth of this component; that's just determined + // by the basic H and V specified for the component + for (y=0; y < z->img_comp[n].v; ++y) { + for (x=0; x < z->img_comp[n].h; ++x) { + int x2 = (i*z->img_comp[n].h + x)*8; + int y2 = (j*z->img_comp[n].v + y)*8; + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*y2+x2, z->img_comp[n].w2, data); + } + } + } + // after all interleaved components, that's an interleaved MCU, + // so now count down the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } + } else { + if (z->scan_n == 1) { + int i,j; + int n = z->order[0]; + // non-interleaved data, we just need to process one block at a time, + // in trivial scanline order + // number of blocks to do just depends on how many actual "pixels" this + // component has, independent of interleaved MCU blocking and such + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); + if (z->spec_start == 0) { + if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) + return 0; + } else { + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block_prog_ac(z, data, &z->huff_ac[ha], z->fast_ac[ha])) + return 0; + } + // every data block is an MCU, so countdown the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } else { // interleaved + int i,j,k,x,y; + for (j=0; j < z->img_mcu_y; ++j) { + for (i=0; i < z->img_mcu_x; ++i) { + // scan an interleaved mcu... process scan_n components in order + for (k=0; k < z->scan_n; ++k) { + int n = z->order[k]; + // scan out an mcu's worth of this component; that's just determined + // by the basic H and V specified for the component + for (y=0; y < z->img_comp[n].v; ++y) { + for (x=0; x < z->img_comp[n].h; ++x) { + int x2 = (i*z->img_comp[n].h + x); + int y2 = (j*z->img_comp[n].v + y); + short *data = z->img_comp[n].coeff + 64 * (x2 + y2 * z->img_comp[n].coeff_w); + if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) + return 0; + } + } + } + // after all interleaved components, that's an interleaved MCU, + // so now count down the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } + } +} + +static void stbi__jpeg_dequantize(short *data, stbi__uint16 *dequant) +{ + int i; + for (i=0; i < 64; ++i) + data[i] *= dequant[i]; +} + +static void stbi__jpeg_finish(stbi__jpeg *z) +{ + if (z->progressive) { + // dequantize and idct the data + int i,j,n; + for (n=0; n < z->s->img_n; ++n) { + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); + stbi__jpeg_dequantize(data, z->dequant[z->img_comp[n].tq]); + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); + } + } + } + } +} + +static int stbi__process_marker(stbi__jpeg *z, int m) +{ + int L; + switch (m) { + case STBI__MARKER_none: // no marker found + return stbi__err("expected marker","Corrupt JPEG"); + + case 0xDD: // DRI - specify restart interval + if (stbi__get16be(z->s) != 4) return stbi__err("bad DRI len","Corrupt JPEG"); + z->restart_interval = stbi__get16be(z->s); + return 1; + + case 0xDB: // DQT - define quantization table + L = stbi__get16be(z->s)-2; + while (L > 0) { + int q = stbi__get8(z->s); + int p = q >> 4, sixteen = (p != 0); + int t = q & 15,i; + if (p != 0 && p != 1) return stbi__err("bad DQT type","Corrupt JPEG"); + if (t > 3) return stbi__err("bad DQT table","Corrupt JPEG"); + + for (i=0; i < 64; ++i) + z->dequant[t][stbi__jpeg_dezigzag[i]] = (stbi__uint16)(sixteen ? stbi__get16be(z->s) : stbi__get8(z->s)); + L -= (sixteen ? 129 : 65); + } + return L==0; + + case 0xC4: // DHT - define huffman table + L = stbi__get16be(z->s)-2; + while (L > 0) { + stbi_uc *v; + int sizes[16],i,n=0; + int q = stbi__get8(z->s); + int tc = q >> 4; + int th = q & 15; + if (tc > 1 || th > 3) return stbi__err("bad DHT header","Corrupt JPEG"); + for (i=0; i < 16; ++i) { + sizes[i] = stbi__get8(z->s); + n += sizes[i]; + } + if(n > 256) return stbi__err("bad DHT header","Corrupt JPEG"); // Loop over i < n would write past end of values! + L -= 17; + if (tc == 0) { + if (!stbi__build_huffman(z->huff_dc+th, sizes)) return 0; + v = z->huff_dc[th].values; + } else { + if (!stbi__build_huffman(z->huff_ac+th, sizes)) return 0; + v = z->huff_ac[th].values; + } + for (i=0; i < n; ++i) + v[i] = stbi__get8(z->s); + if (tc != 0) + stbi__build_fast_ac(z->fast_ac[th], z->huff_ac + th); + L -= n; + } + return L==0; + } + + // check for comment block or APP blocks + if ((m >= 0xE0 && m <= 0xEF) || m == 0xFE) { + L = stbi__get16be(z->s); + if (L < 2) { + if (m == 0xFE) + return stbi__err("bad COM len","Corrupt JPEG"); + else + return stbi__err("bad APP len","Corrupt JPEG"); + } + L -= 2; + + if (m == 0xE0 && L >= 5) { // JFIF APP0 segment + static const unsigned char tag[5] = {'J','F','I','F','\0'}; + int ok = 1; + int i; + for (i=0; i < 5; ++i) + if (stbi__get8(z->s) != tag[i]) + ok = 0; + L -= 5; + if (ok) + z->jfif = 1; + } else if (m == 0xEE && L >= 12) { // Adobe APP14 segment + static const unsigned char tag[6] = {'A','d','o','b','e','\0'}; + int ok = 1; + int i; + for (i=0; i < 6; ++i) + if (stbi__get8(z->s) != tag[i]) + ok = 0; + L -= 6; + if (ok) { + stbi__get8(z->s); // version + stbi__get16be(z->s); // flags0 + stbi__get16be(z->s); // flags1 + z->app14_color_transform = stbi__get8(z->s); // color transform + L -= 6; + } + } + + stbi__skip(z->s, L); + return 1; + } + + return stbi__err("unknown marker","Corrupt JPEG"); +} + +// after we see SOS +static int stbi__process_scan_header(stbi__jpeg *z) +{ + int i; + int Ls = stbi__get16be(z->s); + z->scan_n = stbi__get8(z->s); + if (z->scan_n < 1 || z->scan_n > 4 || z->scan_n > (int) z->s->img_n) return stbi__err("bad SOS component count","Corrupt JPEG"); + if (Ls != 6+2*z->scan_n) return stbi__err("bad SOS len","Corrupt JPEG"); + for (i=0; i < z->scan_n; ++i) { + int id = stbi__get8(z->s), which; + int q = stbi__get8(z->s); + for (which = 0; which < z->s->img_n; ++which) + if (z->img_comp[which].id == id) + break; + if (which == z->s->img_n) return 0; // no match + z->img_comp[which].hd = q >> 4; if (z->img_comp[which].hd > 3) return stbi__err("bad DC huff","Corrupt JPEG"); + z->img_comp[which].ha = q & 15; if (z->img_comp[which].ha > 3) return stbi__err("bad AC huff","Corrupt JPEG"); + z->order[i] = which; + } + + { + int aa; + z->spec_start = stbi__get8(z->s); + z->spec_end = stbi__get8(z->s); // should be 63, but might be 0 + aa = stbi__get8(z->s); + z->succ_high = (aa >> 4); + z->succ_low = (aa & 15); + if (z->progressive) { + if (z->spec_start > 63 || z->spec_end > 63 || z->spec_start > z->spec_end || z->succ_high > 13 || z->succ_low > 13) + return stbi__err("bad SOS", "Corrupt JPEG"); + } else { + if (z->spec_start != 0) return stbi__err("bad SOS","Corrupt JPEG"); + if (z->succ_high != 0 || z->succ_low != 0) return stbi__err("bad SOS","Corrupt JPEG"); + z->spec_end = 63; + } + } + + return 1; +} + +static int stbi__free_jpeg_components(stbi__jpeg *z, int ncomp, int why) +{ + int i; + for (i=0; i < ncomp; ++i) { + if (z->img_comp[i].raw_data) { + STBI_FREE(z->img_comp[i].raw_data); + z->img_comp[i].raw_data = NULL; + z->img_comp[i].data = NULL; + } + if (z->img_comp[i].raw_coeff) { + STBI_FREE(z->img_comp[i].raw_coeff); + z->img_comp[i].raw_coeff = 0; + z->img_comp[i].coeff = 0; + } + if (z->img_comp[i].linebuf) { + STBI_FREE(z->img_comp[i].linebuf); + z->img_comp[i].linebuf = NULL; + } + } + return why; +} + +static int stbi__process_frame_header(stbi__jpeg *z, int scan) +{ + stbi__context *s = z->s; + int Lf,p,i,q, h_max=1,v_max=1,c; + Lf = stbi__get16be(s); if (Lf < 11) return stbi__err("bad SOF len","Corrupt JPEG"); // JPEG + p = stbi__get8(s); if (p != 8) return stbi__err("only 8-bit","JPEG format not supported: 8-bit only"); // JPEG baseline + s->img_y = stbi__get16be(s); if (s->img_y == 0) return stbi__err("no header height", "JPEG format not supported: delayed height"); // Legal, but we don't handle it--but neither does IJG + s->img_x = stbi__get16be(s); if (s->img_x == 0) return stbi__err("0 width","Corrupt JPEG"); // JPEG requires + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + c = stbi__get8(s); + if (c != 3 && c != 1 && c != 4) return stbi__err("bad component count","Corrupt JPEG"); + s->img_n = c; + for (i=0; i < c; ++i) { + z->img_comp[i].data = NULL; + z->img_comp[i].linebuf = NULL; + } + + if (Lf != 8+3*s->img_n) return stbi__err("bad SOF len","Corrupt JPEG"); + + z->rgb = 0; + for (i=0; i < s->img_n; ++i) { + static const unsigned char rgb[3] = { 'R', 'G', 'B' }; + z->img_comp[i].id = stbi__get8(s); + if (s->img_n == 3 && z->img_comp[i].id == rgb[i]) + ++z->rgb; + q = stbi__get8(s); + z->img_comp[i].h = (q >> 4); if (!z->img_comp[i].h || z->img_comp[i].h > 4) return stbi__err("bad H","Corrupt JPEG"); + z->img_comp[i].v = q & 15; if (!z->img_comp[i].v || z->img_comp[i].v > 4) return stbi__err("bad V","Corrupt JPEG"); + z->img_comp[i].tq = stbi__get8(s); if (z->img_comp[i].tq > 3) return stbi__err("bad TQ","Corrupt JPEG"); + } + + if (scan != STBI__SCAN_load) return 1; + + if (!stbi__mad3sizes_valid(s->img_x, s->img_y, s->img_n, 0)) return stbi__err("too large", "Image too large to decode"); + + for (i=0; i < s->img_n; ++i) { + if (z->img_comp[i].h > h_max) h_max = z->img_comp[i].h; + if (z->img_comp[i].v > v_max) v_max = z->img_comp[i].v; + } + + // check that plane subsampling factors are integer ratios; our resamplers can't deal with fractional ratios + // and I've never seen a non-corrupted JPEG file actually use them + for (i=0; i < s->img_n; ++i) { + if (h_max % z->img_comp[i].h != 0) return stbi__err("bad H","Corrupt JPEG"); + if (v_max % z->img_comp[i].v != 0) return stbi__err("bad V","Corrupt JPEG"); + } + + // compute interleaved mcu info + z->img_h_max = h_max; + z->img_v_max = v_max; + z->img_mcu_w = h_max * 8; + z->img_mcu_h = v_max * 8; + // these sizes can't be more than 17 bits + z->img_mcu_x = (s->img_x + z->img_mcu_w-1) / z->img_mcu_w; + z->img_mcu_y = (s->img_y + z->img_mcu_h-1) / z->img_mcu_h; + + for (i=0; i < s->img_n; ++i) { + // number of effective pixels (e.g. for non-interleaved MCU) + z->img_comp[i].x = (s->img_x * z->img_comp[i].h + h_max-1) / h_max; + z->img_comp[i].y = (s->img_y * z->img_comp[i].v + v_max-1) / v_max; + // to simplify generation, we'll allocate enough memory to decode + // the bogus oversized data from using interleaved MCUs and their + // big blocks (e.g. a 16x16 iMCU on an image of width 33); we won't + // discard the extra data until colorspace conversion + // + // img_mcu_x, img_mcu_y: <=17 bits; comp[i].h and .v are <=4 (checked earlier) + // so these muls can't overflow with 32-bit ints (which we require) + z->img_comp[i].w2 = z->img_mcu_x * z->img_comp[i].h * 8; + z->img_comp[i].h2 = z->img_mcu_y * z->img_comp[i].v * 8; + z->img_comp[i].coeff = 0; + z->img_comp[i].raw_coeff = 0; + z->img_comp[i].linebuf = NULL; + z->img_comp[i].raw_data = stbi__malloc_mad2(z->img_comp[i].w2, z->img_comp[i].h2, 15); + if (z->img_comp[i].raw_data == NULL) + return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); + // align blocks for idct using mmx/sse + z->img_comp[i].data = (stbi_uc*) (((size_t) z->img_comp[i].raw_data + 15) & ~15); + if (z->progressive) { + // w2, h2 are multiples of 8 (see above) + z->img_comp[i].coeff_w = z->img_comp[i].w2 / 8; + z->img_comp[i].coeff_h = z->img_comp[i].h2 / 8; + z->img_comp[i].raw_coeff = stbi__malloc_mad3(z->img_comp[i].w2, z->img_comp[i].h2, sizeof(short), 15); + if (z->img_comp[i].raw_coeff == NULL) + return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); + z->img_comp[i].coeff = (short*) (((size_t) z->img_comp[i].raw_coeff + 15) & ~15); + } + } + + return 1; +} + +// use comparisons since in some cases we handle more than one case (e.g. SOF) +#define stbi__DNL(x) ((x) == 0xdc) +#define stbi__SOI(x) ((x) == 0xd8) +#define stbi__EOI(x) ((x) == 0xd9) +#define stbi__SOF(x) ((x) == 0xc0 || (x) == 0xc1 || (x) == 0xc2) +#define stbi__SOS(x) ((x) == 0xda) + +#define stbi__SOF_progressive(x) ((x) == 0xc2) + +static int stbi__decode_jpeg_header(stbi__jpeg *z, int scan) +{ + int m; + z->jfif = 0; + z->app14_color_transform = -1; // valid values are 0,1,2 + z->marker = STBI__MARKER_none; // initialize cached marker to empty + m = stbi__get_marker(z); + if (!stbi__SOI(m)) return stbi__err("no SOI","Corrupt JPEG"); + if (scan == STBI__SCAN_type) return 1; + m = stbi__get_marker(z); + while (!stbi__SOF(m)) { + if (!stbi__process_marker(z,m)) return 0; + m = stbi__get_marker(z); + while (m == STBI__MARKER_none) { + // some files have extra padding after their blocks, so ok, we'll scan + if (stbi__at_eof(z->s)) return stbi__err("no SOF", "Corrupt JPEG"); + m = stbi__get_marker(z); + } + } + z->progressive = stbi__SOF_progressive(m); + if (!stbi__process_frame_header(z, scan)) return 0; + return 1; +} + +static stbi_uc stbi__skip_jpeg_junk_at_end(stbi__jpeg *j) +{ + // some JPEGs have junk at end, skip over it but if we find what looks + // like a valid marker, resume there + while (!stbi__at_eof(j->s)) { + stbi_uc x = stbi__get8(j->s); + while (x == 0xff) { // might be a marker + if (stbi__at_eof(j->s)) return STBI__MARKER_none; + x = stbi__get8(j->s); + if (x != 0x00 && x != 0xff) { + // not a stuffed zero or lead-in to another marker, looks + // like an actual marker, return it + return x; + } + // stuffed zero has x=0 now which ends the loop, meaning we go + // back to regular scan loop. + // repeated 0xff keeps trying to read the next byte of the marker. + } + } + return STBI__MARKER_none; +} + +// decode image to YCbCr format +static int stbi__decode_jpeg_image(stbi__jpeg *j) +{ + int m; + for (m = 0; m < 4; m++) { + j->img_comp[m].raw_data = NULL; + j->img_comp[m].raw_coeff = NULL; + } + j->restart_interval = 0; + if (!stbi__decode_jpeg_header(j, STBI__SCAN_load)) return 0; + m = stbi__get_marker(j); + while (!stbi__EOI(m)) { + if (stbi__SOS(m)) { + if (!stbi__process_scan_header(j)) return 0; + if (!stbi__parse_entropy_coded_data(j)) return 0; + if (j->marker == STBI__MARKER_none ) { + j->marker = stbi__skip_jpeg_junk_at_end(j); + // if we reach eof without hitting a marker, stbi__get_marker() below will fail and we'll eventually return 0 + } + m = stbi__get_marker(j); + if (STBI__RESTART(m)) + m = stbi__get_marker(j); + } else if (stbi__DNL(m)) { + int Ld = stbi__get16be(j->s); + stbi__uint32 NL = stbi__get16be(j->s); + if (Ld != 4) return stbi__err("bad DNL len", "Corrupt JPEG"); + if (NL != j->s->img_y) return stbi__err("bad DNL height", "Corrupt JPEG"); + m = stbi__get_marker(j); + } else { + if (!stbi__process_marker(j, m)) return 1; + m = stbi__get_marker(j); + } + } + if (j->progressive) + stbi__jpeg_finish(j); + return 1; +} + +// static jfif-centered resampling (across block boundaries) + +typedef stbi_uc *(*resample_row_func)(stbi_uc *out, stbi_uc *in0, stbi_uc *in1, + int w, int hs); + +#define stbi__div4(x) ((stbi_uc) ((x) >> 2)) + +static stbi_uc *resample_row_1(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + STBI_NOTUSED(out); + STBI_NOTUSED(in_far); + STBI_NOTUSED(w); + STBI_NOTUSED(hs); + return in_near; +} + +static stbi_uc* stbi__resample_row_v_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate two samples vertically for every one in input + int i; + STBI_NOTUSED(hs); + for (i=0; i < w; ++i) + out[i] = stbi__div4(3*in_near[i] + in_far[i] + 2); + return out; +} + +static stbi_uc* stbi__resample_row_h_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate two samples horizontally for every one in input + int i; + stbi_uc *input = in_near; + + if (w == 1) { + // if only one sample, can't do any interpolation + out[0] = out[1] = input[0]; + return out; + } + + out[0] = input[0]; + out[1] = stbi__div4(input[0]*3 + input[1] + 2); + for (i=1; i < w-1; ++i) { + int n = 3*input[i]+2; + out[i*2+0] = stbi__div4(n+input[i-1]); + out[i*2+1] = stbi__div4(n+input[i+1]); + } + out[i*2+0] = stbi__div4(input[w-2]*3 + input[w-1] + 2); + out[i*2+1] = input[w-1]; + + STBI_NOTUSED(in_far); + STBI_NOTUSED(hs); + + return out; +} + +#define stbi__div16(x) ((stbi_uc) ((x) >> 4)) + +static stbi_uc *stbi__resample_row_hv_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate 2x2 samples for every one in input + int i,t0,t1; + if (w == 1) { + out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); + return out; + } + + t1 = 3*in_near[0] + in_far[0]; + out[0] = stbi__div4(t1+2); + for (i=1; i < w; ++i) { + t0 = t1; + t1 = 3*in_near[i]+in_far[i]; + out[i*2-1] = stbi__div16(3*t0 + t1 + 8); + out[i*2 ] = stbi__div16(3*t1 + t0 + 8); + } + out[w*2-1] = stbi__div4(t1+2); + + STBI_NOTUSED(hs); + + return out; +} + +#if defined(STBI_SSE2) || defined(STBI_NEON) +static stbi_uc *stbi__resample_row_hv_2_simd(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate 2x2 samples for every one in input + int i=0,t0,t1; + + if (w == 1) { + out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); + return out; + } + + t1 = 3*in_near[0] + in_far[0]; + // process groups of 8 pixels for as long as we can. + // note we can't handle the last pixel in a row in this loop + // because we need to handle the filter boundary conditions. + for (; i < ((w-1) & ~7); i += 8) { +#if defined(STBI_SSE2) + // load and perform the vertical filtering pass + // this uses 3*x + y = 4*x + (y - x) + __m128i zero = _mm_setzero_si128(); + __m128i farb = _mm_loadl_epi64((__m128i *) (in_far + i)); + __m128i nearb = _mm_loadl_epi64((__m128i *) (in_near + i)); + __m128i farw = _mm_unpacklo_epi8(farb, zero); + __m128i nearw = _mm_unpacklo_epi8(nearb, zero); + __m128i diff = _mm_sub_epi16(farw, nearw); + __m128i nears = _mm_slli_epi16(nearw, 2); + __m128i curr = _mm_add_epi16(nears, diff); // current row + + // horizontal filter works the same based on shifted vers of current + // row. "prev" is current row shifted right by 1 pixel; we need to + // insert the previous pixel value (from t1). + // "next" is current row shifted left by 1 pixel, with first pixel + // of next block of 8 pixels added in. + __m128i prv0 = _mm_slli_si128(curr, 2); + __m128i nxt0 = _mm_srli_si128(curr, 2); + __m128i prev = _mm_insert_epi16(prv0, t1, 0); + __m128i next = _mm_insert_epi16(nxt0, 3*in_near[i+8] + in_far[i+8], 7); + + // horizontal filter, polyphase implementation since it's convenient: + // even pixels = 3*cur + prev = cur*4 + (prev - cur) + // odd pixels = 3*cur + next = cur*4 + (next - cur) + // note the shared term. + __m128i bias = _mm_set1_epi16(8); + __m128i curs = _mm_slli_epi16(curr, 2); + __m128i prvd = _mm_sub_epi16(prev, curr); + __m128i nxtd = _mm_sub_epi16(next, curr); + __m128i curb = _mm_add_epi16(curs, bias); + __m128i even = _mm_add_epi16(prvd, curb); + __m128i odd = _mm_add_epi16(nxtd, curb); + + // interleave even and odd pixels, then undo scaling. + __m128i int0 = _mm_unpacklo_epi16(even, odd); + __m128i int1 = _mm_unpackhi_epi16(even, odd); + __m128i de0 = _mm_srli_epi16(int0, 4); + __m128i de1 = _mm_srli_epi16(int1, 4); + + // pack and write output + __m128i outv = _mm_packus_epi16(de0, de1); + _mm_storeu_si128((__m128i *) (out + i*2), outv); +#elif defined(STBI_NEON) + // load and perform the vertical filtering pass + // this uses 3*x + y = 4*x + (y - x) + uint8x8_t farb = vld1_u8(in_far + i); + uint8x8_t nearb = vld1_u8(in_near + i); + int16x8_t diff = vreinterpretq_s16_u16(vsubl_u8(farb, nearb)); + int16x8_t nears = vreinterpretq_s16_u16(vshll_n_u8(nearb, 2)); + int16x8_t curr = vaddq_s16(nears, diff); // current row + + // horizontal filter works the same based on shifted vers of current + // row. "prev" is current row shifted right by 1 pixel; we need to + // insert the previous pixel value (from t1). + // "next" is current row shifted left by 1 pixel, with first pixel + // of next block of 8 pixels added in. + int16x8_t prv0 = vextq_s16(curr, curr, 7); + int16x8_t nxt0 = vextq_s16(curr, curr, 1); + int16x8_t prev = vsetq_lane_s16(t1, prv0, 0); + int16x8_t next = vsetq_lane_s16(3*in_near[i+8] + in_far[i+8], nxt0, 7); + + // horizontal filter, polyphase implementation since it's convenient: + // even pixels = 3*cur + prev = cur*4 + (prev - cur) + // odd pixels = 3*cur + next = cur*4 + (next - cur) + // note the shared term. + int16x8_t curs = vshlq_n_s16(curr, 2); + int16x8_t prvd = vsubq_s16(prev, curr); + int16x8_t nxtd = vsubq_s16(next, curr); + int16x8_t even = vaddq_s16(curs, prvd); + int16x8_t odd = vaddq_s16(curs, nxtd); + + // undo scaling and round, then store with even/odd phases interleaved + uint8x8x2_t o; + o.val[0] = vqrshrun_n_s16(even, 4); + o.val[1] = vqrshrun_n_s16(odd, 4); + vst2_u8(out + i*2, o); +#endif + + // "previous" value for next iter + t1 = 3*in_near[i+7] + in_far[i+7]; + } + + t0 = t1; + t1 = 3*in_near[i] + in_far[i]; + out[i*2] = stbi__div16(3*t1 + t0 + 8); + + for (++i; i < w; ++i) { + t0 = t1; + t1 = 3*in_near[i]+in_far[i]; + out[i*2-1] = stbi__div16(3*t0 + t1 + 8); + out[i*2 ] = stbi__div16(3*t1 + t0 + 8); + } + out[w*2-1] = stbi__div4(t1+2); + + STBI_NOTUSED(hs); + + return out; +} +#endif + +static stbi_uc *stbi__resample_row_generic(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // resample with nearest-neighbor + int i,j; + STBI_NOTUSED(in_far); + for (i=0; i < w; ++i) + for (j=0; j < hs; ++j) + out[i*hs+j] = in_near[i]; + return out; +} + +// this is a reduced-precision calculation of YCbCr-to-RGB introduced +// to make sure the code produces the same results in both SIMD and scalar +#define stbi__float2fixed(x) (((int) ((x) * 4096.0f + 0.5f)) << 8) +static void stbi__YCbCr_to_RGB_row(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step) +{ + int i; + for (i=0; i < count; ++i) { + int y_fixed = (y[i] << 20) + (1<<19); // rounding + int r,g,b; + int cr = pcr[i] - 128; + int cb = pcb[i] - 128; + r = y_fixed + cr* stbi__float2fixed(1.40200f); + g = y_fixed + (cr*-stbi__float2fixed(0.71414f)) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); + b = y_fixed + cb* stbi__float2fixed(1.77200f); + r >>= 20; + g >>= 20; + b >>= 20; + if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } + if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } + if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } + out[0] = (stbi_uc)r; + out[1] = (stbi_uc)g; + out[2] = (stbi_uc)b; + out[3] = 255; + out += step; + } +} + +#if defined(STBI_SSE2) || defined(STBI_NEON) +static void stbi__YCbCr_to_RGB_simd(stbi_uc *out, stbi_uc const *y, stbi_uc const *pcb, stbi_uc const *pcr, int count, int step) +{ + int i = 0; + +#ifdef STBI_SSE2 + // step == 3 is pretty ugly on the final interleave, and i'm not convinced + // it's useful in practice (you wouldn't use it for textures, for example). + // so just accelerate step == 4 case. + if (step == 4) { + // this is a fairly straightforward implementation and not super-optimized. + __m128i signflip = _mm_set1_epi8(-0x80); + __m128i cr_const0 = _mm_set1_epi16( (short) ( 1.40200f*4096.0f+0.5f)); + __m128i cr_const1 = _mm_set1_epi16( - (short) ( 0.71414f*4096.0f+0.5f)); + __m128i cb_const0 = _mm_set1_epi16( - (short) ( 0.34414f*4096.0f+0.5f)); + __m128i cb_const1 = _mm_set1_epi16( (short) ( 1.77200f*4096.0f+0.5f)); + __m128i y_bias = _mm_set1_epi8((char) (unsigned char) 128); + __m128i xw = _mm_set1_epi16(255); // alpha channel + + for (; i+7 < count; i += 8) { + // load + __m128i y_bytes = _mm_loadl_epi64((__m128i *) (y+i)); + __m128i cr_bytes = _mm_loadl_epi64((__m128i *) (pcr+i)); + __m128i cb_bytes = _mm_loadl_epi64((__m128i *) (pcb+i)); + __m128i cr_biased = _mm_xor_si128(cr_bytes, signflip); // -128 + __m128i cb_biased = _mm_xor_si128(cb_bytes, signflip); // -128 + + // unpack to short (and left-shift cr, cb by 8) + __m128i yw = _mm_unpacklo_epi8(y_bias, y_bytes); + __m128i crw = _mm_unpacklo_epi8(_mm_setzero_si128(), cr_biased); + __m128i cbw = _mm_unpacklo_epi8(_mm_setzero_si128(), cb_biased); + + // color transform + __m128i yws = _mm_srli_epi16(yw, 4); + __m128i cr0 = _mm_mulhi_epi16(cr_const0, crw); + __m128i cb0 = _mm_mulhi_epi16(cb_const0, cbw); + __m128i cb1 = _mm_mulhi_epi16(cbw, cb_const1); + __m128i cr1 = _mm_mulhi_epi16(crw, cr_const1); + __m128i rws = _mm_add_epi16(cr0, yws); + __m128i gwt = _mm_add_epi16(cb0, yws); + __m128i bws = _mm_add_epi16(yws, cb1); + __m128i gws = _mm_add_epi16(gwt, cr1); + + // descale + __m128i rw = _mm_srai_epi16(rws, 4); + __m128i bw = _mm_srai_epi16(bws, 4); + __m128i gw = _mm_srai_epi16(gws, 4); + + // back to byte, set up for transpose + __m128i brb = _mm_packus_epi16(rw, bw); + __m128i gxb = _mm_packus_epi16(gw, xw); + + // transpose to interleave channels + __m128i t0 = _mm_unpacklo_epi8(brb, gxb); + __m128i t1 = _mm_unpackhi_epi8(brb, gxb); + __m128i o0 = _mm_unpacklo_epi16(t0, t1); + __m128i o1 = _mm_unpackhi_epi16(t0, t1); + + // store + _mm_storeu_si128((__m128i *) (out + 0), o0); + _mm_storeu_si128((__m128i *) (out + 16), o1); + out += 32; + } + } +#endif + +#ifdef STBI_NEON + // in this version, step=3 support would be easy to add. but is there demand? + if (step == 4) { + // this is a fairly straightforward implementation and not super-optimized. + uint8x8_t signflip = vdup_n_u8(0x80); + int16x8_t cr_const0 = vdupq_n_s16( (short) ( 1.40200f*4096.0f+0.5f)); + int16x8_t cr_const1 = vdupq_n_s16( - (short) ( 0.71414f*4096.0f+0.5f)); + int16x8_t cb_const0 = vdupq_n_s16( - (short) ( 0.34414f*4096.0f+0.5f)); + int16x8_t cb_const1 = vdupq_n_s16( (short) ( 1.77200f*4096.0f+0.5f)); + + for (; i+7 < count; i += 8) { + // load + uint8x8_t y_bytes = vld1_u8(y + i); + uint8x8_t cr_bytes = vld1_u8(pcr + i); + uint8x8_t cb_bytes = vld1_u8(pcb + i); + int8x8_t cr_biased = vreinterpret_s8_u8(vsub_u8(cr_bytes, signflip)); + int8x8_t cb_biased = vreinterpret_s8_u8(vsub_u8(cb_bytes, signflip)); + + // expand to s16 + int16x8_t yws = vreinterpretq_s16_u16(vshll_n_u8(y_bytes, 4)); + int16x8_t crw = vshll_n_s8(cr_biased, 7); + int16x8_t cbw = vshll_n_s8(cb_biased, 7); + + // color transform + int16x8_t cr0 = vqdmulhq_s16(crw, cr_const0); + int16x8_t cb0 = vqdmulhq_s16(cbw, cb_const0); + int16x8_t cr1 = vqdmulhq_s16(crw, cr_const1); + int16x8_t cb1 = vqdmulhq_s16(cbw, cb_const1); + int16x8_t rws = vaddq_s16(yws, cr0); + int16x8_t gws = vaddq_s16(vaddq_s16(yws, cb0), cr1); + int16x8_t bws = vaddq_s16(yws, cb1); + + // undo scaling, round, convert to byte + uint8x8x4_t o; + o.val[0] = vqrshrun_n_s16(rws, 4); + o.val[1] = vqrshrun_n_s16(gws, 4); + o.val[2] = vqrshrun_n_s16(bws, 4); + o.val[3] = vdup_n_u8(255); + + // store, interleaving r/g/b/a + vst4_u8(out, o); + out += 8*4; + } + } +#endif + + for (; i < count; ++i) { + int y_fixed = (y[i] << 20) + (1<<19); // rounding + int r,g,b; + int cr = pcr[i] - 128; + int cb = pcb[i] - 128; + r = y_fixed + cr* stbi__float2fixed(1.40200f); + g = y_fixed + cr*-stbi__float2fixed(0.71414f) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); + b = y_fixed + cb* stbi__float2fixed(1.77200f); + r >>= 20; + g >>= 20; + b >>= 20; + if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } + if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } + if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } + out[0] = (stbi_uc)r; + out[1] = (stbi_uc)g; + out[2] = (stbi_uc)b; + out[3] = 255; + out += step; + } +} +#endif + +// set up the kernels +static void stbi__setup_jpeg(stbi__jpeg *j) +{ + j->idct_block_kernel = stbi__idct_block; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_row; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2; + +#ifdef STBI_SSE2 + if (stbi__sse2_available()) { + j->idct_block_kernel = stbi__idct_simd; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; + } +#endif + +#ifdef STBI_NEON + j->idct_block_kernel = stbi__idct_simd; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; +#endif +} + +// clean up the temporary component buffers +static void stbi__cleanup_jpeg(stbi__jpeg *j) +{ + stbi__free_jpeg_components(j, j->s->img_n, 0); +} + +typedef struct +{ + resample_row_func resample; + stbi_uc *line0,*line1; + int hs,vs; // expansion factor in each axis + int w_lores; // horizontal pixels pre-expansion + int ystep; // how far through vertical expansion we are + int ypos; // which pre-expansion row we're on +} stbi__resample; + +// fast 0..255 * 0..255 => 0..255 rounded multiplication +static stbi_uc stbi__blinn_8x8(stbi_uc x, stbi_uc y) +{ + unsigned int t = x*y + 128; + return (stbi_uc) ((t + (t >>8)) >> 8); +} + +static stbi_uc *load_jpeg_image(stbi__jpeg *z, int *out_x, int *out_y, int *comp, int req_comp) +{ + int n, decode_n, is_rgb; + z->s->img_n = 0; // make stbi__cleanup_jpeg safe + + // validate req_comp + if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); + + // load a jpeg image from whichever source, but leave in YCbCr format + if (!stbi__decode_jpeg_image(z)) { stbi__cleanup_jpeg(z); return NULL; } + + // determine actual number of components to generate + n = req_comp ? req_comp : z->s->img_n >= 3 ? 3 : 1; + + is_rgb = z->s->img_n == 3 && (z->rgb == 3 || (z->app14_color_transform == 0 && !z->jfif)); + + if (z->s->img_n == 3 && n < 3 && !is_rgb) + decode_n = 1; + else + decode_n = z->s->img_n; + + // nothing to do if no components requested; check this now to avoid + // accessing uninitialized coutput[0] later + if (decode_n <= 0) { stbi__cleanup_jpeg(z); return NULL; } + + // resample and color-convert + { + int k; + unsigned int i,j; + stbi_uc *output; + stbi_uc *coutput[4] = { NULL, NULL, NULL, NULL }; + + stbi__resample res_comp[4]; + + for (k=0; k < decode_n; ++k) { + stbi__resample *r = &res_comp[k]; + + // allocate line buffer big enough for upsampling off the edges + // with upsample factor of 4 + z->img_comp[k].linebuf = (stbi_uc *) stbi__malloc(z->s->img_x + 3); + if (!z->img_comp[k].linebuf) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } + + r->hs = z->img_h_max / z->img_comp[k].h; + r->vs = z->img_v_max / z->img_comp[k].v; + r->ystep = r->vs >> 1; + r->w_lores = (z->s->img_x + r->hs-1) / r->hs; + r->ypos = 0; + r->line0 = r->line1 = z->img_comp[k].data; + + if (r->hs == 1 && r->vs == 1) r->resample = resample_row_1; + else if (r->hs == 1 && r->vs == 2) r->resample = stbi__resample_row_v_2; + else if (r->hs == 2 && r->vs == 1) r->resample = stbi__resample_row_h_2; + else if (r->hs == 2 && r->vs == 2) r->resample = z->resample_row_hv_2_kernel; + else r->resample = stbi__resample_row_generic; + } + + // can't error after this so, this is safe + output = (stbi_uc *) stbi__malloc_mad3(n, z->s->img_x, z->s->img_y, 1); + if (!output) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } + + // now go ahead and resample + for (j=0; j < z->s->img_y; ++j) { + stbi_uc *out = output + n * z->s->img_x * j; + for (k=0; k < decode_n; ++k) { + stbi__resample *r = &res_comp[k]; + int y_bot = r->ystep >= (r->vs >> 1); + coutput[k] = r->resample(z->img_comp[k].linebuf, + y_bot ? r->line1 : r->line0, + y_bot ? r->line0 : r->line1, + r->w_lores, r->hs); + if (++r->ystep >= r->vs) { + r->ystep = 0; + r->line0 = r->line1; + if (++r->ypos < z->img_comp[k].y) + r->line1 += z->img_comp[k].w2; + } + } + if (n >= 3) { + stbi_uc *y = coutput[0]; + if (z->s->img_n == 3) { + if (is_rgb) { + for (i=0; i < z->s->img_x; ++i) { + out[0] = y[i]; + out[1] = coutput[1][i]; + out[2] = coutput[2][i]; + out[3] = 255; + out += n; + } + } else { + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + } + } else if (z->s->img_n == 4) { + if (z->app14_color_transform == 0) { // CMYK + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + out[0] = stbi__blinn_8x8(coutput[0][i], m); + out[1] = stbi__blinn_8x8(coutput[1][i], m); + out[2] = stbi__blinn_8x8(coutput[2][i], m); + out[3] = 255; + out += n; + } + } else if (z->app14_color_transform == 2) { // YCCK + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + out[0] = stbi__blinn_8x8(255 - out[0], m); + out[1] = stbi__blinn_8x8(255 - out[1], m); + out[2] = stbi__blinn_8x8(255 - out[2], m); + out += n; + } + } else { // YCbCr + alpha? Ignore the fourth channel for now + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + } + } else + for (i=0; i < z->s->img_x; ++i) { + out[0] = out[1] = out[2] = y[i]; + out[3] = 255; // not used if n==3 + out += n; + } + } else { + if (is_rgb) { + if (n == 1) + for (i=0; i < z->s->img_x; ++i) + *out++ = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); + else { + for (i=0; i < z->s->img_x; ++i, out += 2) { + out[0] = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); + out[1] = 255; + } + } + } else if (z->s->img_n == 4 && z->app14_color_transform == 0) { + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + stbi_uc r = stbi__blinn_8x8(coutput[0][i], m); + stbi_uc g = stbi__blinn_8x8(coutput[1][i], m); + stbi_uc b = stbi__blinn_8x8(coutput[2][i], m); + out[0] = stbi__compute_y(r, g, b); + out[1] = 255; + out += n; + } + } else if (z->s->img_n == 4 && z->app14_color_transform == 2) { + for (i=0; i < z->s->img_x; ++i) { + out[0] = stbi__blinn_8x8(255 - coutput[0][i], coutput[3][i]); + out[1] = 255; + out += n; + } + } else { + stbi_uc *y = coutput[0]; + if (n == 1) + for (i=0; i < z->s->img_x; ++i) out[i] = y[i]; + else + for (i=0; i < z->s->img_x; ++i) { *out++ = y[i]; *out++ = 255; } + } + } + } + stbi__cleanup_jpeg(z); + *out_x = z->s->img_x; + *out_y = z->s->img_y; + if (comp) *comp = z->s->img_n >= 3 ? 3 : 1; // report original components, not output + return output; + } +} + +static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + unsigned char* result; + stbi__jpeg* j = (stbi__jpeg*) stbi__malloc(sizeof(stbi__jpeg)); + if (!j) return stbi__errpuc("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + STBI_NOTUSED(ri); + j->s = s; + stbi__setup_jpeg(j); + result = load_jpeg_image(j, x,y,comp,req_comp); + STBI_FREE(j); + return result; +} + +static int stbi__jpeg_test(stbi__context *s) +{ + int r; + stbi__jpeg* j = (stbi__jpeg*)stbi__malloc(sizeof(stbi__jpeg)); + if (!j) return stbi__err("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + j->s = s; + stbi__setup_jpeg(j); + r = stbi__decode_jpeg_header(j, STBI__SCAN_type); + stbi__rewind(s); + STBI_FREE(j); + return r; +} + +static int stbi__jpeg_info_raw(stbi__jpeg *j, int *x, int *y, int *comp) +{ + if (!stbi__decode_jpeg_header(j, STBI__SCAN_header)) { + stbi__rewind( j->s ); + return 0; + } + if (x) *x = j->s->img_x; + if (y) *y = j->s->img_y; + if (comp) *comp = j->s->img_n >= 3 ? 3 : 1; + return 1; +} + +static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp) +{ + int result; + stbi__jpeg* j = (stbi__jpeg*) (stbi__malloc(sizeof(stbi__jpeg))); + if (!j) return stbi__err("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + j->s = s; + result = stbi__jpeg_info_raw(j, x, y, comp); + STBI_FREE(j); + return result; +} +#endif + +// public domain zlib decode v0.2 Sean Barrett 2006-11-18 +// simple implementation +// - all input must be provided in an upfront buffer +// - all output is written to a single output buffer (can malloc/realloc) +// performance +// - fast huffman + +#ifndef STBI_NO_ZLIB + +// fast-way is faster to check than jpeg huffman, but slow way is slower +#define STBI__ZFAST_BITS 9 // accelerate all cases in default tables +#define STBI__ZFAST_MASK ((1 << STBI__ZFAST_BITS) - 1) +#define STBI__ZNSYMS 288 // number of symbols in literal/length alphabet + +// zlib-style huffman encoding +// (jpegs packs from left, zlib from right, so can't share code) +typedef struct +{ + stbi__uint16 fast[1 << STBI__ZFAST_BITS]; + stbi__uint16 firstcode[16]; + int maxcode[17]; + stbi__uint16 firstsymbol[16]; + stbi_uc size[STBI__ZNSYMS]; + stbi__uint16 value[STBI__ZNSYMS]; +} stbi__zhuffman; + +stbi_inline static int stbi__bitreverse16(int n) +{ + n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1); + n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2); + n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4); + n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8); + return n; +} + +stbi_inline static int stbi__bit_reverse(int v, int bits) +{ + STBI_ASSERT(bits <= 16); + // to bit reverse n bits, reverse 16 and shift + // e.g. 11 bits, bit reverse and shift away 5 + return stbi__bitreverse16(v) >> (16-bits); +} + +static int stbi__zbuild_huffman(stbi__zhuffman *z, const stbi_uc *sizelist, int num) +{ + int i,k=0; + int code, next_code[16], sizes[17]; + + // DEFLATE spec for generating codes + memset(sizes, 0, sizeof(sizes)); + memset(z->fast, 0, sizeof(z->fast)); + for (i=0; i < num; ++i) + ++sizes[sizelist[i]]; + sizes[0] = 0; + for (i=1; i < 16; ++i) + if (sizes[i] > (1 << i)) + return stbi__err("bad sizes", "Corrupt PNG"); + code = 0; + for (i=1; i < 16; ++i) { + next_code[i] = code; + z->firstcode[i] = (stbi__uint16) code; + z->firstsymbol[i] = (stbi__uint16) k; + code = (code + sizes[i]); + if (sizes[i]) + if (code-1 >= (1 << i)) return stbi__err("bad codelengths","Corrupt PNG"); + z->maxcode[i] = code << (16-i); // preshift for inner loop + code <<= 1; + k += sizes[i]; + } + z->maxcode[16] = 0x10000; // sentinel + for (i=0; i < num; ++i) { + int s = sizelist[i]; + if (s) { + int c = next_code[s] - z->firstcode[s] + z->firstsymbol[s]; + stbi__uint16 fastv = (stbi__uint16) ((s << 9) | i); + z->size [c] = (stbi_uc ) s; + z->value[c] = (stbi__uint16) i; + if (s <= STBI__ZFAST_BITS) { + int j = stbi__bit_reverse(next_code[s],s); + while (j < (1 << STBI__ZFAST_BITS)) { + z->fast[j] = fastv; + j += (1 << s); + } + } + ++next_code[s]; + } + } + return 1; +} + +// zlib-from-memory implementation for PNG reading +// because PNG allows splitting the zlib stream arbitrarily, +// and it's annoying structurally to have PNG call ZLIB call PNG, +// we require PNG read all the IDATs and combine them into a single +// memory buffer + +typedef struct +{ + stbi_uc *zbuffer, *zbuffer_end; + int num_bits; + int hit_zeof_once; + stbi__uint32 code_buffer; + + char *zout; + char *zout_start; + char *zout_end; + int z_expandable; + + stbi__zhuffman z_length, z_distance; +} stbi__zbuf; + +stbi_inline static int stbi__zeof(stbi__zbuf *z) +{ + return (z->zbuffer >= z->zbuffer_end); +} + +stbi_inline static stbi_uc stbi__zget8(stbi__zbuf *z) +{ + return stbi__zeof(z) ? 0 : *z->zbuffer++; +} + +static void stbi__fill_bits(stbi__zbuf *z) +{ + do { + if (z->code_buffer >= (1U << z->num_bits)) { + z->zbuffer = z->zbuffer_end; /* treat this as EOF so we fail. */ + return; + } + z->code_buffer |= (unsigned int) stbi__zget8(z) << z->num_bits; + z->num_bits += 8; + } while (z->num_bits <= 24); +} + +stbi_inline static unsigned int stbi__zreceive(stbi__zbuf *z, int n) +{ + unsigned int k; + if (z->num_bits < n) stbi__fill_bits(z); + k = z->code_buffer & ((1 << n) - 1); + z->code_buffer >>= n; + z->num_bits -= n; + return k; +} + +static int stbi__zhuffman_decode_slowpath(stbi__zbuf *a, stbi__zhuffman *z) +{ + int b,s,k; + // not resolved by fast table, so compute it the slow way + // use jpeg approach, which requires MSbits at top + k = stbi__bit_reverse(a->code_buffer, 16); + for (s=STBI__ZFAST_BITS+1; ; ++s) + if (k < z->maxcode[s]) + break; + if (s >= 16) return -1; // invalid code! + // code size is s, so: + b = (k >> (16-s)) - z->firstcode[s] + z->firstsymbol[s]; + if (b >= STBI__ZNSYMS) return -1; // some data was corrupt somewhere! + if (z->size[b] != s) return -1; // was originally an assert, but report failure instead. + a->code_buffer >>= s; + a->num_bits -= s; + return z->value[b]; +} + +stbi_inline static int stbi__zhuffman_decode(stbi__zbuf *a, stbi__zhuffman *z) +{ + int b,s; + if (a->num_bits < 16) { + if (stbi__zeof(a)) { + if (!a->hit_zeof_once) { + // This is the first time we hit eof, insert 16 extra padding btis + // to allow us to keep going; if we actually consume any of them + // though, that is invalid data. This is caught later. + a->hit_zeof_once = 1; + a->num_bits += 16; // add 16 implicit zero bits + } else { + // We already inserted our extra 16 padding bits and are again + // out, this stream is actually prematurely terminated. + return -1; + } + } else { + stbi__fill_bits(a); + } + } + b = z->fast[a->code_buffer & STBI__ZFAST_MASK]; + if (b) { + s = b >> 9; + a->code_buffer >>= s; + a->num_bits -= s; + return b & 511; + } + return stbi__zhuffman_decode_slowpath(a, z); +} + +static int stbi__zexpand(stbi__zbuf *z, char *zout, int n) // need to make room for n bytes +{ + char *q; + unsigned int cur, limit, old_limit; + z->zout = zout; + if (!z->z_expandable) return stbi__err("output buffer limit","Corrupt PNG"); + cur = (unsigned int) (z->zout - z->zout_start); + limit = old_limit = (unsigned) (z->zout_end - z->zout_start); + if (UINT_MAX - cur < (unsigned) n) return stbi__err("outofmem", "Out of memory"); + while (cur + n > limit) { + if(limit > UINT_MAX / 2) return stbi__err("outofmem", "Out of memory"); + limit *= 2; + } + q = (char *) STBI_REALLOC_SIZED(z->zout_start, old_limit, limit); + STBI_NOTUSED(old_limit); + if (q == NULL) return stbi__err("outofmem", "Out of memory"); + z->zout_start = q; + z->zout = q + cur; + z->zout_end = q + limit; + return 1; +} + +static const int stbi__zlength_base[31] = { + 3,4,5,6,7,8,9,10,11,13, + 15,17,19,23,27,31,35,43,51,59, + 67,83,99,115,131,163,195,227,258,0,0 }; + +static const int stbi__zlength_extra[31]= +{ 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0,0,0 }; + +static const int stbi__zdist_base[32] = { 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, +257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577,0,0}; + +static const int stbi__zdist_extra[32] = +{ 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13}; + +static int stbi__parse_huffman_block(stbi__zbuf *a) +{ + char *zout = a->zout; + for(;;) { + int z = stbi__zhuffman_decode(a, &a->z_length); + if (z < 256) { + if (z < 0) return stbi__err("bad huffman code","Corrupt PNG"); // error in huffman codes + if (zout >= a->zout_end) { + if (!stbi__zexpand(a, zout, 1)) return 0; + zout = a->zout; + } + *zout++ = (char) z; + } else { + stbi_uc *p; + int len,dist; + if (z == 256) { + a->zout = zout; + if (a->hit_zeof_once && a->num_bits < 16) { + // The first time we hit zeof, we inserted 16 extra zero bits into our bit + // buffer so the decoder can just do its speculative decoding. But if we + // actually consumed any of those bits (which is the case when num_bits < 16), + // the stream actually read past the end so it is malformed. + return stbi__err("unexpected end","Corrupt PNG"); + } + return 1; + } + if (z >= 286) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, length codes 286 and 287 must not appear in compressed data + z -= 257; + len = stbi__zlength_base[z]; + if (stbi__zlength_extra[z]) len += stbi__zreceive(a, stbi__zlength_extra[z]); + z = stbi__zhuffman_decode(a, &a->z_distance); + if (z < 0 || z >= 30) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, distance codes 30 and 31 must not appear in compressed data + dist = stbi__zdist_base[z]; + if (stbi__zdist_extra[z]) dist += stbi__zreceive(a, stbi__zdist_extra[z]); + if (zout - a->zout_start < dist) return stbi__err("bad dist","Corrupt PNG"); + if (len > a->zout_end - zout) { + if (!stbi__zexpand(a, zout, len)) return 0; + zout = a->zout; + } + p = (stbi_uc *) (zout - dist); + if (dist == 1) { // run of one byte; common in images. + stbi_uc v = *p; + if (len) { do *zout++ = v; while (--len); } + } else { + if (len) { do *zout++ = *p++; while (--len); } + } + } + } +} + +static int stbi__compute_huffman_codes(stbi__zbuf *a) +{ + static const stbi_uc length_dezigzag[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; + stbi__zhuffman z_codelength; + stbi_uc lencodes[286+32+137];//padding for maximum single op + stbi_uc codelength_sizes[19]; + int i,n; + + int hlit = stbi__zreceive(a,5) + 257; + int hdist = stbi__zreceive(a,5) + 1; + int hclen = stbi__zreceive(a,4) + 4; + int ntot = hlit + hdist; + + memset(codelength_sizes, 0, sizeof(codelength_sizes)); + for (i=0; i < hclen; ++i) { + int s = stbi__zreceive(a,3); + codelength_sizes[length_dezigzag[i]] = (stbi_uc) s; + } + if (!stbi__zbuild_huffman(&z_codelength, codelength_sizes, 19)) return 0; + + n = 0; + while (n < ntot) { + int c = stbi__zhuffman_decode(a, &z_codelength); + if (c < 0 || c >= 19) return stbi__err("bad codelengths", "Corrupt PNG"); + if (c < 16) + lencodes[n++] = (stbi_uc) c; + else { + stbi_uc fill = 0; + if (c == 16) { + c = stbi__zreceive(a,2)+3; + if (n == 0) return stbi__err("bad codelengths", "Corrupt PNG"); + fill = lencodes[n-1]; + } else if (c == 17) { + c = stbi__zreceive(a,3)+3; + } else if (c == 18) { + c = stbi__zreceive(a,7)+11; + } else { + return stbi__err("bad codelengths", "Corrupt PNG"); + } + if (ntot - n < c) return stbi__err("bad codelengths", "Corrupt PNG"); + memset(lencodes+n, fill, c); + n += c; + } + } + if (n != ntot) return stbi__err("bad codelengths","Corrupt PNG"); + if (!stbi__zbuild_huffman(&a->z_length, lencodes, hlit)) return 0; + if (!stbi__zbuild_huffman(&a->z_distance, lencodes+hlit, hdist)) return 0; + return 1; +} + +static int stbi__parse_uncompressed_block(stbi__zbuf *a) +{ + stbi_uc header[4]; + int len,nlen,k; + if (a->num_bits & 7) + stbi__zreceive(a, a->num_bits & 7); // discard + // drain the bit-packed data into header + k = 0; + while (a->num_bits > 0) { + header[k++] = (stbi_uc) (a->code_buffer & 255); // suppress MSVC run-time check + a->code_buffer >>= 8; + a->num_bits -= 8; + } + if (a->num_bits < 0) return stbi__err("zlib corrupt","Corrupt PNG"); + // now fill header the normal way + while (k < 4) + header[k++] = stbi__zget8(a); + len = header[1] * 256 + header[0]; + nlen = header[3] * 256 + header[2]; + if (nlen != (len ^ 0xffff)) return stbi__err("zlib corrupt","Corrupt PNG"); + if (a->zbuffer + len > a->zbuffer_end) return stbi__err("read past buffer","Corrupt PNG"); + if (a->zout + len > a->zout_end) + if (!stbi__zexpand(a, a->zout, len)) return 0; + memcpy(a->zout, a->zbuffer, len); + a->zbuffer += len; + a->zout += len; + return 1; +} + +static int stbi__parse_zlib_header(stbi__zbuf *a) +{ + int cmf = stbi__zget8(a); + int cm = cmf & 15; + /* int cinfo = cmf >> 4; */ + int flg = stbi__zget8(a); + if (stbi__zeof(a)) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec + if ((cmf*256+flg) % 31 != 0) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec + if (flg & 32) return stbi__err("no preset dict","Corrupt PNG"); // preset dictionary not allowed in png + if (cm != 8) return stbi__err("bad compression","Corrupt PNG"); // DEFLATE required for png + // window = 1 << (8 + cinfo)... but who cares, we fully buffer output + return 1; +} + +static const stbi_uc stbi__zdefault_length[STBI__ZNSYMS] = +{ + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8 +}; +static const stbi_uc stbi__zdefault_distance[32] = +{ + 5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5 +}; +/* +Init algorithm: +{ + int i; // use <= to match clearly with spec + for (i=0; i <= 143; ++i) stbi__zdefault_length[i] = 8; + for ( ; i <= 255; ++i) stbi__zdefault_length[i] = 9; + for ( ; i <= 279; ++i) stbi__zdefault_length[i] = 7; + for ( ; i <= 287; ++i) stbi__zdefault_length[i] = 8; + + for (i=0; i <= 31; ++i) stbi__zdefault_distance[i] = 5; +} +*/ + +static int stbi__parse_zlib(stbi__zbuf *a, int parse_header) +{ + int final, type; + if (parse_header) + if (!stbi__parse_zlib_header(a)) return 0; + a->num_bits = 0; + a->code_buffer = 0; + a->hit_zeof_once = 0; + do { + final = stbi__zreceive(a,1); + type = stbi__zreceive(a,2); + if (type == 0) { + if (!stbi__parse_uncompressed_block(a)) return 0; + } else if (type == 3) { + return 0; + } else { + if (type == 1) { + // use fixed code lengths + if (!stbi__zbuild_huffman(&a->z_length , stbi__zdefault_length , STBI__ZNSYMS)) return 0; + if (!stbi__zbuild_huffman(&a->z_distance, stbi__zdefault_distance, 32)) return 0; + } else { + if (!stbi__compute_huffman_codes(a)) return 0; + } + if (!stbi__parse_huffman_block(a)) return 0; + } + } while (!final); + return 1; +} + +static int stbi__do_zlib(stbi__zbuf *a, char *obuf, int olen, int exp, int parse_header) +{ + a->zout_start = obuf; + a->zout = obuf; + a->zout_end = obuf + olen; + a->z_expandable = exp; + + return stbi__parse_zlib(a, parse_header); +} + +STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(initial_size); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer + len; + if (stbi__do_zlib(&a, p, initial_size, 1, 1)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF char *stbi_zlib_decode_malloc(char const *buffer, int len, int *outlen) +{ + return stbi_zlib_decode_malloc_guesssize(buffer, len, 16384, outlen); +} + +STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(initial_size); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer + len; + if (stbi__do_zlib(&a, p, initial_size, 1, parse_header)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, char const *ibuffer, int ilen) +{ + stbi__zbuf a; + a.zbuffer = (stbi_uc *) ibuffer; + a.zbuffer_end = (stbi_uc *) ibuffer + ilen; + if (stbi__do_zlib(&a, obuffer, olen, 0, 1)) + return (int) (a.zout - a.zout_start); + else + return -1; +} + +STBIDEF char *stbi_zlib_decode_noheader_malloc(char const *buffer, int len, int *outlen) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(16384); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer+len; + if (stbi__do_zlib(&a, p, 16384, 1, 0)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen) +{ + stbi__zbuf a; + a.zbuffer = (stbi_uc *) ibuffer; + a.zbuffer_end = (stbi_uc *) ibuffer + ilen; + if (stbi__do_zlib(&a, obuffer, olen, 0, 0)) + return (int) (a.zout - a.zout_start); + else + return -1; +} +#endif + +// public domain "baseline" PNG decoder v0.10 Sean Barrett 2006-11-18 +// simple implementation +// - only 8-bit samples +// - no CRC checking +// - allocates lots of intermediate memory +// - avoids problem of streaming data between subsystems +// - avoids explicit window management +// performance +// - uses stb_zlib, a PD zlib implementation with fast huffman decoding + +#ifndef STBI_NO_PNG +typedef struct +{ + stbi__uint32 length; + stbi__uint32 type; +} stbi__pngchunk; + +static stbi__pngchunk stbi__get_chunk_header(stbi__context *s) +{ + stbi__pngchunk c; + c.length = stbi__get32be(s); + c.type = stbi__get32be(s); + return c; +} + +static int stbi__check_png_header(stbi__context *s) +{ + static const stbi_uc png_sig[8] = { 137,80,78,71,13,10,26,10 }; + int i; + for (i=0; i < 8; ++i) + if (stbi__get8(s) != png_sig[i]) return stbi__err("bad png sig","Not a PNG"); + return 1; +} + +typedef struct +{ + stbi__context *s; + stbi_uc *idata, *expanded, *out; + int depth; +} stbi__png; + + +enum { + STBI__F_none=0, + STBI__F_sub=1, + STBI__F_up=2, + STBI__F_avg=3, + STBI__F_paeth=4, + // synthetic filter used for first scanline to avoid needing a dummy row of 0s + STBI__F_avg_first +}; + +static stbi_uc first_row_filter[5] = +{ + STBI__F_none, + STBI__F_sub, + STBI__F_none, + STBI__F_avg_first, + STBI__F_sub // Paeth with b=c=0 turns out to be equivalent to sub +}; + +static int stbi__paeth(int a, int b, int c) +{ + // This formulation looks very different from the reference in the PNG spec, but is + // actually equivalent and has favorable data dependencies and admits straightforward + // generation of branch-free code, which helps performance significantly. + int thresh = c*3 - (a + b); + int lo = a < b ? a : b; + int hi = a < b ? b : a; + int t0 = (hi <= thresh) ? lo : c; + int t1 = (thresh <= lo) ? hi : t0; + return t1; +} + +static const stbi_uc stbi__depth_scale_table[9] = { 0, 0xff, 0x55, 0, 0x11, 0,0,0, 0x01 }; + +// adds an extra all-255 alpha channel +// dest == src is legal +// img_n must be 1 or 3 +static void stbi__create_png_alpha_expand8(stbi_uc *dest, stbi_uc *src, stbi__uint32 x, int img_n) +{ + int i; + // must process data backwards since we allow dest==src + if (img_n == 1) { + for (i=x-1; i >= 0; --i) { + dest[i*2+1] = 255; + dest[i*2+0] = src[i]; + } + } else { + STBI_ASSERT(img_n == 3); + for (i=x-1; i >= 0; --i) { + dest[i*4+3] = 255; + dest[i*4+2] = src[i*3+2]; + dest[i*4+1] = src[i*3+1]; + dest[i*4+0] = src[i*3+0]; + } + } +} + +// create the png data from post-deflated data +static int stbi__create_png_image_raw(stbi__png *a, stbi_uc *raw, stbi__uint32 raw_len, int out_n, stbi__uint32 x, stbi__uint32 y, int depth, int color) +{ + int bytes = (depth == 16 ? 2 : 1); + stbi__context *s = a->s; + stbi__uint32 i,j,stride = x*out_n*bytes; + stbi__uint32 img_len, img_width_bytes; + stbi_uc *filter_buf; + int all_ok = 1; + int k; + int img_n = s->img_n; // copy it into a local for later + + int output_bytes = out_n*bytes; + int filter_bytes = img_n*bytes; + int width = x; + + STBI_ASSERT(out_n == s->img_n || out_n == s->img_n+1); + a->out = (stbi_uc *) stbi__malloc_mad3(x, y, output_bytes, 0); // extra bytes to write off the end into + if (!a->out) return stbi__err("outofmem", "Out of memory"); + + // note: error exits here don't need to clean up a->out individually, + // stbi__do_png always does on error. + if (!stbi__mad3sizes_valid(img_n, x, depth, 7)) return stbi__err("too large", "Corrupt PNG"); + img_width_bytes = (((img_n * x * depth) + 7) >> 3); + if (!stbi__mad2sizes_valid(img_width_bytes, y, img_width_bytes)) return stbi__err("too large", "Corrupt PNG"); + img_len = (img_width_bytes + 1) * y; + + // we used to check for exact match between raw_len and img_len on non-interlaced PNGs, + // but issue #276 reported a PNG in the wild that had extra data at the end (all zeros), + // so just check for raw_len < img_len always. + if (raw_len < img_len) return stbi__err("not enough pixels","Corrupt PNG"); + + // Allocate two scan lines worth of filter workspace buffer. + filter_buf = (stbi_uc *) stbi__malloc_mad2(img_width_bytes, 2, 0); + if (!filter_buf) return stbi__err("outofmem", "Out of memory"); + + // Filtering for low-bit-depth images + if (depth < 8) { + filter_bytes = 1; + width = img_width_bytes; + } + + for (j=0; j < y; ++j) { + // cur/prior filter buffers alternate + stbi_uc *cur = filter_buf + (j & 1)*img_width_bytes; + stbi_uc *prior = filter_buf + (~j & 1)*img_width_bytes; + stbi_uc *dest = a->out + stride*j; + int nk = width * filter_bytes; + int filter = *raw++; + + // check filter type + if (filter > 4) { + all_ok = stbi__err("invalid filter","Corrupt PNG"); + break; + } + + // if first row, use special filter that doesn't sample previous row + if (j == 0) filter = first_row_filter[filter]; + + // perform actual filtering + switch (filter) { + case STBI__F_none: + memcpy(cur, raw, nk); + break; + case STBI__F_sub: + memcpy(cur, raw, filter_bytes); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + cur[k-filter_bytes]); + break; + case STBI__F_up: + for (k = 0; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + prior[k]); + break; + case STBI__F_avg: + for (k = 0; k < filter_bytes; ++k) + cur[k] = STBI__BYTECAST(raw[k] + (prior[k]>>1)); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + ((prior[k] + cur[k-filter_bytes])>>1)); + break; + case STBI__F_paeth: + for (k = 0; k < filter_bytes; ++k) + cur[k] = STBI__BYTECAST(raw[k] + prior[k]); // prior[k] == stbi__paeth(0,prior[k],0) + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + stbi__paeth(cur[k-filter_bytes], prior[k], prior[k-filter_bytes])); + break; + case STBI__F_avg_first: + memcpy(cur, raw, filter_bytes); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + (cur[k-filter_bytes] >> 1)); + break; + } + + raw += nk; + + // expand decoded bits in cur to dest, also adding an extra alpha channel if desired + if (depth < 8) { + stbi_uc scale = (color == 0) ? stbi__depth_scale_table[depth] : 1; // scale grayscale values to 0..255 range + stbi_uc *in = cur; + stbi_uc *out = dest; + stbi_uc inb = 0; + stbi__uint32 nsmp = x*img_n; + + // expand bits to bytes first + if (depth == 4) { + for (i=0; i < nsmp; ++i) { + if ((i & 1) == 0) inb = *in++; + *out++ = scale * (inb >> 4); + inb <<= 4; + } + } else if (depth == 2) { + for (i=0; i < nsmp; ++i) { + if ((i & 3) == 0) inb = *in++; + *out++ = scale * (inb >> 6); + inb <<= 2; + } + } else { + STBI_ASSERT(depth == 1); + for (i=0; i < nsmp; ++i) { + if ((i & 7) == 0) inb = *in++; + *out++ = scale * (inb >> 7); + inb <<= 1; + } + } + + // insert alpha=255 values if desired + if (img_n != out_n) + stbi__create_png_alpha_expand8(dest, dest, x, img_n); + } else if (depth == 8) { + if (img_n == out_n) + memcpy(dest, cur, x*img_n); + else + stbi__create_png_alpha_expand8(dest, cur, x, img_n); + } else if (depth == 16) { + // convert the image data from big-endian to platform-native + stbi__uint16 *dest16 = (stbi__uint16*)dest; + stbi__uint32 nsmp = x*img_n; + + if (img_n == out_n) { + for (i = 0; i < nsmp; ++i, ++dest16, cur += 2) + *dest16 = (cur[0] << 8) | cur[1]; + } else { + STBI_ASSERT(img_n+1 == out_n); + if (img_n == 1) { + for (i = 0; i < x; ++i, dest16 += 2, cur += 2) { + dest16[0] = (cur[0] << 8) | cur[1]; + dest16[1] = 0xffff; + } + } else { + STBI_ASSERT(img_n == 3); + for (i = 0; i < x; ++i, dest16 += 4, cur += 6) { + dest16[0] = (cur[0] << 8) | cur[1]; + dest16[1] = (cur[2] << 8) | cur[3]; + dest16[2] = (cur[4] << 8) | cur[5]; + dest16[3] = 0xffff; + } + } + } + } + } + + STBI_FREE(filter_buf); + if (!all_ok) return 0; + + return 1; +} + +static int stbi__create_png_image(stbi__png *a, stbi_uc *image_data, stbi__uint32 image_data_len, int out_n, int depth, int color, int interlaced) +{ + int bytes = (depth == 16 ? 2 : 1); + int out_bytes = out_n * bytes; + stbi_uc *final; + int p; + if (!interlaced) + return stbi__create_png_image_raw(a, image_data, image_data_len, out_n, a->s->img_x, a->s->img_y, depth, color); + + // de-interlacing + final = (stbi_uc *) stbi__malloc_mad3(a->s->img_x, a->s->img_y, out_bytes, 0); + if (!final) return stbi__err("outofmem", "Out of memory"); + for (p=0; p < 7; ++p) { + int xorig[] = { 0,4,0,2,0,1,0 }; + int yorig[] = { 0,0,4,0,2,0,1 }; + int xspc[] = { 8,8,4,4,2,2,1 }; + int yspc[] = { 8,8,8,4,4,2,2 }; + int i,j,x,y; + // pass1_x[4] = 0, pass1_x[5] = 1, pass1_x[12] = 1 + x = (a->s->img_x - xorig[p] + xspc[p]-1) / xspc[p]; + y = (a->s->img_y - yorig[p] + yspc[p]-1) / yspc[p]; + if (x && y) { + stbi__uint32 img_len = ((((a->s->img_n * x * depth) + 7) >> 3) + 1) * y; + if (!stbi__create_png_image_raw(a, image_data, image_data_len, out_n, x, y, depth, color)) { + STBI_FREE(final); + return 0; + } + for (j=0; j < y; ++j) { + for (i=0; i < x; ++i) { + int out_y = j*yspc[p]+yorig[p]; + int out_x = i*xspc[p]+xorig[p]; + memcpy(final + out_y*a->s->img_x*out_bytes + out_x*out_bytes, + a->out + (j*x+i)*out_bytes, out_bytes); + } + } + STBI_FREE(a->out); + image_data += img_len; + image_data_len -= img_len; + } + } + a->out = final; + + return 1; +} + +static int stbi__compute_transparency(stbi__png *z, stbi_uc tc[3], int out_n) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi_uc *p = z->out; + + // compute color-based transparency, assuming we've + // already got 255 as the alpha value in the output + STBI_ASSERT(out_n == 2 || out_n == 4); + + if (out_n == 2) { + for (i=0; i < pixel_count; ++i) { + p[1] = (p[0] == tc[0] ? 0 : 255); + p += 2; + } + } else { + for (i=0; i < pixel_count; ++i) { + if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) + p[3] = 0; + p += 4; + } + } + return 1; +} + +static int stbi__compute_transparency16(stbi__png *z, stbi__uint16 tc[3], int out_n) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi__uint16 *p = (stbi__uint16*) z->out; + + // compute color-based transparency, assuming we've + // already got 65535 as the alpha value in the output + STBI_ASSERT(out_n == 2 || out_n == 4); + + if (out_n == 2) { + for (i = 0; i < pixel_count; ++i) { + p[1] = (p[0] == tc[0] ? 0 : 65535); + p += 2; + } + } else { + for (i = 0; i < pixel_count; ++i) { + if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) + p[3] = 0; + p += 4; + } + } + return 1; +} + +static int stbi__expand_png_palette(stbi__png *a, stbi_uc *palette, int len, int pal_img_n) +{ + stbi__uint32 i, pixel_count = a->s->img_x * a->s->img_y; + stbi_uc *p, *temp_out, *orig = a->out; + + p = (stbi_uc *) stbi__malloc_mad2(pixel_count, pal_img_n, 0); + if (p == NULL) return stbi__err("outofmem", "Out of memory"); + + // between here and free(out) below, exitting would leak + temp_out = p; + + if (pal_img_n == 3) { + for (i=0; i < pixel_count; ++i) { + int n = orig[i]*4; + p[0] = palette[n ]; + p[1] = palette[n+1]; + p[2] = palette[n+2]; + p += 3; + } + } else { + for (i=0; i < pixel_count; ++i) { + int n = orig[i]*4; + p[0] = palette[n ]; + p[1] = palette[n+1]; + p[2] = palette[n+2]; + p[3] = palette[n+3]; + p += 4; + } + } + STBI_FREE(a->out); + a->out = temp_out; + + STBI_NOTUSED(len); + + return 1; +} + +static int stbi__unpremultiply_on_load_global = 0; +static int stbi__de_iphone_flag_global = 0; + +STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply) +{ + stbi__unpremultiply_on_load_global = flag_true_if_should_unpremultiply; +} + +STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert) +{ + stbi__de_iphone_flag_global = flag_true_if_should_convert; +} + +#ifndef STBI_THREAD_LOCAL +#define stbi__unpremultiply_on_load stbi__unpremultiply_on_load_global +#define stbi__de_iphone_flag stbi__de_iphone_flag_global +#else +static STBI_THREAD_LOCAL int stbi__unpremultiply_on_load_local, stbi__unpremultiply_on_load_set; +static STBI_THREAD_LOCAL int stbi__de_iphone_flag_local, stbi__de_iphone_flag_set; + +STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply) +{ + stbi__unpremultiply_on_load_local = flag_true_if_should_unpremultiply; + stbi__unpremultiply_on_load_set = 1; +} + +STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert) +{ + stbi__de_iphone_flag_local = flag_true_if_should_convert; + stbi__de_iphone_flag_set = 1; +} + +#define stbi__unpremultiply_on_load (stbi__unpremultiply_on_load_set \ + ? stbi__unpremultiply_on_load_local \ + : stbi__unpremultiply_on_load_global) +#define stbi__de_iphone_flag (stbi__de_iphone_flag_set \ + ? stbi__de_iphone_flag_local \ + : stbi__de_iphone_flag_global) +#endif // STBI_THREAD_LOCAL + +static void stbi__de_iphone(stbi__png *z) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi_uc *p = z->out; + + if (s->img_out_n == 3) { // convert bgr to rgb + for (i=0; i < pixel_count; ++i) { + stbi_uc t = p[0]; + p[0] = p[2]; + p[2] = t; + p += 3; + } + } else { + STBI_ASSERT(s->img_out_n == 4); + if (stbi__unpremultiply_on_load) { + // convert bgr to rgb and unpremultiply + for (i=0; i < pixel_count; ++i) { + stbi_uc a = p[3]; + stbi_uc t = p[0]; + if (a) { + stbi_uc half = a / 2; + p[0] = (p[2] * 255 + half) / a; + p[1] = (p[1] * 255 + half) / a; + p[2] = ( t * 255 + half) / a; + } else { + p[0] = p[2]; + p[2] = t; + } + p += 4; + } + } else { + // convert bgr to rgb + for (i=0; i < pixel_count; ++i) { + stbi_uc t = p[0]; + p[0] = p[2]; + p[2] = t; + p += 4; + } + } + } +} + +#define STBI__PNG_TYPE(a,b,c,d) (((unsigned) (a) << 24) + ((unsigned) (b) << 16) + ((unsigned) (c) << 8) + (unsigned) (d)) + +static int stbi__parse_png_file(stbi__png *z, int scan, int req_comp) +{ + stbi_uc palette[1024], pal_img_n=0; + stbi_uc has_trans=0, tc[3]={0}; + stbi__uint16 tc16[3]; + stbi__uint32 ioff=0, idata_limit=0, i, pal_len=0; + int first=1,k,interlace=0, color=0, is_iphone=0; + stbi__context *s = z->s; + + z->expanded = NULL; + z->idata = NULL; + z->out = NULL; + + if (!stbi__check_png_header(s)) return 0; + + if (scan == STBI__SCAN_type) return 1; + + for (;;) { + stbi__pngchunk c = stbi__get_chunk_header(s); + switch (c.type) { + case STBI__PNG_TYPE('C','g','B','I'): + is_iphone = 1; + stbi__skip(s, c.length); + break; + case STBI__PNG_TYPE('I','H','D','R'): { + int comp,filter; + if (!first) return stbi__err("multiple IHDR","Corrupt PNG"); + first = 0; + if (c.length != 13) return stbi__err("bad IHDR len","Corrupt PNG"); + s->img_x = stbi__get32be(s); + s->img_y = stbi__get32be(s); + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + z->depth = stbi__get8(s); if (z->depth != 1 && z->depth != 2 && z->depth != 4 && z->depth != 8 && z->depth != 16) return stbi__err("1/2/4/8/16-bit only","PNG not supported: 1/2/4/8/16-bit only"); + color = stbi__get8(s); if (color > 6) return stbi__err("bad ctype","Corrupt PNG"); + if (color == 3 && z->depth == 16) return stbi__err("bad ctype","Corrupt PNG"); + if (color == 3) pal_img_n = 3; else if (color & 1) return stbi__err("bad ctype","Corrupt PNG"); + comp = stbi__get8(s); if (comp) return stbi__err("bad comp method","Corrupt PNG"); + filter= stbi__get8(s); if (filter) return stbi__err("bad filter method","Corrupt PNG"); + interlace = stbi__get8(s); if (interlace>1) return stbi__err("bad interlace method","Corrupt PNG"); + if (!s->img_x || !s->img_y) return stbi__err("0-pixel image","Corrupt PNG"); + if (!pal_img_n) { + s->img_n = (color & 2 ? 3 : 1) + (color & 4 ? 1 : 0); + if ((1 << 30) / s->img_x / s->img_n < s->img_y) return stbi__err("too large", "Image too large to decode"); + } else { + // if paletted, then pal_n is our final components, and + // img_n is # components to decompress/filter. + s->img_n = 1; + if ((1 << 30) / s->img_x / 4 < s->img_y) return stbi__err("too large","Corrupt PNG"); + } + // even with SCAN_header, have to scan to see if we have a tRNS + break; + } + + case STBI__PNG_TYPE('P','L','T','E'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (c.length > 256*3) return stbi__err("invalid PLTE","Corrupt PNG"); + pal_len = c.length / 3; + if (pal_len * 3 != c.length) return stbi__err("invalid PLTE","Corrupt PNG"); + for (i=0; i < pal_len; ++i) { + palette[i*4+0] = stbi__get8(s); + palette[i*4+1] = stbi__get8(s); + palette[i*4+2] = stbi__get8(s); + palette[i*4+3] = 255; + } + break; + } + + case STBI__PNG_TYPE('t','R','N','S'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (z->idata) return stbi__err("tRNS after IDAT","Corrupt PNG"); + if (pal_img_n) { + if (scan == STBI__SCAN_header) { s->img_n = 4; return 1; } + if (pal_len == 0) return stbi__err("tRNS before PLTE","Corrupt PNG"); + if (c.length > pal_len) return stbi__err("bad tRNS len","Corrupt PNG"); + pal_img_n = 4; + for (i=0; i < c.length; ++i) + palette[i*4+3] = stbi__get8(s); + } else { + if (!(s->img_n & 1)) return stbi__err("tRNS with alpha","Corrupt PNG"); + if (c.length != (stbi__uint32) s->img_n*2) return stbi__err("bad tRNS len","Corrupt PNG"); + has_trans = 1; + // non-paletted with tRNS = constant alpha. if header-scanning, we can stop now. + if (scan == STBI__SCAN_header) { ++s->img_n; return 1; } + if (z->depth == 16) { + for (k = 0; k < s->img_n && k < 3; ++k) // extra loop test to suppress false GCC warning + tc16[k] = (stbi__uint16)stbi__get16be(s); // copy the values as-is + } else { + for (k = 0; k < s->img_n && k < 3; ++k) + tc[k] = (stbi_uc)(stbi__get16be(s) & 255) * stbi__depth_scale_table[z->depth]; // non 8-bit images will be larger + } + } + break; + } + + case STBI__PNG_TYPE('I','D','A','T'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (pal_img_n && !pal_len) return stbi__err("no PLTE","Corrupt PNG"); + if (scan == STBI__SCAN_header) { + // header scan definitely stops at first IDAT + if (pal_img_n) + s->img_n = pal_img_n; + return 1; + } + if (c.length > (1u << 30)) return stbi__err("IDAT size limit", "IDAT section larger than 2^30 bytes"); + if ((int)(ioff + c.length) < (int)ioff) return 0; + if (ioff + c.length > idata_limit) { + stbi__uint32 idata_limit_old = idata_limit; + stbi_uc *p; + if (idata_limit == 0) idata_limit = c.length > 4096 ? c.length : 4096; + while (ioff + c.length > idata_limit) + idata_limit *= 2; + STBI_NOTUSED(idata_limit_old); + p = (stbi_uc *) STBI_REALLOC_SIZED(z->idata, idata_limit_old, idata_limit); if (p == NULL) return stbi__err("outofmem", "Out of memory"); + z->idata = p; + } + if (!stbi__getn(s, z->idata+ioff,c.length)) return stbi__err("outofdata","Corrupt PNG"); + ioff += c.length; + break; + } + + case STBI__PNG_TYPE('I','E','N','D'): { + stbi__uint32 raw_len, bpl; + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (scan != STBI__SCAN_load) return 1; + if (z->idata == NULL) return stbi__err("no IDAT","Corrupt PNG"); + // initial guess for decoded data size to avoid unnecessary reallocs + bpl = (s->img_x * z->depth + 7) / 8; // bytes per line, per component + raw_len = bpl * s->img_y * s->img_n /* pixels */ + s->img_y /* filter mode per row */; + z->expanded = (stbi_uc *) stbi_zlib_decode_malloc_guesssize_headerflag((char *) z->idata, ioff, raw_len, (int *) &raw_len, !is_iphone); + if (z->expanded == NULL) return 0; // zlib should set error + STBI_FREE(z->idata); z->idata = NULL; + if ((req_comp == s->img_n+1 && req_comp != 3 && !pal_img_n) || has_trans) + s->img_out_n = s->img_n+1; + else + s->img_out_n = s->img_n; + if (!stbi__create_png_image(z, z->expanded, raw_len, s->img_out_n, z->depth, color, interlace)) return 0; + if (has_trans) { + if (z->depth == 16) { + if (!stbi__compute_transparency16(z, tc16, s->img_out_n)) return 0; + } else { + if (!stbi__compute_transparency(z, tc, s->img_out_n)) return 0; + } + } + if (is_iphone && stbi__de_iphone_flag && s->img_out_n > 2) + stbi__de_iphone(z); + if (pal_img_n) { + // pal_img_n == 3 or 4 + s->img_n = pal_img_n; // record the actual colors we had + s->img_out_n = pal_img_n; + if (req_comp >= 3) s->img_out_n = req_comp; + if (!stbi__expand_png_palette(z, palette, pal_len, s->img_out_n)) + return 0; + } else if (has_trans) { + // non-paletted image with tRNS -> source image has (constant) alpha + ++s->img_n; + } + STBI_FREE(z->expanded); z->expanded = NULL; + // end of PNG chunk, read and skip CRC + stbi__get32be(s); + return 1; + } + + default: + // if critical, fail + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if ((c.type & (1 << 29)) == 0) { + #ifndef STBI_NO_FAILURE_STRINGS + // not threadsafe + static char invalid_chunk[] = "XXXX PNG chunk not known"; + invalid_chunk[0] = STBI__BYTECAST(c.type >> 24); + invalid_chunk[1] = STBI__BYTECAST(c.type >> 16); + invalid_chunk[2] = STBI__BYTECAST(c.type >> 8); + invalid_chunk[3] = STBI__BYTECAST(c.type >> 0); + #endif + return stbi__err(invalid_chunk, "PNG not supported: unknown PNG chunk type"); + } + stbi__skip(s, c.length); + break; + } + // end of PNG chunk, read and skip CRC + stbi__get32be(s); + } +} + +static void *stbi__do_png(stbi__png *p, int *x, int *y, int *n, int req_comp, stbi__result_info *ri) +{ + void *result=NULL; + if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); + if (stbi__parse_png_file(p, STBI__SCAN_load, req_comp)) { + if (p->depth <= 8) + ri->bits_per_channel = 8; + else if (p->depth == 16) + ri->bits_per_channel = 16; + else + return stbi__errpuc("bad bits_per_channel", "PNG not supported: unsupported color depth"); + result = p->out; + p->out = NULL; + if (req_comp && req_comp != p->s->img_out_n) { + if (ri->bits_per_channel == 8) + result = stbi__convert_format((unsigned char *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); + else + result = stbi__convert_format16((stbi__uint16 *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); + p->s->img_out_n = req_comp; + if (result == NULL) return result; + } + *x = p->s->img_x; + *y = p->s->img_y; + if (n) *n = p->s->img_n; + } + STBI_FREE(p->out); p->out = NULL; + STBI_FREE(p->expanded); p->expanded = NULL; + STBI_FREE(p->idata); p->idata = NULL; + + return result; +} + +static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi__png p; + p.s = s; + return stbi__do_png(&p, x,y,comp,req_comp, ri); +} + +static int stbi__png_test(stbi__context *s) +{ + int r; + r = stbi__check_png_header(s); + stbi__rewind(s); + return r; +} + +static int stbi__png_info_raw(stbi__png *p, int *x, int *y, int *comp) +{ + if (!stbi__parse_png_file(p, STBI__SCAN_header, 0)) { + stbi__rewind( p->s ); + return 0; + } + if (x) *x = p->s->img_x; + if (y) *y = p->s->img_y; + if (comp) *comp = p->s->img_n; + return 1; +} + +static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp) +{ + stbi__png p; + p.s = s; + return stbi__png_info_raw(&p, x, y, comp); +} + +static int stbi__png_is16(stbi__context *s) +{ + stbi__png p; + p.s = s; + if (!stbi__png_info_raw(&p, NULL, NULL, NULL)) + return 0; + if (p.depth != 16) { + stbi__rewind(p.s); + return 0; + } + return 1; +} +#endif + +// Microsoft/Windows BMP image + +#ifndef STBI_NO_BMP +static int stbi__bmp_test_raw(stbi__context *s) +{ + int r; + int sz; + if (stbi__get8(s) != 'B') return 0; + if (stbi__get8(s) != 'M') return 0; + stbi__get32le(s); // discard filesize + stbi__get16le(s); // discard reserved + stbi__get16le(s); // discard reserved + stbi__get32le(s); // discard data offset + sz = stbi__get32le(s); + r = (sz == 12 || sz == 40 || sz == 56 || sz == 108 || sz == 124); + return r; +} + +static int stbi__bmp_test(stbi__context *s) +{ + int r = stbi__bmp_test_raw(s); + stbi__rewind(s); + return r; +} + + +// returns 0..31 for the highest set bit +static int stbi__high_bit(unsigned int z) +{ + int n=0; + if (z == 0) return -1; + if (z >= 0x10000) { n += 16; z >>= 16; } + if (z >= 0x00100) { n += 8; z >>= 8; } + if (z >= 0x00010) { n += 4; z >>= 4; } + if (z >= 0x00004) { n += 2; z >>= 2; } + if (z >= 0x00002) { n += 1;/* >>= 1;*/ } + return n; +} + +static int stbi__bitcount(unsigned int a) +{ + a = (a & 0x55555555) + ((a >> 1) & 0x55555555); // max 2 + a = (a & 0x33333333) + ((a >> 2) & 0x33333333); // max 4 + a = (a + (a >> 4)) & 0x0f0f0f0f; // max 8 per 4, now 8 bits + a = (a + (a >> 8)); // max 16 per 8 bits + a = (a + (a >> 16)); // max 32 per 8 bits + return a & 0xff; +} + +// extract an arbitrarily-aligned N-bit value (N=bits) +// from v, and then make it 8-bits long and fractionally +// extend it to full full range. +static int stbi__shiftsigned(unsigned int v, int shift, int bits) +{ + static unsigned int mul_table[9] = { + 0, + 0xff/*0b11111111*/, 0x55/*0b01010101*/, 0x49/*0b01001001*/, 0x11/*0b00010001*/, + 0x21/*0b00100001*/, 0x41/*0b01000001*/, 0x81/*0b10000001*/, 0x01/*0b00000001*/, + }; + static unsigned int shift_table[9] = { + 0, 0,0,1,0,2,4,6,0, + }; + if (shift < 0) + v <<= -shift; + else + v >>= shift; + STBI_ASSERT(v < 256); + v >>= (8-bits); + STBI_ASSERT(bits >= 0 && bits <= 8); + return (int) ((unsigned) v * mul_table[bits]) >> shift_table[bits]; +} + +typedef struct +{ + int bpp, offset, hsz; + unsigned int mr,mg,mb,ma, all_a; + int extra_read; +} stbi__bmp_data; + +static int stbi__bmp_set_mask_defaults(stbi__bmp_data *info, int compress) +{ + // BI_BITFIELDS specifies masks explicitly, don't override + if (compress == 3) + return 1; + + if (compress == 0) { + if (info->bpp == 16) { + info->mr = 31u << 10; + info->mg = 31u << 5; + info->mb = 31u << 0; + } else if (info->bpp == 32) { + info->mr = 0xffu << 16; + info->mg = 0xffu << 8; + info->mb = 0xffu << 0; + info->ma = 0xffu << 24; + info->all_a = 0; // if all_a is 0 at end, then we loaded alpha channel but it was all 0 + } else { + // otherwise, use defaults, which is all-0 + info->mr = info->mg = info->mb = info->ma = 0; + } + return 1; + } + return 0; // error +} + +static void *stbi__bmp_parse_header(stbi__context *s, stbi__bmp_data *info) +{ + int hsz; + if (stbi__get8(s) != 'B' || stbi__get8(s) != 'M') return stbi__errpuc("not BMP", "Corrupt BMP"); + stbi__get32le(s); // discard filesize + stbi__get16le(s); // discard reserved + stbi__get16le(s); // discard reserved + info->offset = stbi__get32le(s); + info->hsz = hsz = stbi__get32le(s); + info->mr = info->mg = info->mb = info->ma = 0; + info->extra_read = 14; + + if (info->offset < 0) return stbi__errpuc("bad BMP", "bad BMP"); + + if (hsz != 12 && hsz != 40 && hsz != 56 && hsz != 108 && hsz != 124) return stbi__errpuc("unknown BMP", "BMP type not supported: unknown"); + if (hsz == 12) { + s->img_x = stbi__get16le(s); + s->img_y = stbi__get16le(s); + } else { + s->img_x = stbi__get32le(s); + s->img_y = stbi__get32le(s); + } + if (stbi__get16le(s) != 1) return stbi__errpuc("bad BMP", "bad BMP"); + info->bpp = stbi__get16le(s); + if (hsz != 12) { + int compress = stbi__get32le(s); + if (compress == 1 || compress == 2) return stbi__errpuc("BMP RLE", "BMP type not supported: RLE"); + if (compress >= 4) return stbi__errpuc("BMP JPEG/PNG", "BMP type not supported: unsupported compression"); // this includes PNG/JPEG modes + if (compress == 3 && info->bpp != 16 && info->bpp != 32) return stbi__errpuc("bad BMP", "bad BMP"); // bitfields requires 16 or 32 bits/pixel + stbi__get32le(s); // discard sizeof + stbi__get32le(s); // discard hres + stbi__get32le(s); // discard vres + stbi__get32le(s); // discard colorsused + stbi__get32le(s); // discard max important + if (hsz == 40 || hsz == 56) { + if (hsz == 56) { + stbi__get32le(s); + stbi__get32le(s); + stbi__get32le(s); + stbi__get32le(s); + } + if (info->bpp == 16 || info->bpp == 32) { + if (compress == 0) { + stbi__bmp_set_mask_defaults(info, compress); + } else if (compress == 3) { + info->mr = stbi__get32le(s); + info->mg = stbi__get32le(s); + info->mb = stbi__get32le(s); + info->extra_read += 12; + // not documented, but generated by photoshop and handled by mspaint + if (info->mr == info->mg && info->mg == info->mb) { + // ?!?!? + return stbi__errpuc("bad BMP", "bad BMP"); + } + } else + return stbi__errpuc("bad BMP", "bad BMP"); + } + } else { + // V4/V5 header + int i; + if (hsz != 108 && hsz != 124) + return stbi__errpuc("bad BMP", "bad BMP"); + info->mr = stbi__get32le(s); + info->mg = stbi__get32le(s); + info->mb = stbi__get32le(s); + info->ma = stbi__get32le(s); + if (compress != 3) // override mr/mg/mb unless in BI_BITFIELDS mode, as per docs + stbi__bmp_set_mask_defaults(info, compress); + stbi__get32le(s); // discard color space + for (i=0; i < 12; ++i) + stbi__get32le(s); // discard color space parameters + if (hsz == 124) { + stbi__get32le(s); // discard rendering intent + stbi__get32le(s); // discard offset of profile data + stbi__get32le(s); // discard size of profile data + stbi__get32le(s); // discard reserved + } + } + } + return (void *) 1; +} + + +static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *out; + unsigned int mr=0,mg=0,mb=0,ma=0, all_a; + stbi_uc pal[256][4]; + int psize=0,i,j,width; + int flip_vertically, pad, target; + stbi__bmp_data info; + STBI_NOTUSED(ri); + + info.all_a = 255; + if (stbi__bmp_parse_header(s, &info) == NULL) + return NULL; // error code already set + + flip_vertically = ((int) s->img_y) > 0; + s->img_y = abs((int) s->img_y); + + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + mr = info.mr; + mg = info.mg; + mb = info.mb; + ma = info.ma; + all_a = info.all_a; + + if (info.hsz == 12) { + if (info.bpp < 24) + psize = (info.offset - info.extra_read - 24) / 3; + } else { + if (info.bpp < 16) + psize = (info.offset - info.extra_read - info.hsz) >> 2; + } + if (psize == 0) { + // accept some number of extra bytes after the header, but if the offset points either to before + // the header ends or implies a large amount of extra data, reject the file as malformed + int bytes_read_so_far = s->callback_already_read + (int)(s->img_buffer - s->img_buffer_original); + int header_limit = 1024; // max we actually read is below 256 bytes currently. + int extra_data_limit = 256*4; // what ordinarily goes here is a palette; 256 entries*4 bytes is its max size. + if (bytes_read_so_far <= 0 || bytes_read_so_far > header_limit) { + return stbi__errpuc("bad header", "Corrupt BMP"); + } + // we established that bytes_read_so_far is positive and sensible. + // the first half of this test rejects offsets that are either too small positives, or + // negative, and guarantees that info.offset >= bytes_read_so_far > 0. this in turn + // ensures the number computed in the second half of the test can't overflow. + if (info.offset < bytes_read_so_far || info.offset - bytes_read_so_far > extra_data_limit) { + return stbi__errpuc("bad offset", "Corrupt BMP"); + } else { + stbi__skip(s, info.offset - bytes_read_so_far); + } + } + + if (info.bpp == 24 && ma == 0xff000000) + s->img_n = 3; + else + s->img_n = ma ? 4 : 3; + if (req_comp && req_comp >= 3) // we can directly decode 3 or 4 + target = req_comp; + else + target = s->img_n; // if they want monochrome, we'll post-convert + + // sanity-check size + if (!stbi__mad3sizes_valid(target, s->img_x, s->img_y, 0)) + return stbi__errpuc("too large", "Corrupt BMP"); + + out = (stbi_uc *) stbi__malloc_mad3(target, s->img_x, s->img_y, 0); + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + if (info.bpp < 16) { + int z=0; + if (psize == 0 || psize > 256) { STBI_FREE(out); return stbi__errpuc("invalid", "Corrupt BMP"); } + for (i=0; i < psize; ++i) { + pal[i][2] = stbi__get8(s); + pal[i][1] = stbi__get8(s); + pal[i][0] = stbi__get8(s); + if (info.hsz != 12) stbi__get8(s); + pal[i][3] = 255; + } + stbi__skip(s, info.offset - info.extra_read - info.hsz - psize * (info.hsz == 12 ? 3 : 4)); + if (info.bpp == 1) width = (s->img_x + 7) >> 3; + else if (info.bpp == 4) width = (s->img_x + 1) >> 1; + else if (info.bpp == 8) width = s->img_x; + else { STBI_FREE(out); return stbi__errpuc("bad bpp", "Corrupt BMP"); } + pad = (-width)&3; + if (info.bpp == 1) { + for (j=0; j < (int) s->img_y; ++j) { + int bit_offset = 7, v = stbi__get8(s); + for (i=0; i < (int) s->img_x; ++i) { + int color = (v>>bit_offset)&0x1; + out[z++] = pal[color][0]; + out[z++] = pal[color][1]; + out[z++] = pal[color][2]; + if (target == 4) out[z++] = 255; + if (i+1 == (int) s->img_x) break; + if((--bit_offset) < 0) { + bit_offset = 7; + v = stbi__get8(s); + } + } + stbi__skip(s, pad); + } + } else { + for (j=0; j < (int) s->img_y; ++j) { + for (i=0; i < (int) s->img_x; i += 2) { + int v=stbi__get8(s),v2=0; + if (info.bpp == 4) { + v2 = v & 15; + v >>= 4; + } + out[z++] = pal[v][0]; + out[z++] = pal[v][1]; + out[z++] = pal[v][2]; + if (target == 4) out[z++] = 255; + if (i+1 == (int) s->img_x) break; + v = (info.bpp == 8) ? stbi__get8(s) : v2; + out[z++] = pal[v][0]; + out[z++] = pal[v][1]; + out[z++] = pal[v][2]; + if (target == 4) out[z++] = 255; + } + stbi__skip(s, pad); + } + } + } else { + int rshift=0,gshift=0,bshift=0,ashift=0,rcount=0,gcount=0,bcount=0,acount=0; + int z = 0; + int easy=0; + stbi__skip(s, info.offset - info.extra_read - info.hsz); + if (info.bpp == 24) width = 3 * s->img_x; + else if (info.bpp == 16) width = 2*s->img_x; + else /* bpp = 32 and pad = 0 */ width=0; + pad = (-width) & 3; + if (info.bpp == 24) { + easy = 1; + } else if (info.bpp == 32) { + if (mb == 0xff && mg == 0xff00 && mr == 0x00ff0000 && ma == 0xff000000) + easy = 2; + } + if (!easy) { + if (!mr || !mg || !mb) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } + // right shift amt to put high bit in position #7 + rshift = stbi__high_bit(mr)-7; rcount = stbi__bitcount(mr); + gshift = stbi__high_bit(mg)-7; gcount = stbi__bitcount(mg); + bshift = stbi__high_bit(mb)-7; bcount = stbi__bitcount(mb); + ashift = stbi__high_bit(ma)-7; acount = stbi__bitcount(ma); + if (rcount > 8 || gcount > 8 || bcount > 8 || acount > 8) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } + } + for (j=0; j < (int) s->img_y; ++j) { + if (easy) { + for (i=0; i < (int) s->img_x; ++i) { + unsigned char a; + out[z+2] = stbi__get8(s); + out[z+1] = stbi__get8(s); + out[z+0] = stbi__get8(s); + z += 3; + a = (easy == 2 ? stbi__get8(s) : 255); + all_a |= a; + if (target == 4) out[z++] = a; + } + } else { + int bpp = info.bpp; + for (i=0; i < (int) s->img_x; ++i) { + stbi__uint32 v = (bpp == 16 ? (stbi__uint32) stbi__get16le(s) : stbi__get32le(s)); + unsigned int a; + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mr, rshift, rcount)); + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mg, gshift, gcount)); + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mb, bshift, bcount)); + a = (ma ? stbi__shiftsigned(v & ma, ashift, acount) : 255); + all_a |= a; + if (target == 4) out[z++] = STBI__BYTECAST(a); + } + } + stbi__skip(s, pad); + } + } + + // if alpha channel is all 0s, replace with all 255s + if (target == 4 && all_a == 0) + for (i=4*s->img_x*s->img_y-1; i >= 0; i -= 4) + out[i] = 255; + + if (flip_vertically) { + stbi_uc t; + for (j=0; j < (int) s->img_y>>1; ++j) { + stbi_uc *p1 = out + j *s->img_x*target; + stbi_uc *p2 = out + (s->img_y-1-j)*s->img_x*target; + for (i=0; i < (int) s->img_x*target; ++i) { + t = p1[i]; p1[i] = p2[i]; p2[i] = t; + } + } + } + + if (req_comp && req_comp != target) { + out = stbi__convert_format(out, target, req_comp, s->img_x, s->img_y); + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + + *x = s->img_x; + *y = s->img_y; + if (comp) *comp = s->img_n; + return out; +} +#endif + +// Targa Truevision - TGA +// by Jonathan Dummer +#ifndef STBI_NO_TGA +// returns STBI_rgb or whatever, 0 on error +static int stbi__tga_get_comp(int bits_per_pixel, int is_grey, int* is_rgb16) +{ + // only RGB or RGBA (incl. 16bit) or grey allowed + if (is_rgb16) *is_rgb16 = 0; + switch(bits_per_pixel) { + case 8: return STBI_grey; + case 16: if(is_grey) return STBI_grey_alpha; + // fallthrough + case 15: if(is_rgb16) *is_rgb16 = 1; + return STBI_rgb; + case 24: // fallthrough + case 32: return bits_per_pixel/8; + default: return 0; + } +} + +static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp) +{ + int tga_w, tga_h, tga_comp, tga_image_type, tga_bits_per_pixel, tga_colormap_bpp; + int sz, tga_colormap_type; + stbi__get8(s); // discard Offset + tga_colormap_type = stbi__get8(s); // colormap type + if( tga_colormap_type > 1 ) { + stbi__rewind(s); + return 0; // only RGB or indexed allowed + } + tga_image_type = stbi__get8(s); // image type + if ( tga_colormap_type == 1 ) { // colormapped (paletted) image + if (tga_image_type != 1 && tga_image_type != 9) { + stbi__rewind(s); + return 0; + } + stbi__skip(s,4); // skip index of first colormap entry and number of entries + sz = stbi__get8(s); // check bits per palette color entry + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) { + stbi__rewind(s); + return 0; + } + stbi__skip(s,4); // skip image x and y origin + tga_colormap_bpp = sz; + } else { // "normal" image w/o colormap - only RGB or grey allowed, +/- RLE + if ( (tga_image_type != 2) && (tga_image_type != 3) && (tga_image_type != 10) && (tga_image_type != 11) ) { + stbi__rewind(s); + return 0; // only RGB or grey allowed, +/- RLE + } + stbi__skip(s,9); // skip colormap specification and image x/y origin + tga_colormap_bpp = 0; + } + tga_w = stbi__get16le(s); + if( tga_w < 1 ) { + stbi__rewind(s); + return 0; // test width + } + tga_h = stbi__get16le(s); + if( tga_h < 1 ) { + stbi__rewind(s); + return 0; // test height + } + tga_bits_per_pixel = stbi__get8(s); // bits per pixel + stbi__get8(s); // ignore alpha bits + if (tga_colormap_bpp != 0) { + if((tga_bits_per_pixel != 8) && (tga_bits_per_pixel != 16)) { + // when using a colormap, tga_bits_per_pixel is the size of the indexes + // I don't think anything but 8 or 16bit indexes makes sense + stbi__rewind(s); + return 0; + } + tga_comp = stbi__tga_get_comp(tga_colormap_bpp, 0, NULL); + } else { + tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3) || (tga_image_type == 11), NULL); + } + if(!tga_comp) { + stbi__rewind(s); + return 0; + } + if (x) *x = tga_w; + if (y) *y = tga_h; + if (comp) *comp = tga_comp; + return 1; // seems to have passed everything +} + +static int stbi__tga_test(stbi__context *s) +{ + int res = 0; + int sz, tga_color_type; + stbi__get8(s); // discard Offset + tga_color_type = stbi__get8(s); // color type + if ( tga_color_type > 1 ) goto errorEnd; // only RGB or indexed allowed + sz = stbi__get8(s); // image type + if ( tga_color_type == 1 ) { // colormapped (paletted) image + if (sz != 1 && sz != 9) goto errorEnd; // colortype 1 demands image type 1 or 9 + stbi__skip(s,4); // skip index of first colormap entry and number of entries + sz = stbi__get8(s); // check bits per palette color entry + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; + stbi__skip(s,4); // skip image x and y origin + } else { // "normal" image w/o colormap + if ( (sz != 2) && (sz != 3) && (sz != 10) && (sz != 11) ) goto errorEnd; // only RGB or grey allowed, +/- RLE + stbi__skip(s,9); // skip colormap specification and image x/y origin + } + if ( stbi__get16le(s) < 1 ) goto errorEnd; // test width + if ( stbi__get16le(s) < 1 ) goto errorEnd; // test height + sz = stbi__get8(s); // bits per pixel + if ( (tga_color_type == 1) && (sz != 8) && (sz != 16) ) goto errorEnd; // for colormapped images, bpp is size of an index + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; + + res = 1; // if we got this far, everything's good and we can return 1 instead of 0 + +errorEnd: + stbi__rewind(s); + return res; +} + +// read 16bit value and convert to 24bit RGB +static void stbi__tga_read_rgb16(stbi__context *s, stbi_uc* out) +{ + stbi__uint16 px = (stbi__uint16)stbi__get16le(s); + stbi__uint16 fiveBitMask = 31; + // we have 3 channels with 5bits each + int r = (px >> 10) & fiveBitMask; + int g = (px >> 5) & fiveBitMask; + int b = px & fiveBitMask; + // Note that this saves the data in RGB(A) order, so it doesn't need to be swapped later + out[0] = (stbi_uc)((r * 255)/31); + out[1] = (stbi_uc)((g * 255)/31); + out[2] = (stbi_uc)((b * 255)/31); + + // some people claim that the most significant bit might be used for alpha + // (possibly if an alpha-bit is set in the "image descriptor byte") + // but that only made 16bit test images completely translucent.. + // so let's treat all 15 and 16bit TGAs as RGB with no alpha. +} + +static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + // read in the TGA header stuff + int tga_offset = stbi__get8(s); + int tga_indexed = stbi__get8(s); + int tga_image_type = stbi__get8(s); + int tga_is_RLE = 0; + int tga_palette_start = stbi__get16le(s); + int tga_palette_len = stbi__get16le(s); + int tga_palette_bits = stbi__get8(s); + int tga_x_origin = stbi__get16le(s); + int tga_y_origin = stbi__get16le(s); + int tga_width = stbi__get16le(s); + int tga_height = stbi__get16le(s); + int tga_bits_per_pixel = stbi__get8(s); + int tga_comp, tga_rgb16=0; + int tga_inverted = stbi__get8(s); + // int tga_alpha_bits = tga_inverted & 15; // the 4 lowest bits - unused (useless?) + // image data + unsigned char *tga_data; + unsigned char *tga_palette = NULL; + int i, j; + unsigned char raw_data[4] = {0}; + int RLE_count = 0; + int RLE_repeating = 0; + int read_next_pixel = 1; + STBI_NOTUSED(ri); + STBI_NOTUSED(tga_x_origin); // @TODO + STBI_NOTUSED(tga_y_origin); // @TODO + + if (tga_height > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (tga_width > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + // do a tiny bit of precessing + if ( tga_image_type >= 8 ) + { + tga_image_type -= 8; + tga_is_RLE = 1; + } + tga_inverted = 1 - ((tga_inverted >> 5) & 1); + + // If I'm paletted, then I'll use the number of bits from the palette + if ( tga_indexed ) tga_comp = stbi__tga_get_comp(tga_palette_bits, 0, &tga_rgb16); + else tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3), &tga_rgb16); + + if(!tga_comp) // shouldn't really happen, stbi__tga_test() should have ensured basic consistency + return stbi__errpuc("bad format", "Can't find out TGA pixelformat"); + + // tga info + *x = tga_width; + *y = tga_height; + if (comp) *comp = tga_comp; + + if (!stbi__mad3sizes_valid(tga_width, tga_height, tga_comp, 0)) + return stbi__errpuc("too large", "Corrupt TGA"); + + tga_data = (unsigned char*)stbi__malloc_mad3(tga_width, tga_height, tga_comp, 0); + if (!tga_data) return stbi__errpuc("outofmem", "Out of memory"); + + // skip to the data's starting position (offset usually = 0) + stbi__skip(s, tga_offset ); + + if ( !tga_indexed && !tga_is_RLE && !tga_rgb16 ) { + for (i=0; i < tga_height; ++i) { + int row = tga_inverted ? tga_height -i - 1 : i; + stbi_uc *tga_row = tga_data + row*tga_width*tga_comp; + stbi__getn(s, tga_row, tga_width * tga_comp); + } + } else { + // do I need to load a palette? + if ( tga_indexed) + { + if (tga_palette_len == 0) { /* you have to have at least one entry! */ + STBI_FREE(tga_data); + return stbi__errpuc("bad palette", "Corrupt TGA"); + } + + // any data to skip? (offset usually = 0) + stbi__skip(s, tga_palette_start ); + // load the palette + tga_palette = (unsigned char*)stbi__malloc_mad2(tga_palette_len, tga_comp, 0); + if (!tga_palette) { + STBI_FREE(tga_data); + return stbi__errpuc("outofmem", "Out of memory"); + } + if (tga_rgb16) { + stbi_uc *pal_entry = tga_palette; + STBI_ASSERT(tga_comp == STBI_rgb); + for (i=0; i < tga_palette_len; ++i) { + stbi__tga_read_rgb16(s, pal_entry); + pal_entry += tga_comp; + } + } else if (!stbi__getn(s, tga_palette, tga_palette_len * tga_comp)) { + STBI_FREE(tga_data); + STBI_FREE(tga_palette); + return stbi__errpuc("bad palette", "Corrupt TGA"); + } + } + // load the data + for (i=0; i < tga_width * tga_height; ++i) + { + // if I'm in RLE mode, do I need to get a RLE stbi__pngchunk? + if ( tga_is_RLE ) + { + if ( RLE_count == 0 ) + { + // yep, get the next byte as a RLE command + int RLE_cmd = stbi__get8(s); + RLE_count = 1 + (RLE_cmd & 127); + RLE_repeating = RLE_cmd >> 7; + read_next_pixel = 1; + } else if ( !RLE_repeating ) + { + read_next_pixel = 1; + } + } else + { + read_next_pixel = 1; + } + // OK, if I need to read a pixel, do it now + if ( read_next_pixel ) + { + // load however much data we did have + if ( tga_indexed ) + { + // read in index, then perform the lookup + int pal_idx = (tga_bits_per_pixel == 8) ? stbi__get8(s) : stbi__get16le(s); + if ( pal_idx >= tga_palette_len ) { + // invalid index + pal_idx = 0; + } + pal_idx *= tga_comp; + for (j = 0; j < tga_comp; ++j) { + raw_data[j] = tga_palette[pal_idx+j]; + } + } else if(tga_rgb16) { + STBI_ASSERT(tga_comp == STBI_rgb); + stbi__tga_read_rgb16(s, raw_data); + } else { + // read in the data raw + for (j = 0; j < tga_comp; ++j) { + raw_data[j] = stbi__get8(s); + } + } + // clear the reading flag for the next pixel + read_next_pixel = 0; + } // end of reading a pixel + + // copy data + for (j = 0; j < tga_comp; ++j) + tga_data[i*tga_comp+j] = raw_data[j]; + + // in case we're in RLE mode, keep counting down + --RLE_count; + } + // do I need to invert the image? + if ( tga_inverted ) + { + for (j = 0; j*2 < tga_height; ++j) + { + int index1 = j * tga_width * tga_comp; + int index2 = (tga_height - 1 - j) * tga_width * tga_comp; + for (i = tga_width * tga_comp; i > 0; --i) + { + unsigned char temp = tga_data[index1]; + tga_data[index1] = tga_data[index2]; + tga_data[index2] = temp; + ++index1; + ++index2; + } + } + } + // clear my palette, if I had one + if ( tga_palette != NULL ) + { + STBI_FREE( tga_palette ); + } + } + + // swap RGB - if the source data was RGB16, it already is in the right order + if (tga_comp >= 3 && !tga_rgb16) + { + unsigned char* tga_pixel = tga_data; + for (i=0; i < tga_width * tga_height; ++i) + { + unsigned char temp = tga_pixel[0]; + tga_pixel[0] = tga_pixel[2]; + tga_pixel[2] = temp; + tga_pixel += tga_comp; + } + } + + // convert to target component count + if (req_comp && req_comp != tga_comp) + tga_data = stbi__convert_format(tga_data, tga_comp, req_comp, tga_width, tga_height); + + // the things I do to get rid of an error message, and yet keep + // Microsoft's C compilers happy... [8^( + tga_palette_start = tga_palette_len = tga_palette_bits = + tga_x_origin = tga_y_origin = 0; + STBI_NOTUSED(tga_palette_start); + // OK, done + return tga_data; +} +#endif + +// ************************************************************************************************* +// Photoshop PSD loader -- PD by Thatcher Ulrich, integration by Nicolas Schulz, tweaked by STB + +#ifndef STBI_NO_PSD +static int stbi__psd_test(stbi__context *s) +{ + int r = (stbi__get32be(s) == 0x38425053); + stbi__rewind(s); + return r; +} + +static int stbi__psd_decode_rle(stbi__context *s, stbi_uc *p, int pixelCount) +{ + int count, nleft, len; + + count = 0; + while ((nleft = pixelCount - count) > 0) { + len = stbi__get8(s); + if (len == 128) { + // No-op. + } else if (len < 128) { + // Copy next len+1 bytes literally. + len++; + if (len > nleft) return 0; // corrupt data + count += len; + while (len) { + *p = stbi__get8(s); + p += 4; + len--; + } + } else if (len > 128) { + stbi_uc val; + // Next -len+1 bytes in the dest are replicated from next source byte. + // (Interpret len as a negative 8-bit int.) + len = 257 - len; + if (len > nleft) return 0; // corrupt data + val = stbi__get8(s); + count += len; + while (len) { + *p = val; + p += 4; + len--; + } + } + } + + return 1; +} + +static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) +{ + int pixelCount; + int channelCount, compression; + int channel, i; + int bitdepth; + int w,h; + stbi_uc *out; + STBI_NOTUSED(ri); + + // Check identifier + if (stbi__get32be(s) != 0x38425053) // "8BPS" + return stbi__errpuc("not PSD", "Corrupt PSD image"); + + // Check file type version. + if (stbi__get16be(s) != 1) + return stbi__errpuc("wrong version", "Unsupported version of PSD image"); + + // Skip 6 reserved bytes. + stbi__skip(s, 6 ); + + // Read the number of channels (R, G, B, A, etc). + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) + return stbi__errpuc("wrong channel count", "Unsupported number of channels in PSD image"); + + // Read the rows and columns of the image. + h = stbi__get32be(s); + w = stbi__get32be(s); + + if (h > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (w > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + // Make sure the depth is 8 bits. + bitdepth = stbi__get16be(s); + if (bitdepth != 8 && bitdepth != 16) + return stbi__errpuc("unsupported bit depth", "PSD bit depth is not 8 or 16 bit"); + + // Make sure the color mode is RGB. + // Valid options are: + // 0: Bitmap + // 1: Grayscale + // 2: Indexed color + // 3: RGB color + // 4: CMYK color + // 7: Multichannel + // 8: Duotone + // 9: Lab color + if (stbi__get16be(s) != 3) + return stbi__errpuc("wrong color format", "PSD is not in RGB color format"); + + // Skip the Mode Data. (It's the palette for indexed color; other info for other modes.) + stbi__skip(s,stbi__get32be(s) ); + + // Skip the image resources. (resolution, pen tool paths, etc) + stbi__skip(s, stbi__get32be(s) ); + + // Skip the reserved data. + stbi__skip(s, stbi__get32be(s) ); + + // Find out if the data is compressed. + // Known values: + // 0: no compression + // 1: RLE compressed + compression = stbi__get16be(s); + if (compression > 1) + return stbi__errpuc("bad compression", "PSD has an unknown compression format"); + + // Check size + if (!stbi__mad3sizes_valid(4, w, h, 0)) + return stbi__errpuc("too large", "Corrupt PSD"); + + // Create the destination image. + + if (!compression && bitdepth == 16 && bpc == 16) { + out = (stbi_uc *) stbi__malloc_mad3(8, w, h, 0); + ri->bits_per_channel = 16; + } else + out = (stbi_uc *) stbi__malloc(4 * w*h); + + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + pixelCount = w*h; + + // Initialize the data to zero. + //memset( out, 0, pixelCount * 4 ); + + // Finally, the image data. + if (compression) { + // RLE as used by .PSD and .TIFF + // Loop until you get the number of unpacked bytes you are expecting: + // Read the next source byte into n. + // If n is between 0 and 127 inclusive, copy the next n+1 bytes literally. + // Else if n is between -127 and -1 inclusive, copy the next byte -n+1 times. + // Else if n is 128, noop. + // Endloop + + // The RLE-compressed data is preceded by a 2-byte data count for each row in the data, + // which we're going to just skip. + stbi__skip(s, h * channelCount * 2 ); + + // Read the RLE data by channel. + for (channel = 0; channel < 4; channel++) { + stbi_uc *p; + + p = out+channel; + if (channel >= channelCount) { + // Fill this channel with default data. + for (i = 0; i < pixelCount; i++, p += 4) + *p = (channel == 3 ? 255 : 0); + } else { + // Read the RLE data. + if (!stbi__psd_decode_rle(s, p, pixelCount)) { + STBI_FREE(out); + return stbi__errpuc("corrupt", "bad RLE data"); + } + } + } + + } else { + // We're at the raw image data. It's each channel in order (Red, Green, Blue, Alpha, ...) + // where each channel consists of an 8-bit (or 16-bit) value for each pixel in the image. + + // Read the data by channel. + for (channel = 0; channel < 4; channel++) { + if (channel >= channelCount) { + // Fill this channel with default data. + if (bitdepth == 16 && bpc == 16) { + stbi__uint16 *q = ((stbi__uint16 *) out) + channel; + stbi__uint16 val = channel == 3 ? 65535 : 0; + for (i = 0; i < pixelCount; i++, q += 4) + *q = val; + } else { + stbi_uc *p = out+channel; + stbi_uc val = channel == 3 ? 255 : 0; + for (i = 0; i < pixelCount; i++, p += 4) + *p = val; + } + } else { + if (ri->bits_per_channel == 16) { // output bpc + stbi__uint16 *q = ((stbi__uint16 *) out) + channel; + for (i = 0; i < pixelCount; i++, q += 4) + *q = (stbi__uint16) stbi__get16be(s); + } else { + stbi_uc *p = out+channel; + if (bitdepth == 16) { // input bpc + for (i = 0; i < pixelCount; i++, p += 4) + *p = (stbi_uc) (stbi__get16be(s) >> 8); + } else { + for (i = 0; i < pixelCount; i++, p += 4) + *p = stbi__get8(s); + } + } + } + } + } + + // remove weird white matte from PSD + if (channelCount >= 4) { + if (ri->bits_per_channel == 16) { + for (i=0; i < w*h; ++i) { + stbi__uint16 *pixel = (stbi__uint16 *) out + 4*i; + if (pixel[3] != 0 && pixel[3] != 65535) { + float a = pixel[3] / 65535.0f; + float ra = 1.0f / a; + float inv_a = 65535.0f * (1 - ra); + pixel[0] = (stbi__uint16) (pixel[0]*ra + inv_a); + pixel[1] = (stbi__uint16) (pixel[1]*ra + inv_a); + pixel[2] = (stbi__uint16) (pixel[2]*ra + inv_a); + } + } + } else { + for (i=0; i < w*h; ++i) { + unsigned char *pixel = out + 4*i; + if (pixel[3] != 0 && pixel[3] != 255) { + float a = pixel[3] / 255.0f; + float ra = 1.0f / a; + float inv_a = 255.0f * (1 - ra); + pixel[0] = (unsigned char) (pixel[0]*ra + inv_a); + pixel[1] = (unsigned char) (pixel[1]*ra + inv_a); + pixel[2] = (unsigned char) (pixel[2]*ra + inv_a); + } + } + } + } + + // convert to desired output format + if (req_comp && req_comp != 4) { + if (ri->bits_per_channel == 16) + out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, 4, req_comp, w, h); + else + out = stbi__convert_format(out, 4, req_comp, w, h); + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + + if (comp) *comp = 4; + *y = h; + *x = w; + + return out; +} +#endif + +// ************************************************************************************************* +// Softimage PIC loader +// by Tom Seddon +// +// See http://softimage.wiki.softimage.com/index.php/INFO:_PIC_file_format +// See http://ozviz.wasp.uwa.edu.au/~pbourke/dataformats/softimagepic/ + +#ifndef STBI_NO_PIC +static int stbi__pic_is4(stbi__context *s,const char *str) +{ + int i; + for (i=0; i<4; ++i) + if (stbi__get8(s) != (stbi_uc)str[i]) + return 0; + + return 1; +} + +static int stbi__pic_test_core(stbi__context *s) +{ + int i; + + if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) + return 0; + + for(i=0;i<84;++i) + stbi__get8(s); + + if (!stbi__pic_is4(s,"PICT")) + return 0; + + return 1; +} + +typedef struct +{ + stbi_uc size,type,channel; +} stbi__pic_packet; + +static stbi_uc *stbi__readval(stbi__context *s, int channel, stbi_uc *dest) +{ + int mask=0x80, i; + + for (i=0; i<4; ++i, mask>>=1) { + if (channel & mask) { + if (stbi__at_eof(s)) return stbi__errpuc("bad file","PIC file too short"); + dest[i]=stbi__get8(s); + } + } + + return dest; +} + +static void stbi__copyval(int channel,stbi_uc *dest,const stbi_uc *src) +{ + int mask=0x80,i; + + for (i=0;i<4; ++i, mask>>=1) + if (channel&mask) + dest[i]=src[i]; +} + +static stbi_uc *stbi__pic_load_core(stbi__context *s,int width,int height,int *comp, stbi_uc *result) +{ + int act_comp=0,num_packets=0,y,chained; + stbi__pic_packet packets[10]; + + // this will (should...) cater for even some bizarre stuff like having data + // for the same channel in multiple packets. + do { + stbi__pic_packet *packet; + + if (num_packets==sizeof(packets)/sizeof(packets[0])) + return stbi__errpuc("bad format","too many packets"); + + packet = &packets[num_packets++]; + + chained = stbi__get8(s); + packet->size = stbi__get8(s); + packet->type = stbi__get8(s); + packet->channel = stbi__get8(s); + + act_comp |= packet->channel; + + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (reading packets)"); + if (packet->size != 8) return stbi__errpuc("bad format","packet isn't 8bpp"); + } while (chained); + + *comp = (act_comp & 0x10 ? 4 : 3); // has alpha channel? + + for(y=0; ytype) { + default: + return stbi__errpuc("bad format","packet has bad compression type"); + + case 0: {//uncompressed + int x; + + for(x=0;xchannel,dest)) + return 0; + break; + } + + case 1://Pure RLE + { + int left=width, i; + + while (left>0) { + stbi_uc count,value[4]; + + count=stbi__get8(s); + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pure read count)"); + + if (count > left) + count = (stbi_uc) left; + + if (!stbi__readval(s,packet->channel,value)) return 0; + + for(i=0; ichannel,dest,value); + left -= count; + } + } + break; + + case 2: {//Mixed RLE + int left=width; + while (left>0) { + int count = stbi__get8(s), i; + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (mixed read count)"); + + if (count >= 128) { // Repeated + stbi_uc value[4]; + + if (count==128) + count = stbi__get16be(s); + else + count -= 127; + if (count > left) + return stbi__errpuc("bad file","scanline overrun"); + + if (!stbi__readval(s,packet->channel,value)) + return 0; + + for(i=0;ichannel,dest,value); + } else { // Raw + ++count; + if (count>left) return stbi__errpuc("bad file","scanline overrun"); + + for(i=0;ichannel,dest)) + return 0; + } + left-=count; + } + break; + } + } + } + } + + return result; +} + +static void *stbi__pic_load(stbi__context *s,int *px,int *py,int *comp,int req_comp, stbi__result_info *ri) +{ + stbi_uc *result; + int i, x,y, internal_comp; + STBI_NOTUSED(ri); + + if (!comp) comp = &internal_comp; + + for (i=0; i<92; ++i) + stbi__get8(s); + + x = stbi__get16be(s); + y = stbi__get16be(s); + + if (y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pic header)"); + if (!stbi__mad3sizes_valid(x, y, 4, 0)) return stbi__errpuc("too large", "PIC image too large to decode"); + + stbi__get32be(s); //skip `ratio' + stbi__get16be(s); //skip `fields' + stbi__get16be(s); //skip `pad' + + // intermediate buffer is RGBA + result = (stbi_uc *) stbi__malloc_mad3(x, y, 4, 0); + if (!result) return stbi__errpuc("outofmem", "Out of memory"); + memset(result, 0xff, x*y*4); + + if (!stbi__pic_load_core(s,x,y,comp, result)) { + STBI_FREE(result); + result=0; + } + *px = x; + *py = y; + if (req_comp == 0) req_comp = *comp; + result=stbi__convert_format(result,4,req_comp,x,y); + + return result; +} + +static int stbi__pic_test(stbi__context *s) +{ + int r = stbi__pic_test_core(s); + stbi__rewind(s); + return r; +} +#endif + +// ************************************************************************************************* +// GIF loader -- public domain by Jean-Marc Lienher -- simplified/shrunk by stb + +#ifndef STBI_NO_GIF +typedef struct +{ + stbi__int16 prefix; + stbi_uc first; + stbi_uc suffix; +} stbi__gif_lzw; + +typedef struct +{ + int w,h; + stbi_uc *out; // output buffer (always 4 components) + stbi_uc *background; // The current "background" as far as a gif is concerned + stbi_uc *history; + int flags, bgindex, ratio, transparent, eflags; + stbi_uc pal[256][4]; + stbi_uc lpal[256][4]; + stbi__gif_lzw codes[8192]; + stbi_uc *color_table; + int parse, step; + int lflags; + int start_x, start_y; + int max_x, max_y; + int cur_x, cur_y; + int line_size; + int delay; +} stbi__gif; + +static int stbi__gif_test_raw(stbi__context *s) +{ + int sz; + if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') return 0; + sz = stbi__get8(s); + if (sz != '9' && sz != '7') return 0; + if (stbi__get8(s) != 'a') return 0; + return 1; +} + +static int stbi__gif_test(stbi__context *s) +{ + int r = stbi__gif_test_raw(s); + stbi__rewind(s); + return r; +} + +static void stbi__gif_parse_colortable(stbi__context *s, stbi_uc pal[256][4], int num_entries, int transp) +{ + int i; + for (i=0; i < num_entries; ++i) { + pal[i][2] = stbi__get8(s); + pal[i][1] = stbi__get8(s); + pal[i][0] = stbi__get8(s); + pal[i][3] = transp == i ? 0 : 255; + } +} + +static int stbi__gif_header(stbi__context *s, stbi__gif *g, int *comp, int is_info) +{ + stbi_uc version; + if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') + return stbi__err("not GIF", "Corrupt GIF"); + + version = stbi__get8(s); + if (version != '7' && version != '9') return stbi__err("not GIF", "Corrupt GIF"); + if (stbi__get8(s) != 'a') return stbi__err("not GIF", "Corrupt GIF"); + + stbi__g_failure_reason = ""; + g->w = stbi__get16le(s); + g->h = stbi__get16le(s); + g->flags = stbi__get8(s); + g->bgindex = stbi__get8(s); + g->ratio = stbi__get8(s); + g->transparent = -1; + + if (g->w > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (g->h > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + + if (comp != 0) *comp = 4; // can't actually tell whether it's 3 or 4 until we parse the comments + + if (is_info) return 1; + + if (g->flags & 0x80) + stbi__gif_parse_colortable(s,g->pal, 2 << (g->flags & 7), -1); + + return 1; +} + +static int stbi__gif_info_raw(stbi__context *s, int *x, int *y, int *comp) +{ + stbi__gif* g = (stbi__gif*) stbi__malloc(sizeof(stbi__gif)); + if (!g) return stbi__err("outofmem", "Out of memory"); + if (!stbi__gif_header(s, g, comp, 1)) { + STBI_FREE(g); + stbi__rewind( s ); + return 0; + } + if (x) *x = g->w; + if (y) *y = g->h; + STBI_FREE(g); + return 1; +} + +static void stbi__out_gif_code(stbi__gif *g, stbi__uint16 code) +{ + stbi_uc *p, *c; + int idx; + + // recurse to decode the prefixes, since the linked-list is backwards, + // and working backwards through an interleaved image would be nasty + if (g->codes[code].prefix >= 0) + stbi__out_gif_code(g, g->codes[code].prefix); + + if (g->cur_y >= g->max_y) return; + + idx = g->cur_x + g->cur_y; + p = &g->out[idx]; + g->history[idx / 4] = 1; + + c = &g->color_table[g->codes[code].suffix * 4]; + if (c[3] > 128) { // don't render transparent pixels; + p[0] = c[2]; + p[1] = c[1]; + p[2] = c[0]; + p[3] = c[3]; + } + g->cur_x += 4; + + if (g->cur_x >= g->max_x) { + g->cur_x = g->start_x; + g->cur_y += g->step; + + while (g->cur_y >= g->max_y && g->parse > 0) { + g->step = (1 << g->parse) * g->line_size; + g->cur_y = g->start_y + (g->step >> 1); + --g->parse; + } + } +} + +static stbi_uc *stbi__process_gif_raster(stbi__context *s, stbi__gif *g) +{ + stbi_uc lzw_cs; + stbi__int32 len, init_code; + stbi__uint32 first; + stbi__int32 codesize, codemask, avail, oldcode, bits, valid_bits, clear; + stbi__gif_lzw *p; + + lzw_cs = stbi__get8(s); + if (lzw_cs > 12) return NULL; + clear = 1 << lzw_cs; + first = 1; + codesize = lzw_cs + 1; + codemask = (1 << codesize) - 1; + bits = 0; + valid_bits = 0; + for (init_code = 0; init_code < clear; init_code++) { + g->codes[init_code].prefix = -1; + g->codes[init_code].first = (stbi_uc) init_code; + g->codes[init_code].suffix = (stbi_uc) init_code; + } + + // support no starting clear code + avail = clear+2; + oldcode = -1; + + len = 0; + for(;;) { + if (valid_bits < codesize) { + if (len == 0) { + len = stbi__get8(s); // start new block + if (len == 0) + return g->out; + } + --len; + bits |= (stbi__int32) stbi__get8(s) << valid_bits; + valid_bits += 8; + } else { + stbi__int32 code = bits & codemask; + bits >>= codesize; + valid_bits -= codesize; + // @OPTIMIZE: is there some way we can accelerate the non-clear path? + if (code == clear) { // clear code + codesize = lzw_cs + 1; + codemask = (1 << codesize) - 1; + avail = clear + 2; + oldcode = -1; + first = 0; + } else if (code == clear + 1) { // end of stream code + stbi__skip(s, len); + while ((len = stbi__get8(s)) > 0) + stbi__skip(s,len); + return g->out; + } else if (code <= avail) { + if (first) { + return stbi__errpuc("no clear code", "Corrupt GIF"); + } + + if (oldcode >= 0) { + p = &g->codes[avail++]; + if (avail > 8192) { + return stbi__errpuc("too many codes", "Corrupt GIF"); + } + + p->prefix = (stbi__int16) oldcode; + p->first = g->codes[oldcode].first; + p->suffix = (code == avail) ? p->first : g->codes[code].first; + } else if (code == avail) + return stbi__errpuc("illegal code in raster", "Corrupt GIF"); + + stbi__out_gif_code(g, (stbi__uint16) code); + + if ((avail & codemask) == 0 && avail <= 0x0FFF) { + codesize++; + codemask = (1 << codesize) - 1; + } + + oldcode = code; + } else { + return stbi__errpuc("illegal code in raster", "Corrupt GIF"); + } + } + } +} + +// this function is designed to support animated gifs, although stb_image doesn't support it +// two back is the image from two frames ago, used for a very specific disposal format +static stbi_uc *stbi__gif_load_next(stbi__context *s, stbi__gif *g, int *comp, int req_comp, stbi_uc *two_back) +{ + int dispose; + int first_frame; + int pi; + int pcount; + STBI_NOTUSED(req_comp); + + // on first frame, any non-written pixels get the background colour (non-transparent) + first_frame = 0; + if (g->out == 0) { + if (!stbi__gif_header(s, g, comp,0)) return 0; // stbi__g_failure_reason set by stbi__gif_header + if (!stbi__mad3sizes_valid(4, g->w, g->h, 0)) + return stbi__errpuc("too large", "GIF image is too large"); + pcount = g->w * g->h; + g->out = (stbi_uc *) stbi__malloc(4 * pcount); + g->background = (stbi_uc *) stbi__malloc(4 * pcount); + g->history = (stbi_uc *) stbi__malloc(pcount); + if (!g->out || !g->background || !g->history) + return stbi__errpuc("outofmem", "Out of memory"); + + // image is treated as "transparent" at the start - ie, nothing overwrites the current background; + // background colour is only used for pixels that are not rendered first frame, after that "background" + // color refers to the color that was there the previous frame. + memset(g->out, 0x00, 4 * pcount); + memset(g->background, 0x00, 4 * pcount); // state of the background (starts transparent) + memset(g->history, 0x00, pcount); // pixels that were affected previous frame + first_frame = 1; + } else { + // second frame - how do we dispose of the previous one? + dispose = (g->eflags & 0x1C) >> 2; + pcount = g->w * g->h; + + if ((dispose == 3) && (two_back == 0)) { + dispose = 2; // if I don't have an image to revert back to, default to the old background + } + + if (dispose == 3) { // use previous graphic + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi]) { + memcpy( &g->out[pi * 4], &two_back[pi * 4], 4 ); + } + } + } else if (dispose == 2) { + // restore what was changed last frame to background before that frame; + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi]) { + memcpy( &g->out[pi * 4], &g->background[pi * 4], 4 ); + } + } + } else { + // This is a non-disposal case eithe way, so just + // leave the pixels as is, and they will become the new background + // 1: do not dispose + // 0: not specified. + } + + // background is what out is after the undoing of the previou frame; + memcpy( g->background, g->out, 4 * g->w * g->h ); + } + + // clear my history; + memset( g->history, 0x00, g->w * g->h ); // pixels that were affected previous frame + + for (;;) { + int tag = stbi__get8(s); + switch (tag) { + case 0x2C: /* Image Descriptor */ + { + stbi__int32 x, y, w, h; + stbi_uc *o; + + x = stbi__get16le(s); + y = stbi__get16le(s); + w = stbi__get16le(s); + h = stbi__get16le(s); + if (((x + w) > (g->w)) || ((y + h) > (g->h))) + return stbi__errpuc("bad Image Descriptor", "Corrupt GIF"); + + g->line_size = g->w * 4; + g->start_x = x * 4; + g->start_y = y * g->line_size; + g->max_x = g->start_x + w * 4; + g->max_y = g->start_y + h * g->line_size; + g->cur_x = g->start_x; + g->cur_y = g->start_y; + + // if the width of the specified rectangle is 0, that means + // we may not see *any* pixels or the image is malformed; + // to make sure this is caught, move the current y down to + // max_y (which is what out_gif_code checks). + if (w == 0) + g->cur_y = g->max_y; + + g->lflags = stbi__get8(s); + + if (g->lflags & 0x40) { + g->step = 8 * g->line_size; // first interlaced spacing + g->parse = 3; + } else { + g->step = g->line_size; + g->parse = 0; + } + + if (g->lflags & 0x80) { + stbi__gif_parse_colortable(s,g->lpal, 2 << (g->lflags & 7), g->eflags & 0x01 ? g->transparent : -1); + g->color_table = (stbi_uc *) g->lpal; + } else if (g->flags & 0x80) { + g->color_table = (stbi_uc *) g->pal; + } else + return stbi__errpuc("missing color table", "Corrupt GIF"); + + o = stbi__process_gif_raster(s, g); + if (!o) return NULL; + + // if this was the first frame, + pcount = g->w * g->h; + if (first_frame && (g->bgindex > 0)) { + // if first frame, any pixel not drawn to gets the background color + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi] == 0) { + g->pal[g->bgindex][3] = 255; // just in case it was made transparent, undo that; It will be reset next frame if need be; + memcpy( &g->out[pi * 4], &g->pal[g->bgindex], 4 ); + } + } + } + + return o; + } + + case 0x21: // Comment Extension. + { + int len; + int ext = stbi__get8(s); + if (ext == 0xF9) { // Graphic Control Extension. + len = stbi__get8(s); + if (len == 4) { + g->eflags = stbi__get8(s); + g->delay = 10 * stbi__get16le(s); // delay - 1/100th of a second, saving as 1/1000ths. + + // unset old transparent + if (g->transparent >= 0) { + g->pal[g->transparent][3] = 255; + } + if (g->eflags & 0x01) { + g->transparent = stbi__get8(s); + if (g->transparent >= 0) { + g->pal[g->transparent][3] = 0; + } + } else { + // don't need transparent + stbi__skip(s, 1); + g->transparent = -1; + } + } else { + stbi__skip(s, len); + break; + } + } + while ((len = stbi__get8(s)) != 0) { + stbi__skip(s, len); + } + break; + } + + case 0x3B: // gif stream termination code + return (stbi_uc *) s; // using '1' causes warning on some compilers + + default: + return stbi__errpuc("unknown code", "Corrupt GIF"); + } + } +} + +static void *stbi__load_gif_main_outofmem(stbi__gif *g, stbi_uc *out, int **delays) +{ + STBI_FREE(g->out); + STBI_FREE(g->history); + STBI_FREE(g->background); + + if (out) STBI_FREE(out); + if (delays && *delays) STBI_FREE(*delays); + return stbi__errpuc("outofmem", "Out of memory"); +} + +static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp) +{ + if (stbi__gif_test(s)) { + int layers = 0; + stbi_uc *u = 0; + stbi_uc *out = 0; + stbi_uc *two_back = 0; + stbi__gif g; + int stride; + int out_size = 0; + int delays_size = 0; + + STBI_NOTUSED(out_size); + STBI_NOTUSED(delays_size); + + memset(&g, 0, sizeof(g)); + if (delays) { + *delays = 0; + } + + do { + u = stbi__gif_load_next(s, &g, comp, req_comp, two_back); + if (u == (stbi_uc *) s) u = 0; // end of animated gif marker + + if (u) { + *x = g.w; + *y = g.h; + ++layers; + stride = g.w * g.h * 4; + + if (out) { + void *tmp = (stbi_uc*) STBI_REALLOC_SIZED( out, out_size, layers * stride ); + if (!tmp) + return stbi__load_gif_main_outofmem(&g, out, delays); + else { + out = (stbi_uc*) tmp; + out_size = layers * stride; + } + + if (delays) { + int *new_delays = (int*) STBI_REALLOC_SIZED( *delays, delays_size, sizeof(int) * layers ); + if (!new_delays) + return stbi__load_gif_main_outofmem(&g, out, delays); + *delays = new_delays; + delays_size = layers * sizeof(int); + } + } else { + out = (stbi_uc*)stbi__malloc( layers * stride ); + if (!out) + return stbi__load_gif_main_outofmem(&g, out, delays); + out_size = layers * stride; + if (delays) { + *delays = (int*) stbi__malloc( layers * sizeof(int) ); + if (!*delays) + return stbi__load_gif_main_outofmem(&g, out, delays); + delays_size = layers * sizeof(int); + } + } + memcpy( out + ((layers - 1) * stride), u, stride ); + if (layers >= 2) { + two_back = out - 2 * stride; + } + + if (delays) { + (*delays)[layers - 1U] = g.delay; + } + } + } while (u != 0); + + // free temp buffer; + STBI_FREE(g.out); + STBI_FREE(g.history); + STBI_FREE(g.background); + + // do the final conversion after loading everything; + if (req_comp && req_comp != 4) + out = stbi__convert_format(out, 4, req_comp, layers * g.w, g.h); + + *z = layers; + return out; + } else { + return stbi__errpuc("not GIF", "Image was not as a gif type."); + } +} + +static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *u = 0; + stbi__gif g; + memset(&g, 0, sizeof(g)); + STBI_NOTUSED(ri); + + u = stbi__gif_load_next(s, &g, comp, req_comp, 0); + if (u == (stbi_uc *) s) u = 0; // end of animated gif marker + if (u) { + *x = g.w; + *y = g.h; + + // moved conversion to after successful load so that the same + // can be done for multiple frames. + if (req_comp && req_comp != 4) + u = stbi__convert_format(u, 4, req_comp, g.w, g.h); + } else if (g.out) { + // if there was an error and we allocated an image buffer, free it! + STBI_FREE(g.out); + } + + // free buffers needed for multiple frame loading; + STBI_FREE(g.history); + STBI_FREE(g.background); + + return u; +} + +static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp) +{ + return stbi__gif_info_raw(s,x,y,comp); +} +#endif + +// ************************************************************************************************* +// Radiance RGBE HDR loader +// originally by Nicolas Schulz +#ifndef STBI_NO_HDR +static int stbi__hdr_test_core(stbi__context *s, const char *signature) +{ + int i; + for (i=0; signature[i]; ++i) + if (stbi__get8(s) != signature[i]) + return 0; + stbi__rewind(s); + return 1; +} + +static int stbi__hdr_test(stbi__context* s) +{ + int r = stbi__hdr_test_core(s, "#?RADIANCE\n"); + stbi__rewind(s); + if(!r) { + r = stbi__hdr_test_core(s, "#?RGBE\n"); + stbi__rewind(s); + } + return r; +} + +#define STBI__HDR_BUFLEN 1024 +static char *stbi__hdr_gettoken(stbi__context *z, char *buffer) +{ + int len=0; + char c = '\0'; + + c = (char) stbi__get8(z); + + while (!stbi__at_eof(z) && c != '\n') { + buffer[len++] = c; + if (len == STBI__HDR_BUFLEN-1) { + // flush to end of line + while (!stbi__at_eof(z) && stbi__get8(z) != '\n') + ; + break; + } + c = (char) stbi__get8(z); + } + + buffer[len] = 0; + return buffer; +} + +static void stbi__hdr_convert(float *output, stbi_uc *input, int req_comp) +{ + if ( input[3] != 0 ) { + float f1; + // Exponent + f1 = (float) ldexp(1.0f, input[3] - (int)(128 + 8)); + if (req_comp <= 2) + output[0] = (input[0] + input[1] + input[2]) * f1 / 3; + else { + output[0] = input[0] * f1; + output[1] = input[1] * f1; + output[2] = input[2] * f1; + } + if (req_comp == 2) output[1] = 1; + if (req_comp == 4) output[3] = 1; + } else { + switch (req_comp) { + case 4: output[3] = 1; /* fallthrough */ + case 3: output[0] = output[1] = output[2] = 0; + break; + case 2: output[1] = 1; /* fallthrough */ + case 1: output[0] = 0; + break; + } + } +} + +static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + char buffer[STBI__HDR_BUFLEN]; + char *token; + int valid = 0; + int width, height; + stbi_uc *scanline; + float *hdr_data; + int len; + unsigned char count, value; + int i, j, k, c1,c2, z; + const char *headerToken; + STBI_NOTUSED(ri); + + // Check identifier + headerToken = stbi__hdr_gettoken(s,buffer); + if (strcmp(headerToken, "#?RADIANCE") != 0 && strcmp(headerToken, "#?RGBE") != 0) + return stbi__errpf("not HDR", "Corrupt HDR image"); + + // Parse header + for(;;) { + token = stbi__hdr_gettoken(s,buffer); + if (token[0] == 0) break; + if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; + } + + if (!valid) return stbi__errpf("unsupported format", "Unsupported HDR format"); + + // Parse width and height + // can't use sscanf() if we're not using stdio! + token = stbi__hdr_gettoken(s,buffer); + if (strncmp(token, "-Y ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); + token += 3; + height = (int) strtol(token, &token, 10); + while (*token == ' ') ++token; + if (strncmp(token, "+X ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); + token += 3; + width = (int) strtol(token, NULL, 10); + + if (height > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); + if (width > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); + + *x = width; + *y = height; + + if (comp) *comp = 3; + if (req_comp == 0) req_comp = 3; + + if (!stbi__mad4sizes_valid(width, height, req_comp, sizeof(float), 0)) + return stbi__errpf("too large", "HDR image is too large"); + + // Read data + hdr_data = (float *) stbi__malloc_mad4(width, height, req_comp, sizeof(float), 0); + if (!hdr_data) + return stbi__errpf("outofmem", "Out of memory"); + + // Load image data + // image data is stored as some number of sca + if ( width < 8 || width >= 32768) { + // Read flat data + for (j=0; j < height; ++j) { + for (i=0; i < width; ++i) { + stbi_uc rgbe[4]; + main_decode_loop: + stbi__getn(s, rgbe, 4); + stbi__hdr_convert(hdr_data + j * width * req_comp + i * req_comp, rgbe, req_comp); + } + } + } else { + // Read RLE-encoded data + scanline = NULL; + + for (j = 0; j < height; ++j) { + c1 = stbi__get8(s); + c2 = stbi__get8(s); + len = stbi__get8(s); + if (c1 != 2 || c2 != 2 || (len & 0x80)) { + // not run-length encoded, so we have to actually use THIS data as a decoded + // pixel (note this can't be a valid pixel--one of RGB must be >= 128) + stbi_uc rgbe[4]; + rgbe[0] = (stbi_uc) c1; + rgbe[1] = (stbi_uc) c2; + rgbe[2] = (stbi_uc) len; + rgbe[3] = (stbi_uc) stbi__get8(s); + stbi__hdr_convert(hdr_data, rgbe, req_comp); + i = 1; + j = 0; + STBI_FREE(scanline); + goto main_decode_loop; // yes, this makes no sense + } + len <<= 8; + len |= stbi__get8(s); + if (len != width) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("invalid decoded scanline length", "corrupt HDR"); } + if (scanline == NULL) { + scanline = (stbi_uc *) stbi__malloc_mad2(width, 4, 0); + if (!scanline) { + STBI_FREE(hdr_data); + return stbi__errpf("outofmem", "Out of memory"); + } + } + + for (k = 0; k < 4; ++k) { + int nleft; + i = 0; + while ((nleft = width - i) > 0) { + count = stbi__get8(s); + if (count > 128) { + // Run + value = stbi__get8(s); + count -= 128; + if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } + for (z = 0; z < count; ++z) + scanline[i++ * 4 + k] = value; + } else { + // Dump + if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } + for (z = 0; z < count; ++z) + scanline[i++ * 4 + k] = stbi__get8(s); + } + } + } + for (i=0; i < width; ++i) + stbi__hdr_convert(hdr_data+(j*width + i)*req_comp, scanline + i*4, req_comp); + } + if (scanline) + STBI_FREE(scanline); + } + + return hdr_data; +} + +static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp) +{ + char buffer[STBI__HDR_BUFLEN]; + char *token; + int valid = 0; + int dummy; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + if (stbi__hdr_test(s) == 0) { + stbi__rewind( s ); + return 0; + } + + for(;;) { + token = stbi__hdr_gettoken(s,buffer); + if (token[0] == 0) break; + if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; + } + + if (!valid) { + stbi__rewind( s ); + return 0; + } + token = stbi__hdr_gettoken(s,buffer); + if (strncmp(token, "-Y ", 3)) { + stbi__rewind( s ); + return 0; + } + token += 3; + *y = (int) strtol(token, &token, 10); + while (*token == ' ') ++token; + if (strncmp(token, "+X ", 3)) { + stbi__rewind( s ); + return 0; + } + token += 3; + *x = (int) strtol(token, NULL, 10); + *comp = 3; + return 1; +} +#endif // STBI_NO_HDR + +#ifndef STBI_NO_BMP +static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp) +{ + void *p; + stbi__bmp_data info; + + info.all_a = 255; + p = stbi__bmp_parse_header(s, &info); + if (p == NULL) { + stbi__rewind( s ); + return 0; + } + if (x) *x = s->img_x; + if (y) *y = s->img_y; + if (comp) { + if (info.bpp == 24 && info.ma == 0xff000000) + *comp = 3; + else + *comp = info.ma ? 4 : 3; + } + return 1; +} +#endif + +#ifndef STBI_NO_PSD +static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp) +{ + int channelCount, dummy, depth; + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + if (stbi__get32be(s) != 0x38425053) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 1) { + stbi__rewind( s ); + return 0; + } + stbi__skip(s, 6); + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) { + stbi__rewind( s ); + return 0; + } + *y = stbi__get32be(s); + *x = stbi__get32be(s); + depth = stbi__get16be(s); + if (depth != 8 && depth != 16) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 3) { + stbi__rewind( s ); + return 0; + } + *comp = 4; + return 1; +} + +static int stbi__psd_is16(stbi__context *s) +{ + int channelCount, depth; + if (stbi__get32be(s) != 0x38425053) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 1) { + stbi__rewind( s ); + return 0; + } + stbi__skip(s, 6); + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) { + stbi__rewind( s ); + return 0; + } + STBI_NOTUSED(stbi__get32be(s)); + STBI_NOTUSED(stbi__get32be(s)); + depth = stbi__get16be(s); + if (depth != 16) { + stbi__rewind( s ); + return 0; + } + return 1; +} +#endif + +#ifndef STBI_NO_PIC +static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp) +{ + int act_comp=0,num_packets=0,chained,dummy; + stbi__pic_packet packets[10]; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) { + stbi__rewind(s); + return 0; + } + + stbi__skip(s, 88); + + *x = stbi__get16be(s); + *y = stbi__get16be(s); + if (stbi__at_eof(s)) { + stbi__rewind( s); + return 0; + } + if ( (*x) != 0 && (1 << 28) / (*x) < (*y)) { + stbi__rewind( s ); + return 0; + } + + stbi__skip(s, 8); + + do { + stbi__pic_packet *packet; + + if (num_packets==sizeof(packets)/sizeof(packets[0])) + return 0; + + packet = &packets[num_packets++]; + chained = stbi__get8(s); + packet->size = stbi__get8(s); + packet->type = stbi__get8(s); + packet->channel = stbi__get8(s); + act_comp |= packet->channel; + + if (stbi__at_eof(s)) { + stbi__rewind( s ); + return 0; + } + if (packet->size != 8) { + stbi__rewind( s ); + return 0; + } + } while (chained); + + *comp = (act_comp & 0x10 ? 4 : 3); + + return 1; +} +#endif + +// ************************************************************************************************* +// Portable Gray Map and Portable Pixel Map loader +// by Ken Miller +// +// PGM: http://netpbm.sourceforge.net/doc/pgm.html +// PPM: http://netpbm.sourceforge.net/doc/ppm.html +// +// Known limitations: +// Does not support comments in the header section +// Does not support ASCII image data (formats P2 and P3) + +#ifndef STBI_NO_PNM + +static int stbi__pnm_test(stbi__context *s) +{ + char p, t; + p = (char) stbi__get8(s); + t = (char) stbi__get8(s); + if (p != 'P' || (t != '5' && t != '6')) { + stbi__rewind( s ); + return 0; + } + return 1; +} + +static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *out; + STBI_NOTUSED(ri); + + ri->bits_per_channel = stbi__pnm_info(s, (int *)&s->img_x, (int *)&s->img_y, (int *)&s->img_n); + if (ri->bits_per_channel == 0) + return 0; + + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + *x = s->img_x; + *y = s->img_y; + if (comp) *comp = s->img_n; + + if (!stbi__mad4sizes_valid(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0)) + return stbi__errpuc("too large", "PNM too large"); + + out = (stbi_uc *) stbi__malloc_mad4(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0); + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + if (!stbi__getn(s, out, s->img_n * s->img_x * s->img_y * (ri->bits_per_channel / 8))) { + STBI_FREE(out); + return stbi__errpuc("bad PNM", "PNM file truncated"); + } + + if (req_comp && req_comp != s->img_n) { + if (ri->bits_per_channel == 16) { + out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, s->img_n, req_comp, s->img_x, s->img_y); + } else { + out = stbi__convert_format(out, s->img_n, req_comp, s->img_x, s->img_y); + } + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + return out; +} + +static int stbi__pnm_isspace(char c) +{ + return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r'; +} + +static void stbi__pnm_skip_whitespace(stbi__context *s, char *c) +{ + for (;;) { + while (!stbi__at_eof(s) && stbi__pnm_isspace(*c)) + *c = (char) stbi__get8(s); + + if (stbi__at_eof(s) || *c != '#') + break; + + while (!stbi__at_eof(s) && *c != '\n' && *c != '\r' ) + *c = (char) stbi__get8(s); + } +} + +static int stbi__pnm_isdigit(char c) +{ + return c >= '0' && c <= '9'; +} + +static int stbi__pnm_getinteger(stbi__context *s, char *c) +{ + int value = 0; + + while (!stbi__at_eof(s) && stbi__pnm_isdigit(*c)) { + value = value*10 + (*c - '0'); + *c = (char) stbi__get8(s); + if((value > 214748364) || (value == 214748364 && *c > '7')) + return stbi__err("integer parse overflow", "Parsing an integer in the PPM header overflowed a 32-bit int"); + } + + return value; +} + +static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp) +{ + int maxv, dummy; + char c, p, t; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + stbi__rewind(s); + + // Get identifier + p = (char) stbi__get8(s); + t = (char) stbi__get8(s); + if (p != 'P' || (t != '5' && t != '6')) { + stbi__rewind(s); + return 0; + } + + *comp = (t == '6') ? 3 : 1; // '5' is 1-component .pgm; '6' is 3-component .ppm + + c = (char) stbi__get8(s); + stbi__pnm_skip_whitespace(s, &c); + + *x = stbi__pnm_getinteger(s, &c); // read width + if(*x == 0) + return stbi__err("invalid width", "PPM image header had zero or overflowing width"); + stbi__pnm_skip_whitespace(s, &c); + + *y = stbi__pnm_getinteger(s, &c); // read height + if (*y == 0) + return stbi__err("invalid width", "PPM image header had zero or overflowing width"); + stbi__pnm_skip_whitespace(s, &c); + + maxv = stbi__pnm_getinteger(s, &c); // read max value + if (maxv > 65535) + return stbi__err("max value > 65535", "PPM image supports only 8-bit and 16-bit images"); + else if (maxv > 255) + return 16; + else + return 8; +} + +static int stbi__pnm_is16(stbi__context *s) +{ + if (stbi__pnm_info(s, NULL, NULL, NULL) == 16) + return 1; + return 0; +} +#endif + +static int stbi__info_main(stbi__context *s, int *x, int *y, int *comp) +{ + #ifndef STBI_NO_JPEG + if (stbi__jpeg_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PNG + if (stbi__png_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_GIF + if (stbi__gif_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_BMP + if (stbi__bmp_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PSD + if (stbi__psd_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PIC + if (stbi__pic_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PNM + if (stbi__pnm_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_HDR + if (stbi__hdr_info(s, x, y, comp)) return 1; + #endif + + // test tga last because it's a crappy test! + #ifndef STBI_NO_TGA + if (stbi__tga_info(s, x, y, comp)) + return 1; + #endif + return stbi__err("unknown image type", "Image not of any known type, or corrupt"); +} + +static int stbi__is_16_main(stbi__context *s) +{ + #ifndef STBI_NO_PNG + if (stbi__png_is16(s)) return 1; + #endif + + #ifndef STBI_NO_PSD + if (stbi__psd_is16(s)) return 1; + #endif + + #ifndef STBI_NO_PNM + if (stbi__pnm_is16(s)) return 1; + #endif + return 0; +} + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_info(char const *filename, int *x, int *y, int *comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result; + if (!f) return stbi__err("can't fopen", "Unable to open file"); + result = stbi_info_from_file(f, x, y, comp); + fclose(f); + return result; +} + +STBIDEF int stbi_info_from_file(FILE *f, int *x, int *y, int *comp) +{ + int r; + stbi__context s; + long pos = ftell(f); + stbi__start_file(&s, f); + r = stbi__info_main(&s,x,y,comp); + fseek(f,pos,SEEK_SET); + return r; +} + +STBIDEF int stbi_is_16_bit(char const *filename) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result; + if (!f) return stbi__err("can't fopen", "Unable to open file"); + result = stbi_is_16_bit_from_file(f); + fclose(f); + return result; +} + +STBIDEF int stbi_is_16_bit_from_file(FILE *f) +{ + int r; + stbi__context s; + long pos = ftell(f); + stbi__start_file(&s, f); + r = stbi__is_16_main(&s); + fseek(f,pos,SEEK_SET); + return r; +} +#endif // !STBI_NO_STDIO + +STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__info_main(&s,x,y,comp); +} + +STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *c, void *user, int *x, int *y, int *comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); + return stbi__info_main(&s,x,y,comp); +} + +STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__is_16_main(&s); +} + +STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *c, void *user) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); + return stbi__is_16_main(&s); +} + +#endif // STB_IMAGE_IMPLEMENTATION + +/* + revision history: + 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs + 2.19 (2018-02-11) fix warning + 2.18 (2018-01-30) fix warnings + 2.17 (2018-01-29) change sbti__shiftsigned to avoid clang -O2 bug + 1-bit BMP + *_is_16_bit api + avoid warnings + 2.16 (2017-07-23) all functions have 16-bit variants; + STBI_NO_STDIO works again; + compilation fixes; + fix rounding in unpremultiply; + optimize vertical flip; + disable raw_len validation; + documentation fixes + 2.15 (2017-03-18) fix png-1,2,4 bug; now all Imagenet JPGs decode; + warning fixes; disable run-time SSE detection on gcc; + uniform handling of optional "return" values; + thread-safe initialization of zlib tables + 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs + 2.13 (2016-11-29) add 16-bit API, only supported for PNG right now + 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes + 2.11 (2016-04-02) allocate large structures on the stack + remove white matting for transparent PSD + fix reported channel count for PNG & BMP + re-enable SSE2 in non-gcc 64-bit + support RGB-formatted JPEG + read 16-bit PNGs (only as 8-bit) + 2.10 (2016-01-22) avoid warning introduced in 2.09 by STBI_REALLOC_SIZED + 2.09 (2016-01-16) allow comments in PNM files + 16-bit-per-pixel TGA (not bit-per-component) + info() for TGA could break due to .hdr handling + info() for BMP to shares code instead of sloppy parse + can use STBI_REALLOC_SIZED if allocator doesn't support realloc + code cleanup + 2.08 (2015-09-13) fix to 2.07 cleanup, reading RGB PSD as RGBA + 2.07 (2015-09-13) fix compiler warnings + partial animated GIF support + limited 16-bpc PSD support + #ifdef unused functions + bug with < 92 byte PIC,PNM,HDR,TGA + 2.06 (2015-04-19) fix bug where PSD returns wrong '*comp' value + 2.05 (2015-04-19) fix bug in progressive JPEG handling, fix warning + 2.04 (2015-04-15) try to re-enable SIMD on MinGW 64-bit + 2.03 (2015-04-12) extra corruption checking (mmozeiko) + stbi_set_flip_vertically_on_load (nguillemot) + fix NEON support; fix mingw support + 2.02 (2015-01-19) fix incorrect assert, fix warning + 2.01 (2015-01-17) fix various warnings; suppress SIMD on gcc 32-bit without -msse2 + 2.00b (2014-12-25) fix STBI_MALLOC in progressive JPEG + 2.00 (2014-12-25) optimize JPG, including x86 SSE2 & NEON SIMD (ryg) + progressive JPEG (stb) + PGM/PPM support (Ken Miller) + STBI_MALLOC,STBI_REALLOC,STBI_FREE + GIF bugfix -- seemingly never worked + STBI_NO_*, STBI_ONLY_* + 1.48 (2014-12-14) fix incorrectly-named assert() + 1.47 (2014-12-14) 1/2/4-bit PNG support, both direct and paletted (Omar Cornut & stb) + optimize PNG (ryg) + fix bug in interlaced PNG with user-specified channel count (stb) + 1.46 (2014-08-26) + fix broken tRNS chunk (colorkey-style transparency) in non-paletted PNG + 1.45 (2014-08-16) + fix MSVC-ARM internal compiler error by wrapping malloc + 1.44 (2014-08-07) + various warning fixes from Ronny Chevalier + 1.43 (2014-07-15) + fix MSVC-only compiler problem in code changed in 1.42 + 1.42 (2014-07-09) + don't define _CRT_SECURE_NO_WARNINGS (affects user code) + fixes to stbi__cleanup_jpeg path + added STBI_ASSERT to avoid requiring assert.h + 1.41 (2014-06-25) + fix search&replace from 1.36 that messed up comments/error messages + 1.40 (2014-06-22) + fix gcc struct-initialization warning + 1.39 (2014-06-15) + fix to TGA optimization when req_comp != number of components in TGA; + fix to GIF loading because BMP wasn't rewinding (whoops, no GIFs in my test suite) + add support for BMP version 5 (more ignored fields) + 1.38 (2014-06-06) + suppress MSVC warnings on integer casts truncating values + fix accidental rename of 'skip' field of I/O + 1.37 (2014-06-04) + remove duplicate typedef + 1.36 (2014-06-03) + convert to header file single-file library + if de-iphone isn't set, load iphone images color-swapped instead of returning NULL + 1.35 (2014-05-27) + various warnings + fix broken STBI_SIMD path + fix bug where stbi_load_from_file no longer left file pointer in correct place + fix broken non-easy path for 32-bit BMP (possibly never used) + TGA optimization by Arseny Kapoulkine + 1.34 (unknown) + use STBI_NOTUSED in stbi__resample_row_generic(), fix one more leak in tga failure case + 1.33 (2011-07-14) + make stbi_is_hdr work in STBI_NO_HDR (as specified), minor compiler-friendly improvements + 1.32 (2011-07-13) + support for "info" function for all supported filetypes (SpartanJ) + 1.31 (2011-06-20) + a few more leak fixes, bug in PNG handling (SpartanJ) + 1.30 (2011-06-11) + added ability to load files via callbacks to accomidate custom input streams (Ben Wenger) + removed deprecated format-specific test/load functions + removed support for installable file formats (stbi_loader) -- would have been broken for IO callbacks anyway + error cases in bmp and tga give messages and don't leak (Raymond Barbiero, grisha) + fix inefficiency in decoding 32-bit BMP (David Woo) + 1.29 (2010-08-16) + various warning fixes from Aurelien Pocheville + 1.28 (2010-08-01) + fix bug in GIF palette transparency (SpartanJ) + 1.27 (2010-08-01) + cast-to-stbi_uc to fix warnings + 1.26 (2010-07-24) + fix bug in file buffering for PNG reported by SpartanJ + 1.25 (2010-07-17) + refix trans_data warning (Won Chun) + 1.24 (2010-07-12) + perf improvements reading from files on platforms with lock-heavy fgetc() + minor perf improvements for jpeg + deprecated type-specific functions so we'll get feedback if they're needed + attempt to fix trans_data warning (Won Chun) + 1.23 fixed bug in iPhone support + 1.22 (2010-07-10) + removed image *writing* support + stbi_info support from Jetro Lauha + GIF support from Jean-Marc Lienher + iPhone PNG-extensions from James Brown + warning-fixes from Nicolas Schulz and Janez Zemva (i.stbi__err. Janez (U+017D)emva) + 1.21 fix use of 'stbi_uc' in header (reported by jon blow) + 1.20 added support for Softimage PIC, by Tom Seddon + 1.19 bug in interlaced PNG corruption check (found by ryg) + 1.18 (2008-08-02) + fix a threading bug (local mutable static) + 1.17 support interlaced PNG + 1.16 major bugfix - stbi__convert_format converted one too many pixels + 1.15 initialize some fields for thread safety + 1.14 fix threadsafe conversion bug + header-file-only version (#define STBI_HEADER_FILE_ONLY before including) + 1.13 threadsafe + 1.12 const qualifiers in the API + 1.11 Support installable IDCT, colorspace conversion routines + 1.10 Fixes for 64-bit (don't use "unsigned long") + optimized upsampling by Fabian "ryg" Giesen + 1.09 Fix format-conversion for PSD code (bad global variables!) + 1.08 Thatcher Ulrich's PSD code integrated by Nicolas Schulz + 1.07 attempt to fix C++ warning/errors again + 1.06 attempt to fix C++ warning/errors again + 1.05 fix TGA loading to return correct *comp and use good luminance calc + 1.04 default float alpha is 1, not 255; use 'void *' for stbi_image_free + 1.03 bugfixes to STBI_NO_STDIO, STBI_NO_HDR + 1.02 support for (subset of) HDR files, float interface for preferred access to them + 1.01 fix bug: possible bug in handling right-side up bmps... not sure + fix bug: the stbi__bmp_load() and stbi__tga_load() functions didn't work at all + 1.00 interface to zlib that skips zlib header + 0.99 correct handling of alpha in palette + 0.98 TGA loader by lonesock; dynamically add loaders (untested) + 0.97 jpeg errors on too large a file; also catch another malloc failure + 0.96 fix detection of invalid v value - particleman@mollyrocket forum + 0.95 during header scan, seek to markers in case of padding + 0.94 STBI_NO_STDIO to disable stdio usage; rename all #defines the same + 0.93 handle jpegtran output; verbose errors + 0.92 read 4,8,16,24,32-bit BMP files of several formats + 0.91 output 24-bit Windows 3.0 BMP files + 0.90 fix a few more warnings; bump version number to approach 1.0 + 0.61 bugfixes due to Marc LeBlanc, Christopher Lloyd + 0.60 fix compiling as c++ + 0.59 fix warnings: merge Dave Moore's -Wall fixes + 0.58 fix bug: zlib uncompressed mode len/nlen was wrong endian + 0.57 fix bug: jpg last huffman symbol before marker was >9 bits but less than 16 available + 0.56 fix bug: zlib uncompressed mode len vs. nlen + 0.55 fix bug: restart_interval not initialized to 0 + 0.54 allow NULL for 'int *comp' + 0.53 fix bug in png 3->4; speedup png decoding + 0.52 png handles req_comp=3,4 directly; minor cleanup; jpeg comments + 0.51 obey req_comp requests, 1-component jpegs return as 1-component, + on 'test' only check type, not whether we support this variant + 0.50 (2006-11-19) + first released version +*/ + + +/* +------------------------------------------------------------------------------ +This software is available under 2 licenses -- choose whichever you prefer. +------------------------------------------------------------------------------ +ALTERNATIVE A - MIT License +Copyright (c) 2017 Sean Barrett +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +------------------------------------------------------------------------------ +ALTERNATIVE B - Public Domain (www.unlicense.org) +This is free and unencumbered software released into the public domain. +Anyone is free to copy, modify, publish, use, compile, sell, or distribute this +software, either in source code form or as a compiled binary, for any purpose, +commercial or non-commercial, and by any means. +In jurisdictions that recognize copyright laws, the author or authors of this +software dedicate any and all copyright interest in the software to the public +domain. We make this dedication for the benefit of the public at large and to +the detriment of our heirs and successors. We intend this dedication to be an +overt act of relinquishment in perpetuity of all present and future rights to +this software under copyright law. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +------------------------------------------------------------------------------ +*/ diff --git a/zig/vendor/stb/stb_image_impl.c b/zig/vendor/stb/stb_image_impl.c new file mode 100644 index 0000000000..908982ad8f --- /dev/null +++ b/zig/vendor/stb/stb_image_impl.c @@ -0,0 +1,10 @@ +// stb_image v2.30 (public domain / MIT, github.com/nothings/stb at 2c980bb59875b0d32144a71867fbdebb2f77cd20): +// the native server's JPEG and PNG decoder for image inputs (zig/src/vision/qwen_image.zig). +#define STB_IMAGE_IMPLEMENTATION +#define STBI_ONLY_JPEG +#define STBI_ONLY_PNG +#define STBI_NO_STDIO +#define STBI_NO_HDR +#define STBI_NO_LINEAR +#define STBI_NO_FAILURE_STRINGS +#include "stb_image.h" From 153d6b438625a65a3c2cf64da767f5519cd74c5b Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:41:58 +0000 Subject: [PATCH 25/46] golden check: the request log compares with LF endings on either side A checkout with core.autocrlf, as GitHub's Windows runners make, gives the frozen request log CRLF line endings while the server writes LF, so the lifecycle/request-log case differed at its first line's end on windows-latest and nowhere else. Both logs compare with CRLF read as LF; the other goldens are parsed as JSON and never saw the endings. --- zig/tests/server/golden_check.zig | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 13dc2d2b13..73d607350f 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -504,8 +504,9 @@ pub fn main(init: std.process.Init) !void { }); } else if (std.mem.eql(u8, kind, "compare_log")) { const golden_log = try std.fmt.allocPrint(init.arena.allocator(), "{s}/{s}", .{ golden_dir, t.get("golden").?.string }); - const want = Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch ""; - const got = Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch ""; + // a checkout with autocrlf gives the frozen log CRLF endings, the server writes LF: both compare as LF + const want = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); + const got = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); try outcomes.append(init.arena.allocator(), .{ .group = "lifecycle", .name = "request-log", From 22d92e4ae48f9e6874ab77676cbb1e867b5e0509 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:36:27 +0000 Subject: [PATCH 26/46] lane host: a warm pass that steps aside is not launched after its cancel; lean rule and zig fmt setup() returned cancel()'s true for a background warm pass that yields to a newer request, so the caller launched the job cancel() had just finished (SEGV in lane_host_warm_test). It now returns false, as for a drop. Comments one line each; lean_check skips zig/vendor/ (stb_image.h, as on #726). lane_host.zig and lanes/engine.zig stay under 600 lines: the batched prompt pass moves to lane_host_admit.zig (many, one) and lanes/open.zig (addStreams, adoptPrefilled), aliased where they were. --- tools/zig/lean_check.py | 4 +- zig/build/vision.zig | 3 +- zig/src/core/lane_host.zig | 77 +++++++------------------------- zig/src/core/lane_host_admit.zig | 50 +++++++++++++++++++++ zig/src/core/lanes/backend.zig | 10 ++--- zig/src/core/lanes/engine.zig | 61 ++----------------------- zig/src/core/lanes/open.zig | 55 +++++++++++++++++++++++ zig/src/core/lanes/stream.zig | 5 +-- zig/src/native/cuda.zig | 33 +++++++------- zig/src/server/chat.zig | 7 ++- zig/src/server/main.zig | 2 +- zig/src/server/media_fetch.zig | 20 +++------ zig/src/server/prompt.zig | 4 +- zig/src/server/vision_inputs.zig | 21 +++------ zig/src/vision/glm_image.zig | 8 +--- zig/src/vision/qwen_image.zig | 24 +++------- zig/src/vision/tf_video.c | 23 ++-------- 17 files changed, 179 insertions(+), 228 deletions(-) diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index a97aa37701..a5928f867b 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -19,6 +19,8 @@ "zig/src/families/flashnext/roles_gen.zig", "zig/kernels/metal/flashnext/sources_gen.zig", ) +# Third-party sources kept as their authors wrote them (stb_image.h) are skipped too. +VENDORED = "zig/vendor/" def generator_output(path: Path) -> bool: @@ -31,7 +33,7 @@ def generator_output(path: Path) -> bool: def problems(path: Path) -> list[str]: """Every rule break in one file, as 'path:line: reason'.""" - if generator_output(path): + if generator_output(path) or path.relative_to(ROOT).as_posix().startswith(VENDORED): return [] lines = path.read_text(errors="replace").splitlines() rel, found = path.relative_to(ROOT), [] diff --git a/zig/build/vision.zig b/zig/build/vision.zig index 11eb50c143..69950154dc 100644 --- a/zig/build/vision.zig +++ b/zig/build/vision.zig @@ -1,5 +1,4 @@ -//! The image-input module (zig/src/vision/qwen_image.zig with stb_image compiled in), shared by every build that -//! carries the engine API (the server reaches it as engine_api.qwen_image) and by the CUDA test runner. +//! The image-input module (vision/qwen_image.zig with stb_image), shared by the engine API and the CUDA test runner. const std = @import("std"); pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 7a557c20c3..9b07dc4d38 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -285,9 +285,8 @@ pub const LaneHost = struct { h.mutex.unlock(h.io); } - /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit yet). - fn takeJob(h: *LaneHost) ?*Job { + /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit). + pub fn takeJob(h: *LaneHost) ?*Job { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); @@ -310,9 +309,12 @@ pub const LaneHost = struct { } /// The admitted request's stream (its kept state, drafter and grammar); false when it was dropped instead. - fn setup(h: *LaneHost, job: *Job) bool { + pub fn setup(h: *LaneHost, job: *Job) bool { const r = job.request; - if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran + if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) { // other work came first: it never ran + _ = h.cancel(job); + return false; + } 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 (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { @@ -349,9 +351,9 @@ pub const LaneHost = struct { .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); { - _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); - return false; - } + _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return false; + } } else null, .images = r.images, .decode_spans = h.spans(r), @@ -369,8 +371,8 @@ pub const LaneHost = struct { return true; } - /// An admitted, set-up request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. - fn launch(h: *LaneHost, job: *Job) bool { + /// An admitted request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + pub fn launch(h: *LaneHost, job: *Job) bool { const r = job.request; const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); @@ -388,55 +390,6 @@ pub const LaneHost = struct { if (h.deliver(job)) h.remove(job); return true; } - fn admitOne(h: *LaneHost) bool { - const job = h.takeJob() orelse return false; - if (!h.setup(job)) return true; - return h.launch(job); - } - - /// Every queued request a lane is free for, their prompts in one pass where the backend takes several - /// (lanes.Engine.addStreams); the engine's special paths (chunked prompts, warm passes, the lone driver) go - /// through admitOne as before. - fn admitMany(h: *LaneHost) bool { - if (!h.core.batchesPrompts() or h.core.fills()) return h.admitOne(); - h.lock(); - const waiting = h.queued.items.len; - const free = h.info_.lanes -| h.admitted.items.len; - h.unlock(); - if (waiting < 2 or free < 2) return h.admitOne(); - var jobs: [lanes.Engine.max_together]*Job = undefined; - var n: usize = 0; - while (n < jobs.len and n < free) { - const job = h.takeJob() orelse break; - if (!h.setup(job)) continue; - if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { - _ = h.launch(job); - continue; - } - jobs[n] = job; - n += 1; - } - if (n == 0) return true; - if (n == 1) return h.launch(jobs[0]); - var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; - for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; - var errs: [lanes.Engine.max_together]?anyerror = undefined; - h.core.addStreams(streams[0..n], errs[0..n]); - for (jobs[0..n], errs[0..n]) |job, err| { - if (err) |e| { - _ = switch (e) { - error.Cancelled => h.cancel(job), - error.LaneMemory => h.requeue(job), - else => h.drop(job, h.words(e)), - }; - continue; - } - h.prefilled(job, job.began); - if (h.deliver(job)) h.remove(job); - } - return true; - } - /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt. fn fillOne(h: *LaneHost) void { @@ -515,7 +468,7 @@ pub const LaneHost = struct { } /// No memory for another lane yet: the job waits at the queue's head, and admission holds until a job ends. - fn requeue(h: *LaneHost, job: *Job) bool { + pub fn requeue(h: *LaneHost, job: *Job) bool { if (job.stream.grammar) |g| g.deinit(h.gpa); // the next admission compiles its own job.stream.deinit(h.gpa); job.proposer.deinit(); @@ -546,7 +499,7 @@ pub const LaneHost = struct { const alone_words = "the engine has no memory for this request's caches with nothing else running: lower --context or --parallel"; /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); h.finish(job, .cancelled, ""); return true; @@ -565,7 +518,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitMany()) {} + while (admit.many(h)) {} h.fillOne(); status.noteLive(h); h.lock(); diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index 016ab4398a..f394a3729e 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/zig/src/core/lane_host_admit.zig @@ -2,6 +2,8 @@ const api = @import("engine_api.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; +const lanes = @import("lanes"); +const lone = @import("lane_host_lone.zig"); /// The queued job to admit (its index), or null to hold; the caller holds the lock and a lane is free. pub fn pick(h: *LaneHost) ?usize { @@ -21,3 +23,51 @@ pub fn pick(h: *LaneHost) ?usize { fn fits(a: api.Admit, job: *const Job) bool { return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens); } + +/// Every queued request a lane is free for, prompts in one pass where the backend takes several; else one at a time. +pub fn many(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return one(h); + h.lock(); + const waiting = h.queued.items.len; + const free = h.info_.lanes -| h.admitted.items.len; + h.unlock(); + if (waiting < 2 or free < 2) return one(h); + var jobs: [lanes.Engine.max_together]*Job = undefined; + var n: usize = 0; + while (n < jobs.len and n < free) { + const job = h.takeJob() orelse break; + if (!h.setup(job)) continue; + if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { + _ = h.launch(job); + continue; + } + jobs[n] = job; + n += 1; + } + if (n == 0) return true; + if (n == 1) return h.launch(jobs[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], errs[0..n]) |job, err| { + if (err) |e| { + _ = switch (e) { + error.Cancelled => h.cancel(job), + error.LaneMemory => h.requeue(job), + else => h.drop(job, h.words(e)), + }; + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + return true; +} + +/// The next queued request admitted alone: its stream set up and launched. +fn one(h: *LaneHost) bool { + const job = h.takeJob() orelse return false; + if (!h.setup(job)) return true; + return h.launch(job); +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 723dc5effc..9d568efb4e 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -1,4 +1,4 @@ -//! What the round loop asks a backend (Metal, CUDA or the fake) to compute; drawn tokens stay on the GPU behind handles. +//! What the round loop asks a backend (Metal, CUDA or the fake); drawn tokens stay on the GPU behind handles. const std = @import("std"); const Stream = @import("stream.zig").Stream; const Shape = @import("shape.zig").Shape; @@ -68,8 +68,7 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, - /// Prefill several streams' prompts in one pass (each then as after `prefill`); false, having done nothing, - /// when they do not go together (each is then prefilled alone). + /// Several streams' prompts in one pass, each as after `prefill`; false, doing nothing, when they don't fit. prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, @@ -105,8 +104,7 @@ pub const Backend = struct { prompt_rows: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, - /// After a failed round: the streams it failed for alone (their caches could not grow), which the caller - /// ends while the round's other streams go on; 0 when the failure is the round's. + /// After a failed round: the streams whose caches could not grow, ended alone; 0 when the round failed. refused: ?*const fn (ptr: *anyopaque, out: []*Stream) usize = null, }; @@ -143,7 +141,7 @@ pub const Backend = struct { b.vtable.release(b.ptr, s); } - /// A driver of its own starts a stream: the prompt pass and its first token, committed; a cancel or failure releases it. + /// A driver of its own starts a stream: prompt pass and first token, committed; cancel or failure releases it. pub fn opening(b: Backend, gpa: std.mem.Allocator, s: *Stream) !u32 { errdefer b.release(s); try b.prefill(s); diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index d5be72c0bb..a3c6230ed9 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -83,69 +83,16 @@ pub const Engine = struct { pub const fills = open.fills; pub const fillStream = open.fillStream; - /// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. - pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { - _ = e.arena.reset(.retain_capacity); - s.context.shrinkRetainingCapacity(s.prompt_len); - s.pending = null; - s.cache_len = s.prompt_len; - if (e.cfg.family_mtp and s.drafts) { - const d: u32 = @intCast(try e.rule.depth(win.who(s))); - try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); - s.dropHeld(e.gpa); - s.next = .{ .count = d }; - if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { - s.next = held; - }; - } - if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row - _ = try s.commit(e.gpa, &.{token}, &.{}); - s.pending = token; - if (s.finished) return e.release(s); - try e.live.append(e.gpa, s); - } - - /// Streams a backend prefills in one pass at most. - pub const max_together = 16; + pub const adoptPrefilled = open.adoptPrefilled; + pub const max_together = open.max_together; + pub const addStreams = open.addStreams; /// Whether the backend prefills several prompts in one pass (addStreams batches them). pub fn batchesPrompts(e: *const Engine) bool { return e.backend.vtable.prefill_many != null; } - /// open.add for several streams: their prompts in one pass when the backend takes them together (else each - /// alone), then each opened as after its own prefill. ``errs[i]``: stream i's error (a cancelled one is released - /// already; the caller discards a failed one). - pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { - for (errs) |*x| x.* = null; - const many = e.backend.vtable.prefill_many orelse return e.addEach(ss, errs); - if (ss.len < 2 or ss.len > max_together) return e.addEach(ss, errs); - _ = e.arena.reset(.retain_capacity); - for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; - const together = many(e.backend.ptr, ss) catch |err| { - for (errs) |*x| x.* = err; - return; - }; - if (!together) return e.addEach(ss, errs); - for (ss, errs) |s, *x| { - if (s.isCancelled()) { // cancelled in its pass: no first token - e.backend.release(s); - x.* = error.Cancelled; - continue; - } - open.opened(e, s) catch |err| { - x.* = err; - }; - } - } - - fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { - for (ss, errs) |s, *x| open.add(e, s) catch |err| { - x.* = err; - }; - } - - /// Take over a stream another driver decoded so far: its cache settled, `pending` and `cache_len` set, no drafts held. + /// Take over a stream another driver decoded: cache settled, `pending` and `cache_len` set, no drafts held. pub fn adopt(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..c75fa81b86 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -88,3 +88,58 @@ pub fn fillStream(e: *Engine, s: *Stream, first: bool) !bool { try opened(e, s); return true; } + +/// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. +pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { + _ = e.arena.reset(.retain_capacity); + s.context.shrinkRetainingCapacity(s.prompt_len); + s.pending = null; + s.cache_len = s.prompt_len; + if (e.cfg.family_mtp and s.drafts) { + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; + }; + } + if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row + _ = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.finished) return e.release(s); + try e.live.append(e.gpa, s); +} + +/// Streams a backend prefills in one pass at most. +pub const max_together = 16; + +/// open.add for several streams, prompts in one pass when the backend allows; ``errs[i]``: stream i's error. +pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return addEach(e, ss, errs); + if (ss.len < 2 or ss.len > max_together) return addEach(e, ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch |err| { + for (errs) |*x| x.* = err; + return; + }; + if (!together) return addEach(e, ss, errs); + for (ss, errs) |s, *x| { + if (s.isCancelled()) { // cancelled in its pass: no first token + e.backend.release(s); + x.* = error.Cancelled; + continue; + } + opened(e, s) catch |err| { + x.* = err; + }; + } +} + +fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| add(e, s) catch |err| { + x.* = err; + }; +} diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index aade9ad50c..6d634c9600 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -76,7 +76,7 @@ pub const CancelCheck = struct { check: *const fn (ptr: *anyopaque) bool, }; -/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and where it keeps new ones. +/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and keeps new ones in. pub const Reuse = struct { saved: ?*anyopaque = null, // the backend's copy of the state after `at` prompt tokens (null: start at 0) at: u32 = 0, @@ -97,8 +97,7 @@ pub const Held = struct { parents: ?[]i32 = null, }; -/// An image input prepared for a vision tower: its patches (fp32 [grid h * w, 1536], merge-block order) and grid -/// [1, h, w]; the prompt holds its placeholder rows. +/// An image for a vision tower: patches (fp32 [h * w, 1536], merge-block order), grid [1, h, w]; prompt placeholders. pub const Image = struct { patches: []const f32, grid: [3]i64 }; pub const Spec = struct { diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a34186a4a3..7fbf38ba3d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,4 +1,5 @@ //! CUDA family registration, device setup, memory admission and the native lane host. + /// The serve flags this build's engine takes past the common ones (server/cli.zig served()): two ranks. pub const serves: []const []const u8 = &.{"--tp"}; @@ -252,8 +253,7 @@ const Host = struct { const refusals = @import("cuda_refusal.zig"); -/// A tensor-parallel rank above 0 (server/engines.zig follow): the family opened as that rank, then its replay loop -/// until rank 0 stops; false with ``problem`` set when it cannot. +/// A tensor-parallel rank above 0: the family opened as that rank, then its replay loop; false with ``problem`` set. pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { const opened = (try open(a, gpa, io, o, problem)) orelse return false; defer opened.close(opened.ctx); @@ -421,19 +421,19 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; if (comptime @hasField(@TypeOf(loaded), "learns")) { // a family with a Living Weights learner if (loaded.learns != null) { - const Slide = nemotron_slide.Adapter(F.slide); - h.dir = try gpa.dupe(u8, o.dir); - const s = try gpa.create(Slide); - s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); - h.slide = s; - h.slide_free = struct { - fn f(al: Allocator, p: *anyopaque) void { - const x: *Slide = @ptrCast(@alignCast(p)); - x.deinit(); - al.destroy(x); - } - }.f; - h.learn = s.hook(); + const Slide = nemotron_slide.Adapter(F.slide); + h.dir = try gpa.dupe(u8, o.dir); + const s = try gpa.create(Slide); + s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); + h.slide = s; + h.slide_free = struct { + fn f(al: Allocator, p: *anyopaque) void { + const x: *Slide = @ptrCast(@alignCast(p)); + x.deinit(); + al.destroy(x); + } + }.f; + h.learn = s.hook(); } } errdefer if (h.slide) |s| { @@ -441,8 +441,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, - .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; if (follow_loop) |f| { diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 82c72d52ad..63eebea817 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -92,8 +92,7 @@ pub const Prepared = struct { preparing: bool, }; -/// --tool-system: the server's instruction as the first message of a chat request that offers tools and sends no -/// system message (the client's own instruction wins); other requests pass unchanged. +/// --tool-system: the server's instruction first in a tools request with no system message; others pass unchanged. fn withToolSystem(srv: *Server, a: Allocator, input: Input) Allocator.Error!Input { const text = srv.config.tool_system orelse return input; if (input.tools.len == 0 or input.prompt != null or input.messages != .array) return input; @@ -318,7 +317,7 @@ const Generation = struct { return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; } - /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. + /// What closes a call the end token left open (``tool_parse.closeCall``); nothing for stop, length or ignore_eos. fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { const t = g.collected.items; if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; @@ -398,7 +397,7 @@ const Generation = struct { }; } - /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. + /// Ends the request and waits for the engine; one ended unfinished counts as a disconnect, as on the Mac. fn cancel(g: *Generation) void { if (!g.engine_done) g.srv.engine.cancel(g.id); g.drain(); diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 0cd1b03c6f..18a6d3a5f7 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -78,7 +78,7 @@ pub fn main(init: std.process.Init) !u8 { } orelse return fail(problem); defer up.text.deinit(); defer up.engine.close(up.engine.ctx); - if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls + if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls follow(up.engine.ctx) catch |e| return fail(@errorName(e)); return 0; } diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig index e175033bf6..c9ef04847f 100644 --- a/zig/src/server/media_fetch.zig +++ b/zig/src/server/media_fetch.zig @@ -1,7 +1,4 @@ -//! Public HTTPS media for --vision-urls, as the Python server fetches it (vision/images_http.py): an https URL on -//! port 443 without credentials or a fragment, its host resolved and every address checked public, a TLS connection -//! verified for the host on an address that was checked, GET with identity encoding, up to 3 redirects (each checked -//! the same way), a declared media type the caller accepts, the bytes within the caller's limit, all inside a deadline. +//! Public HTTPS media for --vision-urls as vision/images_http.py fetches it: checked hosts, TLS, 3 redirects, limits. const std = @import("std"); const Io = std.Io; @@ -22,8 +19,7 @@ fn fail(f: *Failure, text: []const u8) Error { return error.Media; } -/// Python ipaddress: is_global and not multicast or reserved, nor 192.0.0.0/24 or 168.63.129.16; no IPv6 address -/// that embeds an IPv4 one (mapped, 6to4, Teredo). +/// Python ipaddress is_global, less multicast, reserved, 192.0.0.0/24, 168.63.129.16 and IPv6 embedding an IPv4 one. pub fn publicIp(a: net.IpAddress) bool { switch (a) { .ip4 => |v| { @@ -37,9 +33,9 @@ pub fn publicIp(a: net.IpAddress) bool { } }.net4; const blocked = [_]struct { [4]u8, u5 }{ - .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, - .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, - .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, }; for (blocked) |blk| if (in(b, blk[0], blk[1])) return false; @@ -48,8 +44,7 @@ pub fn publicIp(a: net.IpAddress) bool { }, .ip6 => |v| { const b = v.bytes; - // global unicast 2000::/3 only, less documentation (2001:db8::/32), the 2001::/23 IETF block (Teredo - // 2001::/32 among it), 6to4 (2002::/16) + // 2000::/3 only, less 2001:db8::/32, the 2001::/23 IETF block (Teredo among it) and 6to4 2002::/16 if (b[0] & 0xe0 != 0x20) return false; if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; @@ -127,8 +122,7 @@ fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, std.mem.asBytes(&tv)) catch {}; } -/// A TCP connection to the checked address, its connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds a -/// Linux connect too; std's connect has no timeout yet). +/// A TCP connection to the checked address, connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds connect). fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { const posix = std.posix; const linux = std.os.linux; diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index 6b7faaf8e2..f48b0a76dd 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -53,9 +53,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } -/// GLM-5.3's template writes a reasoning-effort system line whatever the request asks; with thinking off the Python -/// frontend renders it without (glm5_next/prompts.py thinking_off: no effort line, an empty think block). The line is -/// its own segment between special tokens, so its ids alone are the ids it renders to; the first is dropped. +/// GLM-5.3's effort system line, dropped with thinking off as glm5_next/prompts.py does; its ids are its own segment. fn withoutEffortLine(srv: *Server, cx: *Cx, ids: []u32) errors.Refused![]u32 { const line = srv.text.encode(cx.a, "<|system|>Reasoning Effort: Max", false) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index fe457dfb34..04637f872b 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -1,9 +1,4 @@ -//! Image and video inputs of a chat request, as the Python server's frontend takes them (tensorfold/vision/images.py, -//! videos.py, qwen_processing.py): each ``image_url`` part (a data: URL; JPEG or PNG) decoded and prepared, each -//! ``video_url`` part (a data: URL; MP4, WebM...) decoded by libtfvideo and cut into frame groups, and each replaced in -//! its message by the text Python's frontend expands it to: ``<|vision_start|>`` + one ``<|image_pad|>`` a visual -//! token + ``<|vision_end|>`` for an image, and ``<|vision_start|>`` + ``<|video_pad|>``s + ``<|vision_end|>`` -//! for each frame group of a video. The template copies the text, and one tokenization gives Python's ids. +//! A chat request's image_url and video_url parts, prepared and replaced by the text Python's frontend expands them to. const std = @import("std"); const media_fetch = @import("media_fetch.zig"); @@ -26,8 +21,7 @@ const max_tokens_one: u64 = 4096; pub const Extracted = struct { messages: Value, images: []api.Image }; -/// The image URL of a content part (an OpenAI ``image_url`` part, or ``image`` / ``input_image`` spellings) and its -/// ``detail``; null for a part that is not an image. +/// The image URL of a content part (``image_url``, ``image`` or ``input_image``) and its ``detail``; else null. const Kind = enum { image, video }; fn mediaPart(part: Value) ?struct { kind: Kind, url: ?[]const u8, detail: []const u8 } { @@ -89,8 +83,7 @@ fn hasImages(messages: Value) bool { const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; -/// The bytes of a ``data:;base64,`` URL, or (--vision-urls) of a public HTTPS URL, fetched before -/// ``deadline`` (images.py _check_source and load_images). +/// The bytes of a ``data:`` URL, or (--vision-urls) a public HTTPS URL fetched before ``deadline`` (images.py). fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, deadline: std.Io.Timestamp) errors.Refused![]u8 { if (!std.mem.startsWith(u8, url, "data:")) { if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse(if (kind == .image) "images require data URLs or public HTTPS URLs" else "videos require data URLs or public HTTPS URLs"); @@ -114,8 +107,7 @@ fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, d return out; } -/// A request's images extracted and prepared, its messages with their image parts as the template's image text; -/// null when it has none. A server without --vision refuses images. +/// A request's images prepared, its image parts as the template's image text; null without images. Needs --vision. pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { const list = messages orelse return null; if (!hasImages(list)) return null; @@ -191,8 +183,7 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte try text.appendSlice(cx.a, vision_end); }, .video => if (srv.info.vision_glm) { - // GLM: <|begin_of_video|>, each frame group as <|begin_of_image|> <|image|>s <|end_of_image|> and its - // first frame's "t.t seconds", <|end_of_video|> (mlx-vlm's glm5_next processor) + // GLM: <|begin_of_video|>, image groups with their "t.t seconds", <|end_of_video|> (glm5_next) var v = try videoGlm(srv, cx, bytes); defer v.deinit(cx.a); try text.appendSlice(cx.a, "<|begin_of_video|>"); @@ -244,7 +235,7 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte return .{ .messages = .{ .array = out }, .images = images.items }; } -// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives -------------------------------------------------- +// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives ------------------------------------------ const TfVideo = extern struct { rgb: ?[*]u8, frames: c_int, height: c_int, width: c_int, indices: ?[*]c_int, fps: f64 }; const SizeFn = *const fn (ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int; diff --git a/zig/src/vision/glm_image.zig b/zig/src/vision/glm_image.zig index b46e75f076..889d3e679b 100644 --- a/zig/src/vision/glm_image.zig +++ b/zig/src/vision/glm_image.zig @@ -1,7 +1,4 @@ -//! GLM-5.3-Flash's image and video preparation (mlx-vlm's glm5_next/processing.py, as TensorFold's GLM frontend runs -//! it): the canvas (smart_resize to multiples of 28 within a visual-token budget), the content resized into it (PIL's -//! antialiased bicubic) and zero-padded at the bottom and right, CLIP's normalization (x * f32(1/255) - mean) / std, -//! and the patches in merge-block order ([c, t, 14, 14] each; an image's one frame twice as the temporal pair). +//! GLM-5.3-Flash's image and video preparation as mlx-vlm's glm5_next/processing.py: canvas, resize, norm, patches. const std = @import("std"); const qi = @import("qwen_image.zig"); @@ -128,8 +125,7 @@ pub fn prepare(gpa: std.mem.Allocator, im: qi.Image, min_tokens: u64, max_tokens return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; } -/// A video's frame groups: each group's patches (two consecutive frames as the temporal pair; an odd last frame -/// repeats) and its time (its first frame's seconds, as the processor's timestamps[::2]). +/// A video's frame groups: patches of two consecutive frames (an odd last repeats) and the first frame's seconds. pub const Video = struct { patches: []f32, // groups x gh x gw x 1176 groups: usize, diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index f918651811..e5331346eb 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -1,9 +1,4 @@ -//! Image inputs for Qwen-style vision towers, as the Python engine prepares them (tensorfold/vision/images.py and -//! qwen_processing.py with transformers' Qwen2VL processor): decode (JPEG, PNG; EXIF orientation applied, alpha over -//! white), smart_resize to multiples of patch x merge within a pixel budget, an antialiased bicubic resize (PIL's -//! filter and bounds with torch's uint8 fixed point: int16 weights at the most precision the axis allows), -//! (x - 127.5) / 127.5, and the patches in merge-block order, -//! each [channel, frame (the image twice), 16, 16] -- 1,536 values. +//! Image inputs for Qwen-style vision towers as Qwen2VL's processor makes them: decode, smart_resize, resize, patches. const std = @import("std"); /// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. @@ -160,7 +155,6 @@ pub fn smartResize(height: u32, width: u32, factor: u32, min_pixels: u64, max_pi return .{ @intFromFloat(hb), @intFromFloat(wb) }; } - fn bicubic(x0: f64) f64 { const a = -0.5; const x = @abs(x0); @@ -171,8 +165,7 @@ fn bicubic(x0: f64) f64 { const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; -/// PIL's precompute_coeffs for an axis of ``in_size`` resampled to ``out_size``, as fixed point the way torch's uint8 -/// antialias kernel takes them: int16 weights at the largest precision (< 22 bits) that keeps the largest in range. +/// PIL's precompute_coeffs for an axis, as int16 fixed point at the most precision (< 22 bits) torch's uint8 takes. fn coeffs(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); const filterscale = @max(scale, 1.0); @@ -222,8 +215,7 @@ fn clip8(v: i64, precision: u5) u8 { return @intCast(std.math.clamp(r, 0, 255)); } -/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits: the horizontal pass, then the vertical, each -/// rounding to bytes). +/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits): horizontal pass, then vertical, rounding to bytes. pub fn resize(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error!Image { var cur = im; var owned = false; @@ -289,8 +281,7 @@ pub const MERGE = 2; pub const TEMPORAL = 2; pub const FACTOR = PATCH * MERGE; -/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels as the Python frontend passes it: one merged patch), -/// normalized and cut into patches. +/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels: one merged patch), normalized and cut into patches. pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepared { const cap = max_tokens * FACTOR * FACTOR; const size = try smartResize(im.h, im.w, FACTOR, @min(FACTOR * FACTOR, cap), cap); @@ -322,8 +313,7 @@ pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepare // -- videos --------------------------------------------------------------------------------------------------------- -/// A video's frame size for the tower (the Python frontend's video_size): Qwen3-VL's smart_resize with its per-frame -/// cap (at most 768 tokens a frame group, at least ~134), the whole video within ``budget`` tokens. +/// A video's frame size: Qwen3-VL's smart_resize, at most 768 tokens a frame group, the video within ``budget``. pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u32 { const factor: f64 = FACTOR; const shortest: f64 = 128 * FACTOR * FACTOR; @@ -355,9 +345,7 @@ pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u return .{ @intFromFloat(hb), @intFromFloat(wb) }; } -/// A decoded video's frame groups for the tower: each group's patches (fp32 [gh * gw, 1536], two consecutive frames -/// as the temporal pair; an odd last frame repeats), normalized as numpy does ((x * f32(1/255) - 0.5) / 0.5), and -/// each group's time (the mean of its two frames' seconds). +/// A video's frame groups: patches of two consecutive frames, normalized as numpy does, and their mean time. pub const Video = struct { patches: []f32, // groups x gh x gw x 1536 groups: usize, diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c index 1049416565..6436b675e4 100644 --- a/zig/src/vision/tf_video.c +++ b/zig/src/vision/tf_video.c @@ -1,19 +1,4 @@ -// libtfvideo: a video's sampled frames for the native server's vision inputs, as the Python engine decodes them -// (tensorfold/vision/videos.py with PyAV): the first video stream; its rate (average, else guessed, else 24) and -// frame count (the stream's, else its duration times the rate); frames sampled at ``rate`` a second over the whole -// video (Qwen3-VL's rule: at least min_frames, at most max_frames, numpy's linspace rounded half to even), decoded -// in order and scaled straight to the size ``size`` picks as RGB the way PyAV's VideoFrame.to_ndarray(rgb24, width, -// height, interpolation="BICUBIC") does it (VideoReformatter: the frame's colour space and range kept, transfer and -// primaries unspecified, sws_scale_frame with SWS_BICUBIC alone). Built against FFmpeg 9 (the FFmpeg PyAV 19.0.1 -// ships, so the frames equal Python's); the server loads it when a request carries a video. -#include -#include -#include -#include -#include -#include -#include -#include +// libtfvideo: a video's sampled frames as videos.py decodes them with PyAV 19.0.1 (FFmpeg 9), RGB via SWS_BICUBIC. typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); @@ -56,8 +41,7 @@ static int64_t mem_seek(void* o, int64_t off, int whence) { return at; } -// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): an SWS_CS_* value names the frame's colour space -// (the frame's own AVColorSpace number passes through it as PyAV passes it) +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): the frame's own AVColorSpace passes through. static void set_colorspace(AVFrame* f, int cs) { if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; @@ -171,8 +155,7 @@ int tf_video_decode(const uint8_t* data, size_t len, double rate, int min_frames if (e < 0) { finished = 1; break; } if (n > last) { finished = 1; av_frame_unref(frame); break; } if (next < count && n == wanted[next]) { - // PyAV's _reformat: the destination copies the frame's properties; its colour space and range pass - // through _set_frame_colorspace with the frame's own values; transfer and primaries unspecified + // PyAV's _reformat: the frame's colour space and range kept; transfer and primaries unspecified av_frame_unref(rgb); av_frame_copy_props(rgb, frame); set_colorspace(frame, (int)frame->colorspace); From ef73ec3d0e8b2eb751bc55c0752b002dfd9357a0 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:36:27 +0000 Subject: [PATCH 27/46] tf_video: the includes back under its one-line header --- zig/src/vision/tf_video.c | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c index 6436b675e4..9cd1290cc7 100644 --- a/zig/src/vision/tf_video.c +++ b/zig/src/vision/tf_video.c @@ -1,4 +1,12 @@ // libtfvideo: a video's sampled frames as videos.py decodes them with PyAV 19.0.1 (FFmpeg 9), RGB via SWS_BICUBIC. +#include +#include +#include +#include +#include +#include +#include +#include typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); From cc5d9d16fe45d3631248fe9461f98cc2e142d968 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:45:12 +0000 Subject: [PATCH 28/46] vision inputs: a server without --vision refuses image parts in the frozen words (the errors/content-image golden) --- zig/src/server/vision_inputs.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index 04637f872b..c63bce5edf 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -111,7 +111,7 @@ fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, d pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { const list = messages orelse return null; if (!hasImages(list)) return null; - if (!srv.info.vision) return cx.refuse("image and video inputs need a server started with --vision"); + if (!srv.info.vision) return cx.refuse("this server accepts text parts only; image, audio and video inputs are unsupported"); // every media part, in prompt order, with where its text goes const Slot = struct { kind: Kind, url: ?[]const u8, detail: []const u8, msg: usize, part: usize }; var slots: std.ArrayList(Slot) = .empty; From aaed024a2edbbfa82e96d91fb185145b7118f236 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:57:12 -0500 Subject: [PATCH 29/46] qwen_image.formatSeconds: t * 10 rounded in integers, not f128, so tensorfold-native takes memcpy from libSystem again (zig build test on macOS) --- zig/src/vision/qwen_image.zig | 53 +++++++++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 6 deletions(-) diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index e5331346eb..0601d3c858 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -406,15 +406,56 @@ pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, h: u /// Python's f"{t:.1f}": one decimal, an exact tie rounded to even. pub fn formatSeconds(buf: []u8, t: f64) ![]const u8 { - const x: f128 = @as(f128, t) * 10; // exact: a double's 53 bits times ten fit f128's 113 - const fl = @floor(x); - const frac = x - fl; - var r = if (frac > 0.5) fl + 1 else fl; - if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; - const v: i64 = @intFromFloat(r); + const v = tenths(t); return std.fmt.bufPrint(buf, "{d}.{d}", .{ @divTrunc(v, 10), @as(u64, @intCast(@mod(v, 10))) }); } +/// t times ten, an exact tie to even, in integers: t is m * 2^e and 10m fits 57 bits (f128 would link compiler-rt). +fn tenths(t: f64) i64 { + const bits: u64 = @bitCast(@abs(t)); + const biased: i32 = @intCast(bits >> 52); + const m: u64 = if (biased == 0) bits & ((1 << 52) - 1) else (bits & ((1 << 52) - 1)) | (1 << 52); + const e = @as(i32, @max(biased, 1)) - 1075; + if (e >= 0) return @intFromFloat(t * 10); // at least 2^52 seconds: never a frame's time + const tm = m * 10; + if (e <= -58) return 0; // t * 10 under a half: tm < 2^57 <= 2^(s - 1) + const s: u6 = @intCast(-e); + const q = tm >> s; + const rem = tm & ((@as(u64, 1) << s) - 1); + const half = @as(u64, 1) << (s - 1); + const r: i64 = @intCast(if (rem > half or (rem == half and q & 1 == 1)) q + 1 else q); + return if (t < 0) -r else r; +} + +test "tenths rounds t * 10 exactly, ties to even, as the f128 product does" { + const ref = struct { + fn f(t: f64) i64 { + const x: f128 = @as(f128, t) * 10; + const fl = @floor(x); + const frac = x - fl; + var r = if (frac > 0.5) fl + 1 else fl; + if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; + return @intFromFloat(r); + } + }.f; + const fixed = [_]f64{ 0, -0.0, 0.05, 0.15, 0.25, 0.35, 0.45, 1.25, 2.5, 0.95, 1.05, 3.0, 123.45, 1e-300, 5e-324, 4503599627370495.5, -0.25, -1.35 }; + for (fixed) |t| try std.testing.expectEqual(ref(t), tenths(t)); + var prng: std.Random.DefaultPrng = .init(0x7e17); + const rand = prng.random(); + for (0..200_000) |i| { + const t: f64 = switch (i % 4) { + 0 => rand.float(f64) * 4000, + 1 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 0, 80_000))) / 20, // halves of tenths + 2 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 100_000))) / @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 121))), + else => @bitCast(@as(u64, rand.uintLessThan(u64, 1075)) << 52 | rand.int(u64) >> 12), // any double below 2^52 + }; + try std.testing.expectEqual(ref(t), tenths(t)); + } + var buf: [32]u8 = undefined; + try std.testing.expectEqualStrings("0.2", try formatSeconds(&buf, 0.25)); + try std.testing.expectEqualStrings("1.4", try formatSeconds(&buf, 1.35)); +} + test { _ = glm; } From 9e6dfcd4c0be11305c304db901337fa9b8cc6d9f Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 07:30:04 +0000 Subject: [PATCH 30/46] lane host: a set-up job's kept entry stays pinned until its restore reports; images wait for the shared decoder; no refused hook - prompt_cache: Entry.pins, Store.pin/unpin. victimBut passes pinned entries over, so a later job's lookup in admit.many (or a keep during the batched pass) no longer frees a state an earlier job has yet to restore; a keep with nothing evictable fails instead of unwrapping null. LaneHost.setup pins, reported and requeue unpin. - vision: zig/vendor/stb dropped and no longer compiled. qwen_image.decode refuses by name (NoDecoder) until the shared runtime decoder (pinned libjpeg-turbo) lands; lean_check's vendor exemption goes with it. - lanes: Backend.VTable.refused, Engine.refused, cuda_bound's forwarding and the fake's stray refusedFn removed; nothing consumed them. --- tools/zig/lean_check.py | 4 +- zig/build/vision.zig | 7 +- zig/src/core/lane_host.zig | 5 +- zig/src/core/lanes/backend.zig | 2 - zig/src/core/lanes/engine.zig | 6 - zig/src/core/lanes/fake.zig | 8 - zig/src/core/prompt_cache.zig | 19 +- zig/src/core/prompt_cache_test.zig | 30 + zig/src/native/cuda_bound.zig | 6 - zig/src/server/vision_inputs.zig | 1 + zig/src/vision/qwen_image.zig | 108 +- zig/vendor/stb/README.md | 3 - zig/vendor/stb/stb_image.h | 7988 ---------------------------- zig/vendor/stb/stb_image_impl.c | 10 - 14 files changed, 56 insertions(+), 8141 deletions(-) delete mode 100644 zig/vendor/stb/README.md delete mode 100644 zig/vendor/stb/stb_image.h delete mode 100644 zig/vendor/stb/stb_image_impl.c diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index a5928f867b..a97aa37701 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -19,8 +19,6 @@ "zig/src/families/flashnext/roles_gen.zig", "zig/kernels/metal/flashnext/sources_gen.zig", ) -# Third-party sources kept as their authors wrote them (stb_image.h) are skipped too. -VENDORED = "zig/vendor/" def generator_output(path: Path) -> bool: @@ -33,7 +31,7 @@ def generator_output(path: Path) -> bool: def problems(path: Path) -> list[str]: """Every rule break in one file, as 'path:line: reason'.""" - if generator_output(path) or path.relative_to(ROOT).as_posix().startswith(VENDORED): + if generator_output(path): return [] lines = path.read_text(errors="replace").splitlines() rel, found = path.relative_to(ROOT), [] diff --git a/zig/build/vision.zig b/zig/build/vision.zig index 69950154dc..0b78064367 100644 --- a/zig/build/vision.zig +++ b/zig/build/vision.zig @@ -1,9 +1,6 @@ -//! The image-input module (vision/qwen_image.zig with stb_image), shared by the engine API and the CUDA test runner. +//! The image-input module (vision/qwen_image.zig), shared by the engine API and the CUDA test runner. const std = @import("std"); pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { - const m = b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); - m.addCSourceFile(.{ .file = b.path("zig/vendor/stb/stb_image_impl.c"), .flags = &.{ "-std=c11", "-O2" } }); - m.addIncludePath(b.path("zig/vendor/stb")); - return m; + return b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); } diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 9b07dc4d38..9ddbf6a8fe 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -85,6 +85,7 @@ pub const LaneHost = struct { const e = job.entry orelse return; job.entry = null; const store = job.host.cache orelse return; + store.unpin(e); if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); @@ -316,9 +317,8 @@ pub const LaneHost = struct { return false; } 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 (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; + job.entry = if (l.entry) |e| store.pin(e) else null; // kept alive until its restore reports 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 } }; @@ -474,6 +474,7 @@ pub const LaneHost = struct { job.proposer.deinit(); job.started = false; job.fill_began = false; + if (job.entry) |e| h.cache.?.unpin(e); job.entry = null; // looked up again at its next admission h.gpa.free(job.marks); job.marks = &.{}; diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 9d568efb4e..23a4efd892 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -104,8 +104,6 @@ pub const Backend = struct { prompt_rows: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, - /// After a failed round: the streams whose caches could not grow, ended alone; 0 when the round failed. - refused: ?*const fn (ptr: *anyopaque, out: []*Stream) usize = null, }; pub fn prefill(b: Backend, s: *Stream) !void { diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index a3c6230ed9..e4c7950781 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -109,12 +109,6 @@ pub const Engine = struct { return false; } - /// After ``step`` failed: the streams the backend refused alone (the caller ends them; the rest go on). - pub fn refused(e: *Engine, out: []*Stream) usize { - const ask = e.backend.vtable.refused orelse return 0; - return ask(e.backend.ptr, out); - } - /// Release a cancelled stream between rounds: no pending draws or drafts. pub fn discard(e: *Engine, s: *Stream) void { s.finished = true; diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index fc90ee7483..9434736c46 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -243,14 +243,6 @@ pub const Fake = struct { return self(ptr).drawn.items[handle]; } - fn refusedFn(ptr: *anyopaque, out: []*Stream) usize { - const x = self(ptr); - const n = @min(out.len, x.refused.items.len); - @memcpy(out[0..n], x.refused.items[0..n]); - x.refused.clearRetainingCapacity(); - return n; - } - fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const x = self(ptr); x.rounds += 1; diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index f1465325e1..6fe20982f7 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -108,6 +108,7 @@ pub const Entry = struct { last: []u32, // the last prompt that kept or resumed it shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded + pins: u32 = 0, // admitted jobs that restore it and have not yet: eviction passes it over }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -161,6 +162,16 @@ pub const Store = struct { s.free(e); } + /// `e` waits for a restore: no lookup or keep evicts it until `unpin` (jobs set up after it run their own first). + pub fn pin(_: *Store, e: *Entry) *Entry { + e.pins += 1; + return e; + } + + pub fn unpin(_: *Store, e: *Entry) void { + e.pins -= 1; + } + /// The family's free storage kept for later saves. pub fn spare(s: *const Store) u64 { return if (s.family.vtable.spare) |f| f(s.family.ptr) else 0; @@ -450,7 +461,7 @@ pub const Store = struct { if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held -| bytes); // spare that cannot take this state goes before any state if (s.spare() < free_) continue; } - s.remove(s.victimBut(null, prompt, born).?); // its storage may come back spare and take this state + s.remove(s.victimBut(null, prompt, born) orelse return s.fail(at, error.RestoresPending)); // spare for this state s.counts.evicted += 1; } if (s.entries.items.len == 0) if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| bytes); @@ -534,17 +545,17 @@ pub const Store = struct { return false; } - /// The entry to free first, never `skip`: a superseded one, no branch point, oldest; else the oldest, rewind last. + /// The entry to free first, never `skip` or pinned: superseded, no branch point, oldest; else oldest, rewind last. fn victimBut(s: *const Store, skip: ?*Entry, prompt: []const u32, born: u32) ?usize { var best: ?usize = null; for (s.entries.items, 0..) |e, i| { - if (e.shared or e.branched or e == skip or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used + if (e.shared or e.branched or e == skip or e.pins > 0 or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used const now = born > e.born and prompt.len > e.tokens.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); // a longer request's tokens extend it if ((now or s.extended(e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; } if (best) |i| return i; var oldest: ?usize = null; - for (s.entries.items, 0..) |e, i| if (e != skip and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { + for (s.entries.items, 0..) |e, i| if (e != skip and e.pins == 0 and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { oldest = i; }; return oldest; diff --git a/zig/src/core/prompt_cache_test.zig b/zig/src/core/prompt_cache_test.zig index 8452036b75..b408b3c7ad 100644 --- a/zig/src/core/prompt_cache_test.zig +++ b/zig/src/core/prompt_cache_test.zig @@ -391,3 +391,33 @@ test "a kept state grows only into what the budget leaves, and evicts nothing to var stranger: u8 = 0; try std.testing.expect(!s.grow(&stranger, 0)); // a state the store doesn't hold } + +test "an entry pinned for a restore still to come outlives the lookups set up after it; unpinned, it goes again" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 220); // fake bytes: 100 + at + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); // 103 bytes, the oldest + const y = [_]u32{ 6, 6, 6, 6 }; + _ = f.pass(&s, &y, try s.begin(a, &y, 3, &.{}, &.{}, null, &.{})); // 103: 206 held + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const l = try s.lookup(a, &xa, 3, &.{}, &.{}, &.{}); // a job set up to resume x + const e = l.entry.?; + _ = s.pin(e); + const z = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7 }; + _ = try s.lookup(a, &z, 7, &.{}, &.{}, &.{}); // the next job's room: y goes, not x + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + try std.testing.expect(s.find(&.{ 6, 6, 6, 6 }, &.{}, &.{}) == null); + const w: [130]u32 = @splat(8); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); // 220 bytes, and nothing but x to evict: it stays + try std.testing.expect(!s.keep(&w, 120, null, &.{}, &.{})); // and a keep with no room is refused + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + s.unpin(e); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(s.entries.items.len, f.live); +} diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 2144981686..fe929ae849 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -114,12 +114,6 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.held.?(x.inner.ptr, s, out); } }.f else null, - .refused = if (v.refused != null) struct { - fn f(p: *anyopaque, out: []*lanes.Stream) usize { - const x = H.bind(p); - return x.inner.vtable.refused.?(x.inner.ptr, out); - } - }.f else null, .prompt_rows = v.prompt_rows, .release = struct { fn f(p: *anyopaque, s: *lanes.Stream) void { diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index c63bce5edf..b1e4476fb1 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -154,6 +154,7 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte var img = qi.decode(cx.a, bytes, .{}) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, error.ImageTooLarge => return cx.refuse("image dimensions exceed the decoded pixel limit"), + error.NoDecoder => return cx.refuse("image inputs need the shared JPEG and PNG decoder, which this server does not have yet"), else => return cx.refuse("image bytes are invalid or unsupported; use JPEG or PNG"), }; total_pixels += @as(u64, img.w) * img.h; diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index 0601d3c858..79e7c4be3a 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -4,10 +4,6 @@ const std = @import("std"); /// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. pub const glm = @import("glm_image.zig"); -extern fn stbi_info_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int) c_int; -extern fn stbi_load_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int, req: c_int) ?[*]u8; -extern fn stbi_image_free(p: ?*anyopaque) void; - pub const Limits = struct { max_dimension: u32 = 8192, max_pixels: u64 = 16 * 1024 * 1024, @@ -24,108 +20,12 @@ pub const Image = struct { } }; -pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio }; +pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio, NoDecoder }; -/// JPEG or PNG bytes to RGB: EXIF orientation applied (JPEG), any alpha composited over white. +/// JPEG or PNG bytes to RGB: the shared runtime decoder's job (pinned libjpeg-turbo); until it lands, refused by name. pub fn decode(gpa: std.mem.Allocator, bytes: []const u8, limits: Limits) Error!Image { - if (bytes.len == 0 or bytes.len > std.math.maxInt(c_int)) return error.UnsupportedImage; - var w: c_int = 0; - var h: c_int = 0; - var comp: c_int = 0; - if (stbi_info_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp) == 0) return error.UnsupportedImage; - if (w <= 0 or h <= 0 or w > limits.max_dimension or h > limits.max_dimension or - @as(u64, @intCast(w)) * @as(u64, @intCast(h)) > limits.max_pixels) return error.ImageTooLarge; - const alpha = comp == 2 or comp == 4; - const want: c_int = if (alpha) 4 else 3; - var got: c_int = 0; - const px = stbi_load_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp, want) orelse return error.UnsupportedImage; - defer stbi_image_free(px); - const n: usize = @intCast(w * h); - _ = &got; - var rgb = try gpa.alloc(u8, n * 3); - errdefer gpa.free(rgb); - if (alpha) { - // PIL's paste of RGBA onto white with the alpha as mask: out = white + (src - white) * a / 255, rounded - for (0..n) |i| { - const a: i32 = px[i * 4 + 3]; - for (0..3) |c| { - const s: i32 = px[i * 4 + c]; - const t: i32 = (s - 255) * a + 128; - rgb[i * 3 + c] = @intCast(255 + ((t + (t >> 8)) >> 8)); - } - } - } else @memcpy(rgb, px[0 .. n * 3]); - var im: Image = .{ .w = @intCast(w), .h = @intCast(h), .rgb = rgb }; - const o = orientation(bytes); - if (o > 1) { - const t = try transpose(gpa, im, o); - im.deinit(gpa); - im = t; - } - return im; -} - -/// A JPEG's EXIF orientation (1-8; 1 when absent or not a JPEG). -fn orientation(b: []const u8) u16 { - if (b.len < 4 or b[0] != 0xFF or b[1] != 0xD8) return 1; - var i: usize = 2; - while (i + 4 <= b.len and b[i] == 0xFF) { - const marker = b[i + 1]; - const len = (@as(usize, b[i + 2]) << 8) | b[i + 3]; - if (marker == 0xDA or len < 2 or i + 2 + len > b.len) break; - if (marker == 0xE1 and len >= 16 and std.mem.eql(u8, b[i + 4 .. i + 10], "Exif\x00\x00")) { - const t = b[i + 10 .. i + 2 + len]; - const le = t[0] == 'I'; - const rd16 = struct { - fn f(x: []const u8, at: usize, little: bool) u16 { - return if (little) std.mem.readInt(u16, x[at..][0..2], .little) else std.mem.readInt(u16, x[at..][0..2], .big); - } - }.f; - const rd32 = struct { - fn f(x: []const u8, at: usize, little: bool) u32 { - return if (little) std.mem.readInt(u32, x[at..][0..4], .little) else std.mem.readInt(u32, x[at..][0..4], .big); - } - }.f; - if (t.len < 8) return 1; - const ifd = rd32(t, 4, le); - if (ifd + 2 > t.len) return 1; - const count = rd16(t, ifd, le); - for (0..count) |e| { - const at = ifd + 2 + e * 12; - if (at + 12 > t.len) return 1; - if (rd16(t, at, le) == 0x0112) { - const v = rd16(t, at + 8, le); - return if (v >= 1 and v <= 8) v else 1; - } - } - return 1; - } - i += 2 + len; - } - return 1; -} - -/// PIL's exif_transpose for orientations 2-8. -fn transpose(gpa: std.mem.Allocator, im: Image, o: u16) Error!Image { - const swap = o >= 5; - const W: u32 = if (swap) im.h else im.w; - const H: u32 = if (swap) im.w else im.h; - const out = try gpa.alloc(u8, im.rgb.len); - for (0..H) |y| for (0..W) |x| { - // the source pixel of output (x, y) - const sx: u32, const sy: u32 = switch (o) { - 2 => .{ im.w - 1 - @as(u32, @intCast(x)), @intCast(y) }, - 3 => .{ im.w - 1 - @as(u32, @intCast(x)), im.h - 1 - @as(u32, @intCast(y)) }, - 4 => .{ @intCast(x), im.h - 1 - @as(u32, @intCast(y)) }, - 5 => .{ @intCast(y), @intCast(x) }, - 6 => .{ @intCast(y), im.h - 1 - @as(u32, @intCast(x)) }, - 7 => .{ im.w - 1 - @as(u32, @intCast(y)), im.h - 1 - @as(u32, @intCast(x)) }, - 8 => .{ im.w - 1 - @as(u32, @intCast(y)), @intCast(x) }, - else => .{ @intCast(x), @intCast(y) }, - }; - @memcpy(out[(y * W + x) * 3 ..][0..3], im.rgb[(sy * im.w + sx) * 3 ..][0..3]); - }; - return .{ .w = W, .h = H, .rgb = out }; + _ = .{ gpa, bytes, limits }; + return error.NoDecoder; } /// Python's round(): halves to even. diff --git a/zig/vendor/stb/README.md b/zig/vendor/stb/README.md deleted file mode 100644 index 11466f4316..0000000000 --- a/zig/vendor/stb/README.md +++ /dev/null @@ -1,3 +0,0 @@ -stb_image.h v2.30 from https://github.com/nothings/stb at commit 2c980bb59875b0d32144a71867fbdebb2f77cd20 -(sha256 594c2fe35d49488b4382dbfaec8f98366defca819d916ac95becf3e75f4200b3), unmodified; dual public domain / MIT -(see the license at the end of the file). stb_image_impl.c compiles it for JPEG and PNG only. diff --git a/zig/vendor/stb/stb_image.h b/zig/vendor/stb/stb_image.h deleted file mode 100644 index 9eedabedc4..0000000000 --- a/zig/vendor/stb/stb_image.h +++ /dev/null @@ -1,7988 +0,0 @@ -/* stb_image - v2.30 - public domain image loader - http://nothings.org/stb - no warranty implied; use at your own risk - - Do this: - #define STB_IMAGE_IMPLEMENTATION - before you include this file in *one* C or C++ file to create the implementation. - - // i.e. it should look like this: - #include ... - #include ... - #include ... - #define STB_IMAGE_IMPLEMENTATION - #include "stb_image.h" - - You can #define STBI_ASSERT(x) before the #include to avoid using assert.h. - And #define STBI_MALLOC, STBI_REALLOC, and STBI_FREE to avoid using malloc,realloc,free - - - QUICK NOTES: - Primarily of interest to game developers and other people who can - avoid problematic images and only need the trivial interface - - JPEG baseline & progressive (12 bpc/arithmetic not supported, same as stock IJG lib) - PNG 1/2/4/8/16-bit-per-channel - - TGA (not sure what subset, if a subset) - BMP non-1bpp, non-RLE - PSD (composited view only, no extra channels, 8/16 bit-per-channel) - - GIF (*comp always reports as 4-channel) - HDR (radiance rgbE format) - PIC (Softimage PIC) - PNM (PPM and PGM binary only) - - Animated GIF still needs a proper API, but here's one way to do it: - http://gist.github.com/urraka/685d9a6340b26b830d49 - - - decode from memory or through FILE (define STBI_NO_STDIO to remove code) - - decode from arbitrary I/O callbacks - - SIMD acceleration on x86/x64 (SSE2) and ARM (NEON) - - Full documentation under "DOCUMENTATION" below. - - -LICENSE - - See end of file for license information. - -RECENT REVISION HISTORY: - - 2.30 (2024-05-31) avoid erroneous gcc warning - 2.29 (2023-05-xx) optimizations - 2.28 (2023-01-29) many error fixes, security errors, just tons of stuff - 2.27 (2021-07-11) document stbi_info better, 16-bit PNM support, bug fixes - 2.26 (2020-07-13) many minor fixes - 2.25 (2020-02-02) fix warnings - 2.24 (2020-02-02) fix warnings; thread-local failure_reason and flip_vertically - 2.23 (2019-08-11) fix clang static analysis warning - 2.22 (2019-03-04) gif fixes, fix warnings - 2.21 (2019-02-25) fix typo in comment - 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs - 2.19 (2018-02-11) fix warning - 2.18 (2018-01-30) fix warnings - 2.17 (2018-01-29) bugfix, 1-bit BMP, 16-bitness query, fix warnings - 2.16 (2017-07-23) all functions have 16-bit variants; optimizations; bugfixes - 2.15 (2017-03-18) fix png-1,2,4; all Imagenet JPGs; no runtime SSE detection on GCC - 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs - 2.13 (2016-12-04) experimental 16-bit API, only for PNG so far; fixes - 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes - 2.11 (2016-04-02) 16-bit PNGS; enable SSE2 in non-gcc x64 - RGB-format JPEG; remove white matting in PSD; - allocate large structures on the stack; - correct channel count for PNG & BMP - 2.10 (2016-01-22) avoid warning introduced in 2.09 - 2.09 (2016-01-16) 16-bit TGA; comments in PNM files; STBI_REALLOC_SIZED - - See end of file for full revision history. - - - ============================ Contributors ========================= - - Image formats Extensions, features - Sean Barrett (jpeg, png, bmp) Jetro Lauha (stbi_info) - Nicolas Schulz (hdr, psd) Martin "SpartanJ" Golini (stbi_info) - Jonathan Dummer (tga) James "moose2000" Brown (iPhone PNG) - Jean-Marc Lienher (gif) Ben "Disch" Wenger (io callbacks) - Tom Seddon (pic) Omar Cornut (1/2/4-bit PNG) - Thatcher Ulrich (psd) Nicolas Guillemot (vertical flip) - Ken Miller (pgm, ppm) Richard Mitton (16-bit PSD) - github:urraka (animated gif) Junggon Kim (PNM comments) - Christopher Forseth (animated gif) Daniel Gibson (16-bit TGA) - socks-the-fox (16-bit PNG) - Jeremy Sawicki (handle all ImageNet JPGs) - Optimizations & bugfixes Mikhail Morozov (1-bit BMP) - Fabian "ryg" Giesen Anael Seghezzi (is-16-bit query) - Arseny Kapoulkine Simon Breuss (16-bit PNM) - John-Mark Allen - Carmelo J Fdez-Aguera - - Bug & warning fixes - Marc LeBlanc David Woo Guillaume George Martins Mozeiko - Christpher Lloyd Jerry Jansson Joseph Thomson Blazej Dariusz Roszkowski - Phil Jordan Dave Moore Roy Eltham - Hayaki Saito Nathan Reed Won Chun - Luke Graham Johan Duparc Nick Verigakis the Horde3D community - Thomas Ruf Ronny Chevalier github:rlyeh - Janez Zemva John Bartholomew Michal Cichon github:romigrou - Jonathan Blow Ken Hamada Tero Hanninen github:svdijk - Eugene Golushkov Laurent Gomila Cort Stratton github:snagar - Aruelien Pocheville Sergio Gonzalez Thibault Reuille github:Zelex - Cass Everitt Ryamond Barbiero github:grim210 - Paul Du Bois Engin Manap Aldo Culquicondor github:sammyhw - Philipp Wiesemann Dale Weiler Oriol Ferrer Mesia github:phprus - Josh Tobin Neil Bickford Matthew Gregan github:poppolopoppo - Julian Raschke Gregory Mullen Christian Floisand github:darealshinji - Baldur Karlsson Kevin Schmidt JR Smith github:Michaelangel007 - Brad Weinberger Matvey Cherevko github:mosra - Luca Sas Alexander Veselov Zack Middleton [reserved] - Ryan C. Gordon [reserved] [reserved] - DO NOT ADD YOUR NAME HERE - - Jacko Dirks - - To add your name to the credits, pick a random blank space in the middle and fill it. - 80% of merge conflicts on stb PRs are due to people adding their name at the end - of the credits. -*/ - -#ifndef STBI_INCLUDE_STB_IMAGE_H -#define STBI_INCLUDE_STB_IMAGE_H - -// DOCUMENTATION -// -// Limitations: -// - no 12-bit-per-channel JPEG -// - no JPEGs with arithmetic coding -// - GIF always returns *comp=4 -// -// Basic usage (see HDR discussion below for HDR usage): -// int x,y,n; -// unsigned char *data = stbi_load(filename, &x, &y, &n, 0); -// // ... process data if not NULL ... -// // ... x = width, y = height, n = # 8-bit components per pixel ... -// // ... replace '0' with '1'..'4' to force that many components per pixel -// // ... but 'n' will always be the number that it would have been if you said 0 -// stbi_image_free(data); -// -// Standard parameters: -// int *x -- outputs image width in pixels -// int *y -- outputs image height in pixels -// int *channels_in_file -- outputs # of image components in image file -// int desired_channels -- if non-zero, # of image components requested in result -// -// The return value from an image loader is an 'unsigned char *' which points -// to the pixel data, or NULL on an allocation failure or if the image is -// corrupt or invalid. The pixel data consists of *y scanlines of *x pixels, -// with each pixel consisting of N interleaved 8-bit components; the first -// pixel pointed to is top-left-most in the image. There is no padding between -// image scanlines or between pixels, regardless of format. The number of -// components N is 'desired_channels' if desired_channels is non-zero, or -// *channels_in_file otherwise. If desired_channels is non-zero, -// *channels_in_file has the number of components that _would_ have been -// output otherwise. E.g. if you set desired_channels to 4, you will always -// get RGBA output, but you can check *channels_in_file to see if it's trivially -// opaque because e.g. there were only 3 channels in the source image. -// -// An output image with N components has the following components interleaved -// in this order in each pixel: -// -// N=#comp components -// 1 grey -// 2 grey, alpha -// 3 red, green, blue -// 4 red, green, blue, alpha -// -// If image loading fails for any reason, the return value will be NULL, -// and *x, *y, *channels_in_file will be unchanged. The function -// stbi_failure_reason() can be queried for an extremely brief, end-user -// unfriendly explanation of why the load failed. Define STBI_NO_FAILURE_STRINGS -// to avoid compiling these strings at all, and STBI_FAILURE_USERMSG to get slightly -// more user-friendly ones. -// -// Paletted PNG, BMP, GIF, and PIC images are automatically depalettized. -// -// To query the width, height and component count of an image without having to -// decode the full file, you can use the stbi_info family of functions: -// -// int x,y,n,ok; -// ok = stbi_info(filename, &x, &y, &n); -// // returns ok=1 and sets x, y, n if image is a supported format, -// // 0 otherwise. -// -// Note that stb_image pervasively uses ints in its public API for sizes, -// including sizes of memory buffers. This is now part of the API and thus -// hard to change without causing breakage. As a result, the various image -// loaders all have certain limits on image size; these differ somewhat -// by format but generally boil down to either just under 2GB or just under -// 1GB. When the decoded image would be larger than this, stb_image decoding -// will fail. -// -// Additionally, stb_image will reject image files that have any of their -// dimensions set to a larger value than the configurable STBI_MAX_DIMENSIONS, -// which defaults to 2**24 = 16777216 pixels. Due to the above memory limit, -// the only way to have an image with such dimensions load correctly -// is for it to have a rather extreme aspect ratio. Either way, the -// assumption here is that such larger images are likely to be malformed -// or malicious. If you do need to load an image with individual dimensions -// larger than that, and it still fits in the overall size limit, you can -// #define STBI_MAX_DIMENSIONS on your own to be something larger. -// -// =========================================================================== -// -// UNICODE: -// -// If compiling for Windows and you wish to use Unicode filenames, compile -// with -// #define STBI_WINDOWS_UTF8 -// and pass utf8-encoded filenames. Call stbi_convert_wchar_to_utf8 to convert -// Windows wchar_t filenames to utf8. -// -// =========================================================================== -// -// Philosophy -// -// stb libraries are designed with the following priorities: -// -// 1. easy to use -// 2. easy to maintain -// 3. good performance -// -// Sometimes I let "good performance" creep up in priority over "easy to maintain", -// and for best performance I may provide less-easy-to-use APIs that give higher -// performance, in addition to the easy-to-use ones. Nevertheless, it's important -// to keep in mind that from the standpoint of you, a client of this library, -// all you care about is #1 and #3, and stb libraries DO NOT emphasize #3 above all. -// -// Some secondary priorities arise directly from the first two, some of which -// provide more explicit reasons why performance can't be emphasized. -// -// - Portable ("ease of use") -// - Small source code footprint ("easy to maintain") -// - No dependencies ("ease of use") -// -// =========================================================================== -// -// I/O callbacks -// -// I/O callbacks allow you to read from arbitrary sources, like packaged -// files or some other source. Data read from callbacks are processed -// through a small internal buffer (currently 128 bytes) to try to reduce -// overhead. -// -// The three functions you must define are "read" (reads some bytes of data), -// "skip" (skips some bytes of data), "eof" (reports if the stream is at the end). -// -// =========================================================================== -// -// SIMD support -// -// The JPEG decoder will try to automatically use SIMD kernels on x86 when -// supported by the compiler. For ARM Neon support, you must explicitly -// request it. -// -// (The old do-it-yourself SIMD API is no longer supported in the current -// code.) -// -// On x86, SSE2 will automatically be used when available based on a run-time -// test; if not, the generic C versions are used as a fall-back. On ARM targets, -// the typical path is to have separate builds for NEON and non-NEON devices -// (at least this is true for iOS and Android). Therefore, the NEON support is -// toggled by a build flag: define STBI_NEON to get NEON loops. -// -// If for some reason you do not want to use any of SIMD code, or if -// you have issues compiling it, you can disable it entirely by -// defining STBI_NO_SIMD. -// -// =========================================================================== -// -// HDR image support (disable by defining STBI_NO_HDR) -// -// stb_image supports loading HDR images in general, and currently the Radiance -// .HDR file format specifically. You can still load any file through the existing -// interface; if you attempt to load an HDR file, it will be automatically remapped -// to LDR, assuming gamma 2.2 and an arbitrary scale factor defaulting to 1; -// both of these constants can be reconfigured through this interface: -// -// stbi_hdr_to_ldr_gamma(2.2f); -// stbi_hdr_to_ldr_scale(1.0f); -// -// (note, do not use _inverse_ constants; stbi_image will invert them -// appropriately). -// -// Additionally, there is a new, parallel interface for loading files as -// (linear) floats to preserve the full dynamic range: -// -// float *data = stbi_loadf(filename, &x, &y, &n, 0); -// -// If you load LDR images through this interface, those images will -// be promoted to floating point values, run through the inverse of -// constants corresponding to the above: -// -// stbi_ldr_to_hdr_scale(1.0f); -// stbi_ldr_to_hdr_gamma(2.2f); -// -// Finally, given a filename (or an open file or memory block--see header -// file for details) containing image data, you can query for the "most -// appropriate" interface to use (that is, whether the image is HDR or -// not), using: -// -// stbi_is_hdr(char *filename); -// -// =========================================================================== -// -// iPhone PNG support: -// -// We optionally support converting iPhone-formatted PNGs (which store -// premultiplied BGRA) back to RGB, even though they're internally encoded -// differently. To enable this conversion, call -// stbi_convert_iphone_png_to_rgb(1). -// -// Call stbi_set_unpremultiply_on_load(1) as well to force a divide per -// pixel to remove any premultiplied alpha *only* if the image file explicitly -// says there's premultiplied data (currently only happens in iPhone images, -// and only if iPhone convert-to-rgb processing is on). -// -// =========================================================================== -// -// ADDITIONAL CONFIGURATION -// -// - You can suppress implementation of any of the decoders to reduce -// your code footprint by #defining one or more of the following -// symbols before creating the implementation. -// -// STBI_NO_JPEG -// STBI_NO_PNG -// STBI_NO_BMP -// STBI_NO_PSD -// STBI_NO_TGA -// STBI_NO_GIF -// STBI_NO_HDR -// STBI_NO_PIC -// STBI_NO_PNM (.ppm and .pgm) -// -// - You can request *only* certain decoders and suppress all other ones -// (this will be more forward-compatible, as addition of new decoders -// doesn't require you to disable them explicitly): -// -// STBI_ONLY_JPEG -// STBI_ONLY_PNG -// STBI_ONLY_BMP -// STBI_ONLY_PSD -// STBI_ONLY_TGA -// STBI_ONLY_GIF -// STBI_ONLY_HDR -// STBI_ONLY_PIC -// STBI_ONLY_PNM (.ppm and .pgm) -// -// - If you use STBI_NO_PNG (or _ONLY_ without PNG), and you still -// want the zlib decoder to be available, #define STBI_SUPPORT_ZLIB -// -// - If you define STBI_MAX_DIMENSIONS, stb_image will reject images greater -// than that size (in either width or height) without further processing. -// This is to let programs in the wild set an upper bound to prevent -// denial-of-service attacks on untrusted data, as one could generate a -// valid image of gigantic dimensions and force stb_image to allocate a -// huge block of memory and spend disproportionate time decoding it. By -// default this is set to (1 << 24), which is 16777216, but that's still -// very big. - -#ifndef STBI_NO_STDIO -#include -#endif // STBI_NO_STDIO - -#define STBI_VERSION 1 - -enum -{ - STBI_default = 0, // only used for desired_channels - - STBI_grey = 1, - STBI_grey_alpha = 2, - STBI_rgb = 3, - STBI_rgb_alpha = 4 -}; - -#include -typedef unsigned char stbi_uc; -typedef unsigned short stbi_us; - -#ifdef __cplusplus -extern "C" { -#endif - -#ifndef STBIDEF -#ifdef STB_IMAGE_STATIC -#define STBIDEF static -#else -#define STBIDEF extern -#endif -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// PRIMARY API - works on images of any type -// - -// -// load image by filename, open file, or memory buffer -// - -typedef struct -{ - int (*read) (void *user,char *data,int size); // fill 'data' with 'size' bytes. return number of bytes actually read - void (*skip) (void *user,int n); // skip the next 'n' bytes, or 'unget' the last -n bytes if negative - int (*eof) (void *user); // returns nonzero if we are at end of file/data -} stbi_io_callbacks; - -//////////////////////////////////// -// -// 8-bits-per-channel interface -// - -STBIDEF stbi_uc *stbi_load_from_memory (stbi_uc const *buffer, int len , int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk , void *user, int *x, int *y, int *channels_in_file, int desired_channels); - -#ifndef STBI_NO_STDIO -STBIDEF stbi_uc *stbi_load (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_uc *stbi_load_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); -// for stbi_load_from_file, file pointer is left pointing immediately after image -#endif - -#ifndef STBI_NO_GIF -STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp); -#endif - -#ifdef STBI_WINDOWS_UTF8 -STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input); -#endif - -//////////////////////////////////// -// -// 16-bits-per-channel interface -// - -STBIDEF stbi_us *stbi_load_16_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); - -#ifndef STBI_NO_STDIO -STBIDEF stbi_us *stbi_load_16 (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_us *stbi_load_from_file_16(FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); -#endif - -//////////////////////////////////// -// -// float-per-channel interface -// -#ifndef STBI_NO_LINEAR - STBIDEF float *stbi_loadf_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); - STBIDEF float *stbi_loadf_from_callbacks (stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); - - #ifndef STBI_NO_STDIO - STBIDEF float *stbi_loadf (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); - STBIDEF float *stbi_loadf_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); - #endif -#endif - -#ifndef STBI_NO_HDR - STBIDEF void stbi_hdr_to_ldr_gamma(float gamma); - STBIDEF void stbi_hdr_to_ldr_scale(float scale); -#endif // STBI_NO_HDR - -#ifndef STBI_NO_LINEAR - STBIDEF void stbi_ldr_to_hdr_gamma(float gamma); - STBIDEF void stbi_ldr_to_hdr_scale(float scale); -#endif // STBI_NO_LINEAR - -// stbi_is_hdr is always defined, but always returns false if STBI_NO_HDR -STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user); -STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len); -#ifndef STBI_NO_STDIO -STBIDEF int stbi_is_hdr (char const *filename); -STBIDEF int stbi_is_hdr_from_file(FILE *f); -#endif // STBI_NO_STDIO - - -// get a VERY brief reason for failure -// on most compilers (and ALL modern mainstream compilers) this is threadsafe -STBIDEF const char *stbi_failure_reason (void); - -// free the loaded image -- this is just free() -STBIDEF void stbi_image_free (void *retval_from_stbi_load); - -// get image dimensions & components without fully decoding -STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp); -STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp); -STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len); -STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *clbk, void *user); - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_info (char const *filename, int *x, int *y, int *comp); -STBIDEF int stbi_info_from_file (FILE *f, int *x, int *y, int *comp); -STBIDEF int stbi_is_16_bit (char const *filename); -STBIDEF int stbi_is_16_bit_from_file(FILE *f); -#endif - - - -// for image formats that explicitly notate that they have premultiplied alpha, -// we just return the colors as stored in the file. set this flag to force -// unpremultiplication. results are undefined if the unpremultiply overflow. -STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply); - -// indicate whether we should process iphone images back to canonical format, -// or just pass them through "as-is" -STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert); - -// flip the image vertically, so the first pixel in the output array is the bottom left -STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip); - -// as above, but only applies to images loaded on the thread that calls the function -// this function is only available if your compiler supports thread-local variables; -// calling it will fail to link if your compiler doesn't -STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply); -STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert); -STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip); - -// ZLIB client - used by PNG, available for other purposes - -STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen); -STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header); -STBIDEF char *stbi_zlib_decode_malloc(const char *buffer, int len, int *outlen); -STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); - -STBIDEF char *stbi_zlib_decode_noheader_malloc(const char *buffer, int len, int *outlen); -STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); - - -#ifdef __cplusplus -} -#endif - -// -// -//// end header file ///////////////////////////////////////////////////// -#endif // STBI_INCLUDE_STB_IMAGE_H - -#ifdef STB_IMAGE_IMPLEMENTATION - -#if defined(STBI_ONLY_JPEG) || defined(STBI_ONLY_PNG) || defined(STBI_ONLY_BMP) \ - || defined(STBI_ONLY_TGA) || defined(STBI_ONLY_GIF) || defined(STBI_ONLY_PSD) \ - || defined(STBI_ONLY_HDR) || defined(STBI_ONLY_PIC) || defined(STBI_ONLY_PNM) \ - || defined(STBI_ONLY_ZLIB) - #ifndef STBI_ONLY_JPEG - #define STBI_NO_JPEG - #endif - #ifndef STBI_ONLY_PNG - #define STBI_NO_PNG - #endif - #ifndef STBI_ONLY_BMP - #define STBI_NO_BMP - #endif - #ifndef STBI_ONLY_PSD - #define STBI_NO_PSD - #endif - #ifndef STBI_ONLY_TGA - #define STBI_NO_TGA - #endif - #ifndef STBI_ONLY_GIF - #define STBI_NO_GIF - #endif - #ifndef STBI_ONLY_HDR - #define STBI_NO_HDR - #endif - #ifndef STBI_ONLY_PIC - #define STBI_NO_PIC - #endif - #ifndef STBI_ONLY_PNM - #define STBI_NO_PNM - #endif -#endif - -#if defined(STBI_NO_PNG) && !defined(STBI_SUPPORT_ZLIB) && !defined(STBI_NO_ZLIB) -#define STBI_NO_ZLIB -#endif - - -#include -#include // ptrdiff_t on osx -#include -#include -#include - -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) -#include // ldexp, pow -#endif - -#ifndef STBI_NO_STDIO -#include -#endif - -#ifndef STBI_ASSERT -#include -#define STBI_ASSERT(x) assert(x) -#endif - -#ifdef __cplusplus -#define STBI_EXTERN extern "C" -#else -#define STBI_EXTERN extern -#endif - - -#ifndef _MSC_VER - #ifdef __cplusplus - #define stbi_inline inline - #else - #define stbi_inline - #endif -#else - #define stbi_inline __forceinline -#endif - -#ifndef STBI_NO_THREAD_LOCALS - #if defined(__cplusplus) && __cplusplus >= 201103L - #define STBI_THREAD_LOCAL thread_local - #elif defined(__GNUC__) && __GNUC__ < 5 - #define STBI_THREAD_LOCAL __thread - #elif defined(_MSC_VER) - #define STBI_THREAD_LOCAL __declspec(thread) - #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_THREADS__) - #define STBI_THREAD_LOCAL _Thread_local - #endif - - #ifndef STBI_THREAD_LOCAL - #if defined(__GNUC__) - #define STBI_THREAD_LOCAL __thread - #endif - #endif -#endif - -#if defined(_MSC_VER) || defined(__SYMBIAN32__) -typedef unsigned short stbi__uint16; -typedef signed short stbi__int16; -typedef unsigned int stbi__uint32; -typedef signed int stbi__int32; -#else -#include -typedef uint16_t stbi__uint16; -typedef int16_t stbi__int16; -typedef uint32_t stbi__uint32; -typedef int32_t stbi__int32; -#endif - -// should produce compiler error if size is wrong -typedef unsigned char validate_uint32[sizeof(stbi__uint32)==4 ? 1 : -1]; - -#ifdef _MSC_VER -#define STBI_NOTUSED(v) (void)(v) -#else -#define STBI_NOTUSED(v) (void)sizeof(v) -#endif - -#ifdef _MSC_VER -#define STBI_HAS_LROTL -#endif - -#ifdef STBI_HAS_LROTL - #define stbi_lrot(x,y) _lrotl(x,y) -#else - #define stbi_lrot(x,y) (((x) << (y)) | ((x) >> (-(y) & 31))) -#endif - -#if defined(STBI_MALLOC) && defined(STBI_FREE) && (defined(STBI_REALLOC) || defined(STBI_REALLOC_SIZED)) -// ok -#elif !defined(STBI_MALLOC) && !defined(STBI_FREE) && !defined(STBI_REALLOC) && !defined(STBI_REALLOC_SIZED) -// ok -#else -#error "Must define all or none of STBI_MALLOC, STBI_FREE, and STBI_REALLOC (or STBI_REALLOC_SIZED)." -#endif - -#ifndef STBI_MALLOC -#define STBI_MALLOC(sz) malloc(sz) -#define STBI_REALLOC(p,newsz) realloc(p,newsz) -#define STBI_FREE(p) free(p) -#endif - -#ifndef STBI_REALLOC_SIZED -#define STBI_REALLOC_SIZED(p,oldsz,newsz) STBI_REALLOC(p,newsz) -#endif - -// x86/x64 detection -#if defined(__x86_64__) || defined(_M_X64) -#define STBI__X64_TARGET -#elif defined(__i386) || defined(_M_IX86) -#define STBI__X86_TARGET -#endif - -#if defined(__GNUC__) && defined(STBI__X86_TARGET) && !defined(__SSE2__) && !defined(STBI_NO_SIMD) -// gcc doesn't support sse2 intrinsics unless you compile with -msse2, -// which in turn means it gets to use SSE2 everywhere. This is unfortunate, -// but previous attempts to provide the SSE2 functions with runtime -// detection caused numerous issues. The way architecture extensions are -// exposed in GCC/Clang is, sadly, not really suited for one-file libs. -// New behavior: if compiled with -msse2, we use SSE2 without any -// detection; if not, we don't use it at all. -#define STBI_NO_SIMD -#endif - -#if defined(__MINGW32__) && defined(STBI__X86_TARGET) && !defined(STBI_MINGW_ENABLE_SSE2) && !defined(STBI_NO_SIMD) -// Note that __MINGW32__ doesn't actually mean 32-bit, so we have to avoid STBI__X64_TARGET -// -// 32-bit MinGW wants ESP to be 16-byte aligned, but this is not in the -// Windows ABI and VC++ as well as Windows DLLs don't maintain that invariant. -// As a result, enabling SSE2 on 32-bit MinGW is dangerous when not -// simultaneously enabling "-mstackrealign". -// -// See https://github.com/nothings/stb/issues/81 for more information. -// -// So default to no SSE2 on 32-bit MinGW. If you've read this far and added -// -mstackrealign to your build settings, feel free to #define STBI_MINGW_ENABLE_SSE2. -#define STBI_NO_SIMD -#endif - -#if !defined(STBI_NO_SIMD) && (defined(STBI__X86_TARGET) || defined(STBI__X64_TARGET)) -#define STBI_SSE2 -#include - -#ifdef _MSC_VER - -#if _MSC_VER >= 1400 // not VC6 -#include // __cpuid -static int stbi__cpuid3(void) -{ - int info[4]; - __cpuid(info,1); - return info[3]; -} -#else -static int stbi__cpuid3(void) -{ - int res; - __asm { - mov eax,1 - cpuid - mov res,edx - } - return res; -} -#endif - -#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name - -#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) -static int stbi__sse2_available(void) -{ - int info3 = stbi__cpuid3(); - return ((info3 >> 26) & 1) != 0; -} -#endif - -#else // assume GCC-style if not VC++ -#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) - -#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) -static int stbi__sse2_available(void) -{ - // If we're even attempting to compile this on GCC/Clang, that means - // -msse2 is on, which means the compiler is allowed to use SSE2 - // instructions at will, and so are we. - return 1; -} -#endif - -#endif -#endif - -// ARM NEON -#if defined(STBI_NO_SIMD) && defined(STBI_NEON) -#undef STBI_NEON -#endif - -#ifdef STBI_NEON -#include -#ifdef _MSC_VER -#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name -#else -#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) -#endif -#endif - -#ifndef STBI_SIMD_ALIGN -#define STBI_SIMD_ALIGN(type, name) type name -#endif - -#ifndef STBI_MAX_DIMENSIONS -#define STBI_MAX_DIMENSIONS (1 << 24) -#endif - -/////////////////////////////////////////////// -// -// stbi__context struct and start_xxx functions - -// stbi__context structure is our basic context used by all images, so it -// contains all the IO context, plus some basic image information -typedef struct -{ - stbi__uint32 img_x, img_y; - int img_n, img_out_n; - - stbi_io_callbacks io; - void *io_user_data; - - int read_from_callbacks; - int buflen; - stbi_uc buffer_start[128]; - int callback_already_read; - - stbi_uc *img_buffer, *img_buffer_end; - stbi_uc *img_buffer_original, *img_buffer_original_end; -} stbi__context; - - -static void stbi__refill_buffer(stbi__context *s); - -// initialize a memory-decode context -static void stbi__start_mem(stbi__context *s, stbi_uc const *buffer, int len) -{ - s->io.read = NULL; - s->read_from_callbacks = 0; - s->callback_already_read = 0; - s->img_buffer = s->img_buffer_original = (stbi_uc *) buffer; - s->img_buffer_end = s->img_buffer_original_end = (stbi_uc *) buffer+len; -} - -// initialize a callback-based context -static void stbi__start_callbacks(stbi__context *s, stbi_io_callbacks *c, void *user) -{ - s->io = *c; - s->io_user_data = user; - s->buflen = sizeof(s->buffer_start); - s->read_from_callbacks = 1; - s->callback_already_read = 0; - s->img_buffer = s->img_buffer_original = s->buffer_start; - stbi__refill_buffer(s); - s->img_buffer_original_end = s->img_buffer_end; -} - -#ifndef STBI_NO_STDIO - -static int stbi__stdio_read(void *user, char *data, int size) -{ - return (int) fread(data,1,size,(FILE*) user); -} - -static void stbi__stdio_skip(void *user, int n) -{ - int ch; - fseek((FILE*) user, n, SEEK_CUR); - ch = fgetc((FILE*) user); /* have to read a byte to reset feof()'s flag */ - if (ch != EOF) { - ungetc(ch, (FILE *) user); /* push byte back onto stream if valid. */ - } -} - -static int stbi__stdio_eof(void *user) -{ - return feof((FILE*) user) || ferror((FILE *) user); -} - -static stbi_io_callbacks stbi__stdio_callbacks = -{ - stbi__stdio_read, - stbi__stdio_skip, - stbi__stdio_eof, -}; - -static void stbi__start_file(stbi__context *s, FILE *f) -{ - stbi__start_callbacks(s, &stbi__stdio_callbacks, (void *) f); -} - -//static void stop_file(stbi__context *s) { } - -#endif // !STBI_NO_STDIO - -static void stbi__rewind(stbi__context *s) -{ - // conceptually rewind SHOULD rewind to the beginning of the stream, - // but we just rewind to the beginning of the initial buffer, because - // we only use it after doing 'test', which only ever looks at at most 92 bytes - s->img_buffer = s->img_buffer_original; - s->img_buffer_end = s->img_buffer_original_end; -} - -enum -{ - STBI_ORDER_RGB, - STBI_ORDER_BGR -}; - -typedef struct -{ - int bits_per_channel; - int num_channels; - int channel_order; -} stbi__result_info; - -#ifndef STBI_NO_JPEG -static int stbi__jpeg_test(stbi__context *s); -static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PNG -static int stbi__png_test(stbi__context *s); -static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__png_is16(stbi__context *s); -#endif - -#ifndef STBI_NO_BMP -static int stbi__bmp_test(stbi__context *s); -static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_TGA -static int stbi__tga_test(stbi__context *s); -static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PSD -static int stbi__psd_test(stbi__context *s); -static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc); -static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__psd_is16(stbi__context *s); -#endif - -#ifndef STBI_NO_HDR -static int stbi__hdr_test(stbi__context *s); -static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PIC -static int stbi__pic_test(stbi__context *s); -static void *stbi__pic_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_GIF -static int stbi__gif_test(stbi__context *s); -static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp); -static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PNM -static int stbi__pnm_test(stbi__context *s); -static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__pnm_is16(stbi__context *s); -#endif - -static -#ifdef STBI_THREAD_LOCAL -STBI_THREAD_LOCAL -#endif -const char *stbi__g_failure_reason; - -STBIDEF const char *stbi_failure_reason(void) -{ - return stbi__g_failure_reason; -} - -#ifndef STBI_NO_FAILURE_STRINGS -static int stbi__err(const char *str) -{ - stbi__g_failure_reason = str; - return 0; -} -#endif - -static void *stbi__malloc(size_t size) -{ - return STBI_MALLOC(size); -} - -// stb_image uses ints pervasively, including for offset calculations. -// therefore the largest decoded image size we can support with the -// current code, even on 64-bit targets, is INT_MAX. this is not a -// significant limitation for the intended use case. -// -// we do, however, need to make sure our size calculations don't -// overflow. hence a few helper functions for size calculations that -// multiply integers together, making sure that they're non-negative -// and no overflow occurs. - -// return 1 if the sum is valid, 0 on overflow. -// negative terms are considered invalid. -static int stbi__addsizes_valid(int a, int b) -{ - if (b < 0) return 0; - // now 0 <= b <= INT_MAX, hence also - // 0 <= INT_MAX - b <= INTMAX. - // And "a + b <= INT_MAX" (which might overflow) is the - // same as a <= INT_MAX - b (no overflow) - return a <= INT_MAX - b; -} - -// returns 1 if the product is valid, 0 on overflow. -// negative factors are considered invalid. -static int stbi__mul2sizes_valid(int a, int b) -{ - if (a < 0 || b < 0) return 0; - if (b == 0) return 1; // mul-by-0 is always safe - // portable way to check for no overflows in a*b - return a <= INT_MAX/b; -} - -#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) -// returns 1 if "a*b + add" has no negative terms/factors and doesn't overflow -static int stbi__mad2sizes_valid(int a, int b, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__addsizes_valid(a*b, add); -} -#endif - -// returns 1 if "a*b*c + add" has no negative terms/factors and doesn't overflow -static int stbi__mad3sizes_valid(int a, int b, int c, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && - stbi__addsizes_valid(a*b*c, add); -} - -// returns 1 if "a*b*c*d + add" has no negative terms/factors and doesn't overflow -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) -static int stbi__mad4sizes_valid(int a, int b, int c, int d, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && - stbi__mul2sizes_valid(a*b*c, d) && stbi__addsizes_valid(a*b*c*d, add); -} -#endif - -#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) -// mallocs with size overflow checking -static void *stbi__malloc_mad2(int a, int b, int add) -{ - if (!stbi__mad2sizes_valid(a, b, add)) return NULL; - return stbi__malloc(a*b + add); -} -#endif - -static void *stbi__malloc_mad3(int a, int b, int c, int add) -{ - if (!stbi__mad3sizes_valid(a, b, c, add)) return NULL; - return stbi__malloc(a*b*c + add); -} - -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) -static void *stbi__malloc_mad4(int a, int b, int c, int d, int add) -{ - if (!stbi__mad4sizes_valid(a, b, c, d, add)) return NULL; - return stbi__malloc(a*b*c*d + add); -} -#endif - -// returns 1 if the sum of two signed ints is valid (between -2^31 and 2^31-1 inclusive), 0 on overflow. -static int stbi__addints_valid(int a, int b) -{ - if ((a >= 0) != (b >= 0)) return 1; // a and b have different signs, so no overflow - if (a < 0 && b < 0) return a >= INT_MIN - b; // same as a + b >= INT_MIN; INT_MIN - b cannot overflow since b < 0. - return a <= INT_MAX - b; -} - -// returns 1 if the product of two ints fits in a signed short, 0 on overflow. -static int stbi__mul2shorts_valid(int a, int b) -{ - if (b == 0 || b == -1) return 1; // multiplication by 0 is always 0; check for -1 so SHRT_MIN/b doesn't overflow - if ((a >= 0) == (b >= 0)) return a <= SHRT_MAX/b; // product is positive, so similar to mul2sizes_valid - if (b < 0) return a <= SHRT_MIN / b; // same as a * b >= SHRT_MIN - return a >= SHRT_MIN / b; -} - -// stbi__err - error -// stbi__errpf - error returning pointer to float -// stbi__errpuc - error returning pointer to unsigned char - -#ifdef STBI_NO_FAILURE_STRINGS - #define stbi__err(x,y) 0 -#elif defined(STBI_FAILURE_USERMSG) - #define stbi__err(x,y) stbi__err(y) -#else - #define stbi__err(x,y) stbi__err(x) -#endif - -#define stbi__errpf(x,y) ((float *)(size_t) (stbi__err(x,y)?NULL:NULL)) -#define stbi__errpuc(x,y) ((unsigned char *)(size_t) (stbi__err(x,y)?NULL:NULL)) - -STBIDEF void stbi_image_free(void *retval_from_stbi_load) -{ - STBI_FREE(retval_from_stbi_load); -} - -#ifndef STBI_NO_LINEAR -static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp); -#endif - -#ifndef STBI_NO_HDR -static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp); -#endif - -static int stbi__vertically_flip_on_load_global = 0; - -STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip) -{ - stbi__vertically_flip_on_load_global = flag_true_if_should_flip; -} - -#ifndef STBI_THREAD_LOCAL -#define stbi__vertically_flip_on_load stbi__vertically_flip_on_load_global -#else -static STBI_THREAD_LOCAL int stbi__vertically_flip_on_load_local, stbi__vertically_flip_on_load_set; - -STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip) -{ - stbi__vertically_flip_on_load_local = flag_true_if_should_flip; - stbi__vertically_flip_on_load_set = 1; -} - -#define stbi__vertically_flip_on_load (stbi__vertically_flip_on_load_set \ - ? stbi__vertically_flip_on_load_local \ - : stbi__vertically_flip_on_load_global) -#endif // STBI_THREAD_LOCAL - -static void *stbi__load_main(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) -{ - memset(ri, 0, sizeof(*ri)); // make sure it's initialized if we add new fields - ri->bits_per_channel = 8; // default is 8 so most paths don't have to be changed - ri->channel_order = STBI_ORDER_RGB; // all current input & output are this, but this is here so we can add BGR order - ri->num_channels = 0; - - // test the formats with a very explicit header first (at least a FOURCC - // or distinctive magic number first) - #ifndef STBI_NO_PNG - if (stbi__png_test(s)) return stbi__png_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_BMP - if (stbi__bmp_test(s)) return stbi__bmp_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_GIF - if (stbi__gif_test(s)) return stbi__gif_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_PSD - if (stbi__psd_test(s)) return stbi__psd_load(s,x,y,comp,req_comp, ri, bpc); - #else - STBI_NOTUSED(bpc); - #endif - #ifndef STBI_NO_PIC - if (stbi__pic_test(s)) return stbi__pic_load(s,x,y,comp,req_comp, ri); - #endif - - // then the formats that can end up attempting to load with just 1 or 2 - // bytes matching expectations; these are prone to false positives, so - // try them later - #ifndef STBI_NO_JPEG - if (stbi__jpeg_test(s)) return stbi__jpeg_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_PNM - if (stbi__pnm_test(s)) return stbi__pnm_load(s,x,y,comp,req_comp, ri); - #endif - - #ifndef STBI_NO_HDR - if (stbi__hdr_test(s)) { - float *hdr = stbi__hdr_load(s, x,y,comp,req_comp, ri); - return stbi__hdr_to_ldr(hdr, *x, *y, req_comp ? req_comp : *comp); - } - #endif - - #ifndef STBI_NO_TGA - // test tga last because it's a crappy test! - if (stbi__tga_test(s)) - return stbi__tga_load(s,x,y,comp,req_comp, ri); - #endif - - return stbi__errpuc("unknown image type", "Image not of any known type, or corrupt"); -} - -static stbi_uc *stbi__convert_16_to_8(stbi__uint16 *orig, int w, int h, int channels) -{ - int i; - int img_len = w * h * channels; - stbi_uc *reduced; - - reduced = (stbi_uc *) stbi__malloc(img_len); - if (reduced == NULL) return stbi__errpuc("outofmem", "Out of memory"); - - for (i = 0; i < img_len; ++i) - reduced[i] = (stbi_uc)((orig[i] >> 8) & 0xFF); // top half of each byte is sufficient approx of 16->8 bit scaling - - STBI_FREE(orig); - return reduced; -} - -static stbi__uint16 *stbi__convert_8_to_16(stbi_uc *orig, int w, int h, int channels) -{ - int i; - int img_len = w * h * channels; - stbi__uint16 *enlarged; - - enlarged = (stbi__uint16 *) stbi__malloc(img_len*2); - if (enlarged == NULL) return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); - - for (i = 0; i < img_len; ++i) - enlarged[i] = (stbi__uint16)((orig[i] << 8) + orig[i]); // replicate to high and low byte, maps 0->0, 255->0xffff - - STBI_FREE(orig); - return enlarged; -} - -static void stbi__vertical_flip(void *image, int w, int h, int bytes_per_pixel) -{ - int row; - size_t bytes_per_row = (size_t)w * bytes_per_pixel; - stbi_uc temp[2048]; - stbi_uc *bytes = (stbi_uc *)image; - - for (row = 0; row < (h>>1); row++) { - stbi_uc *row0 = bytes + row*bytes_per_row; - stbi_uc *row1 = bytes + (h - row - 1)*bytes_per_row; - // swap row0 with row1 - size_t bytes_left = bytes_per_row; - while (bytes_left) { - size_t bytes_copy = (bytes_left < sizeof(temp)) ? bytes_left : sizeof(temp); - memcpy(temp, row0, bytes_copy); - memcpy(row0, row1, bytes_copy); - memcpy(row1, temp, bytes_copy); - row0 += bytes_copy; - row1 += bytes_copy; - bytes_left -= bytes_copy; - } - } -} - -#ifndef STBI_NO_GIF -static void stbi__vertical_flip_slices(void *image, int w, int h, int z, int bytes_per_pixel) -{ - int slice; - int slice_size = w * h * bytes_per_pixel; - - stbi_uc *bytes = (stbi_uc *)image; - for (slice = 0; slice < z; ++slice) { - stbi__vertical_flip(bytes, w, h, bytes_per_pixel); - bytes += slice_size; - } -} -#endif - -static unsigned char *stbi__load_and_postprocess_8bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - stbi__result_info ri; - void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 8); - - if (result == NULL) - return NULL; - - // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. - STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); - - if (ri.bits_per_channel != 8) { - result = stbi__convert_16_to_8((stbi__uint16 *) result, *x, *y, req_comp == 0 ? *comp : req_comp); - ri.bits_per_channel = 8; - } - - // @TODO: move stbi__convert_format to here - - if (stbi__vertically_flip_on_load) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi_uc)); - } - - return (unsigned char *) result; -} - -static stbi__uint16 *stbi__load_and_postprocess_16bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - stbi__result_info ri; - void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 16); - - if (result == NULL) - return NULL; - - // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. - STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); - - if (ri.bits_per_channel != 16) { - result = stbi__convert_8_to_16((stbi_uc *) result, *x, *y, req_comp == 0 ? *comp : req_comp); - ri.bits_per_channel = 16; - } - - // @TODO: move stbi__convert_format16 to here - // @TODO: special case RGB-to-Y (and RGBA-to-YA) for 8-bit-to-16-bit case to keep more precision - - if (stbi__vertically_flip_on_load) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi__uint16)); - } - - return (stbi__uint16 *) result; -} - -#if !defined(STBI_NO_HDR) && !defined(STBI_NO_LINEAR) -static void stbi__float_postprocess(float *result, int *x, int *y, int *comp, int req_comp) -{ - if (stbi__vertically_flip_on_load && result != NULL) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(float)); - } -} -#endif - -#ifndef STBI_NO_STDIO - -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) -STBI_EXTERN __declspec(dllimport) int __stdcall MultiByteToWideChar(unsigned int cp, unsigned long flags, const char *str, int cbmb, wchar_t *widestr, int cchwide); -STBI_EXTERN __declspec(dllimport) int __stdcall WideCharToMultiByte(unsigned int cp, unsigned long flags, const wchar_t *widestr, int cchwide, char *str, int cbmb, const char *defchar, int *used_default); -#endif - -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) -STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input) -{ - return WideCharToMultiByte(65001 /* UTF8 */, 0, input, -1, buffer, (int) bufferlen, NULL, NULL); -} -#endif - -static FILE *stbi__fopen(char const *filename, char const *mode) -{ - FILE *f; -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) - wchar_t wMode[64]; - wchar_t wFilename[1024]; - if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, filename, -1, wFilename, sizeof(wFilename)/sizeof(*wFilename))) - return 0; - - if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, mode, -1, wMode, sizeof(wMode)/sizeof(*wMode))) - return 0; - -#if defined(_MSC_VER) && _MSC_VER >= 1400 - if (0 != _wfopen_s(&f, wFilename, wMode)) - f = 0; -#else - f = _wfopen(wFilename, wMode); -#endif - -#elif defined(_MSC_VER) && _MSC_VER >= 1400 - if (0 != fopen_s(&f, filename, mode)) - f=0; -#else - f = fopen(filename, mode); -#endif - return f; -} - - -STBIDEF stbi_uc *stbi_load(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - unsigned char *result; - if (!f) return stbi__errpuc("can't fopen", "Unable to open file"); - result = stbi_load_from_file(f,x,y,comp,req_comp); - fclose(f); - return result; -} - -STBIDEF stbi_uc *stbi_load_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - unsigned char *result; - stbi__context s; - stbi__start_file(&s,f); - result = stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); - if (result) { - // need to 'unget' all the characters in the IO buffer - fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); - } - return result; -} - -STBIDEF stbi__uint16 *stbi_load_from_file_16(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - stbi__uint16 *result; - stbi__context s; - stbi__start_file(&s,f); - result = stbi__load_and_postprocess_16bit(&s,x,y,comp,req_comp); - if (result) { - // need to 'unget' all the characters in the IO buffer - fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); - } - return result; -} - -STBIDEF stbi_us *stbi_load_16(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - stbi__uint16 *result; - if (!f) return (stbi_us *) stbi__errpuc("can't fopen", "Unable to open file"); - result = stbi_load_from_file_16(f,x,y,comp,req_comp); - fclose(f); - return result; -} - - -#endif //!STBI_NO_STDIO - -STBIDEF stbi_us *stbi_load_16_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); -} - -STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *)clbk, user); - return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); -} - -STBIDEF stbi_uc *stbi_load_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); -} - -STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); -} - -#ifndef STBI_NO_GIF -STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp) -{ - unsigned char *result; - stbi__context s; - stbi__start_mem(&s,buffer,len); - - result = (unsigned char*) stbi__load_gif_main(&s, delays, x, y, z, comp, req_comp); - if (stbi__vertically_flip_on_load) { - stbi__vertical_flip_slices( result, *x, *y, *z, *comp ); - } - - return result; -} -#endif - -#ifndef STBI_NO_LINEAR -static float *stbi__loadf_main(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - unsigned char *data; - #ifndef STBI_NO_HDR - if (stbi__hdr_test(s)) { - stbi__result_info ri; - float *hdr_data = stbi__hdr_load(s,x,y,comp,req_comp, &ri); - if (hdr_data) - stbi__float_postprocess(hdr_data,x,y,comp,req_comp); - return hdr_data; - } - #endif - data = stbi__load_and_postprocess_8bit(s, x, y, comp, req_comp); - if (data) - return stbi__ldr_to_hdr(data, *x, *y, req_comp ? req_comp : *comp); - return stbi__errpf("unknown image type", "Image not of any known type, or corrupt"); -} - -STBIDEF float *stbi_loadf_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} - -STBIDEF float *stbi_loadf_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} - -#ifndef STBI_NO_STDIO -STBIDEF float *stbi_loadf(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - float *result; - FILE *f = stbi__fopen(filename, "rb"); - if (!f) return stbi__errpf("can't fopen", "Unable to open file"); - result = stbi_loadf_from_file(f,x,y,comp,req_comp); - fclose(f); - return result; -} - -STBIDEF float *stbi_loadf_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_file(&s,f); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} -#endif // !STBI_NO_STDIO - -#endif // !STBI_NO_LINEAR - -// these is-hdr-or-not is defined independent of whether STBI_NO_LINEAR is -// defined, for API simplicity; if STBI_NO_LINEAR is defined, it always -// reports false! - -STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len) -{ - #ifndef STBI_NO_HDR - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__hdr_test(&s); - #else - STBI_NOTUSED(buffer); - STBI_NOTUSED(len); - return 0; - #endif -} - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_is_hdr (char const *filename) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result=0; - if (f) { - result = stbi_is_hdr_from_file(f); - fclose(f); - } - return result; -} - -STBIDEF int stbi_is_hdr_from_file(FILE *f) -{ - #ifndef STBI_NO_HDR - long pos = ftell(f); - int res; - stbi__context s; - stbi__start_file(&s,f); - res = stbi__hdr_test(&s); - fseek(f, pos, SEEK_SET); - return res; - #else - STBI_NOTUSED(f); - return 0; - #endif -} -#endif // !STBI_NO_STDIO - -STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user) -{ - #ifndef STBI_NO_HDR - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__hdr_test(&s); - #else - STBI_NOTUSED(clbk); - STBI_NOTUSED(user); - return 0; - #endif -} - -#ifndef STBI_NO_LINEAR -static float stbi__l2h_gamma=2.2f, stbi__l2h_scale=1.0f; - -STBIDEF void stbi_ldr_to_hdr_gamma(float gamma) { stbi__l2h_gamma = gamma; } -STBIDEF void stbi_ldr_to_hdr_scale(float scale) { stbi__l2h_scale = scale; } -#endif - -static float stbi__h2l_gamma_i=1.0f/2.2f, stbi__h2l_scale_i=1.0f; - -STBIDEF void stbi_hdr_to_ldr_gamma(float gamma) { stbi__h2l_gamma_i = 1/gamma; } -STBIDEF void stbi_hdr_to_ldr_scale(float scale) { stbi__h2l_scale_i = 1/scale; } - - -////////////////////////////////////////////////////////////////////////////// -// -// Common code used by all image loaders -// - -enum -{ - STBI__SCAN_load=0, - STBI__SCAN_type, - STBI__SCAN_header -}; - -static void stbi__refill_buffer(stbi__context *s) -{ - int n = (s->io.read)(s->io_user_data,(char*)s->buffer_start,s->buflen); - s->callback_already_read += (int) (s->img_buffer - s->img_buffer_original); - if (n == 0) { - // at end of file, treat same as if from memory, but need to handle case - // where s->img_buffer isn't pointing to safe memory, e.g. 0-byte file - s->read_from_callbacks = 0; - s->img_buffer = s->buffer_start; - s->img_buffer_end = s->buffer_start+1; - *s->img_buffer = 0; - } else { - s->img_buffer = s->buffer_start; - s->img_buffer_end = s->buffer_start + n; - } -} - -stbi_inline static stbi_uc stbi__get8(stbi__context *s) -{ - if (s->img_buffer < s->img_buffer_end) - return *s->img_buffer++; - if (s->read_from_callbacks) { - stbi__refill_buffer(s); - return *s->img_buffer++; - } - return 0; -} - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_HDR) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -stbi_inline static int stbi__at_eof(stbi__context *s) -{ - if (s->io.read) { - if (!(s->io.eof)(s->io_user_data)) return 0; - // if feof() is true, check if buffer = end - // special case: we've only got the special 0 character at the end - if (s->read_from_callbacks == 0) return 1; - } - - return s->img_buffer >= s->img_buffer_end; -} -#endif - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) -// nothing -#else -static void stbi__skip(stbi__context *s, int n) -{ - if (n == 0) return; // already there! - if (n < 0) { - s->img_buffer = s->img_buffer_end; - return; - } - if (s->io.read) { - int blen = (int) (s->img_buffer_end - s->img_buffer); - if (blen < n) { - s->img_buffer = s->img_buffer_end; - (s->io.skip)(s->io_user_data, n - blen); - return; - } - } - s->img_buffer += n; -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_TGA) && defined(STBI_NO_HDR) && defined(STBI_NO_PNM) -// nothing -#else -static int stbi__getn(stbi__context *s, stbi_uc *buffer, int n) -{ - if (s->io.read) { - int blen = (int) (s->img_buffer_end - s->img_buffer); - if (blen < n) { - int res, count; - - memcpy(buffer, s->img_buffer, blen); - - count = (s->io.read)(s->io_user_data, (char*) buffer + blen, n - blen); - res = (count == (n-blen)); - s->img_buffer = s->img_buffer_end; - return res; - } - } - - if (s->img_buffer+n <= s->img_buffer_end) { - memcpy(buffer, s->img_buffer, n); - s->img_buffer += n; - return 1; - } else - return 0; -} -#endif - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) -// nothing -#else -static int stbi__get16be(stbi__context *s) -{ - int z = stbi__get8(s); - return (z << 8) + stbi__get8(s); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) -// nothing -#else -static stbi__uint32 stbi__get32be(stbi__context *s) -{ - stbi__uint32 z = stbi__get16be(s); - return (z << 16) + stbi__get16be(s); -} -#endif - -#if defined(STBI_NO_BMP) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) -// nothing -#else -static int stbi__get16le(stbi__context *s) -{ - int z = stbi__get8(s); - return z + (stbi__get8(s) << 8); -} -#endif - -#ifndef STBI_NO_BMP -static stbi__uint32 stbi__get32le(stbi__context *s) -{ - stbi__uint32 z = stbi__get16le(s); - z += (stbi__uint32)stbi__get16le(s) << 16; - return z; -} -#endif - -#define STBI__BYTECAST(x) ((stbi_uc) ((x) & 255)) // truncate int to byte without warnings - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -////////////////////////////////////////////////////////////////////////////// -// -// generic converter from built-in img_n to req_comp -// individual types do this automatically as much as possible (e.g. jpeg -// does all cases internally since it needs to colorspace convert anyway, -// and it never has alpha, so very few cases ). png can automatically -// interleave an alpha=255 channel, but falls back to this for other cases -// -// assume data buffer is malloced, so malloc a new one and free that one -// only failure mode is malloc failing - -static stbi_uc stbi__compute_y(int r, int g, int b) -{ - return (stbi_uc) (((r*77) + (g*150) + (29*b)) >> 8); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -static unsigned char *stbi__convert_format(unsigned char *data, int img_n, int req_comp, unsigned int x, unsigned int y) -{ - int i,j; - unsigned char *good; - - if (req_comp == img_n) return data; - STBI_ASSERT(req_comp >= 1 && req_comp <= 4); - - good = (unsigned char *) stbi__malloc_mad3(req_comp, x, y, 0); - if (good == NULL) { - STBI_FREE(data); - return stbi__errpuc("outofmem", "Out of memory"); - } - - for (j=0; j < (int) y; ++j) { - unsigned char *src = data + j * x * img_n ; - unsigned char *dest = good + j * x * req_comp; - - #define STBI__COMBO(a,b) ((a)*8+(b)) - #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) - // convert source image with img_n components to one with req_comp components; - // avoid switch per pixel, so use switch per scanline and massive macros - switch (STBI__COMBO(img_n, req_comp)) { - STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=255; } break; - STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=255; } break; - STBI__CASE(2,1) { dest[0]=src[0]; } break; - STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; - STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=255; } break; - STBI__CASE(3,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; - STBI__CASE(3,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = 255; } break; - STBI__CASE(4,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; - STBI__CASE(4,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = src[3]; } break; - STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; - default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return stbi__errpuc("unsupported", "Unsupported format conversion"); - } - #undef STBI__CASE - } - - STBI_FREE(data); - return good; -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) -// nothing -#else -static stbi__uint16 stbi__compute_y_16(int r, int g, int b) -{ - return (stbi__uint16) (((r*77) + (g*150) + (29*b)) >> 8); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) -// nothing -#else -static stbi__uint16 *stbi__convert_format16(stbi__uint16 *data, int img_n, int req_comp, unsigned int x, unsigned int y) -{ - int i,j; - stbi__uint16 *good; - - if (req_comp == img_n) return data; - STBI_ASSERT(req_comp >= 1 && req_comp <= 4); - - good = (stbi__uint16 *) stbi__malloc(req_comp * x * y * 2); - if (good == NULL) { - STBI_FREE(data); - return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); - } - - for (j=0; j < (int) y; ++j) { - stbi__uint16 *src = data + j * x * img_n ; - stbi__uint16 *dest = good + j * x * req_comp; - - #define STBI__COMBO(a,b) ((a)*8+(b)) - #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) - // convert source image with img_n components to one with req_comp components; - // avoid switch per pixel, so use switch per scanline and massive macros - switch (STBI__COMBO(img_n, req_comp)) { - STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=0xffff; } break; - STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=0xffff; } break; - STBI__CASE(2,1) { dest[0]=src[0]; } break; - STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; - STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=0xffff; } break; - STBI__CASE(3,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; - STBI__CASE(3,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = 0xffff; } break; - STBI__CASE(4,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; - STBI__CASE(4,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = src[3]; } break; - STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; - default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return (stbi__uint16*) stbi__errpuc("unsupported", "Unsupported format conversion"); - } - #undef STBI__CASE - } - - STBI_FREE(data); - return good; -} -#endif - -#ifndef STBI_NO_LINEAR -static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp) -{ - int i,k,n; - float *output; - if (!data) return NULL; - output = (float *) stbi__malloc_mad4(x, y, comp, sizeof(float), 0); - if (output == NULL) { STBI_FREE(data); return stbi__errpf("outofmem", "Out of memory"); } - // compute number of non-alpha components - if (comp & 1) n = comp; else n = comp-1; - for (i=0; i < x*y; ++i) { - for (k=0; k < n; ++k) { - output[i*comp + k] = (float) (pow(data[i*comp+k]/255.0f, stbi__l2h_gamma) * stbi__l2h_scale); - } - } - if (n < comp) { - for (i=0; i < x*y; ++i) { - output[i*comp + n] = data[i*comp + n]/255.0f; - } - } - STBI_FREE(data); - return output; -} -#endif - -#ifndef STBI_NO_HDR -#define stbi__float2int(x) ((int) (x)) -static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp) -{ - int i,k,n; - stbi_uc *output; - if (!data) return NULL; - output = (stbi_uc *) stbi__malloc_mad3(x, y, comp, 0); - if (output == NULL) { STBI_FREE(data); return stbi__errpuc("outofmem", "Out of memory"); } - // compute number of non-alpha components - if (comp & 1) n = comp; else n = comp-1; - for (i=0; i < x*y; ++i) { - for (k=0; k < n; ++k) { - float z = (float) pow(data[i*comp+k]*stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f; - if (z < 0) z = 0; - if (z > 255) z = 255; - output[i*comp + k] = (stbi_uc) stbi__float2int(z); - } - if (k < comp) { - float z = data[i*comp+k] * 255 + 0.5f; - if (z < 0) z = 0; - if (z > 255) z = 255; - output[i*comp + k] = (stbi_uc) stbi__float2int(z); - } - } - STBI_FREE(data); - return output; -} -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// "baseline" JPEG/JFIF decoder -// -// simple implementation -// - doesn't support delayed output of y-dimension -// - simple interface (only one output format: 8-bit interleaved RGB) -// - doesn't try to recover corrupt jpegs -// - doesn't allow partial loading, loading multiple at once -// - still fast on x86 (copying globals into locals doesn't help x86) -// - allocates lots of intermediate memory (full size of all components) -// - non-interleaved case requires this anyway -// - allows good upsampling (see next) -// high-quality -// - upsampled channels are bilinearly interpolated, even across blocks -// - quality integer IDCT derived from IJG's 'slow' -// performance -// - fast huffman; reasonable integer IDCT -// - some SIMD kernels for common paths on targets with SSE2/NEON -// - uses a lot of intermediate memory, could cache poorly - -#ifndef STBI_NO_JPEG - -// huffman decoding acceleration -#define FAST_BITS 9 // larger handles more cases; smaller stomps less cache - -typedef struct -{ - stbi_uc fast[1 << FAST_BITS]; - // weirdly, repacking this into AoS is a 10% speed loss, instead of a win - stbi__uint16 code[256]; - stbi_uc values[256]; - stbi_uc size[257]; - unsigned int maxcode[18]; - int delta[17]; // old 'firstsymbol' - old 'firstcode' -} stbi__huffman; - -typedef struct -{ - stbi__context *s; - stbi__huffman huff_dc[4]; - stbi__huffman huff_ac[4]; - stbi__uint16 dequant[4][64]; - stbi__int16 fast_ac[4][1 << FAST_BITS]; - -// sizes for components, interleaved MCUs - int img_h_max, img_v_max; - int img_mcu_x, img_mcu_y; - int img_mcu_w, img_mcu_h; - -// definition of jpeg image component - struct - { - int id; - int h,v; - int tq; - int hd,ha; - int dc_pred; - - int x,y,w2,h2; - stbi_uc *data; - void *raw_data, *raw_coeff; - stbi_uc *linebuf; - short *coeff; // progressive only - int coeff_w, coeff_h; // number of 8x8 coefficient blocks - } img_comp[4]; - - stbi__uint32 code_buffer; // jpeg entropy-coded buffer - int code_bits; // number of valid bits - unsigned char marker; // marker seen while filling entropy buffer - int nomore; // flag if we saw a marker so must stop - - int progressive; - int spec_start; - int spec_end; - int succ_high; - int succ_low; - int eob_run; - int jfif; - int app14_color_transform; // Adobe APP14 tag - int rgb; - - int scan_n, order[4]; - int restart_interval, todo; - -// kernels - void (*idct_block_kernel)(stbi_uc *out, int out_stride, short data[64]); - void (*YCbCr_to_RGB_kernel)(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step); - stbi_uc *(*resample_row_hv_2_kernel)(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs); -} stbi__jpeg; - -static int stbi__build_huffman(stbi__huffman *h, int *count) -{ - int i,j,k=0; - unsigned int code; - // build size list for each symbol (from JPEG spec) - for (i=0; i < 16; ++i) { - for (j=0; j < count[i]; ++j) { - h->size[k++] = (stbi_uc) (i+1); - if(k >= 257) return stbi__err("bad size list","Corrupt JPEG"); - } - } - h->size[k] = 0; - - // compute actual symbols (from jpeg spec) - code = 0; - k = 0; - for(j=1; j <= 16; ++j) { - // compute delta to add to code to compute symbol id - h->delta[j] = k - code; - if (h->size[k] == j) { - while (h->size[k] == j) - h->code[k++] = (stbi__uint16) (code++); - if (code-1 >= (1u << j)) return stbi__err("bad code lengths","Corrupt JPEG"); - } - // compute largest code + 1 for this size, preshifted as needed later - h->maxcode[j] = code << (16-j); - code <<= 1; - } - h->maxcode[j] = 0xffffffff; - - // build non-spec acceleration table; 255 is flag for not-accelerated - memset(h->fast, 255, 1 << FAST_BITS); - for (i=0; i < k; ++i) { - int s = h->size[i]; - if (s <= FAST_BITS) { - int c = h->code[i] << (FAST_BITS-s); - int m = 1 << (FAST_BITS-s); - for (j=0; j < m; ++j) { - h->fast[c+j] = (stbi_uc) i; - } - } - } - return 1; -} - -// build a table that decodes both magnitude and value of small ACs in -// one go. -static void stbi__build_fast_ac(stbi__int16 *fast_ac, stbi__huffman *h) -{ - int i; - for (i=0; i < (1 << FAST_BITS); ++i) { - stbi_uc fast = h->fast[i]; - fast_ac[i] = 0; - if (fast < 255) { - int rs = h->values[fast]; - int run = (rs >> 4) & 15; - int magbits = rs & 15; - int len = h->size[fast]; - - if (magbits && len + magbits <= FAST_BITS) { - // magnitude code followed by receive_extend code - int k = ((i << len) & ((1 << FAST_BITS) - 1)) >> (FAST_BITS - magbits); - int m = 1 << (magbits - 1); - if (k < m) k += (~0U << magbits) + 1; - // if the result is small enough, we can fit it in fast_ac table - if (k >= -128 && k <= 127) - fast_ac[i] = (stbi__int16) ((k * 256) + (run * 16) + (len + magbits)); - } - } - } -} - -static void stbi__grow_buffer_unsafe(stbi__jpeg *j) -{ - do { - unsigned int b = j->nomore ? 0 : stbi__get8(j->s); - if (b == 0xff) { - int c = stbi__get8(j->s); - while (c == 0xff) c = stbi__get8(j->s); // consume fill bytes - if (c != 0) { - j->marker = (unsigned char) c; - j->nomore = 1; - return; - } - } - j->code_buffer |= b << (24 - j->code_bits); - j->code_bits += 8; - } while (j->code_bits <= 24); -} - -// (1 << n) - 1 -static const stbi__uint32 stbi__bmask[17]={0,1,3,7,15,31,63,127,255,511,1023,2047,4095,8191,16383,32767,65535}; - -// decode a jpeg huffman value from the bitstream -stbi_inline static int stbi__jpeg_huff_decode(stbi__jpeg *j, stbi__huffman *h) -{ - unsigned int temp; - int c,k; - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - - // look at the top FAST_BITS and determine what symbol ID it is, - // if the code is <= FAST_BITS - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - k = h->fast[c]; - if (k < 255) { - int s = h->size[k]; - if (s > j->code_bits) - return -1; - j->code_buffer <<= s; - j->code_bits -= s; - return h->values[k]; - } - - // naive test is to shift the code_buffer down so k bits are - // valid, then test against maxcode. To speed this up, we've - // preshifted maxcode left so that it has (16-k) 0s at the - // end; in other words, regardless of the number of bits, it - // wants to be compared against something shifted to have 16; - // that way we don't need to shift inside the loop. - temp = j->code_buffer >> 16; - for (k=FAST_BITS+1 ; ; ++k) - if (temp < h->maxcode[k]) - break; - if (k == 17) { - // error! code not found - j->code_bits -= 16; - return -1; - } - - if (k > j->code_bits) - return -1; - - // convert the huffman code to the symbol id - c = ((j->code_buffer >> (32 - k)) & stbi__bmask[k]) + h->delta[k]; - if(c < 0 || c >= 256) // symbol id out of bounds! - return -1; - STBI_ASSERT((((j->code_buffer) >> (32 - h->size[c])) & stbi__bmask[h->size[c]]) == h->code[c]); - - // convert the id to a symbol - j->code_bits -= k; - j->code_buffer <<= k; - return h->values[c]; -} - -// bias[n] = (-1<code_bits < n) stbi__grow_buffer_unsafe(j); - if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing - - sgn = j->code_buffer >> 31; // sign bit always in MSB; 0 if MSB clear (positive), 1 if MSB set (negative) - k = stbi_lrot(j->code_buffer, n); - j->code_buffer = k & ~stbi__bmask[n]; - k &= stbi__bmask[n]; - j->code_bits -= n; - return k + (stbi__jbias[n] & (sgn - 1)); -} - -// get some unsigned bits -stbi_inline static int stbi__jpeg_get_bits(stbi__jpeg *j, int n) -{ - unsigned int k; - if (j->code_bits < n) stbi__grow_buffer_unsafe(j); - if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing - k = stbi_lrot(j->code_buffer, n); - j->code_buffer = k & ~stbi__bmask[n]; - k &= stbi__bmask[n]; - j->code_bits -= n; - return k; -} - -stbi_inline static int stbi__jpeg_get_bit(stbi__jpeg *j) -{ - unsigned int k; - if (j->code_bits < 1) stbi__grow_buffer_unsafe(j); - if (j->code_bits < 1) return 0; // ran out of bits from stream, return 0s intead of continuing - k = j->code_buffer; - j->code_buffer <<= 1; - --j->code_bits; - return k & 0x80000000; -} - -// given a value that's at position X in the zigzag stream, -// where does it appear in the 8x8 matrix coded as row-major? -static const stbi_uc stbi__jpeg_dezigzag[64+15] = -{ - 0, 1, 8, 16, 9, 2, 3, 10, - 17, 24, 32, 25, 18, 11, 4, 5, - 12, 19, 26, 33, 40, 48, 41, 34, - 27, 20, 13, 6, 7, 14, 21, 28, - 35, 42, 49, 56, 57, 50, 43, 36, - 29, 22, 15, 23, 30, 37, 44, 51, - 58, 59, 52, 45, 38, 31, 39, 46, - 53, 60, 61, 54, 47, 55, 62, 63, - // let corrupt input sample past end - 63, 63, 63, 63, 63, 63, 63, 63, - 63, 63, 63, 63, 63, 63, 63 -}; - -// decode one 64-entry block-- -static int stbi__jpeg_decode_block(stbi__jpeg *j, short data[64], stbi__huffman *hdc, stbi__huffman *hac, stbi__int16 *fac, int b, stbi__uint16 *dequant) -{ - int diff,dc,k; - int t; - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - t = stbi__jpeg_huff_decode(j, hdc); - if (t < 0 || t > 15) return stbi__err("bad huffman code","Corrupt JPEG"); - - // 0 all the ac values now so we can do it 32-bits at a time - memset(data,0,64*sizeof(data[0])); - - diff = t ? stbi__extend_receive(j, t) : 0; - if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta","Corrupt JPEG"); - dc = j->img_comp[b].dc_pred + diff; - j->img_comp[b].dc_pred = dc; - if (!stbi__mul2shorts_valid(dc, dequant[0])) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - data[0] = (short) (dc * dequant[0]); - - // decode AC components, see JPEG spec - k = 1; - do { - unsigned int zig; - int c,r,s; - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - r = fac[c]; - if (r) { // fast-AC path - k += (r >> 4) & 15; // run - s = r & 15; // combined length - if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); - j->code_buffer <<= s; - j->code_bits -= s; - // decode into unzigzag'd location - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) ((r >> 8) * dequant[zig]); - } else { - int rs = stbi__jpeg_huff_decode(j, hac); - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (rs != 0xf0) break; // end block - k += 16; - } else { - k += r; - // decode into unzigzag'd location - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) (stbi__extend_receive(j,s) * dequant[zig]); - } - } - } while (k < 64); - return 1; -} - -static int stbi__jpeg_decode_block_prog_dc(stbi__jpeg *j, short data[64], stbi__huffman *hdc, int b) -{ - int diff,dc; - int t; - if (j->spec_end != 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - - if (j->succ_high == 0) { - // first scan for DC coefficient, must be first - memset(data,0,64*sizeof(data[0])); // 0 all the ac values now - t = stbi__jpeg_huff_decode(j, hdc); - if (t < 0 || t > 15) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - diff = t ? stbi__extend_receive(j, t) : 0; - - if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta", "Corrupt JPEG"); - dc = j->img_comp[b].dc_pred + diff; - j->img_comp[b].dc_pred = dc; - if (!stbi__mul2shorts_valid(dc, 1 << j->succ_low)) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - data[0] = (short) (dc * (1 << j->succ_low)); - } else { - // refinement scan for DC coefficient - if (stbi__jpeg_get_bit(j)) - data[0] += (short) (1 << j->succ_low); - } - return 1; -} - -// @OPTIMIZE: store non-zigzagged during the decode passes, -// and only de-zigzag when dequantizing -static int stbi__jpeg_decode_block_prog_ac(stbi__jpeg *j, short data[64], stbi__huffman *hac, stbi__int16 *fac) -{ - int k; - if (j->spec_start == 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - - if (j->succ_high == 0) { - int shift = j->succ_low; - - if (j->eob_run) { - --j->eob_run; - return 1; - } - - k = j->spec_start; - do { - unsigned int zig; - int c,r,s; - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - r = fac[c]; - if (r) { // fast-AC path - k += (r >> 4) & 15; // run - s = r & 15; // combined length - if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); - j->code_buffer <<= s; - j->code_bits -= s; - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) ((r >> 8) * (1 << shift)); - } else { - int rs = stbi__jpeg_huff_decode(j, hac); - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (r < 15) { - j->eob_run = (1 << r); - if (r) - j->eob_run += stbi__jpeg_get_bits(j, r); - --j->eob_run; - break; - } - k += 16; - } else { - k += r; - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) (stbi__extend_receive(j,s) * (1 << shift)); - } - } - } while (k <= j->spec_end); - } else { - // refinement scan for these AC coefficients - - short bit = (short) (1 << j->succ_low); - - if (j->eob_run) { - --j->eob_run; - for (k = j->spec_start; k <= j->spec_end; ++k) { - short *p = &data[stbi__jpeg_dezigzag[k]]; - if (*p != 0) - if (stbi__jpeg_get_bit(j)) - if ((*p & bit)==0) { - if (*p > 0) - *p += bit; - else - *p -= bit; - } - } - } else { - k = j->spec_start; - do { - int r,s; - int rs = stbi__jpeg_huff_decode(j, hac); // @OPTIMIZE see if we can use the fast path here, advance-by-r is so slow, eh - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (r < 15) { - j->eob_run = (1 << r) - 1; - if (r) - j->eob_run += stbi__jpeg_get_bits(j, r); - r = 64; // force end of block - } else { - // r=15 s=0 should write 16 0s, so we just do - // a run of 15 0s and then write s (which is 0), - // so we don't have to do anything special here - } - } else { - if (s != 1) return stbi__err("bad huffman code", "Corrupt JPEG"); - // sign bit - if (stbi__jpeg_get_bit(j)) - s = bit; - else - s = -bit; - } - - // advance by r - while (k <= j->spec_end) { - short *p = &data[stbi__jpeg_dezigzag[k++]]; - if (*p != 0) { - if (stbi__jpeg_get_bit(j)) - if ((*p & bit)==0) { - if (*p > 0) - *p += bit; - else - *p -= bit; - } - } else { - if (r == 0) { - *p = (short) s; - break; - } - --r; - } - } - } while (k <= j->spec_end); - } - } - return 1; -} - -// take a -128..127 value and stbi__clamp it and convert to 0..255 -stbi_inline static stbi_uc stbi__clamp(int x) -{ - // trick to use a single test to catch both cases - if ((unsigned int) x > 255) { - if (x < 0) return 0; - if (x > 255) return 255; - } - return (stbi_uc) x; -} - -#define stbi__f2f(x) ((int) (((x) * 4096 + 0.5))) -#define stbi__fsh(x) ((x) * 4096) - -// derived from jidctint -- DCT_ISLOW -#define STBI__IDCT_1D(s0,s1,s2,s3,s4,s5,s6,s7) \ - int t0,t1,t2,t3,p1,p2,p3,p4,p5,x0,x1,x2,x3; \ - p2 = s2; \ - p3 = s6; \ - p1 = (p2+p3) * stbi__f2f(0.5411961f); \ - t2 = p1 + p3*stbi__f2f(-1.847759065f); \ - t3 = p1 + p2*stbi__f2f( 0.765366865f); \ - p2 = s0; \ - p3 = s4; \ - t0 = stbi__fsh(p2+p3); \ - t1 = stbi__fsh(p2-p3); \ - x0 = t0+t3; \ - x3 = t0-t3; \ - x1 = t1+t2; \ - x2 = t1-t2; \ - t0 = s7; \ - t1 = s5; \ - t2 = s3; \ - t3 = s1; \ - p3 = t0+t2; \ - p4 = t1+t3; \ - p1 = t0+t3; \ - p2 = t1+t2; \ - p5 = (p3+p4)*stbi__f2f( 1.175875602f); \ - t0 = t0*stbi__f2f( 0.298631336f); \ - t1 = t1*stbi__f2f( 2.053119869f); \ - t2 = t2*stbi__f2f( 3.072711026f); \ - t3 = t3*stbi__f2f( 1.501321110f); \ - p1 = p5 + p1*stbi__f2f(-0.899976223f); \ - p2 = p5 + p2*stbi__f2f(-2.562915447f); \ - p3 = p3*stbi__f2f(-1.961570560f); \ - p4 = p4*stbi__f2f(-0.390180644f); \ - t3 += p1+p4; \ - t2 += p2+p3; \ - t1 += p2+p4; \ - t0 += p1+p3; - -static void stbi__idct_block(stbi_uc *out, int out_stride, short data[64]) -{ - int i,val[64],*v=val; - stbi_uc *o; - short *d = data; - - // columns - for (i=0; i < 8; ++i,++d, ++v) { - // if all zeroes, shortcut -- this avoids dequantizing 0s and IDCTing - if (d[ 8]==0 && d[16]==0 && d[24]==0 && d[32]==0 - && d[40]==0 && d[48]==0 && d[56]==0) { - // no shortcut 0 seconds - // (1|2|3|4|5|6|7)==0 0 seconds - // all separate -0.047 seconds - // 1 && 2|3 && 4|5 && 6|7: -0.047 seconds - int dcterm = d[0]*4; - v[0] = v[8] = v[16] = v[24] = v[32] = v[40] = v[48] = v[56] = dcterm; - } else { - STBI__IDCT_1D(d[ 0],d[ 8],d[16],d[24],d[32],d[40],d[48],d[56]) - // constants scaled things up by 1<<12; let's bring them back - // down, but keep 2 extra bits of precision - x0 += 512; x1 += 512; x2 += 512; x3 += 512; - v[ 0] = (x0+t3) >> 10; - v[56] = (x0-t3) >> 10; - v[ 8] = (x1+t2) >> 10; - v[48] = (x1-t2) >> 10; - v[16] = (x2+t1) >> 10; - v[40] = (x2-t1) >> 10; - v[24] = (x3+t0) >> 10; - v[32] = (x3-t0) >> 10; - } - } - - for (i=0, v=val, o=out; i < 8; ++i,v+=8,o+=out_stride) { - // no fast case since the first 1D IDCT spread components out - STBI__IDCT_1D(v[0],v[1],v[2],v[3],v[4],v[5],v[6],v[7]) - // constants scaled things up by 1<<12, plus we had 1<<2 from first - // loop, plus horizontal and vertical each scale by sqrt(8) so together - // we've got an extra 1<<3, so 1<<17 total we need to remove. - // so we want to round that, which means adding 0.5 * 1<<17, - // aka 65536. Also, we'll end up with -128 to 127 that we want - // to encode as 0..255 by adding 128, so we'll add that before the shift - x0 += 65536 + (128<<17); - x1 += 65536 + (128<<17); - x2 += 65536 + (128<<17); - x3 += 65536 + (128<<17); - // tried computing the shifts into temps, or'ing the temps to see - // if any were out of range, but that was slower - o[0] = stbi__clamp((x0+t3) >> 17); - o[7] = stbi__clamp((x0-t3) >> 17); - o[1] = stbi__clamp((x1+t2) >> 17); - o[6] = stbi__clamp((x1-t2) >> 17); - o[2] = stbi__clamp((x2+t1) >> 17); - o[5] = stbi__clamp((x2-t1) >> 17); - o[3] = stbi__clamp((x3+t0) >> 17); - o[4] = stbi__clamp((x3-t0) >> 17); - } -} - -#ifdef STBI_SSE2 -// sse2 integer IDCT. not the fastest possible implementation but it -// produces bit-identical results to the generic C version so it's -// fully "transparent". -static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) -{ - // This is constructed to match our regular (generic) integer IDCT exactly. - __m128i row0, row1, row2, row3, row4, row5, row6, row7; - __m128i tmp; - - // dot product constant: even elems=x, odd elems=y - #define dct_const(x,y) _mm_setr_epi16((x),(y),(x),(y),(x),(y),(x),(y)) - - // out(0) = c0[even]*x + c0[odd]*y (c0, x, y 16-bit, out 32-bit) - // out(1) = c1[even]*x + c1[odd]*y - #define dct_rot(out0,out1, x,y,c0,c1) \ - __m128i c0##lo = _mm_unpacklo_epi16((x),(y)); \ - __m128i c0##hi = _mm_unpackhi_epi16((x),(y)); \ - __m128i out0##_l = _mm_madd_epi16(c0##lo, c0); \ - __m128i out0##_h = _mm_madd_epi16(c0##hi, c0); \ - __m128i out1##_l = _mm_madd_epi16(c0##lo, c1); \ - __m128i out1##_h = _mm_madd_epi16(c0##hi, c1) - - // out = in << 12 (in 16-bit, out 32-bit) - #define dct_widen(out, in) \ - __m128i out##_l = _mm_srai_epi32(_mm_unpacklo_epi16(_mm_setzero_si128(), (in)), 4); \ - __m128i out##_h = _mm_srai_epi32(_mm_unpackhi_epi16(_mm_setzero_si128(), (in)), 4) - - // wide add - #define dct_wadd(out, a, b) \ - __m128i out##_l = _mm_add_epi32(a##_l, b##_l); \ - __m128i out##_h = _mm_add_epi32(a##_h, b##_h) - - // wide sub - #define dct_wsub(out, a, b) \ - __m128i out##_l = _mm_sub_epi32(a##_l, b##_l); \ - __m128i out##_h = _mm_sub_epi32(a##_h, b##_h) - - // butterfly a/b, add bias, then shift by "s" and pack - #define dct_bfly32o(out0, out1, a,b,bias,s) \ - { \ - __m128i abiased_l = _mm_add_epi32(a##_l, bias); \ - __m128i abiased_h = _mm_add_epi32(a##_h, bias); \ - dct_wadd(sum, abiased, b); \ - dct_wsub(dif, abiased, b); \ - out0 = _mm_packs_epi32(_mm_srai_epi32(sum_l, s), _mm_srai_epi32(sum_h, s)); \ - out1 = _mm_packs_epi32(_mm_srai_epi32(dif_l, s), _mm_srai_epi32(dif_h, s)); \ - } - - // 8-bit interleave step (for transposes) - #define dct_interleave8(a, b) \ - tmp = a; \ - a = _mm_unpacklo_epi8(a, b); \ - b = _mm_unpackhi_epi8(tmp, b) - - // 16-bit interleave step (for transposes) - #define dct_interleave16(a, b) \ - tmp = a; \ - a = _mm_unpacklo_epi16(a, b); \ - b = _mm_unpackhi_epi16(tmp, b) - - #define dct_pass(bias,shift) \ - { \ - /* even part */ \ - dct_rot(t2e,t3e, row2,row6, rot0_0,rot0_1); \ - __m128i sum04 = _mm_add_epi16(row0, row4); \ - __m128i dif04 = _mm_sub_epi16(row0, row4); \ - dct_widen(t0e, sum04); \ - dct_widen(t1e, dif04); \ - dct_wadd(x0, t0e, t3e); \ - dct_wsub(x3, t0e, t3e); \ - dct_wadd(x1, t1e, t2e); \ - dct_wsub(x2, t1e, t2e); \ - /* odd part */ \ - dct_rot(y0o,y2o, row7,row3, rot2_0,rot2_1); \ - dct_rot(y1o,y3o, row5,row1, rot3_0,rot3_1); \ - __m128i sum17 = _mm_add_epi16(row1, row7); \ - __m128i sum35 = _mm_add_epi16(row3, row5); \ - dct_rot(y4o,y5o, sum17,sum35, rot1_0,rot1_1); \ - dct_wadd(x4, y0o, y4o); \ - dct_wadd(x5, y1o, y5o); \ - dct_wadd(x6, y2o, y5o); \ - dct_wadd(x7, y3o, y4o); \ - dct_bfly32o(row0,row7, x0,x7,bias,shift); \ - dct_bfly32o(row1,row6, x1,x6,bias,shift); \ - dct_bfly32o(row2,row5, x2,x5,bias,shift); \ - dct_bfly32o(row3,row4, x3,x4,bias,shift); \ - } - - __m128i rot0_0 = dct_const(stbi__f2f(0.5411961f), stbi__f2f(0.5411961f) + stbi__f2f(-1.847759065f)); - __m128i rot0_1 = dct_const(stbi__f2f(0.5411961f) + stbi__f2f( 0.765366865f), stbi__f2f(0.5411961f)); - __m128i rot1_0 = dct_const(stbi__f2f(1.175875602f) + stbi__f2f(-0.899976223f), stbi__f2f(1.175875602f)); - __m128i rot1_1 = dct_const(stbi__f2f(1.175875602f), stbi__f2f(1.175875602f) + stbi__f2f(-2.562915447f)); - __m128i rot2_0 = dct_const(stbi__f2f(-1.961570560f) + stbi__f2f( 0.298631336f), stbi__f2f(-1.961570560f)); - __m128i rot2_1 = dct_const(stbi__f2f(-1.961570560f), stbi__f2f(-1.961570560f) + stbi__f2f( 3.072711026f)); - __m128i rot3_0 = dct_const(stbi__f2f(-0.390180644f) + stbi__f2f( 2.053119869f), stbi__f2f(-0.390180644f)); - __m128i rot3_1 = dct_const(stbi__f2f(-0.390180644f), stbi__f2f(-0.390180644f) + stbi__f2f( 1.501321110f)); - - // rounding biases in column/row passes, see stbi__idct_block for explanation. - __m128i bias_0 = _mm_set1_epi32(512); - __m128i bias_1 = _mm_set1_epi32(65536 + (128<<17)); - - // load - row0 = _mm_load_si128((const __m128i *) (data + 0*8)); - row1 = _mm_load_si128((const __m128i *) (data + 1*8)); - row2 = _mm_load_si128((const __m128i *) (data + 2*8)); - row3 = _mm_load_si128((const __m128i *) (data + 3*8)); - row4 = _mm_load_si128((const __m128i *) (data + 4*8)); - row5 = _mm_load_si128((const __m128i *) (data + 5*8)); - row6 = _mm_load_si128((const __m128i *) (data + 6*8)); - row7 = _mm_load_si128((const __m128i *) (data + 7*8)); - - // column pass - dct_pass(bias_0, 10); - - { - // 16bit 8x8 transpose pass 1 - dct_interleave16(row0, row4); - dct_interleave16(row1, row5); - dct_interleave16(row2, row6); - dct_interleave16(row3, row7); - - // transpose pass 2 - dct_interleave16(row0, row2); - dct_interleave16(row1, row3); - dct_interleave16(row4, row6); - dct_interleave16(row5, row7); - - // transpose pass 3 - dct_interleave16(row0, row1); - dct_interleave16(row2, row3); - dct_interleave16(row4, row5); - dct_interleave16(row6, row7); - } - - // row pass - dct_pass(bias_1, 17); - - { - // pack - __m128i p0 = _mm_packus_epi16(row0, row1); // a0a1a2a3...a7b0b1b2b3...b7 - __m128i p1 = _mm_packus_epi16(row2, row3); - __m128i p2 = _mm_packus_epi16(row4, row5); - __m128i p3 = _mm_packus_epi16(row6, row7); - - // 8bit 8x8 transpose pass 1 - dct_interleave8(p0, p2); // a0e0a1e1... - dct_interleave8(p1, p3); // c0g0c1g1... - - // transpose pass 2 - dct_interleave8(p0, p1); // a0c0e0g0... - dct_interleave8(p2, p3); // b0d0f0h0... - - // transpose pass 3 - dct_interleave8(p0, p2); // a0b0c0d0... - dct_interleave8(p1, p3); // a4b4c4d4... - - // store - _mm_storel_epi64((__m128i *) out, p0); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p0, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p2); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p2, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p1); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p1, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p3); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p3, 0x4e)); - } - -#undef dct_const -#undef dct_rot -#undef dct_widen -#undef dct_wadd -#undef dct_wsub -#undef dct_bfly32o -#undef dct_interleave8 -#undef dct_interleave16 -#undef dct_pass -} - -#endif // STBI_SSE2 - -#ifdef STBI_NEON - -// NEON integer IDCT. should produce bit-identical -// results to the generic C version. -static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) -{ - int16x8_t row0, row1, row2, row3, row4, row5, row6, row7; - - int16x4_t rot0_0 = vdup_n_s16(stbi__f2f(0.5411961f)); - int16x4_t rot0_1 = vdup_n_s16(stbi__f2f(-1.847759065f)); - int16x4_t rot0_2 = vdup_n_s16(stbi__f2f( 0.765366865f)); - int16x4_t rot1_0 = vdup_n_s16(stbi__f2f( 1.175875602f)); - int16x4_t rot1_1 = vdup_n_s16(stbi__f2f(-0.899976223f)); - int16x4_t rot1_2 = vdup_n_s16(stbi__f2f(-2.562915447f)); - int16x4_t rot2_0 = vdup_n_s16(stbi__f2f(-1.961570560f)); - int16x4_t rot2_1 = vdup_n_s16(stbi__f2f(-0.390180644f)); - int16x4_t rot3_0 = vdup_n_s16(stbi__f2f( 0.298631336f)); - int16x4_t rot3_1 = vdup_n_s16(stbi__f2f( 2.053119869f)); - int16x4_t rot3_2 = vdup_n_s16(stbi__f2f( 3.072711026f)); - int16x4_t rot3_3 = vdup_n_s16(stbi__f2f( 1.501321110f)); - -#define dct_long_mul(out, inq, coeff) \ - int32x4_t out##_l = vmull_s16(vget_low_s16(inq), coeff); \ - int32x4_t out##_h = vmull_s16(vget_high_s16(inq), coeff) - -#define dct_long_mac(out, acc, inq, coeff) \ - int32x4_t out##_l = vmlal_s16(acc##_l, vget_low_s16(inq), coeff); \ - int32x4_t out##_h = vmlal_s16(acc##_h, vget_high_s16(inq), coeff) - -#define dct_widen(out, inq) \ - int32x4_t out##_l = vshll_n_s16(vget_low_s16(inq), 12); \ - int32x4_t out##_h = vshll_n_s16(vget_high_s16(inq), 12) - -// wide add -#define dct_wadd(out, a, b) \ - int32x4_t out##_l = vaddq_s32(a##_l, b##_l); \ - int32x4_t out##_h = vaddq_s32(a##_h, b##_h) - -// wide sub -#define dct_wsub(out, a, b) \ - int32x4_t out##_l = vsubq_s32(a##_l, b##_l); \ - int32x4_t out##_h = vsubq_s32(a##_h, b##_h) - -// butterfly a/b, then shift using "shiftop" by "s" and pack -#define dct_bfly32o(out0,out1, a,b,shiftop,s) \ - { \ - dct_wadd(sum, a, b); \ - dct_wsub(dif, a, b); \ - out0 = vcombine_s16(shiftop(sum_l, s), shiftop(sum_h, s)); \ - out1 = vcombine_s16(shiftop(dif_l, s), shiftop(dif_h, s)); \ - } - -#define dct_pass(shiftop, shift) \ - { \ - /* even part */ \ - int16x8_t sum26 = vaddq_s16(row2, row6); \ - dct_long_mul(p1e, sum26, rot0_0); \ - dct_long_mac(t2e, p1e, row6, rot0_1); \ - dct_long_mac(t3e, p1e, row2, rot0_2); \ - int16x8_t sum04 = vaddq_s16(row0, row4); \ - int16x8_t dif04 = vsubq_s16(row0, row4); \ - dct_widen(t0e, sum04); \ - dct_widen(t1e, dif04); \ - dct_wadd(x0, t0e, t3e); \ - dct_wsub(x3, t0e, t3e); \ - dct_wadd(x1, t1e, t2e); \ - dct_wsub(x2, t1e, t2e); \ - /* odd part */ \ - int16x8_t sum15 = vaddq_s16(row1, row5); \ - int16x8_t sum17 = vaddq_s16(row1, row7); \ - int16x8_t sum35 = vaddq_s16(row3, row5); \ - int16x8_t sum37 = vaddq_s16(row3, row7); \ - int16x8_t sumodd = vaddq_s16(sum17, sum35); \ - dct_long_mul(p5o, sumodd, rot1_0); \ - dct_long_mac(p1o, p5o, sum17, rot1_1); \ - dct_long_mac(p2o, p5o, sum35, rot1_2); \ - dct_long_mul(p3o, sum37, rot2_0); \ - dct_long_mul(p4o, sum15, rot2_1); \ - dct_wadd(sump13o, p1o, p3o); \ - dct_wadd(sump24o, p2o, p4o); \ - dct_wadd(sump23o, p2o, p3o); \ - dct_wadd(sump14o, p1o, p4o); \ - dct_long_mac(x4, sump13o, row7, rot3_0); \ - dct_long_mac(x5, sump24o, row5, rot3_1); \ - dct_long_mac(x6, sump23o, row3, rot3_2); \ - dct_long_mac(x7, sump14o, row1, rot3_3); \ - dct_bfly32o(row0,row7, x0,x7,shiftop,shift); \ - dct_bfly32o(row1,row6, x1,x6,shiftop,shift); \ - dct_bfly32o(row2,row5, x2,x5,shiftop,shift); \ - dct_bfly32o(row3,row4, x3,x4,shiftop,shift); \ - } - - // load - row0 = vld1q_s16(data + 0*8); - row1 = vld1q_s16(data + 1*8); - row2 = vld1q_s16(data + 2*8); - row3 = vld1q_s16(data + 3*8); - row4 = vld1q_s16(data + 4*8); - row5 = vld1q_s16(data + 5*8); - row6 = vld1q_s16(data + 6*8); - row7 = vld1q_s16(data + 7*8); - - // add DC bias - row0 = vaddq_s16(row0, vsetq_lane_s16(1024, vdupq_n_s16(0), 0)); - - // column pass - dct_pass(vrshrn_n_s32, 10); - - // 16bit 8x8 transpose - { -// these three map to a single VTRN.16, VTRN.32, and VSWP, respectively. -// whether compilers actually get this is another story, sadly. -#define dct_trn16(x, y) { int16x8x2_t t = vtrnq_s16(x, y); x = t.val[0]; y = t.val[1]; } -#define dct_trn32(x, y) { int32x4x2_t t = vtrnq_s32(vreinterpretq_s32_s16(x), vreinterpretq_s32_s16(y)); x = vreinterpretq_s16_s32(t.val[0]); y = vreinterpretq_s16_s32(t.val[1]); } -#define dct_trn64(x, y) { int16x8_t x0 = x; int16x8_t y0 = y; x = vcombine_s16(vget_low_s16(x0), vget_low_s16(y0)); y = vcombine_s16(vget_high_s16(x0), vget_high_s16(y0)); } - - // pass 1 - dct_trn16(row0, row1); // a0b0a2b2a4b4a6b6 - dct_trn16(row2, row3); - dct_trn16(row4, row5); - dct_trn16(row6, row7); - - // pass 2 - dct_trn32(row0, row2); // a0b0c0d0a4b4c4d4 - dct_trn32(row1, row3); - dct_trn32(row4, row6); - dct_trn32(row5, row7); - - // pass 3 - dct_trn64(row0, row4); // a0b0c0d0e0f0g0h0 - dct_trn64(row1, row5); - dct_trn64(row2, row6); - dct_trn64(row3, row7); - -#undef dct_trn16 -#undef dct_trn32 -#undef dct_trn64 - } - - // row pass - // vrshrn_n_s32 only supports shifts up to 16, we need - // 17. so do a non-rounding shift of 16 first then follow - // up with a rounding shift by 1. - dct_pass(vshrn_n_s32, 16); - - { - // pack and round - uint8x8_t p0 = vqrshrun_n_s16(row0, 1); - uint8x8_t p1 = vqrshrun_n_s16(row1, 1); - uint8x8_t p2 = vqrshrun_n_s16(row2, 1); - uint8x8_t p3 = vqrshrun_n_s16(row3, 1); - uint8x8_t p4 = vqrshrun_n_s16(row4, 1); - uint8x8_t p5 = vqrshrun_n_s16(row5, 1); - uint8x8_t p6 = vqrshrun_n_s16(row6, 1); - uint8x8_t p7 = vqrshrun_n_s16(row7, 1); - - // again, these can translate into one instruction, but often don't. -#define dct_trn8_8(x, y) { uint8x8x2_t t = vtrn_u8(x, y); x = t.val[0]; y = t.val[1]; } -#define dct_trn8_16(x, y) { uint16x4x2_t t = vtrn_u16(vreinterpret_u16_u8(x), vreinterpret_u16_u8(y)); x = vreinterpret_u8_u16(t.val[0]); y = vreinterpret_u8_u16(t.val[1]); } -#define dct_trn8_32(x, y) { uint32x2x2_t t = vtrn_u32(vreinterpret_u32_u8(x), vreinterpret_u32_u8(y)); x = vreinterpret_u8_u32(t.val[0]); y = vreinterpret_u8_u32(t.val[1]); } - - // sadly can't use interleaved stores here since we only write - // 8 bytes to each scan line! - - // 8x8 8-bit transpose pass 1 - dct_trn8_8(p0, p1); - dct_trn8_8(p2, p3); - dct_trn8_8(p4, p5); - dct_trn8_8(p6, p7); - - // pass 2 - dct_trn8_16(p0, p2); - dct_trn8_16(p1, p3); - dct_trn8_16(p4, p6); - dct_trn8_16(p5, p7); - - // pass 3 - dct_trn8_32(p0, p4); - dct_trn8_32(p1, p5); - dct_trn8_32(p2, p6); - dct_trn8_32(p3, p7); - - // store - vst1_u8(out, p0); out += out_stride; - vst1_u8(out, p1); out += out_stride; - vst1_u8(out, p2); out += out_stride; - vst1_u8(out, p3); out += out_stride; - vst1_u8(out, p4); out += out_stride; - vst1_u8(out, p5); out += out_stride; - vst1_u8(out, p6); out += out_stride; - vst1_u8(out, p7); - -#undef dct_trn8_8 -#undef dct_trn8_16 -#undef dct_trn8_32 - } - -#undef dct_long_mul -#undef dct_long_mac -#undef dct_widen -#undef dct_wadd -#undef dct_wsub -#undef dct_bfly32o -#undef dct_pass -} - -#endif // STBI_NEON - -#define STBI__MARKER_none 0xff -// if there's a pending marker from the entropy stream, return that -// otherwise, fetch from the stream and get a marker. if there's no -// marker, return 0xff, which is never a valid marker value -static stbi_uc stbi__get_marker(stbi__jpeg *j) -{ - stbi_uc x; - if (j->marker != STBI__MARKER_none) { x = j->marker; j->marker = STBI__MARKER_none; return x; } - x = stbi__get8(j->s); - if (x != 0xff) return STBI__MARKER_none; - while (x == 0xff) - x = stbi__get8(j->s); // consume repeated 0xff fill bytes - return x; -} - -// in each scan, we'll have scan_n components, and the order -// of the components is specified by order[] -#define STBI__RESTART(x) ((x) >= 0xd0 && (x) <= 0xd7) - -// after a restart interval, stbi__jpeg_reset the entropy decoder and -// the dc prediction -static void stbi__jpeg_reset(stbi__jpeg *j) -{ - j->code_bits = 0; - j->code_buffer = 0; - j->nomore = 0; - j->img_comp[0].dc_pred = j->img_comp[1].dc_pred = j->img_comp[2].dc_pred = j->img_comp[3].dc_pred = 0; - j->marker = STBI__MARKER_none; - j->todo = j->restart_interval ? j->restart_interval : 0x7fffffff; - j->eob_run = 0; - // no more than 1<<31 MCUs if no restart_interal? that's plenty safe, - // since we don't even allow 1<<30 pixels -} - -static int stbi__parse_entropy_coded_data(stbi__jpeg *z) -{ - stbi__jpeg_reset(z); - if (!z->progressive) { - if (z->scan_n == 1) { - int i,j; - STBI_SIMD_ALIGN(short, data[64]); - int n = z->order[0]; - // non-interleaved data, we just need to process one block at a time, - // in trivial scanline order - // number of blocks to do just depends on how many actual "pixels" this - // component has, independent of interleaved MCU blocking and such - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); - // every data block is an MCU, so countdown the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - // if it's NOT a restart, then just bail, so we get corrupt data - // rather than no data - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } else { // interleaved - int i,j,k,x,y; - STBI_SIMD_ALIGN(short, data[64]); - for (j=0; j < z->img_mcu_y; ++j) { - for (i=0; i < z->img_mcu_x; ++i) { - // scan an interleaved mcu... process scan_n components in order - for (k=0; k < z->scan_n; ++k) { - int n = z->order[k]; - // scan out an mcu's worth of this component; that's just determined - // by the basic H and V specified for the component - for (y=0; y < z->img_comp[n].v; ++y) { - for (x=0; x < z->img_comp[n].h; ++x) { - int x2 = (i*z->img_comp[n].h + x)*8; - int y2 = (j*z->img_comp[n].v + y)*8; - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*y2+x2, z->img_comp[n].w2, data); - } - } - } - // after all interleaved components, that's an interleaved MCU, - // so now count down the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } - } else { - if (z->scan_n == 1) { - int i,j; - int n = z->order[0]; - // non-interleaved data, we just need to process one block at a time, - // in trivial scanline order - // number of blocks to do just depends on how many actual "pixels" this - // component has, independent of interleaved MCU blocking and such - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); - if (z->spec_start == 0) { - if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) - return 0; - } else { - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block_prog_ac(z, data, &z->huff_ac[ha], z->fast_ac[ha])) - return 0; - } - // every data block is an MCU, so countdown the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } else { // interleaved - int i,j,k,x,y; - for (j=0; j < z->img_mcu_y; ++j) { - for (i=0; i < z->img_mcu_x; ++i) { - // scan an interleaved mcu... process scan_n components in order - for (k=0; k < z->scan_n; ++k) { - int n = z->order[k]; - // scan out an mcu's worth of this component; that's just determined - // by the basic H and V specified for the component - for (y=0; y < z->img_comp[n].v; ++y) { - for (x=0; x < z->img_comp[n].h; ++x) { - int x2 = (i*z->img_comp[n].h + x); - int y2 = (j*z->img_comp[n].v + y); - short *data = z->img_comp[n].coeff + 64 * (x2 + y2 * z->img_comp[n].coeff_w); - if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) - return 0; - } - } - } - // after all interleaved components, that's an interleaved MCU, - // so now count down the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } - } -} - -static void stbi__jpeg_dequantize(short *data, stbi__uint16 *dequant) -{ - int i; - for (i=0; i < 64; ++i) - data[i] *= dequant[i]; -} - -static void stbi__jpeg_finish(stbi__jpeg *z) -{ - if (z->progressive) { - // dequantize and idct the data - int i,j,n; - for (n=0; n < z->s->img_n; ++n) { - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); - stbi__jpeg_dequantize(data, z->dequant[z->img_comp[n].tq]); - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); - } - } - } - } -} - -static int stbi__process_marker(stbi__jpeg *z, int m) -{ - int L; - switch (m) { - case STBI__MARKER_none: // no marker found - return stbi__err("expected marker","Corrupt JPEG"); - - case 0xDD: // DRI - specify restart interval - if (stbi__get16be(z->s) != 4) return stbi__err("bad DRI len","Corrupt JPEG"); - z->restart_interval = stbi__get16be(z->s); - return 1; - - case 0xDB: // DQT - define quantization table - L = stbi__get16be(z->s)-2; - while (L > 0) { - int q = stbi__get8(z->s); - int p = q >> 4, sixteen = (p != 0); - int t = q & 15,i; - if (p != 0 && p != 1) return stbi__err("bad DQT type","Corrupt JPEG"); - if (t > 3) return stbi__err("bad DQT table","Corrupt JPEG"); - - for (i=0; i < 64; ++i) - z->dequant[t][stbi__jpeg_dezigzag[i]] = (stbi__uint16)(sixteen ? stbi__get16be(z->s) : stbi__get8(z->s)); - L -= (sixteen ? 129 : 65); - } - return L==0; - - case 0xC4: // DHT - define huffman table - L = stbi__get16be(z->s)-2; - while (L > 0) { - stbi_uc *v; - int sizes[16],i,n=0; - int q = stbi__get8(z->s); - int tc = q >> 4; - int th = q & 15; - if (tc > 1 || th > 3) return stbi__err("bad DHT header","Corrupt JPEG"); - for (i=0; i < 16; ++i) { - sizes[i] = stbi__get8(z->s); - n += sizes[i]; - } - if(n > 256) return stbi__err("bad DHT header","Corrupt JPEG"); // Loop over i < n would write past end of values! - L -= 17; - if (tc == 0) { - if (!stbi__build_huffman(z->huff_dc+th, sizes)) return 0; - v = z->huff_dc[th].values; - } else { - if (!stbi__build_huffman(z->huff_ac+th, sizes)) return 0; - v = z->huff_ac[th].values; - } - for (i=0; i < n; ++i) - v[i] = stbi__get8(z->s); - if (tc != 0) - stbi__build_fast_ac(z->fast_ac[th], z->huff_ac + th); - L -= n; - } - return L==0; - } - - // check for comment block or APP blocks - if ((m >= 0xE0 && m <= 0xEF) || m == 0xFE) { - L = stbi__get16be(z->s); - if (L < 2) { - if (m == 0xFE) - return stbi__err("bad COM len","Corrupt JPEG"); - else - return stbi__err("bad APP len","Corrupt JPEG"); - } - L -= 2; - - if (m == 0xE0 && L >= 5) { // JFIF APP0 segment - static const unsigned char tag[5] = {'J','F','I','F','\0'}; - int ok = 1; - int i; - for (i=0; i < 5; ++i) - if (stbi__get8(z->s) != tag[i]) - ok = 0; - L -= 5; - if (ok) - z->jfif = 1; - } else if (m == 0xEE && L >= 12) { // Adobe APP14 segment - static const unsigned char tag[6] = {'A','d','o','b','e','\0'}; - int ok = 1; - int i; - for (i=0; i < 6; ++i) - if (stbi__get8(z->s) != tag[i]) - ok = 0; - L -= 6; - if (ok) { - stbi__get8(z->s); // version - stbi__get16be(z->s); // flags0 - stbi__get16be(z->s); // flags1 - z->app14_color_transform = stbi__get8(z->s); // color transform - L -= 6; - } - } - - stbi__skip(z->s, L); - return 1; - } - - return stbi__err("unknown marker","Corrupt JPEG"); -} - -// after we see SOS -static int stbi__process_scan_header(stbi__jpeg *z) -{ - int i; - int Ls = stbi__get16be(z->s); - z->scan_n = stbi__get8(z->s); - if (z->scan_n < 1 || z->scan_n > 4 || z->scan_n > (int) z->s->img_n) return stbi__err("bad SOS component count","Corrupt JPEG"); - if (Ls != 6+2*z->scan_n) return stbi__err("bad SOS len","Corrupt JPEG"); - for (i=0; i < z->scan_n; ++i) { - int id = stbi__get8(z->s), which; - int q = stbi__get8(z->s); - for (which = 0; which < z->s->img_n; ++which) - if (z->img_comp[which].id == id) - break; - if (which == z->s->img_n) return 0; // no match - z->img_comp[which].hd = q >> 4; if (z->img_comp[which].hd > 3) return stbi__err("bad DC huff","Corrupt JPEG"); - z->img_comp[which].ha = q & 15; if (z->img_comp[which].ha > 3) return stbi__err("bad AC huff","Corrupt JPEG"); - z->order[i] = which; - } - - { - int aa; - z->spec_start = stbi__get8(z->s); - z->spec_end = stbi__get8(z->s); // should be 63, but might be 0 - aa = stbi__get8(z->s); - z->succ_high = (aa >> 4); - z->succ_low = (aa & 15); - if (z->progressive) { - if (z->spec_start > 63 || z->spec_end > 63 || z->spec_start > z->spec_end || z->succ_high > 13 || z->succ_low > 13) - return stbi__err("bad SOS", "Corrupt JPEG"); - } else { - if (z->spec_start != 0) return stbi__err("bad SOS","Corrupt JPEG"); - if (z->succ_high != 0 || z->succ_low != 0) return stbi__err("bad SOS","Corrupt JPEG"); - z->spec_end = 63; - } - } - - return 1; -} - -static int stbi__free_jpeg_components(stbi__jpeg *z, int ncomp, int why) -{ - int i; - for (i=0; i < ncomp; ++i) { - if (z->img_comp[i].raw_data) { - STBI_FREE(z->img_comp[i].raw_data); - z->img_comp[i].raw_data = NULL; - z->img_comp[i].data = NULL; - } - if (z->img_comp[i].raw_coeff) { - STBI_FREE(z->img_comp[i].raw_coeff); - z->img_comp[i].raw_coeff = 0; - z->img_comp[i].coeff = 0; - } - if (z->img_comp[i].linebuf) { - STBI_FREE(z->img_comp[i].linebuf); - z->img_comp[i].linebuf = NULL; - } - } - return why; -} - -static int stbi__process_frame_header(stbi__jpeg *z, int scan) -{ - stbi__context *s = z->s; - int Lf,p,i,q, h_max=1,v_max=1,c; - Lf = stbi__get16be(s); if (Lf < 11) return stbi__err("bad SOF len","Corrupt JPEG"); // JPEG - p = stbi__get8(s); if (p != 8) return stbi__err("only 8-bit","JPEG format not supported: 8-bit only"); // JPEG baseline - s->img_y = stbi__get16be(s); if (s->img_y == 0) return stbi__err("no header height", "JPEG format not supported: delayed height"); // Legal, but we don't handle it--but neither does IJG - s->img_x = stbi__get16be(s); if (s->img_x == 0) return stbi__err("0 width","Corrupt JPEG"); // JPEG requires - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - c = stbi__get8(s); - if (c != 3 && c != 1 && c != 4) return stbi__err("bad component count","Corrupt JPEG"); - s->img_n = c; - for (i=0; i < c; ++i) { - z->img_comp[i].data = NULL; - z->img_comp[i].linebuf = NULL; - } - - if (Lf != 8+3*s->img_n) return stbi__err("bad SOF len","Corrupt JPEG"); - - z->rgb = 0; - for (i=0; i < s->img_n; ++i) { - static const unsigned char rgb[3] = { 'R', 'G', 'B' }; - z->img_comp[i].id = stbi__get8(s); - if (s->img_n == 3 && z->img_comp[i].id == rgb[i]) - ++z->rgb; - q = stbi__get8(s); - z->img_comp[i].h = (q >> 4); if (!z->img_comp[i].h || z->img_comp[i].h > 4) return stbi__err("bad H","Corrupt JPEG"); - z->img_comp[i].v = q & 15; if (!z->img_comp[i].v || z->img_comp[i].v > 4) return stbi__err("bad V","Corrupt JPEG"); - z->img_comp[i].tq = stbi__get8(s); if (z->img_comp[i].tq > 3) return stbi__err("bad TQ","Corrupt JPEG"); - } - - if (scan != STBI__SCAN_load) return 1; - - if (!stbi__mad3sizes_valid(s->img_x, s->img_y, s->img_n, 0)) return stbi__err("too large", "Image too large to decode"); - - for (i=0; i < s->img_n; ++i) { - if (z->img_comp[i].h > h_max) h_max = z->img_comp[i].h; - if (z->img_comp[i].v > v_max) v_max = z->img_comp[i].v; - } - - // check that plane subsampling factors are integer ratios; our resamplers can't deal with fractional ratios - // and I've never seen a non-corrupted JPEG file actually use them - for (i=0; i < s->img_n; ++i) { - if (h_max % z->img_comp[i].h != 0) return stbi__err("bad H","Corrupt JPEG"); - if (v_max % z->img_comp[i].v != 0) return stbi__err("bad V","Corrupt JPEG"); - } - - // compute interleaved mcu info - z->img_h_max = h_max; - z->img_v_max = v_max; - z->img_mcu_w = h_max * 8; - z->img_mcu_h = v_max * 8; - // these sizes can't be more than 17 bits - z->img_mcu_x = (s->img_x + z->img_mcu_w-1) / z->img_mcu_w; - z->img_mcu_y = (s->img_y + z->img_mcu_h-1) / z->img_mcu_h; - - for (i=0; i < s->img_n; ++i) { - // number of effective pixels (e.g. for non-interleaved MCU) - z->img_comp[i].x = (s->img_x * z->img_comp[i].h + h_max-1) / h_max; - z->img_comp[i].y = (s->img_y * z->img_comp[i].v + v_max-1) / v_max; - // to simplify generation, we'll allocate enough memory to decode - // the bogus oversized data from using interleaved MCUs and their - // big blocks (e.g. a 16x16 iMCU on an image of width 33); we won't - // discard the extra data until colorspace conversion - // - // img_mcu_x, img_mcu_y: <=17 bits; comp[i].h and .v are <=4 (checked earlier) - // so these muls can't overflow with 32-bit ints (which we require) - z->img_comp[i].w2 = z->img_mcu_x * z->img_comp[i].h * 8; - z->img_comp[i].h2 = z->img_mcu_y * z->img_comp[i].v * 8; - z->img_comp[i].coeff = 0; - z->img_comp[i].raw_coeff = 0; - z->img_comp[i].linebuf = NULL; - z->img_comp[i].raw_data = stbi__malloc_mad2(z->img_comp[i].w2, z->img_comp[i].h2, 15); - if (z->img_comp[i].raw_data == NULL) - return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); - // align blocks for idct using mmx/sse - z->img_comp[i].data = (stbi_uc*) (((size_t) z->img_comp[i].raw_data + 15) & ~15); - if (z->progressive) { - // w2, h2 are multiples of 8 (see above) - z->img_comp[i].coeff_w = z->img_comp[i].w2 / 8; - z->img_comp[i].coeff_h = z->img_comp[i].h2 / 8; - z->img_comp[i].raw_coeff = stbi__malloc_mad3(z->img_comp[i].w2, z->img_comp[i].h2, sizeof(short), 15); - if (z->img_comp[i].raw_coeff == NULL) - return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); - z->img_comp[i].coeff = (short*) (((size_t) z->img_comp[i].raw_coeff + 15) & ~15); - } - } - - return 1; -} - -// use comparisons since in some cases we handle more than one case (e.g. SOF) -#define stbi__DNL(x) ((x) == 0xdc) -#define stbi__SOI(x) ((x) == 0xd8) -#define stbi__EOI(x) ((x) == 0xd9) -#define stbi__SOF(x) ((x) == 0xc0 || (x) == 0xc1 || (x) == 0xc2) -#define stbi__SOS(x) ((x) == 0xda) - -#define stbi__SOF_progressive(x) ((x) == 0xc2) - -static int stbi__decode_jpeg_header(stbi__jpeg *z, int scan) -{ - int m; - z->jfif = 0; - z->app14_color_transform = -1; // valid values are 0,1,2 - z->marker = STBI__MARKER_none; // initialize cached marker to empty - m = stbi__get_marker(z); - if (!stbi__SOI(m)) return stbi__err("no SOI","Corrupt JPEG"); - if (scan == STBI__SCAN_type) return 1; - m = stbi__get_marker(z); - while (!stbi__SOF(m)) { - if (!stbi__process_marker(z,m)) return 0; - m = stbi__get_marker(z); - while (m == STBI__MARKER_none) { - // some files have extra padding after their blocks, so ok, we'll scan - if (stbi__at_eof(z->s)) return stbi__err("no SOF", "Corrupt JPEG"); - m = stbi__get_marker(z); - } - } - z->progressive = stbi__SOF_progressive(m); - if (!stbi__process_frame_header(z, scan)) return 0; - return 1; -} - -static stbi_uc stbi__skip_jpeg_junk_at_end(stbi__jpeg *j) -{ - // some JPEGs have junk at end, skip over it but if we find what looks - // like a valid marker, resume there - while (!stbi__at_eof(j->s)) { - stbi_uc x = stbi__get8(j->s); - while (x == 0xff) { // might be a marker - if (stbi__at_eof(j->s)) return STBI__MARKER_none; - x = stbi__get8(j->s); - if (x != 0x00 && x != 0xff) { - // not a stuffed zero or lead-in to another marker, looks - // like an actual marker, return it - return x; - } - // stuffed zero has x=0 now which ends the loop, meaning we go - // back to regular scan loop. - // repeated 0xff keeps trying to read the next byte of the marker. - } - } - return STBI__MARKER_none; -} - -// decode image to YCbCr format -static int stbi__decode_jpeg_image(stbi__jpeg *j) -{ - int m; - for (m = 0; m < 4; m++) { - j->img_comp[m].raw_data = NULL; - j->img_comp[m].raw_coeff = NULL; - } - j->restart_interval = 0; - if (!stbi__decode_jpeg_header(j, STBI__SCAN_load)) return 0; - m = stbi__get_marker(j); - while (!stbi__EOI(m)) { - if (stbi__SOS(m)) { - if (!stbi__process_scan_header(j)) return 0; - if (!stbi__parse_entropy_coded_data(j)) return 0; - if (j->marker == STBI__MARKER_none ) { - j->marker = stbi__skip_jpeg_junk_at_end(j); - // if we reach eof without hitting a marker, stbi__get_marker() below will fail and we'll eventually return 0 - } - m = stbi__get_marker(j); - if (STBI__RESTART(m)) - m = stbi__get_marker(j); - } else if (stbi__DNL(m)) { - int Ld = stbi__get16be(j->s); - stbi__uint32 NL = stbi__get16be(j->s); - if (Ld != 4) return stbi__err("bad DNL len", "Corrupt JPEG"); - if (NL != j->s->img_y) return stbi__err("bad DNL height", "Corrupt JPEG"); - m = stbi__get_marker(j); - } else { - if (!stbi__process_marker(j, m)) return 1; - m = stbi__get_marker(j); - } - } - if (j->progressive) - stbi__jpeg_finish(j); - return 1; -} - -// static jfif-centered resampling (across block boundaries) - -typedef stbi_uc *(*resample_row_func)(stbi_uc *out, stbi_uc *in0, stbi_uc *in1, - int w, int hs); - -#define stbi__div4(x) ((stbi_uc) ((x) >> 2)) - -static stbi_uc *resample_row_1(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - STBI_NOTUSED(out); - STBI_NOTUSED(in_far); - STBI_NOTUSED(w); - STBI_NOTUSED(hs); - return in_near; -} - -static stbi_uc* stbi__resample_row_v_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate two samples vertically for every one in input - int i; - STBI_NOTUSED(hs); - for (i=0; i < w; ++i) - out[i] = stbi__div4(3*in_near[i] + in_far[i] + 2); - return out; -} - -static stbi_uc* stbi__resample_row_h_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate two samples horizontally for every one in input - int i; - stbi_uc *input = in_near; - - if (w == 1) { - // if only one sample, can't do any interpolation - out[0] = out[1] = input[0]; - return out; - } - - out[0] = input[0]; - out[1] = stbi__div4(input[0]*3 + input[1] + 2); - for (i=1; i < w-1; ++i) { - int n = 3*input[i]+2; - out[i*2+0] = stbi__div4(n+input[i-1]); - out[i*2+1] = stbi__div4(n+input[i+1]); - } - out[i*2+0] = stbi__div4(input[w-2]*3 + input[w-1] + 2); - out[i*2+1] = input[w-1]; - - STBI_NOTUSED(in_far); - STBI_NOTUSED(hs); - - return out; -} - -#define stbi__div16(x) ((stbi_uc) ((x) >> 4)) - -static stbi_uc *stbi__resample_row_hv_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate 2x2 samples for every one in input - int i,t0,t1; - if (w == 1) { - out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); - return out; - } - - t1 = 3*in_near[0] + in_far[0]; - out[0] = stbi__div4(t1+2); - for (i=1; i < w; ++i) { - t0 = t1; - t1 = 3*in_near[i]+in_far[i]; - out[i*2-1] = stbi__div16(3*t0 + t1 + 8); - out[i*2 ] = stbi__div16(3*t1 + t0 + 8); - } - out[w*2-1] = stbi__div4(t1+2); - - STBI_NOTUSED(hs); - - return out; -} - -#if defined(STBI_SSE2) || defined(STBI_NEON) -static stbi_uc *stbi__resample_row_hv_2_simd(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate 2x2 samples for every one in input - int i=0,t0,t1; - - if (w == 1) { - out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); - return out; - } - - t1 = 3*in_near[0] + in_far[0]; - // process groups of 8 pixels for as long as we can. - // note we can't handle the last pixel in a row in this loop - // because we need to handle the filter boundary conditions. - for (; i < ((w-1) & ~7); i += 8) { -#if defined(STBI_SSE2) - // load and perform the vertical filtering pass - // this uses 3*x + y = 4*x + (y - x) - __m128i zero = _mm_setzero_si128(); - __m128i farb = _mm_loadl_epi64((__m128i *) (in_far + i)); - __m128i nearb = _mm_loadl_epi64((__m128i *) (in_near + i)); - __m128i farw = _mm_unpacklo_epi8(farb, zero); - __m128i nearw = _mm_unpacklo_epi8(nearb, zero); - __m128i diff = _mm_sub_epi16(farw, nearw); - __m128i nears = _mm_slli_epi16(nearw, 2); - __m128i curr = _mm_add_epi16(nears, diff); // current row - - // horizontal filter works the same based on shifted vers of current - // row. "prev" is current row shifted right by 1 pixel; we need to - // insert the previous pixel value (from t1). - // "next" is current row shifted left by 1 pixel, with first pixel - // of next block of 8 pixels added in. - __m128i prv0 = _mm_slli_si128(curr, 2); - __m128i nxt0 = _mm_srli_si128(curr, 2); - __m128i prev = _mm_insert_epi16(prv0, t1, 0); - __m128i next = _mm_insert_epi16(nxt0, 3*in_near[i+8] + in_far[i+8], 7); - - // horizontal filter, polyphase implementation since it's convenient: - // even pixels = 3*cur + prev = cur*4 + (prev - cur) - // odd pixels = 3*cur + next = cur*4 + (next - cur) - // note the shared term. - __m128i bias = _mm_set1_epi16(8); - __m128i curs = _mm_slli_epi16(curr, 2); - __m128i prvd = _mm_sub_epi16(prev, curr); - __m128i nxtd = _mm_sub_epi16(next, curr); - __m128i curb = _mm_add_epi16(curs, bias); - __m128i even = _mm_add_epi16(prvd, curb); - __m128i odd = _mm_add_epi16(nxtd, curb); - - // interleave even and odd pixels, then undo scaling. - __m128i int0 = _mm_unpacklo_epi16(even, odd); - __m128i int1 = _mm_unpackhi_epi16(even, odd); - __m128i de0 = _mm_srli_epi16(int0, 4); - __m128i de1 = _mm_srli_epi16(int1, 4); - - // pack and write output - __m128i outv = _mm_packus_epi16(de0, de1); - _mm_storeu_si128((__m128i *) (out + i*2), outv); -#elif defined(STBI_NEON) - // load and perform the vertical filtering pass - // this uses 3*x + y = 4*x + (y - x) - uint8x8_t farb = vld1_u8(in_far + i); - uint8x8_t nearb = vld1_u8(in_near + i); - int16x8_t diff = vreinterpretq_s16_u16(vsubl_u8(farb, nearb)); - int16x8_t nears = vreinterpretq_s16_u16(vshll_n_u8(nearb, 2)); - int16x8_t curr = vaddq_s16(nears, diff); // current row - - // horizontal filter works the same based on shifted vers of current - // row. "prev" is current row shifted right by 1 pixel; we need to - // insert the previous pixel value (from t1). - // "next" is current row shifted left by 1 pixel, with first pixel - // of next block of 8 pixels added in. - int16x8_t prv0 = vextq_s16(curr, curr, 7); - int16x8_t nxt0 = vextq_s16(curr, curr, 1); - int16x8_t prev = vsetq_lane_s16(t1, prv0, 0); - int16x8_t next = vsetq_lane_s16(3*in_near[i+8] + in_far[i+8], nxt0, 7); - - // horizontal filter, polyphase implementation since it's convenient: - // even pixels = 3*cur + prev = cur*4 + (prev - cur) - // odd pixels = 3*cur + next = cur*4 + (next - cur) - // note the shared term. - int16x8_t curs = vshlq_n_s16(curr, 2); - int16x8_t prvd = vsubq_s16(prev, curr); - int16x8_t nxtd = vsubq_s16(next, curr); - int16x8_t even = vaddq_s16(curs, prvd); - int16x8_t odd = vaddq_s16(curs, nxtd); - - // undo scaling and round, then store with even/odd phases interleaved - uint8x8x2_t o; - o.val[0] = vqrshrun_n_s16(even, 4); - o.val[1] = vqrshrun_n_s16(odd, 4); - vst2_u8(out + i*2, o); -#endif - - // "previous" value for next iter - t1 = 3*in_near[i+7] + in_far[i+7]; - } - - t0 = t1; - t1 = 3*in_near[i] + in_far[i]; - out[i*2] = stbi__div16(3*t1 + t0 + 8); - - for (++i; i < w; ++i) { - t0 = t1; - t1 = 3*in_near[i]+in_far[i]; - out[i*2-1] = stbi__div16(3*t0 + t1 + 8); - out[i*2 ] = stbi__div16(3*t1 + t0 + 8); - } - out[w*2-1] = stbi__div4(t1+2); - - STBI_NOTUSED(hs); - - return out; -} -#endif - -static stbi_uc *stbi__resample_row_generic(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // resample with nearest-neighbor - int i,j; - STBI_NOTUSED(in_far); - for (i=0; i < w; ++i) - for (j=0; j < hs; ++j) - out[i*hs+j] = in_near[i]; - return out; -} - -// this is a reduced-precision calculation of YCbCr-to-RGB introduced -// to make sure the code produces the same results in both SIMD and scalar -#define stbi__float2fixed(x) (((int) ((x) * 4096.0f + 0.5f)) << 8) -static void stbi__YCbCr_to_RGB_row(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step) -{ - int i; - for (i=0; i < count; ++i) { - int y_fixed = (y[i] << 20) + (1<<19); // rounding - int r,g,b; - int cr = pcr[i] - 128; - int cb = pcb[i] - 128; - r = y_fixed + cr* stbi__float2fixed(1.40200f); - g = y_fixed + (cr*-stbi__float2fixed(0.71414f)) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); - b = y_fixed + cb* stbi__float2fixed(1.77200f); - r >>= 20; - g >>= 20; - b >>= 20; - if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } - if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } - if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } - out[0] = (stbi_uc)r; - out[1] = (stbi_uc)g; - out[2] = (stbi_uc)b; - out[3] = 255; - out += step; - } -} - -#if defined(STBI_SSE2) || defined(STBI_NEON) -static void stbi__YCbCr_to_RGB_simd(stbi_uc *out, stbi_uc const *y, stbi_uc const *pcb, stbi_uc const *pcr, int count, int step) -{ - int i = 0; - -#ifdef STBI_SSE2 - // step == 3 is pretty ugly on the final interleave, and i'm not convinced - // it's useful in practice (you wouldn't use it for textures, for example). - // so just accelerate step == 4 case. - if (step == 4) { - // this is a fairly straightforward implementation and not super-optimized. - __m128i signflip = _mm_set1_epi8(-0x80); - __m128i cr_const0 = _mm_set1_epi16( (short) ( 1.40200f*4096.0f+0.5f)); - __m128i cr_const1 = _mm_set1_epi16( - (short) ( 0.71414f*4096.0f+0.5f)); - __m128i cb_const0 = _mm_set1_epi16( - (short) ( 0.34414f*4096.0f+0.5f)); - __m128i cb_const1 = _mm_set1_epi16( (short) ( 1.77200f*4096.0f+0.5f)); - __m128i y_bias = _mm_set1_epi8((char) (unsigned char) 128); - __m128i xw = _mm_set1_epi16(255); // alpha channel - - for (; i+7 < count; i += 8) { - // load - __m128i y_bytes = _mm_loadl_epi64((__m128i *) (y+i)); - __m128i cr_bytes = _mm_loadl_epi64((__m128i *) (pcr+i)); - __m128i cb_bytes = _mm_loadl_epi64((__m128i *) (pcb+i)); - __m128i cr_biased = _mm_xor_si128(cr_bytes, signflip); // -128 - __m128i cb_biased = _mm_xor_si128(cb_bytes, signflip); // -128 - - // unpack to short (and left-shift cr, cb by 8) - __m128i yw = _mm_unpacklo_epi8(y_bias, y_bytes); - __m128i crw = _mm_unpacklo_epi8(_mm_setzero_si128(), cr_biased); - __m128i cbw = _mm_unpacklo_epi8(_mm_setzero_si128(), cb_biased); - - // color transform - __m128i yws = _mm_srli_epi16(yw, 4); - __m128i cr0 = _mm_mulhi_epi16(cr_const0, crw); - __m128i cb0 = _mm_mulhi_epi16(cb_const0, cbw); - __m128i cb1 = _mm_mulhi_epi16(cbw, cb_const1); - __m128i cr1 = _mm_mulhi_epi16(crw, cr_const1); - __m128i rws = _mm_add_epi16(cr0, yws); - __m128i gwt = _mm_add_epi16(cb0, yws); - __m128i bws = _mm_add_epi16(yws, cb1); - __m128i gws = _mm_add_epi16(gwt, cr1); - - // descale - __m128i rw = _mm_srai_epi16(rws, 4); - __m128i bw = _mm_srai_epi16(bws, 4); - __m128i gw = _mm_srai_epi16(gws, 4); - - // back to byte, set up for transpose - __m128i brb = _mm_packus_epi16(rw, bw); - __m128i gxb = _mm_packus_epi16(gw, xw); - - // transpose to interleave channels - __m128i t0 = _mm_unpacklo_epi8(brb, gxb); - __m128i t1 = _mm_unpackhi_epi8(brb, gxb); - __m128i o0 = _mm_unpacklo_epi16(t0, t1); - __m128i o1 = _mm_unpackhi_epi16(t0, t1); - - // store - _mm_storeu_si128((__m128i *) (out + 0), o0); - _mm_storeu_si128((__m128i *) (out + 16), o1); - out += 32; - } - } -#endif - -#ifdef STBI_NEON - // in this version, step=3 support would be easy to add. but is there demand? - if (step == 4) { - // this is a fairly straightforward implementation and not super-optimized. - uint8x8_t signflip = vdup_n_u8(0x80); - int16x8_t cr_const0 = vdupq_n_s16( (short) ( 1.40200f*4096.0f+0.5f)); - int16x8_t cr_const1 = vdupq_n_s16( - (short) ( 0.71414f*4096.0f+0.5f)); - int16x8_t cb_const0 = vdupq_n_s16( - (short) ( 0.34414f*4096.0f+0.5f)); - int16x8_t cb_const1 = vdupq_n_s16( (short) ( 1.77200f*4096.0f+0.5f)); - - for (; i+7 < count; i += 8) { - // load - uint8x8_t y_bytes = vld1_u8(y + i); - uint8x8_t cr_bytes = vld1_u8(pcr + i); - uint8x8_t cb_bytes = vld1_u8(pcb + i); - int8x8_t cr_biased = vreinterpret_s8_u8(vsub_u8(cr_bytes, signflip)); - int8x8_t cb_biased = vreinterpret_s8_u8(vsub_u8(cb_bytes, signflip)); - - // expand to s16 - int16x8_t yws = vreinterpretq_s16_u16(vshll_n_u8(y_bytes, 4)); - int16x8_t crw = vshll_n_s8(cr_biased, 7); - int16x8_t cbw = vshll_n_s8(cb_biased, 7); - - // color transform - int16x8_t cr0 = vqdmulhq_s16(crw, cr_const0); - int16x8_t cb0 = vqdmulhq_s16(cbw, cb_const0); - int16x8_t cr1 = vqdmulhq_s16(crw, cr_const1); - int16x8_t cb1 = vqdmulhq_s16(cbw, cb_const1); - int16x8_t rws = vaddq_s16(yws, cr0); - int16x8_t gws = vaddq_s16(vaddq_s16(yws, cb0), cr1); - int16x8_t bws = vaddq_s16(yws, cb1); - - // undo scaling, round, convert to byte - uint8x8x4_t o; - o.val[0] = vqrshrun_n_s16(rws, 4); - o.val[1] = vqrshrun_n_s16(gws, 4); - o.val[2] = vqrshrun_n_s16(bws, 4); - o.val[3] = vdup_n_u8(255); - - // store, interleaving r/g/b/a - vst4_u8(out, o); - out += 8*4; - } - } -#endif - - for (; i < count; ++i) { - int y_fixed = (y[i] << 20) + (1<<19); // rounding - int r,g,b; - int cr = pcr[i] - 128; - int cb = pcb[i] - 128; - r = y_fixed + cr* stbi__float2fixed(1.40200f); - g = y_fixed + cr*-stbi__float2fixed(0.71414f) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); - b = y_fixed + cb* stbi__float2fixed(1.77200f); - r >>= 20; - g >>= 20; - b >>= 20; - if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } - if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } - if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } - out[0] = (stbi_uc)r; - out[1] = (stbi_uc)g; - out[2] = (stbi_uc)b; - out[3] = 255; - out += step; - } -} -#endif - -// set up the kernels -static void stbi__setup_jpeg(stbi__jpeg *j) -{ - j->idct_block_kernel = stbi__idct_block; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_row; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2; - -#ifdef STBI_SSE2 - if (stbi__sse2_available()) { - j->idct_block_kernel = stbi__idct_simd; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; - } -#endif - -#ifdef STBI_NEON - j->idct_block_kernel = stbi__idct_simd; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; -#endif -} - -// clean up the temporary component buffers -static void stbi__cleanup_jpeg(stbi__jpeg *j) -{ - stbi__free_jpeg_components(j, j->s->img_n, 0); -} - -typedef struct -{ - resample_row_func resample; - stbi_uc *line0,*line1; - int hs,vs; // expansion factor in each axis - int w_lores; // horizontal pixels pre-expansion - int ystep; // how far through vertical expansion we are - int ypos; // which pre-expansion row we're on -} stbi__resample; - -// fast 0..255 * 0..255 => 0..255 rounded multiplication -static stbi_uc stbi__blinn_8x8(stbi_uc x, stbi_uc y) -{ - unsigned int t = x*y + 128; - return (stbi_uc) ((t + (t >>8)) >> 8); -} - -static stbi_uc *load_jpeg_image(stbi__jpeg *z, int *out_x, int *out_y, int *comp, int req_comp) -{ - int n, decode_n, is_rgb; - z->s->img_n = 0; // make stbi__cleanup_jpeg safe - - // validate req_comp - if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); - - // load a jpeg image from whichever source, but leave in YCbCr format - if (!stbi__decode_jpeg_image(z)) { stbi__cleanup_jpeg(z); return NULL; } - - // determine actual number of components to generate - n = req_comp ? req_comp : z->s->img_n >= 3 ? 3 : 1; - - is_rgb = z->s->img_n == 3 && (z->rgb == 3 || (z->app14_color_transform == 0 && !z->jfif)); - - if (z->s->img_n == 3 && n < 3 && !is_rgb) - decode_n = 1; - else - decode_n = z->s->img_n; - - // nothing to do if no components requested; check this now to avoid - // accessing uninitialized coutput[0] later - if (decode_n <= 0) { stbi__cleanup_jpeg(z); return NULL; } - - // resample and color-convert - { - int k; - unsigned int i,j; - stbi_uc *output; - stbi_uc *coutput[4] = { NULL, NULL, NULL, NULL }; - - stbi__resample res_comp[4]; - - for (k=0; k < decode_n; ++k) { - stbi__resample *r = &res_comp[k]; - - // allocate line buffer big enough for upsampling off the edges - // with upsample factor of 4 - z->img_comp[k].linebuf = (stbi_uc *) stbi__malloc(z->s->img_x + 3); - if (!z->img_comp[k].linebuf) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } - - r->hs = z->img_h_max / z->img_comp[k].h; - r->vs = z->img_v_max / z->img_comp[k].v; - r->ystep = r->vs >> 1; - r->w_lores = (z->s->img_x + r->hs-1) / r->hs; - r->ypos = 0; - r->line0 = r->line1 = z->img_comp[k].data; - - if (r->hs == 1 && r->vs == 1) r->resample = resample_row_1; - else if (r->hs == 1 && r->vs == 2) r->resample = stbi__resample_row_v_2; - else if (r->hs == 2 && r->vs == 1) r->resample = stbi__resample_row_h_2; - else if (r->hs == 2 && r->vs == 2) r->resample = z->resample_row_hv_2_kernel; - else r->resample = stbi__resample_row_generic; - } - - // can't error after this so, this is safe - output = (stbi_uc *) stbi__malloc_mad3(n, z->s->img_x, z->s->img_y, 1); - if (!output) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } - - // now go ahead and resample - for (j=0; j < z->s->img_y; ++j) { - stbi_uc *out = output + n * z->s->img_x * j; - for (k=0; k < decode_n; ++k) { - stbi__resample *r = &res_comp[k]; - int y_bot = r->ystep >= (r->vs >> 1); - coutput[k] = r->resample(z->img_comp[k].linebuf, - y_bot ? r->line1 : r->line0, - y_bot ? r->line0 : r->line1, - r->w_lores, r->hs); - if (++r->ystep >= r->vs) { - r->ystep = 0; - r->line0 = r->line1; - if (++r->ypos < z->img_comp[k].y) - r->line1 += z->img_comp[k].w2; - } - } - if (n >= 3) { - stbi_uc *y = coutput[0]; - if (z->s->img_n == 3) { - if (is_rgb) { - for (i=0; i < z->s->img_x; ++i) { - out[0] = y[i]; - out[1] = coutput[1][i]; - out[2] = coutput[2][i]; - out[3] = 255; - out += n; - } - } else { - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - } - } else if (z->s->img_n == 4) { - if (z->app14_color_transform == 0) { // CMYK - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - out[0] = stbi__blinn_8x8(coutput[0][i], m); - out[1] = stbi__blinn_8x8(coutput[1][i], m); - out[2] = stbi__blinn_8x8(coutput[2][i], m); - out[3] = 255; - out += n; - } - } else if (z->app14_color_transform == 2) { // YCCK - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - out[0] = stbi__blinn_8x8(255 - out[0], m); - out[1] = stbi__blinn_8x8(255 - out[1], m); - out[2] = stbi__blinn_8x8(255 - out[2], m); - out += n; - } - } else { // YCbCr + alpha? Ignore the fourth channel for now - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - } - } else - for (i=0; i < z->s->img_x; ++i) { - out[0] = out[1] = out[2] = y[i]; - out[3] = 255; // not used if n==3 - out += n; - } - } else { - if (is_rgb) { - if (n == 1) - for (i=0; i < z->s->img_x; ++i) - *out++ = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); - else { - for (i=0; i < z->s->img_x; ++i, out += 2) { - out[0] = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); - out[1] = 255; - } - } - } else if (z->s->img_n == 4 && z->app14_color_transform == 0) { - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - stbi_uc r = stbi__blinn_8x8(coutput[0][i], m); - stbi_uc g = stbi__blinn_8x8(coutput[1][i], m); - stbi_uc b = stbi__blinn_8x8(coutput[2][i], m); - out[0] = stbi__compute_y(r, g, b); - out[1] = 255; - out += n; - } - } else if (z->s->img_n == 4 && z->app14_color_transform == 2) { - for (i=0; i < z->s->img_x; ++i) { - out[0] = stbi__blinn_8x8(255 - coutput[0][i], coutput[3][i]); - out[1] = 255; - out += n; - } - } else { - stbi_uc *y = coutput[0]; - if (n == 1) - for (i=0; i < z->s->img_x; ++i) out[i] = y[i]; - else - for (i=0; i < z->s->img_x; ++i) { *out++ = y[i]; *out++ = 255; } - } - } - } - stbi__cleanup_jpeg(z); - *out_x = z->s->img_x; - *out_y = z->s->img_y; - if (comp) *comp = z->s->img_n >= 3 ? 3 : 1; // report original components, not output - return output; - } -} - -static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - unsigned char* result; - stbi__jpeg* j = (stbi__jpeg*) stbi__malloc(sizeof(stbi__jpeg)); - if (!j) return stbi__errpuc("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - STBI_NOTUSED(ri); - j->s = s; - stbi__setup_jpeg(j); - result = load_jpeg_image(j, x,y,comp,req_comp); - STBI_FREE(j); - return result; -} - -static int stbi__jpeg_test(stbi__context *s) -{ - int r; - stbi__jpeg* j = (stbi__jpeg*)stbi__malloc(sizeof(stbi__jpeg)); - if (!j) return stbi__err("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - j->s = s; - stbi__setup_jpeg(j); - r = stbi__decode_jpeg_header(j, STBI__SCAN_type); - stbi__rewind(s); - STBI_FREE(j); - return r; -} - -static int stbi__jpeg_info_raw(stbi__jpeg *j, int *x, int *y, int *comp) -{ - if (!stbi__decode_jpeg_header(j, STBI__SCAN_header)) { - stbi__rewind( j->s ); - return 0; - } - if (x) *x = j->s->img_x; - if (y) *y = j->s->img_y; - if (comp) *comp = j->s->img_n >= 3 ? 3 : 1; - return 1; -} - -static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp) -{ - int result; - stbi__jpeg* j = (stbi__jpeg*) (stbi__malloc(sizeof(stbi__jpeg))); - if (!j) return stbi__err("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - j->s = s; - result = stbi__jpeg_info_raw(j, x, y, comp); - STBI_FREE(j); - return result; -} -#endif - -// public domain zlib decode v0.2 Sean Barrett 2006-11-18 -// simple implementation -// - all input must be provided in an upfront buffer -// - all output is written to a single output buffer (can malloc/realloc) -// performance -// - fast huffman - -#ifndef STBI_NO_ZLIB - -// fast-way is faster to check than jpeg huffman, but slow way is slower -#define STBI__ZFAST_BITS 9 // accelerate all cases in default tables -#define STBI__ZFAST_MASK ((1 << STBI__ZFAST_BITS) - 1) -#define STBI__ZNSYMS 288 // number of symbols in literal/length alphabet - -// zlib-style huffman encoding -// (jpegs packs from left, zlib from right, so can't share code) -typedef struct -{ - stbi__uint16 fast[1 << STBI__ZFAST_BITS]; - stbi__uint16 firstcode[16]; - int maxcode[17]; - stbi__uint16 firstsymbol[16]; - stbi_uc size[STBI__ZNSYMS]; - stbi__uint16 value[STBI__ZNSYMS]; -} stbi__zhuffman; - -stbi_inline static int stbi__bitreverse16(int n) -{ - n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1); - n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2); - n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4); - n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8); - return n; -} - -stbi_inline static int stbi__bit_reverse(int v, int bits) -{ - STBI_ASSERT(bits <= 16); - // to bit reverse n bits, reverse 16 and shift - // e.g. 11 bits, bit reverse and shift away 5 - return stbi__bitreverse16(v) >> (16-bits); -} - -static int stbi__zbuild_huffman(stbi__zhuffman *z, const stbi_uc *sizelist, int num) -{ - int i,k=0; - int code, next_code[16], sizes[17]; - - // DEFLATE spec for generating codes - memset(sizes, 0, sizeof(sizes)); - memset(z->fast, 0, sizeof(z->fast)); - for (i=0; i < num; ++i) - ++sizes[sizelist[i]]; - sizes[0] = 0; - for (i=1; i < 16; ++i) - if (sizes[i] > (1 << i)) - return stbi__err("bad sizes", "Corrupt PNG"); - code = 0; - for (i=1; i < 16; ++i) { - next_code[i] = code; - z->firstcode[i] = (stbi__uint16) code; - z->firstsymbol[i] = (stbi__uint16) k; - code = (code + sizes[i]); - if (sizes[i]) - if (code-1 >= (1 << i)) return stbi__err("bad codelengths","Corrupt PNG"); - z->maxcode[i] = code << (16-i); // preshift for inner loop - code <<= 1; - k += sizes[i]; - } - z->maxcode[16] = 0x10000; // sentinel - for (i=0; i < num; ++i) { - int s = sizelist[i]; - if (s) { - int c = next_code[s] - z->firstcode[s] + z->firstsymbol[s]; - stbi__uint16 fastv = (stbi__uint16) ((s << 9) | i); - z->size [c] = (stbi_uc ) s; - z->value[c] = (stbi__uint16) i; - if (s <= STBI__ZFAST_BITS) { - int j = stbi__bit_reverse(next_code[s],s); - while (j < (1 << STBI__ZFAST_BITS)) { - z->fast[j] = fastv; - j += (1 << s); - } - } - ++next_code[s]; - } - } - return 1; -} - -// zlib-from-memory implementation for PNG reading -// because PNG allows splitting the zlib stream arbitrarily, -// and it's annoying structurally to have PNG call ZLIB call PNG, -// we require PNG read all the IDATs and combine them into a single -// memory buffer - -typedef struct -{ - stbi_uc *zbuffer, *zbuffer_end; - int num_bits; - int hit_zeof_once; - stbi__uint32 code_buffer; - - char *zout; - char *zout_start; - char *zout_end; - int z_expandable; - - stbi__zhuffman z_length, z_distance; -} stbi__zbuf; - -stbi_inline static int stbi__zeof(stbi__zbuf *z) -{ - return (z->zbuffer >= z->zbuffer_end); -} - -stbi_inline static stbi_uc stbi__zget8(stbi__zbuf *z) -{ - return stbi__zeof(z) ? 0 : *z->zbuffer++; -} - -static void stbi__fill_bits(stbi__zbuf *z) -{ - do { - if (z->code_buffer >= (1U << z->num_bits)) { - z->zbuffer = z->zbuffer_end; /* treat this as EOF so we fail. */ - return; - } - z->code_buffer |= (unsigned int) stbi__zget8(z) << z->num_bits; - z->num_bits += 8; - } while (z->num_bits <= 24); -} - -stbi_inline static unsigned int stbi__zreceive(stbi__zbuf *z, int n) -{ - unsigned int k; - if (z->num_bits < n) stbi__fill_bits(z); - k = z->code_buffer & ((1 << n) - 1); - z->code_buffer >>= n; - z->num_bits -= n; - return k; -} - -static int stbi__zhuffman_decode_slowpath(stbi__zbuf *a, stbi__zhuffman *z) -{ - int b,s,k; - // not resolved by fast table, so compute it the slow way - // use jpeg approach, which requires MSbits at top - k = stbi__bit_reverse(a->code_buffer, 16); - for (s=STBI__ZFAST_BITS+1; ; ++s) - if (k < z->maxcode[s]) - break; - if (s >= 16) return -1; // invalid code! - // code size is s, so: - b = (k >> (16-s)) - z->firstcode[s] + z->firstsymbol[s]; - if (b >= STBI__ZNSYMS) return -1; // some data was corrupt somewhere! - if (z->size[b] != s) return -1; // was originally an assert, but report failure instead. - a->code_buffer >>= s; - a->num_bits -= s; - return z->value[b]; -} - -stbi_inline static int stbi__zhuffman_decode(stbi__zbuf *a, stbi__zhuffman *z) -{ - int b,s; - if (a->num_bits < 16) { - if (stbi__zeof(a)) { - if (!a->hit_zeof_once) { - // This is the first time we hit eof, insert 16 extra padding btis - // to allow us to keep going; if we actually consume any of them - // though, that is invalid data. This is caught later. - a->hit_zeof_once = 1; - a->num_bits += 16; // add 16 implicit zero bits - } else { - // We already inserted our extra 16 padding bits and are again - // out, this stream is actually prematurely terminated. - return -1; - } - } else { - stbi__fill_bits(a); - } - } - b = z->fast[a->code_buffer & STBI__ZFAST_MASK]; - if (b) { - s = b >> 9; - a->code_buffer >>= s; - a->num_bits -= s; - return b & 511; - } - return stbi__zhuffman_decode_slowpath(a, z); -} - -static int stbi__zexpand(stbi__zbuf *z, char *zout, int n) // need to make room for n bytes -{ - char *q; - unsigned int cur, limit, old_limit; - z->zout = zout; - if (!z->z_expandable) return stbi__err("output buffer limit","Corrupt PNG"); - cur = (unsigned int) (z->zout - z->zout_start); - limit = old_limit = (unsigned) (z->zout_end - z->zout_start); - if (UINT_MAX - cur < (unsigned) n) return stbi__err("outofmem", "Out of memory"); - while (cur + n > limit) { - if(limit > UINT_MAX / 2) return stbi__err("outofmem", "Out of memory"); - limit *= 2; - } - q = (char *) STBI_REALLOC_SIZED(z->zout_start, old_limit, limit); - STBI_NOTUSED(old_limit); - if (q == NULL) return stbi__err("outofmem", "Out of memory"); - z->zout_start = q; - z->zout = q + cur; - z->zout_end = q + limit; - return 1; -} - -static const int stbi__zlength_base[31] = { - 3,4,5,6,7,8,9,10,11,13, - 15,17,19,23,27,31,35,43,51,59, - 67,83,99,115,131,163,195,227,258,0,0 }; - -static const int stbi__zlength_extra[31]= -{ 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0,0,0 }; - -static const int stbi__zdist_base[32] = { 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, -257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577,0,0}; - -static const int stbi__zdist_extra[32] = -{ 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13}; - -static int stbi__parse_huffman_block(stbi__zbuf *a) -{ - char *zout = a->zout; - for(;;) { - int z = stbi__zhuffman_decode(a, &a->z_length); - if (z < 256) { - if (z < 0) return stbi__err("bad huffman code","Corrupt PNG"); // error in huffman codes - if (zout >= a->zout_end) { - if (!stbi__zexpand(a, zout, 1)) return 0; - zout = a->zout; - } - *zout++ = (char) z; - } else { - stbi_uc *p; - int len,dist; - if (z == 256) { - a->zout = zout; - if (a->hit_zeof_once && a->num_bits < 16) { - // The first time we hit zeof, we inserted 16 extra zero bits into our bit - // buffer so the decoder can just do its speculative decoding. But if we - // actually consumed any of those bits (which is the case when num_bits < 16), - // the stream actually read past the end so it is malformed. - return stbi__err("unexpected end","Corrupt PNG"); - } - return 1; - } - if (z >= 286) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, length codes 286 and 287 must not appear in compressed data - z -= 257; - len = stbi__zlength_base[z]; - if (stbi__zlength_extra[z]) len += stbi__zreceive(a, stbi__zlength_extra[z]); - z = stbi__zhuffman_decode(a, &a->z_distance); - if (z < 0 || z >= 30) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, distance codes 30 and 31 must not appear in compressed data - dist = stbi__zdist_base[z]; - if (stbi__zdist_extra[z]) dist += stbi__zreceive(a, stbi__zdist_extra[z]); - if (zout - a->zout_start < dist) return stbi__err("bad dist","Corrupt PNG"); - if (len > a->zout_end - zout) { - if (!stbi__zexpand(a, zout, len)) return 0; - zout = a->zout; - } - p = (stbi_uc *) (zout - dist); - if (dist == 1) { // run of one byte; common in images. - stbi_uc v = *p; - if (len) { do *zout++ = v; while (--len); } - } else { - if (len) { do *zout++ = *p++; while (--len); } - } - } - } -} - -static int stbi__compute_huffman_codes(stbi__zbuf *a) -{ - static const stbi_uc length_dezigzag[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; - stbi__zhuffman z_codelength; - stbi_uc lencodes[286+32+137];//padding for maximum single op - stbi_uc codelength_sizes[19]; - int i,n; - - int hlit = stbi__zreceive(a,5) + 257; - int hdist = stbi__zreceive(a,5) + 1; - int hclen = stbi__zreceive(a,4) + 4; - int ntot = hlit + hdist; - - memset(codelength_sizes, 0, sizeof(codelength_sizes)); - for (i=0; i < hclen; ++i) { - int s = stbi__zreceive(a,3); - codelength_sizes[length_dezigzag[i]] = (stbi_uc) s; - } - if (!stbi__zbuild_huffman(&z_codelength, codelength_sizes, 19)) return 0; - - n = 0; - while (n < ntot) { - int c = stbi__zhuffman_decode(a, &z_codelength); - if (c < 0 || c >= 19) return stbi__err("bad codelengths", "Corrupt PNG"); - if (c < 16) - lencodes[n++] = (stbi_uc) c; - else { - stbi_uc fill = 0; - if (c == 16) { - c = stbi__zreceive(a,2)+3; - if (n == 0) return stbi__err("bad codelengths", "Corrupt PNG"); - fill = lencodes[n-1]; - } else if (c == 17) { - c = stbi__zreceive(a,3)+3; - } else if (c == 18) { - c = stbi__zreceive(a,7)+11; - } else { - return stbi__err("bad codelengths", "Corrupt PNG"); - } - if (ntot - n < c) return stbi__err("bad codelengths", "Corrupt PNG"); - memset(lencodes+n, fill, c); - n += c; - } - } - if (n != ntot) return stbi__err("bad codelengths","Corrupt PNG"); - if (!stbi__zbuild_huffman(&a->z_length, lencodes, hlit)) return 0; - if (!stbi__zbuild_huffman(&a->z_distance, lencodes+hlit, hdist)) return 0; - return 1; -} - -static int stbi__parse_uncompressed_block(stbi__zbuf *a) -{ - stbi_uc header[4]; - int len,nlen,k; - if (a->num_bits & 7) - stbi__zreceive(a, a->num_bits & 7); // discard - // drain the bit-packed data into header - k = 0; - while (a->num_bits > 0) { - header[k++] = (stbi_uc) (a->code_buffer & 255); // suppress MSVC run-time check - a->code_buffer >>= 8; - a->num_bits -= 8; - } - if (a->num_bits < 0) return stbi__err("zlib corrupt","Corrupt PNG"); - // now fill header the normal way - while (k < 4) - header[k++] = stbi__zget8(a); - len = header[1] * 256 + header[0]; - nlen = header[3] * 256 + header[2]; - if (nlen != (len ^ 0xffff)) return stbi__err("zlib corrupt","Corrupt PNG"); - if (a->zbuffer + len > a->zbuffer_end) return stbi__err("read past buffer","Corrupt PNG"); - if (a->zout + len > a->zout_end) - if (!stbi__zexpand(a, a->zout, len)) return 0; - memcpy(a->zout, a->zbuffer, len); - a->zbuffer += len; - a->zout += len; - return 1; -} - -static int stbi__parse_zlib_header(stbi__zbuf *a) -{ - int cmf = stbi__zget8(a); - int cm = cmf & 15; - /* int cinfo = cmf >> 4; */ - int flg = stbi__zget8(a); - if (stbi__zeof(a)) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec - if ((cmf*256+flg) % 31 != 0) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec - if (flg & 32) return stbi__err("no preset dict","Corrupt PNG"); // preset dictionary not allowed in png - if (cm != 8) return stbi__err("bad compression","Corrupt PNG"); // DEFLATE required for png - // window = 1 << (8 + cinfo)... but who cares, we fully buffer output - return 1; -} - -static const stbi_uc stbi__zdefault_length[STBI__ZNSYMS] = -{ - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8 -}; -static const stbi_uc stbi__zdefault_distance[32] = -{ - 5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5 -}; -/* -Init algorithm: -{ - int i; // use <= to match clearly with spec - for (i=0; i <= 143; ++i) stbi__zdefault_length[i] = 8; - for ( ; i <= 255; ++i) stbi__zdefault_length[i] = 9; - for ( ; i <= 279; ++i) stbi__zdefault_length[i] = 7; - for ( ; i <= 287; ++i) stbi__zdefault_length[i] = 8; - - for (i=0; i <= 31; ++i) stbi__zdefault_distance[i] = 5; -} -*/ - -static int stbi__parse_zlib(stbi__zbuf *a, int parse_header) -{ - int final, type; - if (parse_header) - if (!stbi__parse_zlib_header(a)) return 0; - a->num_bits = 0; - a->code_buffer = 0; - a->hit_zeof_once = 0; - do { - final = stbi__zreceive(a,1); - type = stbi__zreceive(a,2); - if (type == 0) { - if (!stbi__parse_uncompressed_block(a)) return 0; - } else if (type == 3) { - return 0; - } else { - if (type == 1) { - // use fixed code lengths - if (!stbi__zbuild_huffman(&a->z_length , stbi__zdefault_length , STBI__ZNSYMS)) return 0; - if (!stbi__zbuild_huffman(&a->z_distance, stbi__zdefault_distance, 32)) return 0; - } else { - if (!stbi__compute_huffman_codes(a)) return 0; - } - if (!stbi__parse_huffman_block(a)) return 0; - } - } while (!final); - return 1; -} - -static int stbi__do_zlib(stbi__zbuf *a, char *obuf, int olen, int exp, int parse_header) -{ - a->zout_start = obuf; - a->zout = obuf; - a->zout_end = obuf + olen; - a->z_expandable = exp; - - return stbi__parse_zlib(a, parse_header); -} - -STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(initial_size); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer + len; - if (stbi__do_zlib(&a, p, initial_size, 1, 1)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF char *stbi_zlib_decode_malloc(char const *buffer, int len, int *outlen) -{ - return stbi_zlib_decode_malloc_guesssize(buffer, len, 16384, outlen); -} - -STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(initial_size); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer + len; - if (stbi__do_zlib(&a, p, initial_size, 1, parse_header)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, char const *ibuffer, int ilen) -{ - stbi__zbuf a; - a.zbuffer = (stbi_uc *) ibuffer; - a.zbuffer_end = (stbi_uc *) ibuffer + ilen; - if (stbi__do_zlib(&a, obuffer, olen, 0, 1)) - return (int) (a.zout - a.zout_start); - else - return -1; -} - -STBIDEF char *stbi_zlib_decode_noheader_malloc(char const *buffer, int len, int *outlen) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(16384); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer+len; - if (stbi__do_zlib(&a, p, 16384, 1, 0)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen) -{ - stbi__zbuf a; - a.zbuffer = (stbi_uc *) ibuffer; - a.zbuffer_end = (stbi_uc *) ibuffer + ilen; - if (stbi__do_zlib(&a, obuffer, olen, 0, 0)) - return (int) (a.zout - a.zout_start); - else - return -1; -} -#endif - -// public domain "baseline" PNG decoder v0.10 Sean Barrett 2006-11-18 -// simple implementation -// - only 8-bit samples -// - no CRC checking -// - allocates lots of intermediate memory -// - avoids problem of streaming data between subsystems -// - avoids explicit window management -// performance -// - uses stb_zlib, a PD zlib implementation with fast huffman decoding - -#ifndef STBI_NO_PNG -typedef struct -{ - stbi__uint32 length; - stbi__uint32 type; -} stbi__pngchunk; - -static stbi__pngchunk stbi__get_chunk_header(stbi__context *s) -{ - stbi__pngchunk c; - c.length = stbi__get32be(s); - c.type = stbi__get32be(s); - return c; -} - -static int stbi__check_png_header(stbi__context *s) -{ - static const stbi_uc png_sig[8] = { 137,80,78,71,13,10,26,10 }; - int i; - for (i=0; i < 8; ++i) - if (stbi__get8(s) != png_sig[i]) return stbi__err("bad png sig","Not a PNG"); - return 1; -} - -typedef struct -{ - stbi__context *s; - stbi_uc *idata, *expanded, *out; - int depth; -} stbi__png; - - -enum { - STBI__F_none=0, - STBI__F_sub=1, - STBI__F_up=2, - STBI__F_avg=3, - STBI__F_paeth=4, - // synthetic filter used for first scanline to avoid needing a dummy row of 0s - STBI__F_avg_first -}; - -static stbi_uc first_row_filter[5] = -{ - STBI__F_none, - STBI__F_sub, - STBI__F_none, - STBI__F_avg_first, - STBI__F_sub // Paeth with b=c=0 turns out to be equivalent to sub -}; - -static int stbi__paeth(int a, int b, int c) -{ - // This formulation looks very different from the reference in the PNG spec, but is - // actually equivalent and has favorable data dependencies and admits straightforward - // generation of branch-free code, which helps performance significantly. - int thresh = c*3 - (a + b); - int lo = a < b ? a : b; - int hi = a < b ? b : a; - int t0 = (hi <= thresh) ? lo : c; - int t1 = (thresh <= lo) ? hi : t0; - return t1; -} - -static const stbi_uc stbi__depth_scale_table[9] = { 0, 0xff, 0x55, 0, 0x11, 0,0,0, 0x01 }; - -// adds an extra all-255 alpha channel -// dest == src is legal -// img_n must be 1 or 3 -static void stbi__create_png_alpha_expand8(stbi_uc *dest, stbi_uc *src, stbi__uint32 x, int img_n) -{ - int i; - // must process data backwards since we allow dest==src - if (img_n == 1) { - for (i=x-1; i >= 0; --i) { - dest[i*2+1] = 255; - dest[i*2+0] = src[i]; - } - } else { - STBI_ASSERT(img_n == 3); - for (i=x-1; i >= 0; --i) { - dest[i*4+3] = 255; - dest[i*4+2] = src[i*3+2]; - dest[i*4+1] = src[i*3+1]; - dest[i*4+0] = src[i*3+0]; - } - } -} - -// create the png data from post-deflated data -static int stbi__create_png_image_raw(stbi__png *a, stbi_uc *raw, stbi__uint32 raw_len, int out_n, stbi__uint32 x, stbi__uint32 y, int depth, int color) -{ - int bytes = (depth == 16 ? 2 : 1); - stbi__context *s = a->s; - stbi__uint32 i,j,stride = x*out_n*bytes; - stbi__uint32 img_len, img_width_bytes; - stbi_uc *filter_buf; - int all_ok = 1; - int k; - int img_n = s->img_n; // copy it into a local for later - - int output_bytes = out_n*bytes; - int filter_bytes = img_n*bytes; - int width = x; - - STBI_ASSERT(out_n == s->img_n || out_n == s->img_n+1); - a->out = (stbi_uc *) stbi__malloc_mad3(x, y, output_bytes, 0); // extra bytes to write off the end into - if (!a->out) return stbi__err("outofmem", "Out of memory"); - - // note: error exits here don't need to clean up a->out individually, - // stbi__do_png always does on error. - if (!stbi__mad3sizes_valid(img_n, x, depth, 7)) return stbi__err("too large", "Corrupt PNG"); - img_width_bytes = (((img_n * x * depth) + 7) >> 3); - if (!stbi__mad2sizes_valid(img_width_bytes, y, img_width_bytes)) return stbi__err("too large", "Corrupt PNG"); - img_len = (img_width_bytes + 1) * y; - - // we used to check for exact match between raw_len and img_len on non-interlaced PNGs, - // but issue #276 reported a PNG in the wild that had extra data at the end (all zeros), - // so just check for raw_len < img_len always. - if (raw_len < img_len) return stbi__err("not enough pixels","Corrupt PNG"); - - // Allocate two scan lines worth of filter workspace buffer. - filter_buf = (stbi_uc *) stbi__malloc_mad2(img_width_bytes, 2, 0); - if (!filter_buf) return stbi__err("outofmem", "Out of memory"); - - // Filtering for low-bit-depth images - if (depth < 8) { - filter_bytes = 1; - width = img_width_bytes; - } - - for (j=0; j < y; ++j) { - // cur/prior filter buffers alternate - stbi_uc *cur = filter_buf + (j & 1)*img_width_bytes; - stbi_uc *prior = filter_buf + (~j & 1)*img_width_bytes; - stbi_uc *dest = a->out + stride*j; - int nk = width * filter_bytes; - int filter = *raw++; - - // check filter type - if (filter > 4) { - all_ok = stbi__err("invalid filter","Corrupt PNG"); - break; - } - - // if first row, use special filter that doesn't sample previous row - if (j == 0) filter = first_row_filter[filter]; - - // perform actual filtering - switch (filter) { - case STBI__F_none: - memcpy(cur, raw, nk); - break; - case STBI__F_sub: - memcpy(cur, raw, filter_bytes); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + cur[k-filter_bytes]); - break; - case STBI__F_up: - for (k = 0; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + prior[k]); - break; - case STBI__F_avg: - for (k = 0; k < filter_bytes; ++k) - cur[k] = STBI__BYTECAST(raw[k] + (prior[k]>>1)); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + ((prior[k] + cur[k-filter_bytes])>>1)); - break; - case STBI__F_paeth: - for (k = 0; k < filter_bytes; ++k) - cur[k] = STBI__BYTECAST(raw[k] + prior[k]); // prior[k] == stbi__paeth(0,prior[k],0) - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + stbi__paeth(cur[k-filter_bytes], prior[k], prior[k-filter_bytes])); - break; - case STBI__F_avg_first: - memcpy(cur, raw, filter_bytes); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + (cur[k-filter_bytes] >> 1)); - break; - } - - raw += nk; - - // expand decoded bits in cur to dest, also adding an extra alpha channel if desired - if (depth < 8) { - stbi_uc scale = (color == 0) ? stbi__depth_scale_table[depth] : 1; // scale grayscale values to 0..255 range - stbi_uc *in = cur; - stbi_uc *out = dest; - stbi_uc inb = 0; - stbi__uint32 nsmp = x*img_n; - - // expand bits to bytes first - if (depth == 4) { - for (i=0; i < nsmp; ++i) { - if ((i & 1) == 0) inb = *in++; - *out++ = scale * (inb >> 4); - inb <<= 4; - } - } else if (depth == 2) { - for (i=0; i < nsmp; ++i) { - if ((i & 3) == 0) inb = *in++; - *out++ = scale * (inb >> 6); - inb <<= 2; - } - } else { - STBI_ASSERT(depth == 1); - for (i=0; i < nsmp; ++i) { - if ((i & 7) == 0) inb = *in++; - *out++ = scale * (inb >> 7); - inb <<= 1; - } - } - - // insert alpha=255 values if desired - if (img_n != out_n) - stbi__create_png_alpha_expand8(dest, dest, x, img_n); - } else if (depth == 8) { - if (img_n == out_n) - memcpy(dest, cur, x*img_n); - else - stbi__create_png_alpha_expand8(dest, cur, x, img_n); - } else if (depth == 16) { - // convert the image data from big-endian to platform-native - stbi__uint16 *dest16 = (stbi__uint16*)dest; - stbi__uint32 nsmp = x*img_n; - - if (img_n == out_n) { - for (i = 0; i < nsmp; ++i, ++dest16, cur += 2) - *dest16 = (cur[0] << 8) | cur[1]; - } else { - STBI_ASSERT(img_n+1 == out_n); - if (img_n == 1) { - for (i = 0; i < x; ++i, dest16 += 2, cur += 2) { - dest16[0] = (cur[0] << 8) | cur[1]; - dest16[1] = 0xffff; - } - } else { - STBI_ASSERT(img_n == 3); - for (i = 0; i < x; ++i, dest16 += 4, cur += 6) { - dest16[0] = (cur[0] << 8) | cur[1]; - dest16[1] = (cur[2] << 8) | cur[3]; - dest16[2] = (cur[4] << 8) | cur[5]; - dest16[3] = 0xffff; - } - } - } - } - } - - STBI_FREE(filter_buf); - if (!all_ok) return 0; - - return 1; -} - -static int stbi__create_png_image(stbi__png *a, stbi_uc *image_data, stbi__uint32 image_data_len, int out_n, int depth, int color, int interlaced) -{ - int bytes = (depth == 16 ? 2 : 1); - int out_bytes = out_n * bytes; - stbi_uc *final; - int p; - if (!interlaced) - return stbi__create_png_image_raw(a, image_data, image_data_len, out_n, a->s->img_x, a->s->img_y, depth, color); - - // de-interlacing - final = (stbi_uc *) stbi__malloc_mad3(a->s->img_x, a->s->img_y, out_bytes, 0); - if (!final) return stbi__err("outofmem", "Out of memory"); - for (p=0; p < 7; ++p) { - int xorig[] = { 0,4,0,2,0,1,0 }; - int yorig[] = { 0,0,4,0,2,0,1 }; - int xspc[] = { 8,8,4,4,2,2,1 }; - int yspc[] = { 8,8,8,4,4,2,2 }; - int i,j,x,y; - // pass1_x[4] = 0, pass1_x[5] = 1, pass1_x[12] = 1 - x = (a->s->img_x - xorig[p] + xspc[p]-1) / xspc[p]; - y = (a->s->img_y - yorig[p] + yspc[p]-1) / yspc[p]; - if (x && y) { - stbi__uint32 img_len = ((((a->s->img_n * x * depth) + 7) >> 3) + 1) * y; - if (!stbi__create_png_image_raw(a, image_data, image_data_len, out_n, x, y, depth, color)) { - STBI_FREE(final); - return 0; - } - for (j=0; j < y; ++j) { - for (i=0; i < x; ++i) { - int out_y = j*yspc[p]+yorig[p]; - int out_x = i*xspc[p]+xorig[p]; - memcpy(final + out_y*a->s->img_x*out_bytes + out_x*out_bytes, - a->out + (j*x+i)*out_bytes, out_bytes); - } - } - STBI_FREE(a->out); - image_data += img_len; - image_data_len -= img_len; - } - } - a->out = final; - - return 1; -} - -static int stbi__compute_transparency(stbi__png *z, stbi_uc tc[3], int out_n) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi_uc *p = z->out; - - // compute color-based transparency, assuming we've - // already got 255 as the alpha value in the output - STBI_ASSERT(out_n == 2 || out_n == 4); - - if (out_n == 2) { - for (i=0; i < pixel_count; ++i) { - p[1] = (p[0] == tc[0] ? 0 : 255); - p += 2; - } - } else { - for (i=0; i < pixel_count; ++i) { - if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) - p[3] = 0; - p += 4; - } - } - return 1; -} - -static int stbi__compute_transparency16(stbi__png *z, stbi__uint16 tc[3], int out_n) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi__uint16 *p = (stbi__uint16*) z->out; - - // compute color-based transparency, assuming we've - // already got 65535 as the alpha value in the output - STBI_ASSERT(out_n == 2 || out_n == 4); - - if (out_n == 2) { - for (i = 0; i < pixel_count; ++i) { - p[1] = (p[0] == tc[0] ? 0 : 65535); - p += 2; - } - } else { - for (i = 0; i < pixel_count; ++i) { - if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) - p[3] = 0; - p += 4; - } - } - return 1; -} - -static int stbi__expand_png_palette(stbi__png *a, stbi_uc *palette, int len, int pal_img_n) -{ - stbi__uint32 i, pixel_count = a->s->img_x * a->s->img_y; - stbi_uc *p, *temp_out, *orig = a->out; - - p = (stbi_uc *) stbi__malloc_mad2(pixel_count, pal_img_n, 0); - if (p == NULL) return stbi__err("outofmem", "Out of memory"); - - // between here and free(out) below, exitting would leak - temp_out = p; - - if (pal_img_n == 3) { - for (i=0; i < pixel_count; ++i) { - int n = orig[i]*4; - p[0] = palette[n ]; - p[1] = palette[n+1]; - p[2] = palette[n+2]; - p += 3; - } - } else { - for (i=0; i < pixel_count; ++i) { - int n = orig[i]*4; - p[0] = palette[n ]; - p[1] = palette[n+1]; - p[2] = palette[n+2]; - p[3] = palette[n+3]; - p += 4; - } - } - STBI_FREE(a->out); - a->out = temp_out; - - STBI_NOTUSED(len); - - return 1; -} - -static int stbi__unpremultiply_on_load_global = 0; -static int stbi__de_iphone_flag_global = 0; - -STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply) -{ - stbi__unpremultiply_on_load_global = flag_true_if_should_unpremultiply; -} - -STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert) -{ - stbi__de_iphone_flag_global = flag_true_if_should_convert; -} - -#ifndef STBI_THREAD_LOCAL -#define stbi__unpremultiply_on_load stbi__unpremultiply_on_load_global -#define stbi__de_iphone_flag stbi__de_iphone_flag_global -#else -static STBI_THREAD_LOCAL int stbi__unpremultiply_on_load_local, stbi__unpremultiply_on_load_set; -static STBI_THREAD_LOCAL int stbi__de_iphone_flag_local, stbi__de_iphone_flag_set; - -STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply) -{ - stbi__unpremultiply_on_load_local = flag_true_if_should_unpremultiply; - stbi__unpremultiply_on_load_set = 1; -} - -STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert) -{ - stbi__de_iphone_flag_local = flag_true_if_should_convert; - stbi__de_iphone_flag_set = 1; -} - -#define stbi__unpremultiply_on_load (stbi__unpremultiply_on_load_set \ - ? stbi__unpremultiply_on_load_local \ - : stbi__unpremultiply_on_load_global) -#define stbi__de_iphone_flag (stbi__de_iphone_flag_set \ - ? stbi__de_iphone_flag_local \ - : stbi__de_iphone_flag_global) -#endif // STBI_THREAD_LOCAL - -static void stbi__de_iphone(stbi__png *z) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi_uc *p = z->out; - - if (s->img_out_n == 3) { // convert bgr to rgb - for (i=0; i < pixel_count; ++i) { - stbi_uc t = p[0]; - p[0] = p[2]; - p[2] = t; - p += 3; - } - } else { - STBI_ASSERT(s->img_out_n == 4); - if (stbi__unpremultiply_on_load) { - // convert bgr to rgb and unpremultiply - for (i=0; i < pixel_count; ++i) { - stbi_uc a = p[3]; - stbi_uc t = p[0]; - if (a) { - stbi_uc half = a / 2; - p[0] = (p[2] * 255 + half) / a; - p[1] = (p[1] * 255 + half) / a; - p[2] = ( t * 255 + half) / a; - } else { - p[0] = p[2]; - p[2] = t; - } - p += 4; - } - } else { - // convert bgr to rgb - for (i=0; i < pixel_count; ++i) { - stbi_uc t = p[0]; - p[0] = p[2]; - p[2] = t; - p += 4; - } - } - } -} - -#define STBI__PNG_TYPE(a,b,c,d) (((unsigned) (a) << 24) + ((unsigned) (b) << 16) + ((unsigned) (c) << 8) + (unsigned) (d)) - -static int stbi__parse_png_file(stbi__png *z, int scan, int req_comp) -{ - stbi_uc palette[1024], pal_img_n=0; - stbi_uc has_trans=0, tc[3]={0}; - stbi__uint16 tc16[3]; - stbi__uint32 ioff=0, idata_limit=0, i, pal_len=0; - int first=1,k,interlace=0, color=0, is_iphone=0; - stbi__context *s = z->s; - - z->expanded = NULL; - z->idata = NULL; - z->out = NULL; - - if (!stbi__check_png_header(s)) return 0; - - if (scan == STBI__SCAN_type) return 1; - - for (;;) { - stbi__pngchunk c = stbi__get_chunk_header(s); - switch (c.type) { - case STBI__PNG_TYPE('C','g','B','I'): - is_iphone = 1; - stbi__skip(s, c.length); - break; - case STBI__PNG_TYPE('I','H','D','R'): { - int comp,filter; - if (!first) return stbi__err("multiple IHDR","Corrupt PNG"); - first = 0; - if (c.length != 13) return stbi__err("bad IHDR len","Corrupt PNG"); - s->img_x = stbi__get32be(s); - s->img_y = stbi__get32be(s); - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - z->depth = stbi__get8(s); if (z->depth != 1 && z->depth != 2 && z->depth != 4 && z->depth != 8 && z->depth != 16) return stbi__err("1/2/4/8/16-bit only","PNG not supported: 1/2/4/8/16-bit only"); - color = stbi__get8(s); if (color > 6) return stbi__err("bad ctype","Corrupt PNG"); - if (color == 3 && z->depth == 16) return stbi__err("bad ctype","Corrupt PNG"); - if (color == 3) pal_img_n = 3; else if (color & 1) return stbi__err("bad ctype","Corrupt PNG"); - comp = stbi__get8(s); if (comp) return stbi__err("bad comp method","Corrupt PNG"); - filter= stbi__get8(s); if (filter) return stbi__err("bad filter method","Corrupt PNG"); - interlace = stbi__get8(s); if (interlace>1) return stbi__err("bad interlace method","Corrupt PNG"); - if (!s->img_x || !s->img_y) return stbi__err("0-pixel image","Corrupt PNG"); - if (!pal_img_n) { - s->img_n = (color & 2 ? 3 : 1) + (color & 4 ? 1 : 0); - if ((1 << 30) / s->img_x / s->img_n < s->img_y) return stbi__err("too large", "Image too large to decode"); - } else { - // if paletted, then pal_n is our final components, and - // img_n is # components to decompress/filter. - s->img_n = 1; - if ((1 << 30) / s->img_x / 4 < s->img_y) return stbi__err("too large","Corrupt PNG"); - } - // even with SCAN_header, have to scan to see if we have a tRNS - break; - } - - case STBI__PNG_TYPE('P','L','T','E'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (c.length > 256*3) return stbi__err("invalid PLTE","Corrupt PNG"); - pal_len = c.length / 3; - if (pal_len * 3 != c.length) return stbi__err("invalid PLTE","Corrupt PNG"); - for (i=0; i < pal_len; ++i) { - palette[i*4+0] = stbi__get8(s); - palette[i*4+1] = stbi__get8(s); - palette[i*4+2] = stbi__get8(s); - palette[i*4+3] = 255; - } - break; - } - - case STBI__PNG_TYPE('t','R','N','S'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (z->idata) return stbi__err("tRNS after IDAT","Corrupt PNG"); - if (pal_img_n) { - if (scan == STBI__SCAN_header) { s->img_n = 4; return 1; } - if (pal_len == 0) return stbi__err("tRNS before PLTE","Corrupt PNG"); - if (c.length > pal_len) return stbi__err("bad tRNS len","Corrupt PNG"); - pal_img_n = 4; - for (i=0; i < c.length; ++i) - palette[i*4+3] = stbi__get8(s); - } else { - if (!(s->img_n & 1)) return stbi__err("tRNS with alpha","Corrupt PNG"); - if (c.length != (stbi__uint32) s->img_n*2) return stbi__err("bad tRNS len","Corrupt PNG"); - has_trans = 1; - // non-paletted with tRNS = constant alpha. if header-scanning, we can stop now. - if (scan == STBI__SCAN_header) { ++s->img_n; return 1; } - if (z->depth == 16) { - for (k = 0; k < s->img_n && k < 3; ++k) // extra loop test to suppress false GCC warning - tc16[k] = (stbi__uint16)stbi__get16be(s); // copy the values as-is - } else { - for (k = 0; k < s->img_n && k < 3; ++k) - tc[k] = (stbi_uc)(stbi__get16be(s) & 255) * stbi__depth_scale_table[z->depth]; // non 8-bit images will be larger - } - } - break; - } - - case STBI__PNG_TYPE('I','D','A','T'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (pal_img_n && !pal_len) return stbi__err("no PLTE","Corrupt PNG"); - if (scan == STBI__SCAN_header) { - // header scan definitely stops at first IDAT - if (pal_img_n) - s->img_n = pal_img_n; - return 1; - } - if (c.length > (1u << 30)) return stbi__err("IDAT size limit", "IDAT section larger than 2^30 bytes"); - if ((int)(ioff + c.length) < (int)ioff) return 0; - if (ioff + c.length > idata_limit) { - stbi__uint32 idata_limit_old = idata_limit; - stbi_uc *p; - if (idata_limit == 0) idata_limit = c.length > 4096 ? c.length : 4096; - while (ioff + c.length > idata_limit) - idata_limit *= 2; - STBI_NOTUSED(idata_limit_old); - p = (stbi_uc *) STBI_REALLOC_SIZED(z->idata, idata_limit_old, idata_limit); if (p == NULL) return stbi__err("outofmem", "Out of memory"); - z->idata = p; - } - if (!stbi__getn(s, z->idata+ioff,c.length)) return stbi__err("outofdata","Corrupt PNG"); - ioff += c.length; - break; - } - - case STBI__PNG_TYPE('I','E','N','D'): { - stbi__uint32 raw_len, bpl; - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (scan != STBI__SCAN_load) return 1; - if (z->idata == NULL) return stbi__err("no IDAT","Corrupt PNG"); - // initial guess for decoded data size to avoid unnecessary reallocs - bpl = (s->img_x * z->depth + 7) / 8; // bytes per line, per component - raw_len = bpl * s->img_y * s->img_n /* pixels */ + s->img_y /* filter mode per row */; - z->expanded = (stbi_uc *) stbi_zlib_decode_malloc_guesssize_headerflag((char *) z->idata, ioff, raw_len, (int *) &raw_len, !is_iphone); - if (z->expanded == NULL) return 0; // zlib should set error - STBI_FREE(z->idata); z->idata = NULL; - if ((req_comp == s->img_n+1 && req_comp != 3 && !pal_img_n) || has_trans) - s->img_out_n = s->img_n+1; - else - s->img_out_n = s->img_n; - if (!stbi__create_png_image(z, z->expanded, raw_len, s->img_out_n, z->depth, color, interlace)) return 0; - if (has_trans) { - if (z->depth == 16) { - if (!stbi__compute_transparency16(z, tc16, s->img_out_n)) return 0; - } else { - if (!stbi__compute_transparency(z, tc, s->img_out_n)) return 0; - } - } - if (is_iphone && stbi__de_iphone_flag && s->img_out_n > 2) - stbi__de_iphone(z); - if (pal_img_n) { - // pal_img_n == 3 or 4 - s->img_n = pal_img_n; // record the actual colors we had - s->img_out_n = pal_img_n; - if (req_comp >= 3) s->img_out_n = req_comp; - if (!stbi__expand_png_palette(z, palette, pal_len, s->img_out_n)) - return 0; - } else if (has_trans) { - // non-paletted image with tRNS -> source image has (constant) alpha - ++s->img_n; - } - STBI_FREE(z->expanded); z->expanded = NULL; - // end of PNG chunk, read and skip CRC - stbi__get32be(s); - return 1; - } - - default: - // if critical, fail - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if ((c.type & (1 << 29)) == 0) { - #ifndef STBI_NO_FAILURE_STRINGS - // not threadsafe - static char invalid_chunk[] = "XXXX PNG chunk not known"; - invalid_chunk[0] = STBI__BYTECAST(c.type >> 24); - invalid_chunk[1] = STBI__BYTECAST(c.type >> 16); - invalid_chunk[2] = STBI__BYTECAST(c.type >> 8); - invalid_chunk[3] = STBI__BYTECAST(c.type >> 0); - #endif - return stbi__err(invalid_chunk, "PNG not supported: unknown PNG chunk type"); - } - stbi__skip(s, c.length); - break; - } - // end of PNG chunk, read and skip CRC - stbi__get32be(s); - } -} - -static void *stbi__do_png(stbi__png *p, int *x, int *y, int *n, int req_comp, stbi__result_info *ri) -{ - void *result=NULL; - if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); - if (stbi__parse_png_file(p, STBI__SCAN_load, req_comp)) { - if (p->depth <= 8) - ri->bits_per_channel = 8; - else if (p->depth == 16) - ri->bits_per_channel = 16; - else - return stbi__errpuc("bad bits_per_channel", "PNG not supported: unsupported color depth"); - result = p->out; - p->out = NULL; - if (req_comp && req_comp != p->s->img_out_n) { - if (ri->bits_per_channel == 8) - result = stbi__convert_format((unsigned char *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); - else - result = stbi__convert_format16((stbi__uint16 *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); - p->s->img_out_n = req_comp; - if (result == NULL) return result; - } - *x = p->s->img_x; - *y = p->s->img_y; - if (n) *n = p->s->img_n; - } - STBI_FREE(p->out); p->out = NULL; - STBI_FREE(p->expanded); p->expanded = NULL; - STBI_FREE(p->idata); p->idata = NULL; - - return result; -} - -static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi__png p; - p.s = s; - return stbi__do_png(&p, x,y,comp,req_comp, ri); -} - -static int stbi__png_test(stbi__context *s) -{ - int r; - r = stbi__check_png_header(s); - stbi__rewind(s); - return r; -} - -static int stbi__png_info_raw(stbi__png *p, int *x, int *y, int *comp) -{ - if (!stbi__parse_png_file(p, STBI__SCAN_header, 0)) { - stbi__rewind( p->s ); - return 0; - } - if (x) *x = p->s->img_x; - if (y) *y = p->s->img_y; - if (comp) *comp = p->s->img_n; - return 1; -} - -static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp) -{ - stbi__png p; - p.s = s; - return stbi__png_info_raw(&p, x, y, comp); -} - -static int stbi__png_is16(stbi__context *s) -{ - stbi__png p; - p.s = s; - if (!stbi__png_info_raw(&p, NULL, NULL, NULL)) - return 0; - if (p.depth != 16) { - stbi__rewind(p.s); - return 0; - } - return 1; -} -#endif - -// Microsoft/Windows BMP image - -#ifndef STBI_NO_BMP -static int stbi__bmp_test_raw(stbi__context *s) -{ - int r; - int sz; - if (stbi__get8(s) != 'B') return 0; - if (stbi__get8(s) != 'M') return 0; - stbi__get32le(s); // discard filesize - stbi__get16le(s); // discard reserved - stbi__get16le(s); // discard reserved - stbi__get32le(s); // discard data offset - sz = stbi__get32le(s); - r = (sz == 12 || sz == 40 || sz == 56 || sz == 108 || sz == 124); - return r; -} - -static int stbi__bmp_test(stbi__context *s) -{ - int r = stbi__bmp_test_raw(s); - stbi__rewind(s); - return r; -} - - -// returns 0..31 for the highest set bit -static int stbi__high_bit(unsigned int z) -{ - int n=0; - if (z == 0) return -1; - if (z >= 0x10000) { n += 16; z >>= 16; } - if (z >= 0x00100) { n += 8; z >>= 8; } - if (z >= 0x00010) { n += 4; z >>= 4; } - if (z >= 0x00004) { n += 2; z >>= 2; } - if (z >= 0x00002) { n += 1;/* >>= 1;*/ } - return n; -} - -static int stbi__bitcount(unsigned int a) -{ - a = (a & 0x55555555) + ((a >> 1) & 0x55555555); // max 2 - a = (a & 0x33333333) + ((a >> 2) & 0x33333333); // max 4 - a = (a + (a >> 4)) & 0x0f0f0f0f; // max 8 per 4, now 8 bits - a = (a + (a >> 8)); // max 16 per 8 bits - a = (a + (a >> 16)); // max 32 per 8 bits - return a & 0xff; -} - -// extract an arbitrarily-aligned N-bit value (N=bits) -// from v, and then make it 8-bits long and fractionally -// extend it to full full range. -static int stbi__shiftsigned(unsigned int v, int shift, int bits) -{ - static unsigned int mul_table[9] = { - 0, - 0xff/*0b11111111*/, 0x55/*0b01010101*/, 0x49/*0b01001001*/, 0x11/*0b00010001*/, - 0x21/*0b00100001*/, 0x41/*0b01000001*/, 0x81/*0b10000001*/, 0x01/*0b00000001*/, - }; - static unsigned int shift_table[9] = { - 0, 0,0,1,0,2,4,6,0, - }; - if (shift < 0) - v <<= -shift; - else - v >>= shift; - STBI_ASSERT(v < 256); - v >>= (8-bits); - STBI_ASSERT(bits >= 0 && bits <= 8); - return (int) ((unsigned) v * mul_table[bits]) >> shift_table[bits]; -} - -typedef struct -{ - int bpp, offset, hsz; - unsigned int mr,mg,mb,ma, all_a; - int extra_read; -} stbi__bmp_data; - -static int stbi__bmp_set_mask_defaults(stbi__bmp_data *info, int compress) -{ - // BI_BITFIELDS specifies masks explicitly, don't override - if (compress == 3) - return 1; - - if (compress == 0) { - if (info->bpp == 16) { - info->mr = 31u << 10; - info->mg = 31u << 5; - info->mb = 31u << 0; - } else if (info->bpp == 32) { - info->mr = 0xffu << 16; - info->mg = 0xffu << 8; - info->mb = 0xffu << 0; - info->ma = 0xffu << 24; - info->all_a = 0; // if all_a is 0 at end, then we loaded alpha channel but it was all 0 - } else { - // otherwise, use defaults, which is all-0 - info->mr = info->mg = info->mb = info->ma = 0; - } - return 1; - } - return 0; // error -} - -static void *stbi__bmp_parse_header(stbi__context *s, stbi__bmp_data *info) -{ - int hsz; - if (stbi__get8(s) != 'B' || stbi__get8(s) != 'M') return stbi__errpuc("not BMP", "Corrupt BMP"); - stbi__get32le(s); // discard filesize - stbi__get16le(s); // discard reserved - stbi__get16le(s); // discard reserved - info->offset = stbi__get32le(s); - info->hsz = hsz = stbi__get32le(s); - info->mr = info->mg = info->mb = info->ma = 0; - info->extra_read = 14; - - if (info->offset < 0) return stbi__errpuc("bad BMP", "bad BMP"); - - if (hsz != 12 && hsz != 40 && hsz != 56 && hsz != 108 && hsz != 124) return stbi__errpuc("unknown BMP", "BMP type not supported: unknown"); - if (hsz == 12) { - s->img_x = stbi__get16le(s); - s->img_y = stbi__get16le(s); - } else { - s->img_x = stbi__get32le(s); - s->img_y = stbi__get32le(s); - } - if (stbi__get16le(s) != 1) return stbi__errpuc("bad BMP", "bad BMP"); - info->bpp = stbi__get16le(s); - if (hsz != 12) { - int compress = stbi__get32le(s); - if (compress == 1 || compress == 2) return stbi__errpuc("BMP RLE", "BMP type not supported: RLE"); - if (compress >= 4) return stbi__errpuc("BMP JPEG/PNG", "BMP type not supported: unsupported compression"); // this includes PNG/JPEG modes - if (compress == 3 && info->bpp != 16 && info->bpp != 32) return stbi__errpuc("bad BMP", "bad BMP"); // bitfields requires 16 or 32 bits/pixel - stbi__get32le(s); // discard sizeof - stbi__get32le(s); // discard hres - stbi__get32le(s); // discard vres - stbi__get32le(s); // discard colorsused - stbi__get32le(s); // discard max important - if (hsz == 40 || hsz == 56) { - if (hsz == 56) { - stbi__get32le(s); - stbi__get32le(s); - stbi__get32le(s); - stbi__get32le(s); - } - if (info->bpp == 16 || info->bpp == 32) { - if (compress == 0) { - stbi__bmp_set_mask_defaults(info, compress); - } else if (compress == 3) { - info->mr = stbi__get32le(s); - info->mg = stbi__get32le(s); - info->mb = stbi__get32le(s); - info->extra_read += 12; - // not documented, but generated by photoshop and handled by mspaint - if (info->mr == info->mg && info->mg == info->mb) { - // ?!?!? - return stbi__errpuc("bad BMP", "bad BMP"); - } - } else - return stbi__errpuc("bad BMP", "bad BMP"); - } - } else { - // V4/V5 header - int i; - if (hsz != 108 && hsz != 124) - return stbi__errpuc("bad BMP", "bad BMP"); - info->mr = stbi__get32le(s); - info->mg = stbi__get32le(s); - info->mb = stbi__get32le(s); - info->ma = stbi__get32le(s); - if (compress != 3) // override mr/mg/mb unless in BI_BITFIELDS mode, as per docs - stbi__bmp_set_mask_defaults(info, compress); - stbi__get32le(s); // discard color space - for (i=0; i < 12; ++i) - stbi__get32le(s); // discard color space parameters - if (hsz == 124) { - stbi__get32le(s); // discard rendering intent - stbi__get32le(s); // discard offset of profile data - stbi__get32le(s); // discard size of profile data - stbi__get32le(s); // discard reserved - } - } - } - return (void *) 1; -} - - -static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *out; - unsigned int mr=0,mg=0,mb=0,ma=0, all_a; - stbi_uc pal[256][4]; - int psize=0,i,j,width; - int flip_vertically, pad, target; - stbi__bmp_data info; - STBI_NOTUSED(ri); - - info.all_a = 255; - if (stbi__bmp_parse_header(s, &info) == NULL) - return NULL; // error code already set - - flip_vertically = ((int) s->img_y) > 0; - s->img_y = abs((int) s->img_y); - - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - mr = info.mr; - mg = info.mg; - mb = info.mb; - ma = info.ma; - all_a = info.all_a; - - if (info.hsz == 12) { - if (info.bpp < 24) - psize = (info.offset - info.extra_read - 24) / 3; - } else { - if (info.bpp < 16) - psize = (info.offset - info.extra_read - info.hsz) >> 2; - } - if (psize == 0) { - // accept some number of extra bytes after the header, but if the offset points either to before - // the header ends or implies a large amount of extra data, reject the file as malformed - int bytes_read_so_far = s->callback_already_read + (int)(s->img_buffer - s->img_buffer_original); - int header_limit = 1024; // max we actually read is below 256 bytes currently. - int extra_data_limit = 256*4; // what ordinarily goes here is a palette; 256 entries*4 bytes is its max size. - if (bytes_read_so_far <= 0 || bytes_read_so_far > header_limit) { - return stbi__errpuc("bad header", "Corrupt BMP"); - } - // we established that bytes_read_so_far is positive and sensible. - // the first half of this test rejects offsets that are either too small positives, or - // negative, and guarantees that info.offset >= bytes_read_so_far > 0. this in turn - // ensures the number computed in the second half of the test can't overflow. - if (info.offset < bytes_read_so_far || info.offset - bytes_read_so_far > extra_data_limit) { - return stbi__errpuc("bad offset", "Corrupt BMP"); - } else { - stbi__skip(s, info.offset - bytes_read_so_far); - } - } - - if (info.bpp == 24 && ma == 0xff000000) - s->img_n = 3; - else - s->img_n = ma ? 4 : 3; - if (req_comp && req_comp >= 3) // we can directly decode 3 or 4 - target = req_comp; - else - target = s->img_n; // if they want monochrome, we'll post-convert - - // sanity-check size - if (!stbi__mad3sizes_valid(target, s->img_x, s->img_y, 0)) - return stbi__errpuc("too large", "Corrupt BMP"); - - out = (stbi_uc *) stbi__malloc_mad3(target, s->img_x, s->img_y, 0); - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - if (info.bpp < 16) { - int z=0; - if (psize == 0 || psize > 256) { STBI_FREE(out); return stbi__errpuc("invalid", "Corrupt BMP"); } - for (i=0; i < psize; ++i) { - pal[i][2] = stbi__get8(s); - pal[i][1] = stbi__get8(s); - pal[i][0] = stbi__get8(s); - if (info.hsz != 12) stbi__get8(s); - pal[i][3] = 255; - } - stbi__skip(s, info.offset - info.extra_read - info.hsz - psize * (info.hsz == 12 ? 3 : 4)); - if (info.bpp == 1) width = (s->img_x + 7) >> 3; - else if (info.bpp == 4) width = (s->img_x + 1) >> 1; - else if (info.bpp == 8) width = s->img_x; - else { STBI_FREE(out); return stbi__errpuc("bad bpp", "Corrupt BMP"); } - pad = (-width)&3; - if (info.bpp == 1) { - for (j=0; j < (int) s->img_y; ++j) { - int bit_offset = 7, v = stbi__get8(s); - for (i=0; i < (int) s->img_x; ++i) { - int color = (v>>bit_offset)&0x1; - out[z++] = pal[color][0]; - out[z++] = pal[color][1]; - out[z++] = pal[color][2]; - if (target == 4) out[z++] = 255; - if (i+1 == (int) s->img_x) break; - if((--bit_offset) < 0) { - bit_offset = 7; - v = stbi__get8(s); - } - } - stbi__skip(s, pad); - } - } else { - for (j=0; j < (int) s->img_y; ++j) { - for (i=0; i < (int) s->img_x; i += 2) { - int v=stbi__get8(s),v2=0; - if (info.bpp == 4) { - v2 = v & 15; - v >>= 4; - } - out[z++] = pal[v][0]; - out[z++] = pal[v][1]; - out[z++] = pal[v][2]; - if (target == 4) out[z++] = 255; - if (i+1 == (int) s->img_x) break; - v = (info.bpp == 8) ? stbi__get8(s) : v2; - out[z++] = pal[v][0]; - out[z++] = pal[v][1]; - out[z++] = pal[v][2]; - if (target == 4) out[z++] = 255; - } - stbi__skip(s, pad); - } - } - } else { - int rshift=0,gshift=0,bshift=0,ashift=0,rcount=0,gcount=0,bcount=0,acount=0; - int z = 0; - int easy=0; - stbi__skip(s, info.offset - info.extra_read - info.hsz); - if (info.bpp == 24) width = 3 * s->img_x; - else if (info.bpp == 16) width = 2*s->img_x; - else /* bpp = 32 and pad = 0 */ width=0; - pad = (-width) & 3; - if (info.bpp == 24) { - easy = 1; - } else if (info.bpp == 32) { - if (mb == 0xff && mg == 0xff00 && mr == 0x00ff0000 && ma == 0xff000000) - easy = 2; - } - if (!easy) { - if (!mr || !mg || !mb) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } - // right shift amt to put high bit in position #7 - rshift = stbi__high_bit(mr)-7; rcount = stbi__bitcount(mr); - gshift = stbi__high_bit(mg)-7; gcount = stbi__bitcount(mg); - bshift = stbi__high_bit(mb)-7; bcount = stbi__bitcount(mb); - ashift = stbi__high_bit(ma)-7; acount = stbi__bitcount(ma); - if (rcount > 8 || gcount > 8 || bcount > 8 || acount > 8) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } - } - for (j=0; j < (int) s->img_y; ++j) { - if (easy) { - for (i=0; i < (int) s->img_x; ++i) { - unsigned char a; - out[z+2] = stbi__get8(s); - out[z+1] = stbi__get8(s); - out[z+0] = stbi__get8(s); - z += 3; - a = (easy == 2 ? stbi__get8(s) : 255); - all_a |= a; - if (target == 4) out[z++] = a; - } - } else { - int bpp = info.bpp; - for (i=0; i < (int) s->img_x; ++i) { - stbi__uint32 v = (bpp == 16 ? (stbi__uint32) stbi__get16le(s) : stbi__get32le(s)); - unsigned int a; - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mr, rshift, rcount)); - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mg, gshift, gcount)); - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mb, bshift, bcount)); - a = (ma ? stbi__shiftsigned(v & ma, ashift, acount) : 255); - all_a |= a; - if (target == 4) out[z++] = STBI__BYTECAST(a); - } - } - stbi__skip(s, pad); - } - } - - // if alpha channel is all 0s, replace with all 255s - if (target == 4 && all_a == 0) - for (i=4*s->img_x*s->img_y-1; i >= 0; i -= 4) - out[i] = 255; - - if (flip_vertically) { - stbi_uc t; - for (j=0; j < (int) s->img_y>>1; ++j) { - stbi_uc *p1 = out + j *s->img_x*target; - stbi_uc *p2 = out + (s->img_y-1-j)*s->img_x*target; - for (i=0; i < (int) s->img_x*target; ++i) { - t = p1[i]; p1[i] = p2[i]; p2[i] = t; - } - } - } - - if (req_comp && req_comp != target) { - out = stbi__convert_format(out, target, req_comp, s->img_x, s->img_y); - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - - *x = s->img_x; - *y = s->img_y; - if (comp) *comp = s->img_n; - return out; -} -#endif - -// Targa Truevision - TGA -// by Jonathan Dummer -#ifndef STBI_NO_TGA -// returns STBI_rgb or whatever, 0 on error -static int stbi__tga_get_comp(int bits_per_pixel, int is_grey, int* is_rgb16) -{ - // only RGB or RGBA (incl. 16bit) or grey allowed - if (is_rgb16) *is_rgb16 = 0; - switch(bits_per_pixel) { - case 8: return STBI_grey; - case 16: if(is_grey) return STBI_grey_alpha; - // fallthrough - case 15: if(is_rgb16) *is_rgb16 = 1; - return STBI_rgb; - case 24: // fallthrough - case 32: return bits_per_pixel/8; - default: return 0; - } -} - -static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp) -{ - int tga_w, tga_h, tga_comp, tga_image_type, tga_bits_per_pixel, tga_colormap_bpp; - int sz, tga_colormap_type; - stbi__get8(s); // discard Offset - tga_colormap_type = stbi__get8(s); // colormap type - if( tga_colormap_type > 1 ) { - stbi__rewind(s); - return 0; // only RGB or indexed allowed - } - tga_image_type = stbi__get8(s); // image type - if ( tga_colormap_type == 1 ) { // colormapped (paletted) image - if (tga_image_type != 1 && tga_image_type != 9) { - stbi__rewind(s); - return 0; - } - stbi__skip(s,4); // skip index of first colormap entry and number of entries - sz = stbi__get8(s); // check bits per palette color entry - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) { - stbi__rewind(s); - return 0; - } - stbi__skip(s,4); // skip image x and y origin - tga_colormap_bpp = sz; - } else { // "normal" image w/o colormap - only RGB or grey allowed, +/- RLE - if ( (tga_image_type != 2) && (tga_image_type != 3) && (tga_image_type != 10) && (tga_image_type != 11) ) { - stbi__rewind(s); - return 0; // only RGB or grey allowed, +/- RLE - } - stbi__skip(s,9); // skip colormap specification and image x/y origin - tga_colormap_bpp = 0; - } - tga_w = stbi__get16le(s); - if( tga_w < 1 ) { - stbi__rewind(s); - return 0; // test width - } - tga_h = stbi__get16le(s); - if( tga_h < 1 ) { - stbi__rewind(s); - return 0; // test height - } - tga_bits_per_pixel = stbi__get8(s); // bits per pixel - stbi__get8(s); // ignore alpha bits - if (tga_colormap_bpp != 0) { - if((tga_bits_per_pixel != 8) && (tga_bits_per_pixel != 16)) { - // when using a colormap, tga_bits_per_pixel is the size of the indexes - // I don't think anything but 8 or 16bit indexes makes sense - stbi__rewind(s); - return 0; - } - tga_comp = stbi__tga_get_comp(tga_colormap_bpp, 0, NULL); - } else { - tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3) || (tga_image_type == 11), NULL); - } - if(!tga_comp) { - stbi__rewind(s); - return 0; - } - if (x) *x = tga_w; - if (y) *y = tga_h; - if (comp) *comp = tga_comp; - return 1; // seems to have passed everything -} - -static int stbi__tga_test(stbi__context *s) -{ - int res = 0; - int sz, tga_color_type; - stbi__get8(s); // discard Offset - tga_color_type = stbi__get8(s); // color type - if ( tga_color_type > 1 ) goto errorEnd; // only RGB or indexed allowed - sz = stbi__get8(s); // image type - if ( tga_color_type == 1 ) { // colormapped (paletted) image - if (sz != 1 && sz != 9) goto errorEnd; // colortype 1 demands image type 1 or 9 - stbi__skip(s,4); // skip index of first colormap entry and number of entries - sz = stbi__get8(s); // check bits per palette color entry - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; - stbi__skip(s,4); // skip image x and y origin - } else { // "normal" image w/o colormap - if ( (sz != 2) && (sz != 3) && (sz != 10) && (sz != 11) ) goto errorEnd; // only RGB or grey allowed, +/- RLE - stbi__skip(s,9); // skip colormap specification and image x/y origin - } - if ( stbi__get16le(s) < 1 ) goto errorEnd; // test width - if ( stbi__get16le(s) < 1 ) goto errorEnd; // test height - sz = stbi__get8(s); // bits per pixel - if ( (tga_color_type == 1) && (sz != 8) && (sz != 16) ) goto errorEnd; // for colormapped images, bpp is size of an index - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; - - res = 1; // if we got this far, everything's good and we can return 1 instead of 0 - -errorEnd: - stbi__rewind(s); - return res; -} - -// read 16bit value and convert to 24bit RGB -static void stbi__tga_read_rgb16(stbi__context *s, stbi_uc* out) -{ - stbi__uint16 px = (stbi__uint16)stbi__get16le(s); - stbi__uint16 fiveBitMask = 31; - // we have 3 channels with 5bits each - int r = (px >> 10) & fiveBitMask; - int g = (px >> 5) & fiveBitMask; - int b = px & fiveBitMask; - // Note that this saves the data in RGB(A) order, so it doesn't need to be swapped later - out[0] = (stbi_uc)((r * 255)/31); - out[1] = (stbi_uc)((g * 255)/31); - out[2] = (stbi_uc)((b * 255)/31); - - // some people claim that the most significant bit might be used for alpha - // (possibly if an alpha-bit is set in the "image descriptor byte") - // but that only made 16bit test images completely translucent.. - // so let's treat all 15 and 16bit TGAs as RGB with no alpha. -} - -static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - // read in the TGA header stuff - int tga_offset = stbi__get8(s); - int tga_indexed = stbi__get8(s); - int tga_image_type = stbi__get8(s); - int tga_is_RLE = 0; - int tga_palette_start = stbi__get16le(s); - int tga_palette_len = stbi__get16le(s); - int tga_palette_bits = stbi__get8(s); - int tga_x_origin = stbi__get16le(s); - int tga_y_origin = stbi__get16le(s); - int tga_width = stbi__get16le(s); - int tga_height = stbi__get16le(s); - int tga_bits_per_pixel = stbi__get8(s); - int tga_comp, tga_rgb16=0; - int tga_inverted = stbi__get8(s); - // int tga_alpha_bits = tga_inverted & 15; // the 4 lowest bits - unused (useless?) - // image data - unsigned char *tga_data; - unsigned char *tga_palette = NULL; - int i, j; - unsigned char raw_data[4] = {0}; - int RLE_count = 0; - int RLE_repeating = 0; - int read_next_pixel = 1; - STBI_NOTUSED(ri); - STBI_NOTUSED(tga_x_origin); // @TODO - STBI_NOTUSED(tga_y_origin); // @TODO - - if (tga_height > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (tga_width > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - // do a tiny bit of precessing - if ( tga_image_type >= 8 ) - { - tga_image_type -= 8; - tga_is_RLE = 1; - } - tga_inverted = 1 - ((tga_inverted >> 5) & 1); - - // If I'm paletted, then I'll use the number of bits from the palette - if ( tga_indexed ) tga_comp = stbi__tga_get_comp(tga_palette_bits, 0, &tga_rgb16); - else tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3), &tga_rgb16); - - if(!tga_comp) // shouldn't really happen, stbi__tga_test() should have ensured basic consistency - return stbi__errpuc("bad format", "Can't find out TGA pixelformat"); - - // tga info - *x = tga_width; - *y = tga_height; - if (comp) *comp = tga_comp; - - if (!stbi__mad3sizes_valid(tga_width, tga_height, tga_comp, 0)) - return stbi__errpuc("too large", "Corrupt TGA"); - - tga_data = (unsigned char*)stbi__malloc_mad3(tga_width, tga_height, tga_comp, 0); - if (!tga_data) return stbi__errpuc("outofmem", "Out of memory"); - - // skip to the data's starting position (offset usually = 0) - stbi__skip(s, tga_offset ); - - if ( !tga_indexed && !tga_is_RLE && !tga_rgb16 ) { - for (i=0; i < tga_height; ++i) { - int row = tga_inverted ? tga_height -i - 1 : i; - stbi_uc *tga_row = tga_data + row*tga_width*tga_comp; - stbi__getn(s, tga_row, tga_width * tga_comp); - } - } else { - // do I need to load a palette? - if ( tga_indexed) - { - if (tga_palette_len == 0) { /* you have to have at least one entry! */ - STBI_FREE(tga_data); - return stbi__errpuc("bad palette", "Corrupt TGA"); - } - - // any data to skip? (offset usually = 0) - stbi__skip(s, tga_palette_start ); - // load the palette - tga_palette = (unsigned char*)stbi__malloc_mad2(tga_palette_len, tga_comp, 0); - if (!tga_palette) { - STBI_FREE(tga_data); - return stbi__errpuc("outofmem", "Out of memory"); - } - if (tga_rgb16) { - stbi_uc *pal_entry = tga_palette; - STBI_ASSERT(tga_comp == STBI_rgb); - for (i=0; i < tga_palette_len; ++i) { - stbi__tga_read_rgb16(s, pal_entry); - pal_entry += tga_comp; - } - } else if (!stbi__getn(s, tga_palette, tga_palette_len * tga_comp)) { - STBI_FREE(tga_data); - STBI_FREE(tga_palette); - return stbi__errpuc("bad palette", "Corrupt TGA"); - } - } - // load the data - for (i=0; i < tga_width * tga_height; ++i) - { - // if I'm in RLE mode, do I need to get a RLE stbi__pngchunk? - if ( tga_is_RLE ) - { - if ( RLE_count == 0 ) - { - // yep, get the next byte as a RLE command - int RLE_cmd = stbi__get8(s); - RLE_count = 1 + (RLE_cmd & 127); - RLE_repeating = RLE_cmd >> 7; - read_next_pixel = 1; - } else if ( !RLE_repeating ) - { - read_next_pixel = 1; - } - } else - { - read_next_pixel = 1; - } - // OK, if I need to read a pixel, do it now - if ( read_next_pixel ) - { - // load however much data we did have - if ( tga_indexed ) - { - // read in index, then perform the lookup - int pal_idx = (tga_bits_per_pixel == 8) ? stbi__get8(s) : stbi__get16le(s); - if ( pal_idx >= tga_palette_len ) { - // invalid index - pal_idx = 0; - } - pal_idx *= tga_comp; - for (j = 0; j < tga_comp; ++j) { - raw_data[j] = tga_palette[pal_idx+j]; - } - } else if(tga_rgb16) { - STBI_ASSERT(tga_comp == STBI_rgb); - stbi__tga_read_rgb16(s, raw_data); - } else { - // read in the data raw - for (j = 0; j < tga_comp; ++j) { - raw_data[j] = stbi__get8(s); - } - } - // clear the reading flag for the next pixel - read_next_pixel = 0; - } // end of reading a pixel - - // copy data - for (j = 0; j < tga_comp; ++j) - tga_data[i*tga_comp+j] = raw_data[j]; - - // in case we're in RLE mode, keep counting down - --RLE_count; - } - // do I need to invert the image? - if ( tga_inverted ) - { - for (j = 0; j*2 < tga_height; ++j) - { - int index1 = j * tga_width * tga_comp; - int index2 = (tga_height - 1 - j) * tga_width * tga_comp; - for (i = tga_width * tga_comp; i > 0; --i) - { - unsigned char temp = tga_data[index1]; - tga_data[index1] = tga_data[index2]; - tga_data[index2] = temp; - ++index1; - ++index2; - } - } - } - // clear my palette, if I had one - if ( tga_palette != NULL ) - { - STBI_FREE( tga_palette ); - } - } - - // swap RGB - if the source data was RGB16, it already is in the right order - if (tga_comp >= 3 && !tga_rgb16) - { - unsigned char* tga_pixel = tga_data; - for (i=0; i < tga_width * tga_height; ++i) - { - unsigned char temp = tga_pixel[0]; - tga_pixel[0] = tga_pixel[2]; - tga_pixel[2] = temp; - tga_pixel += tga_comp; - } - } - - // convert to target component count - if (req_comp && req_comp != tga_comp) - tga_data = stbi__convert_format(tga_data, tga_comp, req_comp, tga_width, tga_height); - - // the things I do to get rid of an error message, and yet keep - // Microsoft's C compilers happy... [8^( - tga_palette_start = tga_palette_len = tga_palette_bits = - tga_x_origin = tga_y_origin = 0; - STBI_NOTUSED(tga_palette_start); - // OK, done - return tga_data; -} -#endif - -// ************************************************************************************************* -// Photoshop PSD loader -- PD by Thatcher Ulrich, integration by Nicolas Schulz, tweaked by STB - -#ifndef STBI_NO_PSD -static int stbi__psd_test(stbi__context *s) -{ - int r = (stbi__get32be(s) == 0x38425053); - stbi__rewind(s); - return r; -} - -static int stbi__psd_decode_rle(stbi__context *s, stbi_uc *p, int pixelCount) -{ - int count, nleft, len; - - count = 0; - while ((nleft = pixelCount - count) > 0) { - len = stbi__get8(s); - if (len == 128) { - // No-op. - } else if (len < 128) { - // Copy next len+1 bytes literally. - len++; - if (len > nleft) return 0; // corrupt data - count += len; - while (len) { - *p = stbi__get8(s); - p += 4; - len--; - } - } else if (len > 128) { - stbi_uc val; - // Next -len+1 bytes in the dest are replicated from next source byte. - // (Interpret len as a negative 8-bit int.) - len = 257 - len; - if (len > nleft) return 0; // corrupt data - val = stbi__get8(s); - count += len; - while (len) { - *p = val; - p += 4; - len--; - } - } - } - - return 1; -} - -static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) -{ - int pixelCount; - int channelCount, compression; - int channel, i; - int bitdepth; - int w,h; - stbi_uc *out; - STBI_NOTUSED(ri); - - // Check identifier - if (stbi__get32be(s) != 0x38425053) // "8BPS" - return stbi__errpuc("not PSD", "Corrupt PSD image"); - - // Check file type version. - if (stbi__get16be(s) != 1) - return stbi__errpuc("wrong version", "Unsupported version of PSD image"); - - // Skip 6 reserved bytes. - stbi__skip(s, 6 ); - - // Read the number of channels (R, G, B, A, etc). - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) - return stbi__errpuc("wrong channel count", "Unsupported number of channels in PSD image"); - - // Read the rows and columns of the image. - h = stbi__get32be(s); - w = stbi__get32be(s); - - if (h > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (w > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - // Make sure the depth is 8 bits. - bitdepth = stbi__get16be(s); - if (bitdepth != 8 && bitdepth != 16) - return stbi__errpuc("unsupported bit depth", "PSD bit depth is not 8 or 16 bit"); - - // Make sure the color mode is RGB. - // Valid options are: - // 0: Bitmap - // 1: Grayscale - // 2: Indexed color - // 3: RGB color - // 4: CMYK color - // 7: Multichannel - // 8: Duotone - // 9: Lab color - if (stbi__get16be(s) != 3) - return stbi__errpuc("wrong color format", "PSD is not in RGB color format"); - - // Skip the Mode Data. (It's the palette for indexed color; other info for other modes.) - stbi__skip(s,stbi__get32be(s) ); - - // Skip the image resources. (resolution, pen tool paths, etc) - stbi__skip(s, stbi__get32be(s) ); - - // Skip the reserved data. - stbi__skip(s, stbi__get32be(s) ); - - // Find out if the data is compressed. - // Known values: - // 0: no compression - // 1: RLE compressed - compression = stbi__get16be(s); - if (compression > 1) - return stbi__errpuc("bad compression", "PSD has an unknown compression format"); - - // Check size - if (!stbi__mad3sizes_valid(4, w, h, 0)) - return stbi__errpuc("too large", "Corrupt PSD"); - - // Create the destination image. - - if (!compression && bitdepth == 16 && bpc == 16) { - out = (stbi_uc *) stbi__malloc_mad3(8, w, h, 0); - ri->bits_per_channel = 16; - } else - out = (stbi_uc *) stbi__malloc(4 * w*h); - - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - pixelCount = w*h; - - // Initialize the data to zero. - //memset( out, 0, pixelCount * 4 ); - - // Finally, the image data. - if (compression) { - // RLE as used by .PSD and .TIFF - // Loop until you get the number of unpacked bytes you are expecting: - // Read the next source byte into n. - // If n is between 0 and 127 inclusive, copy the next n+1 bytes literally. - // Else if n is between -127 and -1 inclusive, copy the next byte -n+1 times. - // Else if n is 128, noop. - // Endloop - - // The RLE-compressed data is preceded by a 2-byte data count for each row in the data, - // which we're going to just skip. - stbi__skip(s, h * channelCount * 2 ); - - // Read the RLE data by channel. - for (channel = 0; channel < 4; channel++) { - stbi_uc *p; - - p = out+channel; - if (channel >= channelCount) { - // Fill this channel with default data. - for (i = 0; i < pixelCount; i++, p += 4) - *p = (channel == 3 ? 255 : 0); - } else { - // Read the RLE data. - if (!stbi__psd_decode_rle(s, p, pixelCount)) { - STBI_FREE(out); - return stbi__errpuc("corrupt", "bad RLE data"); - } - } - } - - } else { - // We're at the raw image data. It's each channel in order (Red, Green, Blue, Alpha, ...) - // where each channel consists of an 8-bit (or 16-bit) value for each pixel in the image. - - // Read the data by channel. - for (channel = 0; channel < 4; channel++) { - if (channel >= channelCount) { - // Fill this channel with default data. - if (bitdepth == 16 && bpc == 16) { - stbi__uint16 *q = ((stbi__uint16 *) out) + channel; - stbi__uint16 val = channel == 3 ? 65535 : 0; - for (i = 0; i < pixelCount; i++, q += 4) - *q = val; - } else { - stbi_uc *p = out+channel; - stbi_uc val = channel == 3 ? 255 : 0; - for (i = 0; i < pixelCount; i++, p += 4) - *p = val; - } - } else { - if (ri->bits_per_channel == 16) { // output bpc - stbi__uint16 *q = ((stbi__uint16 *) out) + channel; - for (i = 0; i < pixelCount; i++, q += 4) - *q = (stbi__uint16) stbi__get16be(s); - } else { - stbi_uc *p = out+channel; - if (bitdepth == 16) { // input bpc - for (i = 0; i < pixelCount; i++, p += 4) - *p = (stbi_uc) (stbi__get16be(s) >> 8); - } else { - for (i = 0; i < pixelCount; i++, p += 4) - *p = stbi__get8(s); - } - } - } - } - } - - // remove weird white matte from PSD - if (channelCount >= 4) { - if (ri->bits_per_channel == 16) { - for (i=0; i < w*h; ++i) { - stbi__uint16 *pixel = (stbi__uint16 *) out + 4*i; - if (pixel[3] != 0 && pixel[3] != 65535) { - float a = pixel[3] / 65535.0f; - float ra = 1.0f / a; - float inv_a = 65535.0f * (1 - ra); - pixel[0] = (stbi__uint16) (pixel[0]*ra + inv_a); - pixel[1] = (stbi__uint16) (pixel[1]*ra + inv_a); - pixel[2] = (stbi__uint16) (pixel[2]*ra + inv_a); - } - } - } else { - for (i=0; i < w*h; ++i) { - unsigned char *pixel = out + 4*i; - if (pixel[3] != 0 && pixel[3] != 255) { - float a = pixel[3] / 255.0f; - float ra = 1.0f / a; - float inv_a = 255.0f * (1 - ra); - pixel[0] = (unsigned char) (pixel[0]*ra + inv_a); - pixel[1] = (unsigned char) (pixel[1]*ra + inv_a); - pixel[2] = (unsigned char) (pixel[2]*ra + inv_a); - } - } - } - } - - // convert to desired output format - if (req_comp && req_comp != 4) { - if (ri->bits_per_channel == 16) - out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, 4, req_comp, w, h); - else - out = stbi__convert_format(out, 4, req_comp, w, h); - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - - if (comp) *comp = 4; - *y = h; - *x = w; - - return out; -} -#endif - -// ************************************************************************************************* -// Softimage PIC loader -// by Tom Seddon -// -// See http://softimage.wiki.softimage.com/index.php/INFO:_PIC_file_format -// See http://ozviz.wasp.uwa.edu.au/~pbourke/dataformats/softimagepic/ - -#ifndef STBI_NO_PIC -static int stbi__pic_is4(stbi__context *s,const char *str) -{ - int i; - for (i=0; i<4; ++i) - if (stbi__get8(s) != (stbi_uc)str[i]) - return 0; - - return 1; -} - -static int stbi__pic_test_core(stbi__context *s) -{ - int i; - - if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) - return 0; - - for(i=0;i<84;++i) - stbi__get8(s); - - if (!stbi__pic_is4(s,"PICT")) - return 0; - - return 1; -} - -typedef struct -{ - stbi_uc size,type,channel; -} stbi__pic_packet; - -static stbi_uc *stbi__readval(stbi__context *s, int channel, stbi_uc *dest) -{ - int mask=0x80, i; - - for (i=0; i<4; ++i, mask>>=1) { - if (channel & mask) { - if (stbi__at_eof(s)) return stbi__errpuc("bad file","PIC file too short"); - dest[i]=stbi__get8(s); - } - } - - return dest; -} - -static void stbi__copyval(int channel,stbi_uc *dest,const stbi_uc *src) -{ - int mask=0x80,i; - - for (i=0;i<4; ++i, mask>>=1) - if (channel&mask) - dest[i]=src[i]; -} - -static stbi_uc *stbi__pic_load_core(stbi__context *s,int width,int height,int *comp, stbi_uc *result) -{ - int act_comp=0,num_packets=0,y,chained; - stbi__pic_packet packets[10]; - - // this will (should...) cater for even some bizarre stuff like having data - // for the same channel in multiple packets. - do { - stbi__pic_packet *packet; - - if (num_packets==sizeof(packets)/sizeof(packets[0])) - return stbi__errpuc("bad format","too many packets"); - - packet = &packets[num_packets++]; - - chained = stbi__get8(s); - packet->size = stbi__get8(s); - packet->type = stbi__get8(s); - packet->channel = stbi__get8(s); - - act_comp |= packet->channel; - - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (reading packets)"); - if (packet->size != 8) return stbi__errpuc("bad format","packet isn't 8bpp"); - } while (chained); - - *comp = (act_comp & 0x10 ? 4 : 3); // has alpha channel? - - for(y=0; ytype) { - default: - return stbi__errpuc("bad format","packet has bad compression type"); - - case 0: {//uncompressed - int x; - - for(x=0;xchannel,dest)) - return 0; - break; - } - - case 1://Pure RLE - { - int left=width, i; - - while (left>0) { - stbi_uc count,value[4]; - - count=stbi__get8(s); - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pure read count)"); - - if (count > left) - count = (stbi_uc) left; - - if (!stbi__readval(s,packet->channel,value)) return 0; - - for(i=0; ichannel,dest,value); - left -= count; - } - } - break; - - case 2: {//Mixed RLE - int left=width; - while (left>0) { - int count = stbi__get8(s), i; - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (mixed read count)"); - - if (count >= 128) { // Repeated - stbi_uc value[4]; - - if (count==128) - count = stbi__get16be(s); - else - count -= 127; - if (count > left) - return stbi__errpuc("bad file","scanline overrun"); - - if (!stbi__readval(s,packet->channel,value)) - return 0; - - for(i=0;ichannel,dest,value); - } else { // Raw - ++count; - if (count>left) return stbi__errpuc("bad file","scanline overrun"); - - for(i=0;ichannel,dest)) - return 0; - } - left-=count; - } - break; - } - } - } - } - - return result; -} - -static void *stbi__pic_load(stbi__context *s,int *px,int *py,int *comp,int req_comp, stbi__result_info *ri) -{ - stbi_uc *result; - int i, x,y, internal_comp; - STBI_NOTUSED(ri); - - if (!comp) comp = &internal_comp; - - for (i=0; i<92; ++i) - stbi__get8(s); - - x = stbi__get16be(s); - y = stbi__get16be(s); - - if (y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pic header)"); - if (!stbi__mad3sizes_valid(x, y, 4, 0)) return stbi__errpuc("too large", "PIC image too large to decode"); - - stbi__get32be(s); //skip `ratio' - stbi__get16be(s); //skip `fields' - stbi__get16be(s); //skip `pad' - - // intermediate buffer is RGBA - result = (stbi_uc *) stbi__malloc_mad3(x, y, 4, 0); - if (!result) return stbi__errpuc("outofmem", "Out of memory"); - memset(result, 0xff, x*y*4); - - if (!stbi__pic_load_core(s,x,y,comp, result)) { - STBI_FREE(result); - result=0; - } - *px = x; - *py = y; - if (req_comp == 0) req_comp = *comp; - result=stbi__convert_format(result,4,req_comp,x,y); - - return result; -} - -static int stbi__pic_test(stbi__context *s) -{ - int r = stbi__pic_test_core(s); - stbi__rewind(s); - return r; -} -#endif - -// ************************************************************************************************* -// GIF loader -- public domain by Jean-Marc Lienher -- simplified/shrunk by stb - -#ifndef STBI_NO_GIF -typedef struct -{ - stbi__int16 prefix; - stbi_uc first; - stbi_uc suffix; -} stbi__gif_lzw; - -typedef struct -{ - int w,h; - stbi_uc *out; // output buffer (always 4 components) - stbi_uc *background; // The current "background" as far as a gif is concerned - stbi_uc *history; - int flags, bgindex, ratio, transparent, eflags; - stbi_uc pal[256][4]; - stbi_uc lpal[256][4]; - stbi__gif_lzw codes[8192]; - stbi_uc *color_table; - int parse, step; - int lflags; - int start_x, start_y; - int max_x, max_y; - int cur_x, cur_y; - int line_size; - int delay; -} stbi__gif; - -static int stbi__gif_test_raw(stbi__context *s) -{ - int sz; - if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') return 0; - sz = stbi__get8(s); - if (sz != '9' && sz != '7') return 0; - if (stbi__get8(s) != 'a') return 0; - return 1; -} - -static int stbi__gif_test(stbi__context *s) -{ - int r = stbi__gif_test_raw(s); - stbi__rewind(s); - return r; -} - -static void stbi__gif_parse_colortable(stbi__context *s, stbi_uc pal[256][4], int num_entries, int transp) -{ - int i; - for (i=0; i < num_entries; ++i) { - pal[i][2] = stbi__get8(s); - pal[i][1] = stbi__get8(s); - pal[i][0] = stbi__get8(s); - pal[i][3] = transp == i ? 0 : 255; - } -} - -static int stbi__gif_header(stbi__context *s, stbi__gif *g, int *comp, int is_info) -{ - stbi_uc version; - if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') - return stbi__err("not GIF", "Corrupt GIF"); - - version = stbi__get8(s); - if (version != '7' && version != '9') return stbi__err("not GIF", "Corrupt GIF"); - if (stbi__get8(s) != 'a') return stbi__err("not GIF", "Corrupt GIF"); - - stbi__g_failure_reason = ""; - g->w = stbi__get16le(s); - g->h = stbi__get16le(s); - g->flags = stbi__get8(s); - g->bgindex = stbi__get8(s); - g->ratio = stbi__get8(s); - g->transparent = -1; - - if (g->w > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (g->h > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - - if (comp != 0) *comp = 4; // can't actually tell whether it's 3 or 4 until we parse the comments - - if (is_info) return 1; - - if (g->flags & 0x80) - stbi__gif_parse_colortable(s,g->pal, 2 << (g->flags & 7), -1); - - return 1; -} - -static int stbi__gif_info_raw(stbi__context *s, int *x, int *y, int *comp) -{ - stbi__gif* g = (stbi__gif*) stbi__malloc(sizeof(stbi__gif)); - if (!g) return stbi__err("outofmem", "Out of memory"); - if (!stbi__gif_header(s, g, comp, 1)) { - STBI_FREE(g); - stbi__rewind( s ); - return 0; - } - if (x) *x = g->w; - if (y) *y = g->h; - STBI_FREE(g); - return 1; -} - -static void stbi__out_gif_code(stbi__gif *g, stbi__uint16 code) -{ - stbi_uc *p, *c; - int idx; - - // recurse to decode the prefixes, since the linked-list is backwards, - // and working backwards through an interleaved image would be nasty - if (g->codes[code].prefix >= 0) - stbi__out_gif_code(g, g->codes[code].prefix); - - if (g->cur_y >= g->max_y) return; - - idx = g->cur_x + g->cur_y; - p = &g->out[idx]; - g->history[idx / 4] = 1; - - c = &g->color_table[g->codes[code].suffix * 4]; - if (c[3] > 128) { // don't render transparent pixels; - p[0] = c[2]; - p[1] = c[1]; - p[2] = c[0]; - p[3] = c[3]; - } - g->cur_x += 4; - - if (g->cur_x >= g->max_x) { - g->cur_x = g->start_x; - g->cur_y += g->step; - - while (g->cur_y >= g->max_y && g->parse > 0) { - g->step = (1 << g->parse) * g->line_size; - g->cur_y = g->start_y + (g->step >> 1); - --g->parse; - } - } -} - -static stbi_uc *stbi__process_gif_raster(stbi__context *s, stbi__gif *g) -{ - stbi_uc lzw_cs; - stbi__int32 len, init_code; - stbi__uint32 first; - stbi__int32 codesize, codemask, avail, oldcode, bits, valid_bits, clear; - stbi__gif_lzw *p; - - lzw_cs = stbi__get8(s); - if (lzw_cs > 12) return NULL; - clear = 1 << lzw_cs; - first = 1; - codesize = lzw_cs + 1; - codemask = (1 << codesize) - 1; - bits = 0; - valid_bits = 0; - for (init_code = 0; init_code < clear; init_code++) { - g->codes[init_code].prefix = -1; - g->codes[init_code].first = (stbi_uc) init_code; - g->codes[init_code].suffix = (stbi_uc) init_code; - } - - // support no starting clear code - avail = clear+2; - oldcode = -1; - - len = 0; - for(;;) { - if (valid_bits < codesize) { - if (len == 0) { - len = stbi__get8(s); // start new block - if (len == 0) - return g->out; - } - --len; - bits |= (stbi__int32) stbi__get8(s) << valid_bits; - valid_bits += 8; - } else { - stbi__int32 code = bits & codemask; - bits >>= codesize; - valid_bits -= codesize; - // @OPTIMIZE: is there some way we can accelerate the non-clear path? - if (code == clear) { // clear code - codesize = lzw_cs + 1; - codemask = (1 << codesize) - 1; - avail = clear + 2; - oldcode = -1; - first = 0; - } else if (code == clear + 1) { // end of stream code - stbi__skip(s, len); - while ((len = stbi__get8(s)) > 0) - stbi__skip(s,len); - return g->out; - } else if (code <= avail) { - if (first) { - return stbi__errpuc("no clear code", "Corrupt GIF"); - } - - if (oldcode >= 0) { - p = &g->codes[avail++]; - if (avail > 8192) { - return stbi__errpuc("too many codes", "Corrupt GIF"); - } - - p->prefix = (stbi__int16) oldcode; - p->first = g->codes[oldcode].first; - p->suffix = (code == avail) ? p->first : g->codes[code].first; - } else if (code == avail) - return stbi__errpuc("illegal code in raster", "Corrupt GIF"); - - stbi__out_gif_code(g, (stbi__uint16) code); - - if ((avail & codemask) == 0 && avail <= 0x0FFF) { - codesize++; - codemask = (1 << codesize) - 1; - } - - oldcode = code; - } else { - return stbi__errpuc("illegal code in raster", "Corrupt GIF"); - } - } - } -} - -// this function is designed to support animated gifs, although stb_image doesn't support it -// two back is the image from two frames ago, used for a very specific disposal format -static stbi_uc *stbi__gif_load_next(stbi__context *s, stbi__gif *g, int *comp, int req_comp, stbi_uc *two_back) -{ - int dispose; - int first_frame; - int pi; - int pcount; - STBI_NOTUSED(req_comp); - - // on first frame, any non-written pixels get the background colour (non-transparent) - first_frame = 0; - if (g->out == 0) { - if (!stbi__gif_header(s, g, comp,0)) return 0; // stbi__g_failure_reason set by stbi__gif_header - if (!stbi__mad3sizes_valid(4, g->w, g->h, 0)) - return stbi__errpuc("too large", "GIF image is too large"); - pcount = g->w * g->h; - g->out = (stbi_uc *) stbi__malloc(4 * pcount); - g->background = (stbi_uc *) stbi__malloc(4 * pcount); - g->history = (stbi_uc *) stbi__malloc(pcount); - if (!g->out || !g->background || !g->history) - return stbi__errpuc("outofmem", "Out of memory"); - - // image is treated as "transparent" at the start - ie, nothing overwrites the current background; - // background colour is only used for pixels that are not rendered first frame, after that "background" - // color refers to the color that was there the previous frame. - memset(g->out, 0x00, 4 * pcount); - memset(g->background, 0x00, 4 * pcount); // state of the background (starts transparent) - memset(g->history, 0x00, pcount); // pixels that were affected previous frame - first_frame = 1; - } else { - // second frame - how do we dispose of the previous one? - dispose = (g->eflags & 0x1C) >> 2; - pcount = g->w * g->h; - - if ((dispose == 3) && (two_back == 0)) { - dispose = 2; // if I don't have an image to revert back to, default to the old background - } - - if (dispose == 3) { // use previous graphic - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi]) { - memcpy( &g->out[pi * 4], &two_back[pi * 4], 4 ); - } - } - } else if (dispose == 2) { - // restore what was changed last frame to background before that frame; - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi]) { - memcpy( &g->out[pi * 4], &g->background[pi * 4], 4 ); - } - } - } else { - // This is a non-disposal case eithe way, so just - // leave the pixels as is, and they will become the new background - // 1: do not dispose - // 0: not specified. - } - - // background is what out is after the undoing of the previou frame; - memcpy( g->background, g->out, 4 * g->w * g->h ); - } - - // clear my history; - memset( g->history, 0x00, g->w * g->h ); // pixels that were affected previous frame - - for (;;) { - int tag = stbi__get8(s); - switch (tag) { - case 0x2C: /* Image Descriptor */ - { - stbi__int32 x, y, w, h; - stbi_uc *o; - - x = stbi__get16le(s); - y = stbi__get16le(s); - w = stbi__get16le(s); - h = stbi__get16le(s); - if (((x + w) > (g->w)) || ((y + h) > (g->h))) - return stbi__errpuc("bad Image Descriptor", "Corrupt GIF"); - - g->line_size = g->w * 4; - g->start_x = x * 4; - g->start_y = y * g->line_size; - g->max_x = g->start_x + w * 4; - g->max_y = g->start_y + h * g->line_size; - g->cur_x = g->start_x; - g->cur_y = g->start_y; - - // if the width of the specified rectangle is 0, that means - // we may not see *any* pixels or the image is malformed; - // to make sure this is caught, move the current y down to - // max_y (which is what out_gif_code checks). - if (w == 0) - g->cur_y = g->max_y; - - g->lflags = stbi__get8(s); - - if (g->lflags & 0x40) { - g->step = 8 * g->line_size; // first interlaced spacing - g->parse = 3; - } else { - g->step = g->line_size; - g->parse = 0; - } - - if (g->lflags & 0x80) { - stbi__gif_parse_colortable(s,g->lpal, 2 << (g->lflags & 7), g->eflags & 0x01 ? g->transparent : -1); - g->color_table = (stbi_uc *) g->lpal; - } else if (g->flags & 0x80) { - g->color_table = (stbi_uc *) g->pal; - } else - return stbi__errpuc("missing color table", "Corrupt GIF"); - - o = stbi__process_gif_raster(s, g); - if (!o) return NULL; - - // if this was the first frame, - pcount = g->w * g->h; - if (first_frame && (g->bgindex > 0)) { - // if first frame, any pixel not drawn to gets the background color - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi] == 0) { - g->pal[g->bgindex][3] = 255; // just in case it was made transparent, undo that; It will be reset next frame if need be; - memcpy( &g->out[pi * 4], &g->pal[g->bgindex], 4 ); - } - } - } - - return o; - } - - case 0x21: // Comment Extension. - { - int len; - int ext = stbi__get8(s); - if (ext == 0xF9) { // Graphic Control Extension. - len = stbi__get8(s); - if (len == 4) { - g->eflags = stbi__get8(s); - g->delay = 10 * stbi__get16le(s); // delay - 1/100th of a second, saving as 1/1000ths. - - // unset old transparent - if (g->transparent >= 0) { - g->pal[g->transparent][3] = 255; - } - if (g->eflags & 0x01) { - g->transparent = stbi__get8(s); - if (g->transparent >= 0) { - g->pal[g->transparent][3] = 0; - } - } else { - // don't need transparent - stbi__skip(s, 1); - g->transparent = -1; - } - } else { - stbi__skip(s, len); - break; - } - } - while ((len = stbi__get8(s)) != 0) { - stbi__skip(s, len); - } - break; - } - - case 0x3B: // gif stream termination code - return (stbi_uc *) s; // using '1' causes warning on some compilers - - default: - return stbi__errpuc("unknown code", "Corrupt GIF"); - } - } -} - -static void *stbi__load_gif_main_outofmem(stbi__gif *g, stbi_uc *out, int **delays) -{ - STBI_FREE(g->out); - STBI_FREE(g->history); - STBI_FREE(g->background); - - if (out) STBI_FREE(out); - if (delays && *delays) STBI_FREE(*delays); - return stbi__errpuc("outofmem", "Out of memory"); -} - -static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp) -{ - if (stbi__gif_test(s)) { - int layers = 0; - stbi_uc *u = 0; - stbi_uc *out = 0; - stbi_uc *two_back = 0; - stbi__gif g; - int stride; - int out_size = 0; - int delays_size = 0; - - STBI_NOTUSED(out_size); - STBI_NOTUSED(delays_size); - - memset(&g, 0, sizeof(g)); - if (delays) { - *delays = 0; - } - - do { - u = stbi__gif_load_next(s, &g, comp, req_comp, two_back); - if (u == (stbi_uc *) s) u = 0; // end of animated gif marker - - if (u) { - *x = g.w; - *y = g.h; - ++layers; - stride = g.w * g.h * 4; - - if (out) { - void *tmp = (stbi_uc*) STBI_REALLOC_SIZED( out, out_size, layers * stride ); - if (!tmp) - return stbi__load_gif_main_outofmem(&g, out, delays); - else { - out = (stbi_uc*) tmp; - out_size = layers * stride; - } - - if (delays) { - int *new_delays = (int*) STBI_REALLOC_SIZED( *delays, delays_size, sizeof(int) * layers ); - if (!new_delays) - return stbi__load_gif_main_outofmem(&g, out, delays); - *delays = new_delays; - delays_size = layers * sizeof(int); - } - } else { - out = (stbi_uc*)stbi__malloc( layers * stride ); - if (!out) - return stbi__load_gif_main_outofmem(&g, out, delays); - out_size = layers * stride; - if (delays) { - *delays = (int*) stbi__malloc( layers * sizeof(int) ); - if (!*delays) - return stbi__load_gif_main_outofmem(&g, out, delays); - delays_size = layers * sizeof(int); - } - } - memcpy( out + ((layers - 1) * stride), u, stride ); - if (layers >= 2) { - two_back = out - 2 * stride; - } - - if (delays) { - (*delays)[layers - 1U] = g.delay; - } - } - } while (u != 0); - - // free temp buffer; - STBI_FREE(g.out); - STBI_FREE(g.history); - STBI_FREE(g.background); - - // do the final conversion after loading everything; - if (req_comp && req_comp != 4) - out = stbi__convert_format(out, 4, req_comp, layers * g.w, g.h); - - *z = layers; - return out; - } else { - return stbi__errpuc("not GIF", "Image was not as a gif type."); - } -} - -static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *u = 0; - stbi__gif g; - memset(&g, 0, sizeof(g)); - STBI_NOTUSED(ri); - - u = stbi__gif_load_next(s, &g, comp, req_comp, 0); - if (u == (stbi_uc *) s) u = 0; // end of animated gif marker - if (u) { - *x = g.w; - *y = g.h; - - // moved conversion to after successful load so that the same - // can be done for multiple frames. - if (req_comp && req_comp != 4) - u = stbi__convert_format(u, 4, req_comp, g.w, g.h); - } else if (g.out) { - // if there was an error and we allocated an image buffer, free it! - STBI_FREE(g.out); - } - - // free buffers needed for multiple frame loading; - STBI_FREE(g.history); - STBI_FREE(g.background); - - return u; -} - -static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp) -{ - return stbi__gif_info_raw(s,x,y,comp); -} -#endif - -// ************************************************************************************************* -// Radiance RGBE HDR loader -// originally by Nicolas Schulz -#ifndef STBI_NO_HDR -static int stbi__hdr_test_core(stbi__context *s, const char *signature) -{ - int i; - for (i=0; signature[i]; ++i) - if (stbi__get8(s) != signature[i]) - return 0; - stbi__rewind(s); - return 1; -} - -static int stbi__hdr_test(stbi__context* s) -{ - int r = stbi__hdr_test_core(s, "#?RADIANCE\n"); - stbi__rewind(s); - if(!r) { - r = stbi__hdr_test_core(s, "#?RGBE\n"); - stbi__rewind(s); - } - return r; -} - -#define STBI__HDR_BUFLEN 1024 -static char *stbi__hdr_gettoken(stbi__context *z, char *buffer) -{ - int len=0; - char c = '\0'; - - c = (char) stbi__get8(z); - - while (!stbi__at_eof(z) && c != '\n') { - buffer[len++] = c; - if (len == STBI__HDR_BUFLEN-1) { - // flush to end of line - while (!stbi__at_eof(z) && stbi__get8(z) != '\n') - ; - break; - } - c = (char) stbi__get8(z); - } - - buffer[len] = 0; - return buffer; -} - -static void stbi__hdr_convert(float *output, stbi_uc *input, int req_comp) -{ - if ( input[3] != 0 ) { - float f1; - // Exponent - f1 = (float) ldexp(1.0f, input[3] - (int)(128 + 8)); - if (req_comp <= 2) - output[0] = (input[0] + input[1] + input[2]) * f1 / 3; - else { - output[0] = input[0] * f1; - output[1] = input[1] * f1; - output[2] = input[2] * f1; - } - if (req_comp == 2) output[1] = 1; - if (req_comp == 4) output[3] = 1; - } else { - switch (req_comp) { - case 4: output[3] = 1; /* fallthrough */ - case 3: output[0] = output[1] = output[2] = 0; - break; - case 2: output[1] = 1; /* fallthrough */ - case 1: output[0] = 0; - break; - } - } -} - -static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - char buffer[STBI__HDR_BUFLEN]; - char *token; - int valid = 0; - int width, height; - stbi_uc *scanline; - float *hdr_data; - int len; - unsigned char count, value; - int i, j, k, c1,c2, z; - const char *headerToken; - STBI_NOTUSED(ri); - - // Check identifier - headerToken = stbi__hdr_gettoken(s,buffer); - if (strcmp(headerToken, "#?RADIANCE") != 0 && strcmp(headerToken, "#?RGBE") != 0) - return stbi__errpf("not HDR", "Corrupt HDR image"); - - // Parse header - for(;;) { - token = stbi__hdr_gettoken(s,buffer); - if (token[0] == 0) break; - if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; - } - - if (!valid) return stbi__errpf("unsupported format", "Unsupported HDR format"); - - // Parse width and height - // can't use sscanf() if we're not using stdio! - token = stbi__hdr_gettoken(s,buffer); - if (strncmp(token, "-Y ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); - token += 3; - height = (int) strtol(token, &token, 10); - while (*token == ' ') ++token; - if (strncmp(token, "+X ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); - token += 3; - width = (int) strtol(token, NULL, 10); - - if (height > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); - if (width > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); - - *x = width; - *y = height; - - if (comp) *comp = 3; - if (req_comp == 0) req_comp = 3; - - if (!stbi__mad4sizes_valid(width, height, req_comp, sizeof(float), 0)) - return stbi__errpf("too large", "HDR image is too large"); - - // Read data - hdr_data = (float *) stbi__malloc_mad4(width, height, req_comp, sizeof(float), 0); - if (!hdr_data) - return stbi__errpf("outofmem", "Out of memory"); - - // Load image data - // image data is stored as some number of sca - if ( width < 8 || width >= 32768) { - // Read flat data - for (j=0; j < height; ++j) { - for (i=0; i < width; ++i) { - stbi_uc rgbe[4]; - main_decode_loop: - stbi__getn(s, rgbe, 4); - stbi__hdr_convert(hdr_data + j * width * req_comp + i * req_comp, rgbe, req_comp); - } - } - } else { - // Read RLE-encoded data - scanline = NULL; - - for (j = 0; j < height; ++j) { - c1 = stbi__get8(s); - c2 = stbi__get8(s); - len = stbi__get8(s); - if (c1 != 2 || c2 != 2 || (len & 0x80)) { - // not run-length encoded, so we have to actually use THIS data as a decoded - // pixel (note this can't be a valid pixel--one of RGB must be >= 128) - stbi_uc rgbe[4]; - rgbe[0] = (stbi_uc) c1; - rgbe[1] = (stbi_uc) c2; - rgbe[2] = (stbi_uc) len; - rgbe[3] = (stbi_uc) stbi__get8(s); - stbi__hdr_convert(hdr_data, rgbe, req_comp); - i = 1; - j = 0; - STBI_FREE(scanline); - goto main_decode_loop; // yes, this makes no sense - } - len <<= 8; - len |= stbi__get8(s); - if (len != width) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("invalid decoded scanline length", "corrupt HDR"); } - if (scanline == NULL) { - scanline = (stbi_uc *) stbi__malloc_mad2(width, 4, 0); - if (!scanline) { - STBI_FREE(hdr_data); - return stbi__errpf("outofmem", "Out of memory"); - } - } - - for (k = 0; k < 4; ++k) { - int nleft; - i = 0; - while ((nleft = width - i) > 0) { - count = stbi__get8(s); - if (count > 128) { - // Run - value = stbi__get8(s); - count -= 128; - if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } - for (z = 0; z < count; ++z) - scanline[i++ * 4 + k] = value; - } else { - // Dump - if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } - for (z = 0; z < count; ++z) - scanline[i++ * 4 + k] = stbi__get8(s); - } - } - } - for (i=0; i < width; ++i) - stbi__hdr_convert(hdr_data+(j*width + i)*req_comp, scanline + i*4, req_comp); - } - if (scanline) - STBI_FREE(scanline); - } - - return hdr_data; -} - -static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp) -{ - char buffer[STBI__HDR_BUFLEN]; - char *token; - int valid = 0; - int dummy; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - if (stbi__hdr_test(s) == 0) { - stbi__rewind( s ); - return 0; - } - - for(;;) { - token = stbi__hdr_gettoken(s,buffer); - if (token[0] == 0) break; - if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; - } - - if (!valid) { - stbi__rewind( s ); - return 0; - } - token = stbi__hdr_gettoken(s,buffer); - if (strncmp(token, "-Y ", 3)) { - stbi__rewind( s ); - return 0; - } - token += 3; - *y = (int) strtol(token, &token, 10); - while (*token == ' ') ++token; - if (strncmp(token, "+X ", 3)) { - stbi__rewind( s ); - return 0; - } - token += 3; - *x = (int) strtol(token, NULL, 10); - *comp = 3; - return 1; -} -#endif // STBI_NO_HDR - -#ifndef STBI_NO_BMP -static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp) -{ - void *p; - stbi__bmp_data info; - - info.all_a = 255; - p = stbi__bmp_parse_header(s, &info); - if (p == NULL) { - stbi__rewind( s ); - return 0; - } - if (x) *x = s->img_x; - if (y) *y = s->img_y; - if (comp) { - if (info.bpp == 24 && info.ma == 0xff000000) - *comp = 3; - else - *comp = info.ma ? 4 : 3; - } - return 1; -} -#endif - -#ifndef STBI_NO_PSD -static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp) -{ - int channelCount, dummy, depth; - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - if (stbi__get32be(s) != 0x38425053) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 1) { - stbi__rewind( s ); - return 0; - } - stbi__skip(s, 6); - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) { - stbi__rewind( s ); - return 0; - } - *y = stbi__get32be(s); - *x = stbi__get32be(s); - depth = stbi__get16be(s); - if (depth != 8 && depth != 16) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 3) { - stbi__rewind( s ); - return 0; - } - *comp = 4; - return 1; -} - -static int stbi__psd_is16(stbi__context *s) -{ - int channelCount, depth; - if (stbi__get32be(s) != 0x38425053) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 1) { - stbi__rewind( s ); - return 0; - } - stbi__skip(s, 6); - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) { - stbi__rewind( s ); - return 0; - } - STBI_NOTUSED(stbi__get32be(s)); - STBI_NOTUSED(stbi__get32be(s)); - depth = stbi__get16be(s); - if (depth != 16) { - stbi__rewind( s ); - return 0; - } - return 1; -} -#endif - -#ifndef STBI_NO_PIC -static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp) -{ - int act_comp=0,num_packets=0,chained,dummy; - stbi__pic_packet packets[10]; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) { - stbi__rewind(s); - return 0; - } - - stbi__skip(s, 88); - - *x = stbi__get16be(s); - *y = stbi__get16be(s); - if (stbi__at_eof(s)) { - stbi__rewind( s); - return 0; - } - if ( (*x) != 0 && (1 << 28) / (*x) < (*y)) { - stbi__rewind( s ); - return 0; - } - - stbi__skip(s, 8); - - do { - stbi__pic_packet *packet; - - if (num_packets==sizeof(packets)/sizeof(packets[0])) - return 0; - - packet = &packets[num_packets++]; - chained = stbi__get8(s); - packet->size = stbi__get8(s); - packet->type = stbi__get8(s); - packet->channel = stbi__get8(s); - act_comp |= packet->channel; - - if (stbi__at_eof(s)) { - stbi__rewind( s ); - return 0; - } - if (packet->size != 8) { - stbi__rewind( s ); - return 0; - } - } while (chained); - - *comp = (act_comp & 0x10 ? 4 : 3); - - return 1; -} -#endif - -// ************************************************************************************************* -// Portable Gray Map and Portable Pixel Map loader -// by Ken Miller -// -// PGM: http://netpbm.sourceforge.net/doc/pgm.html -// PPM: http://netpbm.sourceforge.net/doc/ppm.html -// -// Known limitations: -// Does not support comments in the header section -// Does not support ASCII image data (formats P2 and P3) - -#ifndef STBI_NO_PNM - -static int stbi__pnm_test(stbi__context *s) -{ - char p, t; - p = (char) stbi__get8(s); - t = (char) stbi__get8(s); - if (p != 'P' || (t != '5' && t != '6')) { - stbi__rewind( s ); - return 0; - } - return 1; -} - -static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *out; - STBI_NOTUSED(ri); - - ri->bits_per_channel = stbi__pnm_info(s, (int *)&s->img_x, (int *)&s->img_y, (int *)&s->img_n); - if (ri->bits_per_channel == 0) - return 0; - - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - *x = s->img_x; - *y = s->img_y; - if (comp) *comp = s->img_n; - - if (!stbi__mad4sizes_valid(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0)) - return stbi__errpuc("too large", "PNM too large"); - - out = (stbi_uc *) stbi__malloc_mad4(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0); - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - if (!stbi__getn(s, out, s->img_n * s->img_x * s->img_y * (ri->bits_per_channel / 8))) { - STBI_FREE(out); - return stbi__errpuc("bad PNM", "PNM file truncated"); - } - - if (req_comp && req_comp != s->img_n) { - if (ri->bits_per_channel == 16) { - out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, s->img_n, req_comp, s->img_x, s->img_y); - } else { - out = stbi__convert_format(out, s->img_n, req_comp, s->img_x, s->img_y); - } - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - return out; -} - -static int stbi__pnm_isspace(char c) -{ - return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r'; -} - -static void stbi__pnm_skip_whitespace(stbi__context *s, char *c) -{ - for (;;) { - while (!stbi__at_eof(s) && stbi__pnm_isspace(*c)) - *c = (char) stbi__get8(s); - - if (stbi__at_eof(s) || *c != '#') - break; - - while (!stbi__at_eof(s) && *c != '\n' && *c != '\r' ) - *c = (char) stbi__get8(s); - } -} - -static int stbi__pnm_isdigit(char c) -{ - return c >= '0' && c <= '9'; -} - -static int stbi__pnm_getinteger(stbi__context *s, char *c) -{ - int value = 0; - - while (!stbi__at_eof(s) && stbi__pnm_isdigit(*c)) { - value = value*10 + (*c - '0'); - *c = (char) stbi__get8(s); - if((value > 214748364) || (value == 214748364 && *c > '7')) - return stbi__err("integer parse overflow", "Parsing an integer in the PPM header overflowed a 32-bit int"); - } - - return value; -} - -static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp) -{ - int maxv, dummy; - char c, p, t; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - stbi__rewind(s); - - // Get identifier - p = (char) stbi__get8(s); - t = (char) stbi__get8(s); - if (p != 'P' || (t != '5' && t != '6')) { - stbi__rewind(s); - return 0; - } - - *comp = (t == '6') ? 3 : 1; // '5' is 1-component .pgm; '6' is 3-component .ppm - - c = (char) stbi__get8(s); - stbi__pnm_skip_whitespace(s, &c); - - *x = stbi__pnm_getinteger(s, &c); // read width - if(*x == 0) - return stbi__err("invalid width", "PPM image header had zero or overflowing width"); - stbi__pnm_skip_whitespace(s, &c); - - *y = stbi__pnm_getinteger(s, &c); // read height - if (*y == 0) - return stbi__err("invalid width", "PPM image header had zero or overflowing width"); - stbi__pnm_skip_whitespace(s, &c); - - maxv = stbi__pnm_getinteger(s, &c); // read max value - if (maxv > 65535) - return stbi__err("max value > 65535", "PPM image supports only 8-bit and 16-bit images"); - else if (maxv > 255) - return 16; - else - return 8; -} - -static int stbi__pnm_is16(stbi__context *s) -{ - if (stbi__pnm_info(s, NULL, NULL, NULL) == 16) - return 1; - return 0; -} -#endif - -static int stbi__info_main(stbi__context *s, int *x, int *y, int *comp) -{ - #ifndef STBI_NO_JPEG - if (stbi__jpeg_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PNG - if (stbi__png_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_GIF - if (stbi__gif_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_BMP - if (stbi__bmp_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PSD - if (stbi__psd_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PIC - if (stbi__pic_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PNM - if (stbi__pnm_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_HDR - if (stbi__hdr_info(s, x, y, comp)) return 1; - #endif - - // test tga last because it's a crappy test! - #ifndef STBI_NO_TGA - if (stbi__tga_info(s, x, y, comp)) - return 1; - #endif - return stbi__err("unknown image type", "Image not of any known type, or corrupt"); -} - -static int stbi__is_16_main(stbi__context *s) -{ - #ifndef STBI_NO_PNG - if (stbi__png_is16(s)) return 1; - #endif - - #ifndef STBI_NO_PSD - if (stbi__psd_is16(s)) return 1; - #endif - - #ifndef STBI_NO_PNM - if (stbi__pnm_is16(s)) return 1; - #endif - return 0; -} - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_info(char const *filename, int *x, int *y, int *comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result; - if (!f) return stbi__err("can't fopen", "Unable to open file"); - result = stbi_info_from_file(f, x, y, comp); - fclose(f); - return result; -} - -STBIDEF int stbi_info_from_file(FILE *f, int *x, int *y, int *comp) -{ - int r; - stbi__context s; - long pos = ftell(f); - stbi__start_file(&s, f); - r = stbi__info_main(&s,x,y,comp); - fseek(f,pos,SEEK_SET); - return r; -} - -STBIDEF int stbi_is_16_bit(char const *filename) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result; - if (!f) return stbi__err("can't fopen", "Unable to open file"); - result = stbi_is_16_bit_from_file(f); - fclose(f); - return result; -} - -STBIDEF int stbi_is_16_bit_from_file(FILE *f) -{ - int r; - stbi__context s; - long pos = ftell(f); - stbi__start_file(&s, f); - r = stbi__is_16_main(&s); - fseek(f,pos,SEEK_SET); - return r; -} -#endif // !STBI_NO_STDIO - -STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__info_main(&s,x,y,comp); -} - -STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *c, void *user, int *x, int *y, int *comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); - return stbi__info_main(&s,x,y,comp); -} - -STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__is_16_main(&s); -} - -STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *c, void *user) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); - return stbi__is_16_main(&s); -} - -#endif // STB_IMAGE_IMPLEMENTATION - -/* - revision history: - 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs - 2.19 (2018-02-11) fix warning - 2.18 (2018-01-30) fix warnings - 2.17 (2018-01-29) change sbti__shiftsigned to avoid clang -O2 bug - 1-bit BMP - *_is_16_bit api - avoid warnings - 2.16 (2017-07-23) all functions have 16-bit variants; - STBI_NO_STDIO works again; - compilation fixes; - fix rounding in unpremultiply; - optimize vertical flip; - disable raw_len validation; - documentation fixes - 2.15 (2017-03-18) fix png-1,2,4 bug; now all Imagenet JPGs decode; - warning fixes; disable run-time SSE detection on gcc; - uniform handling of optional "return" values; - thread-safe initialization of zlib tables - 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs - 2.13 (2016-11-29) add 16-bit API, only supported for PNG right now - 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes - 2.11 (2016-04-02) allocate large structures on the stack - remove white matting for transparent PSD - fix reported channel count for PNG & BMP - re-enable SSE2 in non-gcc 64-bit - support RGB-formatted JPEG - read 16-bit PNGs (only as 8-bit) - 2.10 (2016-01-22) avoid warning introduced in 2.09 by STBI_REALLOC_SIZED - 2.09 (2016-01-16) allow comments in PNM files - 16-bit-per-pixel TGA (not bit-per-component) - info() for TGA could break due to .hdr handling - info() for BMP to shares code instead of sloppy parse - can use STBI_REALLOC_SIZED if allocator doesn't support realloc - code cleanup - 2.08 (2015-09-13) fix to 2.07 cleanup, reading RGB PSD as RGBA - 2.07 (2015-09-13) fix compiler warnings - partial animated GIF support - limited 16-bpc PSD support - #ifdef unused functions - bug with < 92 byte PIC,PNM,HDR,TGA - 2.06 (2015-04-19) fix bug where PSD returns wrong '*comp' value - 2.05 (2015-04-19) fix bug in progressive JPEG handling, fix warning - 2.04 (2015-04-15) try to re-enable SIMD on MinGW 64-bit - 2.03 (2015-04-12) extra corruption checking (mmozeiko) - stbi_set_flip_vertically_on_load (nguillemot) - fix NEON support; fix mingw support - 2.02 (2015-01-19) fix incorrect assert, fix warning - 2.01 (2015-01-17) fix various warnings; suppress SIMD on gcc 32-bit without -msse2 - 2.00b (2014-12-25) fix STBI_MALLOC in progressive JPEG - 2.00 (2014-12-25) optimize JPG, including x86 SSE2 & NEON SIMD (ryg) - progressive JPEG (stb) - PGM/PPM support (Ken Miller) - STBI_MALLOC,STBI_REALLOC,STBI_FREE - GIF bugfix -- seemingly never worked - STBI_NO_*, STBI_ONLY_* - 1.48 (2014-12-14) fix incorrectly-named assert() - 1.47 (2014-12-14) 1/2/4-bit PNG support, both direct and paletted (Omar Cornut & stb) - optimize PNG (ryg) - fix bug in interlaced PNG with user-specified channel count (stb) - 1.46 (2014-08-26) - fix broken tRNS chunk (colorkey-style transparency) in non-paletted PNG - 1.45 (2014-08-16) - fix MSVC-ARM internal compiler error by wrapping malloc - 1.44 (2014-08-07) - various warning fixes from Ronny Chevalier - 1.43 (2014-07-15) - fix MSVC-only compiler problem in code changed in 1.42 - 1.42 (2014-07-09) - don't define _CRT_SECURE_NO_WARNINGS (affects user code) - fixes to stbi__cleanup_jpeg path - added STBI_ASSERT to avoid requiring assert.h - 1.41 (2014-06-25) - fix search&replace from 1.36 that messed up comments/error messages - 1.40 (2014-06-22) - fix gcc struct-initialization warning - 1.39 (2014-06-15) - fix to TGA optimization when req_comp != number of components in TGA; - fix to GIF loading because BMP wasn't rewinding (whoops, no GIFs in my test suite) - add support for BMP version 5 (more ignored fields) - 1.38 (2014-06-06) - suppress MSVC warnings on integer casts truncating values - fix accidental rename of 'skip' field of I/O - 1.37 (2014-06-04) - remove duplicate typedef - 1.36 (2014-06-03) - convert to header file single-file library - if de-iphone isn't set, load iphone images color-swapped instead of returning NULL - 1.35 (2014-05-27) - various warnings - fix broken STBI_SIMD path - fix bug where stbi_load_from_file no longer left file pointer in correct place - fix broken non-easy path for 32-bit BMP (possibly never used) - TGA optimization by Arseny Kapoulkine - 1.34 (unknown) - use STBI_NOTUSED in stbi__resample_row_generic(), fix one more leak in tga failure case - 1.33 (2011-07-14) - make stbi_is_hdr work in STBI_NO_HDR (as specified), minor compiler-friendly improvements - 1.32 (2011-07-13) - support for "info" function for all supported filetypes (SpartanJ) - 1.31 (2011-06-20) - a few more leak fixes, bug in PNG handling (SpartanJ) - 1.30 (2011-06-11) - added ability to load files via callbacks to accomidate custom input streams (Ben Wenger) - removed deprecated format-specific test/load functions - removed support for installable file formats (stbi_loader) -- would have been broken for IO callbacks anyway - error cases in bmp and tga give messages and don't leak (Raymond Barbiero, grisha) - fix inefficiency in decoding 32-bit BMP (David Woo) - 1.29 (2010-08-16) - various warning fixes from Aurelien Pocheville - 1.28 (2010-08-01) - fix bug in GIF palette transparency (SpartanJ) - 1.27 (2010-08-01) - cast-to-stbi_uc to fix warnings - 1.26 (2010-07-24) - fix bug in file buffering for PNG reported by SpartanJ - 1.25 (2010-07-17) - refix trans_data warning (Won Chun) - 1.24 (2010-07-12) - perf improvements reading from files on platforms with lock-heavy fgetc() - minor perf improvements for jpeg - deprecated type-specific functions so we'll get feedback if they're needed - attempt to fix trans_data warning (Won Chun) - 1.23 fixed bug in iPhone support - 1.22 (2010-07-10) - removed image *writing* support - stbi_info support from Jetro Lauha - GIF support from Jean-Marc Lienher - iPhone PNG-extensions from James Brown - warning-fixes from Nicolas Schulz and Janez Zemva (i.stbi__err. Janez (U+017D)emva) - 1.21 fix use of 'stbi_uc' in header (reported by jon blow) - 1.20 added support for Softimage PIC, by Tom Seddon - 1.19 bug in interlaced PNG corruption check (found by ryg) - 1.18 (2008-08-02) - fix a threading bug (local mutable static) - 1.17 support interlaced PNG - 1.16 major bugfix - stbi__convert_format converted one too many pixels - 1.15 initialize some fields for thread safety - 1.14 fix threadsafe conversion bug - header-file-only version (#define STBI_HEADER_FILE_ONLY before including) - 1.13 threadsafe - 1.12 const qualifiers in the API - 1.11 Support installable IDCT, colorspace conversion routines - 1.10 Fixes for 64-bit (don't use "unsigned long") - optimized upsampling by Fabian "ryg" Giesen - 1.09 Fix format-conversion for PSD code (bad global variables!) - 1.08 Thatcher Ulrich's PSD code integrated by Nicolas Schulz - 1.07 attempt to fix C++ warning/errors again - 1.06 attempt to fix C++ warning/errors again - 1.05 fix TGA loading to return correct *comp and use good luminance calc - 1.04 default float alpha is 1, not 255; use 'void *' for stbi_image_free - 1.03 bugfixes to STBI_NO_STDIO, STBI_NO_HDR - 1.02 support for (subset of) HDR files, float interface for preferred access to them - 1.01 fix bug: possible bug in handling right-side up bmps... not sure - fix bug: the stbi__bmp_load() and stbi__tga_load() functions didn't work at all - 1.00 interface to zlib that skips zlib header - 0.99 correct handling of alpha in palette - 0.98 TGA loader by lonesock; dynamically add loaders (untested) - 0.97 jpeg errors on too large a file; also catch another malloc failure - 0.96 fix detection of invalid v value - particleman@mollyrocket forum - 0.95 during header scan, seek to markers in case of padding - 0.94 STBI_NO_STDIO to disable stdio usage; rename all #defines the same - 0.93 handle jpegtran output; verbose errors - 0.92 read 4,8,16,24,32-bit BMP files of several formats - 0.91 output 24-bit Windows 3.0 BMP files - 0.90 fix a few more warnings; bump version number to approach 1.0 - 0.61 bugfixes due to Marc LeBlanc, Christopher Lloyd - 0.60 fix compiling as c++ - 0.59 fix warnings: merge Dave Moore's -Wall fixes - 0.58 fix bug: zlib uncompressed mode len/nlen was wrong endian - 0.57 fix bug: jpg last huffman symbol before marker was >9 bits but less than 16 available - 0.56 fix bug: zlib uncompressed mode len vs. nlen - 0.55 fix bug: restart_interval not initialized to 0 - 0.54 allow NULL for 'int *comp' - 0.53 fix bug in png 3->4; speedup png decoding - 0.52 png handles req_comp=3,4 directly; minor cleanup; jpeg comments - 0.51 obey req_comp requests, 1-component jpegs return as 1-component, - on 'test' only check type, not whether we support this variant - 0.50 (2006-11-19) - first released version -*/ - - -/* ------------------------------------------------------------------------------- -This software is available under 2 licenses -- choose whichever you prefer. ------------------------------------------------------------------------------- -ALTERNATIVE A - MIT License -Copyright (c) 2017 Sean Barrett -Permission is hereby granted, free of charge, to any person obtaining a copy of -this software and associated documentation files (the "Software"), to deal in -the Software without restriction, including without limitation the rights to -use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies -of the Software, and to permit persons to whom the Software is furnished to do -so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. ------------------------------------------------------------------------------- -ALTERNATIVE B - Public Domain (www.unlicense.org) -This is free and unencumbered software released into the public domain. -Anyone is free to copy, modify, publish, use, compile, sell, or distribute this -software, either in source code form or as a compiled binary, for any purpose, -commercial or non-commercial, and by any means. -In jurisdictions that recognize copyright laws, the author or authors of this -software dedicate any and all copyright interest in the software to the public -domain. We make this dedication for the benefit of the public at large and to -the detriment of our heirs and successors. We intend this dedication to be an -overt act of relinquishment in perpetuity of all present and future rights to -this software under copyright law. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN -ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION -WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ------------------------------------------------------------------------------- -*/ diff --git a/zig/vendor/stb/stb_image_impl.c b/zig/vendor/stb/stb_image_impl.c deleted file mode 100644 index 908982ad8f..0000000000 --- a/zig/vendor/stb/stb_image_impl.c +++ /dev/null @@ -1,10 +0,0 @@ -// stb_image v2.30 (public domain / MIT, github.com/nothings/stb at 2c980bb59875b0d32144a71867fbdebb2f77cd20): -// the native server's JPEG and PNG decoder for image inputs (zig/src/vision/qwen_image.zig). -#define STB_IMAGE_IMPLEMENTATION -#define STBI_ONLY_JPEG -#define STBI_ONLY_PNG -#define STBI_NO_STDIO -#define STBI_NO_HDR -#define STBI_NO_LINEAR -#define STBI_NO_FAILURE_STRINGS -#include "stb_image.h" From 8570e0688ae6118e0e90eee0655940aeef7614bf Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:16:46 +0000 Subject: [PATCH 31/46] qwen3_5: attention chunks follow each row's own keys, not the forward's widest stream The chunk width came from the longest stream or window in the forward, so a 201-key row cut its keys into 128-key chunks alone and 256-key chunks beside a stream past 8192, and a window crossing 8192 recut its earlier rows; either changed the fp32 sums. Each row now takes chunk_keys(its logical keys), the width its one-row step uses, in both the partial and the merge kernels, and the partials sit at a fixed max_chunks stride. tf-qwen35-cuda takes TF_Q35_PREFILL=N (start after N cycled ids) and TF_Q35_BESIDE=N (step an N-token stream in each forward), which cover both cases: before this change, windows from 8151 keys kept 36 of 120 rows and trees 35, and a 151-key stream beside one from 8189 kept 4 to 11 of 120; after it all of them equal one-row steps on the 2B and 27B. --- zig/kernels/cuda/qwen35.cu | 25 +++++++----- zig/kernels/cuda/qwen35.cuh | 1 + zig/src/families/qwen3_5/cuda_engine.zig | 15 +++---- zig/tests/qwen35_cuda_main.zig | 50 ++++++++++++++++++++++-- 4 files changed, 68 insertions(+), 23 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 558baa7908..9217a19fa3 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -310,6 +310,9 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i // A row's ancestors-and-self: its tree window's row bits, or with the top bit set a chain row's index in its window. __device__ __forceinline__ int rows_seen(unsigned anc) { return (anc >> 31) ? (int)(anc & 0x7fffffffu) + 1 : __popc(anc); } +// A row's chunk width from its own logical keys alone, as its one-row step's: at most MAXCH chunks, 128-key multiples. +__device__ __forceinline__ int chunk_keys(int keys) { return max(128, (keys + MAXCH * 128 - 1) / (MAXCH * 128) * 128); } + // The window row holding a row's nth seen window key. __device__ __forceinline__ int seen_row(unsigned anc, int n) { if (anc >> 31) return n; @@ -342,24 +345,24 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo } } -// Causal attention over the cache in P[0] chunks of P[1] keys, a warp a query head; partials by (row, head, chunk). +// Causal attention in chunks of chunk_keys(the row's keys), a warp a query head; partials by (row, head, chunk). extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, const int* __restrict__ P, float scale, float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; - const int nch = P[0], ck = P[1]; const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; - if (c >= nch || past_rows(P, blockIdx.z)) return; + if (c >= MAXCH || past_rows(P, blockIdx.z)) return; const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; const unsigned anc = rowanc[m]; // chunks cut at logical positions (a window key's: the stream's length plus its ancestors), as a step's are - const int k0 = c * ck, k1 = min(k0 + ck, sg.len + rows_seen(anc)); - const size_t slot = (m * QH + qh) * nch + c; + const int keys = sg.len + rows_seen(anc), ck = chunk_keys(keys); + const int k0 = c * ck, k1 = min(k0 + ck, keys); + const size_t slot = (m * QH + qh) * MAXCH + c; if (k0 >= k1) { if (lane == 0) { PM[slot] = -INFINITY; @@ -401,17 +404,19 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; } -// The P[0] chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +// A row's chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, - const float* __restrict__ PO, const int* __restrict__ P, - __nv_bfloat16* __restrict__ OUT) { + const float* __restrict__ PO, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, + const int* __restrict__ P, __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; - const int nch = P[0]; const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; if (past_rows(P, blockIdx.y)) return; const size_t m = blockIdx.y; - const size_t base = (m * QH + qh) * nch; + const int keys = segs[rowseg[m]].len + rows_seen(rowanc[m]); + const int ck = chunk_keys(keys), nch = (keys + ck - 1) / ck; + const size_t base = (m * QH + qh) * MAXCH; float mx = -INFINITY; for (int c = 0; c < nch; c++) if (PL[base + c] > 0.0f) mx = fmaxf(mx, PM[base + c]); diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index b4f789814f..5df37e308e 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -11,6 +11,7 @@ constexpr int TAPS = 4; // causal convolution taps constexpr int HD = 256; // attention head width constexpr int MAXQ = 32; // attention query heads at most constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) +constexpr int MAXCH = 64; // attention chunks a row at most (cuda_engine.zig's max_chunks) // The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. struct Geo { diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 389b32f766..15975b8d42 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -22,7 +22,7 @@ pub const max_segs = 32; pub const most_verify = 32; /// A window's rows at most (exact: one-row steps' bits; qwen35.cuh's MAXW, the rows a stream's log holds). pub const most_window = 16; -/// Partial attention chunks a row at most: the chunk width grows with the window instead. +/// Partial attention chunks a row at most (qwen35.cuh's MAXCH): each row's chunk width grows with its own keys. const max_chunks = 64; /// Device allocations are 256-byte aligned, as the driver's are. @@ -147,7 +147,7 @@ pub const Engine = struct { const rules_at = rowanc_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows, record + const params_at = tokens_at + most_verify * 4; // two unused words, the rows, the windows, record const tables_len = params_at + 8 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { @@ -343,12 +343,10 @@ pub const Engine = struct { pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { const g = e.g; var total: usize = 0; - var widest: usize = 0; for (parts) |p| { try p.cache.settle(e); if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; total += p.rows; - widest = @max(widest, p.cache.len + p.rows); } if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; const record = keep != .last; @@ -378,11 +376,10 @@ pub const Engine = struct { m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 } } - // attention chunks: at most max_chunks a row, each a multiple of 128 keys - const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); + // words 0 and 1 are reserved (each row's attention chunks follow its own keys: qwen35.cu's chunk_keys) const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..32]); - params[0] = @intCast((widest + ck - 1) / ck); - params[1] = @intCast(ck); + params[0] = 0; + params[1] = 0; params[2] = @intCast(total); params[3] = @intCast(parts.len); params[4] = @intFromBool(record); @@ -497,7 +494,7 @@ pub const Engine = struct { try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); - try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_segs, t_rows, t_anc, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 0d1f711252..e0bb3ac025 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -20,15 +20,23 @@ pub fn main(init: std.process.Init) !u8 { defer ctx.deinit(); const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); defer e.deinit(); - var cache = try E.Cache.init(init.gpa, e, 4096); + // TF_Q35_PREFILL=N starts the stream after N cycled ids; TF_Q35_BESIDE=N steps an N-token stream beside it + const prefill = try envCount("TF_Q35_PREFILL"); + const beside = try envCount("TF_Q35_BESIDE"); + var cache = try E.Cache.init(init.gpa, e, 4096 + prefill); defer cache.deinit(init.gpa); + try fill(e, &cache, ids.items, prefill); + var other = try E.Cache.init(init.gpa, e, 4096 + beside); + defer other.deinit(init.gpa); + try fill(e, &other, ids.items, beside); + var f: Fwd = .{ .e = e, .other = if (beside > 0) &other else null, .ids = ids.items }; var out: std.ArrayList(u8) = .empty; if (std.mem.eql(u8, args[2], "step")) { var first: i64 = 0; var rest: i64 = 0; for (ids.items, 0..) |id, i| { const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); - try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, &.{id}, .logits); + try f.run(.{ .cache = &cache, .rows = 1 }, &.{id}); const dt: i64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); if (i == 0) first = dt else rest += dt; @memcpy(cache.logits, e.rowLogits(0)); @@ -42,7 +50,7 @@ pub fn main(init: std.process.Init) !u8 { var w: usize = 1; while (at < ids.items.len) : (w = w % e.window_rows + 1) { const n = @min(w, ids.items.len - at); - try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .logits); + try f.run(.{ .cache = &cache, .rows = n }, ids.items[at..][0..n]); for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); at += n; } @@ -76,7 +84,7 @@ pub fn main(init: std.process.Init) !u8 { parents[n] = 1; n += 1; } - try e.forward(&.{.{ .cache = &cache, .rows = n, .parents = parents[0..n] }}, toks[0..n], .logits); + try f.run(.{ .cache = &cache, .rows = n, .parents = parents[0..n] }, toks[0..n]); for (path[0..kept]) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); if (n > 1) try cache.keep(e, path[0..kept]); at += kept; @@ -93,3 +101,37 @@ pub fn main(init: std.process.Init) !u8 { try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = out.items }); return 0; } + +/// One forward keeping logits: the stream's part first (its rows' logits come first), then the second stream's step. +const Fwd = struct { + e: *E.Engine, + other: ?*E.Cache, + ids: []const u32, + n: usize = 0, + + fn run(f: *Fwd, part: E.Part, toks: []const u32) !void { + const o = f.other orelse return f.e.forward(&.{part}, toks, .logits); + var all: [E.most_window + 1]u32 = undefined; + @memcpy(all[0..toks.len], toks); + all[toks.len] = f.ids[f.n % f.ids.len]; + f.n += 1; + try f.e.forward(&.{ part, .{ .cache = o, .rows = 1 } }, all[0 .. toks.len + 1], .logits); + } +}; + +fn envCount(name: [:0]const u8) !usize { + const v = std.c.getenv(name) orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10); +} + +/// `n` positions of the ids cycled, through prompt forwards. +fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { + var buf: [E.max_rows]u32 = undefined; + var at: usize = 0; + while (at < n) { + const k = @min(E.max_rows, n - at); + for (buf[0..k], at..) |*t, i| t.* = ids[i % ids.len]; + try e.forward(&.{.{ .cache = cache, .rows = k }}, buf[0..k], .last); + at += k; + } +} From 5e1c3feaa79e40f7659537d96b939129abe39df9 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:16:46 +0000 Subject: [PATCH 32/46] qwen3_5: the DFlash2 config's divisors must be positive, and staging fits its selector rank The integer reader accepted zero, so a conv_group_size or num_key_value_heads of 0 reached a remainder at admission; sizes and divisors are now positive (mask_token_id alone may be 0). The host staging area held 16 rows at rank 256, so a larger selector_rank overran it; it is now sized for the admitted rank. Config.parse takes the bytes, so a host test covers both. --- zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_draft.zig | 42 +++++++++++++++++++++++-- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 435ebe162e..47158adda9 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -10,4 +10,5 @@ test { _ = geometry; _ = @import("draw.zig"); _ = @import("cuda_draft_weights.zig"); + _ = draft; } diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index cfb8670cf5..7f0e4e61be 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -58,6 +58,10 @@ pub const Config = struct { defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); defer gpa.free(bytes); + return parse(gpa, bytes, target); + } + + fn parse(gpa: std.mem.Allocator, bytes: []const u8, target: engine.Geometry) !Config { const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); defer doc.deinit(); if (doc.value != .object) return error.BadDraftConfig; @@ -73,7 +77,7 @@ pub const Config = struct { .kv_heads = try int(o, "num_key_value_heads"), .head_dim = try int(o, "head_dim"), .block = try int(d.object, "block_size"), - .mask = @intCast(try int(d.object, "mask_token_id")), + .mask = @intCast(try whole(d.object, "mask_token_id")), .window = try int(o, "sliding_window"), .group = try int(d.object, "conv_group_size"), .rank = try int(d.object, "selector_rank"), @@ -98,7 +102,15 @@ pub const Config = struct { } }; +/// A size, count or divisor: positive, so a zero can't reach a remainder or an empty allocation. fn int(o: std.json.ObjectMap, key: []const u8) !usize { + const v = try whole(o, key); + if (v == 0) return error.BadDraftConfig; + return v; +} + +/// A token id: any whole number. +fn whole(o: std.json.ObjectMap, key: []const u8) !usize { const v = o.get(key) orelse return error.BadDraftConfig; if (v != .integer or v.integer < 0) return error.BadDraftConfig; return @intCast(v.integer); @@ -176,7 +188,10 @@ pub const Draft = struct { const words_at = 64; const hp_at = words_at + 16 * lp_words * 4; - const staged_len = hp_at + 16 * 256 * 2; + /// The host staging area: the words then 16 rows of `rank` selector projections (the block's mask rows fit). + fn stagedLen(rank: usize) usize { + return hp_at + 16 * rank * 2; + } pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8) !*Draft { const c = try Config.read(gpa, io, dir, e.g); @@ -193,7 +208,7 @@ pub const Draft = struct { try d.scratch.fill8(0, null); d.at = 0; _ = try d.layout(true); - d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); + d.staged = try cuda.HostBuffer.alloc(e.d, stagedLen(c.rank)); errdefer d.staged.free(); try d.picksZero(); // the first context, and a hold timed on it (the quickest of several: the first ones pay the driver's stalls) @@ -546,3 +561,24 @@ pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); std.debug.print("[timing] {s} {d:.2} ms at {d:.1}\n", .{ what, ns / 1e6, at / 1e6 }); } + +test "a drafter config with a zero divisor is refused, and any admitted rank's selector rows fit the staging area" { + const gpa = std.testing.allocator; + const target: engine.Geometry = .{ .hidden = 5120, .intermediate = 17408, .layers = 64, .vocab = 248320, .k_heads = 16, .v_heads = 48, .q_heads = 24, .kv_heads = 4, .eps = 1e-6, .tied = false, .context = 262144, .linear = @splat(false) }; + const text = + \\{"dflash_config":{"block_size":8,"conv_group_size":16,"conv_kernel_size":2,"mask_token_id":248070,"selector_rank":256, + \\ "selector_top_k":16,"target_layer_ids":[5,19,33,47,61]},"head_dim":128,"hidden_size":5120,"intermediate_size":17408, + \\ "num_attention_heads":32,"num_hidden_layers":5,"num_key_value_heads":8,"rms_norm_eps":1e-06, + \\ "rope_parameters":{"rope_theta":10000000},"sliding_window":2048,"vocab_size":248320} + ; + try std.testing.expectEqual(@as(usize, 256), (try Config.parse(gpa, text, target)).rank); + for ([_][2][]const u8{ .{ "\"conv_group_size\":16", "\"conv_group_size\":0" }, .{ "\"num_key_value_heads\":8", "\"num_key_value_heads\":0" } }) |r| { + const bad = try std.mem.replaceOwned(u8, gpa, text, r[0], r[1]); + defer gpa.free(bad); + try std.testing.expectError(error.BadDraftConfig, Config.parse(gpa, bad, target)); + } + const wide = try std.mem.replaceOwned(u8, gpa, text, "\"selector_rank\":256", "\"selector_rank\":1024"); + defer gpa.free(wide); + const c = try Config.parse(gpa, wide, target); + try std.testing.expect(Draft.stagedLen(c.rank) >= Draft.hp_at + (c.block - 1) * c.rank * 2); +} From 52b82fc59954ec299ebd597ced7aaa6a32eea759 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 06:15:42 -0500 Subject: [PATCH 33/46] prompt cache: an entry removed while pinned is dead to lookups and freed only at its last unpin; a job reports resumed before it unpins --- zig/src/core/lane_host.zig | 2 +- zig/src/core/prompt_cache.zig | 13 +++++---- zig/src/core/prompt_cache_test.zig | 47 ++++++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 6 deletions(-) diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 9ddbf6a8fe..73937e925c 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -85,7 +85,7 @@ pub const LaneHost = struct { const e = job.entry orelse return; job.entry = null; const store = job.host.cache orelse return; - store.unpin(e); + defer store.unpin(e); // after resumed, which reads e if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index 6fe20982f7..81d5cf4ed4 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -109,6 +109,7 @@ pub const Entry = struct { shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded pins: u32 = 0, // admitted jobs that restore it and have not yet: eviction passes it over + dead: bool = false, // removed while pinned: no lookup finds it, and its last unpin frees it }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -149,6 +150,7 @@ pub const Store = struct { } fn free(s: *Store, e: *Entry) void { + s.held -= e.bytes; s.family.vtable.drop(s.family.ptr, e.saved); s.gpa.free(e.tokens); s.gpa.free(e.decode_spans); @@ -158,18 +160,19 @@ pub const Store = struct { fn remove(s: *Store, i: usize) void { const e = s.entries.orderedRemove(i); - s.held -= e.bytes; - s.free(e); + e.dead = true; + if (e.pins == 0) s.free(e); } - /// `e` waits for a restore: no lookup or keep evicts it until `unpin` (jobs set up after it run their own first). + /// `e` waits for a restore: nothing frees it until its last `unpin` (removed meanwhile, no lookup finds it). pub fn pin(_: *Store, e: *Entry) *Entry { e.pins += 1; return e; } - pub fn unpin(_: *Store, e: *Entry) void { + pub fn unpin(s: *Store, e: *Entry) void { e.pins -= 1; + if (e.pins == 0 and e.dead) s.free(e); } /// The family's free storage kept for later saves. @@ -484,11 +487,11 @@ pub const Store = struct { const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes; e.* = .{ .tokens = tokens, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared }; owns_decoded = false; + s.held += charged; // free takes it back s.entries.append(s.gpa, e) catch { s.free(e); return s.fail(at, error.OutOfMemory); }; - s.held += charged; s.counts.kept += 1; if (shared) learned.learn(s, e, starts); if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held); // spare storage only inside what the budget leaves diff --git a/zig/src/core/prompt_cache_test.zig b/zig/src/core/prompt_cache_test.zig index b408b3c7ad..cde079b745 100644 --- a/zig/src/core/prompt_cache_test.zig +++ b/zig/src/core/prompt_cache_test.zig @@ -421,3 +421,50 @@ test "an entry pinned for a restore still to come outlives the lookups set up af try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); try std.testing.expectEqual(s.entries.items.len, f.live); } + +test "two jobs pin one entry and the first restore fails: lookups miss it, and it lives until the second job unpins" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const xb = [_]u32{ 5, 5, 5, 5, 8 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + try std.testing.expect(s.pin((try s.lookup(a, &xb, 3, &.{}, &.{}, &.{})).entry.?) == e); + const held = s.held; + s.resumed(e, &xa, false); // the first job's restore failed: it prefilled from 0 + s.unpin(e); + try std.testing.expect(s.find(&xb, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 0), s.entries.items.len); + try std.testing.expectEqual(@as(usize, 1), f.live); // the second job's restore still reads it + try std.testing.expectEqual(held, s.held); + try std.testing.expectEqual(@as(u32, 1), e.pins); + s.resumed(e, &xb, true); // no longer kept: its report changes nothing + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); + try std.testing.expectEqual(@as(u64, 0), s.held); +} + +test "forget leaves a pinned entry to its restore" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + s.forget(&xa, 3, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 1), f.live); + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); +} From 73f06f14661267f91ada96efd7b48270bf33c7e1 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:19:23 +0000 Subject: [PATCH 34/46] qwen3_5: image rows and M-RoPE positions in the CUDA forward A prompt chunk can carry an image prompt's rows (Part.images): the tower's features replace its placeholder rows after the embedding (tf_q35_splice), and tf_q35_rope turns each row by its t, h and w positions (interleaved sections 11, 11, 10) from a staged table. Rows after an image prompt rotate at their position plus the stream's rope_delta (Seg.delta), so decode, verify windows and trees keep their captured graphs. Text rows rotate at one position on all three axes, the same arithmetic as before: their logits keep their bits. --- zig/kernels/cuda/qwen35.cu | 26 +++++++++++++---- zig/kernels/cuda/qwen35.cuh | 2 ++ zig/src/families/qwen3_5/cuda_cache.zig | 7 +++-- zig/src/families/qwen3_5/cuda_engine.zig | 37 +++++++++++++++++++++--- 4 files changed, 61 insertions(+), 11 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 9217a19fa3..df81a2103f 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -24,6 +24,15 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __rest } } +// Image rows into a prompt chunk's embeddings: row m of H (D wide) takes feature row src[m] where src[m] >= 0. +extern "C" __global__ void __launch_bounds__(256) tf_q35_splice(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ F, + const int* __restrict__ src, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const int s = src[blockIdx.x]; + if (s < 0) return; + for (int i = threadIdx.x; i < D; i += blockDim.x) H[(size_t)blockIdx.x * D + i] = F[(size_t)s * D + i]; +} + // h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)); D a multiple of 256 up to 8192; a 256-thread block a row. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, @@ -296,11 +305,14 @@ extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfl } } -// Text RoPE on the first 64 of HD: pairs (i, i + 32) turned by position * theta^(-i / 32). -__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, int position) { +// M-RoPE's interleaved sections (mrope_section 11, 11, 10): frequency i turns by the t, h or w position. +__device__ __forceinline__ int mrope_axis(int i) { return (i % 3 == 1 && i < 33) ? 1 : ((i % 3 == 2 && i < 30) ? 2 : 0); } + +// RoPE on the first 64 of HD: pairs (i, i + 32) turned by pos[axis(i)] * theta^(-i / 32); text has one position. +__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, const int pos[3]) { if (d >= 64) return x[d]; const int i = d % 32; - const float angle = (float)position * powf(10000000.0f, -(float)i / 32.0f); + const float angle = (float)pos[mrope_axis(i)] * powf(10000000.0f, -(float)i / 32.0f); float sn, cs; sincosf(angle, &sn, &cs); const float a = bf(x[i]), b = bf(x[i + 32]); @@ -324,13 +336,17 @@ __device__ __forceinline__ int seen_row(unsigned anc, int n) { extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, - const int* __restrict__ P) { + const int* __restrict__ P, const int* __restrict__ M) { const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; - const int pos = sg.len + rows_seen(rowanc[m]) - 1; + // an image prompt's chunk reads each row's t, h and w positions from M ([3][MAXR]); other rows are shifted text + const int text = sg.len + rows_seen(rowanc[m]) - 1 + sg.delta; + int pos[3] = {text, text, text}; + if (M != nullptr) + for (int a = 0; a < 3; a++) pos[a] = M[a * MAXR + m]; const int slot = sg.len + (int)m - sg.row0; __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index 5df37e308e..4dffc534ce 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -12,6 +12,7 @@ constexpr int HD = 256; // attention head width constexpr int MAXQ = 32; // attention query heads at most constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) constexpr int MAXCH = 64; // attention chunks a row at most (cuda_engine.zig's max_chunks) +constexpr int MAXR = 128; // a forward's rows at most (cuda_engine.zig's max_rows): an image chunk's position rows // The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. struct Geo { @@ -31,6 +32,7 @@ struct Seg { unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each unsigned long long log; // the window's rows' records (MAXW rows of nl layers of log_bytes), for keep int cap, len, row0, rows, tap_rows, taps; + int delta; // an image prompt's rotary shift: rows past it rotate at position + delta }; // A window row's record of one linear layer (Seg.log). diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index b4376c9490..7b482fa9a0 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -11,7 +11,7 @@ const DeviceBuffer = cuda.DeviceBuffer; pub const tap_rows = 2048 + engine.max_rows; /// qwen35.cuh's Seg and Path. -pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int, delta: c_int }; pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; pub const Cache = struct { @@ -28,6 +28,8 @@ pub const Cache = struct { logits: []u16, /// The newest verify window of more than one row, until its path is kept: where it started, its rows. last: ?struct { start: usize, rows: usize } = null, + /// An image prompt's rotary shift (get_rope_index's rope_delta): rows after its prompt rotate at position + delta. + delta: i32 = 0, pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); @@ -53,6 +55,7 @@ pub const Cache = struct { self.len = 0; self.written = 0; self.last = null; + self.delta = 0; try self.state.fill8(0, e.stream.handle); } @@ -67,7 +70,7 @@ pub const Cache = struct { /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { const g = e.g; - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps), .delta = self.delta }; } /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 15975b8d42..ea6437f0da 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -60,13 +60,21 @@ const Fns = struct { mlp_act: cuda.Function, argmax: cuda.Function, tap: cuda.Function, + splice: cuda.Function, }; /// qwen35_keep.cu's kernels (Cache.keep). pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: cuda.Function }; /// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). -pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; +pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null, images: ?*const Images = null }; + +/// An image prompt's rows for its prompt chunks: device feature rows, each prompt row's feature row, its positions. +pub const Images = struct { + features: u64, // bf16 [image rows, hidden] + src: []const i32, // [prompt]: a feature row, or -1 for text + pos: []const i32, // [3][prompt]: t, h and w rotary positions (get_rope_index) +}; /// What a forward keeps: a prompt chunk its last row's logits, a verify window every row's draw or logits and log. pub const Keep = union(enum) { last, draws: []const Draw, logits }; @@ -138,6 +146,8 @@ pub const Engine = struct { /// A drafter's tapped layers (`setTaps`): each row's residual after them lands in its stream's tap ring. taps: [8]u32 = undefined, n_taps: usize = 0, + /// During an image prompt chunk's forward: its feature rows (0: none); the chunk's src and positions are staged. + features: u64 = 0, const ids_at = 0; const segs_at = max_rows * 4; @@ -148,7 +158,9 @@ pub const Engine = struct { const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; const params_at = tokens_at + most_verify * 4; // two unused words, the rows, the windows, record - const tables_len = params_at + 8 * 4; + const src_at = params_at + 8 * 4; // an image chunk's rows' feature rows + const mpos_at = src_at + max_rows * 4; // and their t, h and w positions, max_rows apart (qwen35.cuh's MAXR) + const tables_len = mpos_at + 3 * max_rows * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -166,6 +178,7 @@ pub const Engine = struct { e.arena_at = 0; e.graphs = .empty; e.n_taps = 0; + e.features = 0; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); @@ -355,6 +368,7 @@ pub const Engine = struct { if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; for (parts) |p| if (p.parents != null and !record) return error.InvalidQwenParents; + for (parts) |p| if (p.images != null and record) return error.InvalidQwenImages; const heads: usize = if (record) total else 1; // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps const prompt = !record; @@ -383,7 +397,17 @@ pub const Engine = struct { params[2] = @intCast(total); params[3] = @intCast(parts.len); params[4] = @intFromBool(record); - try e.tables.uploadAsync(0, st, e.stream.handle); + // an image prompt's chunk (one part, its rows at prompt positions len..): feature rows and t, h, w positions + const images: ?*const Images = if (record) null else parts[0].images; + if (images) |im| { + const p0 = parts[0].cache.len; + if (p0 + total > im.src.len or im.pos.len != 3 * im.src.len) return error.InvalidQwenImages; + const n = im.src.len; + @memcpy(std.mem.bytesAsSlice(i32, st[src_at..][0 .. total * 4]), im.src[p0..][0..total]); + const mpos = std.mem.bytesAsSlice(i32, st[mpos_at..][0 .. 3 * max_rows * 4]); + for (0..3) |a| @memcpy(mpos[a * max_rows ..][0..total], im.pos[a * n + p0 ..][0..total]); + } + try e.tables.uploadAsync(0, st[0..if (images != null) tables_len else src_at], e.stream.handle); // text: as before // a captured forward covers a row bucket and every window; kernels read the true counts; a lone step is exact const lone = total == 1 and parts.len == 1; const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; @@ -391,6 +415,8 @@ pub const Engine = struct { const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; if (prompt) { // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length's graph would pay the stall + e.features = if (images) |im| im.features else 0; + defer e.features = 0; try e.body(shape); } else { const gop = try e.graphs.getOrPut(e.gpa, shape.key()); @@ -467,6 +493,9 @@ pub const Engine = struct { const group = g.q_heads / g.kv_heads; try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)), t_params }); + // an image chunk's rows take the tower's features, and rotate at their staged t, h, w positions + const mpos: u64 = if (e.features != 0) e.tables.ptr + mpos_at else 0; + if (e.features != 0) try e.run(e.fns.splice, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.features, e.tables.ptr + src_at, @as(c_int, @intCast(g.hidden)), t_params }); var di: usize = 0; var ai: c_int = 0; for (e.layers, 0..) |l, i| { @@ -492,7 +521,7 @@ pub const Engine = struct { try e.proj(m.v, e.x, e.vraw, total, prompt); try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params, mpos }); try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_segs, t_rows, t_anc, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); From a10b4deb5fdc06cf413588cdb1ed3bd273000a8d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:19:23 +0000 Subject: [PATCH 35/46] qwen3_5: the vision tower on CUDA, and --vision on the Qwen engine The checkpoint's Qwen3_5VisionModel in bf16 (#726's tower and kernels, its Linears on cuBLASLt): width, heads, depth and MLP read from vision_config, so the 2B's tower (24 blocks of 1024, heads of 64) runs beside the 27B's and the 35B-A3B's. Its weights and one 16,384-patch call's scratch are one allocation made at startup and counted as model memory. At prefill a request's images go through the tower; their rows and get_rope_index's positions feed the prompt chunks, and the stream keeps rope_delta. Video parts are refused by name. --- docs/vision.md | 16 +- zig/build/cuda.zig | 1 + zig/kernels/cuda/qwen_vision.cu | 202 ++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda.zig | 3 + zig/src/families/qwen3_5/cuda_native.zig | 32 +- zig/src/families/qwen3_5/cuda_vision.zig | 472 ++++++++++++++++++ .../families/qwen3_5/cuda_vision_prompt.zig | 175 +++++++ 8 files changed, 897 insertions(+), 5 deletions(-) create mode 100644 zig/kernels/cuda/qwen_vision.cu create mode 100644 zig/src/families/qwen3_5/cuda_vision.zig create mode 100644 zig/src/families/qwen3_5/cuda_vision_prompt.zig diff --git a/docs/vision.md b/docs/vision.md index 486c283663..3b61b0464a 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -1,8 +1,8 @@ # Image input This page describes image input in the Python engine, frozen at 0.6.6. Install that engine from the -[`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6) to serve images. The native binary refuses every -`--vision` flag at startup, reads no vision tower, and answers image parts with HTTP 400. +[`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6) to serve images. The native binary reads images +on its Qwen3.5/3.8 CUDA engine ([below](#native-cuda-qwen3538)); elsewhere it refuses `--vision` at startup. Parts of it came from pull requests: GLM-5.3-Flash images by @mgoldwasser ([#101](https://github.com/ashhart/TensorFold/pull/101)), Flash Next images on CUDA by @shantanugoel ([#146](https://github.com/ashhart/TensorFold/pull/146)) from @MiaAI-Lab's patch, video by @MiaAI-Lab @@ -149,3 +149,15 @@ Tests cover input validation, bounded fetching, expanded prompt accounting, cach Hardware qualification is separate from these tests: each backend and chip needs real image understanding, drafted/serial equality, concurrency, chunked-prefill and memory checks before a release claim. Checkpoint metadata must describe the decoder separately from MTP: the `mlp_layer_types` list must match `num_hidden_layers` for Transformers validation. Preserve the separate MTP configuration and weights. Image input stays experimental until the hardware matrix is complete, and makes no throughput claim. + +## Native CUDA: Qwen3.5/3.8 + +`tensorfold-native serve --vision` loads the checkpoint's vision tower beside the language model on the +Qwen CUDA engine (`families/qwen3_5/cuda_vision.zig`, kernels in `qwen_vision.cu`, its Linears on cuBLASLt). The +tower and the scratch of one 16,384-patch call are allocated at startup and counted as model memory, so the +stream budget never has to make room for an image. A request's images go through the tower before its prompt; +their rows replace the `<|image_pad|>` placeholders in the prompt chunks, which rotate at M-RoPE's t, h and w +positions (interleaved sections 11, 11, 10), and the reply's rows rotate at their position plus `rope_delta`, as +transformers' `get_rope_index` lays them out. Text-only requests keep their bits. Images arrive through the native +server's image parts (decoded by the shared JPEG and PNG decoder); video parts are refused for now. + diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 05c378b553..51c7dbdfef 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -36,6 +36,7 @@ const kernels = [_]Kernel{ .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) .{ .name = "qwen35_keep", .flags = &.{"-O3"} }, // ours: a verified window's kept path applied to its stream's state + .{ .name = "qwen_vision", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the Qwen3.5 vision tower's glue and attention .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen_vision.cu b/zig/kernels/cuda/qwen_vision.cu new file mode 100644 index 0000000000..7281b7a943 --- /dev/null +++ b/zig/kernels/cuda/qwen_vision.cu @@ -0,0 +1,202 @@ +// Qwen3.5 vision tower kernels (0.6.6's --vision in bf16); epilogues add the bias to cuBLASLt's fp32, round once. +#include +#include + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float rnd(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// Patch embedding plus the resampled position grid: x = bf16(acc + b) + bf16(sum_t pos[idx_t] * w_t). +extern "C" __global__ void __launch_bounds__(256) tf_qv_embed(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + const __nv_bfloat16* __restrict__ pos, + const int* __restrict__ idx, + const float* __restrict__ wts, + __nv_bfloat16* __restrict__ out, int D) { + const long long row = blockIdx.x; + const int* ri = idx + row * 4; + const float* rw = wts + row * 4; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float h = rnd(acc[row * D + i] + bf(bias[i])); + float p = 0.0f; + for (int t = 0; t < 4; ++t) p += bf(pos[(long long)ri[t] * D + i]) * rw[t]; + out[row * D + i] = __float2bfloat16_rn(h + rnd(p)); + } +} + +// LayerNorm of bf16 rows in fp32 (mean, then the biased variance about it), bf16 out; a block a row. +extern "C" __global__ void __launch_bounds__(256) tf_qv_layernorm(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ y, int D, float eps) { + __shared__ float part[32]; + __shared__ float stat; + const long long row = blockIdx.x; + const __nv_bfloat16* xr = x + row * D; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5; + float s = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) s += bf(xr[i]); + for (int o = 16; o; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if (lane == 0) part[warp] = s; + __syncthreads(); + if (threadIdx.x == 0) { + float t = 0.0f; + for (int i = 0; i < warps; ++i) t += part[i]; + stat = t / D; + } + __syncthreads(); + const float mean = stat; + float v = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float d = bf(xr[i]) - mean; + v += d * d; + } + for (int o = 16; o; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + if (threadIdx.x == 0) { + float t = 0.0f; + for (int i = 0; i < warps; ++i) t += part[i]; + stat = rsqrtf(t / D + eps); + } + __syncthreads(); + const float r = stat; + for (int i = threadIdx.x; i < D; i += blockDim.x) + y[row * D + i] = __float2bfloat16_rn((bf(xr[i]) - mean) * r * bf(w[i]) + bf(b[i])); +} + +// qkv = bf16(acc + b) split into q, k, v [heads][rows][hd]; q and k get the axial 2-D rotary in fp32 (h, then w). +extern "C" __global__ void __launch_bounds__(128) tf_qv_qkv_rope(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + const int* __restrict__ hw, + const float* __restrict__ inv, + __nv_bfloat16* __restrict__ q, + __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ v, int rows, int heads, + int hd) { + const long long row = blockIdx.x; + const int head = blockIdx.y; + const int D = heads * hd; + const int half = hd / 2, quarter = hd / 4; + const float ph = static_cast(hw[row * 2]), pw = static_cast(hw[row * 2 + 1]); + const float* a = acc + row * 3 * D; + const long long o = (static_cast(head) * rows + row) * hd; + for (int j = threadIdx.x; j < hd; j += blockDim.x) { + const int c = head * hd + j; + // rotate_half pairs j with j +- hd/2; angle h * inv[jj] for the first quarter's frequencies, w * inv after + const int pj = j < half ? j + half : j - half; + const int pc = head * hd + pj; + const int jj = j % half; + const float ang = jj < quarter ? ph * inv[jj] : pw * inv[jj - quarter]; + const float cs = cosf(ang), sn = sinf(ang); + const float sign = j < half ? -1.0f : 1.0f; + const float qx = rnd(a[c] + bf(bias[c])), qp = rnd(a[pc] + bf(bias[pc])); + const float kx = rnd(a[D + c] + bf(bias[D + c])), kp = rnd(a[D + pc] + bf(bias[D + pc])); + q[o + j] = __float2bfloat16_rn(qx * cs + sign * qp * sn); + k[o + j] = __float2bfloat16_rn(kx * cs + sign * kp * sn); + v[o + j] = __float2bfloat16_rn(a[2 * D + c] + bf(bias[2 * D + c])); + } +} + +// Attention within one image's rows [s0, s0 + n); a thread a query row, online softmax in fp32, 64-row key tiles. +template +__device__ void qv_attention(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ out, int rows, + int heads, int s0, int n, float scale) { + constexpr int T = 64; + __shared__ float ks[T][HD]; + __shared__ float vs[T][HD]; + const int head = blockIdx.y; + const int qi = blockIdx.x * T + threadIdx.x; + const bool live = qi < n; + const long long base = static_cast(head) * rows + s0; + float qr[HD], acc[HD]; + float m = -INFINITY, l = 0.0f; + if (live) { + const __nv_bfloat16* qp = q + (base + qi) * HD; + for (int d = 0; d < HD; ++d) { + qr[d] = bf(qp[d]) * scale; + acc[d] = 0.0f; + } + } + for (int t0 = 0; t0 < n; t0 += T) { + const int count = min(T, n - t0); + __syncthreads(); + for (int i = threadIdx.x; i < count * HD; i += blockDim.x) { + const int r = i / HD, d = i % HD; + ks[r][d] = bf(k[(base + t0 + r) * HD + d]); + vs[r][d] = bf(v[(base + t0 + r) * HD + d]); + } + __syncthreads(); + if (!live) continue; + for (int r = 0; r < count; ++r) { + float s = 0.0f; +#pragma unroll + for (int d = 0; d < HD; ++d) s += qr[d] * ks[r][d]; + const float mn = fmaxf(m, s); + const float corr = __expf(m - mn); + const float p = __expf(s - mn); + l = l * corr + p; +#pragma unroll + for (int d = 0; d < HD; ++d) acc[d] = acc[d] * corr + p * vs[r][d]; + m = mn; + } + } + if (!live) return; + const float il = 1.0f / l; + __nv_bfloat16* op = out + static_cast(s0 + qi) * heads * HD + head * HD; + for (int d = 0; d < HD; ++d) op[d] = __float2bfloat16_rn(acc[d] * il); +} + +extern "C" __global__ void __launch_bounds__(64) tf_qv_attention72(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ out, int rows, + int heads, int s0, int n, float scale) { + qv_attention<72>(q, k, v, out, rows, heads, s0, n, scale); +} + +extern "C" __global__ void __launch_bounds__(64) tf_qv_attention64(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ out, int rows, + int heads, int s0, int n, float scale) { + qv_attention<64>(q, k, v, out, rows, heads, s0, n, scale); +} + +// y = gelu(bf16(acc + b)), the tanh form (gelu_pytorch_tanh) or the erf form (nn.GELU), computed in fp32, bf16 out. +extern "C" __global__ void __launch_bounds__(256) tf_qv_bias_gelu(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ y, int D, int tanh_form) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float x = rnd(acc[row * D + i] + bf(bias[i])); + float g; + if (tanh_form) { + const float kBeta = 0.7978845608028654f; // sqrt(2 / pi) + g = 0.5f * x * (1.0f + tanhf(kBeta * (x + 0.044715f * x * x * x))); + } else { + g = 0.5f * x * (1.0f + erff(x * 0.7071067811865476f)); + } + y[row * D + i] = __float2bfloat16_rn(g); + } +} + +// x += bf16(acc + b), the residual add in bf16 (both operands bf16, the sum rounded). +extern "C" __global__ void __launch_bounds__(256) tf_qv_bias_residual(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ x, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float y = rnd(acc[row * D + i] + bf(bias[i])); + x[row * D + i] = __float2bfloat16_rn(bf(x[row * D + i]) + y); + } +} + +// y = bf16(acc + b): the merger's last Linear. +extern "C" __global__ void __launch_bounds__(256) tf_qv_bias_cast(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ y, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) y[row * D + i] = __float2bfloat16_rn(acc[row * D + i] + bf(bias[i])); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 756e535a76..08b94a05d5 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,6 +33,7 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; pub const qwen35_keep: []const u8 = if (available) &Blob("fatbin_qwen35_keep").bytes else &.{}; +pub const qwen_vision: []const u8 = if (available) &Blob("fatbin_qwen_vision").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 47158adda9..f537b4c49e 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -4,6 +4,7 @@ pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); pub const draft = @import("cuda_draft.zig"); +pub const vision = @import("cuda_vision.zig"); test { _ = config; @@ -11,4 +12,6 @@ test { _ = @import("draw.zig"); _ = @import("cuda_draft_weights.zig"); _ = draft; + _ = vision; + _ = @import("cuda_vision_prompt.zig"); } diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 8773dffc71..af7305483f 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -6,6 +6,8 @@ const c = @import("config.zig"); const E = @import("cuda_engine.zig"); const draws = @import("draw.zig"); const Draft = @import("cuda_draft.zig").Draft; +const vision = @import("cuda_vision.zig"); +const vision_prompt = @import("cuda_vision_prompt.zig"); const be = lanes.backend; pub const model_type = "qwen3_5"; @@ -16,7 +18,7 @@ pub const prompt_rows: u32 = E.max_rows; /// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. pub const takes_drafter = true; -pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, vision: bool = false }; /// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. pub const Loaded = struct { @@ -37,8 +39,15 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ const e = try E.Engine.init(gpa, io, ctx, dir); errdefer e.deinit(); if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary + // --vision: the checkpoint's tower, its weights and one call's scratch loaded beside the model's + const tower: ?*vision.Tower = if (!o.vision) null else blk: { + const vc = try vision.Config.read(gpa, io, dir) orelse return error.QwenCheckpointHasNoVision; + if (vc.out_hidden != e.g.hidden) return error.UnsupportedQwenVision; + break :blk try vision.Tower.load(gpa, io, ctx.d, e.stream, dir, vc); + }; + errdefer if (tower) |t| t.deinit(); const l = try gpa.create(Lanes); - l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context }; + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .tower = tower }; return .{ .backend = l.backend(), .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, @@ -107,6 +116,10 @@ fn closeDrafter(p: *anyopaque) void { pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.InvalidQwenPrompt => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.QwenVideoUnsupported => "video inputs are not supported by the Qwen CUDA engine yet: send images", + error.BadImagePlaceholders => "the image placeholders do not match the prepared images", + error.BadImageGrid, error.BadImagePatches => "an image's patches do not fit the vision tower", + error.QwenServedWithoutVision => "image inputs need the server started with --vision", error.OutOfMemory => "the GPU had no memory left for this request's caches: retry once another request ends", else => null, }; @@ -129,6 +142,8 @@ const Lanes = struct { io: std.Io, e: *E.Engine, capacity: usize, + /// --vision's tower: an image prompt's images encoded at its prefill. + tower: ?*vision.Tower = null, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. pool: std.ArrayList(*E.Cache) = .empty, @@ -139,6 +154,7 @@ const Lanes = struct { self.caches.deinit(self.gpa); for (self.pool.items) |cache| self.drop(cache); self.pool.deinit(self.gpa); + if (self.tower) |t| t.deinit(); self.e.deinit(); self.gpa.destroy(self); } @@ -199,6 +215,14 @@ const Lanes = struct { const self = cast(ptr); const ids = stream.prompt(); if (ids.len == 0 or ids.len + stream.max_new > self.capacity or self.caches.contains(stream)) return error.InvalidQwenPrompt; + // an image prompt: its images through the tower first, then every chunk at its t, h, w positions + var media: ?vision_prompt.Media = null; + if (stream.images.len > 0) { + const tower = self.tower orelse return error.QwenServedWithoutVision; + media = try vision_prompt.Media.init(self.gpa, tower, ids, stream.images); + } + defer if (media) |*m| m.deinit(); + const images: ?E.Images = if (media) |m| .{ .features = m.features.ptr, .src = m.src, .pos = m.pos } else null; const cache = try self.take(); errdefer self.give(cache); var at: usize = 0; @@ -208,9 +232,11 @@ const Lanes = struct { end = @min(end, boundary); break; }; - try self.e.forward(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], .last); + try self.e.forward(&.{.{ .cache = cache, .rows = end - at, .images = if (images) |*im| im else null }}, ids[at..end], .last); at = end; } + // the reply's rows rotate past the image prompt's positions (rope_delta) + if (media) |m| cache.delta = m.delta; try self.caches.put(self.gpa, stream, cache); } diff --git a/zig/src/families/qwen3_5/cuda_vision.zig b/zig/src/families/qwen3_5/cuda_vision.zig new file mode 100644 index 0000000000..e790373d85 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_vision.zig @@ -0,0 +1,472 @@ +//! The Qwen3.5/3.8 vision tower on CUDA (transformers' Qwen3_5VisionModel, bf16); prompts: cuda_vision_prompt.zig. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const lt = cuda.cublaslt; +const DeviceBuffer = cuda.DeviceBuffer; +const prefix = "vision_tower."; + +/// What every Qwen3.5/3.8 tower shares (vision_config); its width, heads, depth and MLP are the checkpoint's. +pub const patch = 16; +pub const temporal = 2; +pub const merge = 2; +pub const channels = 3; +pub const patch_width = channels * temporal * patch * patch; // 1536 +pub const grid_side = 48; // sqrt(num_position_embeddings 2304) +pub const eps: f32 = 1e-6; +/// Blocks a tower has at most; its heads are 64 or 72 wide (qwen_vision.cu's attention instances). +pub const max_depth = 32; +/// One tower call's patches at most: a 4,096-token image's. +pub const max_patches = 16384; +const workspace_bytes = 32 << 20; + +/// What config.json says about the tower and its placeholders. +pub const Config = struct { + hidden: usize, + heads: usize, + depth: usize, + intermediate: usize, + out_hidden: usize, + image_token: u32, + video_token: u32, + vision_start: u32, + vision_end: u32, + + pub fn headDim(c: Config) usize { + return c.hidden / c.heads; + } + + /// The merger's width: four merge-ordered rows read as one. + pub fn merged(c: Config) usize { + return c.hidden * merge * merge; + } + + /// A checkpoint without a vision_config is null; a tower or M-RoPE this module does not run is refused. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !?Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !?Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const v = (root.get("vision_config") orelse return null).object; + const want = [_]struct { []const u8, i64 }{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", channels }, .{ "num_position_embeddings", grid_side * grid_side } }; + for (want) |w| if (num(v.get(w[0])) != w[1]) return error.UnsupportedQwenVision; + const act = v.get("hidden_act") orelse return error.UnsupportedQwenVision; + if (act != .string or !std.mem.eql(u8, act.string, "gelu_pytorch_tanh")) return error.UnsupportedQwenVision; + if (v.get("deepstack_visual_indexes")) |ds| if (ds != .array or ds.array.items.len != 0) return error.UnsupportedQwenVision; + // qwen35.cu's rotate: interleaved M-RoPE over sections 11, 11 and 10 + const text = if (root.get("text_config")) |t| t.object else root; + const rope = (text.get("rope_parameters") orelse return error.UnsupportedQwenVision).object; + const inter = rope.get("mrope_interleaved") orelse return error.UnsupportedQwenVision; + const sec = rope.get("mrope_section") orelse return error.UnsupportedQwenVision; + if (inter != .bool or !inter.bool or sec != .array or sec.array.items.len != 3) return error.UnsupportedQwenVision; + for (sec.array.items, [_]i64{ 11, 11, 10 }) |s, x| if (num(s) != x) return error.UnsupportedQwenVision; + const out = num(v.get("out_hidden_size")); + const hid = num(v.get("hidden_size")); + const nh = num(v.get("num_heads")); + const dep = num(v.get("depth")); + const mlp = num(v.get("intermediate_size")); + if (out <= 0 or hid <= 0 or nh <= 0 or mlp <= 0 or dep <= 0 or dep > max_depth) return error.UnsupportedQwenVision; + if (@mod(hid, nh) != 0 or (@divExact(hid, nh) != 64 and @divExact(hid, nh) != 72)) return error.UnsupportedQwenVision; + return .{ .hidden = @intCast(hid), .heads = @intCast(nh), .depth = @intCast(dep), .intermediate = @intCast(mlp), .out_hidden = @intCast(out), .image_token = try id(root, "image_token_id"), .video_token = try id(root, "video_token_id"), .vision_start = try id(root, "vision_start_token_id"), .vision_end = try id(root, "vision_end_token_id") }; + } + + fn num(x: ?std.json.Value) i64 { + const v = x orelse return -1; + return if (v == .integer) v.integer else -1; + } + + fn id(o: std.json.ObjectMap, name: []const u8) !u32 { + const v = num(o.get(name)); + if (v < 0) return error.UnsupportedQwenVision; + return @intCast(v); + } +}; + +/// One image's patch grid (rows and columns of 16x16 patches, both even). +pub const Grid = struct { h: u32, w: u32 }; + +const Block = struct { norm1_w: u64, norm1_b: u64, qkv_w: u64, qkv_b: u64, proj_w: u64, proj_b: u64, norm2_w: u64, norm2_b: u64, fc1_w: u64, fc1_b: u64, fc2_w: u64, fc2_b: u64 }; + +const Fns = struct { + embed: cuda.Function, + layernorm: cuda.Function, + qkv_rope: cuda.Function, + attention: cuda.Function, + bias_gelu: cuda.Function, + bias_residual: cuda.Function, + bias_cast: cuda.Function, +}; + +/// The tower's weights and one call's scratch in one allocation, its kernels; `encode` turns patches into model rows. +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + s: cuda.Stream, + c: Config, + mod: cuda.Module, + f: Fns, + lib: lt.Library, + arena: DeviceBuffer, + at: usize = 0, + scratch: u64 = 0, + patch_w: u64 = 0, + patch_b: u64 = 0, + pos: u64 = 0, + blocks: [max_depth]Block = undefined, + merger_norm_w: u64 = 0, + merger_norm_b: u64 = 0, + fc1_w: u64 = 0, + fc1_b: u64 = 0, + fc2_w: u64 = 0, + fc2_b: u64 = 0, + inv: u64 = 0, + + /// The checkpoint's vision_tower tensors (bf16) on stream `s`; the patch convolution's weight as [C, T, y, x] rows. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8, c: Config) !*Tower { + const t = try gpa.create(Tower); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .d = d, .s = s, .c = c, .mod = undefined, .f = undefined, .lib = undefined, .arena = undefined }; + t.lib = try lt.Library.open(); + errdefer t.lib.close(); + t.mod = try cuda.Module.load(d, cuda.kernels.qwen_vision); + errdefer t.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(t.f, name) = try t.mod.function(if (comptime std.mem.eql(u8, name, "attention")) "tf_qv_attention72" else "tf_qv_" ++ name); + if (c.headDim() == 64) t.f.attention = try t.mod.function("tf_qv_attention64"); + const hidden = c.hidden; + const intermediate = c.intermediate; + const merged = c.merged(); + t.arena = try DeviceBuffer.alloc(d, weightBytes(c) + scratchBytes(c, max_patches)); + errdefer t.arena.free(); + try t.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, prefix); + defer ck.close(); + { // [hidden, T, P, P, C] as stored (MLX) to [hidden, C, T, P, P] + const tn = try ck.expect(prefix ++ "patch_embed.proj.weight", .bf16, &.{ hidden, temporal, patch, patch, channels }); + const host = try gpa.alloc(u16, hidden * patch_width); + defer gpa.free(host); + const stored: []align(1) const u16 = std.mem.bytesAsSlice(u16, tn.bytes); + for (0..hidden) |o| for (0..channels) |ch| for (0..temporal) |tt| for (0..patch) |y| for (0..patch) |x| { + host[o * patch_width + ((ch * temporal + tt) * patch + y) * patch + x] = stored[(((o * temporal + tt) * patch + y) * patch + x) * channels + ch]; + }; + t.patch_w = try t.put(std.mem.sliceAsBytes(host)); + } + t.patch_b = try t.get(&ck, "patch_embed.proj.bias", &.{hidden}); + t.pos = try t.get(&ck, "pos_embed.weight", &.{ grid_side * grid_side, hidden }); + for (t.blocks[0..c.depth], 0..) |*b, i| { + const shapes = .{ + .{ "norm1_w", "norm1.weight", .{hidden} }, .{ "norm1_b", "norm1.bias", .{hidden} }, + .{ "qkv_w", "attn.qkv.weight", .{ 3 * hidden, hidden } }, .{ "qkv_b", "attn.qkv.bias", .{3 * hidden} }, + .{ "proj_w", "attn.proj.weight", .{ hidden, hidden } }, .{ "proj_b", "attn.proj.bias", .{hidden} }, + .{ "norm2_w", "norm2.weight", .{hidden} }, .{ "norm2_b", "norm2.bias", .{hidden} }, + .{ "fc1_w", "mlp.linear_fc1.weight", .{ intermediate, hidden } }, .{ "fc1_b", "mlp.linear_fc1.bias", .{intermediate} }, + .{ "fc2_w", "mlp.linear_fc2.weight", .{ hidden, intermediate } }, .{ "fc2_b", "mlp.linear_fc2.bias", .{hidden} }, + }; + inline for (shapes) |sh| { + var nb: [96]u8 = undefined; + const shape: [sh[2].len]usize = sh[2]; + @field(b, sh[0]) = try t.get(&ck, try std.fmt.bufPrint(&nb, "blocks.{d}.{s}", .{ i, sh[1] }), &shape); + } + } + t.merger_norm_w = try t.get(&ck, "merger.norm.weight", &.{hidden}); + t.merger_norm_b = try t.get(&ck, "merger.norm.bias", &.{hidden}); + t.fc1_w = try t.get(&ck, "merger.linear_fc1.weight", &.{ merged, merged }); + t.fc1_b = try t.get(&ck, "merger.linear_fc1.bias", &.{merged}); + t.fc2_w = try t.get(&ck, "merger.linear_fc2.weight", &.{ c.out_hidden, merged }); + t.fc2_b = try t.get(&ck, "merger.linear_fc2.bias", &.{c.out_hidden}); + if (ck.unused() != 0) return error.UnusedVisionTensors; + var inv: [72 / 4]f32 = undefined; + invFreq(inv[0 .. c.headDim() / 4], c.headDim()); + t.inv = try t.put(std.mem.sliceAsBytes(inv[0 .. c.headDim() / 4])); + t.scratch = try t.take(scratchBytes(c, max_patches)); + return t; + } + + pub fn deinit(t: *Tower) void { + t.arena.free(); + t.mod.unload(); + t.lib.close(); + t.gpa.destroy(t); + } + + /// The weights' bytes in the arena (each allocation 256-byte aligned). + pub fn weightBytes(c: Config) usize { + const a = aligned; + const hidden = c.hidden; + const intermediate = c.intermediate; + const merged = c.merged(); + const out_hidden = c.out_hidden; + var n: usize = a(hidden * patch_width * 2) + a(hidden * 2) + a(grid_side * grid_side * hidden * 2); + n += c.depth * (4 * a(hidden * 2) + a(3 * hidden * hidden * 2) + a(3 * hidden * 2) + a(hidden * hidden * 2) + a(hidden * 2) + a(intermediate * hidden * 2) + a(intermediate * 2) + a(hidden * intermediate * 2) + a(hidden * 2)); + n += 2 * a(hidden * 2) + a(merged * merged * 2) + a(merged * 2) + a(out_hidden * merged * 2) + a(out_hidden * 2); + return n + a(c.headDim() / 4 * 4); + } + + /// The scratch one tower call of `n` patches needs, in bytes (the patches' bf16 copy included). + pub fn scratchBytes(c: Config, n: usize) usize { + const hidden = c.hidden; + const intermediate = c.intermediate; + return n * patch_width * 2 + // patches + 2 * n * hidden * 2 + // x, h + n * @max(3 * hidden, intermediate) * 4 + // fp32 Linear out + 3 * n * hidden * 2 + // q, k, v + n * hidden * 2 + // attention out + n * @max(intermediate, hidden) * 2 + // gelu rows + n * (2 + 4) * 4 + n * 4 * 4 + // hw, idx, weights + 16 * 256 + workspace_bytes; + } + + fn take(t: *Tower, len: usize) !u64 { + const at = aligned(t.at); + if (at + len > t.arena.len) return error.ArenaExhausted; + t.at = at + len; + return t.arena.ptr + at; + } + + fn put(t: *Tower, bytes: []const u8) !u64 { + const p = try t.take(bytes.len); + try t.arena.upload(p - t.arena.ptr, bytes); + return p; + } + + fn get(t: *Tower, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + var buf: [96]u8 = undefined; + const tn = try ck.expect(try std.fmt.bufPrint(&buf, prefix ++ "{s}", .{name}), .bf16, shape); + return t.put(tn.bytes); + } + + fn upload(t: *Tower, dst: u64, bytes: []const u8) !void { + if (bytes.len == 0) return; + try t.d.check(t.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, t.s.handle), "cuMemcpyHtoDAsync"); + } + + /// Images' patches (fp32, merge order) to merged bf16 [sum(h * w) / 4, out_hidden] at `out`, each image on its own. + pub fn encode(t: *Tower, pixels: []const f32, grids: []const Grid, out: u64) !void { + var n: usize = 0; + for (grids) |g| { + if (g.h == 0 or g.w == 0 or g.h % merge != 0 or g.w % merge != 0) return error.BadImageGrid; + n += @as(usize, g.h) * g.w; + } + if (n == 0 or n > max_patches or pixels.len != n * patch_width) return error.BadImagePatches; + const gpa = t.gpa; + const c = t.c; + const out_hidden = c.out_hidden; + const hidden = c.hidden; + const heads = c.heads; + const head_dim = c.headDim(); + const intermediate = c.intermediate; + const merged = c.merged(); + // host side: bf16 patches, each patch's (row, column) and its 4 position-grid taps + const px16 = try gpa.alloc(u16, n * patch_width); + defer gpa.free(px16); + for (px16, pixels) |*o, v| o.* = bf16(v); + const hw = try gpa.alloc(i32, n * 2); + defer gpa.free(hw); + const idx = try gpa.alloc(i32, n * 4); + defer gpa.free(idx); + const wts = try gpa.alloc(f32, n * 4); + defer gpa.free(wts); + var at: usize = 0; + for (grids) |g| { + const len = @as(usize, g.h) * g.w; + positions(g, hw[at * 2 ..][0 .. len * 2], idx[at * 4 ..][0 .. len * 4], wts[at * 4 ..][0 .. len * 4]); + at += len; + } + var cur: u64 = t.scratch; + const end = t.scratch + scratchBytes(c, max_patches); + const take_ = struct { + fn f(p: *u64, len: usize) u64 { + const r = p.*; + p.* += std.mem.alignForward(usize, len, 256); + return r; + } + }.f; + const d_px = take_(&cur, n * patch_width * 2); + const x = take_(&cur, n * hidden * 2); + const h = take_(&cur, n * hidden * 2); + const acc = take_(&cur, n * @max(3 * hidden, intermediate) * 4); + const q = take_(&cur, n * hidden * 2); + const k = take_(&cur, n * hidden * 2); + const v = take_(&cur, n * hidden * 2); + const att = take_(&cur, n * hidden * 2); + const mid = take_(&cur, n * @max(intermediate, hidden) * 2); + const d_hw = take_(&cur, n * 2 * 4); + const d_idx = take_(&cur, n * 4 * 4); + const d_wts = take_(&cur, n * 4 * 4); + const ws = cur; + const ws_len = end - cur; + try t.upload(d_px, std.mem.sliceAsBytes(px16)); + try t.upload(d_hw, std.mem.sliceAsBytes(hw)); + try t.upload(d_idx, std.mem.sliceAsBytes(idx)); + try t.upload(d_wts, std.mem.sliceAsBytes(wts)); + + var mm: Gemms = .{ .lib = &t.lib }; + defer mm.deinit(); + const m4 = n / (merge * merge); + try mm.run(n, hidden, patch_width, d_px, t.patch_w, acc, ws, ws_len, t.s); + try t.go(t.f.embed, n, 1, 256, .{ acc, t.patch_b, t.pos, d_idx, d_wts, x, int(hidden) }); + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(head_dim))); + for (t.blocks[0..c.depth]) |b| { + try t.go(t.f.layernorm, n, 1, 256, .{ x, b.norm1_w, b.norm1_b, h, int(hidden), eps }); + try mm.run(n, 3 * hidden, hidden, h, b.qkv_w, acc, ws, ws_len, t.s); + try t.go(t.f.qkv_rope, n, heads, 128, .{ acc, b.qkv_b, d_hw, t.inv, q, k, v, int(n), int(heads), int(head_dim) }); + var s0: usize = 0; + for (grids) |g| { + const len = @as(usize, g.h) * g.w; + try t.go(t.f.attention, (len + 63) / 64, heads, 64, .{ q, k, v, att, int(n), int(heads), int(s0), int(len), scale }); + s0 += len; + } + try mm.run(n, hidden, hidden, att, b.proj_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_residual, n, 1, 256, .{ acc, b.proj_b, x, int(hidden) }); + try t.go(t.f.layernorm, n, 1, 256, .{ x, b.norm2_w, b.norm2_b, h, int(hidden), eps }); + try mm.run(n, intermediate, hidden, h, b.fc1_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_gelu, n, 1, 256, .{ acc, b.fc1_b, mid, int(intermediate), int(1) }); + try mm.run(n, hidden, intermediate, mid, b.fc2_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_residual, n, 1, 256, .{ acc, b.fc2_b, x, int(hidden) }); + } + // merger: LayerNorm over 1152, four merge-ordered rows read as one of 4608 + try t.go(t.f.layernorm, n, 1, 256, .{ x, t.merger_norm_w, t.merger_norm_b, h, int(hidden), eps }); + try mm.run(m4, merged, merged, h, t.fc1_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_gelu, m4, 1, 256, .{ acc, t.fc1_b, mid, int(merged), int(0) }); + try mm.run(m4, out_hidden, merged, mid, t.fc2_w, acc, ws, ws_len, t.s); + try t.go(t.f.bias_cast, m4, 1, 256, .{ acc, t.fc2_b, out, int(out_hidden) }); + try t.s.synchronize(); // the host's patch rows and tables are freed on return + } + + fn go(t: *Tower, f: cuda.Function, gx: usize, gy: usize, block: u32, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(gx), .y = @intCast(gy) }, .block = .{ .x = block } }, t.s, &a); + } +}; + +/// cuBLASLt plans for the shapes one tower call runs, made on first use. +const Gemms = struct { + lib: *const lt.Library, + plans: [8]?struct { m: usize, n: usize, k: usize, lin: lt.Linear } = @splat(null), + + fn run(g: *Gemms, m: usize, n: usize, k: usize, x: u64, w: u64, out: u64, ws: u64, ws_len: usize, s: cuda.Stream) !void { + for (&g.plans) |*p| { + if (p.*) |*have| { + if (have.m == m and have.n == n and have.k == k) return have.lin.run(x, w, out, ws, ws_len, s.handle); + continue; + } + p.* = .{ .m = m, .n = n, .k = k, .lin = try lt.Linear.init(g.lib, m, n, k, .f32, ws_len) }; + return p.*.?.lin.run(x, w, out, ws, ws_len, s.handle); + } + return error.TooManyShapes; + } + + fn deinit(g: *Gemms) void { + for (&g.plans) |*p| if (p.*) |*have| have.lin.deinit(); + } +}; + +fn aligned(n: usize) usize { + return std.mem.alignForward(usize, n, 256); +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +/// fp32 to bf16, round to nearest even (torch's .to(bfloat16)). +fn bf16(v: f32) u16 { + const b: u32 = @bitCast(v); + if (std.math.isNan(v)) return 0x7fc0; + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// The axial rotary's frequencies: 1 / 10000 ** (i / (head_dim / 2)) for even i, in fp32 as torch computes them. +pub fn invFreq(out: []f32, head_dim: usize) void { + const dim: f32 = @floatFromInt(head_dim / 2); + for (out, 0..) |*f, j| { + const e: f32 = @as(f32, @floatFromInt(2 * j)) / dim; + f.* = 1.0 / std.math.pow(f32, 10000.0, e); + } +} + +/// One image's patches in merge order: each one's (row, column) for the rotary and 4 bilinear taps into the 48x48 grid. +pub fn positions(g: Grid, hw: []i32, idx: []i32, wts: []f32) void { + const bw = g.w / merge; + const n = @as(usize, g.h) * g.w; + for (0..n) |i| { + const in_col: u32 = @intCast(i % merge); + const in_row: u32 = @intCast((i / merge) % merge); + const block_col: u32 = @intCast((i / (merge * merge)) % bw); + const block_row: u32 = @intCast(i / (merge * merge * bw)); + const r = block_row * merge + in_row; + const c = block_col * merge + in_col; + hw[i * 2] = @intCast(r); + hw[i * 2 + 1] = @intCast(c); + const ht = taps(r, g.h); + const wt = taps(c, g.w); + for (0..2) |a| for (0..2) |b| { + idx[i * 4 + a * 2 + b] = ht.i[a] * grid_side + wt.i[b]; + wts[i * 4 + a * 2 + b] = ht.w[a] * wt.w[b]; + }; + } +} + +const Taps = struct { i: [2]i32, w: [2]f32 }; + +fn taps(index: u32, size: u32) Taps { + const side_m1: f32 = grid_side - 1; + const denom: f32 = @floatFromInt(@max(size, 2) - 1); + const src = @as(f32, @floatFromInt(index)) * side_m1 / denom; + const fl = @floor(src); + var out: Taps = undefined; + for (0..2) |o| { + const raw: i32 = @as(i32, @intFromFloat(fl)) + @as(i32, @intCast(o)); + out.i[o] = std.math.clamp(raw, 0, grid_side - 1); + const dist = @abs(src - fl - @as(f32, @floatFromInt(o))); + out.w[o] = @max(1.0 - dist, 0.0); + } + return out; +} + +test "bilinear taps match align_corners resampling at the ends and middle" { + const a = taps(0, 10); + try std.testing.expectEqual(@as(i32, 0), a.i[0]); + try std.testing.expectEqual(@as(f32, 1), a.w[0]); + const b = taps(9, 10); + try std.testing.expectEqual(@as(i32, 47), b.i[0]); + try std.testing.expectEqual(@as(f32, 1), b.w[0]); +} + +test "a vision_config's tower is read; other widths, depths and M-RoPE are refused; a text-only config has none" { + const gpa = std.testing.allocator; + try std.testing.expectEqual(@as(?Config, null), try Config.parse(gpa, "{\"hidden_size\":2048}")); + const good = + \\{"vision_config":{"depth":27,"hidden_size":1152,"num_heads":16,"intermediate_size":4304,"patch_size":16, + \\"temporal_patch_size":2,"spatial_merge_size":2,"in_channels":3,"num_position_embeddings":2304, + \\"hidden_act":"gelu_pytorch_tanh","deepstack_visual_indexes":[],"out_hidden_size":5120}, + \\"text_config":{"rope_parameters":{"mrope_interleaved":true,"mrope_section":[11,11,10]}}, + \\"image_token_id":248056,"video_token_id":248057,"vision_start_token_id":248053,"vision_end_token_id":248054} + ; + const c = (try Config.parse(gpa, good)).?; + try std.testing.expectEqual(@as(usize, 5120), c.out_hidden); + try std.testing.expectEqual(@as(u32, 248056), c.image_token); + const deeper = try std.mem.replaceOwned(u8, gpa, good, "\"depth\":27", "\"depth\":33"); + defer gpa.free(deeper); + try std.testing.expectError(error.UnsupportedQwenVision, Config.parse(gpa, deeper)); + const wide = try std.mem.replaceOwned(u8, gpa, good, "\"num_heads\":16", "\"num_heads\":12"); + defer gpa.free(wide); + try std.testing.expectError(error.UnsupportedQwenVision, Config.parse(gpa, wide)); + // the 2B's tower: 24 blocks of 1024, heads of 64 + const small = try std.mem.replaceOwned(u8, gpa, good, "\"depth\":27,\"hidden_size\":1152,\"num_heads\":16,\"intermediate_size\":4304", "\"depth\":24,\"hidden_size\":1024,\"num_heads\":16,\"intermediate_size\":4096"); + defer gpa.free(small); + const c2 = (try Config.parse(gpa, small)).?; + try std.testing.expectEqual(@as(usize, 64), c2.headDim()); + try std.testing.expectEqual(@as(usize, 24), c2.depth); + const sections = try std.mem.replaceOwned(u8, gpa, good, "[11,11,10]", "[16,24,24]"); + defer gpa.free(sections); + try std.testing.expectError(error.UnsupportedQwenVision, Config.parse(gpa, sections)); +} diff --git a/zig/src/families/qwen3_5/cuda_vision_prompt.zig b/zig/src/families/qwen3_5/cuda_vision_prompt.zig new file mode 100644 index 0000000000..3aa5f3b4c4 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_vision_prompt.zig @@ -0,0 +1,175 @@ +//! An image prompt's rows for the CUDA engine: its images through the tower, its M-RoPE positions (get_rope_index). +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const vision = @import("cuda_vision.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Config = vision.Config; +const Tower = vision.Tower; +const Grid = vision.Grid; +const merge = vision.merge; +const patch_width = vision.patch_width; +const max_patches = vision.max_patches; + +/// An image prompt's encoded rows: the tower's features on the device, each prompt row's feature row and positions. +pub const Media = struct { + gpa: std.mem.Allocator, + features: DeviceBuffer, // bf16 [image rows, out_hidden] + src: []i32, // [prompt]: feature row or -1 + pos: []i32, // [3][prompt]: t, h, w + delta: i32, + + /// `images` (the request's, in prompt order) through the tower; the prompt's positions as get_rope_index lays them. + pub fn init(gpa: std.mem.Allocator, tower: *Tower, prompt: []const u32, images: []const lanes.Image) !Media { + const c = tower.c; + const n = prompt.len; + const runs = try gpa.alloc(Run, images.len); + defer gpa.free(runs); + try placeholders(c, prompt, images, runs); + var rows: usize = 0; + for (runs) |r| rows += r.tokens(); + var m: Media = .{ .gpa = gpa, .features = try DeviceBuffer.alloc(tower.d, @max(rows, 1) * c.out_hidden * 2), .src = &.{}, .pos = &.{}, .delta = 0 }; + errdefer m.features.free(); + m.src = try gpa.alloc(i32, n); + errdefer gpa.free(m.src); + m.pos = try gpa.alloc(i32, 3 * n); + errdefer gpa.free(m.pos); + // the tower: runs of whole images within one call's patches, in prompt order + var done: usize = 0; // feature rows written + var i: usize = 0; + while (i < images.len) { + var j = i; + var patches: usize = 0; + while (j < images.len and (j == i or patches + runs[j].patches() <= max_patches)) : (j += 1) patches += runs[j].patches(); + const grids = try gpa.alloc(Grid, j - i); + defer gpa.free(grids); + for (runs[i..j], grids) |r, *g| g.* = .{ .h = r.gh, .w = r.gw }; + if (j - i == 1) { + try tower.encode(images[i].patches, grids, m.features.ptr + done * c.out_hidden * 2); + } else { + const px = try gpa.alloc(f32, patches * patch_width); + defer gpa.free(px); + var at: usize = 0; + for (images[i..j]) |img| { + @memcpy(px[at..][0..img.patches.len], img.patches); + at += img.patches.len; + } + try tower.encode(px, grids, m.features.ptr + done * c.out_hidden * 2); + } + for (runs[i..j]) |r| done += r.tokens(); + i = j; + } + m.delta = layout(n, runs, m.src, m.pos); + return m; + } + + pub fn deinit(m: *Media) void { + m.features.free(); + m.gpa.free(m.src); + m.gpa.free(m.pos); + } +}; + +/// An image's placeholder run in the prompt: its first row and its patch grid. +pub const Run = struct { + at: usize, + gh: u32, + gw: u32, + + pub fn tokens(r: Run) usize { + return @as(usize, r.gh / merge) * (r.gw / merge); + } + + pub fn patches(r: Run) usize { + return @as(usize, r.gh) * r.gw; + } +}; + +/// Each image's run of image_token rows between vision_start and vision_end, in order; videos are refused by name. +pub fn placeholders(c: Config, prompt: []const u32, images: []const lanes.Image, runs: []Run) !void { + var k: usize = 0; + var r: usize = 0; + while (r < prompt.len) : (r += 1) { + const t = prompt[r]; + if (t == c.video_token) return error.QwenVideoUnsupported; + if (t != c.image_token) continue; + if (r == 0 or prompt[r - 1] != c.vision_start) return error.BadImagePlaceholders; + if (k == images.len) return error.BadImagePlaceholders; + const img = images[k]; + if (img.grid[0] != 1) return error.QwenVideoUnsupported; + if (img.grid[1] <= 0 or img.grid[2] <= 0 or img.grid[1] > 4096 or img.grid[2] > 4096) return error.BadImageGrid; + const run: Run = .{ .at = r, .gh = @intCast(img.grid[1]), .gw = @intCast(img.grid[2]) }; + if (run.gh % merge != 0 or run.gw % merge != 0 or run.patches() > max_patches) return error.BadImageGrid; + if (img.patches.len != run.patches() * patch_width) return error.BadImagePatches; + const end = r + run.tokens(); + if (end >= prompt.len or prompt[end] != c.vision_end) return error.BadImagePlaceholders; + for (prompt[r..end]) |x| if (x != c.image_token) return error.BadImagePlaceholders; + runs[k] = run; + k += 1; + r = end; + } + if (k != images.len) return error.BadImagePlaceholders; +} + +/// The prompt's rotary positions ([3][n]) and feature rows (src, -1 for text) per get_rope_index; returns rope_delta. +pub fn layout(n: usize, runs: []const Run, src: []i32, pos: []i32) i32 { + @memset(src, -1); + var next: i32 = 0; + var cursor: usize = 0; + var row: i32 = 0; + for (runs) |img| { + const begin = img.at; + const h: usize = img.gh / merge; + const w: usize = img.gw / merge; + for (cursor..begin) |r| { + const v = next + @as(i32, @intCast(r - cursor)); + for (0..3) |a| pos[a * n + r] = v; + } + const base = next + @as(i32, @intCast(begin - cursor)); + for (0..h * w) |k| { + const r = begin + k; + pos[r] = base; + pos[n + r] = base + @as(i32, @intCast(k / w)); + pos[2 * n + r] = base + @as(i32, @intCast(k % w)); + src[r] = row; + row += 1; + } + next = base + @as(i32, @intCast(@max(h, w))); + cursor = begin + h * w; + } + for (cursor..n) |r| { + const v = next + @as(i32, @intCast(r - cursor)); + for (0..3) |a| pos[a * n + r] = v; + } + return next - @as(i32, @intCast(cursor)); +} + +const test_config: Config = .{ .hidden = 1152, .heads = 16, .depth = 27, .intermediate = 4304, .out_hidden = 8, .image_token = 50, .video_token = 51, .vision_start = 52, .vision_end = 53 }; + +test "media positions: text counts up, an image sits at one position with its grid, text resumes past its side" { + const prompt = [_]u32{ 7, 52, 50, 50, 50, 50, 50, 50, 53, 9 }; + const px: [24 * patch_width]f32 = @splat(0); + var runs: [1]Run = undefined; + try placeholders(test_config, &prompt, &.{.{ .patches = &px, .grid = .{ 1, 4, 6 } }}, &runs); + try std.testing.expectEqual(@as(usize, 2), runs[0].at); + var src: [10]i32 = undefined; + var pos: [30]i32 = undefined; + const delta = layout(prompt.len, &runs, &src, &pos); + try std.testing.expectEqual(@as(i32, -3), delta); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 2, 2, 2, 2, 2, 5, 6 }, pos[0..10]); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 2, 2, 3, 3, 3, 5, 6 }, pos[10..20]); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 4, 2, 3, 4, 5, 6 }, pos[20..30]); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 0, 1, 2, 3, 4, 5, -1, -1 }, &src); +} + +test "placeholders refuse a run of the wrong length, a missing image, an extra image and a video" { + const px: [24 * patch_width]f32 = @splat(0); + const img: lanes.Image = .{ .patches = &px, .grid = .{ 1, 4, 6 } }; + var runs: [2]Run = undefined; + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 52, 50, 50, 53 }, &.{img}, runs[0..1])); + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 52, 50, 50, 50, 50, 50, 50, 53 }, &.{ img, img }, &runs)); + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 7, 8 }, &.{img}, runs[0..1])); + try std.testing.expectError(error.QwenVideoUnsupported, placeholders(test_config, &.{ 52, 51, 53 }, &.{}, runs[0..0])); + try std.testing.expectError(error.BadImagePlaceholders, placeholders(test_config, &.{ 7, 50 }, &.{img}, runs[0..1])); +} From ac94f8faecfd950e5f62edba181eae0859040750 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:19:23 +0000 Subject: [PATCH 36/46] tf-qwen35-cuda: vision and mrope checks vision: each 0.6.6 fixture's patches through the tower, rows' cosine to 0.6.6's (tools/qwen_vision_fixtures.py). mrope: a prompt through the image table with text positions must equal prompt's logits bit for bit; TF_Q35_SHIFT=k moves the h axis, a control whose logits must differ. --- zig/tests/qwen35_cuda_main.zig | 122 +++++++++++++++++++++++++++++++-- 1 file changed, 117 insertions(+), 5 deletions(-) diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index e0bb3ac025..a9945e7d9c 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,4 +1,5 @@ -//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. +//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. +//! `tf-qwen35-cuda vision `: the tower's rows for each fixture against 0.6.6's. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -8,9 +9,10 @@ pub fn main(init: std.process.Init) !u8 { const a = init.arena.allocator(); const args = try init.minimal.args.toSlice(a); if (args.len != 5) { - std.debug.print("usage: tf-qwen35-cuda \n", .{}); + std.debug.print("usage: tf-qwen35-cuda \n", .{}); return 2; } + if (std.mem.eql(u8, args[2], "vision")) return vision(init, a, args[1], args[3], args[4]); var ids: std.ArrayList(u32) = .empty; var it = std.mem.tokenizeAny(u8, args[3], ", "); while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); @@ -90,10 +92,29 @@ pub fn main(init: std.process.Init) !u8 { at += kept; } } else { + // mrope: the chunks through an image prompt's table (no image rows, text positions), which must equal prompt's bits + const table = std.mem.eql(u8, args[2], "mrope"); + const n_ids = ids.items.len; + const src = try a.alloc(i32, n_ids); + @memset(src, -1); + const pos = try a.alloc(i32, 3 * n_ids); + // TF_Q35_SHIFT=k moves the h axis by k: a control whose logits must differ, as the table is read + const shift = try envCount("TF_Q35_SHIFT"); + for (0..3) |ax| for (0..n_ids) |i| { + pos[ax * n_ids + i] = @intCast(cache.len + i + if (ax == 1) shift else 0); + }; + var none = try cuda.DeviceBuffer.alloc(ctx.d, 256); + defer none.free(); + const im: E.Images = .{ .features = none.ptr, .src = src, .pos = pos }; + const p0 = cache.len; var at: usize = 0; - while (at < ids.items.len) { - const n = @min(E.max_rows, ids.items.len - at); - try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); + while (at < n_ids) { + const n = @min(E.max_rows, n_ids - at); + if (table) { + // the table's rows are the prompt's from position p0: shift it so row `at` is the cache's next + const part: E.Images = .{ .features = im.features, .src = try padded(a, src, p0), .pos = try paddedPos(a, pos, p0, n_ids) }; + try e.forward(&.{.{ .cache = &cache, .rows = n, .images = &part }}, ids.items[at..][0..n], .last); + } else try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); at += n; } try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); @@ -135,3 +156,94 @@ fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { at += k; } } + +/// Each fixture's patches (manifest.json, tools/zig/qwen_vision_fixtures.py) through the tower; rows' cosine to 0.6.6's. +fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: []const u8, out_path: []const u8) !u8 { + const V = q.vision; + const io = init.io; + const cfg = try V.Config.read(a, io, model) orelse return error.QwenCheckpointHasNoVision; + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var stream = try cuda.Stream.init(ctx.d, true); + defer stream.deinit(); + const tower = try V.Tower.load(init.gpa, io, ctx.d, stream, model, cfg); + defer tower.deinit(); + const cwd = std.Io.Dir.cwd(); + const manifest = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "manifest.json" }), a, .limited(1 << 20)); + const parsed = try std.json.parseFromSliceLeaky(std.json.Value, a, manifest, .{}); + var out: std.ArrayList(u8) = .empty; + var worst: f64 = 1; + for (parsed.object.get("images").?.array.items) |im| { + const o = im.object; + const grid = o.get("grid").?.array.items; + const g: V.Grid = .{ .h = @intCast(grid[0].integer), .w = @intCast(grid[1].integer) }; + const tokens: usize = @intCast(o.get("tokens").?.integer); + const px_bytes = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, o.get("pixels").?.string }), a, .limited(1 << 30)); + const want = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, o.get("features").?.string }), a, .limited(1 << 30)); + const px = try a.alloc(f32, px_bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(px), px_bytes); + var features = try cuda.DeviceBuffer.alloc(ctx.d, tokens * cfg.out_hidden * 2); + defer features.free(); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + try tower.encode(px, &.{g}, features.ptr); + const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e6; + const got = try a.alloc(u8, tokens * cfg.out_hidden * 2); + try features.download(0, got); + try out.appendSlice(a, got); + if (want.len != got.len) { + std.debug.print("{s}: 0.6.6 has {d} bytes of rows, the tower {d}\n", .{ o.get("name").?.string, want.len, got.len }); + return 1; + } + const gw: []align(1) const u16 = std.mem.bytesAsSlice(u16, got); + const ww: []align(1) const u16 = std.mem.bytesAsSlice(u16, want); + var sum: f64 = 0; + var least: f64 = 1; + var same: usize = 0; + for (0..tokens) |r| { + var dot: f64 = 0; + var ng: f64 = 0; + var nw: f64 = 0; + for (0..cfg.out_hidden) |i| { + const x = bf(gw[r * cfg.out_hidden + i]); + const y = bf(ww[r * cfg.out_hidden + i]); + dot += x * y; + ng += x * x; + nw += y * y; + } + const cos = dot / @sqrt(ng * nw); + sum += cos; + least = @min(least, cos); + if (std.mem.eql(u16, @as([]const u16, @ptrCast(@alignCast(gw[r * cfg.out_hidden ..][0..cfg.out_hidden]))), @ptrCast(@alignCast(ww[r * cfg.out_hidden ..][0..cfg.out_hidden])))) same += 1; + } + worst = @min(worst, least); + std.debug.print("{s}: grid {d}x{d}, {d} rows, row cosine mean {d:.5} min {d:.5}, {d} rows identical, {d:.1} ms\n", .{ o.get("name").?.string, g.h, g.w, tokens, sum / @as(f64, @floatFromInt(tokens)), least, same, ms }); + } + try cwd.writeFile(io, .{ .sub_path = out_path, .data = out.items }); + std.debug.print("worst row cosine {d:.5}\n", .{worst}); + return 0; +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +/// `src` behind `p0` rows of text (-1): the table indexed by the cache's position, as a prompt from position 0 is. +fn padded(a: std.mem.Allocator, src: []const i32, p0: usize) ![]i32 { + const out = try a.alloc(i32, p0 + src.len); + @memset(out[0..p0], -1); + @memcpy(out[p0..], src); + return out; +} + +/// `pos` ([3][n]) behind `p0` rows, each axis: the positions indexed by the cache's position. +fn paddedPos(a: std.mem.Allocator, pos: []const i32, p0: usize, n: usize) ![]i32 { + const m = p0 + n; + const out = try a.alloc(i32, 3 * m); + for (0..3) |ax| { + for (0..p0) |i| out[ax * m + i] = @intCast(i); + @memcpy(out[ax * m + p0 ..][0..n], pos[ax * n ..][0..n]); + } + return out; +} From 7ad45b698f3479dcc154a10280410d5928897c33 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:40:15 +0000 Subject: [PATCH 37/46] native/cuda: drop the doc comment learnBytes left behind on toGib The #853 merge removed learnBytes (#836 reads learn_bytes in place); its doc line stayed above toGib. --- zig/src/native/cuda.zig | 1 - 1 file changed, 1 deletion(-) diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 4050f4ac79..131bdb872b 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -473,7 +473,6 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } -/// A family's bytes kept for --slide (0 for families without a learner). fn toGib(bytes: u64) f64 { return @as(f64, @floatFromInt(bytes)) / gib; } From 69bbf2c9c365058c56d88a9cc8cbc00c39b16c20 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:40:25 +0000 Subject: [PATCH 38/46] lane host: the prompt cache keys a request by cache_key, from #726 Two image prompts with the same text render the same ids, so keyed by tokens alone one could resume the other's state. Request.cache_key holds the prompt with each image's placeholder rows as ids from its content; lookups, keeps and resumes take cacheKey() (the prompt when empty). The test is #726's. --- zig/src/core/engine_api.zig | 7 +++ zig/src/core/lane_host.zig | 8 +-- zig/src/core/lane_host_reuse_test.zig | 88 +++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 4 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 3c633640ad..82fe572e94 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -93,6 +93,13 @@ pub const Request = struct { prompt_logprobs: ?u8 = null, /// An image prompt's inputs in prompt order (Info.vision engines only; the prompt holds their placeholder rows). images: []const lanes.Image = &.{}, + /// The prompt as the prompt cache keys it: each image's placeholder rows as ids from its content (empty: prompt). + cache_key: []const u32 = &.{}, + + /// The ids the prompt cache keys this request's states by: placeholders stand for their images. + pub fn cacheKey(r: *const Request) []const u32 { + return if (r.cache_key.len > 0) r.cache_key else r.prompt; + } }; pub const Reason = enum { stop, length, cancelled, failed }; diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 73937e925c..acb7632595 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -77,7 +77,7 @@ pub const LaneHost = struct { 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.keepBorn(job.request.prompt, born(job.request), at, s, job.request.chunks, job.host.spans(job.request)); + if (job.host.cache) |store| _ = store.keepBorn(job.request.cacheKey(), born(job.request), at, s, job.request.chunks, job.host.spans(job.request)); } /// Report a restored prefix or its failed copy; an untouched prefix remains kept. @@ -88,7 +88,7 @@ pub const LaneHost = struct { defer store.unpin(e); // after resumed, which reads e if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it - if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); + if (job.stream.reuse_failed) store.resumed(e, job.request.cacheKey(), false) else if (job.stream.cached == e.at) store.resumed(e, job.request.cacheKey(), true); } }; @@ -317,7 +317,7 @@ pub const LaneHost = struct { return false; } var reuse: lanes.stream.Reuse = .{}; - if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { + if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { job.entry = if (l.entry) |e| store.pin(e) else null; // kept alive until its restore reports job.kept0 = store.counts.kept; job.marks = l.marks; @@ -448,7 +448,7 @@ pub const LaneHost = struct { return h.cancel(job); }; h.prefilled(job, began); - if (h.cache) |store| _ = if (!job.request.untaught and store.rules.lookahead == 0) store.keepBorn(job.request.prompt, born(job.request), @intCast(job.request.prompt.len), &job.stream, job.request.chunks, h.spans(job.request)) else false; + if (h.cache) |store| _ = if (!job.request.untaught and store.rules.lookahead == 0) store.keepBorn(job.request.cacheKey(), born(job.request), @intCast(job.request.prompt.len), &job.stream, job.request.chunks, h.spans(job.request)) else false; be.release(&job.stream); h.remove(job); h.finish(job, .length, ""); diff --git a/zig/src/core/lane_host_reuse_test.zig b/zig/src/core/lane_host_reuse_test.zig index 01b1fc6bec..69fca5ae1f 100644 --- a/zig/src/core/lane_host_reuse_test.zig +++ b/zig/src/core/lane_host_reuse_test.zig @@ -106,6 +106,94 @@ test "a lane host resumes a conversation from its kept prompt state and reports try std.testing.expectEqual(@as(usize, 2), store.entries.items.len); } +test "an image's placeholders are cached by its content: another image under the same ids resumes none of it" { + 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.backend(), clock.clock()); + defer core.deinit(); + const Snaps = struct { + fn bytes(_: *anyopaque, at: u32) u64 { + return 64 + 4 * @as(u64, at); + } + fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!pc.Saved { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + return f.save(@ptrCast(@alignCast(owner.?)), at); + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: pc.Saved) anyerror!void { + return error.BackendRestores; // lane-core backends restore inside their own prompt pass + } + fn drop(ptr: *anyopaque, saved: pc.Saved) void { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + f.drop(saved); + } + }; + var store = pc.Store.init(gpa, .{ .ptr = &target, .vtable = &.{ .bytes = Snaps.bytes, .save = Snaps.save, .restore = Snaps.restore, .drop = Snaps.drop } }, .{ .min_gap = 4, .min_prompt = 0 }, 1 << 20); + defer store.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + host.cache = &store; + try host.start(); + defer host.stop(); + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + cached: ?u32 = null, + 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.*) { + .prefilled => |c| b.cached = c, + .logprobs, .prompt_logprobs => {}, + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + } + } + 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 e = host.engine(); + // prompt ids 3..5 are one image's placeholders (token 9); cache_key holds ids derived from each image's content + const t1 = [_]u32{ 3, 1, 9, 9, 9, 9, 2, 6, 5, 3 }; + var key_a = t1; + var key_b = t1; + for (2..6) |k| { + key_a[k] = 0x8000_0000 + @as(u32, @intCast(k)); + key_b[k] = 0x9000_0000 + @as(u32, @intCast(k)); + } + var b1: Box = .{}; + defer b1.tokens.deinit(gpa); + const r1: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_a }; + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b1.wait()); + try std.testing.expectEqual(@as(usize, 1), store.entries.items.len); // image A's state at 8, past its image + // the same text and placeholder ids with image B: the kept state is A's, so B resumes none of it + var b2: Box = .{}; + defer b2.tokens.deinit(gpa); + const r2: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_b }; + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(?u32, 0), b2.cached); + // image A again resumes its own state + var b3: Box = .{}; + defer b3.tokens.deinit(gpa); + const r3: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_a }; + try e.submit(3, &r3, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(?u32, 8), b3.cached); +} + test "a turn's mark evicting the state it resumed from: the hit is reported first and the entry freed once" { const gpa = std.testing.allocator; var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); From a78d48094fec5cef6374403eb99396308dfb7513 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:40:25 +0000 Subject: [PATCH 39/46] server: an image request's cache key, its placeholder rows as ids from each image's content image_key.cacheKey: placeholder run k (Qwen's image or video pads, GLM's image token) takes ids hashed from image k's grid and patches and the row's place in its run; text keys as itself, so a change to a later image keeps the prefix before it. Runs that do not pair with the images key by all of them. --- zig/src/server/chat.zig | 16 +++++++ zig/src/server/image_key.zig | 83 ++++++++++++++++++++++++++++++++++++ zig/src/server/root.zig | 1 + 3 files changed, 100 insertions(+) create mode 100644 zig/src/server/image_key.zig diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 63eebea817..e1e868b662 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -13,6 +13,7 @@ const ids = @import("ids.zig"); const clock = @import("clock.zig"); const Server = @import("server.zig").Server; const chunk_plan = @import("chunk_plan.zig"); +const image_key = @import("image_key.zig"); const Mailbox = @import("mailbox.zig").Mailbox; const Value = json.Value; const Cx = errors.Cx; @@ -168,6 +169,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input_in: Input, gone: anytype) Failure!Pr .tools_json = if (input.tools.len > 0) try json.stringify(a, .{ .array = @constCast(input.tools) }, .{ .ascii = false }) else "", .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, }; + if (input.images.len > 0) request.cache_key = try imageCacheKey(srv, a, rendered.ids, input.images); try srv.checkFeatures(cx, f, input.tools.len > 0, thinking, rendered.ids, input.tools, &request); if (thinking) { const budget_field = f.get("thinking_budget"); @@ -577,3 +579,17 @@ pub fn deltaOf(a: Allocator, key: []const u8, text: []const u8) Allocator.Error! try o.put(a, key, .{ .string = text }); return .{ .object = o }; } + +/// An image prompt's cache key: its placeholder rows (Qwen's image and video pads, GLM's image token) by content. +fn imageCacheKey(srv: *Server, a: std.mem.Allocator, prompt: []const u32, images: []const api.Image) ![]const u32 { + var pads: [3]u32 = undefined; + var n: usize = 0; + for ([_][]const u8{ "<|image_pad|>", "<|video_pad|>", "<|image|>" }) |name| if (srv.text.tokenId(name)) |t| { + pads[n] = t; + n += 1; + }; + const key = try image_key.cacheKey(a, prompt, images, pads[0..n]); + // no placeholder token known: the key's first id is the images', so no other prompt shares any state + if (n == 0 and key.len > 0) key[0] = (try image_key.cacheKey(a, &.{0}, images, &.{0}))[0]; + return key; +} diff --git a/zig/src/server/image_key.zig b/zig/src/server/image_key.zig new file mode 100644 index 0000000000..175b8d0809 --- /dev/null +++ b/zig/src/server/image_key.zig @@ -0,0 +1,83 @@ +//! An image prompt's prompt-cache key: each image's placeholder rows as ids drawn from that image's content. +const std = @import("std"); +const api = @import("engine_api"); + +/// Ids the tokenizer never makes: the top bit set. +const marked: u32 = 0x8000_0000; + +/// One image's content: its grid and its patches' bits. +fn imageHash(im: api.Image) u64 { + var h = std.hash.Wyhash.init(0x51ce_7a91); + h.update(std.mem.asBytes(&im.grid)); + h.update(std.mem.sliceAsBytes(im.patches)); + return h.final(); +} + +/// `ids` with placeholder run k (ids in `pads`) as ids from image k's content; unpaired runs key by all the images. +pub fn cacheKey(a: std.mem.Allocator, ids: []const u32, images: []const api.Image, pads: []const u32) ![]u32 { + const key = try a.dupe(u32, ids); + var runs: usize = 0; + for (ids, 0..) |t, i| { + if (isPad(t, pads) and (i == 0 or !isPad(ids[i - 1], pads))) runs += 1; + } + var all = std.hash.Wyhash.init(0x51ce_7a92); + for (images) |im| all.update(std.mem.asBytes(&imageHash(im))); + const together = all.final(); + var run: usize = 0; + var row: u64 = 0; + var h: u64 = together; + for (key, 0..) |*t, i| { + if (!isPad(ids[i], pads)) continue; + if (i == 0 or !isPad(ids[i - 1], pads)) { + h = if (runs == images.len) imageHash(images[run]) else together; + run += 1; + row = 0; + } + t.* = marked | @as(u32, @truncate(std.hash.Wyhash.hash(h, std.mem.asBytes(&row)) & 0x7fff_ffff)); + row += 1; + } + return key; +} + +fn isPad(t: u32, pads: []const u32) bool { + return std.mem.indexOfScalar(u32, pads, t) != null; +} + +test "two images under the same placeholder ids key apart; the text around them keys as itself" { + const a = std.testing.allocator; + const px_a: [4]f32 = .{ 0.1, 0.2, 0.3, 0.4 }; + const px_b: [4]f32 = .{ 0.1, 0.2, 0.3, 0.5 }; + const ids = [_]u32{ 7, 52, 9, 9, 53, 8 }; + const ka = try cacheKey(a, &ids, &.{.{ .patches = &px_a, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(ka); + const kb = try cacheKey(a, &ids, &.{.{ .patches = &px_b, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(kb); + const ka2 = try cacheKey(a, &ids, &.{.{ .patches = &px_a, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(ka2); + try std.testing.expectEqualSlices(u32, &.{ 7, 52 }, ka[0..2]); + try std.testing.expectEqualSlices(u32, &.{ 53, 8 }, ka[4..6]); + try std.testing.expect(ka[2] >= marked and ka[2] != ka[3]); + try std.testing.expect(ka[2] != kb[2] and ka[3] != kb[3]); + try std.testing.expectEqualSlices(u32, ka, ka2); + // the same pixels in another grid are another image + const kg = try cacheKey(a, &ids, &.{.{ .patches = &px_a, .grid = .{ 1, 1, 4 } }}, &.{9}); + defer a.free(kg); + try std.testing.expect(kg[2] != ka[2]); +} + +test "a second image's change keys only its own run apart; runs that do not pair with images key by all of them" { + const a = std.testing.allocator; + const p1: [2]f32 = .{ 1, 2 }; + const p2: [2]f32 = .{ 3, 4 }; + const p3: [2]f32 = .{ 3, 5 }; + const ids = [_]u32{ 52, 9, 9, 53, 7, 52, 9, 53 }; + const k12 = try cacheKey(a, &ids, &.{ .{ .patches = &p1, .grid = .{ 1, 2, 2 } }, .{ .patches = &p2, .grid = .{ 1, 2, 2 } } }, &.{9}); + defer a.free(k12); + const k13 = try cacheKey(a, &ids, &.{ .{ .patches = &p1, .grid = .{ 1, 2, 2 } }, .{ .patches = &p3, .grid = .{ 1, 2, 2 } } }, &.{9}); + defer a.free(k13); + try std.testing.expectEqualSlices(u32, k12[0..6], k13[0..6]); + try std.testing.expect(k12[6] != k13[6]); + const one = try cacheKey(a, &ids, &.{.{ .patches = &p1, .grid = .{ 1, 2, 2 } }}, &.{9}); + defer a.free(one); + try std.testing.expect(one[1] != k12[1]); // two runs, one image: neither run keys as image 1 alone +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 9cfa2dc1d4..6673db74cd 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -34,6 +34,7 @@ test { _ = @import("listener.zig"); _ = @import("hub.zig"); _ = @import("chunk_plan.zig"); + _ = @import("image_key.zig"); _ = @import("status_routes.zig"); _ = @import("memory_routes.zig"); _ = @import("process_memory.zig"); From 32170634fb4bf12de7af4557ad7ea7cdb982fc0b Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:06:42 +0000 Subject: [PATCH 40/46] CI on #885: the server's image paths build on Windows; lean and a trailing blank line - vision_inputs opens libtfvideo through sys.DynLib (LoadLibraryW on Windows, as tfvideo.dll): std.DynLib has no Windows. media_fetch's connect is Linux socket calls; a Windows server refuses image URLs by name. #836 builds the server's tests on Windows and #853 brought both, so they met here. - tf-qwen35-cuda's header and two comments within one line and 120 columns. - docs/vision.md ends in one newline (git diff --check). --- docs/vision.md | 1 - zig/src/server/media_fetch.zig | 6 ++++++ zig/src/server/vision_inputs.zig | 4 +++- zig/tests/qwen35_cuda_main.zig | 7 +++---- 4 files changed, 12 insertions(+), 6 deletions(-) diff --git a/docs/vision.md b/docs/vision.md index 3b61b0464a..d33e8b5162 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -160,4 +160,3 @@ their rows replace the `<|image_pad|>` placeholders in the prompt chunks, which positions (interleaved sections 11, 11, 10), and the reply's rows rotate at their position plus `rope_delta`, as transformers' `get_rope_index` lays them out. Text-only requests keep their bits. Images arrive through the native server's image parts (decoded by the shared JPEG and PNG decoder); video parts are refused for now. - diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig index c9ef04847f..a236d97152 100644 --- a/zig/src/server/media_fetch.zig +++ b/zig/src/server/media_fetch.zig @@ -1,6 +1,7 @@ //! Public HTTPS media for --vision-urls as vision/images_http.py fetches it: checked hosts, TLS, 3 redirects, limits. const std = @import("std"); +const builtin = @import("builtin"); const Io = std.Io; const net = std.Io.net; @@ -124,6 +125,11 @@ fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { /// A TCP connection to the checked address, connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds connect). fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + // the socket calls below are Linux's; a Windows server refuses image URLs by name + return if (comptime builtin.os.tag == .windows) fail(f, "image URLs are not supported on Windows yet") else connectLinux(addr, ms, f); +} + +fn connectLinux(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { const posix = std.posix; const linux = std.os.linux; const family: u32 = switch (addr) { diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index b1e4476fb1..07e87a1cae 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -1,6 +1,8 @@ //! A chat request's image_url and video_url parts, prepared and replaced by the text Python's frontend expands them to. const std = @import("std"); +const builtin = @import("builtin"); +const sys = @import("sys"); const media_fetch = @import("media_fetch.zig"); const api = @import("engine_api"); const json = @import("json.zig"); @@ -250,7 +252,7 @@ fn library(io: std.Io) ?@TypeOf(lib.?) { lib_mutex.lockUncancelable(io); defer lib_mutex.unlock(io); if (lib == null) { - var dl = std.DynLib.open("libtfvideo.so") catch return null; + var dl = sys.DynLib.open(if (builtin.os.tag == .windows) "tfvideo.dll" else "libtfvideo.so") catch return null; const d = dl.lookup(DecodeFn, "tf_video_decode") orelse return null; const f = dl.lookup(FreeFn, "tf_video_free") orelse return null; lib = .{ .decode = d, .free = f }; // kept open for the server's life diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index a9945e7d9c..cf6738e627 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,5 +1,4 @@ -//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. -//! `tf-qwen35-cuda vision `: the tower's rows for each fixture against 0.6.6's. +//! `tf-qwen35-cuda `, or `vision `: rows to compare. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -92,7 +91,7 @@ pub fn main(init: std.process.Init) !u8 { at += kept; } } else { - // mrope: the chunks through an image prompt's table (no image rows, text positions), which must equal prompt's bits + // mrope: chunks through an image prompt's table (text positions, no image rows): prompt's bits const table = std.mem.eql(u8, args[2], "mrope"); const n_ids = ids.items.len; const src = try a.alloc(i32, n_ids); @@ -157,7 +156,7 @@ fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { } } -/// Each fixture's patches (manifest.json, tools/zig/qwen_vision_fixtures.py) through the tower; rows' cosine to 0.6.6's. +/// Each fixture's patches (tools/zig/qwen_vision_fixtures.py) through the tower: rows' cosine to 0.6.6's. fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: []const u8, out_path: []const u8) !u8 { const V = q.vision; const io = init.io; From b0aab1f7246db401cd873a11ecfc8d855bb3990d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:40:25 +0000 Subject: [PATCH 41/46] qwen3_5: kept prompt states for the shared cache, in pinned host memory (#726's reuse) cuda_reuse.zig is the family's Snapshots target: a save copies a stream's recurrences and conv tails whole, each attention layer's key and value rows up to the mark, and its rope_delta, into pinned host memory; restore copies them back. The prompt pass restores a kept state and calls the lane host at the store's marks only where a cold pass cuts a chunk (every max_rows rows from 0, or a request's chunk start), so every resumed reply keeps a cold pass's bits. The native host takes any family's snap_* functions; a snap_host family's cache is --prompt-cache-gib of host memory (8 GiB unasked, at most half of MemAvailable), on planned rules over its prompt grid. --- zig/src/families/qwen3_5/cuda.zig | 2 + zig/src/families/qwen3_5/cuda_native.zig | 69 +++++++++-- zig/src/families/qwen3_5/cuda_reuse.zig | 139 +++++++++++++++++++++++ zig/src/native/cuda.zig | 25 ++-- zig/src/native/cuda_memory.zig | 8 ++ 5 files changed, 224 insertions(+), 19 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_reuse.zig diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index f537b4c49e..59b70cf0aa 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -14,4 +14,6 @@ test { _ = draft; _ = vision; _ = @import("cuda_vision_prompt.zig"); + _ = @import("cuda_reuse.zig"); + _ = native; } diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index af7305483f..11e0d9ea9d 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -8,6 +8,7 @@ const draws = @import("draw.zig"); const Draft = @import("cuda_draft.zig").Draft; const vision = @import("cuda_vision.zig"); const vision_prompt = @import("cuda_vision_prompt.zig"); +const reuse = @import("cuda_reuse.zig"); const be = lanes.backend; pub const model_type = "qwen3_5"; @@ -18,6 +19,14 @@ pub const prompt_rows: u32 = E.max_rows; /// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. pub const takes_drafter = true; +/// Kept prompt states for the shared cache (cuda_reuse.zig), in pinned host memory, on the zero-anchored chunk grid. +pub const snap_bytes = reuse.bytes; +pub const snap_save = reuse.save; +pub const snap_restore = reuse.restore; +pub const snap_drop = reuse.drop; +pub const snap_host = true; +pub const snap_lookahead: u32 = 0; + pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, vision: bool = false }; /// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. @@ -32,6 +41,8 @@ pub const Loaded = struct { deinit: *const fn (*anyopaque) void, /// Called with the admitted stream count: that many caches allocated and touched now, as memory commits on touch. reserve: *const fn (*anyopaque, u32) anyerror!void, + /// What the prompt cache's saves copy: the engine and the streams' caches. + target: *reuse.Target, }; pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { @@ -47,7 +58,8 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ }; errdefer if (tower) |t| t.deinit(); const l = try gpa.create(Lanes); - l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .tower = tower }; + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .tower = tower, .target = undefined }; + l.target = .{ .e = e, .caches = &l.caches }; return .{ .backend = l.backend(), .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, @@ -56,6 +68,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .ctx = l, .deinit = release, .reserve = reserve, + .target = &l.target, }; } @@ -125,6 +138,32 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { }; } +/// Where a prompt pass from `at` cuts its chunk: max_rows rows on, or the request's next chunk start before that. +fn nextCut(len: usize, chunks: []const u32, at: usize) usize { + var end: usize = @min(len, at + E.max_rows); + for (chunks) |boundary| if (boundary > at) { + end = @min(end, boundary); + break; + }; + return end; +} + +/// Whether a cold pass over `len` rows (from 0) cuts a chunk at `pos`. +fn coldCut(len: usize, chunks: []const u32, pos: usize) bool { + var at: usize = 0; + while (at < pos and at < len) at = nextCut(len, chunks, at); + return at == pos; +} + +test "a pass resumes only where a cold pass cuts: every max_rows rows from 0 or from the request's chunk starts" { + try std.testing.expect(coldCut(1000, &.{}, E.max_rows * 3)); + try std.testing.expect(!coldCut(1000, &.{}, E.max_rows * 3 + 1)); + try std.testing.expect(coldCut(1000, &.{200}, 200)); + try std.testing.expect(coldCut(1000, &.{200}, 200 + E.max_rows)); + try std.testing.expect(!coldCut(1000, &.{200}, 2 * E.max_rows)); // cut at 200 instead, then 328 + try std.testing.expectEqual(@as(usize, 200), nextCut(1000, &.{200}, E.max_rows)); +} + fn release(p: *anyopaque) void { const l: *Lanes = @ptrCast(@alignCast(p)); l.deinit(); @@ -144,6 +183,7 @@ const Lanes = struct { capacity: usize, /// --vision's tower: an image prompt's images encoded at its prefill. tower: ?*vision.Tower = null, + target: reuse.Target, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. pool: std.ArrayList(*E.Cache) = .empty, @@ -225,19 +265,34 @@ const Lanes = struct { const images: ?E.Images = if (media) |m| .{ .features = m.features.ptr, .src = m.src, .pos = m.pos } else null; const cache = try self.take(); errdefer self.give(cache); + // the stream's cache is found by a save at a mark (cuda_reuse.zig) while its pass runs + try self.caches.put(self.gpa, stream, cache); + errdefer _ = self.caches.remove(stream); + // a kept state resumes only where a cold pass cuts a chunk, so the chunks after it keep a cold pass's bits var at: usize = 0; + stream.cached = 0; + if (stream.reuse.saved) |saved| { + if (stream.reuse.at > 0 and stream.reuse.at < ids.len and coldCut(ids.len, stream.chunks, stream.reuse.at)) { + if (reuse.restoreCache(self.e, cache, saved)) { + at = stream.reuse.at; + stream.cached = stream.reuse.at; + } else |_| { + try cache.reset(self.e); + stream.reuse_failed = true; + } + } else stream.reuse_failed = true; + } while (at < ids.len) { - var end: usize = @min(ids.len, at + E.max_rows); - for (stream.chunks) |boundary| if (boundary > at) { - end = @min(end, boundary); - break; - }; + const end = nextCut(ids.len, stream.chunks, at); try self.e.forward(&.{.{ .cache = cache, .rows = end - at, .images = if (images) |*im| im else null }}, ids[at..end], .last); at = end; + // a mark the store asked for, where this pass cut anyway: the lane host copies the state here + if (end < ids.len and std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(end)) != null) { + if (stream.reuse.hook) |k| k.at(k.ptr, stream, @intCast(end)); + } } // the reply's rows rotate past the image prompt's positions (rope_delta) if (media) |m| cache.delta = m.delta; - try self.caches.put(self.gpa, stream, cache); } fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig new file mode 100644 index 0000000000..69f4e76867 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -0,0 +1,139 @@ +//! Qwen3.5 CUDA prompt states for the shared cache in pinned host memory: recurrences, tails, keys, values, rope shift. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Engine = engine.Engine; +const Cache = engine.Cache; +const Geometry = geometry.Geometry; +const head_dim = geometry.head_dim; + +/// One kept state: its bytes on the host, the position it stands at, its rotary shift. +pub const Copy = struct { + host: cuda.HostBuffer, + at: usize, + delta: i32, +}; + +/// The engine and the streams' caches a save copies (cuda_native.zig's Lanes owns the map). +pub const Target = struct { + e: *Engine, + caches: *std.AutoHashMapUnmanaged(*lanes.Stream, *Cache), +}; + +/// One piece of a stream's state: where it lives on the device, its length. +const Piece = struct { dev: u64, len: usize }; + +/// Key heads a kept state copies at most, and its pieces: the state, then a key and a value block a (layer, head). +const max_kv_heads = 8; +const max_pieces = 1 + 2 * geometry.max_layers * max_kv_heads; + +/// A state's pieces: recurrences and tails whole, then each (layer, key/value, head)'s first `at` rows; null: lengths. +fn pieces(g: Geometry, cache: ?*const Cache, at: usize, out: *[max_pieces]Piece) []const Piece { + const row = head_dim * 2; + out[0] = .{ .dev = if (cache) |c| c.state.ptr else 0, .len = g.stateBytes() }; + var n: usize = 1; + for (0..g.attentions()) |l| for (0..2) |kv| for (0..g.kv_heads) |h| { + // qwen35.cu's layout: a layer's keys then values, kv_heads blocks of cap rows each + const block: usize = if (cache) |c| ((l * 2 + kv) * g.kv_heads + h) * c.cap * row else 0; + out[n] = .{ .dev = if (cache) |c| c.kv.ptr + block else 0, .len = at * row }; + n += 1; + }; + return out[0..n]; +} + +/// A state's bytes at `at` tokens. +fn total(g: Geometry, at: usize) u64 { + var buf: [max_pieces]Piece = undefined; + var n: u64 = 0; + for (pieces(g, null, at, &buf)) |p| n += p.len; + return n; +} + +/// Whether this engine's states fit a kept state's pieces (cuda_native.zig offers no cache otherwise). +pub fn fits(g: Geometry) bool { + return g.kv_heads <= max_kv_heads; +} + +/// Host bytes a save at `at` takes. +pub fn bytes(ptr: *anyopaque, at: u32) u64 { + const t: *Target = @ptrCast(@alignCast(ptr)); + return total(t.e.g, at); +} + +/// Copy the stream's state after `at` prompt tokens; its prompt pass stands there. +pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const t: *Target = @ptrCast(@alignCast(ptr)); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + if (!fits(t.e.g)) return error.UnsupportedQwenCache; + const cache = t.caches.get(s) orelse return error.UnknownStream; + if (cache.len != at or cache.last != null) return error.NotAtMark; + const e = t.e; + const copy = try e.gpa.create(Copy); + errdefer e.gpa.destroy(copy); + copy.* = .{ .host = try cuda.HostBuffer.alloc(e.d, total(e.g, at)), .at = at, .delta = cache.delta }; + errdefer copy.host.free(); + var buf: [max_pieces]Piece = undefined; + var off: usize = 0; + for (pieces(e.g, cache, at, &buf)) |p| { + if (p.len > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off, p.dev, p.len, e.stream.handle), "cuMemcpyDtoHAsync"); + off += p.len; + } + try e.stream.synchronize(); + return copy; +} + +/// Make `cache` the kept state's: its next prompt chunk starts at the copy's position. +pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque) !void { + const copy: *Copy = @ptrCast(@alignCast(saved)); + if (copy.at >= cache.cap) return error.BadReuse; + var buf: [max_pieces]Piece = undefined; + var off: usize = 0; + for (pieces(e.g, cache, copy.at, &buf)) |p| { + if (p.len > 0) try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(p.dev, copy.host.bytes.ptr + off, p.len, e.stream.handle), "cuMemcpyHtoDAsync"); + off += p.len; + } + try e.stream.synchronize(); // the host copy may be dropped once this returns + cache.len = copy.at; + cache.written = copy.at; + cache.last = null; + cache.delta = copy.delta; +} + +/// The store's restore: the owner's cache made the kept state's. +pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const t: *Target = @ptrCast(@alignCast(ptr)); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + return restoreCache(t.e, t.caches.get(s) orelse return error.UnknownStream, saved); +} + +/// Free a kept state's host copy. +pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { + const t: *Target = @ptrCast(@alignCast(ptr)); + const copy: *Copy = @ptrCast(@alignCast(saved)); + copy.host.free(); + t.e.gpa.destroy(copy); +} + +fn testGeometry() Geometry { + var g: Geometry = undefined; + g.hidden = 5120; + g.k_heads = 16; + g.v_heads = 48; + g.q_heads = 24; + g.kv_heads = 4; + g.layers = 8; + @memset(&g.linear, false); + for (0..8) |i| g.linear[i] = i % 4 != 3; // the 27B's pattern: three linear layers, then attention + return g; +} + +test "a kept state holds the recurrences and tails whole and each attention head's rows it read" { + const g = testGeometry(); + var buf: [max_pieces]Piece = undefined; + try std.testing.expectEqual(@as(usize, 1 + 2 * 2 * 4), pieces(g, null, 7, &buf).len); + try std.testing.expectEqual(@as(u64, g.stateBytes()), total(g, 0)); + try std.testing.expectEqual(@as(u64, g.stateBytes() + 2 * 2 * 4 * 1000 * head_dim * 2), total(g, 1000)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 131bdb872b..ae52d9cd9d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -19,13 +19,12 @@ const drafted = @import("cuda_drafted.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{ nemotron.native, qwen35.native }; -/// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. -const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ - .bytes = nemotron.native.snap_bytes, - .save = nemotron.native.snap_save, - .restore = nemotron.native.snap_restore, - .drop = nemotron.native.snap_drop, -}; +/// A family's prompt-cache copies. A comptime value per family, so the store's pointer stays valid for the process. +fn Snaps(comptime F: type) type { + return struct { + const vtable: api.prompt_cache.Snapshots.VTable = .{ .bytes = F.snap_bytes, .save = F.snap_save, .restore = F.snap_restore, .drop = F.snap_drop }; + }; +} pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { @@ -426,17 +425,19 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = served.backend, .served = served, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; - const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, served.stream_bytes * (streams -| free)); + const cache = if (!@hasField(L, "target")) 0 else if (@hasDecl(F, "snap_host")) budget.hostCacheBytes(a, io, o.prompt_cache_gib) else budget.cacheBytes(o.prompt_cache_gib, room, served.stream_bytes * (streams -| free)); h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(served.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", - if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(served.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache{s}", .{ toGib(cache), if (@hasDecl(F, "snap_host")) " in host memory" else "" }) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); if (@hasField(L, "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); - // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet - store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &nemotron_snaps }, .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }, cache - try nemotron.native.snap_ready(loaded.target, cache)); + // Nemotron keeps states at any row, its head one token on; a host-memory family on its zero-anchored chunk grid + const rules: pc.Rules = if (@hasDecl(F, "snap_host")) .{ .lookahead = F.snap_lookahead, .planned = true, .grid = F.prompt_rows } else .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }; + const ready: u64 = if (@hasDecl(F, "snap_ready")) try F.snap_ready(loaded.target, cache) else 0; + store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &Snaps(F).vtable }, rules, cache - ready); h.store = store; } errdefer if (h.store) |store| { diff --git a/zig/src/native/cuda_memory.zig b/zig/src/native/cuda_memory.zig index 563b78bba1..cf34402888 100644 --- a/zig/src/native/cuda_memory.zig +++ b/zig/src/native/cuda_memory.zig @@ -30,6 +30,14 @@ pub fn meminfo(text: []const u8) ?MemInfo { return if (t > 0 and a <= t) .{ .total = t, .available = a } else null; } +/// --prompt-cache-gib in pinned host memory (8 GiB unasked), at most half of what /proc/meminfo has available. +pub fn hostCacheBytes(a: std.mem.Allocator, io: std.Io, asked: ?f64) u64 { + const want: u64 = if (asked) |g| (if (g > 0) std.math.lossyCast(u64, g * (1 << 30)) else 0) else 8 << 30; + const text = procText(a, io, "/proc/meminfo") orelse return 0; + const m = meminfo(text) orelse return 0; + return @min(want, m.available / 2); +} + /// A /proc file's text, streamed: procfs reports size 0, and a positional read (readFileAlloc) stops there. pub fn procText(a: std.mem.Allocator, io: std.Io, path: []const u8) ?[]u8 { var file = std.Io.Dir.cwd().openFile(io, path, .{}) catch return null; From 55c47f9c298757ba59d8185a9746f71e3db1ff1f Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:17:16 +0000 Subject: [PATCH 42/46] tf-qwen35-cuda vision: the tower's rows bounded, exit 1 outside them Each image's rows must reach a mean cosine of 0.999 and a worst row of 0.94 against 0.6.6's bf16 rows (both settable). A fixture with 0.6.6's fp32 rows (qwen_vision_fixtures.py --fp32) judges the tower against those instead: it may not sit further from them than 0.6.6's own bf16 tower, beyond 0.001 on the mean. The 2B's tower is fragile in bf16 (0.6.6 against itself drops to 0.45 on a row), so only the reference can bound it. --- zig/tests/qwen35_cuda_main.zig | 97 ++++++++++++++++++++++++---------- 1 file changed, 69 insertions(+), 28 deletions(-) diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index cf6738e627..7f76e40314 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,4 +1,5 @@ -//! `tf-qwen35-cuda `, or `vision `: rows to compare. +//! `tf-qwen35-cuda `, or `vision [min mean] [min row]`: +//! rows to compare; vision exits 1 when the tower's rows fall under its bounds. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -7,11 +8,19 @@ const E = q.engine; pub fn main(init: std.process.Init) !u8 { const a = init.arena.allocator(); const args = try init.minimal.args.toSlice(a); - if (args.len != 5) { - std.debug.print("usage: tf-qwen35-cuda \n", .{}); + const is_vision = args.len >= 3 and std.mem.eql(u8, args[2], "vision"); + if (args.len != 5 and !(is_vision and args.len <= 7)) { + std.debug.print("usage: tf-qwen35-cuda \n" ++ + " tf-qwen35-cuda vision [min mean cosine, 0.999] [min row cosine, 0.94]\n", .{}); return 2; } - if (std.mem.eql(u8, args[2], "vision")) return vision(init, a, args[1], args[3], args[4]); + if (is_vision) { + const bounds: Bounds = .{ + .mean = if (args.len > 5) try std.fmt.parseFloat(f64, args[5]) else 0.999, + .row = if (args.len > 6) try std.fmt.parseFloat(f64, args[6]) else 0.94, + }; + return vision(init, a, args[1], args[3], args[4], bounds); + } var ids: std.ArrayList(u32) = .empty; var it = std.mem.tokenizeAny(u8, args[3], ", "); while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); @@ -156,8 +165,41 @@ fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { } } -/// Each fixture's patches (tools/zig/qwen_vision_fixtures.py) through the tower: rows' cosine to 0.6.6's. -fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: []const u8, out_path: []const u8) !u8 { +/// What the tower's rows must reach against 0.6.6's bf16 rows, each image: their mean cosine and their worst row's. +/// With a fixture's fp32 reference rows (qwen_vision_fixtures.py --fp32), the tower must also be as close to them as +/// 0.6.6's bf16 tower is, within `slack` on the mean, so a fragile tower (the 2B's) is judged against the reference. +const Bounds = struct { mean: f64, row: f64, slack: f64 = 0.001 }; + +/// Rows' cosine, mean and worst, of bf16 `got` against `want` (bf16, or fp32 when `want32`), and the identical rows. +const Cosines = struct { mean: f64, least: f64, same: usize }; + +fn cosines(got: []align(1) const u16, want: []const u8, want32: bool, rows: usize, width: usize) Cosines { + var sum: f64 = 0; + var least: f64 = 1; + var same: usize = 0; + for (0..rows) |r| { + var dot: f64 = 0; + var ng: f64 = 0; + var nw: f64 = 0; + var equal = !want32; + for (0..width) |i| { + const x = bf(got[r * width + i]); + const y: f64 = if (want32) @as(f64, @as(f32, @bitCast(std.mem.readInt(u32, want[(r * width + i) * 4 ..][0..4], .little)))) else bf(std.mem.readInt(u16, want[(r * width + i) * 2 ..][0..2], .little)); + if (!want32 and got[r * width + i] != std.mem.readInt(u16, want[(r * width + i) * 2 ..][0..2], .little)) equal = false; + dot += x * y; + ng += x * x; + nw += y * y; + } + const cos = dot / @sqrt(ng * nw); + sum += cos; + least = @min(least, cos); + same += @intFromBool(equal); + } + return .{ .mean = sum / @as(f64, @floatFromInt(rows)), .least = least, .same = same }; +} + +/// Each fixture's patches (tools/zig/qwen_vision_fixtures.py) through the tower: rows' cosine to 0.6.6's, bounded. +fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: []const u8, out_path: []const u8, bounds: Bounds) !u8 { const V = q.vision; const io = init.io; const cfg = try V.Config.read(a, io, model) orelse return error.QwenCheckpointHasNoVision; @@ -174,6 +216,7 @@ fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: const parsed = try std.json.parseFromSliceLeaky(std.json.Value, a, manifest, .{}); var out: std.ArrayList(u8) = .empty; var worst: f64 = 1; + var failed: usize = 0; for (parsed.object.get("images").?.array.items) |im| { const o = im.object; const grid = o.get("grid").?.array.items; @@ -196,32 +239,30 @@ fn vision(init: std.process.Init, a: std.mem.Allocator, model: []const u8, dir: return 1; } const gw: []align(1) const u16 = std.mem.bytesAsSlice(u16, got); - const ww: []align(1) const u16 = std.mem.bytesAsSlice(u16, want); - var sum: f64 = 0; - var least: f64 = 1; - var same: usize = 0; - for (0..tokens) |r| { - var dot: f64 = 0; - var ng: f64 = 0; - var nw: f64 = 0; - for (0..cfg.out_hidden) |i| { - const x = bf(gw[r * cfg.out_hidden + i]); - const y = bf(ww[r * cfg.out_hidden + i]); - dot += x * y; - ng += x * x; - nw += y * y; + const name = o.get("name").?.string; + const c = cosines(gw, want, false, tokens, cfg.out_hidden); + worst = @min(worst, c.least); + std.debug.print("{s}: grid {d}x{d}, {d} rows, row cosine mean {d:.5} min {d:.5}, {d} rows identical, {d:.1} ms", .{ name, g.h, g.w, tokens, c.mean, c.least, c.same, ms }); + var ok = c.mean >= bounds.mean and c.least >= bounds.row; + if (o.get("reference")) |ref| { + // both bf16 towers against 0.6.6's fp32 rows: this one may not be the further, beyond the slack + const r32 = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, ref.string }), a, .limited(1 << 31)); + if (r32.len != tokens * cfg.out_hidden * 4) { + std.debug.print("\n{s}: the fp32 reference has {d} bytes, {d} rows need {d}\n", .{ name, r32.len, tokens, tokens * cfg.out_hidden * 4 }); + return 1; } - const cos = dot / @sqrt(ng * nw); - sum += cos; - least = @min(least, cos); - if (std.mem.eql(u16, @as([]const u16, @ptrCast(@alignCast(gw[r * cfg.out_hidden ..][0..cfg.out_hidden]))), @ptrCast(@alignCast(ww[r * cfg.out_hidden ..][0..cfg.out_hidden])))) same += 1; + const ours = cosines(gw, r32, true, tokens, cfg.out_hidden); + const theirs = cosines(std.mem.bytesAsSlice(u16, want), r32, true, tokens, cfg.out_hidden); + std.debug.print("; against fp32: mean {d:.5} (0.6.6 bf16 {d:.5})", .{ ours.mean, theirs.mean }); + // with a reference, the absolute bounds give way to it: 0.6.6's own bf16 rows can sit under them + ok = ours.mean + bounds.slack >= theirs.mean; } - worst = @min(worst, least); - std.debug.print("{s}: grid {d}x{d}, {d} rows, row cosine mean {d:.5} min {d:.5}, {d} rows identical, {d:.1} ms\n", .{ o.get("name").?.string, g.h, g.w, tokens, sum / @as(f64, @floatFromInt(tokens)), least, same, ms }); + std.debug.print("{s}\n", .{if (ok) "" else " FAIL"}); + failed += @intFromBool(!ok); } try cwd.writeFile(io, .{ .sub_path = out_path, .data = out.items }); - std.debug.print("worst row cosine {d:.5}\n", .{worst}); - return 0; + std.debug.print("worst row cosine {d:.5}; {d} image{s} outside the bounds (mean {d}, row {d}, fp32 slack {d})\n", .{ worst, failed, if (failed == 1) "" else "s", bounds.mean, bounds.row, bounds.slack }); + return if (failed == 0) 0 else 1; } fn bf(h: u16) f64 { From 6e29ebed00912f435dc0dc31e642e683c1fb757b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:21:41 +0000 Subject: [PATCH 43/46] lanes: an image's BLAKE3 digest; the cache key's placeholder ids come from it lanes.Image.digest covers the grid and every patch bit. image_key seeds each run's ids from it and digests each image once. The 31-bit ids only find a kept state: the backend compares whole digests before it restores one. --- zig/src/core/lanes/stream.zig | 24 +++++++++++++++++++++- zig/src/server/image_key.zig | 38 ++++++++++++++++++++++++----------- 2 files changed, 49 insertions(+), 13 deletions(-) diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index 6d634c9600..bc0574e1ca 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -98,7 +98,20 @@ pub const Held = struct { }; /// An image for a vision tower: patches (fp32 [h * w, 1536], merge-block order), grid [1, h, w]; prompt placeholders. -pub const Image = struct { patches: []const f32, grid: [3]i64 }; +pub const Image = struct { + patches: []const f32, + grid: [3]i64, + + /// BLAKE3 of the grid (little-endian) and the patches' bits: the image's identity for kept prompt states. + pub fn digest(im: Image) [32]u8 { + var h = std.crypto.hash.Blake3.init(.{}); + for (im.grid) |g| h.update(&std.mem.toBytes(std.mem.nativeToLittle(i64, g))); + h.update(std.mem.sliceAsBytes(im.patches)); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; pub const Spec = struct { id: []const u8, @@ -453,3 +466,12 @@ test "loop guard fires, closes thinking, and latches once" { _ = try s.commit(gpa, &.{8}, &.{}); try std.testing.expect(!s.finished and s.loop_period.? == 1); } + +test "an image's digest covers its grid and every patch bit" { + const a: [3]f32 = .{ 0.1, 0.2, 0.3 }; + const b: [3]f32 = .{ 0.1, 0.2, -0.3 }; + const da = (Image{ .patches = &a, .grid = .{ 1, 2, 2 } }).digest(); + try std.testing.expect(!std.mem.eql(u8, &da, &(Image{ .patches = &b, .grid = .{ 1, 2, 2 } }).digest())); + try std.testing.expect(!std.mem.eql(u8, &da, &(Image{ .patches = &a, .grid = .{ 1, 1, 4 } }).digest())); + try std.testing.expectEqualSlices(u8, &da, &(Image{ .patches = &a, .grid = .{ 1, 2, 2 } }).digest()); +} diff --git a/zig/src/server/image_key.zig b/zig/src/server/image_key.zig index 175b8d0809..0819d481d4 100644 --- a/zig/src/server/image_key.zig +++ b/zig/src/server/image_key.zig @@ -1,27 +1,26 @@ -//! An image prompt's prompt-cache key: each image's placeholder rows as ids drawn from that image's content. +//! An image prompt's cache key: placeholder rows as ids from each image's digest; restores compare whole digests. const std = @import("std"); const api = @import("engine_api"); /// Ids the tokenizer never makes: the top bit set. const marked: u32 = 0x8000_0000; -/// One image's content: its grid and its patches' bits. -fn imageHash(im: api.Image) u64 { - var h = std.hash.Wyhash.init(0x51ce_7a91); - h.update(std.mem.asBytes(&im.grid)); - h.update(std.mem.sliceAsBytes(im.patches)); - return h.final(); -} - -/// `ids` with placeholder run k (ids in `pads`) as ids from image k's content; unpaired runs key by all the images. +/// `ids` with placeholder run k (ids in `pads`) as ids from image k's digest; unpaired runs key by all the images. pub fn cacheKey(a: std.mem.Allocator, ids: []const u32, images: []const api.Image, pads: []const u32) ![]u32 { const key = try a.dupe(u32, ids); + errdefer a.free(key); var runs: usize = 0; for (ids, 0..) |t, i| { if (isPad(t, pads) and (i == 0 or !isPad(ids[i - 1], pads))) runs += 1; } + const seeds = try a.alloc(u64, images.len); + defer a.free(seeds); var all = std.hash.Wyhash.init(0x51ce_7a92); - for (images) |im| all.update(std.mem.asBytes(&imageHash(im))); + for (images, seeds) |im, *seed| { + const d = im.digest(); + all.update(&d); + seed.* = std.mem.readInt(u64, d[0..8], .little); + } const together = all.final(); var run: usize = 0; var row: u64 = 0; @@ -29,7 +28,7 @@ pub fn cacheKey(a: std.mem.Allocator, ids: []const u32, images: []const api.Imag for (key, 0..) |*t, i| { if (!isPad(ids[i], pads)) continue; if (i == 0 or !isPad(ids[i - 1], pads)) { - h = if (runs == images.len) imageHash(images[run]) else together; + h = if (runs == images.len) seeds[run] else together; run += 1; row = 0; } @@ -81,3 +80,18 @@ test "a second image's change keys only its own run apart; runs that do not pair defer a.free(one); try std.testing.expect(one[1] != k12[1]); // two runs, one image: neither run keys as image 1 alone } + +test "two different images behind one placeholder key apart" { + const a = std.testing.allocator; + const ids = [_]u32{ 7, 52, 9, 53, 8 }; + var seen: [64]u32 = undefined; + // 64 one-row images that differ in one patch value each: 64 different keys at the placeholder + for (&seen, 0..) |*k, i| { + const px: [2]f32 = .{ 0.5, @floatFromInt(i) }; + const key = try cacheKey(a, &ids, &.{.{ .patches = &px, .grid = .{ 1, 1, 1 } }}, &.{9}); + defer a.free(key); + try std.testing.expectEqualSlices(u32, &.{ 7, 52 }, key[0..2]); + k.* = key[2]; + for (seen[0..i]) |o| try std.testing.expect(o != k.*); + } +} From e64191ae9ee4e379b8014ac03ac110295bbd949b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:21:41 +0000 Subject: [PATCH 44/46] qwen3_5: a kept state is checked before its restore: any geometry, its cache format, its images cuda_reuse.zig's pieces are computed by index from the geometry, so no layer or head count passes a fixed array (33 attention layers of 32 key heads did). A kept state records its cache format, and its byte layout comes from it: bf16 rows, or with an engine that has --kv-cache fp8, e4m3 rows and their scales as qwen35_kv8.cu lays them out. It also records the digest of every image whose placeholders start before its position. A restore of another format, length or image set is refused before the stream's cache is touched, and the prompt is prefilled from the start. --- zig/src/families/qwen3_5/cuda_native.zig | 28 ++- zig/src/families/qwen3_5/cuda_reuse.zig | 255 ++++++++++++++++++----- 2 files changed, 225 insertions(+), 58 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 11e0d9ea9d..306d84df4e 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -59,7 +59,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer if (tower) |t| t.deinit(); const l = try gpa.create(Lanes); l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .tower = tower, .target = undefined }; - l.target = .{ .e = e, .caches = &l.caches }; + l.target = .{ .e = e, .caches = &l.caches, .seen = &l.seen }; return .{ .backend = l.backend(), .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, @@ -185,6 +185,8 @@ const Lanes = struct { tower: ?*vision.Tower = null, target: reuse.Target, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + /// An image prompt's images as its kept states name them (cuda_reuse.zig), from its prefill to its release. + seen: std.AutoHashMapUnmanaged(*lanes.Stream, []reuse.Seen) = .empty, /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. pool: std.ArrayList(*E.Cache) = .empty, @@ -192,6 +194,9 @@ const Lanes = struct { var it = self.caches.valueIterator(); while (it.next()) |cache| self.drop(cache.*); self.caches.deinit(self.gpa); + var seen = self.seen.valueIterator(); + while (seen.next()) |v| self.gpa.free(v.*); + self.seen.deinit(self.gpa); for (self.pool.items) |cache| self.drop(cache); self.pool.deinit(self.gpa); if (self.tower) |t| t.deinit(); @@ -263,6 +268,12 @@ const Lanes = struct { } defer if (media) |*m| m.deinit(); const images: ?E.Images = if (media) |m| .{ .features = m.features.ptr, .src = m.src, .pos = m.pos } else null; + const seen = try self.seenOf(ids, stream.images); + self.seen.put(self.gpa, stream, seen) catch |err| { + self.gpa.free(seen); + return err; + }; + errdefer if (self.seen.fetchRemove(stream)) |kv| self.gpa.free(kv.value); const cache = try self.take(); errdefer self.give(cache); // the stream's cache is found by a save at a mark (cuda_reuse.zig) while its pass runs @@ -273,7 +284,7 @@ const Lanes = struct { stream.cached = 0; if (stream.reuse.saved) |saved| { if (stream.reuse.at > 0 and stream.reuse.at < ids.len and coldCut(ids.len, stream.chunks, stream.reuse.at)) { - if (reuse.restoreCache(self.e, cache, saved)) { + if (reuse.restoreCache(self.e, cache, saved, seen)) { at = stream.reuse.at; stream.cached = stream.reuse.at; } else |_| { @@ -415,5 +426,18 @@ const Lanes = struct { fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { const self = cast(ptr); if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); + if (self.seen.fetchRemove(stream)) |entry| self.gpa.free(entry.value); + } + + /// Each image's first placeholder row and digest: a kept state resumes only the same images. + fn seenOf(self: *Lanes, ids: []const u32, images: []const lanes.Image) ![]reuse.Seen { + if (images.len == 0) return self.gpa.alloc(reuse.Seen, 0); + const tower = self.tower orelse return error.QwenServedWithoutVision; + const runs = try self.gpa.alloc(vision_prompt.Run, images.len); + defer self.gpa.free(runs); + try vision_prompt.placeholders(tower.c, ids, images, runs); + const seen = try self.gpa.alloc(reuse.Seen, images.len); + for (seen, runs, images) |*s, r, im| s.* = .{ .at = @intCast(r.at), .digest = im.digest() }; + return seen; } }; diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index 69f4e76867..e14cdb50a3 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -10,89 +10,153 @@ const Cache = engine.Cache; const Geometry = geometry.Geometry; const head_dim = geometry.head_dim; -/// One kept state: its bytes on the host, the position it stands at, its rotary shift. +/// A cache's attention rows: bf16, or e4m3 with an fp32 scale each (--kv-cache fp8, where the engine has it). +pub const Format = enum(u8) { bf16, fp8 }; + +/// The engine's cache format; an engine without fp8 caches keeps bf16 rows. +pub fn formatOf(e: *const Engine) Format { + if (comptime @hasField(Engine, "kv8")) return if (e.kv8) .fp8 else .bf16; + return .bf16; +} + +/// An image a kept state's prompt holds: the row its placeholders start at, its digest (lanes.Image.digest). +pub const Seen = struct { at: u32, digest: [32]u8 }; + +/// One kept state: its bytes on the host, the position it stands at, its rotary shift, its cache format and images. pub const Copy = struct { host: cuda.HostBuffer, + len: u64, at: usize, delta: i32, + format: Format, + images: []Seen, }; -/// The engine and the streams' caches a save copies (cuda_native.zig's Lanes owns the map). +/// The engine, the streams' caches a save copies, and their prompts' images (cuda_native.zig's Lanes owns the maps). pub const Target = struct { e: *Engine, caches: *std.AutoHashMapUnmanaged(*lanes.Stream, *Cache), + seen: *std.AutoHashMapUnmanaged(*lanes.Stream, []Seen), }; -/// One piece of a stream's state: where it lives on the device, its length. -const Piece = struct { dev: u64, len: usize }; - -/// Key heads a kept state copies at most, and its pieces: the state, then a key and a value block a (layer, head). -const max_kv_heads = 8; -const max_pieces = 1 + 2 * geometry.max_layers * max_kv_heads; - -/// A state's pieces: recurrences and tails whole, then each (layer, key/value, head)'s first `at` rows; null: lengths. -fn pieces(g: Geometry, cache: ?*const Cache, at: usize, out: *[max_pieces]Piece) []const Piece { - const row = head_dim * 2; - out[0] = .{ .dev = if (cache) |c| c.state.ptr else 0, .len = g.stateBytes() }; - var n: usize = 1; - for (0..g.attentions()) |l| for (0..2) |kv| for (0..g.kv_heads) |h| { - // qwen35.cu's layout: a layer's keys then values, kv_heads blocks of cap rows each - const block: usize = if (cache) |c| ((l * 2 + kv) * g.kv_heads + h) * c.cap * row else 0; - out[n] = .{ .dev = if (cache) |c| c.kv.ptr + block else 0, .len = at * row }; - n += 1; - }; - return out[0..n]; -} +/// One piece of a stream's state: its offset in the state or key/value buffer, its length. +const Piece = struct { kv: bool, off: usize, len: usize }; -/// A state's bytes at `at` tokens. -fn total(g: Geometry, at: usize) u64 { - var buf: [max_pieces]Piece = undefined; - var n: u64 = 0; - for (pieces(g, null, at, &buf)) |p| n += p.len; - return n; -} +/// A state's pieces by index, any geometry: the state, then each (layer, key/value, head)'s first `at` rows. +const Layout = struct { + g: Geometry, + format: Format, + cap: usize, + at: usize, -/// Whether this engine's states fit a kept state's pieces (cuda_native.zig offers no cache otherwise). -pub fn fits(g: Geometry) bool { - return g.kv_heads <= max_kv_heads; + fn parts(l: Layout) usize { + return if (l.format == .fp8) 2 else 1; + } + + fn count(l: Layout) usize { + return 1 + l.g.attentions() * 2 * l.g.kv_heads * l.parts(); + } + + /// Piece `i`: qwen35.cu's layout (a layer's keys then values, kv_heads blocks of cap rows), or qwen35_kv8.cu's. + fn piece(l: Layout, i: usize) Piece { + if (i == 0) return .{ .kv = false, .off = 0, .len = l.g.stateBytes() }; + const j = (i - 1) / l.parts(); + const kvh = l.g.kv_heads; + const block = j % (2 * kvh); // key heads, then value heads + const layer = j / (2 * kvh); + return switch (l.format) { + .bf16 => .{ .kv = true, .off = (layer * 2 * kvh + block) * l.cap * head_dim * 2, .len = l.at * head_dim * 2 }, + .fp8 => blk: { + const base = layer * 2 * kvh * l.cap * (head_dim + 4); + break :blk if ((i - 1) % 2 == 0) + .{ .kv = true, .off = base + block * l.cap * head_dim, .len = l.at * head_dim } + else + .{ .kv = true, .off = base + 2 * kvh * l.cap * head_dim + block * l.cap * 4, .len = l.at * 4 }; + }, + }; + } + + /// The bytes a save at `at` takes. + fn total(l: Layout) u64 { + var n: u64 = 0; + for (0..l.count()) |i| n += l.piece(i).len; + return n; + } + + /// The key/value buffer's bytes at `cap` (Geometry.kvBytes). + fn kvBytes(l: Layout) usize { + return l.g.attentions() * 2 * l.g.kv_heads * l.cap * @as(usize, if (l.format == .fp8) head_dim + 4 else head_dim * 2); + } +}; + +/// Whether a request's images match a kept state's: the same digests at the same rows, for every image before `at`. +pub fn sameImages(kept: []const Seen, request: []const Seen, at: usize) bool { + var n: usize = 0; + for (request) |r| { + if (r.at >= at) continue; + if (n >= kept.len or kept[n].at != r.at or !std.mem.eql(u8, &kept[n].digest, &r.digest)) return false; + n += 1; + } + return n == kept.len; } /// Host bytes a save at `at` takes. pub fn bytes(ptr: *anyopaque, at: u32) u64 { const t: *Target = @ptrCast(@alignCast(ptr)); - return total(t.e.g, at); + return (Layout{ .g = t.e.g, .format = formatOf(t.e), .cap = 0, .at = at }).total(); } /// Copy the stream's state after `at` prompt tokens; its prompt pass stands there. pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { const t: *Target = @ptrCast(@alignCast(ptr)); const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - if (!fits(t.e.g)) return error.UnsupportedQwenCache; const cache = t.caches.get(s) orelse return error.UnknownStream; - if (cache.len != at or cache.last != null) return error.NotAtMark; + if (cache.len != at or cache.last != null or at > cache.cap) return error.NotAtMark; const e = t.e; + const l: Layout = .{ .g = e.g, .format = formatOf(e), .cap = cache.cap, .at = at }; + const seen: []const Seen = t.seen.get(s) orelse &.{}; + var n: usize = 0; + for (seen) |im| n += @intFromBool(im.at < at); const copy = try e.gpa.create(Copy); errdefer e.gpa.destroy(copy); - copy.* = .{ .host = try cuda.HostBuffer.alloc(e.d, total(e.g, at)), .at = at, .delta = cache.delta }; + const images = try e.gpa.alloc(Seen, n); + errdefer e.gpa.free(images); + n = 0; + for (seen) |im| if (im.at < at) { + images[n] = im; + n += 1; + }; + copy.* = .{ .host = try cuda.HostBuffer.alloc(e.d, l.total()), .len = l.total(), .at = at, .delta = cache.delta, .format = l.format, .images = images }; errdefer copy.host.free(); - var buf: [max_pieces]Piece = undefined; var off: usize = 0; - for (pieces(e.g, cache, at, &buf)) |p| { - if (p.len > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off, p.dev, p.len, e.stream.handle), "cuMemcpyDtoHAsync"); + for (0..l.count()) |i| { + const p = l.piece(i); + const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; + if (p.len > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off, dev, p.len, e.stream.handle), "cuMemcpyDtoHAsync"); off += p.len; } try e.stream.synchronize(); return copy; } -/// Make `cache` the kept state's: its next prompt chunk starts at the copy's position. -pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque) !void { +/// Whether `cache` can take the kept state: the same cache format and bytes, a position inside it, the same images. +pub fn restorable(e: *const Engine, cache: *const Cache, copy: *const Copy, seen: []const Seen) bool { + if (copy.format != formatOf(e) or copy.at >= cache.cap) return false; + const l: Layout = .{ .g = e.g, .format = copy.format, .cap = cache.cap, .at = copy.at }; + return copy.len == l.total() and copy.host.bytes.len >= copy.len and sameImages(copy.images, seen, copy.at); +} + +/// Make `cache` the kept state's: its next prompt chunk starts at the copy's position; refused before any copy. +pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque, seen: []const Seen) !void { const copy: *Copy = @ptrCast(@alignCast(saved)); - if (copy.at >= cache.cap) return error.BadReuse; - var buf: [max_pieces]Piece = undefined; + if (copy.format != formatOf(e)) return error.KeptCacheFormat; + if (!restorable(e, cache, copy, seen)) return error.BadReuse; + const l: Layout = .{ .g = e.g, .format = copy.format, .cap = cache.cap, .at = copy.at }; var off: usize = 0; - for (pieces(e.g, cache, copy.at, &buf)) |p| { - if (p.len > 0) try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(p.dev, copy.host.bytes.ptr + off, p.len, e.stream.handle), "cuMemcpyHtoDAsync"); + for (0..l.count()) |i| { + const p = l.piece(i); + const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; + if (p.len > 0) try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dev, copy.host.bytes.ptr + off, p.len, e.stream.handle), "cuMemcpyHtoDAsync"); off += p.len; } try e.stream.synchronize(); // the host copy may be dropped once this returns @@ -106,7 +170,7 @@ pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque) !void { pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { const t: *Target = @ptrCast(@alignCast(ptr)); const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - return restoreCache(t.e, t.caches.get(s) orelse return error.UnknownStream, saved); + return restoreCache(t.e, t.caches.get(s) orelse return error.UnknownStream, saved, t.seen.get(s) orelse &.{}); } /// Free a kept state's host copy. @@ -114,26 +178,105 @@ pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { const t: *Target = @ptrCast(@alignCast(ptr)); const copy: *Copy = @ptrCast(@alignCast(saved)); copy.host.free(); + t.e.gpa.free(copy.images); t.e.gpa.destroy(copy); } -fn testGeometry() Geometry { +fn testGeometry(layers: usize, attention_every: usize, kv_heads: usize) Geometry { var g: Geometry = undefined; g.hidden = 5120; g.k_heads = 16; g.v_heads = 48; - g.q_heads = 24; - g.kv_heads = 4; - g.layers = 8; + g.q_heads = 32; + g.kv_heads = kv_heads; + g.layers = layers; @memset(&g.linear, false); - for (0..8) |i| g.linear[i] = i % 4 != 3; // the 27B's pattern: three linear layers, then attention + for (0..layers) |i| g.linear[i] = i % attention_every != attention_every - 1; return g; } test "a kept state holds the recurrences and tails whole and each attention head's rows it read" { - const g = testGeometry(); - var buf: [max_pieces]Piece = undefined; - try std.testing.expectEqual(@as(usize, 1 + 2 * 2 * 4), pieces(g, null, 7, &buf).len); - try std.testing.expectEqual(@as(u64, g.stateBytes()), total(g, 0)); - try std.testing.expectEqual(@as(u64, g.stateBytes() + 2 * 2 * 4 * 1000 * head_dim * 2), total(g, 1000)); + const g = testGeometry(8, 4, 4); // the 27B's pattern: three linear layers, then attention + const l: Layout = .{ .g = g, .format = .bf16, .cap = 2048, .at = 1000 }; + try std.testing.expectEqual(@as(usize, 1 + 2 * 2 * 4), l.count()); + try std.testing.expectEqual(@as(u64, g.stateBytes()), (Layout{ .g = g, .format = .bf16, .cap = 2048, .at = 0 }).total()); + try std.testing.expectEqual(@as(u64, g.stateBytes() + 2 * 2 * 4 * 1000 * head_dim * 2), l.total()); + // fp8: e4m3 rows and a scale each, 260 bytes a position and head + const l8: Layout = .{ .g = g, .format = .fp8, .cap = 2048, .at = 1000 }; + try std.testing.expectEqual(@as(usize, 1 + 2 * 2 * 2 * 4), l8.count()); + try std.testing.expectEqual(@as(u64, g.stateBytes() + 2 * 2 * 4 * 1000 * (head_dim + 4)), l8.total()); +} + +test "every geometry's pieces lie inside its allocation, apart, with no fixed limit on layers or heads" { + const a = std.testing.allocator; + // 33 attention layers of 32 key heads passed the old fixed array; the widest the engine reads is 128 layers + for ([_][3]usize{ .{ 8, 4, 4 }, .{ 99, 3, 32 }, .{ 33, 1, 32 }, .{ 128, 1, 32 }, .{ 24, 4, 2 } }) |shape| { + const g = testGeometry(shape[0], shape[1], shape[2]); + for ([_]Format{ .bf16, .fp8 }) |f| for ([_]usize{ 1, 7, 64 }) |at| { + const l: Layout = .{ .g = g, .format = f, .cap = 64, .at = at }; + const end = l.kvBytes(); + if (comptime @typeInfo(@TypeOf(Geometry.kvBytes)).@"fn".param_types.len == 3) try std.testing.expectEqual(g.kvBytes(64, f == .fp8), end); + const spans = try a.alloc([2]usize, l.count() - 1); + defer a.free(spans); + try std.testing.expectEqual(@as(usize, g.stateBytes()), l.piece(0).len); + for (spans, 1..) |*sp, i| { + const p = l.piece(i); + try std.testing.expect(p.kv and p.len > 0 and p.off + p.len <= end); + sp.* = .{ p.off, p.off + p.len }; + } + std.mem.sort([2]usize, spans, {}, struct { + fn lt(_: void, x: [2]usize, y: [2]usize) bool { + return x[0] < y[0]; + } + }.lt); + for (spans[1..], spans[0 .. spans.len - 1]) |sp, prev| try std.testing.expect(sp[0] >= prev[1]); + }; + } +} + +test "fp8 pieces follow qwen35_kv8.cu: a layer's key then value rows, then their scales" { + const g = testGeometry(4, 4, 2); // one attention layer + const l: Layout = .{ .g = g, .format = .fp8, .cap = 100, .at = 10 }; + // (key, head 1): rows at 1 * cap * 256, its scales after both blocks' rows, at 2 * 2 * cap * 256 + 1 * cap * 4 + try std.testing.expectEqual(Piece{ .kv = true, .off = 100 * 256, .len = 10 * 256 }, l.piece(1 + 2 * 1)); + try std.testing.expectEqual(Piece{ .kv = true, .off = 4 * 100 * 256 + 100 * 4, .len = 10 * 4 }, l.piece(2 + 2 * 1)); + // (value, head 0): the third block's rows and scales + try std.testing.expectEqual(Piece{ .kv = true, .off = 2 * 100 * 256, .len = 10 * 256 }, l.piece(1 + 2 * 2)); + try std.testing.expectEqual(Piece{ .kv = true, .off = 4 * 100 * 256 + 2 * 100 * 4, .len = 10 * 4 }, l.piece(2 + 2 * 2)); +} + +test "a state resumes only a request with the same images before its position" { + const d = [_][32]u8{ @splat(1), @splat(2), @splat(3) }; + const kept = [_]Seen{.{ .at = 5, .digest = d[0] }}; + try std.testing.expect(sameImages(&kept, &.{ .{ .at = 5, .digest = d[0] }, .{ .at = 40, .digest = d[2] } }, 20)); // a later image is the pass's + try std.testing.expect(!sameImages(&kept, &.{.{ .at = 5, .digest = d[1] }}, 20)); // another image under the same ids + try std.testing.expect(!sameImages(&kept, &.{.{ .at = 6, .digest = d[0] }}, 20)); + try std.testing.expect(!sameImages(&kept, &.{}, 20)); // a text prompt never resumes an image's state + try std.testing.expect(!sameImages(&.{}, &.{.{ .at = 5, .digest = d[0] }}, 20)); // nor an image prompt a text state past its image + try std.testing.expect(sameImages(&.{}, &.{.{ .at = 5, .digest = d[0] }}, 5)); // a text prefix before it +} + +test "a restore refuses another cache format, another length, a position past the cache and other images" { + var e: Engine = undefined; + e.g = testGeometry(8, 4, 4); + if (comptime @hasField(Engine, "kv8")) e.kv8 = false; + var cache: Cache = undefined; + cache.cap = 64; + const d: [32]u8 = @splat(9); + var images = [_]Seen{.{ .at = 2, .digest = d }}; + const l: Layout = .{ .g = e.g, .format = .bf16, .cap = 64, .at = 10 }; + const host = try std.testing.allocator.alignedAlloc(u8, .@"16", @intCast(l.total())); + defer std.testing.allocator.free(host); + var copy: Copy = .{ .host = .{ .d = undefined, .bytes = host }, .len = l.total(), .at = 10, .delta = 0, .format = .bf16, .images = &images }; + const same = [_]Seen{.{ .at = 2, .digest = d }}; + try std.testing.expect(restorable(&e, &cache, ©, &same)); + try std.testing.expect(!restorable(&e, &cache, ©, &.{.{ .at = 2, .digest = @splat(8) }})); + copy.format = .fp8; // a state kept from fp8 rows never lands in bf16 ones, nor the other way + try std.testing.expect(!restorable(&e, &cache, ©, &same)); + copy.format = .bf16; + copy.len -= 1; + try std.testing.expect(!restorable(&e, &cache, ©, &same)); + copy.len += 1; + cache.cap = 10; + try std.testing.expect(!restorable(&e, &cache, ©, &same)); } From 78a222b8db53c26cd6178974c5a2fe6c62628153 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:29:32 +0000 Subject: [PATCH 45/46] tf-qwen35-cuda: resume from a kept copy (TF_Q35_RESUME) and a rope shift (TF_Q35_DELTA) TF_Q35_RESUME=1 restores the stream after its prefill into a fresh, poisoned cache through cuda_reuse.zig's save and restore before step, window or tree; TF_Q35_RESUME=K does it at position K of a prompt or mrope pass. TF_Q35_DELTA=d rotates the rows after the prefill at position + d, as an image prompt's reply. cuda_reuse.zig's saveCache and dropCopy take a cache. --- zig/src/families/qwen3_5/cuda.zig | 3 ++- zig/src/families/qwen3_5/cuda_reuse.zig | 16 +++++++++++----- zig/tests/qwen35_cuda_main.zig | 23 ++++++++++++++++++++++- 3 files changed, 35 insertions(+), 7 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 59b70cf0aa..5e27358447 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -5,6 +5,7 @@ pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); pub const draft = @import("cuda_draft.zig"); pub const vision = @import("cuda_vision.zig"); +pub const reuse = @import("cuda_reuse.zig"); test { _ = config; @@ -14,6 +15,6 @@ test { _ = draft; _ = vision; _ = @import("cuda_vision_prompt.zig"); - _ = @import("cuda_reuse.zig"); + _ = reuse; _ = native; } diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index e14cdb50a3..41f93ee298 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -111,10 +111,13 @@ pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { const t: *Target = @ptrCast(@alignCast(ptr)); const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); const cache = t.caches.get(s) orelse return error.UnknownStream; + return @ptrCast(try saveCache(t.e, cache, at, t.seen.get(s) orelse &.{})); +} + +/// A host copy of `cache` at `at` rows, naming the prompt's images that start before it. +pub fn saveCache(e: *Engine, cache: *const Cache, at: usize, seen: []const Seen) !*Copy { if (cache.len != at or cache.last != null or at > cache.cap) return error.NotAtMark; - const e = t.e; const l: Layout = .{ .g = e.g, .format = formatOf(e), .cap = cache.cap, .at = at }; - const seen: []const Seen = t.seen.get(s) orelse &.{}; var n: usize = 0; for (seen) |im| n += @intFromBool(im.at < at); const copy = try e.gpa.create(Copy); @@ -176,10 +179,13 @@ pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror! /// Free a kept state's host copy. pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { const t: *Target = @ptrCast(@alignCast(ptr)); - const copy: *Copy = @ptrCast(@alignCast(saved)); + dropCopy(t.e, @ptrCast(@alignCast(saved))); +} + +pub fn dropCopy(e: *Engine, copy: *Copy) void { copy.host.free(); - t.e.gpa.free(copy.images); - t.e.gpa.destroy(copy); + e.gpa.free(copy.images); + e.gpa.destroy(copy); } fn testGeometry(layers: usize, attention_every: usize, kv_heads: usize) Geometry { diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 7f76e40314..f69a8175b4 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,5 +1,5 @@ //! `tf-qwen35-cuda `, or `vision [min mean] [min row]`: -//! rows to compare; vision exits 1 when the tower's rows fall under its bounds. +//! rows to compare; vision exits 1 when the tower's rows fall under its bounds. TF_Q35_RESUME resumes from a kept copy. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -36,6 +36,11 @@ pub fn main(init: std.process.Init) !u8 { var cache = try E.Cache.init(init.gpa, e, 4096 + prefill); defer cache.deinit(init.gpa); try fill(e, &cache, ids.items, prefill); + // TF_Q35_DELTA=d: the rows after the prefill rotate at position + d, as an image prompt's reply does + if (std.c.getenv("TF_Q35_DELTA")) |v| cache.delta = try std.fmt.parseInt(i32, std.mem.span(v), 10); + // TF_Q35_RESUME=1: the stream goes on from a kept copy of its state, restored into another (poisoned) cache + const resume_at = try envCount("TF_Q35_RESUME"); + if (resume_at > 0 and !std.mem.eql(u8, args[2], "prompt") and !std.mem.eql(u8, args[2], "mrope")) try resumeFrom(init.gpa, e, &cache); var other = try E.Cache.init(init.gpa, e, 4096 + beside); defer other.deinit(init.gpa); try fill(e, &other, ids.items, beside); @@ -124,6 +129,8 @@ pub fn main(init: std.process.Init) !u8 { try e.forward(&.{.{ .cache = &cache, .rows = n, .images = &part }}, ids.items[at..][0..n], .last); } else try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); at += n; + // TF_Q35_RESUME=K in prompt and mrope: the pass goes on from a kept copy at position K + if (resume_at > 0 and cache.len == resume_at) try resumeFrom(init.gpa, e, &cache); } try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); } @@ -131,6 +138,20 @@ pub fn main(init: std.process.Init) !u8 { return 0; } +/// `cache` replaced by a fresh one, poisoned, that takes a kept copy of its state (cuda_reuse.zig's save and restore). +fn resumeFrom(gpa: std.mem.Allocator, e: *E.Engine, cache: *E.Cache) !void { + const copy = try q.reuse.saveCache(e, cache, cache.len, &.{}); + defer q.reuse.dropCopy(e, copy); + var fresh = try E.Cache.init(gpa, e, cache.cap); + errdefer fresh.deinit(gpa); + try fresh.state.fill8(0x5a, e.stream.handle); + try fresh.kv.fill8(0x5a, e.stream.handle); + try q.reuse.restoreCache(e, &fresh, copy, &.{}); + if (fresh.len != cache.len or fresh.delta != cache.delta) return error.ResumeMismatch; + cache.deinit(gpa); + cache.* = fresh; +} + /// One forward keeping logits: the stream's part first (its rows' logits come first), then the second stream's step. const Fwd = struct { e: *E.Engine, From 5296eb19f735a544a324da7cabd34e8c45484429 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:59:41 +0000 Subject: [PATCH 46/46] tf-qwen35-cuda: one-line comments within 120 columns (the lean check) --- zig/tests/qwen35_cuda_main.zig | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 7f76e40314..947ea9927a 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,5 +1,4 @@ -//! `tf-qwen35-cuda `, or `vision [min mean] [min row]`: -//! rows to compare; vision exits 1 when the tower's rows fall under its bounds. +//! `tf-qwen35-cuda ...`: rows to compare; vision exits 1 off its bounds. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -165,9 +164,7 @@ fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { } } -/// What the tower's rows must reach against 0.6.6's bf16 rows, each image: their mean cosine and their worst row's. -/// With a fixture's fp32 reference rows (qwen_vision_fixtures.py --fp32), the tower must also be as close to them as -/// 0.6.6's bf16 tower is, within `slack` on the mean, so a fragile tower (the 2B's) is judged against the reference. +/// Each image's mean and worst row cosine to 0.6.6's bf16 rows; with fp32 rows, no further from them than 0.6.6's. const Bounds = struct { mean: f64, row: f64, slack: f64 = 0.001 }; /// Rows' cosine, mean and worst, of bf16 `got` against `want` (bf16, or fp32 when `want32`), and the identical rows.