From 39de0393a2a3d52844d8d0d7714504d0261aedfb Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 18:36:00 +0000 Subject: [PATCH 01/11] 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 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 02/11] 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 03/11] 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 04/11] 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 05/11] 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 06/11] 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 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 07/11] 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 67ff102b37d82b02ab8650acf6b50419f5f69cf8 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:54:23 +0000 Subject: [PATCH 08/11] cuda: ggml IQ2_XXS and Q8_0 blocks, as grouped experts and dense rows (Nemotron 3.5 Bolt's export) NVIDIA's Nemotron-3.5-Bolt-74B-A9B-24GB stores its experts as ggml IQ2_XXS (66 bytes per 256 weights) and its latent, attention and head projections as Q8_0 (34 bytes per 32), U8 blocks in safetensors. cuda.ggml packs them on the GPU into 16-column tiles (a byte permutation, the same size) and reads them as: - Grouped experts (Format.iq2xxs): decode and prompt, relu^2 fused into up, down to fp32 or bf16, on the existing plan contract. - Dense rows (qlinear .q8_0 and .iq2xxs): decode at 16 rows or fewer, the prompt GEMM for any count. The IQ2_XXS grid and sign tables are ggml's (MIT). tools/zig/bolt_ggml.py decodes bit-identically to libggml on the checkpoint; tf-cuda-test ggml checks both forms against an fp64 matmul of it on real tensors (experts, the shared expert, fc1_latent, q_proj, lm_head rows): fp32 out within 5.3e-7 of the largest value. lm_head's Q8_0 GEMV reads at 1.65-1.68 TB/s on an RTX 5090. --- tools/zig/bolt_ggml.py | 95 +++++ tools/zig/bolt_ggml_fixture.py | 101 +++++ zig/build/cuda.zig | 1 + zig/kernels/cuda/ggml.cu | 470 ++++++++++++++++++++++++ zig/kernels/cuda/ggml_iq2xxs_tables.cuh | 84 +++++ zig/src/cuda/experts.zig | 30 +- zig/src/cuda/ggml.zig | 218 +++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/cuda/qlinear.zig | 9 + zig/src/cuda/root.zig | 2 + zig/tests/cuda/ggml_tests.zig | 336 +++++++++++++++++ zig/tests/cuda/main.zig | 5 + 12 files changed, 1347 insertions(+), 5 deletions(-) create mode 100644 tools/zig/bolt_ggml.py create mode 100644 tools/zig/bolt_ggml_fixture.py create mode 100644 zig/kernels/cuda/ggml.cu create mode 100644 zig/kernels/cuda/ggml_iq2xxs_tables.cuh create mode 100644 zig/src/cuda/ggml.zig create mode 100644 zig/tests/cuda/ggml_tests.zig diff --git a/tools/zig/bolt_ggml.py b/tools/zig/bolt_ggml.py new file mode 100644 index 0000000000..04d5d88e28 --- /dev/null +++ b/tools/zig/bolt_ggml.py @@ -0,0 +1,95 @@ +"""Reference decode of Nemotron 3.5 Bolt's ggml-packed tensors (IQ2_XXS, Q8_0). + +Bolt's ModelOpt export stores ggml blocks in safetensors as U8 [rows, blocks, bytes]: +IQ2_XXS is 66 bytes per 256 weights (fp16 d, 32 u16 qs), Q8_0 is 34 bytes per 32 +(fp16 d, 32 int8). The IQ2_XXS tables are read from the CUDA header the kernels +include, so this file and the GPU decode the same grid. + + python3 tools/zig/bolt_ggml.py MODEL_DIR TENSOR_NAME [ROWS] +""" + +import json +import re +import struct +import sys +from pathlib import Path + +import numpy as np + +HEADER = Path(__file__).resolve().parents[2] / "zig/kernels/cuda/ggml_iq2xxs_tables.cuh" +IQ2_BYTES, IQ2_ROW = 66, 256 +Q8_BYTES, Q8_ROW = 34, 32 + + +def _table(text, name): + body = text.split(name, 1)[1].split("};", 1)[0] + return [int(v, 0) for v in re.findall(r"0x[0-9a-fA-F]+|\b\d+\b", body.split("=", 1)[1])] + + +def tables(): + text = HEADER.read_text() + signs = np.array(_table(text, "tf_ksigns_iq2xs[128]"), dtype=np.uint8) + grid = np.array(_table(text, "tf_iq2xxs_grid[256]"), dtype=np.uint64) + assert signs.shape == (128,) and grid.shape == (256,) + return signs, grid.view(np.uint8).reshape(256, 8) + + +def iq2_xxs(blocks): + """[..., 66] uint8 -> [..., 256] float32, as ggml's dequantize_row_iq2_xxs.""" + signs, grid = tables() + lead = blocks.shape[:-1] + b = blocks.reshape(-1, IQ2_BYTES) + d = b[:, :2].copy().view(np.float16).astype(np.float32)[:, 0] + q = b[:, 2:].copy().view(np.uint32).reshape(-1, 8, 2) # 8 sub-blocks of 32 weights + idx = np.ascontiguousarray(q[:, :, 0]).view(np.uint8).reshape(-1, 8, 4) + aux = q[:, :, 1] + db = d[:, None] * (0.5 + (aux >> 28).astype(np.float32)) * 0.25 + out = np.empty((b.shape[0], 8, 4, 8), dtype=np.float32) + bits = (1 << np.arange(8)).astype(np.uint8) + for l in range(4): + s = signs[(aux >> (7 * l)) & 127] + neg = (s[..., None] & bits) != 0 + out[:, :, l, :] = grid[idx[:, :, l]] * np.where(neg, -1.0, 1.0) * db[..., None] + return out.reshape(*lead, IQ2_ROW) + + +def q8_0(blocks): + """[..., 34] uint8 -> [..., 32] float32.""" + lead = blocks.shape[:-1] + b = blocks.reshape(-1, Q8_BYTES) + d = b[:, :2].copy().view(np.float16).astype(np.float32) + return (b[:, 2:].view(np.int8).astype(np.float32) * d).reshape(*lead, Q8_ROW) + + +def read(model, name): + """The raw U8 array and its ggml format (from the export's quantized_layers).""" + model = Path(model) + index = json.loads((model / "model.safetensors.index.json").read_text())["weight_map"] + path = model / index[name] + with open(path, "rb") as fh: + n = struct.unpack(" 3 else None) + print(w.shape, f"mean {w.mean():.6g} std {w.std():.6g} absmax {np.abs(w).max():.6g}") diff --git a/tools/zig/bolt_ggml_fixture.py b/tools/zig/bolt_ggml_fixture.py new file mode 100644 index 0000000000..440f0c4128 --- /dev/null +++ b/tools/zig/bolt_ggml_fixture.py @@ -0,0 +1,101 @@ +"""Fixture for `tf-cuda-test ggml`: real Bolt tensors, random rows and float64 references of the decoded weights. + +The kernels get the checkpoint's raw blocks (packed on the GPU); the references come from bolt_ggml.py's decode, +which is bit-identical to libggml's, times the same bf16 rows in float64. + + python3 tools/zig/bolt_ggml_fixture.py MODEL_DIR OUT_DIR +""" + +import json +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import bolt_ggml # noqa: E402 + +PRE = "language_model.backbone.layers." +EXPERTS = [0, 1, 2, 3, 17, 64, 100, 128, 200, 238, 239] # of layer 1; the stacked layer's experts 0 .. 10 +SLOTS = 9 +DECODE_ROWS = [1, 3] +PROMPT_ROWS = [37, 300] +DENSE = [ # (case, tensor, rows of the weight kept or None) + ("shared_up", PRE + "1.mixer.shared_experts.up_proj.weight", None), + ("shared_down", PRE + "1.mixer.shared_experts.down_proj.weight", None), + ("fc1", PRE + "1.mixer.fc1_latent_proj.weight", None), + ("q7", PRE + "7.mixer.q_proj.weight", None), + ("lm_head", "language_model.lm_head.weight", 2048), +] +DENSE_ROWS = [1, 5, 16, 100] + + +def bf16(a): + """float32 -> bf16 bits (round to nearest even) and their float64 values.""" + u = np.ascontiguousarray(a, dtype=np.float32).view(np.uint32).astype(np.uint64) + bits = ((u + 0x7FFF + ((u >> 16) & 1)) >> 16).astype(np.uint16) + return bits, (bits.astype(np.uint32) << 16).view(np.float32).astype(np.float64) + + +def main(model, out): + out = Path(out) + out.mkdir(parents=True, exist_ok=True) + arrays, params = {}, {} + rng = np.random.default_rng(1234) + + def put(name, a): + a = np.ascontiguousarray(a) + (out / f"{name}.bin").write_bytes(a.tobytes()) + arrays[name] = {"file": f"{name}.bin", "dtype": str(a.dtype), "shape": list(a.shape)} + + # grouped experts: up [W, 1024] and down [1024, W] of each, raw blocks back to back + ups, downs, wu, wd = [], [], [], [] + for e in EXPERTS: + for kind, raws, ws in (("up", ups, wu), ("down", downs, wd)): + raw, algo = bolt_ggml.read(model, f"{PRE}1.mixer.experts.{e}.{kind}_proj.weight") + assert algo == "IQ2_XXS" + raws.append(raw) + ws.append(bolt_ggml.iq2_xxs(raw).reshape(raw.shape[0], -1).astype(np.float64)) + width, dims = wu[0].shape + params.update(experts=len(EXPERTS), width=width, dims=dims, slots=SLOTS) + params["decode"] = ",".join(map(str, DECODE_ROWS)) + params["prompt"] = ",".join(map(str, PROMPT_ROWS)) + put("up_raw", np.stack(ups)) + put("down_raw", np.stack(downs)) + for rows in DECODE_ROWS + PROMPT_ROWS: + picks = np.stack([rng.choice(len(EXPERTS), SLOTS, replace=False) for _ in range(rows)]).astype(np.int32) + xb, x = bf16(rng.standard_normal((rows, dims)).astype(np.float32)) + hb, h = bf16((rng.standard_normal((rows * SLOTS, width)) * 0.05).astype(np.float32)) + flat = picks.reshape(-1) + up = np.stack([wu[e] @ x[p // SLOTS] for p, e in enumerate(flat)]) + down = np.stack([wd[e] @ h[p] for p, e in enumerate(flat)]) + put(f"picks{rows}", flat) + put(f"x{rows}", xb) + put(f"h{rows}", hb) + put(f"up{rows}", up) + put(f"down{rows}", down) + + # dense projections + names = [] + for case, tensor, keep in DENSE: + raw, algo = bolt_ggml.read(model, tensor) + if keep is not None: + raw = raw[:keep] + w = bolt_ggml.weight(model, tensor, keep).astype(np.float64) + n, k = w.shape + names.append(case) + params[f"{case}_n"], params[f"{case}_k"] = n, k + params[f"{case}_kind"] = algo + put(f"{case}_raw", raw) + for rows in DENSE_ROWS: + xb, x = bf16(rng.standard_normal((rows, k)).astype(np.float32)) + put(f"{case}_x{rows}", xb) + put(f"{case}_y{rows}", x @ w.T) + params["dense"] = ",".join(names) + params["dense_rows"] = ",".join(map(str, DENSE_ROWS)) + (out / "manifest.json").write_text(json.dumps({"params": params, "arrays": arrays}, indent=1)) + print(f"wrote {len(arrays)} arrays to {out}") + + +if __name__ == "__main__": + main(sys.argv[1], sys.argv[2]) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4b4fdbb2dc..dce26ba712 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -43,6 +43,7 @@ const kernels = [_]Kernel{ .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts + .{ .name = "ggml", .flags = &.{"-O3"} }, // ours: Bolt's IQ2_XXS and Q8_0 blocks, grouped experts and dense rows .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, diff --git a/zig/kernels/cuda/ggml.cu b/zig/kernels/cuda/ggml.cu new file mode 100644 index 0000000000..eb401dfed9 --- /dev/null +++ b/zig/kernels/cuda/ggml.cu @@ -0,0 +1,470 @@ +// Bolt's ggml blocks (IQ2_XXS, Q8_0) on mma.sync: grouped experts and dense projections, decode rows and prompt rows. +// +// The weight is mma's A operand (16 output columns a warp tile), the rows are B (8 rows an n8 tile), so a lane's +// output columns are the ones it decoded and each 32-input group's scale stays in the lane. The "k" order is ours, +// the same in both operands: in a 32-input group, lane t takes inputs 8t .. 8t + 3 in step 0 and 8t + 4 .. 8t + 7 in +// step 1, so an IQ2_XXS lane decodes one whole grid entry a column. Codes become exact bf16 integers; each group +// sums into fp32 by mma, then takes its column's fp32 scale. +// +// Weights are repacked at load (tf_ggml_pack), a pure byte permutation of the checkpoint's rows, same size: +// IQ2_XXS [n/16][k/256][1056]: 8 groups of [gq 8][col gq, gq + 8][8 bytes qs], then [gq 8][2] fp16 d +// Q8_0 [n/16][k/32][544]: [lane 32][(step 0, col gq), (0, gq + 8), (1, gq), (1, gq + 8)][4 int8], then d as above +// (lane (gq, t) step s holds inputs 8t + 4s .. + 3) +// A matrix's 16-column tiles follow each other; a grouped layer's experts follow each other. + +#include +#include +#include + +#include "experts.cuh" +#include "ggml_iq2xxs_tables.cuh" + +namespace tf_ggml { + +__device__ __forceinline__ uint4 ld_nc(const uint8_t* p) { + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +__device__ __forceinline__ uint32_t ld_nc32(const uint8_t* p) { + uint32_t r; + asm volatile("ld.global.nc.L1::no_allocate.u32 %0, [%1];\n" : "=r"(r) : "l"(p)); + return r; +} + +__device__ __forceinline__ float half_f(uint32_t bits) { return __half2float(__ushort_as_half((unsigned short)bits)); } + +// The sign bits of a bf16 pair from bits 0 and 1 of sb (< 2^16): bit 0 lands on 15, bit 1 on 31, no other bit there. +__device__ __forceinline__ uint32_t sign_mask(uint32_t sb) { return (sb * 0x40008000u) & 0x80008000u; } + +// Two grid bytes of g (selector sel) as exact bf16 integers, signs from bits 0 and 1 of sb. +__device__ __forceinline__ uint32_t grid_pair(uint32_t g, uint32_t sel, uint32_t sb) { + uint32_t v = __byte_perm(g, 0x43u, sel); + const uint32_t k = K128; + __nv_bfloat162 r = __hsub2(*reinterpret_cast<__nv_bfloat162*>(&v), *reinterpret_cast(&k)); + return *reinterpret_cast(&r) ^ sign_mask(sb); +} + +// IQ2_XXS: a group is a 32-weight sub-block (4 grid indices, 4 x 7 sign bits, a 4-bit scale); 8 groups a block. +struct Iq2 { + static constexpr int G = 8, BYTES = 1056, Q = 4; // Q: groups a decode load brings, a lane's 16 bytes each + static constexpr bool TABLE = true; + // lane (gq, t) loads group 4q + t's bytes of columns gq, gq + 8: a warp's load is 512 distinct bytes + static __device__ __forceinline__ const uint8_t* uptr(const uint8_t* tile, int q, int lane) { + return tile + (q >> 1) * BYTES + (q & 1) * 512 + (lane & 3) * 128 + (lane >> 2) * 16; + } + static __device__ __forceinline__ const uint8_t* wptr(const uint8_t* tile, int g, int lane) { + return tile + (g >> 3) * BYTES + (g & 7) * 128 + (lane >> 2) * 16; + } + static __device__ __forceinline__ const uint8_t* dptr(const uint8_t* tile, int g, int lane) { + return tile + (g >> 3) * BYTES + 1024 + (lane >> 2) * 4; + } + static __device__ __forceinline__ void decode(const uint4& w, uint32_t dp, int t, const uint2* tab, + uint32_t (&a)[2][4], float (&sc)[2]) { + sc[0] = half_f(dp & 0xFFFFu) * (0.5f + (float)(w.y >> 28)) * 0.25f; + sc[1] = half_f(dp >> 16) * (0.5f + (float)(w.w >> 28)) * 0.25f; +#pragma unroll + for (int c = 0; c < 2; ++c) { // lane t decodes grid entry t of the group: 8 weights, both steps + const uint32_t aux0 = c ? w.z : w.x, aux1 = c ? w.w : w.y; + const uint2 g = tab[(aux0 >> (8 * t)) & 255u]; + const uint32_t si = (aux1 >> (7 * t)) & 127u; + const uint32_t sb = si | ((__popc(si) & 1u) << 7); // ksigns_iq2xs: even parity in bit 7 +#pragma unroll + for (int s = 0; s < 2; ++s) { + const uint32_t gs = s ? g.y : g.x, bs = sb >> (4 * s); + a[s][c] = grid_pair(gs, 0x4140u, bs); + a[s][2 + c] = grid_pair(gs, 0x4342u, bs >> 2); + } + } + } +}; + +// Two int8 of u (bytes 2h, 2h + 1) as exact bf16. +__device__ __forceinline__ uint32_t i8_pair(uint32_t u, int h) { + const uint32_t v = u ^ 0x80808080u; + const float f0 = __uint_as_float(__byte_perm(v, 0x4B000000u, 0x7440u | (2 * h))) - 8388736.f; + const float f1 = __uint_as_float(__byte_perm(v, 0x4B000000u, 0x7440u | (2 * h + 1))) - 8388736.f; + __nv_bfloat162 r = __floats2bfloat162_rn(f0, f1); + return *reinterpret_cast(&r); +} + +// Q8_0: a group is one block (fp16 d, 32 int8). +struct Q8 { + static constexpr int G = 1, BYTES = 544, Q = 1; + static constexpr bool TABLE = false; + static __device__ __forceinline__ const uint8_t* uptr(const uint8_t* tile, int q, int lane) { + return tile + q * BYTES + lane * 16; + } + static __device__ __forceinline__ const uint8_t* wptr(const uint8_t* tile, int g, int lane) { + return tile + g * BYTES + lane * 16; + } + static __device__ __forceinline__ const uint8_t* dptr(const uint8_t* tile, int g, int lane) { + return tile + g * BYTES + 512 + (lane >> 2) * 4; + } + static __device__ __forceinline__ void decode(const uint4& w, uint32_t dp, int, const uint2*, + uint32_t (&a)[2][4], float (&sc)[2]) { + sc[0] = half_f(dp & 0xFFFFu); + sc[1] = half_f(dp >> 16); +#pragma unroll + for (int s = 0; s < 2; ++s) +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint32_t u = comp(w, 2 * s + c); + a[s][c] = i8_pair(u, 0); + a[s][2 + c] = i8_pair(u, 1); + } + } +}; + +// IQ2_XXS's grid in shared memory, 256 entries of 8 bytes: indices diverge across a warp. Signs and the bf16 +// conversion are arithmetic; tables of them cost more in bank conflicts than they save. +constexpr int TABLE_WORDS = 256; + +template +__device__ __forceinline__ void load_tables(uint2* tab) { + if constexpr (F::TABLE) { + for (int i = threadIdx.x; i < TABLE_WORDS; i += blockDim.x) { + const uint64_t g = tf_iq2xxs_grid[i]; + tab[i] = make_uint2((uint32_t)g, (uint32_t)(g >> 32)); + } + __syncthreads(); + } +} + +// One item of rows: grouped (expert, first member, count) from the plan, or a dense run of rows (members null). +struct Item { + int e, first, cnt; +}; + +__device__ __forceinline__ int pair_of(const int* __restrict__ members, const Item& it, int r) { + return members ? __ldg(members + it.first + r) : it.first + r; +} + +template +__device__ __forceinline__ void store(void* out, size_t at, float v) { + if constexpr (EPI == 0) { + reinterpret_cast(out)[at] = v; + } else { + reinterpret_cast<__nv_bfloat16*>(out)[at] = __float2bfloat16_rn(EPI == 1 ? relu2(v) : v); + } +} + +// One 32-input group into acc: NT n8 row tiles, B fragments by `ldb(nt, s)` (uint2: b0, b1). +template +__device__ __forceinline__ void group(float (&acc)[NT][4], const uint4& w, uint32_t dp, int t, const uint2* tab, + LDB ldb) { + uint32_t a[2][4]; + float sc[2]; + F::decode(w, dp, t, tab, a, sc); + float p[NT][4]; +#pragma unroll + for (int nt = 0; nt < NT; ++nt) p[nt][0] = p[nt][1] = p[nt][2] = p[nt][3] = 0.f; +#pragma unroll + for (int s = 0; s < 2; ++s) +#pragma unroll + for (int nt = 0; nt < NT; ++nt) { + const uint2 b = ldb(nt, s); + mma(p[nt], a[s][0], a[s][1], a[s][2], a[s][3], b.x, b.y); + } +#pragma unroll + for (int nt = 0; nt < NT; ++nt) { + acc[nt][0] = fmaf(p[nt][0], sc[0], acc[nt][0]); + acc[nt][1] = fmaf(p[nt][1], sc[0], acc[nt][1]); + acc[nt][2] = fmaf(p[nt][2], sc[1], acc[nt][2]); + acc[nt][3] = fmaf(p[nt][3], sc[1], acc[nt][3]); + } +} + +// acc[nt] holds columns col0 + gq (0, 1) and col0 + gq + 8 (2, 3) of rows nt * 8 + 2t (0, 2) and + 1 (1, 3). +template +__device__ __forceinline__ void epilogue_rows(const float (&acc)[NT][4], void* out, int N, int col0, int lane, + const int* __restrict__ members, const Item& it) { + const int gq = lane >> 2, t = lane & 3; +#pragma unroll + for (int nt = 0; nt < NT; ++nt) +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int r = nt * 8 + 2 * t + j; + if (r >= it.cnt) continue; + const size_t row = (size_t)pair_of(members, it, r) * N; + store(out, row + col0 + gq, acc[nt][j]); + store(out, row + col0 + gq + 8, acc[nt][2 + j]); + } +} + +constexpr int U = 2; // loads a warp has in flight, twice over + +// Decode rows: a CTA per TPC (item, 16-column tile) units, KW warps each over K's loads (Q groups a load), summed in +// a fixed order. +template +__device__ __forceinline__ void gemv_body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint8_t* __restrict__ W, int KG, int NT16, + const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, int rows, + int skip) { + __shared__ uint2 tab[F::TABLE ? TABLE_WORDS : 1]; + __shared__ float red[TPC][KW > 1 ? KW - 1 : 1][8][32]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, gq = lane >> 2, t = lane & 3; + const int sub = warp / KW, kw = warp - sub * KW; + const int units = (items ? __ldg(counts) : 1) * NT16; + if (blockIdx.x * TPC >= units) return; + const int unit = blockIdx.x * TPC + sub, iu = unit / NT16, tile = unit - iu * NT16; + Item it = {0, 0, 0}; + if (unit < units) it = items ? Item{__ldg(items + 3 * iu), __ldg(items + 3 * iu + 1), __ldg(items + 3 * iu + 2)} + : Item{0, 0, rows}; + const bool live = unit < units && it.e != skip; + const uint8_t* wt = W + ((size_t)it.e * NT16 + tile) * (size_t)(KG / F::G) * F::BYTES; + const int KQ = KG / F::Q; + uint4 w[U]; + uint32_t dp[U]; + auto fetch = [&](int q0) { +#pragma unroll + for (int u = 0; u < U; ++u) + if (q0 + u * KW < KQ) { + w[u] = ld_nc(F::uptr(wt, q0 + u * KW, lane)); + dp[u] = ld_nc32(F::dptr(wt, (q0 + u * KW) * F::Q, lane)); + } + }; + if (live) fetch(kw); // the weights are on their way while the tables load + load_tables(tab); + const bool v0 = gq < it.cnt, v1 = gq + 8 < it.cnt, two = it.cnt > 8; + const int p0 = v0 ? pair_of(members, it, gq) : 0, p1 = v1 ? pair_of(members, it, gq + 8) : 0; + const __nv_bfloat16* x0 = X + (size_t)(slots ? p0 / slots : p0) * x_stride + 8 * t; + const __nv_bfloat16* x1 = X + (size_t)(slots ? p1 / slots : p1) * x_stride + 8 * t; + float acc[2][4]; +#pragma unroll + for (int nt = 0; nt < 2; ++nt) acc[nt][0] = acc[nt][1] = acc[nt][2] = acc[nt][3] = 0.f; + // a load's rows come in at once, Q groups' b fragments (steps 0 and 1) for row tiles 0 and 1 + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + auto xrows = [&](const __nv_bfloat16* x, bool v, int g) { + return v ? __ldg(reinterpret_cast(x + g * 32)) : zero; + }; + for (int q = live ? kw : KQ; q < KQ; q += U * KW) { + uint4 cw[U]; + uint32_t cd[U]; +#pragma unroll + for (int u = 0; u < U; ++u) { + cw[u] = w[u]; + cd[u] = dp[u]; + } + fetch(q + U * KW); +#pragma unroll + for (int u = 0; u < U; ++u) { + const int qu = q + u * KW; + if (qu >= KQ) break; + uint4 xr[F::Q][2]; +#pragma unroll + for (int j = 0; j < F::Q; ++j) { + xr[j][0] = xrows(x0, v0, qu * F::Q + j); + xr[j][1] = two ? xrows(x1, v1, qu * F::Q + j) : zero; + } +#pragma unroll + for (int j = 0; j < F::Q; ++j) { + uint4 wv = cw[u]; + if constexpr (F::Q > 1) { // group 4qu + j of columns gq, gq + 8 sits in lane (gq, j) + const int src = (lane & ~3) | j; + wv = make_uint4(__shfl_sync(~0u, wv.x, src), __shfl_sync(~0u, wv.y, src), __shfl_sync(~0u, wv.z, src), + __shfl_sync(~0u, wv.w, src)); + } + auto ldb = [&](int nt, int s) { + const uint4& v = xr[j][nt]; + return s ? make_uint2(v.z, v.w) : make_uint2(v.x, v.y); + }; + if (two) { + group(acc, wv, cd[u], t, tab, ldb); + } else { + float(&a1)[1][4] = *reinterpret_cast(&acc[0]); + group(a1, wv, cd[u], t, tab, ldb); + } + } + } + } + if constexpr (KW > 1) { + if (kw > 0) +#pragma unroll + for (int q = 0; q < 8; ++q) red[sub][kw - 1][q][lane] = acc[q >> 2][q & 3]; + __syncthreads(); + if (kw > 0) return; + for (int wi = 0; wi < KW - 1; ++wi) +#pragma unroll + for (int q = 0; q < 8; ++q) acc[q >> 2][q & 3] += red[sub][wi][q][lane]; + } + if (!live) return; + epilogue_rows(acc, out, N, tile * 16, lane, members, it); +} + +// Prompt rows: a CTA per (item of up to 64 rows, WN 16-column tiles), rows staged 128 inputs a stage. +constexpr int BM = 64, SK = 128, XC = SK / 8, STAGES = 3; // XC: 16-byte chunks a staged row; 48 KiB of stages + +__device__ __forceinline__ int chunk(int r, int c) { return r * XC + (c ^ (r & 7)); } + +template +__device__ __forceinline__ void prompt_tiles(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint8_t* __restrict__ W, int KG, int NT16, + const int* __restrict__ members, void* __restrict__ out, int N, + const Item& it, int ct, const uint2* tab) { + constexpr int THREADS = WN * 32, XPT = BM * XC / THREADS; + extern __shared__ uint4 sm[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, gq = lane >> 2, t = lane & 3; + const int tile = ct * WN + warp; + const bool active = tile < NT16; + const __nv_bfloat16* xsrc[XPT]; + int xdst[XPT]; +#pragma unroll + for (int i = 0; i < XPT; ++i) { + const int q = tid + i * THREADS, r = q / XC, c = q - r * XC; + xdst[i] = chunk(r, c); + xsrc[i] = nullptr; + if (r < it.cnt) { + const int p = pair_of(members, it, r); + xsrc[i] = X + (size_t)(slots ? p / slots : p) * x_stride + 8 * c; + } + } + auto stage = [&](int s, int k) { + uint4* xs = sm + s * BM * XC; +#pragma unroll + for (int i = 0; i < XPT; ++i) + if (xsrc[i]) cp16(xs + xdst[i], xsrc[i] + (size_t)k * SK); + }; + const uint8_t* wt = W + ((size_t)it.e * NT16 + (active ? tile : 0)) * (size_t)(KG / F::G) * F::BYTES; + constexpr int GS = SK / 32; // groups a stage + float acc[NT][4]; +#pragma unroll + for (int nt = 0; nt < NT; ++nt) acc[nt][0] = acc[nt][1] = acc[nt][2] = acc[nt][3] = 0.f; + const int steps = KG / GS; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < steps) stage(s, s); + cp_commit(); + } + uint4 w[GS]; + uint32_t dp[GS]; + if (active) +#pragma unroll + for (int gi = 0; gi < GS; ++gi) { + w[gi] = ld_nc(F::wptr(wt, gi, lane)); + dp[gi] = ld_nc32(F::dptr(wt, gi, lane)); + } + for (int k = 0; k < steps; ++k) { + cp_wait(); + __syncthreads(); + if (k + STAGES - 1 < steps) stage((k + STAGES - 1) % STAGES, k + STAGES - 1); + cp_commit(); + if (!active) continue; + uint4 wn[GS]; + uint32_t dn[GS]; + if (k + 1 < steps) +#pragma unroll + for (int gi = 0; gi < GS; ++gi) { + wn[gi] = ld_nc(F::wptr(wt, (k + 1) * GS + gi, lane)); + dn[gi] = ld_nc32(F::dptr(wt, (k + 1) * GS + gi, lane)); + } + const uint4* xs = sm + (k % STAGES) * BM * XC; +#pragma unroll + for (int gi = 0; gi < GS; ++gi) + group(acc, w[gi], dp[gi], t, tab, [&](int nt, int s) { + const int r = nt * 8 + gq; + return reinterpret_cast(xs + chunk(r, gi * 4 + t))[s]; + }); +#pragma unroll + for (int gi = 0; gi < GS; ++gi) { + w[gi] = wn[gi]; + dp[gi] = dn[gi]; + } + } + cp_wait<0>(); + if (active) epilogue_rows(acc, out, N, tile * 16, lane, members, it); +} + +template +__device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint8_t* __restrict__ W, int KG, int NT16, + const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, int rows, + int skip) { + __shared__ uint2 tab[F::TABLE ? TABLE_WORDS : 1]; + const int nct = (NT16 + WN - 1) / WN, iu = blockIdx.x / nct, ct = blockIdx.x - iu * nct; + Item it; + if (items) { + if (iu >= __ldg(counts)) return; + it = {__ldg(items + 3 * iu), __ldg(items + 3 * iu + 1), __ldg(items + 3 * iu + 2)}; + if (it.e == skip) return; + } else { + it = {0, iu * BM, min(BM, rows - iu * BM)}; + } + load_tables(tab); + const int nt = min(8, (it.cnt + 7) >> 3); +#define TF_TILES(NT) \ + prompt_tiles(X, x_stride, slots, W, KG, NT16, members, out, N, it, ct, tab) + switch (nt) { + case 8: TF_TILES(8); break; + case 7: TF_TILES(7); break; + case 6: TF_TILES(6); break; + case 5: TF_TILES(5); break; + case 4: TF_TILES(4); break; + case 3: TF_TILES(3); break; + case 2: TF_TILES(2); break; + default: TF_TILES(1); + } +#undef TF_TILES +} + +} // namespace tf_ggml + +// The entries: format x (decode with 4 warps over K a tile, decode with a warp a tile and 4 tiles a CTA, prompt) x +// epilogue (0 fp32, 1 relu^2 bf16, 3 bf16); items null for dense rows. +#define TF_GGML_GEMV(NAME, F, EPI, KW, TPC) \ + extern "C" __global__ void __launch_bounds__(KW * TPC * 32) NAME( \ + const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint8_t* __restrict__ W, int KG, int NT16, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int rows, int skip) { \ + tf_ggml::gemv_body(X, x_stride, slots, W, KG, NT16, items, counts, members, out, N, \ + rows, skip); \ + } +#define TF_GGML_ENTRIES(FMT, F, EPI, SUFFIX) \ + TF_GGML_GEMV(tf_ggml_##FMT##_gemv_##SUFFIX, F, EPI, 4, 1) \ + TF_GGML_GEMV(tf_ggml_##FMT##_gemv4_##SUFFIX, F, EPI, 1, 4) \ + extern "C" __global__ void __launch_bounds__(256) tf_ggml_##FMT##_prompt_##SUFFIX( \ + const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint8_t* __restrict__ W, int KG, int NT16, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int rows, int skip) { \ + tf_ggml::prompt_body(X, x_stride, slots, W, KG, NT16, items, counts, members, out, N, rows, \ + skip); \ + } +TF_GGML_ENTRIES(iq2, Iq2, 1, relu2) +TF_GGML_ENTRIES(iq2, Iq2, 0, f32) +TF_GGML_ENTRIES(iq2, Iq2, 3, bf16) +TF_GGML_ENTRIES(q8, Q8, 1, relu2) +TF_GGML_ENTRIES(q8, Q8, 0, f32) +TF_GGML_ENTRIES(q8, Q8, 3, bf16) +#undef TF_GGML_ENTRIES +#undef TF_GGML_GEMV + +// The load repack: `count` matrices [n][k / KB][B] of checkpoint blocks into the 16-column tiles above, a byte a thread. +extern "C" __global__ void tf_ggml_pack(const uint8_t* __restrict__ src, uint8_t* __restrict__ dst, int iq2, int n, + int kb, long long total) { + const int B = iq2 ? 66 : 34, TB = iq2 ? 1056 : 544; + const long long per = (long long)n * kb * B; + for (long long o = blockIdx.x * (long long)blockDim.x + threadIdx.x; o < total; o += (long long)gridDim.x * blockDim.x) { + const long long m = o / per, rem = o - m * per; + const long long tb = rem / TB; + const int r = (int)(rem - tb * TB), tile = (int)(tb / kb), b = (int)(tb - (long long)tile * kb); + int col, at; + if (r >= TB - 32) { + const int q = r - (TB - 32); + col = tile * 16 + (q >> 2) + 8 * ((q >> 1) & 1); + at = q & 1; + } else if (iq2) { + col = tile * 16 + ((r >> 4) & 7) + 8 * ((r >> 3) & 1); + at = 2 + (r >> 7) * 8 + (r & 7); + } else { + const int ln = r >> 4, q = (r >> 2) & 3; + col = tile * 16 + (ln >> 2) + 8 * (q & 1); + at = 2 + 8 * (ln & 3) + 4 * (q >> 1) + (r & 3); + } + dst[o] = src[m * per + ((long long)col * kb + b) * B + at]; + } +} diff --git a/zig/kernels/cuda/ggml_iq2xxs_tables.cuh b/zig/kernels/cuda/ggml_iq2xxs_tables.cuh new file mode 100644 index 0000000000..687b457d33 --- /dev/null +++ b/zig/kernels/cuda/ggml_iq2xxs_tables.cuh @@ -0,0 +1,84 @@ +// IQ2_XXS lookup tables, copied from ggml (llama.cpp ggml/src/ggml-common.h), MIT License, +// Copyright (c) 2023-2026 The ggml authors. Read by the Bolt expert kernels and by +// tools/zig/bolt_ggml.py, so the reference and the GPU decode the same bytes. +// The grid sits in global memory: kernels copy it to shared memory with coalesced loads, not constant-cache ones. +#pragma once +#include + +static __device__ __constant__ const uint8_t tf_ksigns_iq2xs[128] = { + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, +}; + +static __device__ const uint64_t tf_iq2xxs_grid[256] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, + 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, + 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, + 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, + 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, + 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, + 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, + 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, + 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, + 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, + 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, + 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, + 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, + 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, + 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, + 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, + 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, + 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, + 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, + 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, + 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, + 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, + 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, + 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, + 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, + 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, + 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, + 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, + 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, + 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, + 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, + 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, + 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, + 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, + 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, + 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, + 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, + 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, + 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, + 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, + 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, + 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, + 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, + 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, +}; diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..8bafabc226 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -6,6 +6,7 @@ const stream_ = @import("stream.zig"); const kernels = @import("kernels.zig"); const grouped = @import("grouped.zig"); const driver = @import("driver.zig"); +const ggml = @import("ggml.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -13,20 +14,20 @@ const Function = module.Function; const Stream = stream_.Stream; const Plan = grouped.Plan; -pub const Format = enum { affine4, fp8g, nvfp4 }; +pub const Format = enum { affine4, fp8g, nvfp4, iq2xxs }; /// One layer's experts on the device as each format's make() lays them out; a family stacks its shared ones last. pub const Layer = struct { format: Format, - up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144] - down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144] + up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig) + down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; iq2xxs [E][D/16][NI/256][1056] up_scale: u64 = 0, // fp8g [E][2][NI/128][D/128] fp32; nvfp4 [E][2] global scales; affine4's are in the blocks down_scale: u64 = 0, // fp8g [E][1][D/128][NI/128]; nvfp4 [E][1] width: usize, // NI dims: usize, // D experts: usize, limit: f32 = 0.0, // fp8g, nvfp4: the SwiGLU clamp, 0 for none - relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up + relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up; iq2xxs: always }; /// A call's epilogue: up with its format's activation (affine4 relu^2, else SwiGLU), down to fp32 or to bf16 sums. @@ -40,6 +41,7 @@ pub const Grouped = struct { affine4: ?*const Affine4Experts = null, fp8g: ?*const Fp8Experts = null, nvfp4: ?*const Nvfp4Experts = null, + iq2xxs: ?*const Iq2xxsExperts = null, /// Decode rows: `items` bounds the plan's items (grouped.maxItems at tile 16); `skip` an expert left out, or -1. pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { @@ -55,6 +57,7 @@ pub const Grouped = struct { }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, false, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), + .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, false, epi, in, l, p, out, items, skip), } } @@ -72,6 +75,7 @@ pub const Grouped = struct { }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, true, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), + .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, true, epi, in, l, p, out, items, skip), } } }; @@ -79,10 +83,26 @@ pub const Grouped = struct { /// Pairs a prompt item holds, per format (the Python engines' prompt tiles). pub fn promptTile(f: Format) usize { return switch (f) { - .affine4, .fp8g, .nvfp4 => 64, + .affine4, .fp8g, .nvfp4, .iq2xxs => 64, }; } +/// IQ2_XXS experts (ggml.cu): relu^2 up, down to fp32 or bf16; experts packed by Iq2xxsExperts.pack from the raw blocks. +pub const Iq2xxsExperts = ggml.Kernels; + +fn iq2Run(k: *const Iq2xxsExperts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + if (l.limit != 0) return error.OptionNotBuilt; + const up = epi == .up; + const o: ggml.Out = switch (epi) { + .up => .relu2, + .down_f32 => .f32, + .down_bf16 => .bf16, + }; + const n = if (up) l.width else l.dims; + const kk = if (up) l.dims else l.width; + return k.experts(s, prompt, .iq2xxs, o, in.x, in.stride, if (up) in.slots else 0, if (up) l.up else l.down, n, kk, p, out, items, skip); +} + fn int(x: usize) c_int { return @intCast(x); } diff --git a/zig/src/cuda/ggml.zig b/zig/src/cuda/ggml.zig new file mode 100644 index 0000000000..bd7be5c7e7 --- /dev/null +++ b/zig/src/cuda/ggml.zig @@ -0,0 +1,218 @@ +//! Bolt's ggml blocks (ggml.cu): IQ2_XXS and Q8_0 repacked into 16-column tiles, for grouped experts and dense rows. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const grouped = @import("grouped.zig"); + +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// The checkpoint's block formats: IQ2_XXS (66 bytes a 256 inputs), Q8_0 (34 bytes a 32). +pub const Kind = enum { iq2xxs, q8_0 }; + +/// Inputs and bytes of one checkpoint block. +pub fn block(kind: Kind) struct { k: usize, bytes: usize } { + return switch (kind) { + .iq2xxs => .{ .k = 256, .bytes = 66 }, + .q8_0 => .{ .k = 32, .bytes = 34 }, + }; +} + +/// Bytes of 16 columns' blocks side by side (a packed tile's step along K). +pub fn tileBytes(kind: Kind) usize { + return 16 * block(kind).bytes; +} + +/// Bytes of an (n, k) matrix, the checkpoint's and the packed layout's alike. +pub fn bytes(kind: Kind, n: usize, k: usize) usize { + const b = block(kind); + return n * (k / b.k) * b.bytes; +} + +/// One packed projection on the device: n outputs from k inputs. +pub const Weight = struct { kind: Kind, w: u64, n: usize, k: usize }; + +pub fn q8Weight(w: u64, n: usize, k: usize) Weight { + return .{ .kind = .q8_0, .w = w, .n = n, .k = k }; +} + +pub fn iq2Weight(w: u64, n: usize, k: usize) Weight { + return .{ .kind = .iq2xxs, .w = w, .n = n, .k = k }; +} + +/// What a call writes: fp32, bf16 of relu(x)^2 (Nemotron's expert activation), or bf16. +pub const Out = enum { f32, relu2, bf16 }; + +pub const prompt_rows = 64; // rows a prompt item holds (the plan's prompt tile) +pub const decode_rows = 16; // rows a decode item holds at most (the plan's decode tile) +const prompt_warps = 8; // 16-column tiles a prompt CTA +/// Decode CTAs: 4 warps split a tile's K, or (short K: fewer than `split_loads` loads a tile) a warp a tile, 4 tiles. +const split_loads = 16; +const prompt_smem: u32 = 3 * 64 * 128 * 2; // three stages of 64 rows' 128 inputs + +/// Where byte `o` of a packed matrix comes from in the checkpoint's [n][k / KB][B] (tf_ggml_pack's map). +pub fn source(kind: Kind, n: usize, k: usize, o: usize) usize { + _ = n; + const b = block(kind); + const kb = k / b.k; + const tb_bytes = tileBytes(kind); + const tb = o / tb_bytes; + const r = o % tb_bytes; + const tile = tb / kb; + const blk = tb % kb; + var col: usize = undefined; + var at: usize = undefined; + if (r >= tb_bytes - 32) { + const q = r - (tb_bytes - 32); + col = tile * 16 + (q >> 2) + 8 * ((q >> 1) & 1); + at = q & 1; + } else if (kind == .iq2xxs) { + col = tile * 16 + ((r >> 4) & 7) + 8 * ((r >> 3) & 1); + at = 2 + (r >> 7) * 8 + (r & 7); + } else { + const ln = r >> 4; + const q = (r >> 2) & 3; + col = tile * 16 + (ln >> 2) + 8 * (q & 1); + at = 2 + 8 * (ln & 3) + 4 * (q >> 1) + (r & 3); + } + return (col * kb + blk) * b.bytes + at; +} + +/// The host repack of one (n, k) matrix, as tf_ggml_pack does it on the device. +pub fn packHost(kind: Kind, out: []u8, raw: []const u8, n: usize, k: usize) void { + std.debug.assert(out.len == bytes(kind, n, k) and raw.len == out.len and n % 16 == 0); + for (out, 0..) |*o, i| o.* = raw[source(kind, n, k, i)]; +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +/// The ggml module's kernels: decode (GEMV) and prompt entries per format and output, and the repack. +pub const Kernels = struct { + mod: Module, + gemv: [2][2][3]Function, // by Kind, K split or 4 tiles a CTA, then Out + prompt: [2][3]Function, + pack_fn: Function, + + /// Entry names: `form` is "gemv" (K split), "gemv4" (4 tiles a CTA) or "prompt". + pub fn symbol(comptime kind: Kind, comptime form: []const u8, comptime out: Out) [:0]const u8 { + return "tf_ggml_" ++ (if (kind == .iq2xxs) "iq2" else "q8") ++ "_" ++ form ++ "_" ++ @tagName(out); + } + + pub fn load(d: *const Driver) !Kernels { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.ggml); + errdefer mod.unload(); + var g: Kernels = .{ .mod = mod, .gemv = undefined, .prompt = undefined, .pack_fn = try mod.function("tf_ggml_pack") }; + inline for (.{ Kind.iq2xxs, Kind.q8_0 }, 0..) |kind, ki| inline for (.{ Out.f32, Out.relu2, Out.bf16 }, 0..) |out, oi| { + g.gemv[ki][0][oi] = try mod.function(symbol(kind, "gemv", out)); + g.gemv[ki][1][oi] = try mod.function(symbol(kind, "gemv4", out)); + g.prompt[ki][oi] = try mod.function(symbol(kind, "prompt", out)); + try g.prompt[ki][oi].allowDynamicShared(prompt_smem); + }; + return g; + } + + pub fn unload(g: *Kernels) void { + g.mod.unload(); + } + + /// `count` checkpoint matrices [n][k / KB][B] back to back at `raw` into packed tiles at `out` (as many bytes). + pub fn pack(g: *const Kernels, s: Stream, kind: Kind, raw: u64, out: u64, n: usize, k: usize, count: usize) !void { + if (n % 16 != 0 or k % block(kind).k != 0) return error.ShapeNotTiled; + const total = count * bytes(kind, n, k); + if (total == 0) return; + var a: launch_.Args = .{}; + a.add(raw); + a.add(out); + a.add(@as(c_int, @intFromBool(kind == .iq2xxs))); + a.add(int(n)); + a.add(int(k / block(kind).k)); + a.add(@as(i64, @intCast(total))); + try launch_.launch(g.pack_fn, .{ .grid = .{ .x = @intCast(@min((total + 255) / 256, 1 << 20)), .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + + fn check(kind: Kind, n: usize, k: usize, prompt: bool) !void { + if (n % 16 != 0 or k % block(kind).k != 0 or (prompt and k % 128 != 0)) return error.ShapeNotTiled; + } + + /// The common launch: `items`, `counts` and `members` 0 for dense rows (`rows` of them), else a plan's. + fn go(g: *const Kernels, s: Stream, prompt: bool, out_kind: Out, kind: Kind, x: u64, stride: usize, slots: usize, w: u64, n: usize, k: usize, p: ?grouped.Plan, out: u64, units: usize, rows: usize, skip: i32) !void { + try check(kind, n, k, prompt); + const ki = @intFromEnum(kind); + const oi = @intFromEnum(out_kind); + const nt16 = n / 16; + if (units == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(int(stride)); + a.add(int(slots)); + a.add(w); + a.add(int(k / 32)); + a.add(int(nt16)); + for ([_]u64{ if (p) |q| q.items else 0, if (p) |q| q.counts else 0, if (p) |q| q.members else 0, out }) |v| a.add(v); + a.add(int(n)); + a.add(int(rows)); + a.add(skip); + if (prompt) { + const grid = units * ((nt16 + prompt_warps - 1) / prompt_warps); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 32 * prompt_warps }, .shared = prompt_smem }; + return launch_.launch(g.prompt[ki][oi], cfg, s, &a); + } + const loads = k / 32 / @as(usize, if (kind == .iq2xxs) 4 else 1); + const tiles: usize = if (loads < split_loads) 4 else 1; + const grid = (units * nt16 + tiles - 1) / tiles; // items past the plan's count return at once + try launch_.launch(g.gemv[ki][@intFromBool(tiles > 1)][oi], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } + + /// Dense decode rows (at most 16): out (rows, n) = x (rows, k; rows `ldx` apart, 0 packed) times w. + pub fn decode(g: *const Kernels, s: Stream, w: Weight, x: u64, ldx: usize, rows: usize, out: u64, out_kind: Out) !void { + if (rows > decode_rows) return error.TooManyRows; + return g.go(s, false, out_kind, w.kind, x, if (ldx == 0) w.k else ldx, 0, w.w, w.n, w.k, null, out, @intFromBool(rows > 0), rows, -1); + } + + /// Dense prompt rows (any count): 64 rows by 128 columns a CTA. + pub fn promptRows(g: *const Kernels, s: Stream, w: Weight, x: u64, ldx: usize, rows: usize, out: u64, out_kind: Out) !void { + return g.go(s, true, out_kind, w.kind, x, if (ldx == 0) w.k else ldx, 0, w.w, w.n, w.k, null, out, (rows + prompt_rows - 1) / prompt_rows, rows, -1); + } + + /// Either form by row count: decode at 16 rows or fewer, else prompt. + pub fn matmul(g: *const Kernels, s: Stream, w: Weight, x: u64, ldx: usize, rows: usize, out: u64, out_kind: Out) !void { + return if (rows <= decode_rows) g.decode(s, w, x, ldx, rows, out, out_kind) else g.promptRows(s, w, x, ldx, rows, out, out_kind); + } + + /// Grouped experts on a plan of at most `items` items (16 pairs decode, 64 prompt): stacked [E] packed matrices + /// (n, k) at `w`; x rows `stride` apart, token rows when `slots` > 0 (pair / slots), else pair rows; out pair rows. + pub fn experts(g: *const Kernels, s: Stream, prompt: bool, kind: Kind, out_kind: Out, x: u64, stride: usize, slots: usize, w: u64, n: usize, k: usize, p: grouped.Plan, out: u64, items: usize, skip: i32) !void { + return g.go(s, prompt, out_kind, kind, x, stride, slots, w, n, k, p, out, items, 0, skip); + } +}; + +test "the repack is a permutation of each matrix's bytes" { + const gpa = std.testing.allocator; + inline for (.{ Kind.iq2xxs, Kind.q8_0 }) |kind| { + const n = 32; + const k = 512; + const len = bytes(kind, n, k); + const seen = try gpa.alloc(bool, len); + defer gpa.free(seen); + @memset(seen, false); + for (0..len) |o| { + const src = source(kind, n, k, o); + try std.testing.expect(!seen[src]); + seen[src] = true; + } + } + // q8_0: tile 1, block 3, lane 13 (gq 3, t 1), step 1, column gq + 8, byte 2: column 27, input 8 + 4 + 2 + const o = (1 * 16 + 3) * 544 + 13 * 16 + 3 * 4 + 2; + try std.testing.expectEqual((27 * 16 + 3) * 34 + 2 + 8 + 4 + 2, source(.q8_0, 32, 512, o)); + // iq2xxs: tile 0, block 1, group 5, gq 6, column gq + 8, byte 7; then column 6's d byte 1 + try std.testing.expectEqual((14 * 2 + 1) * 66 + 2 + 40 + 7, source(.iq2xxs, 32, 512, 1056 + 5 * 128 + 6 * 16 + 8 + 7)); + try std.testing.expectEqual((6 * 2 + 1) * 66 + 1, source(.iq2xxs, 32, 512, 1056 + 1024 + 6 * 4 + 1)); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..bc9491ba5e 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 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 &.{}; +pub const ggml: []const u8 = if (available) &Blob("fatbin_ggml").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..1f7d42510e 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -7,6 +7,7 @@ const memory = @import("memory.zig"); const stream_ = @import("stream.zig"); const qmmf = @import("qmmf.zig"); const nvfp4 = @import("nvfp4.zig"); +const ggml = @import("ggml.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -36,11 +37,14 @@ pub const Weight = union(enum) { nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values + q8_0: ggml.Weight, // ggml Q8_0 packed by ggml.Kernels.pack + iq2xxs: ggml.Weight, // ggml IQ2_XXS packed by ggml.Kernels.pack pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .q8_0, .iq2xxs => |q| q.n, }; } @@ -48,6 +52,7 @@ pub const Weight = union(enum) { return switch (w) { .affine4 => |q| q.k, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .q8_0, .iq2xxs => |q| q.k, }; } }; @@ -60,6 +65,7 @@ pub const Linear = struct { affine4: ?*const Affine4Kernels = null, lane: ?*const qmmf.Lane = null, // decode rows of every lane format; fp8g and bf16 prompt rows too nvfp4_prompt: ?*const nvfp4.Prompt = null, // the prompt GEMM: nvfp4 and fp8 prompt rows (linear.py's prefill) + ggml: ?*const ggml.Kernels = null, // q8_0 and iq2xxs rows of both kinds /// Decode rows: out (rows, n) bf16, each row's bits the same at every row count the format takes. pub fn decode(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { @@ -73,6 +79,7 @@ pub const Linear = struct { .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .q8_0, .iq2xxs => |q| return (l.ggml orelse return error.FormatNotLoaded).decode(s, q, in.x, in.ldx, rows, out, .bf16), } } @@ -88,6 +95,8 @@ pub const Linear = struct { .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + // prompt rows always take the prompt GEMM: a row's bits must not depend on where a prompt's chunks are cut + .q8_0, .iq2xxs => |q| return (l.ggml orelse return error.FormatNotLoaded).promptRows(s, q, in.x, in.ldx, rows, out, .bf16), } } diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..5c1a4932e6 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -36,6 +36,7 @@ pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); +pub const ggml = @import("ggml.zig"); test { _ = @import("memory.zig"); @@ -54,4 +55,5 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); + _ = @import("ggml.zig"); } diff --git a/zig/tests/cuda/ggml_tests.zig b/zig/tests/cuda/ggml_tests.zig new file mode 100644 index 0000000000..849349b707 --- /dev/null +++ b/zig/tests/cuda/ggml_tests.zig @@ -0,0 +1,336 @@ +//! Bolt's ggml formats on the GPU (ggml.cu): real tensors' repack, grouped IQ2_XXS experts and dense rows against fp64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Fixture = @import("fixture.zig").Fixture; +const Gpu = check.Gpu; + +const ggml = cuda.ggml; +const grouped = cuda.grouped; +const ex_ = cuda.experts; + +fn arr(fx: Fixture, comptime fmt: []const u8, args: anytype) ![]u8 { + var name: [48]u8 = undefined; + return fx.bytes(try std.fmt.bufPrint(&name, fmt, args)); +} + +fn bf(v: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, v) << 16))); +} + +fn relu2(v: f64) f64 { + const u = @max(v, 0); + return u * u; +} + +/// Max |got - want| and max |want| over outputs fp32 or bf16; fails past `tol` of the largest reference. +fn compare(gpu: Gpu, what: []const u8, b: cuda.DeviceBuffer, ref: []const u8, out: ggml.Out, tol: f64) !void { + const got = try check.download(gpu, b); + defer gpu.gpa.free(got); + const want = std.mem.bytesAsSlice(f64, @as([]align(1) const u8, ref)); + var err: f64 = 0; + var top: f64 = 0; + for (want, 0..) |w0, i| { + const w = if (out == .relu2) relu2(w0) else w0; + const g: f64 = if (out == .f32) @floatCast(std.mem.bytesToValue(f32, got[4 * i ..][0..4])) else bf(std.mem.bytesToValue(u16, got[2 * i ..][0..2])); + if (!(g == g)) return check.expect(false, "{s}: NaN at {d}", .{ what, i }); + err = @max(err, @abs(g - w)); + top = @max(top, @abs(w)); + } + try check.expect(err <= tol * top, "{s}: max abs error {e:.3} is {e:.3} of max |ref| {e:.3}", .{ what, err, err / top, top }); + check.pass("{s}: max abs error {e:.3}, {e:.3} of max |ref| {e:.3}", .{ what, err, err / top, top }); +} + +/// fp32 sums; bf16 rounds once (2^-8 relative); relu^2 rounds its input to bf16 too, then squares it. +fn tolerance(out: ggml.Out) f64 { + return switch (out) { + .f32 => 1e-5, + .bf16 => 1.0 / 256.0, + .relu2 => 3.0 / 256.0, + }; +} + +/// Raw blocks packed on the device equal the host repack. +fn packed_(gpu: Gpu, k: *const ggml.Kernels, s: cuda.Stream, kind: ggml.Kind, raw: []const u8, n: usize, kk: usize, count: usize) !cuda.DeviceBuffer { + var draw = try cuda.DeviceBuffer.fromHost(gpu.d, raw); + defer draw.free(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, raw.len); + errdefer out.free(); + try k.pack(s, kind, draw.ptr, out.ptr, n, kk, count); + try s.synchronize(); + const got = try check.download(gpu, out); + defer gpu.gpa.free(got); + const want = try gpu.gpa.alloc(u8, raw.len); + defer gpu.gpa.free(want); + const one = ggml.bytes(kind, n, kk); + for (0..count) |m| ggml.packHost(kind, want[m * one ..][0..one], raw[m * one ..][0..one], n, kk); + try check.sameBytes("ggml pack", got, want); + return out; +} + +pub fn run(gpu: Gpu, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var k = try ggml.Kernels.load(gpu.d); + defer k.unload(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: launches wait for fromHost's legacy-stream copies + defer stream.deinit(); + try experts(gpu, fx, &k, stream); + try dense(gpu, fx, &k, stream); +} + +fn experts(gpu: Gpu, fx: Fixture, k: *const ggml.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + const e: usize = @intCast(try fx.int("experts")); + const d: usize = @intCast(try fx.int("dims")); + const ni: usize = @intCast(try fx.int("width")); + const slots: usize = @intCast(try fx.int("slots")); + const up_raw = try fx.bytes("up_raw"); + defer a.free(up_raw); + const down_raw = try fx.bytes("down_raw"); + defer a.free(down_raw); + var up = try packed_(gpu, k, s, .iq2xxs, up_raw, ni, d, e); + defer up.free(); + var down = try packed_(gpu, k, s, .iq2xxs, down_raw, d, ni, e); + defer down.free(); + check.pass("ggml pack: {d} real IQ2_XXS experts' up [{d}, {d}] and down on the device equal the host repack", .{ e, ni, d }); + const layer: ex_.Layer = .{ .format = .iq2xxs, .up = up.ptr, .down = down.ptr, .width = ni, .dims = d, .experts = e, .relu2 = true }; + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + const g: ex_.Grouped = .{ .iq2xxs = k }; + inline for (.{ .{ "decode", false }, .{ "prompt", true } }) |set| { + var it = std.mem.tokenizeScalar(u8, try fx.string(set[0]), ','); + while (it.next()) |tok| { + const rows = try std.fmt.parseInt(usize, tok, 10); + try expertCase(gpu, fx, router, g, k, s, layer, rows, slots, set[1]); + if (set[1]) try expertCase(gpu, fx, router, g, k, s, layer, rows, slots, false); // decode form at prompt width + } + } +} + +fn expertCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, k: *const ggml.Kernels, s: cuda.Stream, l: ex_.Layer, rows: usize, slots: usize, prompt: bool) !void { + const a = gpu.gpa; + const pairs = rows * slots; + const tile: usize = if (prompt) ex_.promptTile(.iq2xxs) else 16; + const picks = try arr(fx, "picks{d}", .{rows}); + defer a.free(picks); + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, picks); + defer dpicks.free(); + const big = pairs > grouped.small; + const bound = grouped.maxItems(pairs, l.experts, tile); + var bufs: [5]cuda.DeviceBuffer = undefined; + const sizes = [5]usize{ pairs * 4, bound * 12, 8, (if (big) pairs else 1) * 4, (if (big) (pairs + 1023) / 1024 * l.experts else 1) * 4 }; + var made: usize = 0; + defer for (bufs[0..made]) |*b| b.free(); + for (sizes, &bufs) |n, *b| { + b.* = try cuda.DeviceBuffer.alloc(gpu.d, n); + made += 1; + } + const plan: grouped.Plan = .{ .members = bufs[0].ptr, .items = bufs[1].ptr, .counts = bufs[2].ptr, .rank = bufs[3].ptr, .hist = bufs[4].ptr }; + try router.route(s, dpicks.ptr, pairs, l.experts, tile, plan); + + const x = try arr(fx, "x{d}", .{rows}); + defer a.free(x); + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, x); + defer dx.free(); + const h = try arr(fx, "h{d}", .{rows}); + defer a.free(h); + var dh = try cuda.DeviceBuffer.fromHost(gpu.d, h); + defer dh.free(); + const ref_up = try arr(fx, "up{d}", .{rows}); + defer a.free(ref_up); + const ref_down = try arr(fx, "down{d}", .{rows}); + defer a.free(ref_down); + var y = try cuda.DeviceBuffer.alloc(gpu.d, pairs * @max(l.width, l.dims) * 4); + defer y.free(); + const form = if (prompt) "prompt" else "decode"; + var name: [96]u8 = undefined; + + // up: relu^2 through experts.Grouped, then fp32 sums by the same kernels with no activation + try y.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes + const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode; + try call(g, s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, y.ptr, bound, -1); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "iq2 up relu^2 {s} {d} rows x {d}", .{ form, rows, slots }), y, ref_up, .relu2, tolerance(.relu2)); + try y.fill8(0xff, s.handle); + try k.experts(s, prompt, .iq2xxs, .f32, dx.ptr, l.dims, slots, l.up, l.width, l.dims, plan, y.ptr, bound, -1); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "iq2 up fp32 {s} {d} rows x {d}", .{ form, rows, slots }), y, ref_up, .f32, tolerance(.f32)); + // down from pair rows, fp32 and bf16 + inline for (.{ .{ ex_.Epi.down_f32, ggml.Out.f32 }, .{ ex_.Epi.down_bf16, ggml.Out.bf16 } }) |c| { + try y.fill8(0xff, s.handle); + try call(g, s, c[0], .{ .x = dh.ptr, .stride = l.width }, l, plan, y.ptr, bound, -1); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "iq2 down {t} {s} {d} rows x {d}", .{ c[1], form, rows, slots }), y, ref_down, c[1], tolerance(c[1])); + } +} + +fn dense(gpu: Gpu, fx: Fixture, k: *const ggml.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + var cases = std.mem.tokenizeScalar(u8, try fx.string("dense"), ','); + while (cases.next()) |case| { + var name: [96]u8 = undefined; + const n: usize = @intCast(try fx.int(try std.fmt.bufPrint(&name, "{s}_n", .{case}))); + const kk: usize = @intCast(try fx.int(try std.fmt.bufPrint(&name, "{s}_k", .{case}))); + const kind: ggml.Kind = if (std.mem.eql(u8, try fx.string(try std.fmt.bufPrint(&name, "{s}_kind", .{case})), "Q8_0")) .q8_0 else .iq2xxs; + const raw = try arr(fx, "{s}_raw", .{case}); + defer a.free(raw); + var w_buf = try packed_(gpu, k, s, kind, raw, n, kk, 1); + defer w_buf.free(); + const w = ggml.Weight{ .kind = kind, .w = w_buf.ptr, .n = n, .k = kk }; + var rows_it = std.mem.tokenizeScalar(u8, try fx.string("dense_rows"), ','); + while (rows_it.next()) |tok| { + const rows = try std.fmt.parseInt(usize, tok, 10); + const x = try arr(fx, "{s}_x{d}", .{ case, rows }); + defer a.free(x); + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, x); + defer dx.free(); + const ref = try arr(fx, "{s}_y{d}", .{ case, rows }); + defer a.free(ref); + var y = try cuda.DeviceBuffer.alloc(gpu.d, rows * n * 4); + defer y.free(); + const forms: []const bool = if (rows <= ggml.decode_rows) &.{ false, true } else &.{true}; + for (forms) |prompt| for ([_]ggml.Out{ .f32, .bf16 }) |out| { + try y.fill8(0xff, s.handle); + if (prompt) try k.promptRows(s, w, dx.ptr, 0, rows, y.ptr, out) else try k.decode(s, w, dx.ptr, 0, rows, y.ptr, out); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "{t} {s} [{d}, {d}] {s} {d} rows {t}", .{ kind, case, n, kk, if (prompt) "prompt" else "decode", rows, out }), y, ref, out, tolerance(out)); + }; + } + } +} + +/// Decode GB/s: 9 IQ2_XXS experts' up and down at width 2560 (fresh experts each pass, past L2), and lm_head rows. +pub fn bench(gpu: Gpu) !void { + var k = try ggml.Kernels.load(gpu.d); + defer k.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + + // experts: 16 sets of 9 (144 experts, 194 MB), a pass on each set in turn + const width = 2560; + const dims = 1024; + const sets = 16; + const e = 9 * sets; + const one = ggml.bytes(.iq2xxs, width, dims); + var up = try cuda.DeviceBuffer.alloc(gpu.d, e * one); + defer up.free(); + var down = try cuda.DeviceBuffer.alloc(gpu.d, e * one); + defer down.free(); + try up.fill8(0x11, s.handle); + try down.fill8(0x11, s.handle); + var x = try cuda.DeviceBuffer.alloc(gpu.d, 512 * 4096 * 2); + defer x.free(); + try x.fill8(0x3c, s.handle); // bf16 ~0.01 + var h = try cuda.DeviceBuffer.alloc(gpu.d, 9 * width * 2); + defer h.free(); + try h.fill8(0x3c, s.handle); + var y = try cuda.DeviceBuffer.alloc(gpu.d, 512 * 131072 * 2); + defer y.free(); + var plans = try cuda.DeviceBuffer.alloc(gpu.d, sets * 256); + defer plans.free(); + var picks: [sets * 9]i32 = undefined; + for (&picks, 0..) |*p, i| p.* = @intCast(i); + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&picks)); + defer dpicks.free(); + var ps: [sets]grouped.Plan = undefined; + for (&ps, 0..) |*p, i| { + const base = plans.ptr + i * 256; + p.* = .{ .members = base, .items = base + 64, .counts = base + 192, .rank = 0, .hist = 0 }; + try router.route(s, dpicks.ptr + i * 9 * 4, 9, e, 16, p.*); + } + const l: ex_.Layer = .{ .format = .iq2xxs, .up = up.ptr, .down = down.ptr, .width = width, .dims = dims, .experts = e, .relu2 = true }; + const g: ex_.Grouped = .{ .iq2xxs = &k }; + const reps = 20; + const bound = grouped.maxItems(9, e, 16); + const parts = [_][]const u8{ "up", "down", "up+down" }; + for (parts, 0..) |what, pi| { + for (0..2) |warm| { + try t0.record(s); + for (0..reps) |_| for (ps) |p| { + if (pi != 1) try g.decode(s, .up, .{ .x = x.ptr, .stride = dims, .slots = 9 }, l, p, h.ptr, bound, -1); + if (pi != 0) try g.decode(s, .down_f32, .{ .x = h.ptr, .stride = width }, l, p, y.ptr, bound, -1); + }; + try t1.record(s); + try t1.synchronize(); + if (warm == 0) continue; + const us = try cuda.Event.elapsedMs(t0, t1) * 1e3 / (reps * sets); + const moved: f64 = @floatFromInt(9 * one * @as(usize, if (pi == 2) 2 else 1)); + std.debug.print("RESULT iq2xxs experts decode, 1 token x 9 experts {s} at width {d}: {d:.2} us, {d:.0} GB/s\n", .{ what, width, us, moved / us / 1e3 }); + } + } + + // prompt rows: 2048 tokens x 9 slots over the 144 experts, up then down, as the prompt plan groups them + { + const tokens = 2048; + const pairs = tokens * 9; + const many = try gpu.gpa.alloc(i32, pairs); + defer gpu.gpa.free(many); + var prng = std.Random.DefaultPrng.init(7); + for (0..tokens) |tk| { + const first = prng.random().uintLessThan(usize, e - 9); + for (0..9) |j| many[tk * 9 + j] = @intCast(first + j); + } + var dmany = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(many)); + defer dmany.free(); + const pb = grouped.maxItems(pairs, e, 64); + const hist = (pairs + 1023) / 1024 * e * 4; + var scratch = try cuda.DeviceBuffer.alloc(gpu.d, pairs * 8 + pb * 12 + 16 + hist); + defer scratch.free(); + const base = scratch.ptr; + const pp: grouped.Plan = .{ .members = base, .items = base + pairs * 4, .counts = base + pairs * 4 + pb * 12, .rank = base + pairs * 4 + pb * 12 + 16, .hist = base + pairs * 8 + pb * 12 + 16 }; + try router.route(s, dmany.ptr, pairs, e, 64, pp); + var xp = try cuda.DeviceBuffer.alloc(gpu.d, tokens * dims * 2); + defer xp.free(); + try xp.fill8(0x3c, s.handle); + var hp = try cuda.DeviceBuffer.alloc(gpu.d, pairs * width * 2); + defer hp.free(); + var yp = try cuda.DeviceBuffer.alloc(gpu.d, pairs * dims * 2); + defer yp.free(); + var t2 = try cuda.Event.init(gpu.d, true); + defer t2.deinit(); + for (0..2) |warm| { + try t0.record(s); + try g.prompt(s, .up, .{ .x = xp.ptr, .stride = dims, .slots = 9 }, l, pp, hp.ptr, pb, -1); + try t1.record(s); + try g.prompt(s, .down_bf16, .{ .x = hp.ptr, .stride = width }, l, pp, yp.ptr, pb, -1); + try t2.record(s); + try t2.synchronize(); + if (warm == 0) continue; + const up_us = try cuda.Event.elapsedMs(t0, t1) * 1e3; + const down_us = try cuda.Event.elapsedMs(t1, t2) * 1e3; + const flops: f64 = 2.0 * pairs * width * dims; + std.debug.print("RESULT iq2xxs experts prompt, {d} tokens x 9 over {d} experts at width {d}: up {d:.0} us ({d:.0} TFLOP/s), down {d:.0} us ({d:.0} TFLOP/s)\n", .{ tokens, e, width, up_us, flops / up_us / 1e6, down_us, flops / down_us / 1e6 }); + } + } + + // lm_head: Q8_0 [131072, 4096] + const n = 131072; + const kk = 4096; + const lb = ggml.bytes(.q8_0, n, kk); + var lm = try cuda.DeviceBuffer.alloc(gpu.d, lb); + defer lm.free(); + try lm.fill8(0x11, s.handle); + const w = ggml.q8Weight(lm.ptr, n, kk); + for ([_]usize{ 1, 4, 16, 64, 512 }) |rows| for ([_]ggml.Out{ .f32, .bf16 }) |out| { + if (rows > 16 and out == .f32) continue; + const r2: usize = if (rows > 16) 3 else 10; + for (0..2) |warm| { + try t0.record(s); + for (0..r2) |_| try k.matmul(s, w, x.ptr, 0, rows, y.ptr, out); + try t1.record(s); + try t1.synchronize(); + if (warm == 0) continue; + const us = try cuda.Event.elapsedMs(t0, t1) * 1e3 / @as(f32, @floatFromInt(r2)); + std.debug.print("RESULT q8_0 lm_head [{d}, {d}] {d} rows {t}: {d:.1} us, {d:.0} GB/s\n", .{ n, kk, rows, out, us, @as(f64, @floatFromInt(lb)) / us / 1e3 }); + } + }; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..f8e95d6012 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -29,6 +29,7 @@ const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); +const ggml_tests = @import("ggml_tests.zig"); const usage = \\usage: tf-cuda-test @@ -64,6 +65,8 @@ const usage = \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes + \\ ggml Bolt's IQ2_XXS experts and Q8_0 / IQ2_XXS dense rows on real tensors against fp64 (tools/zig/bolt_ggml_fixture.py) + \\ ggml-bench IQ2_XXS expert and Q8_0 lm_head decode GB/s (synthetic blocks) \\ modelopt ModelOpt repack, FP8 and BF16 lane rows and relu^2 NVFP4 experts against fp64 \\ exlr8 EXLR8 projections: every codebook state, FP64, row windows, the Metal frozen digests \\ exlr8-bench DIR [MANIFEST] [LAYER] [SETS] a real EXLR8 layer: FP64, row windows, then GB/s by stage @@ -124,6 +127,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); + if (std.mem.eql(u8, cmd, "ggml")) return ggml_tests.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ggml-bench")) return ggml_tests.bench(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); From 5f0945436137bad14c1c145d42e79850f3306162 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:56:27 +0000 Subject: [PATCH 09/11] nemotron: Nemotron 3.5 Bolt 74B-A9B on CUDA: puzzle layers, LatentMoE, ggml weights, image input, host prompt states NVIDIA's Nemotron-3.5-Bolt-74B-A9B-24GB in the Nemotron family on one card: - Config: the nemotron_h_omni wrapper's llm_config and its nemotron_h_puzzle block_configs (88 layers; each MoE layer its own 240-496 experts, 1536-2560 wide, top 9), read beside the uniform Lightning config. - LatentMoE: rows to the 1024-wide latent through fc1_latent_proj (Q8_0), routed IQ2_XXS experts there, their weighted sum in fp32 (tf_latent_sum), back through fc2_latent_proj, plus the full-width shared expert as a one-expert layer. The router takes up to 512 experts and eight K slices at 4096 (a block under 48 KiB). - Weights: a tensor's format from its shape (ggml blocks are rank 3); the checkpoint's language_model. scope; ggml prompt rows always on the prompt GEMM, so a resumed chunk keeps its bits. - Image input (--vision): the C-RADIO tower and projector on CUDA (radio_vision.cu, cuda_vision.zig), the processor's one-tile preparation (vision/nemotron_image.zig), ... placeholders, the rows spliced over them as each prompt chunk embeds. Image prompts are never resumed or kept. - Kept prompt states in pinned host memory on a card with its own memory (a state is 166 MiB of Mamba states plus its rows; GB10 keeps device copies). No MTP head: Bolt decodes serially. Structured output reads vocab_size from llm_config too, and the banner says drafts: off when a family decodes one row. --- zig/build/cuda.zig | 7 + zig/kernels/cuda/nemotron_ops.cu | 30 ++ zig/kernels/cuda/nemotron_route.cu | 2 +- zig/kernels/cuda/radio_vision.cu | 275 ++++++++++ zig/src/cli/hub.zig | 1 + zig/src/core/checkpoint.zig | 43 +- zig/src/core/engine_api.zig | 2 + zig/src/cuda/kernels.zig | 1 + zig/src/families/nemotron/config.zig | 160 +++++- zig/src/families/nemotron/cuda.zig | 3 + zig/src/families/nemotron/cuda_engine.zig | 15 +- zig/src/families/nemotron/cuda_forward.zig | 58 ++- zig/src/families/nemotron/cuda_glue.zig | 25 +- zig/src/families/nemotron/cuda_kernels.zig | 52 +- zig/src/families/nemotron/cuda_lanes.zig | 21 +- zig/src/families/nemotron/cuda_modelopt.zig | 120 ++++- zig/src/families/nemotron/cuda_native.zig | 27 +- zig/src/families/nemotron/cuda_prompt.zig | 2 +- zig/src/families/nemotron/cuda_reuse.zig | 80 ++- zig/src/families/nemotron/cuda_state.zig | 15 +- zig/src/families/nemotron/cuda_train_fwd.zig | 2 +- zig/src/families/nemotron/cuda_vision.zig | 511 +++++++++++++++++++ zig/src/families/nemotron/cuda_weights.zig | 14 +- zig/src/families/nemotron/glue_ref.zig | 8 +- zig/src/native/cuda.zig | 43 +- zig/src/native/cuda_memory.zig | 20 + zig/src/server/server.zig | 4 +- zig/src/server/vision_inputs.zig | 18 +- zig/src/vision/fixtures/nemotron_resize.bin | Bin 0 -> 23808 bytes zig/src/vision/nemotron_image.zig | 278 ++++++++++ zig/src/vision/qwen_image.zig | 3 + zig/tests/bolt_vision_main.zig | 183 +++++++ zig/tests/cuda/glue_tests.zig | 31 +- 33 files changed, 1957 insertions(+), 97 deletions(-) create mode 100644 zig/kernels/cuda/radio_vision.cu create mode 100644 zig/src/families/nemotron/cuda_vision.zig create mode 100644 zig/src/vision/fixtures/nemotron_resize.bin create mode 100644 zig/src/vision/nemotron_image.zig create mode 100644 zig/tests/bolt_vision_main.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index dce26ba712..004e23f056 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -33,6 +33,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 = "radio_vision", .flags = &.{ "-O3", "--fmad=false" } }, // ours: Nemotron 3.5 Bolt's C-RADIO tower and projector .{ .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 }, @@ -142,6 +143,12 @@ 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 })); + // `zig build bolt-vision`: Bolt's image tower against the reference's dumps (zig/tests/bolt_vision_main.zig) + const bolt = b.createModule(.{ .root_source_file = b.path("zig/tests/bolt_vision_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); + bolt.addImport("cuda", cuda); + bolt.addImport("nemotron", mods.nemotron); + bolt.addImport("qwen_image", vision_build.module(b, target, optimize)); + b.step("bolt-vision", "tf-bolt-vision: Bolt's image tower and projector against bolt_vision_reference.py's dumps").dependOn(&b.addInstallArtifact(b.addExecutable(.{ .name = "tf-bolt-vision", .root_module = bolt }), .{}).step); nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; } diff --git a/zig/kernels/cuda/nemotron_ops.cu b/zig/kernels/cuda/nemotron_ops.cu index b8bd0603de..d4304e8f7e 100644 --- a/zig/kernels/cuda/nemotron_ops.cu +++ b/zig/kernels/cuda/nemotron_ops.cu @@ -291,3 +291,33 @@ extern "C" __global__ void tf_add_row(void* __restrict__ y, int y_f32, const flo *p = __float2bfloat16(__bfloat162float(*p) + x[i]); } } + +// LatentMoE (Bolt): a row's routed outputs in the latent width, weighted and summed in fp32, then bf16 like the +// reference's index_add_ sum. Y [R * K, L] fp32 (pair rows in slot order), WT [R * K], Z [R, L]; one block a row. +extern "C" __global__ void tf_latent_sum(const float* __restrict__ Y, const float* __restrict__ WT, __nv_bfloat16* __restrict__ Z, + int K, int L) { + const int64_t r = blockIdx.x; + for (int i = threadIdx.x; i < L; i += blockDim.x) { + float acc = 0.0f; + for (int k = 0; k < K; ++k) acc = __fmaf_rn(Y[(r * K + k) * L + i], WT[r * K + k], acc); + Z[r * L + i] = __float2bfloat16(acc); + } +} + +// out = a + b in bf16 (b bf16, or fp32 when b_f32): the latent experts' projection plus the shared expert, n values. +extern "C" __global__ void tf_add_bf16(__nv_bfloat16* __restrict__ out, const __nv_bfloat16* __restrict__ a, const void* __restrict__ b, + int b_f32, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const float y = b_f32 ? static_cast(b)[i] : __bfloat162float(static_cast(b)[i]); + out[i] = __float2bfloat16(__bfloat162float(a[i]) + y); +} + +// An image prompt's chunk: rows whose SRC names a feature row take it in place of their token's embedding. +// H [rows, D] bf16, F [image rows, D] bf16, SRC [rows] (-1: a text row); one block a row. +extern "C" __global__ void tf_splice_rows(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ F, const int* __restrict__ SRC, int D) { + const int64_t r = blockIdx.x; + const int s = SRC[r]; + if (s < 0) return; + for (int i = threadIdx.x; i < D; i += blockDim.x) H[r * D + i] = F[static_cast(s) * D + i]; +} diff --git a/zig/kernels/cuda/nemotron_route.cu b/zig/kernels/cuda/nemotron_route.cu index 23c6b78f35..4538440a5e 100644 --- a/zig/kernels/cuda/nemotron_route.cu +++ b/zig/kernels/cuda/nemotron_route.cu @@ -38,7 +38,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_nemo_router(const __nv_bflo namespace { -constexpr int MAXC = 8; // experts a lane holds: up to 256 experts +constexpr int MAXC = 16; // experts a lane holds: up to 512 experts (Bolt's widest layer has 496) // The larger of two (value, index) picks: a higher value, or the lower index on a tie (values are never NaN). __device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return v > bv || (v == bv && i < bi); } diff --git a/zig/kernels/cuda/radio_vision.cu b/zig/kernels/cuda/radio_vision.cu new file mode 100644 index 0000000000..dcfffded2c --- /dev/null +++ b/zig/kernels/cuda/radio_vision.cu @@ -0,0 +1,275 @@ +// Nemotron 3.5 Bolt's vision tower (C-RADIO ViT-H/16) and projector: glue around cuBLASLt's bf16 GEMMs, and attention. +// The residual stream stays fp32 (its rows grow to |x| ~1700); norms, softmax and epilogues compute in fp32, GEMM +// inputs are rounded to bf16 once. Oracle: tools/zig/bolt_vision_reference.py (fp32 PyTorch). +#include +#include + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } + +__device__ __forceinline__ uint32_t pack2(float lo, float hi) { + __nv_bfloat162 h = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&h); +} + +// fp32 block sum over blockDim.x threads (a multiple of 32, at most 1024); every thread gets the total. +__device__ __forceinline__ float block_sum(float v, float* part) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5; + for (int o = 16; o; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + float t = 0.0f; + for (int i = 0; i < warps; ++i) t += part[i]; + return t; +} + +// x [prefix + n, D] fp32: rows below `prefix` the cls/register tokens, then acc (the patch Linear) plus the CPE +// position rows, bilinear from the 128x128 table: (p0 w0 + p1 w1) h0 + (p2 w0 + p3 w1) h1 as aten's interpolate. +extern "C" __global__ void __launch_bounds__(256) tf_rv_embed(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ pos, + const int* __restrict__ idx, + const float* __restrict__ wts, + const __nv_bfloat16* __restrict__ cls, + float* __restrict__ x, int prefix, int D) { + const long long row = blockIdx.x; + if (row < prefix) { + for (int i = threadIdx.x; i < D; i += blockDim.x) x[row * D + i] = bf(cls[row * D + i]); + return; + } + const long long p = row - prefix; + const int* ri = idx + p * 4; + const float* rw = wts + p * 4; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float a = bf(pos[(long long)ri[0] * D + i]) * rw[0] + bf(pos[(long long)ri[1] * D + i]) * rw[1]; + const float b = bf(pos[(long long)ri[2] * D + i]) * rw[0] + bf(pos[(long long)ri[3] * D + i]) * rw[1]; + x[row * D + i] = acc[p * D + i] + (a * rw[2] + b * rw[3]); + } +} + +// LayerNorm of fp32 rows (mean, then the biased variance about it), bf16 out for the next GEMM; a block a row. +extern "C" __global__ void __launch_bounds__(256) tf_rv_layernorm(const float* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ y, int D, float eps) { + __shared__ float part[32]; + const long long row = blockIdx.x; + const float* xr = x + row * D; + float s = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) s += xr[i]; + const float mean = block_sum(s, part) / D; + float v = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float d = xr[i] - mean; + v += d * d; + } + const float r = rsqrtf(block_sum(v, part) / D + eps); + for (int i = threadIdx.x; i < D; i += blockDim.x) + y[row * D + i] = __float2bfloat16_rn((xr[i] - mean) * r * bf(w[i]) + bf(b[i])); +} + +// acc [rows, 3D] + bias split into q, k, v [heads][rows][hd] bf16 (the reference's view(T, 3, nh, hd)). +extern "C" __global__ void __launch_bounds__(256) tf_rv_qkv(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __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 D = heads * hd; + const float* a = acc + row * 3 * D; + for (int c = threadIdx.x; c < 3 * D; c += blockDim.x) { + const int s = c / D, h = (c % D) / hd, j = c % hd; + const __nv_bfloat16 val = __float2bfloat16_rn(a[c] + bf(bias[c])); + __nv_bfloat16* dst = s == 0 ? q : (s == 1 ? k : v); + dst[((long long)h * rows + row) * hd + j] = val; + } +} + +__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +// Full (bidirectional) attention over all `rows` tokens, heads of 80: flash style, one warp 16 queries, 64-key +// tiles in shared memory, bf16 mma.sync with fp32 scores, an online softmax in fp32 (exp2 of pre-scaled scores) and +// fp32 output accumulators; out [rows, heads * 80] bf16 (the reference's transpose(0, 1).reshape(T, C)). +#define RV_HD 80 +#define RV_BK 64 +extern "C" __global__ void __launch_bounds__(128) tf_rv_attention(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ out, int rows, int heads, + float scale_log2) { + __shared__ __align__(16) __nv_bfloat16 ks[RV_BK][RV_HD + 8]; + __shared__ __align__(16) __nv_bfloat16 vt[RV_HD][RV_BK + 8]; + const int head = blockIdx.y, warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const long long base = (long long)head * rows; + const int q0 = blockIdx.x * 64 + warp * 16; + uint32_t qa[RV_HD / 16][4]; + { + const uint32_t* p0 = reinterpret_cast(q + (base + min(q0 + g, rows - 1)) * RV_HD); + const uint32_t* p1 = reinterpret_cast(q + (base + min(q0 + g + 8, rows - 1)) * RV_HD); +#pragma unroll + for (int kc = 0; kc < RV_HD / 16; ++kc) { + qa[kc][0] = p0[kc * 8 + t]; + qa[kc][1] = p1[kc * 8 + t]; + qa[kc][2] = p0[kc * 8 + 4 + t]; + qa[kc][3] = p1[kc * 8 + 4 + t]; + } + } + float o[RV_HD / 8][4]; +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) o[d][0] = o[d][1] = o[d][2] = o[d][3] = 0.0f; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + for (int k0 = 0; k0 < rows; k0 += RV_BK) { + __syncthreads(); + for (int i = threadIdx.x; i < RV_BK * RV_HD / 2; i += blockDim.x) { + const int r = i / (RV_HD / 2), c = i % (RV_HD / 2); + uint32_t kv = 0, vv = 0; + if (k0 + r < rows) { + kv = reinterpret_cast(k + (base + k0 + r) * RV_HD)[c]; + vv = reinterpret_cast(v + (base + k0 + r) * RV_HD)[c]; + } + *reinterpret_cast(&ks[r][2 * c]) = kv; + const __nv_bfloat162 vh = *reinterpret_cast<__nv_bfloat162*>(&vv); + vt[2 * c][r] = vh.x; + vt[2 * c + 1][r] = vh.y; + } + __syncthreads(); + float s[RV_BK / 8][4]; +#pragma unroll + for (int nt = 0; nt < RV_BK / 8; ++nt) { + s[nt][0] = s[nt][1] = s[nt][2] = s[nt][3] = 0.0f; + const uint32_t* kr = reinterpret_cast(&ks[nt * 8 + g][0]); +#pragma unroll + for (int kc = 0; kc < RV_HD / 16; ++kc) mma16816(s[nt], qa[kc], kr[kc * 8 + t], kr[kc * 8 + 4 + t]); + } + float mx0 = -INFINITY, mx1 = -INFINITY; +#pragma unroll + for (int nt = 0; nt < RV_BK / 8; ++nt) { +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int key = k0 + nt * 8 + 2 * t + (e & 1); + s[nt][e] = key < rows ? s[nt][e] * scale_log2 : -INFINITY; + } + mx0 = fmaxf(mx0, fmaxf(s[nt][0], s[nt][1])); + mx1 = fmaxf(mx1, fmaxf(s[nt][2], s[nt][3])); + } + for (int off = 1; off < 4; off <<= 1) { + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, off)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, off)); + } + const float n0 = fmaxf(m0, mx0), n1 = fmaxf(m1, mx1); + const float c0 = exp2f(m0 - n0), c1 = exp2f(m1 - n1); + m0 = n0; + m1 = n1; + float sum0 = 0.0f, sum1 = 0.0f; +#pragma unroll + for (int nt = 0; nt < RV_BK / 8; ++nt) { + s[nt][0] = exp2f(s[nt][0] - n0); + s[nt][1] = exp2f(s[nt][1] - n0); + s[nt][2] = exp2f(s[nt][2] - n1); + s[nt][3] = exp2f(s[nt][3] - n1); + sum0 += s[nt][0] + s[nt][1]; + sum1 += s[nt][2] + s[nt][3]; + } + l0 = l0 * c0 + sum0; + l1 = l1 * c1 + sum1; +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) { + o[d][0] *= c0; + o[d][1] *= c0; + o[d][2] *= c1; + o[d][3] *= c1; + } +#pragma unroll + for (int kc = 0; kc < RV_BK / 16; ++kc) { + uint32_t pa[4]; + pa[0] = pack2(s[2 * kc][0], s[2 * kc][1]); + pa[1] = pack2(s[2 * kc][2], s[2 * kc][3]); + pa[2] = pack2(s[2 * kc + 1][0], s[2 * kc + 1][1]); + pa[3] = pack2(s[2 * kc + 1][2], s[2 * kc + 1][3]); +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) { + const uint32_t* vr = reinterpret_cast(&vt[d * 8 + g][0]); + mma16816(o[d], pa, vr[kc * 8 + t], vr[kc * 8 + 4 + t]); + } + } + } + for (int off = 1; off < 4; off <<= 1) { + l0 += __shfl_xor_sync(0xffffffffu, l0, off); + l1 += __shfl_xor_sync(0xffffffffu, l1, off); + } + const float i0 = 1.0f / l0, i1 = 1.0f / l1; + const int D = heads * RV_HD; + const int r0 = q0 + g, r1 = q0 + g + 8; +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) { + const int col = head * RV_HD + d * 8 + 2 * t; + if (r0 < rows) *reinterpret_cast(out + (long long)r0 * D + col) = pack2(o[d][0] * i0, o[d][1] * i0); + if (r1 < rows) *reinterpret_cast(out + (long long)r1 * D + col) = pack2(o[d][2] * i1, o[d][3] * i1); + } +} + +// y = gelu(acc + b), the exact erf form (nn.GELU), fp32, bf16 out. +extern "C" __global__ void __launch_bounds__(256) tf_rv_bias_gelu(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) { + const float x = acc[row * D + i] + bf(bias[i]); + y[row * D + i] = __float2bfloat16_rn(0.5f * x * (1.0f + erff(x * 0.7071067811865476f))); + } +} + +// x += acc + b, the residual in fp32. +extern "C" __global__ void __launch_bounds__(256) tf_rv_bias_residual(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + float* __restrict__ x, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) x[row * D + i] += acc[row * D + i] + bf(bias[i]); +} + +// Pixel shuffle v2 (2x2 patches to one row of 4D, channel (dy * 2 + dx) * D + c) of the patch rows behind `prefix`, +// then RMSNorm(4D, eps) with weight w in fp32; bf16 out. A block an output row (grid gh/2 x gw/2, row-major). +extern "C" __global__ void __launch_bounds__(256) tf_rv_shuffle_rms(const float* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, int prefix, int gw, + int D, float eps) { + __shared__ float part[32]; + const int r = blockIdx.x; + const int by = r / (gw / 2), bx = r % (gw / 2); + const int W = 4 * D; + float s = 0.0f; + for (int ch = threadIdx.x; ch < W; ch += blockDim.x) { + const int sub = ch / D, c = ch % D; + const long long tok = prefix + (long long)(2 * by + sub / 2) * gw + 2 * bx + sub % 2; + const float v = x[tok * D + c]; + s += v * v; + } + const float rs = rsqrtf(block_sum(s, part) / W + eps); + for (int ch = threadIdx.x; ch < W; ch += blockDim.x) { + const int sub = ch / D, c = ch % D; + const long long tok = prefix + (long long)(2 * by + sub / 2) * gw + 2 * bx + sub % 2; + y[(long long)r * W + ch] = __float2bfloat16_rn(x[tok * D + c] * rs * bf(w[ch])); + } +} + +// y = relu(acc)^2, bf16 out. +extern "C" __global__ void __launch_bounds__(256) tf_rv_relu2(const float* __restrict__ acc, __nv_bfloat16* __restrict__ y, + int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float a = fmaxf(acc[row * D + i], 0.0f); + y[row * D + i] = __float2bfloat16_rn(a * a); + } +} + +// y = bf16(acc): the projector's last Linear. +extern "C" __global__ void __launch_bounds__(256) tf_rv_cast(const float* __restrict__ acc, __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]); +} diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 64347ae50c..975db3e9a9 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -10,6 +10,7 @@ pub const Family = struct { model_type: []const u8, title: []const u8, formats: const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "nemotron_h", .title = "Nemotron 3.5 Lightning" }, + .{ .model_type = "nemotron_h_omni", .title = "Nemotron 3.5 Bolt" }, .{ .model_type = "qwen4_exp", .title = "Qwen3.8 Flash Next" }, .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, diff --git a/zig/src/core/checkpoint.zig b/zig/src/core/checkpoint.zig index 1c2110a685..8e62ac1575 100644 --- a/zig/src/core/checkpoint.zig +++ b/zig/src/core/checkpoint.zig @@ -59,6 +59,14 @@ pub const Checkpoint = struct { io: Io, files: std.ArrayList(st.File) = .empty, used: std.StringHashMapUnmanaged(void) = .empty, + /// Prepended to every name get, has and skip take: a wrapper's part read with the names it has alone + /// ("language_model." in Nemotron 3.5 Bolt, beside its vision tower's). + scope: []const u8 = "", + + fn scoped(self: *const Checkpoint, buf: *[320]u8, name: []const u8) ![]const u8 { + if (self.scope.len == 0) return name; + return std.fmt.bufPrint(buf, "{s}{s}", .{ self.scope, name }); + } /// Every shard of `dir`, mapped. pub fn openModel(gpa: std.mem.Allocator, io: Io, dir: []const u8) !Checkpoint { @@ -90,7 +98,9 @@ pub const Checkpoint = struct { } /// The tensor named `name`, marked as used; a missing name is an error naming it. - pub fn get(self: *Checkpoint, name: []const u8) !Tensor { + pub fn get(self: *Checkpoint, short: []const u8) !Tensor { + var buf: [320]u8 = undefined; + const name = try self.scoped(&buf, short); for (self.files.items) |*f| if (f.get(name)) |t| { try self.used.put(self.gpa, f.names.getKey(name).?, {}); return t; @@ -110,13 +120,17 @@ pub const Checkpoint = struct { } /// Whether any file holds `name` (it is not marked as used). - pub fn has(self: *const Checkpoint, name: []const u8) bool { + pub fn has(self: *const Checkpoint, short: []const u8) bool { + var buf: [320]u8 = undefined; + const name = self.scoped(&buf, short) catch return false; for (self.files.items) |*f| if (f.get(name) != null) return true; return false; } /// Marks every name starting with `prefix` as used: a part the loader leaves out on purpose. - pub fn skip(self: *Checkpoint, prefix: []const u8) !void { + pub fn skip(self: *Checkpoint, short: []const u8) !void { + var buf: [320]u8 = undefined; + const prefix = try self.scoped(&buf, short); for (self.files.items) |*f| for (f.names.keys()) |k| { if (std.mem.startsWith(u8, k, prefix)) try self.used.put(self.gpa, k, {}); }; @@ -224,3 +238,26 @@ test "has finds a name without taking it; skip marks a prefix used" { _ = try ck.get("backbone.x"); try std.testing.expectEqual(@as(usize, 0), ck.unused()); } + +test "a scope reads a wrapper's part by its own names" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + for ([_][]const u8{ "language_model.backbone.x", "vision_model.y" }, 0..) |name, i| { + const image = try testImage(gpa, name, 4); + defer gpa.free(image); + var file: [32]u8 = undefined; + try tmp.dir.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&file, "model-{d}.safetensors", .{i}), .data = image }); + } + var path: [64]u8 = undefined; + var ck = try Checkpoint.openModel(gpa, io, try std.fmt.bufPrint(&path, ".zig-cache/tmp/{s}", .{tmp.sub_path})); + defer ck.close(); + ck.scope = "language_model."; + try std.testing.expect(ck.has("backbone.x") and !ck.has("vision_model.y")); + _ = try ck.get("backbone.x"); + try std.testing.expectEqual(@as(u32, 1), ck.used.count()); // vision_model.y is left + ck.scope = ""; + try ck.skip("vision_model."); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); +} diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 3c633640ad..1a86bbf3b4 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -171,6 +171,8 @@ pub const Info = struct { 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, + /// Its tower takes Nemotron 3.5 Bolt's preparation (qwen_image.nemotron) and ... placeholders. + vision_nemotron: 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, }; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index bc9491ba5e..d8f77b0dfd 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -40,6 +40,7 @@ pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_point pub const torch_indexing: []const u8 = if (available) &Blob("fatbin_torch_indexing").bytes else &.{}; pub const torch_movement: []const u8 = if (available) &Blob("fatbin_torch_movement").bytes else &.{}; pub const torch_nemotron_constants: []const u8 = if (available) &Blob("fatbin_torch_nemotron_constants").bytes else &.{}; +pub const radio_vision: []const u8 = if (available) &Blob("fatbin_radio_vision").bytes else &.{}; pub const nemotron_norms: []const u8 = if (available) &Blob("fatbin_nemotron_norms").bytes else &.{}; pub const nemotron_route: []const u8 = if (available) &Blob("fatbin_nemotron_route").bytes else &.{}; pub const nemotron_mamba: []const u8 = if (available) &Blob("fatbin_nemotron_mamba").bytes else &.{}; diff --git a/zig/src/families/nemotron/config.zig b/zig/src/families/nemotron/config.zig index aec6e207f0..befe4439ac 100644 --- a/zig/src/families/nemotron/config.zig +++ b/zig/src/families/nemotron/config.zig @@ -3,7 +3,7 @@ const std = @import("std"); pub const Kind = enum { mamba, moe, attention }; -pub const max_layers = 64; +pub const max_layers = 96; // Nemotron 3.5 Bolt has 88 /// The MTP head's acceptance at depth 1, 2, ... given the ones before it, until a stream has its own (model.py draft_prior). pub const draft_prior = [_]f64{ 0.8, 0.72, 0.68, 0.62, 0.58, 0.55, 0.5, 0.5 }; @@ -31,6 +31,9 @@ pub const pocket_key = "tensorfold_learned_neurons"; /// How the checkpoint stores its matrices: MLX affine 4-bit, or NVIDIA ModelOpt (FP8, NVFP4 or bf16, read per tensor). pub const Format = enum { mlx, modelopt }; +/// One MoE layer's routed experts and their width. +pub const Moe = struct { experts: usize = 0, width: usize = 0 }; + pub const Config = struct { hidden: usize, vocab: usize, @@ -44,9 +47,12 @@ pub const Config = struct { heads: usize, kv_heads: usize, head_dim: usize, - experts: usize, + experts: usize, // the widest MoE layer's routed experts (every layer's, unless `moe` says otherwise) top_k: usize, - expert_width: usize, + expert_width: usize, // the widest MoE layer's routed expert width + moe: [max_layers]Moe = @splat(.{}), // each MoE layer's own experts and width (block_configs); zero: the uniform ones + omni: bool = false, // nemotron_h_omni: the language model's tensors under language_model., a vision tower beside it + latent: usize = 0, // moe_latent_size: routed experts read and write this width through fc1/fc2_latent_proj (LatentMoE) shared_width: usize, // the shared expert's own neurons, learned ones (pocket) left out pocket: usize = 0, // tensorfold_learned_neurons: learned neurons after the shared expert's own, in the same tensors routed_scaling: f32, @@ -60,6 +66,17 @@ pub const Config = struct { eos: [4]u32 = .{ 0, 0, 0, 0 }, eos_count: usize = 0, + /// MoE layer `i`'s routed experts and width. + pub fn layerMoe(self: Config, i: usize) Moe { + const m = self.moe[i]; + return if (m.experts == 0) .{ .experts = self.experts, .width = self.expert_width } else m; + } + + /// The width routed experts read and write: the latent one, else hidden. + pub fn expertDims(self: Config) usize { + return if (self.latent > 0) self.latent else self.hidden; + } + pub fn inner(self: Config) usize { return self.mamba_heads * self.mamba_head_dim; } @@ -76,9 +93,9 @@ pub const Config = struct { return (self.heads + 2 * self.kv_heads) * self.head_dim; } - /// A token's expert pairs: its routed experts, then the shared expert's two halves. + /// A token's expert pairs: its routed experts, then the shared expert's two halves (LatentMoE runs it on its own). pub fn slots(self: Config) usize { - return self.top_k + 2; + return if (self.latent > 0) self.top_k else self.top_k + 2; } pub fn count(self: Config, kind: Kind) usize { @@ -96,7 +113,7 @@ pub const Config = struct { 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 text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 25)); defer gpa.free(text); var why: Why = .{}; return parse(gpa, text, &why) catch |e| { @@ -118,6 +135,8 @@ fn number(v: std.json.Value) !f64 { return switch (v) { .float => |f| f, .integer => |i| @floatFromInt(i), + // transformers 5 writes non-finite floats as {"__float__": "Infinity"} + .object => |o| if (o.get("__float__")) |f| (if (f == .string) std.fmt.parseFloat(f64, f.string) catch error.BadConfig else error.BadConfig) else error.BadConfig, else => error.BadConfig, }; } @@ -225,8 +244,11 @@ fn learned(path: []const u8) bool { pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config { var parsed = try std.json.parseFromSlice(std.json.Value, allocator, text, .{}); defer parsed.deinit(); - const o = parsed.value.object; - if (!std.mem.eql(u8, (o.get("model_type") orelse return error.BadConfig).string, "nemotron_h")) return error.NotNemotronH; + const root = parsed.value.object; + // nemotron_h_omni (Bolt) wraps the language model's config in llm_config, its quantization at the top + const wrapped = isType(root, "nemotron_h_omni"); + const o = if (wrapped) (llmConfig(root) orelse return error.BadConfig) else root; + if (!isType(o, "nemotron_h") and !isType(o, "nemotron_h_puzzle")) return error.NotNemotronH; var c = Config{ .hidden = try int(o, "hidden_size"), .vocab = try int(o, "vocab_size"), @@ -239,9 +261,10 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config .heads = try int(o, "num_attention_heads"), .kv_heads = try int(o, "num_key_value_heads"), .head_dim = 0, - .experts = try int(o, "n_routed_experts"), - .top_k = try int(o, "num_experts_per_tok"), - .expert_width = try int(o, "moe_intermediate_size"), + // a puzzle config (Bolt) names these per block + .experts = if (o.get("block_configs") != null) 0 else try int(o, "n_routed_experts"), + .top_k = if (o.get("block_configs") != null) 0 else try int(o, "num_experts_per_tok"), + .expert_width = if (o.get("block_configs") != null) 0 else try int(o, "moe_intermediate_size"), .shared_width = try int(o, "moe_shared_expert_intermediate_size"), .routed_scaling = if (o.get("routed_scaling_factor")) |v| @floatCast(try number(v)) else 1.0, .eps = if (o.get("layer_norm_epsilon")) |v| @floatCast(try number(v)) else 1e-5, @@ -253,8 +276,13 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config c.shared_width -= c.pocket; } if (o.get("norm_topk_prob")) |v| c.norm_topk = v.bool; + c.omni = wrapped; + if (o.get("moe_latent_size")) |v| if (v == .integer) { + c.latent = @intCast(v.integer); + }; + if (o.get("block_configs")) |v| if (v == .array and v.array.items.len > 0) try puzzle(&c, v.array.items, why); // hybrid_override_pattern (one character a block) wins over layers_block_type, as the Python loader reads them - if (o.get("hybrid_override_pattern")) |v| if (v == .string and v.string.len > 0) { + if (c.layers == 0) if (o.get("hybrid_override_pattern")) |v| if (v == .string and v.string.len > 0) { if (v.string.len > max_layers) return error.BadConfig; for (v.string, 0..) |ch, i| c.kinds[i] = try kindOf(&.{ch}); c.layers = v.string.len; @@ -265,25 +293,66 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config for (blocks, 0..) |b, i| c.kinds[i] = try kindOf(b.string); c.layers = blocks.len; } - if (o.get("eos_token_id")) |e| switch (e) { - .integer => |i| { - c.eos[0] = @intCast(i); - c.eos_count = 1; - }, - .array => |a| for (a.items[0..@min(a.items.len, 4)]) |x| { - c.eos[c.eos_count] = @intCast(x.integer); - c.eos_count += 1; - }, - else => {}, - }; + // generation_config's [2, 11] is the omni wrapper's eos plus the language model's + if (wrapped) if (root.get("eos_token_id")) |e| eos(&c, e); + if (o.get("eos_token_id")) |e| eos(&c, e); if (o.get("time_step_limit")) |v| if (v == .array and v.array.items.len == 2) { c.dt_min = @floatCast(try number(v.array.items[0])); c.dt_max = @floatCast(try number(v.array.items[1])); }; - try quantization(&c, o, why); + try quantization(&c, root, why); return c; } +fn isType(o: std.json.ObjectMap, t: []const u8) bool { + const v = o.get("model_type") orelse return false; + return v == .string and std.mem.eql(u8, v.string, t); +} + +fn llmConfig(o: std.json.ObjectMap) ?std.json.ObjectMap { + const v = o.get("llm_config") orelse return null; + return if (v == .object) v.object else null; +} + +/// An eos id, or a list of them, added once each. +fn eos(c: *Config, e: std.json.Value) void { + const items: []const std.json.Value = switch (e) { + .integer => &.{e}, + .array => |a| a.items, + else => &.{}, + }; + for (items) |x| if (x == .integer and !c.isEos(@intCast(x.integer)) and c.eos_count < c.eos.len) { + c.eos[c.eos_count] = @intCast(x.integer); + c.eos_count += 1; + }; +} + +/// nemotron_h_puzzle's block_configs: a block type per layer, each MoE layer with its own experts and width (one top-k). +fn puzzle(c: *Config, blocks: []const std.json.Value, why: *Why) !void { + if (blocks.len > max_layers) { + why.set("config.json has {d} layers; the native Nemotron engine reads up to {d}", .{ blocks.len, max_layers }); + return error.BadConfig; + } + c.experts = 0; + c.expert_width = 0; + for (blocks, 0..) |b, i| { + if (b != .object) return error.BadConfig; + c.kinds[i] = try kindOf((b.object.get("block_type") orelse return error.BadConfig).string); + if (c.kinds[i] != .moe) continue; + const m: Moe = .{ .experts = try int(b.object, "n_routed_experts"), .width = try int(b.object, "moe_intermediate_size") }; + const k = try int(b.object, "num_experts_per_tok"); + if (c.top_k == 0) c.top_k = k; + if (k != c.top_k) { + why.set("layer {d} routes to a different number of experts than the others; the engine reads one top-k", .{i}); + return error.BadConfig; + } + c.moe[i] = m; + c.experts = @max(c.experts, m.experts); + c.expert_width = @max(c.expert_width, m.width); + } + c.layers = blocks.len; +} + /// The shapes our kernels serve (Nemotron 3.5 Lightning 30B-A3B, 4-bit groups of 64): Metal's generated sources, CUDA's captured cubins. pub fn checkShapes(c: Config) !void { const ok = c.hidden == 2688 and c.vocab == 131072 and c.mamba_heads == 64 and c.mamba_head_dim == 64 and @@ -372,3 +441,46 @@ test "parse the Lightning config" { try std.testing.expectEqual(std.math.inf(f32), c.dt_max); try checkShapes(c); } + +test "a Bolt config: omni wrapper, puzzle blocks with their own experts and widths, LatentMoE" { + const bolt = + \\{"model_type": "nemotron_h_omni", "eos_token_id": 11, "llm_config": {"model_type": "nemotron_h_puzzle", + \\ "hidden_size": 4096, "vocab_size": 131072, "mamba_num_heads": 128, "mamba_head_dim": 64, "n_groups": 8, + \\ "ssm_state_size": 128, "conv_kernel": 4, "num_attention_heads": 32, "num_key_value_heads": 2, "head_dim": 128, + \\ "num_experts": 512, "moe_latent_size": 1024, "moe_shared_expert_intermediate_size": 5376, + \\ "routed_scaling_factor": 5.0, "eos_token_id": 2, "time_step_limit": [0.0, {"__float__": "Infinity"}], + \\ "block_configs": [{"block_type": "mamba"}, + \\ {"block_type": "moe", "moe_intermediate_size": 2048, "n_routed_experts": 240, "num_experts_per_tok": 9}, + \\ {"block_type": "attention"}, + \\ {"block_type": "moe", "moe_intermediate_size": 2560, "n_routed_experts": 496, "num_experts_per_tok": 9}]}, + \\ "quantization_config": {"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION"}} + ; + var why: Why = .{}; + const c = try parse(std.testing.allocator, bolt, &why); + try std.testing.expectEqual(@as(usize, 4), c.layers); + try std.testing.expectEqual(Kind.attention, c.kinds[2]); + try std.testing.expectEqual(@as(usize, 9), c.top_k); + try std.testing.expectEqual(@as(usize, 9), c.slots()); + try std.testing.expectEqual(@as(usize, 1024), c.expertDims()); + try std.testing.expectEqual(Moe{ .experts = 240, .width = 2048 }, c.layerMoe(1)); + try std.testing.expectEqual(Moe{ .experts = 496, .width = 2560 }, c.layerMoe(3)); + try std.testing.expectEqual(@as(usize, 496), c.experts); + try std.testing.expectEqual(@as(usize, 2560), c.expert_width); + try std.testing.expectEqual(@as(usize, 18560), c.projDim()); + try std.testing.expect(c.isEos(2) and c.isEos(11)); + try std.testing.expectEqual(std.math.inf(f32), c.dt_max); + try std.testing.expectEqual(Format.modelopt, c.format); + try std.testing.expect(c.omni); +} + +test "the Bolt checkpoint's own config.json, when it is on this machine" { + const io = std.testing.io; + const dir = "/storage/models/nemotron-bolt-74b-24gb"; + std.Io.Dir.cwd().access(io, dir ++ "/config.json", .{}) catch return error.SkipZigTest; + const c = try Config.read(std.testing.allocator, io, dir); + try std.testing.expectEqual(@as(usize, 88), c.layers); + try std.testing.expectEqual(@as(usize, 40), c.count(.moe)); + try std.testing.expectEqual(@as(usize, 8), c.count(.attention)); + try std.testing.expectEqual(@as(usize, 496), c.experts); + try std.testing.expectEqual(Moe{ .experts = 240, .width = 2048 }, c.layerMoe(1)); +} diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index cb2fa934a1..a91bf0b8ca 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -18,6 +18,8 @@ pub const Head = @import("cuda_mtp.zig").Head; pub const Lanes = @import("cuda_lanes.zig").Cuda; pub const native = @import("cuda_native.zig"); pub const train = @import("cuda_train.zig"); +/// Bolt's image tower (C-RADIO) and projector: one image's patches to the rows its `` placeholders take. +pub const vision = @import("cuda_vision.zig"); pub const sites = @import("cuda_sites.zig"); pub const train_ops = @import("cuda_train_ops.zig"); pub const glue = @import("cuda_glue.zig"); @@ -27,6 +29,7 @@ pub const reuse = @import("cuda_reuse.zig"); pub const slide = @import("cuda_slide.zig"); test { + _ = vision; _ = @import("config.zig"); _ = @import("draft_ids.zig"); _ = kernels; diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..0e2b8fbb9e 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -123,6 +123,9 @@ pub const Engine = struct { armed: marks.Armed = .{}, // the marks inside the prompt chunk about to run (cuda_marks.zig) spares: marks.Spares = .{}, // free Mamba buffers the prompt cache set aside for the next marks last_chunk: usize = 0, // where the last prompt pass's last chunk started (the prompt's last row: len - 1 - this) + snap_host: bool = false, // kept prompt states in pinned host memory (cuda_reuse.zig), set when the host opens the engine + features: u64 = 0, // during an image prompt's prefill: the tower's rows, spliced over its placeholders (Bolt) + feature_src: u64 = 0, // and each prompt row's feature row, -1 for text /// Loads the checkpoint into the Python engine's layouts and sizes the caches; `triton_dir` null: our own glue. pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, model_dir: []const u8, triton_dir: ?[]const u8, opts: Options) !*Engine { @@ -151,12 +154,18 @@ pub const Engine = struct { if (opts.slide) return error.SlideNeedsMlx; // the learner reads MLX 4-bit projections try e.k.loadModelopt(ctx); } + // Bolt (LatentMoE, ggml blocks): no MTP head in the checkpoint, and our own glue (captured sets are Lightning's) + const mtp = opts.mtp and e.c.latent == 0; + if (e.c.latent > 0) { + if (e.k.triton != null) return error.CapturedGlueIsLightnings; + try e.k.loadGgml(ctx); + } const chunks = e.max_len / state.chunk_keys; e.nch = if (e.k.triton) |*t| @intCast(t.smallestConst("_chunk", "NCH", @intCast(chunks)) orelse { std.log.err("no captured attention kernel covers {d} chunks of 512 keys; capture one for this context", .{chunks}); return error.ContextUnsupported; }) else chunks; - e.w = try weights.load(gpa, io, e.ops(), model_dir, e.c, opts.mtp); + e.w = try weights.load(gpa, io, e.ops(), model_dir, e.c, mtp); errdefer e.w.deinit(); e.b = try state.Buffers.init(ctx.d, e.c, e.max_len, e.nch, e.carve); errdefer e.b.deinit(); @@ -175,7 +184,7 @@ pub const Engine = struct { e.sampled_ready = try cuda.Event.init(ctx.d, false); try e.copied.record(e.stream); try e.setSampling(opts.sampling); - if (opts.graphs) try e.capture(opts.mtp); + if (opts.graphs) try e.capture(mtp); e.heat_gate = heat.Gate.fromEnv() catch |err| { std.log.err("TF_HEAT_HIGH and TF_HEAT_LOW: both or neither, in degrees, with low at or under high", .{}); return err; @@ -342,6 +351,8 @@ pub const Engine = struct { f.dump = dump; f.slot = e.slot; f.untaught = e.bound.untaught; + f.features = e.features; + f.feature_src = e.feature_src; if (dump == null and e.c.experts <= 128 and e.c.slots() <= 8) f.side = e.side; // tf_plan_routed's bounds return f; } diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index a0e0452e86..65467cb1e1 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -60,6 +60,8 @@ pub const Forward = struct { untaught: bool = false, // the pocket leaves rows (the model as shipped): `segs`' untaught sequences', else all segs: []const Seg = &.{}, // a window's sequences, as windowSegs runs them captures: []const kept.Mark = &.{}, // marks inside the prompt chunk: their Mamba states are copied out as it runs + features: u64 = 0, // an image prompt's tower rows (bf16 [image rows, hidden]), 0 for a text prompt + feature_src: u64 = 0, // and each prompt row's feature row (i32 [prompt], -1 for text) pub fn init(c: Config, w: *const weights.Weights, b: *const state.Buffers, ops: kern.Ops, max_len: usize, nch: usize, sampled: bool) Forward { return .{ @@ -212,7 +214,7 @@ pub const Forward = struct { }, .moe => { try f.experts(blk.moe, rows, false); - delta = .moe; + delta = if (blk.moe.latent != null) .dense else .moe; }, } } @@ -245,8 +247,9 @@ pub const Forward = struct { const o = f.ops; const ex = m.experts; const pairs = rows * c.slots(); + if (m.latent) |lt| return f.latentExperts(m, lt, rows, prompt); if (!prompt) if (f.side) |sd| return f.forked(m, rows, sd); - try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.routed_scaling, c.norm_topk); + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); try f.mark(.route); const tile: usize = if (prompt) cuda.experts.promptTile(ex.format) else 16; try o.plan(b.pick, pairs, ex.experts, tile, b.plan); @@ -267,6 +270,52 @@ pub const Forward = struct { try f.sharedRest(o, m, rows, !prompt); } + /// LatentMoE (Bolt), into b.delta: fc2_latent(sum of the routed experts on fc1_latent(y)) + the shared expert on y. + fn latentExperts(f: *const Forward, m: weights.MoE, lt: weights.Latent, rows: usize, prompt: bool) !void { + const c = f.c; + const b = f.b; + const o = f.ops; + const ex = m.experts; + const K = c.top_k; + const pairs = rows * K; + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, ex.experts, K, K, c.routed_scaling, c.norm_topk); + try f.mark(.route); + if (prompt) try o.prefillDense(b.y, lt.fc1, b.lat, rows) else try o.dense(b.y, b.xs, lt.fc1, b.lat, rows); + const tile: usize = if (prompt) cuda.experts.promptTile(ex.format) else 16; + const g = o.grouped(); + try o.plan(b.pick, pairs, ex.experts, tile, b.plan); + try f.mark(.plan); + var items = kern.maxItems(pairs, ex.experts, tile); + const up_in: cuda.experts.Rows = .{ .x = b.lat, .stride = c.latent, .slots = K }; + const down_in: cuda.experts.Rows = .{ .x = b.act, .stride = ex.width }; + if (prompt) { + try g.prompt(o.s, .up, up_in, ex, b.plan, b.act, items, -1); + try g.prompt(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1); + } else { + try g.decode(o.s, .up, up_in, ex, b.plan, b.act, items, -1); + try f.mark(.experts_up); + try g.decode(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1); + try f.mark(.experts_down); + } + try o.latentSum(b.ymoe, b.wts, b.z, rows, K, c.latent); + if (prompt) try o.prefillDense(b.z, lt.fc2, b.delta, rows) else try o.dense(b.z, 0, lt.fc2, b.delta, rows); + // the shared expert: every row's pair with expert 0 of a one-expert layer, relu^2 between as the routed ones + const sh = lt.shared; + try o.plan(b.zeros, rows, 1, tile, b.plan); + items = kern.maxItems(rows, 1, tile); + const sh_up: cuda.experts.Rows = .{ .x = b.y, .stride = c.hidden, .slots = 1 }; + const sh_down: cuda.experts.Rows = .{ .x = b.act, .stride = sh.width }; + if (prompt) { + try g.prompt(o.s, .up, sh_up, sh, b.plan, b.act, items, -1); + try g.prompt(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1); + } else { + try g.decode(o.s, .up, sh_up, sh, b.plan, b.act, items, -1); + try g.decode(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1); + } + try o.addBf16(b.delta, b.delta, b.ymoe, true, rows * c.hidden); + try f.mark(.shared_wait); + } + /// The pocket over the rows of sequences that read what was learned: every row, or the runs between untaught ones'. fn pocket(f: *const Forward, o: kern.Ops, p: weights.Pocket, d: usize, y_f32: bool, rows: usize) !void { if (!f.untaught) return f.pocketRun(o, p, d, y_f32, .{ 0, rows }); @@ -312,7 +361,7 @@ pub const Forward = struct { try g.decode(sd.s, .down_f32, in[1], ex, sp, b.ymoe, 2, -1); try f.sharedRest(so, m, rows, true); try sd.join.record(sd.s); - try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.routed_scaling, c.norm_topk); + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); try f.mark(.route); try o.planRouted(b.pick, rows, c.slots(), c.top_k, c.experts, 16, b.plan); try f.mark(.plan); @@ -341,6 +390,7 @@ pub const Forward = struct { pub fn chunkBegin(f: *const Forward, w: *Walk) !void { const b = f.b; try f.embed(b.p_ids, b.emb, w.rows); + if (f.features != 0) try f.ops.spliceRows(b.emb, f.features, f.feature_src + @as(u64, w.pos) * 4, w.rows, f.c.hidden); w.x = b.emb; } @@ -427,7 +477,7 @@ pub const Forward = struct { }, .moe => { try f.experts(blk.moe, w.rows, true); - w.delta = .moe; + w.delta = if (blk.moe.latent != null) .dense else .moe; }, } } diff --git a/zig/src/families/nemotron/cuda_glue.zig b/zig/src/families/nemotron/cuda_glue.zig index 7d3f0c5ef5..2bf49b46a9 100644 --- a/zig/src/families/nemotron/cuda_glue.zig +++ b/zig/src/families/nemotron/cuda_glue.zig @@ -87,11 +87,20 @@ pub const AttnSeg = struct { kc: u64, vc: u64, meta: u64, row0: usize, rows: usi const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, row0: [max_segs]c_int, n: c_int }; /// The router's K slices and the shared bytes its block stages: 16 rows and 16 experts of one slice, padded a word. +/// Wide rows (Bolt's 4096) take more slices, so a block stays within the 48 KiB launches get without an opt-in. pub fn routerShape(d: usize) struct { sk: usize, shared: u32 } { - const sk = triton.divisor(d / 64, 6); - return .{ .sk = sk, .shared = @intCast(2 * 16 * (d / sk / 2 + 1) * 4) }; + var sk = triton.divisor(d / 64, 6); + while (staged(d, sk) > 48 << 10 and (d / 64) % (sk * 2) == 0) sk *= 2; + return .{ .sk = sk, .shared = @intCast(staged(d, sk)) }; } +fn staged(d: usize, sk: usize) usize { + return 2 * 16 * (d / sk / 2 + 1) * 4; +} + +/// The most K slices routerShape takes for a row of `d`: the router's scratch is that many [rows, E] fp32 sums. +pub const max_router_slices = 8; + /// Every launch of a window or prompt chunk's glue, on one stream. pub const Glue = struct { set: ?*const cuda.aot.Set, // a captured Triton set (GB10's qualified one), else our own kernels @@ -151,10 +160,10 @@ pub const Glue = struct { try g.go(g.f.add_moe_norm, .{ rows, 1, 1 }, 256, 0, &a); } - /// fp32 router logits over K slices, then each row's top-k and the shared halves' slots. - pub fn route(g: Glue, x: u64, w: u64, bias: u64, part: u64, ids: u64, wts: u64, rows: usize, d: usize, e: usize, top_k: usize, scaling: f32, norm: bool) !void { + /// fp32 router logits over K slices, then each row's top-k and the shared halves' slots (`slots` past top-k). + pub fn route(g: Glue, x: u64, w: u64, bias: u64, part: u64, ids: u64, wts: u64, rows: usize, d: usize, e: usize, top_k: usize, slots: usize, scaling: f32, norm: bool) !void { if (g.tri()) |t| return t.route(x, w, bias, part, ids, wts, rows, d, e, top_k, scaling, norm); - if (e > 256 or top_k + 2 > 8) return error.RouterTooWide; + if (e > 512 or top_k > 16 or slots < top_k or slots > top_k + 2) return error.RouterTooWide; // the top-k kernel's MAXC const shape = routerShape(d); var a: cuda.Args = .{}; ptrs(&a, &.{ x, w, part }); @@ -164,7 +173,7 @@ pub const Glue = struct { ptrs(&b, &.{ part, bias, ids, wts }); b.add(int(rows)); b.add(scaling); - ints(&b, &.{ e, shape.sk, top_k, top_k + 2 }); + ints(&b, &.{ e, shape.sk, top_k, slots }); b.add(@as(c_int, @intFromBool(norm))); try g.go(g.f.topk, .{ rows, 1, 1 }, 32, 0, &b); } @@ -353,4 +362,8 @@ test "the router stages one slice of 16 rows and 16 experts" { const s = routerShape(2688); try std.testing.expectEqual(@as(usize, 6), s.sk); try std.testing.expectEqual(@as(u32, 2 * 16 * 225 * 4), s.shared); + // Bolt's 4096-wide rows: four slices would stage 66 KiB, eight stay under 48 + const b = routerShape(4096); + try std.testing.expectEqual(@as(usize, 8), b.sk); + try std.testing.expect(b.shared <= 48 << 10 and b.sk <= max_router_slices); } diff --git a/zig/src/families/nemotron/cuda_kernels.zig b/zig/src/families/nemotron/cuda_kernels.zig index 7932f1c437..77e041bcee 100644 --- a/zig/src/families/nemotron/cuda_kernels.zig +++ b/zig/src/families/nemotron/cuda_kernels.zig @@ -61,11 +61,15 @@ pub const Kernels = struct { pocket_in: [2]cuda.Function, // the pocket's two halves a block row a row, then tiled for a prompt chunk's rows pocket_out: [2]cuda.Function, add_row: cuda.Function, + latent_sum: cuda.Function, // LatentMoE: a row's routed latent outputs weighted and summed (Bolt) + add_bf16: cuda.Function, + splice_rows: cuda.Function, // an image prompt's rows take the vision tower's features (Bolt) draws: cuda.sample.Draws, // sample.cu: the keyed draws and logprob rows grammar_mask: cuda.Function, // structured output: disallowed tokens -inf before the draw torch: torch_ops.Functions, train: train_ops.Kernels, // Sliding Weights' learning step, and the learned change a forward adds while lessons run modelopt: ?Modelopt = null, // loaded by loadModelopt for a ModelOpt checkpoint, else never + ggml: ?cuda.ggml.Kernels = null, // loaded by loadGgml for a checkpoint with ggml blocks (Bolt's IQ2_XXS and Q8_0) gb10: bool, discrete: bool, // the card has its own memory: checkpoint bytes reach it through page-locked slots @@ -93,6 +97,9 @@ pub const Kernels = struct { k.pocket_in = .{ try k.mods[7].function("tf_pocket_in"), try k.mods[7].function("tf_pocket_in_rows") }; k.pocket_out = .{ try k.mods[7].function("tf_pocket_out"), try k.mods[7].function("tf_pocket_out_rows") }; k.add_row = try k.mods[7].function("tf_add_row"); + k.latent_sum = try k.mods[7].function("tf_latent_sum"); + k.add_bf16 = try k.mods[7].function("tf_add_bf16"); + k.splice_rows = try k.mods[7].function("tf_splice_rows"); k.torch = try torch_ops.Functions.resolve(k.mods[8..14]); k.draws = try cuda.sample.Draws.resolve(k.mods[14]); k.glue = try glue.Fns.resolve(k.mods[16..21]); @@ -102,6 +109,7 @@ pub const Kernels = struct { errdefer k.norms.unload(); k.grammar_mask = try k.mods[14].function("tf_grammar_mask"); k.modelopt = null; + k.ggml = null; k.triton = if (triton_dir) |dir| try cuda.aot.Set.load(gpa, io, d, ctx.device, dir) else null; errdefer if (k.triton) |*t| t.deinit(); k.train = try train_ops.Kernels.load(d); @@ -126,6 +134,7 @@ pub const Kernels = struct { m.experts.unload(); m.pack.unload(); } + if (k.ggml) |*g| g.unload(); k.affine.deinit(); k.train.deinit(); if (k.triton) |*t| t.deinit(); @@ -146,6 +155,13 @@ pub const Kernels = struct { errdefer experts.unload(); k.modelopt = .{ .lane = lane, .gemm = gemm, .experts = experts, .pack = try cuda.modelopt.Packer.load(k.d) }; } + + /// The ggml block kernels (IQ2_XXS experts, IQ2_XXS and Q8_0 projections, their repack) Bolt's export needs. + pub fn loadGgml(k: *Kernels, ctx: *const cuda.Context) !void { + if (k.ggml != null) return; + _ = ctx; + k.ggml = try cuda.ggml.Kernels.load(k.d); + } }; /// qmm.split_k (cuda/qlinear.zig). @@ -204,12 +220,44 @@ pub const Ops = struct { /// Every projection format the kernels loaded (cuda/qlinear.zig). pub fn linear(o: Ops) cuda.qlinear.Linear { const m = if (o.k.modelopt) |*x| x else null; - return .{ .affine4 = &o.k.affine, .lane = if (m) |x| &x.lane else null, .nvfp4_prompt = if (m) |x| &x.gemm else null }; + return .{ .affine4 = &o.k.affine, .lane = if (m) |x| &x.lane else null, .nvfp4_prompt = if (m) |x| &x.gemm else null, .ggml = if (o.k.ggml) |*g| g else null }; } /// Every grouped-expert format the kernels loaded (cuda/experts.zig). pub fn grouped(o: Ops) cuda.experts.Grouped { - return .{ .affine4 = &o.k.experts, .nvfp4 = if (o.k.modelopt) |*x| &x.experts else null }; + return .{ .affine4 = &o.k.experts, .nvfp4 = if (o.k.modelopt) |*x| &x.experts else null, .iq2xxs = if (o.k.ggml) |*g| g else null }; + } + + /// LatentMoE: z[r] = sum_k wt[r, k] y[r * k_slots + k], rows of `l` bf16 from fp32 pair rows (tf_latent_sum). + pub fn latentSum(o: Ops, y: u64, wt: u64, z: u64, rows: usize, k_slots: usize, l: usize) !void { + var a: cuda.Args = .{}; + a.add(y); + a.add(wt); + a.add(z); + a.add(int(k_slots)); + a.add(int(l)); + try o.go(o.k.latent_sum, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// Rows of h whose src (i32, one a row) names a feature row take that row of `features` ([*, d] bf16). + pub fn spliceRows(o: Ops, h: u64, features: u64, src: u64, rows: usize, d: usize) !void { + var a: cuda.Args = .{}; + a.add(h); + a.add(features); + a.add(src); + a.add(int(d)); + try o.go(o.k.splice_rows, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// out = a + b, n bf16 values (b bf16, or fp32 when `b_f32`). + pub fn addBf16(o: Ops, out: u64, a_: u64, b: u64, b_f32: bool, n: usize) !void { + var a: cuda.Args = .{}; + a.add(out); + a.add(a_); + a.add(b); + a.add(@as(c_int, @intFromBool(b_f32))); + a.add(int(n)); + try o.go(o.k.add_bf16, .{ (n + 255) / 256, 1, 1 }, 256, 0, &a); } /// Decode rows of a projection on this stream (the affine-4 paths' bits unchanged); xs: x's 64-group sums. diff --git a/zig/src/families/nemotron/cuda_lanes.zig b/zig/src/families/nemotron/cuda_lanes.zig index 203ccd74cb..baf7e98a4d 100644 --- a/zig/src/families/nemotron/cuda_lanes.zig +++ b/zig/src/families/nemotron/cuda_lanes.zig @@ -1,6 +1,7 @@ //! Nemotron CUDA streams own their sequences and share window rounds. const std = @import("std"); +const vision = @import("cuda_vision.zig"); const cuda = @import("cuda"); const lanes = @import("lanes"); const core = @import("core"); @@ -48,6 +49,7 @@ pub const Cuda = struct { measured: ?core.draft_depth.Costs = null, // a lone stream's depth rule prices its rounds by these bits: ?cuda.HostBuffer = null, // structured output's row bits, device-mapped (row_slots rows), made at first use bits_dev: u64 = 0, + tower: ?*vision.Tower = null, // --vision (Bolt): an image prompt's images are encoded at its prefill pub fn init(gpa: std.mem.Allocator, e: *Engine, head: ?*Head) !Cuda { var pinned = try cuda.HostBuffer.alloc(e.ctx.d, engine.max_streams * state.max_rows * 4); @@ -63,6 +65,7 @@ pub const Cuda = struct { var it = self.lanes.valueIterator(); while (it.next()) |l| if (!l.own) self.e.freeSeq(l.seq); self.lanes.deinit(self.gpa); + if (self.tower) |t| t.deinit(); self.pinned.free(); self.rows_out.free(); if (self.bits) |*b| b.free(); @@ -254,8 +257,22 @@ pub const Cuda = struct { e.decide = if (s.decide != null and e.lessons() != null) s.decide.? else lesson_route.none; try e.setSampling(s.sampling); s.cached = 0; - const from = if (s.prompt_logprobs != null) 0 else self.restoreKept(s, ids.len); // prompt rows need every row - const keep: ?engine.Keep = if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; + // an image prompt: its images through the tower, their rows spliced over the placeholders as each chunk embeds; + // never resumed or kept (every image's placeholders are the same ids) + var media: ?vision.Media = null; + if (s.images.len > 0) { + const tower = self.tower orelse return error.NemotronServedWithoutVision; + media = try vision.Media.init(self.gpa, tower, e.stream, ids, s.images); + e.features = media.?.features.ptr; + e.feature_src = media.?.src.ptr; + } + defer if (media) |*m| { + e.features = 0; + e.feature_src = 0; + m.deinit(); + }; + const from = if (s.prompt_logprobs != null or media != null) 0 else self.restoreKept(s, ids.len); // prompt rows need every row + const keep: ?engine.Keep = if (s.reuse.marks.len == 0 or media != null) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; var sink: PromptSink = .{ .gpa = self.gpa, .s = s }; s.prompt_rows.clearRetainingCapacity(); if (s.prompt_logprobs) |k| { diff --git a/zig/src/families/nemotron/cuda_modelopt.zig b/zig/src/families/nemotron/cuda_modelopt.zig index 546b861b39..63d369560e 100644 --- a/zig/src/families/nemotron/cuda_modelopt.zig +++ b/zig/src/families/nemotron/cuda_modelopt.zig @@ -13,9 +13,24 @@ const Loader = W.Loader; const Tensor = core.checkpoint.Tensor; const Nvfp4Experts = cuda.experts.Nvfp4Experts; -const Kind = enum { nvfp4, fp8, bf16 }; +const Kind = enum { nvfp4, fp8, bf16, q8_0, iq2xxs }; + +/// Bytes and weights of one ggml block: Q8_0 34 per 32, IQ2_XXS 66 per 256. +fn ggmlBlock(k: Kind) [2]usize { + return switch (k) { + .q8_0 => .{ 34, 32 }, + .iq2xxs => .{ 66, 256 }, + else => unreachable, + }; +} fn kindOf(t: Tensor) !Kind { + // ggml blocks (Bolt) keep each row's blocks as a third axis: [rows, blocks, block bytes]; NVFP4 packs pairs as [rows, k / 2] + if (t.dtype == .u8 and t.rank == 3) return switch (t.dim(2)) { + 34 => .q8_0, + 66 => .iq2xxs, + else => error.UnsupportedQuantization, + }; return switch (t.dtype) { .u8 => .nvfp4, .f8_e4m3 => .fp8, @@ -42,16 +57,22 @@ fn scalar(ck: *core.Checkpoint, n: []const u8) !f32 { fn part(ck: *core.Checkpoint, pre: []const u8) !Part { var buf: [192]u8 = undefined; const w = try ck.get(try name(&buf, "{s}.weight", .{pre})); - if (w.rank != 2) return error.UnexpectedTensor; + if (w.rank != 2 and w.rank != 3) return error.UnexpectedTensor; // we keep bf16 activations and a bf16 KV cache, more precise than the W8A8 and FP8-cache math these scales are for for ([_][]const u8{ "input_scale", "k_scale", "v_scale" }) |extra| { const n = try name(&buf, "{s}.{s}", .{ pre, extra }); if (ck.has(n)) _ = try ck.get(n); } const kind = try kindOf(w); - var p: Part = .{ .w = w, .kind = kind, .n = w.dim(0), .k = if (kind == .nvfp4) w.dim(1) * 2 else w.dim(1) }; + if ((w.rank == 3) != (kind == .q8_0 or kind == .iq2xxs)) return error.UnexpectedTensor; + const k = switch (kind) { + .nvfp4 => w.dim(1) * 2, + .q8_0, .iq2xxs => w.dim(1) * ggmlBlock(kind)[1], + else => w.dim(1), + }; + var p: Part = .{ .w = w, .kind = kind, .n = w.dim(0), .k = k }; switch (kind) { - .bf16 => {}, + .bf16, .q8_0, .iq2xxs => {}, .fp8 => p.scale = try scalar(ck, try name(&buf, "{s}.weight_scale", .{pre})), .nvfp4 => { p.blocks = try ck.expect(try name(&buf, "{s}.weight_scale", .{pre}), .f8_e4m3, &.{ p.n, p.k / 16 }); @@ -81,6 +102,7 @@ pub fn dense(L: *Loader, ck: *core.Checkpoint, label: []const u8, prefixes: []co var buf: [128]u8 = undefined; const s = L.ops.s; switch (first.kind) { + .q8_0, .iq2xxs => return ggmlDense(L, label, parts[0..prefixes.len], n, k), .bf16, .fp8 => { const size: usize = if (first.kind == .bf16) 2 else 1; const base = try L.tmp(n * k * size); @@ -115,9 +137,97 @@ pub fn dense(L: *Loader, ck: *core.Checkpoint, label: []const u8, prefixes: []co } } +/// Row-stacked ggml parts (q, k, v: rows after rows) uploaded as stored, then tiled on the GPU (a byte permutation). +fn ggmlDense(L: *Loader, label: []const u8, parts: []const Part, n: usize, k: usize) !kern.QLinear { + const g = try ggmlKernels(L); + const kind = parts[0].kind; + const base = try L.tmp(n * k / ggmlBlock(kind)[1] * ggmlBlock(kind)[0]); + var at: usize = 0; + for (parts) |p| { + try L.src.upload(base + at, p.w.bytes); + at += p.w.bytes.len; + } + var buf: [128]u8 = undefined; + const gk: cuda.ggml.Kind = if (kind == .q8_0) .q8_0 else .iq2xxs; + const out = try L.alloc(try name(&buf, "{s}.weight", .{label}), cuda.ggml.bytes(gk, n, k)); + try L.src.flush(); + try g.pack(L.ops.s, gk, base, out, n, k, 1); + return if (kind == .q8_0) .{ .q8_0 = cuda.ggml.q8Weight(out, n, k) } else .{ .iq2xxs = cuda.ggml.iq2Weight(out, n, k) }; +} + +fn ggmlKernels(L: *Loader) !*const cuda.ggml.Kernels { + return if (L.ops.k.ggml) |*g| g else error.GgmlNotLoaded; +} + +/// IQ2_XXS grouped experts from `count` tensors named by `names`, [n, k] each, uploaded back to back and tiled. +fn iq2Experts(L: *Loader, ck: *core.Checkpoint, label: []const u8, count: usize, n: usize, k: usize, names: anytype) !u64 { + const g = try ggmlKernels(L); + const one = n * k / 256 * 66; + const base = try L.tmp(count * one); + var b: [192]u8 = undefined; + for (0..count) |i| { + const p = try part(ck, try names.at(&b, i)); + if (p.kind != .iq2xxs or p.n != n or p.k != k) { + std.log.err("{s}: {t} [{d}, {d}], expected IQ2_XXS [{d}, {d}]", .{ try names.at(&b, i), p.kind, p.n, p.k, n, k }); + return error.UnexpectedTensor; + } + try L.src.upload(base + i * one, p.w.bytes); + } + const out = try L.alloc(label, cuda.ggml.bytes(.iq2xxs, n, k) * count); + try L.src.flush(); + try g.pack(L.ops.s, .iq2xxs, base, out, n, k, count); + return out; +} + +/// A LatentMoE layer (Bolt): the router and bias over this layer's experts, fc1/fc2_latent_proj, routed IQ2_XXS +/// experts in the latent width, and the shared expert at full width as a one-expert layer of its own. +pub fn latent(L: *Loader, ck: *core.Checkpoint, label: []const u8, pre: []const u8, c: Config, layer: usize) !W.MoE { + const lm = c.layerMoe(layer); + const E = lm.experts; + const Lw = c.latent; + var a: [128]u8 = undefined; + var b: [192]u8 = undefined; + const r = try routerOf(L, try name(&a, "{s}.router", .{label}), try ck.get(try name(&b, "{s}gate.weight", .{pre})), E, c.hidden); + const bias = try ck.expect(try name(&b, "{s}gate.e_score_correction_bias", .{pre}), .f32, &.{E}); + var m: W.MoE = .{ .router = r, .bias = try L.raw(try name(&a, "{s}.bias", .{label}), bias), .experts = undefined }; + var ex: kern.Experts = .{ .format = .iq2xxs, .up = 0, .down = 0, .width = lm.width, .dims = Lw, .experts = E, .relu2 = true }; + const Names = struct { + pre: []const u8, + leaf: []const u8, + fn at(x: @This(), buf: []u8, i: usize) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}experts.{d}.{s}", .{ x.pre, i, x.leaf }); + } + }; + ex.up = try iq2Experts(L, ck, try name(&a, "{s}.experts.up", .{label}), E, lm.width, Lw, Names{ .pre = pre, .leaf = "up_proj" }); + ex.down = try iq2Experts(L, ck, try name(&a, "{s}.experts.down", .{label}), E, Lw, lm.width, Names{ .pre = pre, .leaf = "down_proj" }); + m.experts = ex; + const Shared = struct { + pre: []const u8, + leaf: []const u8, + fn at(x: @This(), buf: []u8, _: usize) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}shared_experts.{s}", .{ x.pre, x.leaf }); + } + }; + var sh: kern.Experts = .{ .format = .iq2xxs, .up = 0, .down = 0, .width = c.shared_width, .dims = c.hidden, .experts = 1, .relu2 = true }; + sh.up = try iq2Experts(L, ck, try name(&a, "{s}.shared.up", .{label}), 1, c.shared_width, c.hidden, Shared{ .pre = pre, .leaf = "up_proj" }); + sh.down = try iq2Experts(L, ck, try name(&a, "{s}.shared.down", .{label}), 1, c.hidden, c.shared_width, Shared{ .pre = pre, .leaf = "down_proj" }); + var p1: [192]u8 = undefined; + var p2: [192]u8 = undefined; + m.latent = .{ + .fc1 = try dense(L, ck, try name(&a, "{s}.fc1_latent", .{label}), &.{try name(&p1, "{s}fc1_latent_proj", .{pre})}), + .fc2 = try dense(L, ck, try name(&b, "{s}.fc2_latent", .{label}), &.{try name(&p2, "{s}fc2_latent_proj", .{pre})}), + .shared = sh, + }; + return m; +} + /// The router as the kernels read it: bf16 [E, D]; an fp32 router is taken only when every value is a bf16. fn router(L: *Loader, label: []const u8, t: Tensor, c: Config) !u64 { - if (t.rank != 2 or t.dim(0) != c.experts or t.dim(1) != c.hidden) return error.UnexpectedTensor; + return routerOf(L, label, t, c.experts, c.hidden); +} + +fn routerOf(L: *Loader, label: []const u8, t: Tensor, experts_: usize, hidden: usize) !u64 { + if (t.rank != 2 or t.dim(0) != experts_ or t.dim(1) != hidden) return error.UnexpectedTensor; if (t.dtype == .bf16) return L.raw(label, t); if (t.dtype != .f32) return error.UnexpectedTensor; const raw = try L.gpa.alloc(u8, t.bytes.len); diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 64594342c4..4469f9151e 100644 --- a/zig/src/families/nemotron/cuda_native.zig +++ b/zig/src/families/nemotron/cuda_native.zig @@ -20,12 +20,18 @@ const core = @import("core"); pub const slide = @import("cuda_slide.zig"); pub const model_type = "nemotron_h"; -pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-nvfp4" }; +/// Nemotron 3.5 Bolt's config: the language model under llm_config, beside its vision tower's. +pub const also_model_types: []const []const u8 = &.{"nemotron_h_omni"}; +pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-nvfp4", "modelopt-ggml" }; pub const default_context: i64 = engine.default_context; pub const max_segments: u32 = @import("cuda_segments.zig").MAX; pub const prompt_rows: u32 = state.pass_rows; -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 }; + +/// --vision on a Bolt checkpoint: its C-RADIO tower, images prepared and placed as qwen_image.nemotron lays them. +pub const vision_nemotron = true; +const vision = @import("cuda_vision.zig"); /// Thinking budgets the lone driver takes (cuda_lone.zig): above it, the stream goes to the lane core before its cut. pub const lone_min_budget = @import("cuda_lone.zig").min_budget; @@ -51,6 +57,7 @@ pub const Loaded = struct { target: *reuse.Target, // the prompt cache copies this engine (cuda_reuse.zig) learns: ?*engine.Engine = null, // --slide: the engine its learner trains, whose sites it attached at load learn_bytes: usize = 0, // --slide: what its trainer allocates at the first lesson, kept out of the streams' room + host_snaps: bool = false, // kept prompt states in pinned host memory (a discrete card): the cache's budget is the host's }; /// The first shard, index or config.json in `dir` that a write in place would reach through a link, or null. @@ -71,19 +78,27 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer if (carve) |c| closeCarveout(gpa, c); const e = try engine.Engine.init(gpa, io, ctx, dir, kernels, .{ .context = o.context, .mtp = o.drafts, .graphs = true, .sampling = null, .segments = o.segments, .carveout = carve, .slide = o.slide }); errdefer e.deinit(); - const head: ?*Head = if (o.drafts) try Head.init(e) else null; + // Bolt's checkpoint has no MTP head: it decodes serially whatever --drafts says + const drafts = o.drafts and e.c.latent == 0; + const head: ?*Head = if (drafts) try Head.init(e) else null; errdefer if (head) |h| h.deinit(); if (head) |h| try h.capture(); - if (!o.drafts) try e.captureWindows(); // the one-row window: Engine.init captures windows only when drafting + if (!drafts) try e.captureWindows(); // the one-row window: Engine.init captures windows only when drafting // the sampled mode's sets: a placeholder rule now, each stream's own uploaded at its prefill try e.setSampling(.{ .seed = 0, .temperature = 1.0 }); try e.captureWindows(); if (head) |h| try h.capture(); try e.setSampling(null); + e.snap_host = e.k.discrete; // a card with its own memory keeps prompt states in the host's (GB10 shares one) const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.* = .{ .gpa = gpa, .e = e, .head = head, .lanes = try Lanes.init(gpa, e, head), .target = .{ .e = e, .head = head }, .carve = carve }; errdefer own.lanes.deinit(); + // --vision: the tower and its scratch for the largest image, loaded beside the model (Bolt alone has one) + if (o.vision) { + if (!e.c.omni) return error.NemotronCheckpointHasNoVision; + own.lanes.tower = try vision.Tower.load(gpa, io, ctx.d, dir); + } own.living = try slide.loadLiving(gpa, io, e, dir); // before the lanes' costs are measured: they decode through it errdefer if (own.living) |*s| s.deinit(); try own.lanes.measure(io, dir); @@ -98,6 +113,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .lone = if (head != null) loneRun else null, .learns = if (o.slide) e else null, .learn_bytes = if (o.slide) @import("cuda_train.zig").Trainer.memory(e.c) else 0, + .host_snaps = e.snap_host, }; } @@ -117,6 +133,9 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.NemotronServedWithoutVision => "image inputs need the server started with --vision", + 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", else => null, }; } diff --git a/zig/src/families/nemotron/cuda_prompt.zig b/zig/src/families/nemotron/cuda_prompt.zig index ba5408db8d..7765bab806 100644 --- a/zig/src/families/nemotron/cuda_prompt.zig +++ b/zig/src/families/nemotron/cuda_prompt.zig @@ -39,7 +39,7 @@ pub fn prefillFrom(e: *Engine, prompt: []const u32, from: usize, dump: ?*Dump, h if (head) |h| try h.reset(); } else if (e.pos != from) return error.BadReuse; const waited = e.heat_gate.waited_s; - if (from == 0 and keep == null and e.segments > 1 and dump == null and e.prompt_rows == null and prompt.len > state.prefill_rows and e.decide == lesson_route.none and !Cancel.background(cancel)) { + if (from == 0 and keep == null and e.features == 0 and e.segments > 1 and dump == null and e.prompt_rows == null and prompt.len > state.prefill_rows and e.decide == lesson_route.none and !Cancel.background(cancel)) { try segs.prefill(e, try e.segmentSet(), prompt, head, e.segments, cancel); } else try e.serialChunks(prompt, from, dump, head, cancel, keep); if (e.heat_gate.waited_s > waited) heat.note(e.heat_gate.waited_s - waited); diff --git a/zig/src/families/nemotron/cuda_reuse.zig b/zig/src/families/nemotron/cuda_reuse.zig index e73851a23c..55084e774f 100644 --- a/zig/src/families/nemotron/cuda_reuse.zig +++ b/zig/src/families/nemotron/cuda_reuse.zig @@ -1,4 +1,5 @@ //! Nemotron CUDA prompt states for the shared cache: a prompt's key and value rows, the Mamba states, the head's rows. +//! On the device, or (a discrete card's `host`) in pinned host memory: the key and value rows, then the Mamba states. const std = @import("std"); const cuda = @import("cuda"); @@ -10,9 +11,11 @@ const Head = @import("cuda_mtp.zig").Head; const marks = @import("cuda_marks.zig"); /// One kept state: its Mamba states in one device buffer, its key and value rows in another, where the engine stood. +/// A host state keeps both in `host` (key and value rows first); `mamba` is then an armed row's staging until `finish`. pub const Copy = struct { mamba: cuda.DeviceBuffer, kv: cuda.DeviceBuffer, + host: ?cuda.HostBuffer = null, head: bool, pos: usize, parity: usize, @@ -24,6 +27,11 @@ pub const Copy = struct { pub const Target = struct { e: *Engine, head: ?*Head, + + /// Kept states in pinned host memory (Engine.snap_host, a discrete card): the budget is the host's. + fn host(t: *const Target) bool { + return t.e.snap_host; + } }; /// `len` bytes at `live` in the bound sequence's planes; `mamba`: the SSM or conv state, which an armed chunk writes. @@ -94,9 +102,28 @@ fn move(t: *const Target, at: usize, copy: *const Copy, into: bool, mamba: bool) return off[0] + off[1]; } -/// Bytes a save at `at` takes, as the budget counts them: what `save` allocates, spare or new. +/// Copies between the bound sequence and a host copy (`into`: live to host), key and value rows first; `mamba`: the +/// Mamba states too, else from `staging` (an armed row's) when given. Queued on the engine's stream. +fn hostMove(t: *const Target, at: usize, host: []u8, into: bool, mamba: bool, staging: ?u64) !void { + const o = t.e.ops(); + var store: [max_parts]Part = undefined; + var off = [2]usize{ 0, kvBytes(t.e.c, at, t.head != null) }; + for (parts(t.e.c, t.e.max_len, planes(t), at, &store)) |p| { + const at_host = off[@intFromBool(p.mamba)]; + const buf = host[at_host..][0..p.len]; + if (!p.mamba) { + if (into) try o.download(buf, p.live) else try o.upload(p.live, buf); + } else if (mamba) { + if (into) try o.download(buf, p.live) else try o.upload(p.live, buf); + } else if (staging) |st| try o.download(buf, st + (at_host - kvBytes(t.e.c, at, t.head != null))); // SSM, then conv + off[@intFromBool(p.mamba)] += p.len; + } +} + +/// Bytes a save at `at` takes, as the budget counts them: what `save` allocates, spare or new (host: its bytes). pub fn bytes(ptr: *anyopaque, at: u32) u64 { const t: *Target = @ptrCast(@alignCast(ptr)); + if (t.host()) return stateBytes(t.e.c, at, t.head != null); return mambaBytes(t.e.c) + allocBytes(kvBytes(t.e.c, at, t.head != null)); } @@ -104,9 +131,10 @@ pub fn bytes(ptr: *anyopaque, at: u32) u64 { pub fn ready(t: *Target, budget: u64) !u64 { const e = t.e; const len = mambaBytes(e.c); - e.spares.want = @min(marks.spares_max, budget / 4 / len); + // host states: one device staging buffer for an armed row, out of the device's room rather than the host budget + e.spares.want = if (t.host()) @min(marks.spares_max, 1) else @min(marks.spares_max, budget / 4 / len); e.spares.fill(e.ctx.d, len); - return e.spares.want * len; + return if (t.host()) 0 else e.spares.want * len; } /// Copy the live state after `at` prompt tokens. The pass is standing there. @@ -119,6 +147,15 @@ pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { if (e.armed.open) |span| if (at > span[0] and at < span[1]) return arm(t, at); if (e.pos != at) return error.NotAtMark; if (t.head) |h| if (h.pos != at) return error.NotAtMark; + if (t.host()) { + const copy = try e.gpa.create(Copy); + errdefer e.gpa.destroy(copy); + copy.* = .{ .mamba = .{ .ptr = 0, .len = 0, .d = e.ctx.d }, .kv = .{ .ptr = 0, .len = 0, .d = e.ctx.d }, .head = t.head != null, .pos = at, .parity = e.parity, .rows = e.rows }; + copy.host = try cuda.HostBuffer.alloc(e.ctx.d, stateBytes(e.c, at, copy.head)); + errdefer copy.host.?.free(); + try hostMove(t, at, copy.host.?.bytes, true, true, null); + return copy; + } const copy = try alloc(t, at, true); errdefer drop(ptr, copy); if (try move(t, at, copy, true, true) != stateBytes(e.c, at, t.head != null)) return error.SnapshotSize; @@ -132,7 +169,7 @@ fn alloc(t: *const Target, at: usize, kv: bool) !*Copy { errdefer e.gpa.destroy(copy); var m = e.spares.take() orelse try cuda.DeviceBuffer.alloc(e.ctx.d, mambaBytes(e.c)); errdefer if (!e.spares.give(m)) m.free(); - const rows = try cuda.DeviceBuffer.alloc(e.ctx.d, if (kv) allocBytes(kvBytes(e.c, at, t.head != null)) else 0); + const rows = try cuda.DeviceBuffer.alloc(e.ctx.d, if (kv and !t.host()) allocBytes(kvBytes(e.c, at, t.head != null)) else 0); copy.* = .{ .mamba = m, .kv = rows, .head = t.head != null, .pos = at, .parity = e.parity, .rows = e.rows }; return copy; } @@ -153,6 +190,20 @@ pub fn finish(t: Target) !void { const e = t.e; for (e.armed.marks()) |m| { const copy: *Copy = @ptrCast(@alignCast(m.copy)); + if (t.host()) { + // the chunk wrote this row's Mamba states into the staging buffer: rows and states to the host, then the + // staging back (the stream orders its next writer after these copies) + var h = cuda.HostBuffer.alloc(e.ctx.d, stateBytes(e.c, m.at, copy.head)) catch continue; // stays not whole + hostMove(&t, m.at, h.bytes, true, false, copy.mamba.ptr) catch { + h.free(); + continue; + }; + copy.host = h; + if (!e.spares.give(copy.mamba)) copy.mamba.free(); + copy.mamba = .{ .ptr = 0, .len = 0, .d = e.ctx.d }; + copy.whole = true; + continue; + } copy.kv = cuda.DeviceBuffer.alloc(e.ctx.d, allocBytes(kvBytes(e.c, m.at, copy.head))) catch continue; // stays not whole _ = try move(&t, m.at, copy, true, false); copy.whole = true; @@ -163,12 +214,17 @@ pub fn finish(t: Target) !void { /// What a kept state holds on the device, both buffers. pub fn held(saved: *anyopaque) usize { const copy: *Copy = @ptrCast(@alignCast(saved)); - return copy.mamba.len + copy.kv.len; + return copy.mamba.len + copy.kv.len + if (copy.host) |h| h.bytes.len else 0; } /// A kept state's bytes in `out` (its key and value rows, then its Mamba states), after the engine's stream. pub fn download(t: *const Target, saved: *anyopaque, out: []u8) !void { const copy: *Copy = @ptrCast(@alignCast(saved)); + if (copy.host) |h| { + try t.e.stream.synchronize(); + @memcpy(out[0..h.bytes.len], h.bytes); + return; + } const kv = kvBytes(t.e.c, copy.pos, copy.head); try t.e.ops().download(out[0..kv], copy.kv.ptr); try t.e.ops().download(out[kv..], copy.mamba.ptr); @@ -183,7 +239,10 @@ pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror! const e = t.e; if (copy.head != (t.head != null)) return error.NoHead; // checked before any copy: nothing half restored if (!copy.whole) return error.Incomplete; // its chunk failed before `finish` - _ = try move(t, copy.pos, copy, false, true); + if (copy.host) |h| { + try hostMove(t, copy.pos, h.bytes, false, true, null); + try e.stream.synchronize(); // the store may drop the host copy once this returns + } else _ = try move(t, copy.pos, copy, false, true); e.pos = copy.pos; e.parity = copy.parity; e.prev_keep = 0; @@ -196,8 +255,13 @@ pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { const t: *Target = @ptrCast(@alignCast(ptr)); const copy: *Copy = @ptrCast(@alignCast(saved)); t.e.armed.forget(saved); // a copy evicted before its chunk ran: the chunk must not write it - if (!t.e.spares.give(copy.mamba)) copy.mamba.free(); - copy.kv.free(); + if (copy.host) |*h| { + // queued copies may still read or write it: they finish before the host memory goes + t.e.stream.synchronize() catch {}; + h.free(); + } + if (copy.mamba.len > 0 and !t.e.spares.give(copy.mamba)) copy.mamba.free(); + if (copy.kv.len > 0) copy.kv.free(); t.e.gpa.destroy(copy); } diff --git a/zig/src/families/nemotron/cuda_state.zig b/zig/src/families/nemotron/cuda_state.zig index 92bac057b9..c65e23c804 100644 --- a/zig/src/families/nemotron/cuda_state.zig +++ b/zig/src/families/nemotron/cuda_state.zig @@ -5,6 +5,7 @@ const lesson_route = @import("lesson_route.zig"); const cuda = @import("cuda"); const Config = @import("config.zig").Config; const kern = @import("cuda_kernels.zig"); +const glue = @import("cuda_glue.zig"); const sampler = @import("cuda_sampler.zig"); const seq_graphs = @import("cuda_seq_graphs.zig"); @@ -81,7 +82,7 @@ fn seqSizes(c: Config, max_len: usize) [seq_fields.len]usize { return .{ kv, kv, m[0], m[1], nm * 2 * W * cd * 2, nm * 2 * W * cd * 2, nm * 2 * W * c.mamba_heads * 4, 16, W * 4, W * c.hidden * 2, W * 4, @sizeOf(sampler.Rule) }; } -const scratch_count = 36; +const scratch_count = 39; /// Scratch windows and chunks share on one stream: window logits, chunk io, activations, the expert plan. fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { @@ -94,6 +95,7 @@ fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { const xd: usize = c.inner(); const pairs = R * ns; const items = kern.maxItems(@intCast(pairs), c.experts + 2, 16); + const lat: usize = @max(c.latent, 1); // LatentMoE's rows in the latent width (Bolt), else unused return .{ W * c.vocab * 2, 16, R * 4, R * D * 2, c.vocab * 2, 16, @@ -103,10 +105,11 @@ fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { R * xd * 2, R * xd / 64 * 4, R * c.qkvDim() * 2, R * qd * 2, R * qd * 2, R * qd / 64 * 4, W * nch * qd * 4, W * nch * c.heads * 4, W * nch * c.heads * 4, - 6 * R * c.experts * 4, R * ns * 4, R * ns * 4, + glue.max_router_slices * R * c.experts * 4, R * ns * 4, R * ns * 4, pairs * 4, items * 12, 8, pairs * 4, (pairs + 1023) / 1024 * (c.experts + 2) * 4, pairs * c.expert_width * 2, pairs * D * 4, shared_words * 4, R * kern.pocket_width * 4, + R * lat * 2, R * lat * 2, R * 4, }; } @@ -218,6 +221,9 @@ pub const Buffers = struct { rule: u64, // the sequence's sampled rule, as sample.cu reads it shared: u64, // the shared expert's fixed plans: 32 members, then (items, counts) for 1..16 rows pocket: u64, // the pocket's activations, fp32 [rows, pocket_width] + lat: u64, // LatentMoE: the rows fc1_latent_proj makes, bf16 [rows, latent] + z: u64, // LatentMoE: the routed experts' weighted sum, bf16 [rows, latent], fc2_latent_proj's input + zeros: u64, // LatentMoE: every pair's expert 0, the shared expert's picks as a one-expert layer /// Sizes every buffer for `max_len` cache rows and `nch` attention chunk partials a row; the own sequence's first. pub fn init(d: *const cuda.Driver, c: Config, max_len: usize, nch: usize, carve: ?*cuda.Carveout) !Buffers { @@ -263,6 +269,7 @@ pub const Buffers = struct { &b.xs, &b.delta, &b.proj, &b.p_xc, &b.sy, &b.g, &b.gxs, &b.qkv, &b.q, &b.att, &b.axs, &b.po, &b.pm, &b.pl, &b.part, &b.pick, &b.wts, &b.plan.members, &b.plan.items, &b.plan.counts, &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, + &b.lat, &b.z, &b.zeros, }; } @@ -285,6 +292,10 @@ pub const Buffers = struct { w[at..][0..8].* = .{ @intCast(c.experts), 0, n, @intCast(c.experts + 1), max_rows, n, 2, 2 }; } try ops.upload(b.shared, std.mem.sliceAsBytes(&w)); + if (c.latent > 0) { + const z: [pass_rows]u32 = @splat(0); + try ops.upload(b.zeros, std.mem.sliceAsBytes(&z)); + } try ops.s.synchronize(); } diff --git a/zig/src/families/nemotron/cuda_train_fwd.zig b/zig/src/families/nemotron/cuda_train_fwd.zig index 3b39c30976..e681d19845 100644 --- a/zig/src/families/nemotron/cuda_train_fwd.zig +++ b/zig/src/families/nemotron/cuda_train_fwd.zig @@ -86,7 +86,7 @@ pub fn experts(f: *const fwd.Forward, m: weights.MoE, rows: usize, scale: Scale, const o = f.ops; const ex = m.experts; const pairs = rows * c.slots(); - try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.routed_scaling, c.norm_topk); + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); const tile = cuda.experts.promptTile(ex.format); try o.plan(b.pick, pairs, ex.experts, tile, b.plan); const items = kern.maxItems(pairs, ex.experts, tile); diff --git a/zig/src/families/nemotron/cuda_vision.zig b/zig/src/families/nemotron/cuda_vision.zig new file mode 100644 index 0000000000..2dd6c4b625 --- /dev/null +++ b/zig/src/families/nemotron/cuda_vision.zig @@ -0,0 +1,511 @@ +//! Nemotron 3.5 Bolt's image tower on CUDA: C-RADIO ViT-H/16 (vision_model.radio_model.) and the mlp1 projector, +//! bf16 weights, cuBLASLt GEMMs with fp32 accumulation, an fp32 residual stream; kernels in radio_vision.cu. +//! +//! `encode` turns one image's patches (vision/nemotron_image.zig's `prepare`) into the bf16 [tokens, 4096] rows that +//! replace the prompt's `` placeholders: `` (19), then `` (18) x tokens, then `` (20), the +//! image directly before the question. Oracle: tools/zig/bolt_vision_reference.py. +//! +//! VRAM (one arena, made at load): weights 1.681 GB (tower 1.303 GB: blocks 1.259 GB, the 128x128 position table +//! 42 MB, patch Linear and prefix tokens 2 MB; projector 377 MB) plus scratch sized for the largest image (13,312 +//! patches) 0.702 GB: 2.383 GB in all, and the caller's output rows (3,328 x 4096 bf16 at most, 27 MB). +//! The video embedder (video_embedder.weight, 3.9 MB) is checked but not loaded: video is not this module's. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const lt = cuda.cublaslt; +const DeviceBuffer = cuda.DeviceBuffer; +const lanes_ = @import("lanes"); +const radio = "vision_model.radio_model."; +const model = radio ++ "model."; + +pub const hidden = 1280; +pub const heads = 16; +pub const head_dim = 80; // radio_vision.cu's attention is built for 80 +pub const depth = 32; +pub const mlp = 5120; +pub const patch = 16; +pub const patch_width = 3 * patch * patch; // 768, (c, dy, dx) +pub const prefix_tokens = 10; // 3 cls + 7 registers, dropped after the last block +pub const cpe_side = 128; // pos_embed's 16384 rows +pub const shuffled = 4 * hidden; // pixel shuffle v2: 2x2 patches a row +pub const projector_hidden = 20480; +pub const out_hidden = 4096; // the language model's width +pub const ln_eps: f32 = 1e-6; +pub const rms_eps: f32 = 1e-5; +/// One call's patches at most: max_num_patches. +pub const max_patches = 13312; +/// The prompt's placeholder ids (tokenizer.json): `` repeated, between `` and ``. +pub const image_token: u32 = 18; +pub const image_start: u32 = 19; +pub const image_end: u32 = 20; +const workspace_bytes = 32 << 20; +/// The processor's normalization (vision/nemotron_image.zig's MEAN and STD): input_conditioner must hold the same. +const clip_mean = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +const clip_std = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; + +/// What config.json must say for this module to run the checkpoint's tower; anything else is refused. +pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + // Bolt's quantization_config makes it ~6.5 MB + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(bytes); + return checkConfig(gpa, bytes); +} + +pub fn checkConfig(gpa: std.mem.Allocator, bytes: []const u8) !void { + 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 error.UnsupportedBoltVision).object; + const ints = [_]struct { std.json.ObjectMap, []const u8, i64 }{ + .{ v, "hidden_size", hidden }, .{ v, "num_attention_heads", heads }, .{ v, "num_hidden_layers", depth }, + .{ v, "patch_size", patch }, .{ v, "num_cls_tokens", 3 }, .{ v, "num_registers", 7 }, + .{ root, "projector_hidden_size", projector_hidden }, .{ root, "img_context_token_id", image_token }, + .{ root, "vit_hidden_size", hidden }, + }; + for (ints) |w| if (num(w[0].get(w[1])) != w[2]) return error.UnsupportedBoltVision; + const llm = (root.get("llm_config") orelse return error.UnsupportedBoltVision).object; + if (num(llm.get("hidden_size")) != out_hidden) return error.UnsupportedBoltVision; + if (!isString(v.get("hidden_act"), "gelu") or !isString(root.get("ps_version"), "v2")) return error.UnsupportedBoltVision; + if (!isBool(v.get("use_swiglu_ffn"), false) or !isBool(root.get("use_thumbnail"), false) or !isBool(root.get("vision_final_layernorm"), false)) return error.UnsupportedBoltVision; + if (!isFloat(root.get("downsample_ratio"), 0.5) or !isFloat(v.get("layer_norm_eps"), 1e-6) or !isFloat(v.get("mlp_ratio"), 4.0)) return error.UnsupportedBoltVision; +} + +fn num(x: ?std.json.Value) i64 { + const v = x orelse return -1; + return if (v == .integer) v.integer else -1; +} + +fn isString(x: ?std.json.Value, want: []const u8) bool { + const v = x orelse return false; + return v == .string and std.mem.eql(u8, v.string, want); +} + +fn isBool(x: ?std.json.Value, want: bool) bool { + const v = x orelse return !want; // an absent flag is false + return v == .bool and v.bool == want; +} + +fn isFloat(x: ?std.json.Value, want: f64) bool { + const v = x orelse return false; + const f: f64 = switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => return false, + }; + return @abs(f - want) <= 1e-12 * @abs(want); +} + +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: cuda.Function, + attention: cuda.Function, + bias_gelu: cuda.Function, + bias_residual: cuda.Function, + shuffle_rms: cuda.Function, + relu2: cuda.Function, + cast: cuda.Function, +}; + +/// Where one call's stages live in the scratch (device pointers), for tests that read them back after `encode`. +pub const Stages = struct { + /// fp32 [prefix_tokens + n, 1280]: the residual stream, after `encode` the last block's output (tower_out). + x: u64, +}; + +/// The tower's and projector's weights and one call's scratch in one allocation, its kernels and GEMM plans. +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + mod: cuda.Module, + f: Fns, + lib: lt.Library, + arena: DeviceBuffer, + at: usize = 0, + scratch: u64 = 0, + embed_w: u64 = 0, + pos: u64 = 0, + cls: u64 = 0, + blocks: [depth]Block = undefined, + norm_w: u64 = 0, + proj1_w: u64 = 0, + proj2_w: u64 = 0, + gemms: Gemms = undefined, + /// Blocks `encode` runs (tests: 0 leaves the embedded rows in Stages.x); the projector runs only after all 32. + run_depth: usize = depth, + + /// The checkpoint's vision_model. and mlp1. tensors (bf16) from `dir` (Bolt's own directory: the language + /// model's loader skips them). Uploads synchronously; the arena is sized for `max_patches`. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !*Tower { + try check(gpa, io, dir); + const t = try gpa.create(Tower); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .d = d, .mod = undefined, .f = undefined, .lib = undefined, .arena = undefined }; + t.lib = try lt.Library.open(); + errdefer t.lib.close(); + t.gemms = .{ .lib = &t.lib }; + errdefer t.gemms.deinit(); + t.mod = try cuda.Module.load(d, cuda.kernels.radio_vision); + errdefer t.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(t.f, name) = try t.mod.function("tf_rv_" ++ name); + t.arena = try DeviceBuffer.alloc(d, weightBytes() + scratchBytes(max_patches)); + errdefer t.arena.free(); + try t.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + var vk = try core.Checkpoint.openModelPrefix(gpa, io, dir, "vision_model."); + defer vk.close(); + var pk = try core.Checkpoint.openModelPrefix(gpa, io, dir, "mlp1."); + defer pk.close(); + { // the input conditioner is the processor's normalization, already applied on the host: it must agree + const mean = try vk.expect(radio ++ "input_conditioner.norm_mean", .f32, &.{ 3, 1, 1 }); + const stdv = try vk.expect(radio ++ "input_conditioner.norm_std", .f32, &.{ 3, 1, 1 }); + for (0..3) |c| { + const m: f32 = @bitCast(std.mem.readInt(u32, mean.bytes[c * 4 ..][0..4], .little)); + const s: f32 = @bitCast(std.mem.readInt(u32, stdv.bytes[c * 4 ..][0..4], .little)); + if (m != clip_mean[c] or s != clip_std[c]) return error.UnsupportedBoltVision; + } + } + _ = try vk.expect("vision_model.summary_idxs", .i64, &.{2}); // the summary heads' cls rows: not used for images + _ = try vk.expect(model ++ "patch_generator.video_embedder.weight", .bf16, &.{ hidden, 2 * patch_width }); // video + t.embed_w = try t.get(&vk, model ++ "patch_generator.embedder.weight", &.{ hidden, patch_width }); + t.pos = try t.get(&vk, model ++ "patch_generator.pos_embed", &.{ 1, cpe_side * cpe_side, hidden }); + t.cls = try t.get(&vk, model ++ "patch_generator.cls_token.token", &.{ prefix_tokens, hidden }); + for (&t.blocks, 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.fc1.weight", .{ mlp, hidden } }, .{ "fc1_b", "mlp.fc1.bias", .{mlp} }, + .{ "fc2_w", "mlp.fc2.weight", .{ hidden, mlp } }, .{ "fc2_b", "mlp.fc2.bias", .{hidden} }, + }; + inline for (shapes) |sh| { + var nb: [128]u8 = undefined; + const shape: [sh[2].len]usize = sh[2]; + @field(b, sh[0]) = try t.get(&vk, try std.fmt.bufPrint(&nb, model ++ "blocks.{d}.{s}", .{ i, sh[1] }), &shape); + } + } + t.norm_w = try t.get(&pk, "mlp1.0.weight", &.{shuffled}); + t.proj1_w = try t.get(&pk, "mlp1.1.weight", &.{ projector_hidden, shuffled }); + t.proj2_w = try t.get(&pk, "mlp1.3.weight", &.{ out_hidden, projector_hidden }); + if (vk.unused() != 0 or pk.unused() != 0) return error.UnusedVisionTensors; + t.scratch = try t.take(scratchBytes(max_patches)); + return t; + } + + pub fn deinit(t: *Tower) void { + t.gemms.deinit(); + t.arena.free(); + t.mod.unload(); + t.lib.close(); + t.gpa.destroy(t); + } + + /// The weights' bytes in the arena (each allocation 256-byte aligned): 1,680,788,480. + pub fn weightBytes() usize { + const a = aligned; + var n: usize = a(hidden * patch_width * 2) + a(cpe_side * cpe_side * hidden * 2) + a(prefix_tokens * hidden * 2); + n += depth * (4 * a(hidden * 2) + a(3 * hidden * hidden * 2) + a(3 * hidden * 2) + a(hidden * hidden * 2) + a(hidden * 2) + a(mlp * hidden * 2) + a(mlp * 2) + a(hidden * mlp * 2) + a(hidden * 2)); + return n + a(shuffled * 2) + a(projector_hidden * shuffled * 2) + a(out_hidden * projector_hidden * 2); + } + + /// The scratch one call of `n` patches needs, in bytes (aligned pieces; 13,312 patches: 702 MB). + pub fn scratchBytes(n: usize) usize { + const a = aligned; + const rows = prefix_tokens + n; + const m4 = n / 4; + return a(n * patch_width * 2) + // bf16 patches + a(rows * hidden * 4) + // x, the fp32 residual + a(@max(rows * hidden, m4 * shuffled) * 2) + // h: LayerNorm out, then the shuffled rows + a(@max(rows * @max(3 * hidden, mlp), m4 * projector_hidden) * 4) + // fp32 GEMM out + 3 * a(rows * hidden * 2) + // q, k, v + a(rows * hidden * 2) + // attention out + a(@max(rows * mlp, m4 * projector_hidden) * 2) + // GELU / relu^2 rows + a(n * 4 * 4) * 2 + // CPE taps and weights + 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 get(t: *Tower, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + const tn = try ck.expect(name, .bf16, shape); + const p = try t.take(tn.bytes.len); + try t.arena.upload(p - t.arena.ptr, tn.bytes); + return p; + } + + /// One image's patches (fp32 [gh * gw, 768], nemotron_image.prepare's) to bf16 [gh * gw / 4, 4096] rows at `out`, + /// on stream `s`. Returns once the host's inputs are copied (the stream is synchronized); `out` is then written. + pub fn encode(t: *Tower, s: cuda.Stream, patches: []const f32, gh: u32, gw: u32, out: u64) !Stages { + const n = @as(usize, gh) * gw; + if (gh == 0 or gw == 0 or gh % 2 != 0 or gw % 2 != 0 or n > max_patches) return error.BadImageGrid; + if (patches.len != n * patch_width) return error.BadImagePatches; + const gpa = t.gpa; + const rows = prefix_tokens + n; + const m4 = n / 4; + const px16 = try gpa.alloc(u16, n * patch_width); + defer gpa.free(px16); + for (px16, patches) |*o, v| o.* = bf16(v); + const idx = try gpa.alloc(i32, n * 4); + defer gpa.free(idx); + const wts = try gpa.alloc(f32, n * 4); + defer gpa.free(wts); + cpeTaps(gh, gw, idx, wts); + + var cur = t.scratch; + const end = t.scratch + scratchBytes(max_patches); + const d_px = carve(&cur, n * patch_width * 2); + const x = carve(&cur, rows * hidden * 4); + const h = carve(&cur, @max(rows * hidden, m4 * shuffled) * 2); + const acc = carve(&cur, @max(rows * @max(3 * hidden, mlp), m4 * projector_hidden) * 4); + const q = carve(&cur, rows * hidden * 2); + const k = carve(&cur, rows * hidden * 2); + const v = carve(&cur, rows * hidden * 2); + const att = carve(&cur, rows * hidden * 2); + const mid = carve(&cur, @max(rows * mlp, m4 * projector_hidden) * 2); + const d_idx = carve(&cur, n * 4 * 4); + const d_wts = carve(&cur, n * 4 * 4); + const ws = cur; + const ws_len = end - cur; + try t.upload(s, d_px, std.mem.sliceAsBytes(px16)); + try t.upload(s, d_idx, std.mem.sliceAsBytes(idx)); + try t.upload(s, d_wts, std.mem.sliceAsBytes(wts)); + + const mm = &t.gemms; + try mm.run(n, hidden, patch_width, d_px, t.embed_w, acc, ws, ws_len, s); + try go(t.f.embed, s, rows, 1, 256, .{ acc, t.pos, d_idx, d_wts, t.cls, x, int(prefix_tokens), int(hidden) }); + const scale_log2: f32 = std.math.log2e / @sqrt(@as(f32, head_dim)); + for (t.blocks[0..t.run_depth]) |b| { + try go(t.f.layernorm, s, rows, 1, 256, .{ x, b.norm1_w, b.norm1_b, h, int(hidden), ln_eps }); + try mm.run(rows, 3 * hidden, hidden, h, b.qkv_w, acc, ws, ws_len, s); + try go(t.f.qkv, s, rows, 1, 256, .{ acc, b.qkv_b, q, k, v, int(rows), int(heads), int(head_dim) }); + try go(t.f.attention, s, (rows + 63) / 64, heads, 128, .{ q, k, v, att, int(rows), int(heads), scale_log2 }); + try mm.run(rows, hidden, hidden, att, b.proj_w, acc, ws, ws_len, s); + try go(t.f.bias_residual, s, rows, 1, 256, .{ acc, b.proj_b, x, int(hidden) }); + try go(t.f.layernorm, s, rows, 1, 256, .{ x, b.norm2_w, b.norm2_b, h, int(hidden), ln_eps }); + try mm.run(rows, mlp, hidden, h, b.fc1_w, acc, ws, ws_len, s); + try go(t.f.bias_gelu, s, rows, 1, 256, .{ acc, b.fc1_b, mid, int(mlp) }); + try mm.run(rows, hidden, mlp, mid, b.fc2_w, acc, ws, ws_len, s); + try go(t.f.bias_residual, s, rows, 1, 256, .{ acc, b.fc2_b, x, int(hidden) }); + } + if (t.run_depth < depth) { + try s.synchronize(); + return .{ .x = x }; + } + // projector: the patch rows (prefix dropped) shuffled 2x2, RMSNorm, Linear, relu^2, Linear + try go(t.f.shuffle_rms, s, m4, 1, 256, .{ x, t.norm_w, h, int(prefix_tokens), int(gw), int(hidden), rms_eps }); + try mm.run(m4, projector_hidden, shuffled, h, t.proj1_w, acc, ws, ws_len, s); + try go(t.f.relu2, s, m4, 1, 256, .{ acc, mid, int(projector_hidden) }); + try mm.run(m4, out_hidden, projector_hidden, mid, t.proj2_w, acc, ws, ws_len, s); + try go(t.f.cast, s, m4, 1, 256, .{ acc, out, int(out_hidden) }); + try s.synchronize(); // the host's patch rows and tables are freed on return + return .{ .x = x }; + } + + fn upload(t: *Tower, s: cuda.Stream, 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, s.handle), "cuMemcpyHtoDAsync"); + } +}; + +fn go(f: cuda.Function, s: cuda.Stream, 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 } }, s, &a); +} + +fn carve(p: *u64, len: usize) u64 { + const r = p.*; + p.* += aligned(len); + return r; +} + +/// cuBLASLt plans for the shapes a call runs, made on first use and kept (a new image size adds plans). +const Gemms = struct { + lib: *const lt.Library, + plans: [48]?struct { m: usize, n: usize, k: usize, lin: lt.Linear } = @splat(null), + next: usize = 0, + + 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); + }; + // a full table drops its oldest plan (8 shapes an image size: six images' worth kept) + const slot = &g.plans[g.next]; + if (slot.*) |*old| old.lin.deinit(); + slot.* = .{ .m = m, .n = n, .k = k, .lin = try lt.Linear.init(g.lib, m, n, k, .f32, ws_len) }; + g.next = (g.next + 1) % g.plans.len; + return slot.*.?.lin.run(x, w, out, ws, ws_len, s.handle); + } + + fn deinit(g: *Gemms) void { + for (&g.plans) |*p| if (p.*) |*have| { + have.lin.deinit(); + p.* = null; + }; + } +}; + +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); +} + +/// CPE: the 128x128 table resized (bilinear, align_corners=False) to max(gh, gw) square and cropped to gh x gw. +/// Each patch's 4 table rows (y0 x0, y0 x1, y1 x0, y1 x1) and weights (w0, w1, h0, h1), as aten computes them in fp32. +pub fn cpeTaps(gh: u32, gw: u32, idx: []i32, wts: []f32) void { + const m = @max(gh, gw); + for (0..gh) |i| { + const hy = axis(@intCast(i), m); + for (0..gw) |j| { + const wx = axis(@intCast(j), m); + const p = i * gw + j; + idx[p * 4 + 0] = hy.lo * cpe_side + wx.lo; + idx[p * 4 + 1] = hy.lo * cpe_side + wx.hi; + idx[p * 4 + 2] = hy.hi * cpe_side + wx.lo; + idx[p * 4 + 3] = hy.hi * cpe_side + wx.hi; + wts[p * 4 + 0] = wx.wlo; + wts[p * 4 + 1] = wx.whi; + wts[p * 4 + 2] = hy.wlo; + wts[p * 4 + 3] = hy.whi; + } + } +} + +const Axis = struct { lo: i32, hi: i32, wlo: f32, whi: f32 }; + +/// aten's area_pixel_compute_source_index (linear, align_corners=False) and lambdas for output `o` of `size`. +fn axis(o: u32, size: u32) Axis { + const scale: f32 = @as(f32, cpe_side) / @as(f32, @floatFromInt(size)); + const src = @max(scale * (@as(f32, @floatFromInt(o)) + 0.5) - 0.5, 0.0); + const lo: i32 = @intFromFloat(src); + const hi: i32 = if (lo < cpe_side - 1) lo + 1 else lo; + const whi = @min(@max(src - @as(f32, @floatFromInt(lo)), 0.0), 1.0); + return .{ .lo = lo, .hi = hi, .wlo = 1.0 - whi, .whi = whi }; +} + +test "CPE taps: a 128 grid reads the table as is; an upsampled one interpolates and clamps at the far edge" { + var idx: [4 * 4]i32 = undefined; + var wts: [4 * 4]f32 = undefined; + const a = axis(5, cpe_side); + try std.testing.expectEqual(@as(i32, 5), a.lo); + try std.testing.expectEqual(@as(f32, 0), a.whi); + const last = axis(255, 256); + try std.testing.expectEqual(@as(i32, 127), last.lo); + try std.testing.expectEqual(@as(i32, 127), last.hi); + cpeTaps(2, 2, &idx, &wts); + // m = 2: scale 64, source 31.5 for row/col 0 + try std.testing.expectEqual(@as(i32, 31 * cpe_side + 31), idx[0]); + try std.testing.expectEqual(@as(f32, 0.5), wts[0]); +} + +test "Bolt's config is accepted; another tower or projector is refused" { + const gpa = std.testing.allocator; + const good = + \\{"vision_config":{"hidden_size":1280,"num_attention_heads":16,"num_hidden_layers":32,"patch_size":16, + \\"num_cls_tokens":3,"num_registers":7,"hidden_act":"gelu","use_swiglu_ffn":false,"layer_norm_eps":1e-06, + \\"mlp_ratio":4.0},"projector_hidden_size":20480,"img_context_token_id":18,"vit_hidden_size":1280, + \\"llm_config":{"hidden_size":4096},"ps_version":"v2","use_thumbnail":false,"vision_final_layernorm":false, + \\"downsample_ratio":0.5} + ; + try checkConfig(gpa, good); + const cases = [_][2][]const u8{ .{ "\"num_registers\":7", "\"num_registers\":4" }, .{ "\"ps_version\":\"v2\"", "\"ps_version\":\"v1\"" }, .{ "\"use_swiglu_ffn\":false", "\"use_swiglu_ffn\":true" }, .{ "\"downsample_ratio\":0.5", "\"downsample_ratio\":0.25" } }; + for (cases) |c| { + const bad = try std.mem.replaceOwned(u8, gpa, good, c[0], c[1]); + defer gpa.free(bad); + try std.testing.expectError(error.UnsupportedBoltVision, checkConfig(gpa, bad)); + } +} + +/// An image prompt's rows for the prompt pass: its images' features, and each prompt row's feature row (-1: text). +pub const Media = struct { + gpa: std.mem.Allocator, + features: DeviceBuffer, // bf16 [image rows, out_hidden] + src: DeviceBuffer, // i32 [prompt] + + /// `images` (the request's, in prompt order) through the tower on `s`; each fills its ... run. + pub fn init(gpa: std.mem.Allocator, t: *Tower, s: cuda.Stream, prompt: []const u32, images: []const lanes_.Image) !Media { + const host = try gpa.alloc(i32, prompt.len); + defer gpa.free(host); + const rows = try placeholders(prompt, images, host); + var m: Media = .{ .gpa = gpa, .features = try DeviceBuffer.alloc(t.d, @max(rows, 1) * out_hidden * 2), .src = undefined }; + errdefer m.features.free(); + m.src = try DeviceBuffer.alloc(t.d, prompt.len * 4); + errdefer m.src.free(); + var at: usize = 0; + for (images) |im| { + const gh: u32 = @intCast(im.grid[1]); + const gw: u32 = @intCast(im.grid[2]); + _ = try t.encode(s, im.patches, gh, gw, m.features.ptr + at * out_hidden * 2); + at += @as(usize, gh) * gw / 4; + } + try m.src.upload(0, std.mem.sliceAsBytes(host)); + return m; + } + + pub fn deinit(m: *Media) void { + m.features.free(); + m.src.free(); + } +}; + +/// Each image's run of image_token rows between image_start and image_end, in order, its length the image's tokens; +/// `src` gets each row's feature row (-1 for text). Returns the feature rows. +pub fn placeholders(prompt: []const u32, images: []const lanes_.Image, src: []i32) !usize { + @memset(src, -1); + var next: usize = 0; // feature rows handed out + var k: usize = 0; // images placed + var i: usize = 0; + while (i < prompt.len) : (i += 1) { + if (prompt[i] != image_token) continue; + if (k == images.len or i == 0 or prompt[i - 1] != image_start) return error.BadImagePlaceholders; + const im = images[k]; + if (im.grid[0] != 1 or im.grid[1] <= 0 or im.grid[2] <= 0 or @mod(im.grid[1] * im.grid[2], 4) != 0) return error.BadImageGrid; + const n: usize = @intCast(@divExact(im.grid[1] * im.grid[2], 4)); + if (im.patches.len != n * 4 * patch_width) return error.BadImagePatches; + if (i + n >= prompt.len or prompt[i + n] != image_end) return error.BadImagePlaceholders; + for (prompt[i..][0..n], src[i..][0..n], 0..) |id, *x, j| { + if (id != image_token) return error.BadImagePlaceholders; + x.* = @intCast(next + j); + } + next += n; + k += 1; + i += n; + } + if (k != images.len) return error.BadImagePlaceholders; + return next; +} + +test "placeholder runs map to feature rows in order; a short run is refused" { + const px: [8 * patch_width]f32 = @splat(0); + const ims = [_]lanes_.Image{ .{ .patches = &px, .grid = .{ 1, 2, 2 } }, .{ .patches = &px, .grid = .{ 1, 2, 4 } } }; + const p = [_]u32{ 1, image_start, image_token, image_end, 7, image_start, image_token, image_token, image_end, 9 }; + var src: [p.len]i32 = undefined; + // the first image is 4 patches (1 token) but its patches slice is 8: refused; give each its own size + const one: [4 * patch_width]f32 = @splat(0); + const good = [_]lanes_.Image{ .{ .patches = &one, .grid = .{ 1, 2, 2 } }, .{ .patches = &px, .grid = .{ 1, 2, 4 } } }; + try std.testing.expectEqual(@as(usize, 3), try placeholders(&p, &good, &src)); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 0, -1, -1, -1, 1, 2, -1, -1 }, &src); + try std.testing.expectError(error.BadImagePatches, placeholders(&p, &ims, &src)); + const short = [_]u32{ image_start, image_token, image_end, image_start, image_token, image_end }; + try std.testing.expectError(error.BadImagePlaceholders, placeholders(&short, &good, &src)); +} diff --git a/zig/src/families/nemotron/cuda_weights.zig b/zig/src/families/nemotron/cuda_weights.zig index ff25a8da37..086fbc07a1 100644 --- a/zig/src/families/nemotron/cuda_weights.zig +++ b/zig/src/families/nemotron/cuda_weights.zig @@ -28,7 +28,11 @@ pub const Attention = struct { qkv: QLinear, o: QLinear, o_rest: u64 = 0, adapte /// The pocket kernels' neurons (config's pocket): input weights and outputs, [n, D] fp32 each, from the 4-bit codes. pub const Pocket = kern.Pocket; /// The shared expert's down projection rest in two halves (experts E and E + 1); its learned change reads both. -pub const MoE = struct { router: u64, bias: u64, experts: Experts, rest: [2]u64 = .{ 0, 0 }, adapter: ?Adapter = null, pocket: ?Pocket = null }; +pub const MoE = struct { router: u64, bias: u64, experts: Experts, rest: [2]u64 = .{ 0, 0 }, adapter: ?Adapter = null, pocket: ?Pocket = null, latent: ?Latent = null }; + +/// LatentMoE's own parts (Bolt): rows go to the latent width and back around the routed experts, and the shared +/// expert reads the full-width rows on its own (`experts` then holds the routed experts alone). +pub const Latent = struct { fc1: QLinear, fc2: QLinear, shared: Experts }; pub const Block = struct { kind: Kind, norm: u64, mamba: Mamba = undefined, attn: Attention = undefined, moe: MoE = undefined }; pub const Mtp = struct { enorm: u64, hnorm: u64, eh_proj: QLinear, attn_norm: u64, attn: Attention, moe_norm: u64, moe: MoE, final_norm: u64 }; @@ -458,6 +462,12 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, var ck = try core.Checkpoint.openModel(gpa, io, dir); defer ck.close(); try src.add(&ck); + if (c.omni) { + // Bolt: the language model under language_model., its vision tower and projector beside it + try ck.skip("vision_model."); + try ck.skip("mlp1."); + ck.scope = "language_model."; + } w.blocks = try gpa.alloc(Block, c.layers); for (c.kinds[0..c.layers], 0..) |kind, i| { var nm: [128]u8 = undefined; @@ -470,7 +480,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, switch (kind) { .mamba => b.mamba = try L.mamba(&ck, try join(&nb, name, ".mamba"), mixer), .attention => b.attn = try L.attention(&ck, try join(&nb, name, ".attn"), mixer), - .moe => b.moe = try L.moe(&ck, try join(&nb, name, ".moe"), mixer, c), + .moe => b.moe = if (c.latent > 0) try modelopt.latent(&L, &ck, try join(&nb, name, ".moe"), mixer, c, i) else try L.moe(&ck, try join(&nb, name, ".moe"), mixer, c), } w.blocks[i] = b; } diff --git a/zig/src/families/nemotron/glue_ref.zig b/zig/src/families/nemotron/glue_ref.zig index f4fc4550b9..d201f98be9 100644 --- a/zig/src/families/nemotron/glue_ref.zig +++ b/zig/src/families/nemotron/glue_ref.zig @@ -159,8 +159,8 @@ pub fn router(x: []const u16, w: []const u16, part: []f32, rows: usize, d: usize /// Sum slices in order, rank sigmoid scores plus bias (lowest id on ties), then add shared slots. pub fn topk(part: []const f32, bias: []const f32, idx: []i32, wt: []f32, rows: usize, scaling: f32, e: usize, sk: usize, top_k: usize, ns: usize, norm: bool) void { - var score: [256]f32 = undefined; - var sel: [256]f32 = undefined; + var score: [512]f32 = undefined; // the top-k kernel's MAXC: up to 512 experts + var sel: [512]f32 = undefined; for (0..rows) |r| { for (0..e) |ex| { var logit: f32 = 0; @@ -170,8 +170,8 @@ pub fn topk(part: []const f32, bias: []const f32, idx: []i32, wt: []f32, rows: u sel[ex] = if (std.math.isNan(v)) -std.math.inf(f32) else v; } var total: f32 = 0; - var probs: [8]f32 = undefined; - var ids: [8]i32 = undefined; + var probs: [16]f32 = undefined; + var ids: [16]i32 = undefined; for (0..top_k) |k| { var bi: usize = 0; for (1..e) |ex| if (sel[ex] > sel[bi]) { diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 7fbf38ba3d..0bed892d5d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -26,12 +26,29 @@ const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { - var out: [registry.len]api.Family = undefined; - for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats }; + var n: usize = 0; + for (registry) |F| n += 1 + aliases(F).len; + var out: [n]api.Family = undefined; + var i: usize = 0; + for (registry) |F| for ([_][]const u8{F.model_type} ++ aliases(F)) |t| { + out[i] = .{ .model_type = t, .formats = F.formats }; + i += 1; + }; const final = out; break :blk &final; }; +/// The other config model_types a family reads (a wrapper's, as Bolt's nemotron_h_omni). +fn aliases(comptime F: type) []const []const u8 { + return if (@hasDecl(F, "also_model_types")) F.also_model_types else &.{}; +} + +fn reads(comptime F: type, model_type: []const u8) bool { + if (std.mem.eql(u8, model_type, F.model_type)) return true; + for (aliases(F)) |t| if (std.mem.eql(u8, model_type, t)) return true; + return false; +} + const gib: f64 = 1 << 30; /// The chip class gate entries name ("nvidia-sm121" for a GB10) of the device TF_CUDA_DEVICE picks; null without one. @@ -71,10 +88,11 @@ fn setting(a: Allocator, flag: ?u32, name: [:0]const u8, default: u32, problem: /// The model's window (config.json's max_position_embeddings, text_config's first), 0 when it names none. fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; - const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(32 << 20)) catch return 0; const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; if (doc != .object) return 0; - const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; + const inner = doc.object.get("text_config") orelse doc.object.get("llm_config"); // Bolt's is llm_config + const text = if (inner) |t| (if (t == .object) t else doc) else doc; const limit = text.object.get("max_position_embeddings") orelse doc.object.get("max_position_embeddings") orelse return 0; return if (limit == .integer and limit.integer > 0) limit.integer else 0; } @@ -275,7 +293,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } inline for (registry) |F| { - if (std.mem.eql(u8, o.model_type, F.model_type)) return openWith(F, a, gpa, io, o, problem); + if (reads(F, o.model_type)) return openWith(F, a, gpa, io, o, problem); } problem.* = try std.fmt.allocPrint(a, "the native CUDA engine has no backend for {s} checkpoints yet", .{o.model_type}); return null; @@ -402,10 +420,11 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O 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 }; // 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; + const host_snaps = @hasField(@TypeOf(loaded), "host_snaps") and loaded.host_snaps; + const cache = if (!@hasField(@TypeOf(loaded), "target")) 0 else if (host_snaps) budget.hostCacheBytes(a, io, o.prompt_cache_gib) 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 {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", + if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache{s}", .{ toGib(cache), if (host_snaps) " 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(@TypeOf(loaded), "target") and cache > 0) { @@ -441,7 +460,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, .vision_nemotron = @hasDecl(F, "vision_nemotron") and F.vision_nemotron, .plain_only = loaded.rows == 1, .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| { @@ -481,9 +500,13 @@ test "chip classes name the compute capability as gate entries do" { try std.testing.expectEqualStrings("nvidia-sm121", name); } -test "every registered family is listed for capabilities" { - try std.testing.expectEqual(@as(usize, registry.len), families.len); +test "every registered family is listed for capabilities, with the wrappers it also reads" { + var n: usize = 0; + inline for (registry) |F| n += 1 + aliases(F).len; + try std.testing.expectEqual(n, families.len); try std.testing.expectEqualStrings("nemotron_h", families[0].model_type); + try std.testing.expectEqualStrings("nemotron_h_omni", families[1].model_type); + try std.testing.expect(reads(nemotron.native, "nemotron_h_omni") and !reads(nemotron.native, "qwen3_5")); } test "variables: unset or empty means the default, anything but a whole number is refused" { diff --git a/zig/src/native/cuda_memory.zig b/zig/src/native/cuda_memory.zig index 563b78bba1..1d3a3b438c 100644 --- a/zig/src/native/cuda_memory.zig +++ b/zig/src/native/cuda_memory.zig @@ -102,6 +102,26 @@ pub fn cacheBytes(asked: ?f64, room: u64, streams: u64) u64 { return @min(room -| streams, want); } +/// Kept prompt states' budget in pinned host memory: --prompt-cache-gib (0: none), else 8 GiB, never past half of +/// what /proc/meminfo calls 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 * gib) else 0) else 8 << 30; + const text = procText(a, io, "/proc/meminfo") orelse return 0; + const m = meminfo(text) orelse return 0; + return hostShare(want, m.available); +} + +fn hostShare(want: u64, available: u64) u64 { + return @min(want, available / 2); +} + +test "host prompt cache: the ask, at most half of what is available" { + const g: u64 = 1 << 30; + try std.testing.expectEqual(8 * g, hostShare(8 * g, 64 * g)); + try std.testing.expectEqual(3 * g, hostShare(8 * g, 6 * g)); + try std.testing.expectEqual(@as(u64, 0), hostShare(0, 64 * g)); +} + /// Integrated GPUs require validated host counts; discrete devices use CUDA's counts. pub fn counts(unified: bool, card: MemInfo, text: ?[]const u8) error{HostMemoryUnavailable}!MemInfo { if (unified) return meminfo(text orelse return error.HostMemoryUnavailable) orelse error.HostMemoryUnavailable; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 0407030322..36e4812af1 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -216,7 +216,7 @@ pub const Server = struct { return .{ .kind = kind, .text = text, .after = if (thinking) srv.text.tokenId(srv.markers.close) else null, .compiler = c }; } - /// The model's grammar compiler: tokenizer.json, config.json's vocab_size (text_config first), the eos ids. + /// The model's grammar compiler: tokenizer.json, config.json's vocab_size (text_config or llm_config first), the eos ids. fn grammarCompiler(srv: *Server) ?*api.grammar.Compiler { srv.grammar_mutex.lockUncancelable(srv.io); defer srv.grammar_mutex.unlock(srv.io); @@ -232,7 +232,7 @@ pub const Server = struct { .err => break :blk "config.json is not JSON", }; var vocab: ?u32 = null; - for ([_]?json.Value{ cfg.get("text_config"), cfg }) |part| if (vocab == null) if (part) |p| if (p.get("vocab_size")) |v| if (v.int64()) |n| { + for ([_]?json.Value{ cfg.get("text_config"), cfg.get("llm_config"), cfg }) |part| if (vocab == null) if (part) |p| if (p.get("vocab_size")) |v| if (v.int64()) |n| { vocab = @intCast(n); }; if (vocab == null) break :blk "structured output needs the checkpoint's config.json vocab_size"; diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index b1e4476fb1..2477926109 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -160,6 +160,20 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte 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_nemotron) { // Bolt: one tile of its own size, ... before the text + const p = qi.nemotron.prepare(cx.a, img) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the image's size is outside what the vision tower reads"), + }; + img.deinit(cx.a); + if (p.tokens() > per) return cx.fail(.request, "an image needs {d} visual tokens and this request's images have {d} each: send fewer images, or start the server with a larger --vision-image-tokens", .{ p.tokens(), per }); + try images.append(cx.a, .{ .patches = p.patches, .grid = .{ 1, @intCast(p.gh), @intCast(p.gw) } }); + try text.appendSlice(cx.a, ""); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, ""); + try text.appendSlice(cx.a, ""); + texts[k] = text.items; + continue; + } 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, @@ -183,7 +197,9 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image_pad|>"); try text.appendSlice(cx.a, vision_end); }, - .video => if (srv.info.vision_glm) { + .video => if (srv.info.vision_nemotron) { + return cx.refuse("video inputs are not supported by the Nemotron CUDA engine yet: send images"); + } else if (srv.info.vision_glm) { // 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); diff --git a/zig/src/vision/fixtures/nemotron_resize.bin b/zig/src/vision/fixtures/nemotron_resize.bin new file mode 100644 index 0000000000000000000000000000000000000000..d708a5eb5db75c79ff5561eb3711c3658d53a481 GIT binary patch literal 23808 zcmWifXT{3niI{MCzTSVL3Q?Orc$X? zk`f`5j3J@-yr0kcuz%09pS{+;?rXg|)=tNiUeH{FC!}LgMyn0-Xsba6?KeoG69!Rq z$smy247@17z>(GkoT7sPd+AcZS_%kQLg@ihs4`$ERnP0;((2V*TRnjfJYvT$kZDr=M^az`v62{t;OW1AYE~cd*4>g7HkWpBG_k;COKX@my1|LI( zvmFYZ-I40-hg;6!2zE}u?S{LEA6kU>e_z4rRUQ16cjDaR{z4rIY2m!iM*8{gF|8I9 z(~)D@bTu-C?$yT7Q>kG3v)G%)TR4)_l9SZ-VK=4u=}}n691__p(aw(kWDx$0*O{kq zQ*}rF?#p~GEuAR}`}O0#$v;0q`!f}RL|~=BdE*Ku*z3t6007`6;T1FQNQ9xis)$3f+1VO&eYY(3g+y6enj(zx~arXYzWwqN_>uBULEO zc@V`!Nl;BmHJ2D4#Ro+07Ul00@?pNyNkv=NrS34?X5P!(X0_fxzv@@)Vt`-g&b63{a*9q#`5FqW)@YF;C* z==CCQX#-BD9ARgqe^TG}xwz$|LOV?d<|VxPK-J6B>CVwaN-+Q{Bl#dSo%xm*AcXvcf+}jC!l}+Djp|yW zDPndYZIN-MYDsJQ^~;#XEn~D`tTJsY8%&2Im1(+EKeFs8;ujkIxW?+8{NdOd(No84 zo#R^zbXwk>)^`0*kyQ<`WAR_J*t`+lEV4imLE4|$??rOZo;Cxj8Y|JIvkxP_Uc~S< z-pKn90LfBs=pPTo^KD6xw8(uLniAVb+g=p+ATt!3+|i1F(uF(FKTNQds>Q3GPH-bhp`+?f^vsyb zOfw@`P+t~vSG_~hgWu8D6{$4-c^q9`o=TpP38au7Kt)sCsdsdV;CjBbqLv@I@5nFesPQqmjk&*zrFG`|ea{`e_PIcP@@7_cDurojKWCTw z?z8#z?^ycW5!h_F7-nHcNKg^roAD;-9<)Y=nLBdd-9Y!#M8r8HK;I!AMSFAM=I{`b z1uyZ|uNG$eLUH|2G17nUpiOxL^L8(ILnjV(QRKD?Dw0YkU8`7fO!KEs#YlQ}Fo4GH za-`EQtfm8E#PlaWCn!c52+Gh!V`_VAboPUji z==Kw=?`S?dwe6I9;v2O?jI4 zD6J@wVzm9~t(!BgZuKV3*-o@p+LZR6TuG&VI+U?Wl}cW`=L$J@ckQ>&MosLOik031-bh(2j@;!P%VdP)(UXgF&3Gg`ZynFj6dmn(AVz- z6h*EmRItNu=7oR7g|x zE#$PNfTl=g(&U$Eba!MZ-8^@lzQkOmg|RN=S9F97cH7XXHZwBatW6tbRH(Lm49#@? z#_LW6@k2F+{N5T*-hc0B(OZ|}I=jBu3*4pM1#P}IY}bTrw#vJb?R9&?u8xqvlBBt? z`V35-qKmNXT}bkG!o?a#{JQB3iSI$kS$_?`u0-HOZUzGP7r}X3Hk2PeLyu}N?({ao z!f_FFM!MjU{S11R-Y;)WTPeNimq)VS6KH2$IITYvPH~RcXlkD|**2b{-7=?X*U7ci zU%bEat(MXGs1anV`JQJ;cJWOPiCnEviN8*nF7md{64@-fAvolhBiOjERq!Lymp%XU zp6PbJX0dVou=eK^IQ|-q@3LAbJ!+0iDt7oX@-+0vU&Y7!K3Li7hlR5eQR)?qt&daT z^|=&#jyB?pO)Xr!n$g?v3%eIM;QaM>WIE;<^(B?jFYP?C@JuFYxqGyaVn{dKlk|;j zsDqrSW11;hTJ9uCyRGzY;%qwSzl4H%l&Pq$iAOCB zK;rck_G79KoA@}41<&sXCBZK?@6uR|{y7I}Gv*`Y>~>sF+Xh9~(`bL`j)pUi@aYJ_ zzq?_mJ{*mGC$nMsd*%RNJ>~?o0%dq*u(zgx8edWnmICKfzR<6ZvD-%%X zDGdGVfCb84P_zz&`tL|=_esL*y%{+9AP*{o%5ZQ^5vqrFAUs-5=vMdxarMdU*S~x8 zu)dFs>YC`r!YuOLcbni9L+ZPOsp08Wnm)#vYNXCl=;s4;SJ{waR0Z@`Vmeh~1l9QW z@xa#&yw&d=mwFP-2T0B0sin_EOJO0J_3@oRbG%NxNHPx`mk%er?R4JryQ-g4y zDxQniy{vOHa{LbCY|j;R+g?GedkC5kg(K@@F)}a@t+hGmYOg`-#dcU#zr+L6Hgs$! z%q%5owdU0e}a*YNk_>$B&O(t*izW8YIgBKaAg85_VaoE*pR&!5K+O%w9Z7wU3# zH{}YF{rm*KXeTquC}$bpdzp+t79-o2Bl5s<95pt8n!hL4 zB3$<#LQ?Kx)}dS|tgS?2Og+}w4iHvUj)i5YkUl#oP-Dbz?rj}KD^|Rs+TE4p*ceV~ zheB!8qhR{6*OSioo}`5%&e5JhN6FxWAywqhCMhK)`ZiaY@|jEv>VAMWrP z(Q4lJQ(AO0ezC}B_YRSaMwj5^+;gmNsTJ!yna0+b$RTKm8q6!^;OWjSV4IIZ+2}Y9 zZ@h^9_pd|dU^t$C4@bc91f0{!fpX=2d{fAW&h=&tJNpN5dOxuCrYpiiH__Xd1M||C z|D>#zU!*+aA??*qq%}PWv_Zq4R&R18IiD-!+F?)2HFuDff-b4%gOt4Ik(rDfjrABr zM`C{Qt-J5?`7I{A{`Dz7_}>g3($|@LyEj~5n;s`P))6b1e#Vi>^_H^4(+}7anI86E z+gu!-z5(*9HX!=?CL|;~;m77{m_PhF*1q(|wz>DvotKMw`#4h5^5Ofe2vMhBqUY2r zsJA!a@W&aj-`M@(zL)T@Q>x+_=O}-e!6ZFANFN3zhNNB|9(gpTvO9xOEh$tNBCy; z%6=Dnv#)|#q{!o~+914IHy&5bH{tcO)3`bO9CRkz!u3CAOc@x8Wj|A(G5-$w*Jk0u zlN{VGE`|Ndw+Oi?AzVCdw9q+tDx#z^$Y$h4I@Oy=Yt|Q1^UDl+aW;*rC*GyVZ~m0K z(1~Ir94Ky@EnS{_ij0j{)2Ra!C~V0Dy0mNv4RIJqwd0@h=mCfM>4J@1$72{DQTai1 z{`zUr+svsKwriH;Oey`SRr_Rt!|&9wLE22$IhNv$b^X>V~Amutx5Sq*92 zq1u%X%^1QDm^=}M&AKbPI^~)6yTCbYms14uiA!Zo(FH8kw1M66?O-;QKUvU?dC(nh z4C%gIXnc799gDp%>t-Z8onz4vehb52-o;V9doZ|B1gpw-xM0_XhWEeWy}uL38!hmr z{tcaPmCD;0*Fb9iwOGaqiFk)zueZ)l}|3)!-YW$`LDDtk#pEQ(T0?h+Li`G1fuWSjNQm%)=K58yO_O?lzgMizBbq;dFRfAk_@;p(C$cN%4t2 z)fb+pwey8^#dQ>2dp(G}@Ry%RYvn=5i}-@EH~GfYE4*jPR{qgpE+697mb+-oPNurl znd#X@vIX(P96VZ>sOUEfij_z6hcQT;Hyx!%H1RlD2z}?XaM8Ywzz%;zs`#O2?LBPk zEQ5dAWAwE=g{@-?E~m%}JGD8?8ue*$=sa3>w;%PZHy5AR)nsCPk7`>J$fiAtT29}f zYR&6pBJWP$Uf9#X!>36{U`Fqqx6{U8C2~%!<4?bFK0_~zSE*m%3Tw6atc%}8TA8+@ zltz=Koj?8)9N#jIX-0Z76S-t|D!YO$U*E}8J`aGy;gL|ypNgNO#p^z*hhw6xuqZM` z@JnyRk4Zq=g=DnPyNScAM37!ui2#RZaAZF)Swlv6i0)&+JsG~OO(Ji>fS=S=&_Vn} z9o^}dPOa;c>9}?b^-A9$#~L3R8seJxLg*5V^6y07_#0^%& za@+rGxz4i1TvfYCRMb91R1-H!P!rrDn6qFhOWkI}G~Ghj#F%_`)8#E&nJk5wqm(f2 z@=OePr-Lts8_*cH4-aOZf{-~NUEJR;cS=T4K^peh-N*D-1-QKBG1`aM!ls1dy8k=G zd^^m?TpF0?6Z?dohL=%aLLT+($e=_2Zqet0ND6lfpvkK|DMrpg{G3nFw(+}Zcg!j} zIbb0bzZg$fh7TdrQ@{CxQ?*=KDUqvBvg5n9tMlHcxgs;URk_~vZUUWcqnY)r!z}e- zC_B3B0UM?9nQgs32x}E4Vx{XMV89yaT;2nX-~U0i)B!uXJn?6bKeh_P(OQ&%vO}p* zzEO&vmbVyHEhm(XE<(-p7+NWPlXeZjC(UPViv@7H*ANMPtN3IR!YJ)1d;juJc{Lzjdv6{h4{tFW+-=3$_ ze@~fU%zm(lau1gF=05v$@hw}YA`dg0DH!*0DLnt%h#P(82=%yx1`jvvsP)5(F_9=8 zmW+TonK0da3rlnEVD=mi)sP2BkCzhm`|lr{edrRecFN{2Lf_FZyraVwCFEq4NR#)5 z(cV$li0yNvo#Sn3`$&6=y>pm+mTsgo7Mi3!d>n;{`je{Q3IFmhiu-l%;7$#L_;g=q zelm44-_}ts8ogmp&R;1mIDL5;TWjOPqT7mC!ogDsS zpT4Ct-MO@S>1|RjjG(S{esp8H7Y)j>CHE>*ir=}Ll-3`j*qH`2OnDxiI;~8xy7HvE z?K!{E9LG(@?dPhOT16eoCAr0-0o?tHgGj={NT6pjp4BKiv4xu1EPl>6R=GzJSB2wY zlB0!DYqw*?<-J%Ex(Pk+j=^fAGm@&k(QjZNo;kxKDVE z9AP=mJJ?}{Y7>QhuO=gzB68cAPDvl*sLnf(HV$y1AA=mIF3ywcUf7WQn!Oa}zMP&v zR-;ja<;ZwWH9uzjl)tF$;!1K!{OpedJoR0hXhi?vI$ztX1!q{i_Mw51g0z_%Smn4# zHaqGWTP-h#-$gUAwqr3wG7}-)rGu9*cc3HU48rHQ!OS}lI{FDH6!*0SBcsrDB^7=@ z^N^!jg=eNOU^x2|=0AIZi6^Cmg?qKIRsJI#*waiuI*MttUJ6Nnil9OVAKF%MkuqPN zqiESnBx`Y){yV2f@BU1p&R8XiE>t1!^F3TH^)9z>JJ0n7@8LEcj@(*z0589OP0KjE zKihfBg56s-gsD9;Wha~OF{@)=S>w0SSiXoMrD{2T%P}-R--t77PQmM^JNVISc--v; z+kxTu_($BYoag9kFM!YQTx8EGgY`y9;q{B@*mZ3K&8X>sM41`_G^6NP&odHU&8DT4 zL^4gmls4Rr&M2Iu$I@mLy>35UmE1xFi{{a#c~dEVvNX->_l`Gfyx^DY0=VNB8y@Pc zz>Qq0bLX-Txn4`g=$LDC3Kn`gGOZaRwolHN1(+2u8!ai!iBiK-F{^D_u?^N?#_&ur z!|oZ*Q25UsIxH9sr(*D+DHcJ=nUFomv0zU*d^?_@+UpstXLO-TLX~b)@;NoxWU8CXYeMl+HgzL?`0oHc8)-ySS5G4g(Lho^p-M~p zWavA=AQSTiZc5S2&{b9GNTv21Ve7rW#kuB(-{( zJSM{+Ne^>tEb#oIF%}CBAJgbOn+zRQ9HFWKWFDeRj*7A=`#j;OyJ#r4~OgGW$*-WEOuZdfnphk-`nC=5!# zrI%?~JcHxk&JrYDYsQ1#-`G?kA$0DP7FzxqC%pJE6?#8k(c;7gdRSLP@lXkU>>(Cup+W+k>U#*tIJ7>G|^&i#w;u3L> z*!w7Jy~AC>fYAADj*b(1nsAq07~R5pk4s}}!x$_{ABu4*6LCaJ3oCUtV8}@`d``H4 zeV<)WF7FG?Y2k>}O~k{+>Bv{g$HDs*7%b-GC(aKPDntY$cQf$`TEqFcuw9}R$+GNE zhZRpbDrw`RvT?~f&(Advckk5_GWhlJ0|$Yy6=pJ-g0$}P&B~4jeC$S z{@WA5YM4lxA<54MKZm78ZpiPa2 z6DTrNipJH|@u(dU{E^#M?)>JNXouAFoYeAeL6r0_!Q?S>1)i5Du{vB~^-8&{^YLdE ztTr5NjdP&jxe6JZ%^)}Y0v7u1!DfYHFx}$}$*k)LR*b}S%iGvhnu8|AQXE@Qk86Q# z`0wmD#I=6I_)G2hb6i?@FLDsAH2z2OXUpl>Mo#U&?$D<2Xc~UXkM5s$rLUUQWkBcy*uaAd6&^Ny}myk zI9JI9Bd%~BJdj_wP@OIPt4q*mP$ZD@t`KZWab#C-XEM1dee6s2KpeLC%ig}21i>f) zn)CLewCD^JW*@{AT!QyWZ(OvGgq?c|e%Qp|eMbtKKINmqxdweEZP>Z(527DG#jN#N zsO>1iTd#Q9W!^9ELr*>Z6+R%*nGD*!K7m@K0;%hd8!cAzr6BS5DYN+m@vG~oarHcM zys?6!vehZ|x-6al-o!hnaNc^rfk&Qf5ly-|L@?#67UONd1#Vu8*_by8EH1i>NpCA- zvse6P<9sKfAa@1M|2vGB(Ys+&Yl#I}ZWx*!gnhewklb|x_Ud<0rebbE(7Z*)*4;hbqwFWfRF4@}wA8$fcgT@Zs84+EnmaK@^CNgFilrPmxYI zs2GZiUt*y%{3hOu?x8325mrAb$Cr!>tgwHHVfUWm!Qnp4%zTQI#wW3Jmn|lo3!uwo zpD4RJi{>WekY`{FIYs%BRK6FDxaCP{2kc4j#a4GL_*+?Gn5WjWvG5Bx;EKM5mP^twE#a~gM+6gnr z3k#;qg8bJmn*OhelHQA`{z?X^xnz)GV;I?p@147yT&P3EfihO^A%zf2(w?=Jdc%~d zE3=my92-RQn+y46$qYW`-7P-q`V1~#zKc8A)QjSrlLT{C>|=s2VAiYSS>jw-cqPoh zj$dPO#$qkfW^cqJjqNa)vCi119s=IN=JaVZ~rN7Z8NnQ|B| ze}!?;zfe*vBkXD)Ce(KyC){FZ&!o;IQODW>3VK;YT4Q2q)S^&oJRL&QUF^wD!I6wC zED2$oX~gnPwCl-C+7Ug2>XH>G*r#61C2Z zq3ISu)Ud>xW=?XV?8~R9cltikEHR+cYE3#NoJNcx|FPC#z1V66Ia!B=r+J5?*4nZ13&4#<3BW|ybH{GKLmP1456Z~x)+cuPFU zx`O4UfjEBG7pKMR&d9!lVB;LTQLKW5ec1|RP0eSqJ0+rf=a%+wj>c_G*7ai^gCvlBCx-p9Tu=rij9>1?fj2OB)#DKpU@ zhbeihp=rJsep!d$bHy9(8@v#$9*J)uX|S1f2lCfa5N%bAT@Pw-UAhh(@}Dtk)HiV+ zEh}8>l7a0~oQ@1ACr^i5ntVKqKCHe^0?izf_KT+e-u|?IlrLq;pCD=mAgcNtm&K4xExhpt(#*}d_5vv)V z#*}s?vR{Y$*yJMd_gOjyBjo`Z{bRTz>xdJzu8`@yiQdec=&HPhSfgxwJ5~XW6=e`* zJjTDs7W^pc#)X|jgw0!3gp(bvLnU)F^^<);e=fhH>)8c#S}KOhmfj?z0r6tz=Bk*Z z+tc5xHe|kDNW)DSRi-bdc2#B4c6`Sh{(a>YhR?aBMFdZqGMjJTep{6HyGr!__w(Gf zHzZh8#!B|welGj(Xe`@R+RFm2bhApuFDz~LR5;zx!y?z^D5*3>%po_Z#{{A!ECdSQ z@8DIb2t8FC$-j#*Yn1Md3Mn6QM;EDza+g2V`RX@^Jy!^3!-d|=FfW~{w*O_>I=ejRW0W1NHq! zacZ;`gsUC!a84+ci;}VXNiw`-?;tY$5vD{oLRj?-`x;*0)S*5E?Ei@e<0OPnT8hyp zTT2oVZ>jqnk>s*;%GrOH9?SUB8h=;H`sqTGUptY=!;$73+(LPOHRwfv21!QDql@;R zxyjND9%7!#FSo|=$14g%=VqDbjxisflRl?Nr{AC90)=1JY|YX8?CSF@b~vV#?fIgD z2sKT}e$vDOO#$WzfX9=LE-)ET*B)GR-|{hER$@%fl#SO^Qr2avF>z=yl9uqv$& z7WW4U9aU6?efE-UQFT12x}2iSxkLjt5GjvNr?OQyDY)B>>@DqxrJbYEvHwx%2@84` zX-2my=Fzx~^0ac;K(gxUM-xB3VM{_}Ct-4KNk&Z8xmw2}Ado2z(rL6Vl>7 zrQ>xi%2_i^YX9Ki@_xd8?}i9tkGJ5+zCQx}9)D)r>rC79e$WJiO8V%0hkDQ6qS962 z^t#!V9#$Tu8F2^6wQn~~+PR6ok6lehGj(b9A~`y<>H+`pDu<`MiQ_AFhjO2wP5h$6 zd{OYE7Qxn{!OY!eIupLR#KLq7*;$WzmJs-Y2{sOd_1pz;5Ho%qGgGVxwZKIE)3_0T z8TWqKO`c=hiymA)E9RatLAX_P1B+Bd*d99+3Ynj1?wodN zoL)??gOh3f)i~N^7eV2sAv9_86$%pVp=}~#QrxtT(o%$UaOy%@S2T@ysw!>QD&@;> z9pf6#hJ5I5Req_kPGsX8E@}^V)w$MnnH7WtvxK#G*zku{Y}nNwOz19yyy78P@&uY1^b^(&R}^-76r%C! zBbwIGNYDQ(rsVJWlpaza?&K7as;!5Pfn%V+Ww$i5NzrABk#=m-fN(a( zDvdp||HlNa44&Qwkn%G|=7_^6?>-CP90%+j>JE1B8tQU_U_Uq#wMKEUl_^AV%rh7q zYr!$skEmWbi5Z>JAX{%u3ip3X!@k#3fPXm^&*xN_8&4Y@BIxqyAj;qCLCT3|>A&aV zytZc>{oA{WT8C-T0oCcGX*8M?ou%pN$qqhjVHtlv^bQ~4k}4`KEEY7yDzbz1o0*-Y z2Qzz>%#sZ%S^S_*rdK8nwK=2k^5R_F-n{{T>`k!C@HqMpu!V2v6&x6O6^$1H@!@3z zev7%>7?az_*$75f&))_PLJLwj0ia-0kmchXrIJ!LtrvBP~{Ztl(V`c`v|EnOnh-K#`yC$n|N_f`t-xacsWy)JC?#Z2`|q&~{Xp)?vPW0Qo<^Pr7C$G|oDM;#v|(LqD8EDc5P0k_+9GI7hr@A07CxhJ0pg zP|i{%%GZ)0H=UjpeVB% z3l=uvnWVV(Opp_5uTc;xFBmAyR2(PlTjmA_*?W}r{uOyg7160jx9H38NE&#=mz0aH zP>iiLNj~36D-}VTE2ffOm@MsW?nm}-2GHXjwfyvuNPhSECjRYtiHP}6TDrclT+lti zfIapNVLG`*Oy_|;t9lg9(z2ei`AJgv)2oUmgC+1a-iQUomKeOv5ed>hFp&*M_k|dw zXG9|T-)$U@%fa@Ir7)Chz^TP;n5;Ea7*gVfK_iax@l&7bv`;A)-CGu!9c zVN%#ybir-tKWxz`NBqSGK46wy zp2Ca|WFPaB?*IErYFbaoBrJsjG@@vCvmgBt*Zeh0j#5aRA$?O>K+#Hy6!l7)uHIIr z?q`2^taKh%{^-c{Glui679F|Uc03gfeV#2SZqyMpx-4Nwjt8&}UNubr<3Px{O+@?$ z87$Wwhj~@nXx(cJ`>)nGD{#YuAvZ9gCK4J7{J|OPNo!G^CL@;mDc&X2+!Vh9 zUs_Ufh2}_~qr%W#x`?u7Qa4B!HBdqSoX*q%M*;CYh!^`g;(JGDG-ig zR%`X?E)u4s!EyV2JXlx(^XM9sH@-ww<2SS{e2J9(f0@lh2`o!|!K1u->G!wyG@-qW zPKKmW$?6|IO|wseZnP@VgPHGm%j*a}()S_XR1?L& z#*E~C`gysAUav$M^*3`}&JAWSqkY+|t`}^*Dlv_tb?lV>Nc2q6MqY&)Xo41^znUVj z(H3rt#r2 zUS5{U+hg>a28_(5qhAwf*_9_Z_ZZ?5PN;&O!<$wBlt>357G9`-P#}L-V*3IPG&<+oM0brhqIlrB`od6 zM>ce|91Mz7a0c^{n7slGO53rp(GrVi*+EU+9bcdN!plDlQ#Qw=L@^aRs&eqvqZpoR zsxV_@fP?47v%GeDSem*!wIQXn7zOAXDN71*a|B6UDbiX$`^UWhiTmEznh}d2S01g+LR#p{A@a# z@caU+GQGu;k2SHRYyD9fJQ@>k&c@}mWoV?WP=90z{Uf$Ww{wSesxO}HxeDJYflxgg zjdlN$aoRBpiE@SL6P2T4bPHznmk@?tkQUCmFj9E_V+B_GP8Hc~cq7p1AI~~Rh4HUy zvGnn=7ugJSr0@YJDSFgyvUFNaEN>$nsaj6OI;p*Ftpyn%3&+tL?*=*dF z{>oCewa|j15}Fs5O{?3($g9ViPKw=|6A^1^!}djV-%Odte~_W!Kc`WGiYooQBteF1 z&$+#XFYgK5#K&BJBsy5_o_%8BLxG#YYIb{L0J|Al%EHZiS&YPJ*oqcnoWvq5{X7wG zR2e4T+lGMu^SCBQOuFI&(>Y-XsZGQv>nzMEEI|D7ba*L>kTaniX0;7)4}OhrdR^$Y zk`s;|l#CCf^Qo&On~r!y(=*QqTH+H=p(Ae4_xWB_m0(91CKl8syOHLK#4`xTRY^KP zlC&N*aziy4`jY*fhup~I0S7(#aq9_uUbLHNRIRR{`=<)K66(ZkLUY(KwSVlRQTC?fwQA}TlwdUyaQuNur$EWy&3 zl~A{Oj^}3IkWn#AxU=5}9kYfFwCclhdibw|qJ%kAZE=fi&2P}ddoE;ZXH7m9+i05Q zW?I?3lRRp*XzDywiu%`|#x8rrBkP0tzV+pN=)Giq>A&rK-HqR(=JgM9qdI&9L(}xx z6Z0_kc6cMRd#V71P!%+aRB&5%Jcgu-J!gfzVz2B!1SVd@U>{HMj6^u5u8qaO(J?q@ zdJ|FdnFza3g3|-*(Q5n(NiFp_R@wrk-Tj68ruP@xT%3j@1)dbV>jh;t6;j`fMDl$b zLO5 zo@NsD-lOKz;S?&2qNi`YXl0r`-JfxQ>`rW_mf;5|%3qhpYtA75(*x<3!wbGsRf0A= zY~Yg=gE-nlC&J3wk^Wj zam(>fQ*-1c-=#TYmV4MsKTeS&A8=WhZwCEObF@0 zs!loKl8PQ|PgD|KHWl~7JL4#CXfQSPgwual1L>2sGo7}yAPL{|G%foet;^S?$7z#D z^yv>zwH`=S$-UhCU=lxnY9EjD^5kbOY~xc4w~M@|y9h>g`sU2H4;5T$vScQK$?S78 zXZ_#yu!~w#u}WbX1o6u;*J;ThU*1Krl1+OYwTa znZ9LgcxY`pO6_AfzdsFh?sqlH6E;>FZri2c~loK zpHNKaB%J8>%B%Fp_5x+<8q<^g24wrsfD~6xpibHT^tJROw@Z4@e}CqDTJr|p(@%mg z+|w&k|6H1Td&*eG>?g5nRTG(EcQ`AKcVGrb)7e5}X>__zh2Emca2mM+W4ca2aMT%3 zw%EdYfFD->5#NE91R~>35U2e~;Ogz~Bjuz%~vxL?B_lI>3 znT^GN*TP(BB_wv5V?w0|_RSB$1#Lg}dqxPq1uVw3ymGE`ag=aU_!xTork3tZyG?-~uafNHa60DbMx*_9kbD27Wb$AY zWV8){=)^>rPTPdqytBC4dk_bcoUyw;0G*E`pfWsBoE6@~>mMSl z4|ssHPpYwVV7pwZ#ui7Uy|psSy__oh<5d3XdL)nv^NR!kK`jwI0n$K|<8 zMtTc4hOx-Ydzn^kAlq6iV&}fUWn1q{qkr639POQhY0su&z&3q+zF>tsn~N|$=7=61 zgxMxh@HiWTukSg$7UsaIQ|z#qwP1Q|D`al;V*jmD{8LYXqG~zhLmMD%KZ`m7o{-Cs z1k$gMr-NJl$T`uP1S@vX0^b89l(L|{A~EV!m_wFYvuXdoPM&|_CpQnv<~!CL;AKA@ z`Kjg%k#dtjH1JN6&M27$Y>~k{)^{zI9qlP#*Rvk5XKG_{ynQ&PCojPMk=x+ewH6;A z=;5cS6+XE;And0NI`;Zul-VsD7w0NWF&oxZwRn892B}_ch)wK)p|6DSM2dniHh!e& zcW)_ay@(}04|mGjA4J+8qbcEn*r8lzPl_pGFV1r>**^lcwydBTVm9>fhzcpTO406q z5@K!|#4ihj`J{IXc;m;$+z$&kiIkd-iZVS&VBT%SPA}iVex0jj!NPa!b=z~+@7pKl z$>-o`#B50Et;0FJ6G%RD1pBrg!goU#G|L2FWRfo~iFuagY!Ui89-^$d0%qqvV109c z;eyB!!rkl<-p+l*z1K(bp2t4)Holr9-lWm1+C=g%52F?@PtpiEN0(Pw&~)K(Y6;aR z%Yl=qL~knHE>fc%k>B`9%H#8O@ABVC-$lFnzZETZOcTBOa74$nZ7R!^8OFBu7qC3t z09KpU%p!jGN9^Cf?D0rN7;o0ZK5ISf{wYLizk~R5>I%GjoG`Y-5n~mD@wz<$CMV-D zbVVkNyDJfU;yI#jyh2?@H{v?3V7`_Gw!hfLmJNPMkK~?EPf;fIRfkcW@PR_-asSlA9`VtM#zhI0-4=nrj!p3Sl^)9Rv%!`7B6;Y){}jhYgj+=eD4TU$d1O_V=6E_Faawj8e!@A zW0a?? z>D?%y_clfo%$|^5z1Zgq2&Ss#*Xa5&Z|d{2rLlz@>GtEbWKyL=Q`WE$8keb2_A12D|pa+fO%P;XO5~q ztW_ioh3u6!xctBt0y7V^rA3G_Ln?+1%S6+U zY)l)F2ZxR#Z0%P9+v=}aoGB;l7c3|Il&2s}JO7Yhb`d)XNyYRvl_>vDIIYt3rE7~^ zX~X)RG`DX9SsYtQPZc#OO>zo_x+qcO{h?$XTg}~5toWLo9sFGKEPmPRr)ZLGps3;I zhg`j-k=hH@_c6^35pz28h>g*H&Dz8Mu-=>jP`Eb^Wx~y{Z{3RmHB0;*dLE0GIOBwj z2h2Wt^(hep+R`xFd6aPT?|UM} z&zbaJi4O^D-6;Hm1O51NmTV&|Xy@LYRCj3|g-WfU?xiXeb*zLRHMzsjZ}j6IT#xbn zNelSc%RQpc8{I_sygj#P>U2Sze1YKOX?a#wpUFxt4M2X|U_2`whoyIC;Pt7+$m1*F zBDV!0KTOdt>ImkKb-{cVhe%-(X1__rE1hhJ3i8n#`4B67t8o2FJtEFG<3+@4ES(}L zTpw*A(t2G+oxkHKY)2@sR@p=sG~K9cku{w&I7l-4wvwEY9(^Cam{dEbQ+%%ysa%w! zYY$)ZX}@lACn-1nO?ETy?o#5*pNT}rH~p`H^M2&A>*Ki8jiNNPq)8fDQYgOXtu%@D ztzjf7L_>=tD=Xu+A|qsEB~g<2o=*}L5@}FGdMK498lva=2dY8GBez;Gj{|-z$3+sehR(-`sY~G@Yk=aF= ze}#VkE~m(i!d_qcC>@TnrR;PQ>L^@E(T@|n!E|bfW`L_mfGh2`&4qPqY8bCxyXE z)Hhm&9?TQRwDPJ2fjQF*1Sf#GR+=EABMm+(8;O zYa1!3ET^nFb4Y*j7-|qn(3`*~TrNJHyO}%lm4VZ^N07g`>F`J0(=Ht%-@tWDZgCiE zQRb{@4&Or{HV4Hl(}_@TYtsCX0dO1`{l3cEH*X9*}$+h-qE-rdc|tols>76q*-$q)-O{^4FlSiATEjx2qOckZp|32eu(j&?Xtlh$kcJ6>;gP#JcuwBX^_ zOQ`qt9g13am0DC%Xp6&nDr^cRQ*~cbtvW(}6&6&KC?cCLGbreVEJ@_Q;cu_!aJN?{ zc(O=_U#!&;yI+yj)fKmjEGli-ydkIAV5`ln_*fDPn$*CIkM^BPK9nlyqv?)?N zfkMp(QP<)su0A({zgeTrCtM2=|1#Yz`dp^%A4S`-L3^K)2Pk%Hq@c}T9l4Tq{~+$w$r zJ(V_?_BY`Awznvp^%D{ra(aIhf5W=UlQa}^Y0u^aGN?=^;Vq#V2hNgOqZ4@HDCYE9imHkf^o9Kk{g20SB{~4xonqmyoDNmN2Qq73 z8kFJ+a5APGp&+E?R z(Rxn&%b$H*Z~iCop{jb_`ao57wK|gRw%yI9_X_y~(|7FO&m*7|^^3)&kHi3#B~X_$ z$DNn%kecHNd43AZ{ljq4IR%%W#bAVYGKOs}z^(Q&RAm(l_1AsmjeLpf%^y+osSg>h z-MC#Kr`OY*hnbHbQu+243V&2g-j8C*YhM6K%Q(`zIc{Vu=|sr`Hc;m2*_8TThoWE2 zrA6zdXwmFyzU=&0zO1X9%ill11Iv2FEv`@Lw0i;79*!c*%!Blm&7(=auXxoAaLXkT zIujOKvY#QJS@2;^1eTq^vq`b&T6P;7)+a+|#X(GboJTt}{-d3h8^~|MPzrY}<%49) z#67iJS%G;9`!sbRl-o97ZDkOYey6~CTRl!KJW7TXxoxpHn*3Ypo~8rmwtB zB7`?oPZGDiFkwNbTG{M=Rh%g|MUZm@RthzV{H_*s6qK=Pq4!9wp`7Ndiz35cr$|;b zh(>jK@jYc$#UDl=WCr~|S?p;YY^*qeiu4d9D&4^7U2T|GVguEVM|9%gb@CPZ2$yIF zda-E%4fC_&m6tq4Gs*&3)ri6Pt)&Bpf$sRcEm~kPZy?gC4KX{T+0vSORCqj>s+R=N zTy=ALZZn58V(xH@)?2!!cOsYyAA*q6tKe(pg^{xoQ73SyLy9|*_;(pRZJ*JfJz1pd z>PzXbH&J)RNOFxNUZgvZyWPCN1}P}w`IH^FC-irAyV5aDq6+hJ5+Lx$B2|YnTD>@t zbmP3}Z1f6R?lGLwLR0ys3nV^x_O)o$#xbZa-7dVXfiUdO!mROiu-AKq+`snJzpItP zug6nyq6>8g>QLq;Ap^2Hid%MQi|f8mV5i+5GF6{FsM;Ec0Oza7sC^Er$tl?K;4lqW zyG=4fggSco5z?BYORuy(aKaHj)GV=B>tZOIp!Jx|pQQ<@j7TWO+`z5+M${axhxYUc z`uc@vz|?3m$v#O&`IG5;!zZ3EX~i|(D2vY9gtNrjes=NT41~*jpdh{g#qPaG`&7xw zQZCb(IX7r+S|nZGdWx)XXj99XJpL?Ok!#wkiatmE#|CZaWyU)5(bBmW58otVMcNxQ zeNw@&T?G{Lk*MD(l&oYeDQ3!K+Nto4NAH%+5p9H0m~MaGwF0b6cR(>4&=D z9Ng`F4sVr5=-XXGYEB6>`=u-8=&qs{>m;bXG=VG2tq~Vq8^$UhK4dRz)!|WThJ@>( zXd81Csp?fody<80l>qUl$YhE=!^nlFWh1LE)l~rDkIbf zj0i)A?ZodKA1obm1+p7kuuH24J>j#+NwzKyi z9v11V20rZ=)P9G;SlHPXF6xBGp1b%kej)ja?vnnv3$%34VIf0fNUEK2d`h#jh~_%7 zsf9mTosdm?-+L0Lm7~8^dIHQbO@`LGZtQnnAR-)KFFZsGh zOL@(l^Q?NuSbVoJhPL3n`Lr$>8Vl~jrLPP0q!mX-3R!T)Li$)L3tBt9~n+w z$5QwP@)U128O};V1s*lo1bbY<;JE1;OgtV!ewB<~a=Rk)34cbj{>`Vcb3ADLHyvuv zkfDj=eEG~K3vte~m27S6OLi{Y`^P5_@>(kB{Q}X}Tiw!{^!2aU(GP-%U^`KLeHR(b(N{5A9DQM6#*1bfi~I z2IoU)|5ht<{57Avum0kjjm@}@rW!L(5VPsaM?hiM2ADT{;q#6QXkAo@G`lW*h!@_W zWq0VWU4+2)?IkB|70RqE;Zo}F#M#@Aie8@aW3&HC!*ZUX;2HCQPf;3TJZn(8Ktj)d zbS9oGx<#&6{mE$ZN*Z)rlD21E=5inH#QO7mSiW%?o4Rl;My|6(uVpm+&g4MnVH;le zq+yk`8EvVJ5Zn#^ROPNkhA;a0`(tOhLg8%PBV`-eKKt*iP^1Q<9f#nR6A!f*AtQh3 zoiG=;3H_;N)UvUb<{M>_s%R~JKG4Ry=AGh`CBus1Ui-4Y|45+yvL@Ce8 z;e<@^AyCr32ienC6wCi2b64P-Qoy zMyKMVz*@`-dV;+AXE3!Bi+=3#rw8GYG*e2C45thw*TEb-LM{$&}GhG9W z^FhO?3oz;=#0h<%(SoOld1^xEjBn9y9dFuwK%blpMo``S7_MfZ!WBYenA+B zB}G1vzbp7`WJ{s(tqqbVWc1EI`XqY2FPC(jT_{^{T^+d8&-ee0=WC0+#2>FnvKHa} zOKJHRQcwNyF)1DU_B_DI*6)xWmk8a>nJhOffd0BUQ(mDKskwdS=dEq|-N*ZNf0Qt$ zSn`Z@+_b>oGvSzOk&QPJwOE$b3D2Ai%rWt#YezE3e6tnZX;Y&`yPxvI&n-BT+I3e} zoMdw>U$MUJ+UT$ehLXUH$z5-N*y$6d8=uFKu}(w+|CQQ!fIj8UqYS6pmJ-_kgGIa^;+oYy{sklqvIWHA)lg2vn zyYacAD!m*gV=)BXuB%X5<%TxHXbk$^fva5wu$p2;hk|m+ajY*X2wsjTx#6_w)jj@X zSDQF@k2HJr`Znt|8-cJRMyUMkjn_f(_!}#R5IMt6wFb66H9VX~@(7HA3zfyMD0G;OWGjK-IkYIc`x?>R~dMOL&(W&DAU4vP)9{lkoG%bC%9JtQ1IjQEQon4*3KA0JoXM9?49N_FDbHfPd#=}Bu( zFQU(W8w&$I~8lB|g4Ljh(t*z~%;y!9$5{*m}Yb@jH{S*P~z@qKjkGcv!%F`o}$w4ywk``d+(ENYpRbO%f*9|KykJ_2 zXz&Gr85w+r>_)Alittfnt^0=i4o>E~4i4b1&Bt|>Jr;s#&GA=8K0<$41-YgVFRRlWG6r#WA+DTrp@E}$KmIh3|ClGcBC&oy&Je9-y` z@#7y?M3Xwa*!%+xEJ|@WX67;&O76wLfdQClk&J<&Yq*&B3KdRSXi74`=$a6*ZOCEL zIb1=_{(^7rov=5Jai>e38|k~VHU*i!<*IiYL~bV&n7i8m?Eh~jR<1X~x{@R4TOEl* z=X3C?ts3jo^KsV12S)$sV!GP8;`5n7R66n&4c~W#{2W4Po4{J#mC&V@nF{2+<{7W( zcI6RU)wqJ;5RuG(YB;F0NATi#;NRv%xR{mUMWk?FD%`-|1NL~qMNHn;j^4N5rQD=z zByXBbn!|l)mH7th>6t8e3nb`x%Vj=Lem(aaZCd>8=5*Gj>B00W%2}}K6d~)FhV`Rs z@GGhjS#=k%dYC^8*}seBTLn>}M*&@RyG9F#ou@ZN9#r3AOu2)k==sOzd`o@=mumVW zJ~UlQHd%^~woK3;y`lYz|!vXOI3CNq}fTO>K{B(;VDGRmauVE$hY}6&% zX>*q3zHg($&*xF!ncqBPRw92{6vl_PwTp)o2xlYPn_ZgvjI})+i#B&n^qsLrb+b3_ zNG4%RVhJR??m=NOZL%7FLAtqDP2TEk)oYM^T*FC3gR@JC(Pd zrLMy^q;085{Uk*Ock}u8R4u;QV4c_?PmUdzjAPGjB=E3jF1-C0LDtX$C$9So?u#pc zTP0k=8sTPGfeCCV?AjO5tcRTY``DA8*)`g0d7egh$k3t_kz8HFUmP|zSfnRw$h4y? z*n8h0uwJ2qu;Oia>EMaaZxW%`a|ut3>yW1Z4%LExBXVFBQ&leFb&5fBcWfp-H_RZh zNjQZx?4%dmfLxZ1Cbxc5QEx&P^Bee<<;zV(vAF@3s2xZ4n;`girbF>p6)Ynf@u2)A zIxQ>l#GsfB?SIT&OK;LC&wNsEkEi~d9`sN4YEsgjKoQ^CIZsUH1~W|gBh8Z{NuOz$ zm%9z7fj(#(5+~$+3m~EXRB)RK&Xv1OP%V?i{l0@dRo#IKBE3j$#Z_V79Zvg%{$Py$ zVv1`XN~@(f-)d{eTPFvLC+sX0O^Z6q?#kygrz7$>ZkLRbFI7mN^%%6W32jMNpsXz; zaCv8V=P~`C#8ZuQb){=GSzTTd%ab06 z-9l!izjqCCZ`vcO^BlfxNknVOdN$5yr7)jdL)R#i&L6uo`l3Lbg#>gerY2`$ym@GlgKm5j$^7>SiFGF_qb$n=15U;Ae zCf0QRPt-HRosH?>>_^W(ka1iNujwcI*N<|ctJd`G#e#ncnbh*MpRqk%2ESgiGfx9 z`#NkK`yLWLRj^oBhRrEgnY8&A-V?Z(7x=kTm+(#RSDWa1ff7x0YU4xpCGtHN<+%BD zAKl=_ne2>K1e?}f&Sp5OqQfZ;#hPLizIlvW)9bJ+I~s%c=dmB{k!<Y>%^cjoU1!QVLV=&cjLV?Kt_~0lk?a z(7aNEr1)Ey5I7t*%2u?XAeJUwh^4q8fpo(~sPowjs!@8&GuP(x&7l_jquwlC<$-@i zwtJ%4>y#emVW)@(pO#}w_8#0CeipOk5|JL51^>OJh?9RUs zcZ8tDUVy6GF#Xpbn0~cHUH%Np5*#ygKWfvPY-w_3IlTV(a(>%5qxeKendt1=Af{p5 zz{XFShyhd9!Yp$=%m%r@X5Tre?94{uf0Zb%ZiUPCCju*X4eRbpi?8c9^Yq*u{E@9E zJ!lFPc7H{D5+BTG-M%XFmN#M&Khjz17D>n?%tV997SxP%#M{*YP!zcIt 0.5) return f + 1; + if (d < 0.5) return f; + return if (@mod(f, 2) == 0) f else f + 1; +} + +/// The patch grid (rows, columns) of the single tile: target_patches in the reference, both sides even. +pub fn grid(width: u32, height: u32) Error![2]u32 { + if (width == 0 or height == 0) return error.UnsupportedImage; + const r: f64 = MERGE; + const mn: f64 = MIN_PATCHES; + const avail: f64 = @max(@min(@as(f64, MODEL_LEN_BUDGET), MAX_PATCHES), mn); + const ph = pyRound(@as(f64, @floatFromInt(height)) / PATCH + 0.5); + const pw = pyRound(@as(f64, @floatFromInt(width)) / PATCH + 0.5); + const f = @min(@sqrt(avail / (ph * pw)), 1.0); + var th = @floor(f * ph); + var tw = @floor(f * pw); + if (avail > mn and th * tw < mn) { + if (th * tw == 0) return error.BadAspectRatio; // Python divides by zero here + const up = @sqrt(mn / (th * tw)); + th = @ceil(up * th); + tw = @ceil(up * tw); + } + if (@mod(th, r) != 0) { + const up = th + (r - @mod(th, r)); + th = if (up * tw <= avail) up else @max(r, th - @mod(th, r)); + } + if (@mod(tw, r) != 0) { + const up = tw + (r - @mod(tw, r)); + tw = if (th * up <= avail) up else @max(r, tw - @mod(tw, r)); + } + return .{ @intFromFloat(th), @intFromFloat(tw) }; +} + +/// torch's antialiased bicubic filter (Keys, a = -0.5), in fp32 as aten's aa_filter computes it. +fn aaCubic(x0: f32) f32 { + const a: f32 = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0; + if (x < 2.0) return ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a; + return 0.0; +} + +/// One axis' taps: output i reads input [first[i], first[i] + count[i]) with weights w[i * ksize ..]. +const Taps = struct { + ksize: usize, + first: []usize, + count: []usize, + w: []f32, + + fn deinit(t: *Taps, gpa: std.mem.Allocator) void { + gpa.free(t.first); + gpa.free(t.count); + gpa.free(t.w); + } +}; + +/// aten's _compute_weights_aa for an axis (align_corners=False, no scale factor): fp32 throughout, then normalized. +fn taps(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Taps { + const scale: f32 = @as(f32, @floatFromInt(in_size)) / @as(f32, @floatFromInt(out_size)); + const support: f32 = if (scale >= 1.0) 2.0 * scale else 2.0; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const invscale: f32 = if (scale >= 1.0) 1.0 / scale else 1.0; + var t: Taps = .{ .ksize = ksize, .first = try gpa.alloc(usize, out_size), .count = undefined, .w = undefined }; + errdefer gpa.free(t.first); + t.count = try gpa.alloc(usize, out_size); + errdefer gpa.free(t.count); + t.w = try gpa.alloc(f32, out_size * ksize); + @memset(t.w, 0); + for (0..out_size) |i| { + const center: f32 = scale * (@as(f32, @floatFromInt(i)) + 0.5); + const lo: i64 = @max(@as(i64, @intFromFloat(center - support + 0.5)), 0); + const hi: i64 = @min(@as(i64, @intFromFloat(center + support + 0.5)), @as(i64, @intCast(in_size))); + const n: usize = @intCast(hi - lo); + const w = t.w[i * ksize ..][0..n]; + var total: f32 = 0; + for (w, 0..) |*v, j| { + const x: f32 = (@as(f32, @floatFromInt(@as(i64, @intCast(j)) + lo)) - center + 0.5) * invscale; + v.* = aaCubic(x); + total += v.*; + } + if (total != 0) for (w) |*v| { + v.* /= total; + }; + t.first[i] = @intCast(lo); + t.count[i] = n; + } + return t; +} + +/// F.interpolate(mode="bicubic", align_corners=False, antialias=True) of fp32 planes [3, h, w] to [3, oh, ow]: +/// the width pass first, then the height pass (aten's separable order); a side already the size is not resampled. +pub fn resizePlanes(gpa: std.mem.Allocator, src: []const f32, w: u32, h: u32, ow: u32, oh: u32) Error![]f32 { + var cur = try gpa.dupe(f32, src); + var cw: usize = w; + if (ow != w) { + var t = try taps(gpa, w, ow); + defer t.deinit(gpa); + const out = try gpa.alloc(f32, 3 * @as(usize, h) * ow); + for (0..3 * @as(usize, h)) |row| { + const in = cur[row * cw ..][0..cw]; + for (0..ow) |x| { + const k = t.w[x * t.ksize ..]; + const s = in[t.first[x]..]; + var acc: f32 = s[0] * k[0]; + for (1..t.count[x]) |j| acc += s[j] * k[j]; + out[row * ow + x] = acc; + } + } + gpa.free(cur); + cur = out; + cw = ow; + } + if (oh != h) { + var t = try taps(gpa, h, oh); + defer t.deinit(gpa); + const out = try gpa.alloc(f32, 3 * @as(usize, oh) * cw); + for (0..3) |c| { + const plane = cur[c * h * cw ..][0 .. h * cw]; + for (0..oh) |y| { + const k = t.w[y * t.ksize ..]; + const o = out[(c * oh + y) * cw ..][0..cw]; + for (o, 0..) |*v, x| { + var acc: f32 = plane[t.first[y] * cw + x] * k[0]; + for (1..t.count[y]) |j| acc += plane[(t.first[y] + j) * cw + x] * k[j]; + v.* = acc; + } + } + } + gpa.free(cur); + cur = out; + } + return cur; +} + +/// RGB bytes resized to [3, oh, ow] fp32 and normalized as the reference does: (x / 255 - mean) / std in fp32. +pub fn pixelValues(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error![]f32 { + const n = @as(usize, im.w) * im.h; + if (im.rgb.len != n * 3) return error.UnsupportedImage; + const planes = try gpa.alloc(f32, 3 * n); + defer gpa.free(planes); + for (0..n) |i| for (0..3) |c| { + planes[c * n + i] = @floatFromInt(im.rgb[i * 3 + c]); + }; + const out = try resizePlanes(gpa, planes, im.w, im.h, ow, oh); + for (0..3) |c| for (out[c * @as(usize, ow) * oh ..][0 .. @as(usize, ow) * oh]) |*v| { + v.* = (v.* / 255.0 - MEAN[c]) / STD[c]; + }; + return out; +} + +/// One image ready for the tower: its patches (fp32 [gh * gw, 768], row-major over the grid, each patch (c, dy, dx)). +pub const Prepared = struct { + patches: []f32, + gh: u32, + gw: u32, + + pub fn count(p: Prepared) usize { + return @as(usize, p.gh) * p.gw; + } + + /// The `` placeholders the prompt carries: one a 2x2 group of patches. + pub fn tokens(p: Prepared) usize { + return p.count() / (MERGE * MERGE); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +/// Planes [3, gh * 16, gw * 16] cut into patches, (c, dy, dx) within each (the reference's reshape and permute). +pub fn patchify(gpa: std.mem.Allocator, px: []const f32, gh: u32, gw: u32) Error![]f32 { + const H = @as(usize, gh) * PATCH; + const W = @as(usize, gw) * PATCH; + if (px.len != 3 * H * W) return error.UnsupportedImage; + const out = try gpa.alloc(f32, @as(usize, gh) * gw * PER); + for (0..gh) |py| for (0..gw) |pxi| { + const dst = out[(py * gw + pxi) * PER ..][0..PER]; + for (0..3) |c| for (0..PATCH) |dy| { + @memcpy(dst[(c * PATCH + dy) * PATCH ..][0..PATCH], px[(c * H + py * PATCH + dy) * W + pxi * PATCH ..][0..PATCH]); + }; + }; + return out; +} + +/// An RGB image to the tower's patches: grid, resize, normalize, patchify. +pub fn prepare(gpa: std.mem.Allocator, im: Image) Error!Prepared { + const g = try grid(im.w, im.h); + const px = try pixelValues(gpa, im, g[1] * PATCH, g[0] * PATCH); + defer gpa.free(px); + return .{ .patches = try patchify(gpa, px, g[0], g[1]), .gh = g[0], .gw = g[1] }; +} + +const fixture = @embedFile("fixtures/nemotron_resize.bin"); + +fn testImage(gpa: std.mem.Allocator, w: u32, h: u32) !Image { + const rgb = try gpa.alloc(u8, @as(usize, w) * h * 3); + for (0..h) |y| for (0..w) |x| for (0..3) |c| { + rgb[(y * w + x) * 3 + c] = @intCast((7 * x + 13 * y + 29 * c + (x * y) % 17) % 256); + }; + return .{ .w = w, .h = h, .rgb = rgb }; +} + +test "the patch grid matches target_patches (tools/zig/bolt_image_fixture.py's table)" { + const cases = [_][4]u32{ + .{ 2172, 724, 136, 46 }, .{ 640, 480, 40, 30 }, .{ 37, 23, 40, 28 }, .{ 4000, 3000, 132, 100 }, + .{ 16, 16, 32, 32 }, .{ 100, 5000, 8, 314 }, .{ 1024, 1024, 64, 64 }, .{ 512, 512, 32, 32 }, + .{ 3000, 100, 188, 8 }, .{ 1920, 1080, 120, 68 }, .{ 800, 600, 50, 38 }, + }; + for (cases) |c| { + const g = try grid(c[0], c[1]); + try std.testing.expectEqual(c[3], g[0]); + try std.testing.expectEqual(c[2], g[1]); + try std.testing.expect(@as(usize, g[0]) * g[1] <= MAX_PATCHES); + } +} + +test "resize and normalization match torch's bicubic antialias (up and down) to 5e-6" { + const gpa = std.testing.allocator; + var im = try testImage(gpa, 37, 23); + defer im.deinit(gpa); + const want: []align(1) const f32 = std.mem.bytesAsSlice(f32, fixture); + var at: usize = 0; + var worst: f32 = 0; + for ([_][2]u32{ .{ 48, 32 }, .{ 16, 16 }, .{ 24, 8 } }) |s| { + const got = try pixelValues(gpa, im, s[0], s[1]); + defer gpa.free(got); + for (got, want[at..][0..got.len]) |g, w| worst = @max(worst, @abs(g - w)); + at += got.len; + } + try std.testing.expectEqual(want.len, at); + try std.testing.expect(worst < 5e-6); +} + +test "patches are (c, dy, dx) over a row-major grid; prepare's count and tokens" { + const gpa = std.testing.allocator; + const px = try gpa.alloc(f32, 3 * 32 * 48); + defer gpa.free(px); + for (px, 0..) |*v, i| v.* = @floatFromInt(i); + const p = try patchify(gpa, px, 2, 3); + defer gpa.free(p); + // patch (1, 2), channel 2, dy 5, dx 7 is pixel (2, 16 + 5, 32 + 7) + try std.testing.expectEqual(@as(f32, @floatFromInt((2 * 32 + 21) * 48 + 39)), p[(1 * 3 + 2) * PER + (2 * PATCH + 5) * PATCH + 7]); + var im = try testImage(gpa, 37, 23); + defer im.deinit(gpa); + var prep = try prepare(gpa, im); + defer prep.deinit(gpa); + try std.testing.expectEqual(@as(usize, 28 * 40), prep.count()); + try std.testing.expectEqual(@as(usize, 280), prep.tokens()); + try std.testing.expectEqual(prep.count() * PER, prep.patches.len); +} diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index 79e7c4be3a..774a22ccac 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -3,6 +3,8 @@ 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"); +/// Nemotron 3.5 Bolt's preparation (one dynamic-resolution tile, torch's fp32 resize, C-RADIO's patches). +pub const nemotron = @import("nemotron_image.zig"); pub const Limits = struct { max_dimension: u32 = 8192, @@ -358,4 +360,5 @@ test "tenths rounds t * 10 exactly, ties to even, as the f128 product does" { test { _ = glm; + _ = nemotron; } diff --git a/zig/tests/bolt_vision_main.zig b/zig/tests/bolt_vision_main.zig new file mode 100644 index 0000000000..a9d8fb1a83 --- /dev/null +++ b/zig/tests/bolt_vision_main.zig @@ -0,0 +1,183 @@ +//! `tf-bolt-vision [reps]`: Nemotron 3.5 Bolt's image path end to end against the fp32 reference. +//! +//! The dump directory is tools/zig/bolt_vision_reference.py's (`--vision-only --dump DIR`): image_rgb.npy (the decoded +//! RGB the reference resized), meta.json, pixel_values.npy, tower_embed.npy, tower_out.npy, projector_out.npy. Host +//! preparation, the embedded rows, the last block's rows and the projector's rows are each compared (max |d|, the +//! relative Frobenius error, cosine over the whole tensor and the worst row's), then `reps` timed calls of `encode`. +//! Exits 1 when the projector rows' cosine is below 0.999 or the host pixels are off by more than 1e-4. +const std = @import("std"); +const cuda = @import("cuda"); +const nemotron = @import("nemotron"); +const qi = @import("qwen_image"); + +const V = nemotron.vision; +const NI = qi.nemotron; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + f32_arena = a; + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len < 3 or args.len > 4) { + std.debug.print("usage: tf-bolt-vision [reps]\n", .{}); + return 2; + } + const model = args[1]; + const dir = args[2]; + const reps: usize = if (args.len == 4) try std.fmt.parseInt(usize, args[3], 10) else 3; + const cwd = std.Io.Dir.cwd(); + + // host: decode is the shared decoder's; the reference's RGB is the input here + const rgb = try readNpy(a, io, dir, "image_rgb.npy"); + if (rgb.shape.len != 3 or rgb.shape[2] != 3 or !std.mem.eql(u8, rgb.descr, "|u1")) return error.BadFixture; + const im: NI.Image = .{ .w = @intCast(rgb.shape[1]), .h = @intCast(rgb.shape[0]), .rgb = @constCast(rgb.data) }; + const meta = try std.json.parseFromSliceLeaky(std.json.Value, a, try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "meta.json" }), a, .limited(1 << 20)), .{}); + const ghw = meta.object.get("grid_hw").?.array.items; + var t0 = now(io); + const g = try NI.grid(im.w, im.h); + const px = try NI.pixelValues(a, im, g[1] * NI.PATCH, g[0] * NI.PATCH); + const patches = try NI.patchify(a, px, g[0], g[1]); + const host_ms = ms(io, t0); + std.debug.print("image {d}x{d} -> grid {d}x{d} (reference {d}x{d}), {d} patches, {d} image tokens; host prep {d:.1} ms\n", .{ im.w, im.h, g[0], g[1], ghw[0].integer, ghw[1].integer, @as(usize, g[0]) * g[1], @as(usize, g[0]) * g[1] / 4, host_ms }); + if (g[0] != ghw[0].integer or g[1] != ghw[1].integer) return 1; + var fail = false; + { + const want = try readNpy(a, io, dir, "pixel_values.npy"); + const st = compare(f32Slice(want.data), px, 3 * @as(usize, g[1]) * NI.PATCH); + report("pixel_values", st); + if (st.max_abs > 1e-4) fail = true; + } + + 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 before = try ctx.memInfo(); + t0 = now(io); + const tower = try V.Tower.load(init.gpa, io, ctx.d, model); + defer tower.deinit(); + const after = try ctx.memInfo(); + std.debug.print("loaded in {d:.0} ms: arena {d:.3} GB (weights {d:.3} GB, scratch for {d} patches {d:.3} GB); device use +{d:.3} GB\n", .{ ms(io, t0), gb(tower.arena.len), gb(V.Tower.weightBytes()), V.max_patches, gb(V.Tower.scratchBytes(V.max_patches)), gb(before.free -| after.free) }); + + const n = @as(usize, g[0]) * g[1]; + const rows = V.prefix_tokens + n; + const tokens = n / 4; + var out = try cuda.DeviceBuffer.alloc(ctx.d, tokens * V.out_hidden * 2); + defer out.free(); + const xs = try a.alloc(f32, rows * V.hidden); + + // the embedded rows: no blocks run + tower.run_depth = 0; + var stages = try tower.encode(stream, patches, g[0], g[1], out.ptr); + try dl(ctx.d, stages.x, std.mem.sliceAsBytes(xs)); + report("tower_embed", compare(f32Slice((try readNpy(a, io, dir, "tower_embed.npy")).data), xs, V.hidden)); + tower.run_depth = V.depth; + + var times: std.ArrayList(f64) = .empty; + for (0..@max(reps, 1)) |_| { + t0 = now(io); + stages = try tower.encode(stream, patches, g[0], g[1], out.ptr); + try times.append(a, ms(io, t0)); + } + try dl(ctx.d, stages.x, std.mem.sliceAsBytes(xs)); + const tower_want = f32Slice((try readNpy(a, io, dir, "tower_out.npy")).data); + report("tower_out (all rows)", compare(tower_want, xs, V.hidden)); + report("tower_features (patch rows)", compare(tower_want[V.prefix_tokens * V.hidden ..], xs[V.prefix_tokens * V.hidden ..], V.hidden)); + const got16 = try a.alloc(u16, tokens * V.out_hidden); + try out.download(0, std.mem.sliceAsBytes(got16)); + const got = try a.alloc(f32, got16.len); + for (got, got16) |*o, h| o.* = @bitCast(@as(u32, h) << 16); + const proj = compare(f32Slice((try readNpy(a, io, dir, "projector_out.npy")).data), got, V.out_hidden); + report("projector_out", proj); + if (proj.cosine < 0.999) fail = true; + std.debug.print("encode ms ({d} patches, {d} tokens; first call makes the GEMM plans):", .{ n, tokens }); + for (times.items) |t| std.debug.print(" {d:.1}", .{t}); + std.debug.print("\n{s}\n", .{if (fail) "FAIL" else "PASS"}); + return if (fail) 1 else 0; +} + +const Stats = struct { max_abs: f64, max_ref: f64, rel: f64, cosine: f64, row_min: f64, rows: usize }; + +/// `got` against `want`, both row-major with rows of `width`. +fn compare(want: []const f32, got: []const f32, width: usize) Stats { + std.debug.assert(want.len == got.len); + var s: Stats = .{ .max_abs = 0, .max_ref = 0, .rel = 0, .cosine = 0, .row_min = 1, .rows = got.len / width }; + var dot: f64 = 0; + var nw: f64 = 0; + var ng: f64 = 0; + var dd: f64 = 0; + for (0..s.rows) |r| { + var rd: f64 = 0; + var rw: f64 = 0; + var rg: f64 = 0; + for (want[r * width ..][0..width], got[r * width ..][0..width]) |w0, g0| { + const w: f64 = w0; + const gg: f64 = g0; + s.max_abs = @max(s.max_abs, @abs(w - gg)); + s.max_ref = @max(s.max_ref, @abs(w)); + rd += w * gg; + rw += w * w; + rg += gg * gg; + dd += (w - gg) * (w - gg); + } + dot += rd; + nw += rw; + ng += rg; + if (rw > 0 and rg > 0) s.row_min = @min(s.row_min, rd / @sqrt(rw * rg)); + } + s.cosine = dot / @sqrt(nw * ng); + s.rel = @sqrt(dd / nw); + return s; +} + +fn report(name: []const u8, s: Stats) void { + std.debug.print("{s}: {d} rows, max|d| {e:.3} (max|ref| {d:.3}), rel {e:.3}, cosine {d:.6}, worst row {d:.6}\n", .{ name, s.rows, s.max_abs, s.max_ref, s.rel, s.cosine, s.row_min }); +} + +const Npy = struct { descr: []const u8, shape: []usize, data: []const u8 }; + +/// A little-endian C-order .npy (version 1 or 2 header). +fn readNpy(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) !Npy { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(a, &.{ dir, name }), a, .limited(1 << 31)); + if (bytes.len < 10 or !std.mem.eql(u8, bytes[0..6], "\x93NUMPY")) return error.BadNpy; + const v2 = bytes[6] >= 2; + const hl: usize = if (v2) std.mem.readInt(u32, bytes[8..12], .little) else std.mem.readInt(u16, bytes[8..10], .little); + const start: usize = if (v2) 12 else 10; + const header = bytes[start..][0..hl]; + if (std.mem.indexOf(u8, header, "'fortran_order': True") != null) return error.BadNpy; + const d0 = (std.mem.indexOf(u8, header, "'descr': '") orelse return error.BadNpy) + 10; + const d1 = std.mem.indexOfScalarPos(u8, header, d0, '\'') orelse return error.BadNpy; + const s0 = (std.mem.indexOf(u8, header, "'shape': (") orelse return error.BadNpy) + 10; + const s1 = std.mem.indexOfScalarPos(u8, header, s0, ')') orelse return error.BadNpy; + var shape: std.ArrayList(usize) = .empty; + var it = std.mem.tokenizeAny(u8, header[s0..s1], ", "); + while (it.next()) |tok| try shape.append(a, try std.fmt.parseInt(usize, tok, 10)); + return .{ .descr = header[d0..d1], .shape = shape.items, .data = bytes[start + hl ..] }; +} + +var f32_arena: std.mem.Allocator = undefined; + +/// The npy payload as f32 (copied: the file's bytes need not be 4-aligned in memory). +fn f32Slice(b: []const u8) []const f32 { + const out = f32_arena.alloc(f32, b.len / 4) catch @panic("out of memory"); + @memcpy(std.mem.sliceAsBytes(out), b[0 .. out.len * 4]); + return out; +} + +fn dl(d: *const cuda.Driver, src: u64, out: []u8) !void { + try d.check(d.api.cuMemcpyDtoH_v2(out.ptr, src, out.len), "cuMemcpyDtoH"); +} + +fn now(io: std.Io) i96 { + return std.Io.Clock.awake.now(io).toNanoseconds(); +} + +fn ms(io: std.Io, t0: i96) f64 { + return @as(f64, @floatFromInt(now(io) - t0)) / 1e6; +} + +fn gb(b: usize) f64 { + return @as(f64, @floatFromInt(b)) / 1e9; +} diff --git a/zig/tests/cuda/glue_tests.zig b/zig/tests/cuda/glue_tests.zig index 7af2fd42b9..331398e564 100644 --- a/zig/tests/cuda/glue_tests.zig +++ b/zig/tests/cuda/glue_tests.zig @@ -177,20 +177,25 @@ fn norms(r: *Rig) !void { } fn route(r: *Rig) !void { - const sk = glue.routerShape(D).sk; + try routeShape(r, D, E, top_k, NS); + try routeShape(r, 4096, 496, 9, 9); // Bolt: 4096-wide rows (eight K slices), 496 experts, top 9, no shared slots +} + +fn routeShape(r: *Rig, d: usize, e: usize, k: usize, ns: usize) !void { + const sk = glue.routerShape(d).sk; for ([_]usize{ 1, 16, 300 }) |rows| { - const x = try r.bfs(rows * D, 1.0); - const w = try r.bfs(E * D, 0.05); - const bias = try r.f32s(E, -0.05, 0.05); - const part = try r.a.alloc(f32, sk * rows * E); - const idx = try r.a.alloc(i32, rows * NS); - const wt = try r.a.alloc(f32, rows * NS); - ref.router(x, w, part, rows, D, E, sk); - ref.topk(part, bias, idx, wt, rows, 2.5, E, sk, top_k, NS, true); - const partd = try r.zeros(f32, sk * rows * E); - const idxd = try r.zeros(i32, rows * NS); - const wtd = try r.zeros(f32, rows * NS); - try r.g().route(try r.dev(u16, x), try r.dev(u16, w), try r.dev(f32, bias), partd, idxd, wtd, rows, D, E, top_k, 2.5, true); + const x = try r.bfs(rows * d, 1.0); + const w = try r.bfs(e * d, 0.05); + const bias = try r.f32s(e, -0.05, 0.05); + const part = try r.a.alloc(f32, sk * rows * e); + const idx = try r.a.alloc(i32, rows * ns); + const wt = try r.a.alloc(f32, rows * ns); + ref.router(x, w, part, rows, d, e, sk); + ref.topk(part, bias, idx, wt, rows, 2.5, e, sk, k, ns, true); + const partd = try r.zeros(f32, sk * rows * e); + const idxd = try r.zeros(i32, rows * ns); + const wtd = try r.zeros(f32, rows * ns); + try r.g().route(try r.dev(u16, x), try r.dev(u16, w), try r.dev(f32, bias), partd, idxd, wtd, rows, d, e, k, ns, 2.5, true); try r.same("router partials", f32, partd, part); try r.same("top-k ids", i32, idxd, idx); try r.same("top-k weights", f32, wtd, wt); From b7428a2320f7c62d0d7503d6164a6e6c80d6758f Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 21:56:27 +0000 Subject: [PATCH 10/11] tools: Bolt's fp32 PyTorch references (language model, image path) and the image fixture writer bolt_reference.py streams the language model a layer at a time on the CPU from the checkpoint's own bytes (ggml decode, ModelOpt NVFP4); it matches vLLM's top-1 on every prompt position once vLLM's bf16 ties are counted, and dumps each layer. bolt_vision_reference.py adds the processor, C-RADIO and the projector; both image fixtures' prompt counts and first tokens equal vLLM's. --- tools/zig/bolt_image_fixture.py | 49 ++++ tools/zig/bolt_reference.py | 415 +++++++++++++++++++++++++++++ tools/zig/bolt_vision_reference.py | 319 ++++++++++++++++++++++ 3 files changed, 783 insertions(+) create mode 100644 tools/zig/bolt_image_fixture.py create mode 100644 tools/zig/bolt_reference.py create mode 100644 tools/zig/bolt_vision_reference.py diff --git a/tools/zig/bolt_image_fixture.py b/tools/zig/bolt_image_fixture.py new file mode 100644 index 0000000000..93cefc13f8 --- /dev/null +++ b/tools/zig/bolt_image_fixture.py @@ -0,0 +1,49 @@ +"""Fixtures for zig/src/vision/nemotron_image.zig: torch's bicubic antialiased resize and CLIP normalization. + +A synthetic 37x23 RGB image (pixel = (7x + 13y + 29c + (x*y) % 17) % 256) resized to 48x32 (up) and to 16x16 and +24x8 (down, the antialias path), normalized as bolt_vision_reference.preprocess does; fp32 [3, H, W] each, little +endian, concatenated in that order. Also prints target_patches for a table of sizes (the Zig test's expectations). + + python3 tools/zig/bolt_image_fixture.py zig/src/vision/fixtures/nemotron_resize.bin +""" + +import sys +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from bolt_vision_reference import target_patches # noqa: E402 + +MEAN = [0.48145466, 0.4578275, 0.40821073] +STD = [0.26862954, 0.26130258, 0.27577711] +CFG = {"patch_size": 16, "min_num_patches": 1024, "max_num_patches": 13312, "downsample_ratio": 0.5, + "max_model_len": 16384} + + +def image(w, h): + y, x, c = np.meshgrid(np.arange(h), np.arange(w), np.arange(3), indexing="ij") + return ((7 * x + 13 * y + 29 * c + (x * y) % 17) % 256).astype(np.uint8) + + +def prep(rgb, ow, oh): + t = torch.from_numpy(rgb).permute(2, 0, 1)[None].float() + t = F.interpolate(t, size=(oh, ow), mode="bicubic", align_corners=False, antialias=True) + mean = torch.tensor(MEAN).view(1, 3, 1, 1) + std = torch.tensor(STD).view(1, 3, 1, 1) + return ((t / 255.0 - mean) / std)[0].numpy().astype(" f32 [rows, blocks * 256]; bit-identical to bolt_ggml.iq2_xxs.""" + grid, sign = _iq2_tables() + rows = raw.shape[0] + b = raw.reshape(-1, 66) + d = b[:, :2].contiguous().view(torch.float16).float() # [n, 1] + qs = b[:, 2:].contiguous() # 8 sub-blocks of (4 grid indices, u32 aux) + pair = qs.view(-1, 8, 8) + idx = pair[:, :, :4].long() # [n, 8, 4] + aux = pair[:, :, 4:].contiguous().view(torch.int32).long().squeeze(-1) & 0xFFFFFFFF # [n, 8] + db = d * (0.5 + (aux >> 28).float()) * 0.25 # [n, 8] + shifts = torch.arange(4) * 7 + sidx = (aux[..., None] >> shifts) & 127 # [n, 8, 4] + out = grid[idx] * sign[sidx] * db[..., None, None] # [n, 8, 4, 8] + return out.reshape(rows, -1) + + +def q8_0(raw): + """U8 [rows, blocks, 34] -> f32 [rows, blocks * 32].""" + rows = raw.shape[0] + b = raw.reshape(-1, 34) + d = b[:, :2].contiguous().view(torch.float16).float() + q = b[:, 2:].contiguous().view(torch.int8).float() + return (q * d).reshape(rows, -1) + + +_E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], dtype=torch.float32) + + +def nvfp4(packed, scale, global_scale): + """ModelOpt NVFP4: U8 [N, K/2] (column 2j low nibble, 2j+1 high), e4m3 [N, K/16], f32 [] -> f32 [N, K].""" + n = packed.shape[0] + codes = torch.stack([packed & 0xF, packed >> 4], dim=-1).reshape(n, -1).long() + w = _E2M1[codes].view(n, -1, 16) * scale.float()[..., None] + return (w * global_scale.float()).reshape(n, -1) + + +class Checkpoint: + def __init__(self, model): + self.dir = Path(model) + self.map = json.loads((self.dir / "model.safetensors.index.json").read_text())["weight_map"] + cfg = json.loads((self.dir / "config.json").read_text()) + self.algo = {k: v["quant_algo"] for k, v in cfg["quantization_config"]["quantized_layers"].items()} + self.cfg = cfg["llm_config"] + self.files = {} + + def raw(self, name): + f = self.map[name] + if f not in self.files: + self.files[f] = safe_open(str(self.dir / f), framework="pt") + return self.files[f].get_tensor(name) + + def has(self, name): + return name in self.map + + def dense(self, name): + """A Linear weight [out, in] in f32, by the format the export records for it.""" + algo = self.algo.get(name) + if algo == "IQ2_XXS": + return iq2_xxs(self.raw(name + ".weight")) + if algo == "Q8_0": + return q8_0(self.raw(name + ".weight")) + if algo == "W4A16_NVFP4": + return nvfp4(self.raw(name + ".weight"), self.raw(name + ".weight_scale"), self.raw(name + ".weight_scale_2")) + if algo is None: + return self.raw(name + ".weight").float() + raise SystemExit(f"{name}: unsupported {algo}") + + def vec(self, name): + return self.raw(name).float() + + +# ---------------------------------------------------------------- blocks + + +def rmsnorm(x, w, eps): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * w + + +class State: + """Per-layer recurrent state: Mamba conv window + ssm state, attention K/V.""" + + def __init__(self): + self.conv, self.ssm, self.k, self.v = {}, {}, {}, {} + self.pos = 0 + + +def mamba(ck, cfg, li, x, st): + p = f"{PRE}layers.{li}.mixer." + H, P, N, G = cfg["mamba_num_heads"], cfg["mamba_head_dim"], cfg["ssm_state_size"], cfg["n_groups"] + inner = H * P + T = x.shape[0] + zxbcdt = x @ ck.dense(p + "in_proj").T # [T, inner + conv_dim + H] + conv_dim = inner + 2 * G * N + z, xbc, dt = zxbcdt.split([inner, conv_dim, H], dim=-1) + + # causal depthwise conv over [previous K-1 inputs, these T], + bias, SiLU + cw = ck.vec(p + "conv1d.weight")[:, 0, :] # [conv_dim, K] + K = cw.shape[1] + prev = st.conv.get(li, torch.zeros(K - 1, conv_dim)) + seq = torch.cat([prev, xbc], 0) # [K-1+T, conv_dim] + st.conv[li] = seq[-(K - 1):].clone() + win = seq.unfold(0, K, 1) # [T, conv_dim, K] + xbc = F.silu((win * cw).sum(-1) + ck.vec(p + "conv1d.bias")) + xs, B, C = xbc.split([inner, G * N, G * N], dim=-1) + + # dt = softplus(dt + dt_bias), limited to time_step_limit = (0, inf): no clamp (the kernel path) + dt = F.softplus(dt + ck.vec(p + "dt_bias")) # [T, H] + lo, hi = cfg["time_step_limit"][0], cfg["time_step_limit"][1] + hi = float(hi["__float__"]) if isinstance(hi, dict) else float(hi) + if lo > 0 or hi != math.inf: + dt = dt.clamp(lo, hi) + A = -torch.exp(ck.vec(p + "A_log")) # [H] + D = ck.vec(p + "D") + xs = xs.view(T, H, P) + rep = H // G + B = B.view(T, G, N).repeat_interleave(rep, dim=1) # [T, H, N] + C = C.view(T, G, N).repeat_interleave(rep, dim=1) + dA = torch.exp(dt * A) # [T, H] + h = st.ssm.get(li, torch.zeros(H, P, N)) + ys = torch.empty(T, H, P) + for t in range(T): + h = h * dA[t][:, None, None] + (dt[t][:, None] * xs[t])[:, :, None] * B[t][:, None, :] + ys[t] = torch.einsum("hpn,hn->hp", h, C[t]) + st.ssm[li] = h + y = (ys + xs * D[:, None]).reshape(T, inner) + + # gated RMSNorm, norm_before_gate=False: normalize y * silu(z) per group of inner / n_groups + g = (y * F.silu(z)).view(T, G, inner // G) + g = g * torch.rsqrt(g.pow(2).mean(-1, keepdim=True) + cfg["layer_norm_epsilon"]) + y = g.reshape(T, inner) * ck.vec(p + "norm.weight") + return y @ ck.dense(p + "out_proj").T + + +def attention(ck, cfg, li, x, st): + p = f"{PRE}layers.{li}.mixer." + nh, nkv, hd = cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + T = x.shape[0] + q = (x @ ck.dense(p + "q_proj").T).view(T, nh, hd) + k = (x @ ck.dense(p + "k_proj").T).view(T, nkv, hd) + v = (x @ ck.dense(p + "v_proj").T).view(T, nkv, hd) + if li in st.k: + k = torch.cat([st.k[li], k], 0) + v = torch.cat([st.v[li], v], 0) + st.k[li], st.v[li] = k, v + S = k.shape[0] + rep = nh // nkv + kk = k.repeat_interleave(rep, dim=1).transpose(0, 1) # [nh, S, hd] + vv = v.repeat_interleave(rep, dim=1).transpose(0, 1) + s = torch.einsum("thd,hsd->hts", q, kk) * hd**-0.5 + qpos = torch.arange(S - T, S)[:, None] + s = s.masked_fill(torch.arange(S)[None, :] > qpos, float("-inf")) + o = torch.einsum("hts,hsd->thd", torch.softmax(s, -1), vv).reshape(T, nh * hd) + return o @ ck.dense(p + "o_proj").T + + +def moe(ck, cfg, bc, li, x, routes): + p = f"{PRE}layers.{li}.mixer." + E, k = bc["n_routed_experts"], bc["num_experts_per_tok"] + T = x.shape[0] + # router: sigmoid(x W^T) in fp32; choice by score + e_score_correction_bias (n_group = 1, so no + # group masking); weights from the un-biased scores, normalized, times routed_scaling_factor + assert cfg["n_group"] == 1 and cfg["topk_group"] == 1 + scores = torch.sigmoid(x @ ck.vec(p + "gate.weight").T) # [T, E] + choice = scores + ck.vec(p + "gate.e_score_correction_bias") + idx = torch.topk(choice, k, dim=-1).indices # [T, k], by descending choice + w = scores.gather(1, idx) + if cfg["norm_topk_prob"]: + w = w / (w.sum(-1, keepdim=True) + 1e-20) + w = w * cfg["routed_scaling_factor"] + routes[li] = (idx.clone(), w.clone()) + + lat = x @ ck.dense(p + "fc1_latent_proj").T # [T, latent] + acc = torch.zeros_like(lat) + for e in torch.unique(idx).tolist(): + tok, slot = torch.where(idx == e) + up = ck.dense(f"{p}experts.{e}.up_proj") # [inter, latent] + down = ck.dense(f"{p}experts.{e}.down_proj") # [latent, inter] + hdn = F.relu(lat[tok] @ up.T).square() + acc.index_add_(0, tok, (hdn @ down.T) * w[tok, slot][:, None]) + routed = acc @ ck.dense(p + "fc2_latent_proj").T + sh = F.relu(x @ ck.dense(p + "shared_experts.up_proj").T).square() + return routed + sh @ ck.dense(p + "shared_experts.down_proj").T + + +def lm_head(ck, x, rows=16384): + raw = ck.raw("language_model.lm_head.weight") + return torch.cat([x @ q8_0(raw[i:i + rows]).T for i in range(0, raw.shape[0], rows)], -1) + + +def forward(ck, ids, st=None, dump=None, dump_last=8, logits_last=None, verbose=True, splice=None): + """Run `ids` (continuing `st`); returns (logits [T or logits_last, vocab] f32, state, routes). + + `splice` = (positions LongTensor [n], rows f32 [n, hidden]) overwrites those positions' token + embeddings (the multimodal path: image-context positions take projected vision embeddings).""" + cfg = ck.cfg + st = st or State() + eps = cfg["layer_norm_epsilon"] + ids_t = torch.tensor(ids, dtype=torch.long) + emb = ck.raw(PRE + "embeddings.weight") + h = emb[ids_t].float() + if splice is not None: + pos, rows = splice + h[pos] = rows.float() + routes = {} + keep = slice(-dump_last, None) if dump_last else slice(None) + if dump: + np.save(dump / "h_embed.npy", h[keep].numpy()) + t0 = time.time() + for li, bc in enumerate(cfg["block_configs"]): + t1 = time.time() + x = rmsnorm(h, ck.vec(f"{PRE}layers.{li}.norm.weight"), eps) + kind = bc["block_type"] + if kind == "mamba": + out = mamba(ck, cfg, li, x, st) + elif kind == "attention": + out = attention(ck, cfg, li, x, st) + elif kind == "moe": + out = moe(ck, cfg, bc, li, x, routes) + else: + raise SystemExit(f"layer {li}: {kind}") + h = h + out + if dump: + np.save(dump / f"h_layer{li:02d}.npy", h[keep].numpy()) + np.save(dump / f"mixer_layer{li:02d}.npy", out[keep].numpy()) + if li in routes: + np.save(dump / f"route_idx_layer{li:02d}.npy", routes[li][0][keep].numpy()) + np.save(dump / f"route_w_layer{li:02d}.npy", routes[li][1][keep].numpy()) + if verbose: + print(f" layer {li:2d} {kind:9s} |h|={h[-1].norm():9.3f} {time.time() - t1:5.1f}s", file=sys.stderr, flush=True) + st.pos += len(ids) + h = rmsnorm(h, ck.vec(PRE + "norm_f.weight"), eps) + if logits_last: + h = h[-logits_last:] + logits = lm_head(ck, h) + if dump: + np.save(dump / "h_final_norm.npy", h[keep].numpy()) + np.save(dump / "logits.npy", logits[keep].numpy()) + if verbose: + print(f" forward {len(ids)} tokens: {time.time() - t0:.1f}s", file=sys.stderr) + return logits, st, routes + + +# ---------------------------------------------------------------- fixtures / scoring + + +def load_fixture(name, path=FIXTURES): + for line in open(path): + d = json.loads(line) + if d["name"] == name: + return d + raise SystemExit(f"no fixture {name!r} in {path}") + + +def score(lp, tops, label): + """lp [n, vocab] log-probs vs vLLM dicts (token id -> logprob) for the same n positions.""" + strict = tie = 0 + diffs = [] + rows = [] + for i, top in enumerate(tops): + ids = torch.tensor([int(t) for t in top]) + ref = torch.tensor(list(top.values())) + mine = int(lp[i].argmax()) + best = ref.max() + strict += mine == int(ids[ref.argmax()]) + tie += bool((ids[ref >= best - 1e-6] == mine).any()) + err = (lp[i][ids] - ref).abs() + d = err.max().item() + diffs.append(d) + j = int(err.argmax()) + rows.append((d, i, int(ids[j]), ref[j].item(), lp[i][ids[j]].item())) + n = len(tops) + diffs = np.array(diffs) + print(f"{label}: top-1 {strict}/{n} = {100 * strict / n:.1f}% (counting vLLM ties: {tie}/{n} = {100 * tie / n:.1f}%) | " + f"max|dlogprob| over vLLM top-20: mean {diffs.mean():.3f} median {np.median(diffs):.3f} max {diffs.max():.3f}") + for d, i, t, r, m in sorted(rows, reverse=True)[:3]: + print(f" worst: position {i} id {t}: vLLM {r:.3f} ours {m:.3f} (|d| {d:.3f})") + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + src = ap.add_mutually_exclusive_group(required=True) + src.add_argument("--ids", help="file of whitespace/comma separated token ids (or a JSON list)") + src.add_argument("--fixture", help="fixture name in text.jsonl (short, code, think, multiturn, long, german)") + ap.add_argument("--model", default=MODEL) + ap.add_argument("--fixtures", default=FIXTURES) + ap.add_argument("--dump", help="directory for per-layer npy dumps") + ap.add_argument("--dump-last", type=int, default=8, help="positions kept in dumps (0 = all)") + ap.add_argument("--compare", action="store_true", help="score prompt logprobs against the fixture") + ap.add_argument("--generate", type=int, default=0, help="greedy tokens to generate after the prompt") + ap.add_argument("--threads", type=int, default=16) + ap.add_argument("--quiet", action="store_true") + a = ap.parse_args() + torch.set_num_threads(a.threads) + torch.set_grad_enabled(False) + ck = Checkpoint(a.model) + + fx = None + if a.fixture: + fx = load_fixture(a.fixture, a.fixtures) + ids = fx["prompt_ids"] + else: + text = Path(a.ids).read_text().strip() + ids = json.loads(text) if text.startswith("[") else [int(t) for t in text.replace(",", " ").split()] + dump = None + if a.dump: + dump = Path(a.dump) + dump.mkdir(parents=True, exist_ok=True) + (dump / "ids.json").write_text(json.dumps(ids)) + + t0 = time.time() + logits, st, _ = forward(ck, ids, dump=dump, dump_last=a.dump_last, verbose=not a.quiet) + print(f"prompt {len(ids)} tokens in {time.time() - t0:.1f}s", flush=True) + lp = torch.log_softmax(logits, -1) + if a.compare: + if fx is None: + raise SystemExit("--compare needs --fixture") + # prompt_top[i] is vLLM's distribution for position i given ids[:i], i.e. our logits at i-1 + score(lp[:-1], fx["prompt_top"][1:], f"{a.fixture} prompt ({len(ids) - 1} positions)") + + if a.generate: + from tokenizers import Tokenizer + tok = Tokenizer.from_file(str(Path(a.model) / "tokenizer.json")) + out = [int(logits[-1].argmax())] + steps = [lp[-1]] + for _ in range(a.generate - 1): + lg, st, _ = forward(ck, [out[-1]], st, verbose=False) + steps.append(torch.log_softmax(lg[-1], -1)) + out.append(int(lg[-1].argmax())) + print(f"generated {len(out)} tokens in {time.time() - t0:.1f}s total") + print("ids:", out) + print("text:", repr(tok.decode(out))) + if fx is not None: + ref = fx["gen_tokens"][: len(out)] + mine = [tok.decode([i]) for i in out] + match = 0 + for m, r in zip(mine, ref): + if m != r: + break + match += 1 + print(f"greedy prefix match vs vLLM: {match}/{len(ref)}") + print("vLLM:", repr("".join(ref))) + # logprob diffs at the generated positions (vLLM keys are token strings) + diffs = [] + for s_lp, top in zip(steps[:match + 1], fx["gen_top"]): + d = 0.0 + for i in torch.topk(s_lp, 40).indices.tolist(): + t = tok.decode([i]) + if t in top: + d = max(d, abs(s_lp[i].item() - top[t])) + diffs.append(d) + print(f"gen max|dlogprob| (shared top ids, first {len(diffs)} steps): mean {np.mean(diffs):.3f} max {np.max(diffs):.3f}") + + +if __name__ == "__main__": + main() diff --git a/tools/zig/bolt_vision_reference.py b/tools/zig/bolt_vision_reference.py new file mode 100644 index 0000000000..bb8c70f59f --- /dev/null +++ b/tools/zig/bolt_vision_reference.py @@ -0,0 +1,319 @@ +"""fp32 PyTorch reference of Nemotron 3.5 Bolt's IMAGE path, spliced into the LM reference. + +A debugging oracle for the native engine, written from the checkpoint's +image_processing_nemotron_h_omni.py / modeling_radio.py / modeling_nemotron_h_omni.py and checked +against vLLM's nano_nemotron_vl.py + radio.py (the code that produced the fixtures): + + preprocess one dynamic-resolution tile per image (no InternVL tiling, no thumbnail): the patch + grid is ceil-ish(H/16) x ceil-ish(W/16), scaled down to <= max_num_patches (13312), + up to >= min_num_patches (1024), each side rounded to even; the RGB uint8 image is + resized with torch bicubic (align_corners=False, antialias=True) in fp32, then + (x/255 - mean)/std with the CLIP mean/std + prompt + * (hp*wp/4) + replaces the chat template's "" + tower C-RADIO ViT-H/16: patch linear (768->1280, no bias) + CPE position embedding + (128x128 table bilinearly resized to max(hp,wp)^2, cropped to hp x wp) + 10 prefix + tokens (3 cls + 7 registers); 32 pre-LN blocks (LayerNorm eps 1e-6, 16 heads x 80, + fused qkv with bias, exact-erf GELU MLP 1280->5120->1280, no layer scale); no final + norm; the 10 prefix tokens are dropped, only patch tokens are used + projector pixel shuffle v2 (2x2 -> 5120 channels ordered [dy, dx, c]), RMSNorm(5120, eps 1e-5), + Linear 5120->20480, ReLU^2, Linear 20480->4096 (no biases) + splice projector rows overwrite the token embeddings at the (id 18) positions + + python3 tools/zig/bolt_vision_reference.py --case 0 --generate 8 + python3 tools/zig/bolt_vision_reference.py --case 1 --vision-only --dump out/ +""" + +import argparse +import copy +import json +import math +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import bolt_reference as lm # noqa: E402 + +FIXTURES = "/storage/models/bolt-fixtures/vision.jsonl" +VPRE = "vision_model.radio_model." + + +# ---------------------------------------------------------------- preprocessing + + +def target_patches(w, h, cfg): + """(w_patches, h_patches) of the single dynamic-resolution tile (image path, one image).""" + ps, mn, mx = cfg["patch_size"], cfg["min_num_patches"], cfg["max_num_patches"] + r = int(round(1 / cfg["downsample_ratio"])) + budget = (cfg["max_model_len"] - 4) * r * r # vLLM subtracts the text length too; capped anyway + avail = max(min(budget, mx), mn) + ph = round(h / ps + 0.5) + pw = round(w / ps + 0.5) + f = min(math.sqrt(avail / (ph * pw)), 1.0) + th, tw = math.floor(f * ph), math.floor(f * pw) + if avail > mn and th * tw < mn: + up = math.sqrt(mn / (th * tw)) + th, tw = math.ceil(up * th), math.ceil(up * tw) + if th % r: + th = th + (r - th % r) if (th + r - th % r) * tw <= avail else max(r, th - th % r) + if tw % r: + tw = tw + (r - tw % r) if th * (tw + r - tw % r) <= avail else max(r, tw - tw % r) + return tw, th + + +def preprocess(path, cfg): + img = Image.open(path) + if img.mode != "RGB": + img = img.convert("RGB") + wp, hp = target_patches(img.width, img.height, cfg) + ps = cfg["patch_size"] + t = torch.from_numpy(np.array(img, dtype=np.uint8)).permute(2, 0, 1)[None].float() + if t.shape[-2:] != (hp * ps, wp * ps): + t = F.interpolate(t, size=(hp * ps, wp * ps), mode="bicubic", align_corners=False, antialias=True) + mean = torch.tensor(cfg["norm_mean"]).view(1, 3, 1, 1) + std = torch.tensor(cfg["norm_std"]).view(1, 3, 1, 1) + return ((t / 255.0 - mean) / std)[0], (img.width, img.height), (hp, wp) + + +# ---------------------------------------------------------------- C-RADIO tower + + +def pos_embed(table, hp, wp): + """CPE: [1, R*R, C] table -> [hp*wp, C]; bilinear (align_corners=False) to max side, crop.""" + n = int(math.isqrt(table.shape[1])) + pos = table.reshape(1, n, n, -1).permute(0, 3, 1, 2).float() + m = max(hp, wp) + pos = F.interpolate(pos, size=(m, m), mode="bilinear", align_corners=False)[..., :hp, :wp] + return pos.flatten(2)[0].T + + +def radio(ck, vcfg, px, dump=None, dump_layers=False, verbose=True): + """px [3, H, W] normalized -> (all tokens [10 + hp*wp, 1280], patch features [hp*wp, 1280]).""" + ps, C, nh = vcfg["patch_size"], vcfg["hidden_size"], vcfg["num_attention_heads"] + eps = vcfg["layer_norm_eps"] + _, H, W = px.shape + hp, wp = H // ps, W // ps + m = VPRE + "model." + g = lambda n: ck.raw(n).float() # noqa: E731 + patches = px.reshape(3, hp, ps, wp, ps).permute(1, 3, 0, 2, 4).reshape(hp * wp, 3 * ps * ps) + x = patches @ g(m + "patch_generator.embedder.weight").T + x = x + pos_embed(g(m + "patch_generator.pos_embed"), hp, wp) + prefix = g(m + "patch_generator.cls_token.token") # [3 cls + 7 registers, C] + x = torch.cat([prefix, x], 0) + if dump: + np.save(dump / "tower_embed.npy", x.numpy()) + hd = C // nh + T = x.shape[0] + for i in range(vcfg["num_hidden_layers"]): + t0 = time.time() + b = f"{m}blocks.{i}." + y = F.layer_norm(x, (C,), g(b + "norm1.weight"), g(b + "norm1.bias"), eps) + qkv = (y @ g(b + "attn.qkv.weight").T + g(b + "attn.qkv.bias")).view(T, 3, nh, hd).permute(1, 2, 0, 3) + o = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2]) # [nh, T, hd], scale hd^-0.5 + o = o.transpose(0, 1).reshape(T, C) + x = x + o @ g(b + "attn.proj.weight").T + g(b + "attn.proj.bias") + y = F.layer_norm(x, (C,), g(b + "norm2.weight"), g(b + "norm2.bias"), eps) + y = F.gelu(y @ g(b + "mlp.fc1.weight").T + g(b + "mlp.fc1.bias")) + x = x + y @ g(b + "mlp.fc2.weight").T + g(b + "mlp.fc2.bias") + if dump and dump_layers: + np.save(dump / f"tower_layer{i:02d}.npy", x.numpy()) + if verbose: + print(f" radio block {i:2d} |x|={x[10:].norm(dim=-1).mean():8.3f} {time.time() - t0:5.1f}s", + file=sys.stderr, flush=True) + skip = vcfg["num_cls_tokens"] + vcfg["num_registers"] + return x, x[skip:] + + +def pixel_shuffle(feat, hp, wp, r=2): + """ps_version v2: [hp*wp, C] -> [(hp/r)*(wp/r), r*r*C], channel index (dy*r + dx)*C + c.""" + C = feat.shape[-1] + x = feat.view(hp // r, r, wp // r, r, C).permute(0, 2, 1, 3, 4) + return x.reshape((hp // r) * (wp // r), r * r * C) + + +def projector(ck, x): + w0 = ck.raw("mlp1.0.weight").float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5) * w0 + x = F.relu(x @ ck.raw("mlp1.1.weight").float().T).square() + return x @ ck.raw("mlp1.3.weight").float().T + + +# ---------------------------------------------------------------- prompt + + +def build_prompt(question, n_img, image_first=True): + """The chat template (no system content, thinking off) with the processor's image expansion.""" + img = "" + "" * n_img + "" + content = img + question if image_first else question + img + return ("<|im_start|>system\n<|im_end|>\n<|im_start|>user\n" + content + + "<|im_end|>\n<|im_start|>assistant\n") + + +# ---------------------------------------------------------------- scoring + + +def compare_step(tok, lp, top, cache): + """lp [vocab] vs vLLM {token string: logprob}; returns (n mapped, max |d|, worst).""" + if "rev" not in cache or cache["rev_lp"] is not lp: + rev = {} + for i in torch.topk(lp, 4000).indices.tolist(): + rev.setdefault(tok.decode([i]), i) + cache["rev"], cache["rev_lp"] = rev, lp + rev = cache["rev"] + d, worst, n = 0.0, None, 0 + for s, ref in top.items(): + i = rev.get(s) + if i is None: + ids = tok.encode(s, add_special_tokens=False).ids + if len(ids) != 1: + continue + i = ids[0] + n += 1 + e = abs(lp[i].item() - ref) + if e > d: + d, worst = e, (s, ref, lp[i].item()) + return n, d, worst + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--case", type=int, default=0, help="line index in vision.jsonl") + ap.add_argument("--image", help="image file (overrides the fixture's; needs --question)") + ap.add_argument("--question") + ap.add_argument("--model", default=lm.MODEL) + ap.add_argument("--fixtures", default=FIXTURES) + ap.add_argument("--text-first", action="store_true", help="question before the image in the user turn") + ap.add_argument("--vision-only", action="store_true", help="stop after the projector") + ap.add_argument("--generate", type=int, default=8, help="free-running greedy tokens") + ap.add_argument("--forced", type=int, default=16, help="teacher-forced steps scored against the fixture") + ap.add_argument("--dump", help="directory for npy dumps (pixels, tower, shuffle, projector)") + ap.add_argument("--dump-tower-layers", action="store_true", help="also every RADIO block output") + ap.add_argument("--dump-lm", action="store_true", help="also the LM per-layer dumps (last 8 positions)") + ap.add_argument("--threads", type=int, default=16) + a = ap.parse_args() + torch.set_num_threads(a.threads) + torch.set_grad_enabled(False) + + root = json.loads((Path(a.model) / "config.json").read_text()) + pcfg = json.loads((Path(a.model) / "preprocessor_config.json").read_text()) + vcfg = root["vision_config"] + assert root["ps_version"] == "v2" and not root["use_thumbnail"] and not root.get("vision_final_layernorm") + assert not vcfg["use_swiglu_ffn"] and vcfg["hidden_act"] == "gelu" + + fx = None + if a.image: + img_path, question = a.image, a.question + else: + fx = json.loads(open(a.fixtures).readlines()[a.case]) + img_path, question = str(Path(a.fixtures).parent / fx["image"]), fx["question"] + + dump = None + if a.dump: + dump = Path(a.dump) + dump.mkdir(parents=True, exist_ok=True) + + t0 = time.time() + px, (w, h), (hp, wp) = preprocess(img_path, pcfg) + n_img = hp * wp // 4 + print(f"image {Path(img_path).name} {w}x{h} -> {wp * 16}x{hp * 16} px, patch grid {hp}x{wp} = {hp * wp}, " + f"{n_img} image tokens", flush=True) + + from tokenizers import Tokenizer + tok = Tokenizer.from_file(str(Path(a.model) / "tokenizer.json")) + ids = tok.encode(build_prompt(question, n_img, not a.text_first), add_special_tokens=False).ids + img_pos = torch.tensor([i for i, t in enumerate(ids) if t == root["img_context_token_id"]]) + assert len(img_pos) == n_img + line = f"prompt tokens {len(ids)}" + if fx: + line += f" (vLLM usage.prompt_tokens {fx['usage']['prompt_tokens']}: " \ + f"{'MATCH' if len(ids) == fx['usage']['prompt_tokens'] else 'MISMATCH'})" + print(line, flush=True) + + ck = lm.Checkpoint(a.model) + allx, feat = radio(ck, vcfg, px, dump, a.dump_tower_layers) + shuf = pixel_shuffle(feat, hp, wp, int(round(1 / root["downsample_ratio"]))) + proj = projector(ck, shuf) + print(f"vision: tower {tuple(allx.shape)} -> features {tuple(feat.shape)} -> shuffle {tuple(shuf.shape)} " + f"-> projector {tuple(proj.shape)} |row| mean {proj.norm(dim=-1).mean():.3f} {time.time() - t0:.1f}s", + flush=True) + if dump: + np.save(dump / "pixel_values.npy", px.numpy()) + np.save(dump / "image_rgb.npy", np.array(Image.open(img_path).convert("RGB"), dtype=np.uint8)) + np.save(dump / "tower_out.npy", allx.numpy()) + np.save(dump / "tower_features.npy", feat.numpy()) + np.save(dump / "pixel_shuffle.npy", shuf.numpy()) + np.save(dump / "projector_out.npy", proj.numpy()) + (dump / "ids.json").write_text(json.dumps(ids)) + (dump / "meta.json").write_text(json.dumps({"image": str(img_path), "orig_wh": [w, h], + "grid_hw": [hp, wp], "n_image_tokens": n_img, + "image_positions": [int(img_pos[0]), int(img_pos[-1])]})) + if a.vision_only: + return + + t1 = time.time() + logits, st, _ = lm.forward(ck, ids, splice=(img_pos, proj), dump=dump if a.dump_lm else None, logits_last=1) + print(f"LM prompt {len(ids)} tokens in {time.time() - t1:.1f}s", flush=True) + lp = torch.log_softmax(logits[-1], -1) + st0 = copy.deepcopy(st) + out = [int(lp.argmax())] + steps = [lp] + for _ in range(a.generate - 1): + lg, st, _ = lm.forward(ck, [out[-1]], st, verbose=False) + steps.append(torch.log_softmax(lg[-1], -1)) + out.append(int(lg[-1].argmax())) + mine = [tok.decode([i]) for i in out] + print("ours:", repr("".join(mine)), out) + if fx is None: + print(f"total {time.time() - t0:.1f}s") + return + ref = [g["t"] for g in fx["gen"]] + print("vLLM:", repr("".join(ref[:len(out)]))) + match = 0 + for m_, r_ in zip(mine, ref): + if m_ != r_: + break + match += 1 + print(f"free-running greedy prefix match vs vLLM: {match}/{len(out)}") + top = torch.topk(lp, 5) + print(" step 0 ours top-5:", [(tok.decode([i]), round(v, 3)) for v, i in zip(top.values.tolist(), top.indices.tolist())]) + rank = sorted(fx["gen"][0]["top"], key=fx["gen"][0]["top"].get, reverse=True) + print(" step 0 vLLM top-5:", [(r, round(fx["gen"][0]["top"][r], 3)) for r in rank[:5]]) + + # teacher-forced: feed vLLM's tokens, score every step (a vLLM near-tie can't derail the rest) + n = min(max(a.forced, 1), len(fx["gen"])) + st, s_lp, cache = st0, lp, {} + strict = tie = 0 + diffs = [] + for k in range(n): + g = fx["gen"][k] + cnt, d, worst = compare_step(tok, s_lp, g["top"], cache) + o = int(s_lp.argmax()) + os_ = tok.decode([o]) + best = max(g["top"].values()) + is_tie = os_ in g["top"] and g["top"][os_] >= best - 1e-4 + strict += os_ == g["t"] + tie += is_tie + diffs.append(d) + rank = sorted(g["top"], key=g["top"].get, reverse=True) + print(f" forced step {k:2d}: ours {os_!r} {s_lp[o].item():7.3f} | vLLM {g['t']!r} {g['lp']:7.3f} " + f"(2nd {rank[1]!r} {g['top'][rank[1]]:7.3f}) | mapped {cnt}/20 max|d| {d:.3f}" + + ("" if os_ == g["t"] else f" {'TIE' if is_tie else 'MISS'} worst {worst}")) + if k + 1 < n: + nxt = cache["rev"].get(g["t"]) + if nxt is None: + nxt = tok.encode(g["t"], add_special_tokens=False).ids[0] + lg, st, _ = lm.forward(ck, [nxt], st, verbose=False) + s_lp = torch.log_softmax(lg[-1], -1) + diffs = np.array(diffs) + print(f"teacher-forced {n} steps: top-1 {strict}/{n} (counting vLLM ties {tie}/{n}) | max|dlogprob| over " + f"vLLM top-20: mean {diffs.mean():.3f} median {np.median(diffs):.3f} max {diffs.max():.3f}") + print(f"total {time.time() - t0:.1f}s") + + +if __name__ == "__main__": + main() From 42b0da79fcd57f45cbd84f2487626e208b15a9c8 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 22:17:04 +0000 Subject: [PATCH 11/11] CI on #897: zig fmt; lean (one-line comments, native/cuda.zig's family aliases to cuda_alias.zig) --- tools/zig/bolt_ggml.py | 10 +---- tools/zig/bolt_ggml_fixture.py | 8 +--- tools/zig/bolt_image_fixture.py | 9 +---- tools/zig/bolt_reference.py | 30 ++------------ tools/zig/bolt_vision_reference.py | 25 +----------- zig/kernels/cuda/ggml.cu | 26 +++---------- zig/kernels/cuda/ggml_iq2xxs_tables.cuh | 6 +-- zig/kernels/cuda/nemotron_ops.cu | 6 +-- zig/kernels/cuda/radio_vision.cu | 14 ++----- zig/src/core/checkpoint.zig | 3 +- zig/src/cuda/experts.zig | 2 +- zig/src/cuda/ggml.zig | 7 ++-- zig/src/families/nemotron/config.zig | 2 +- zig/src/families/nemotron/cuda_engine.zig | 2 +- zig/src/families/nemotron/cuda_glue.zig | 3 +- zig/src/families/nemotron/cuda_lanes.zig | 3 +- zig/src/families/nemotron/cuda_modelopt.zig | 5 +-- zig/src/families/nemotron/cuda_reuse.zig | 10 ++--- zig/src/families/nemotron/cuda_state.zig | 27 +++++++------ zig/src/families/nemotron/cuda_vision.zig | 43 +++++++-------------- zig/src/families/nemotron/cuda_weights.zig | 3 +- zig/src/native/cuda.zig | 15 ++----- zig/src/native/cuda_alias.zig | 13 +++++++ zig/src/native/cuda_memory.zig | 3 +- zig/src/server/server.zig | 2 +- zig/src/vision/nemotron_image.zig | 13 ++----- zig/tests/bolt_vision_main.zig | 8 +--- 27 files changed, 86 insertions(+), 212 deletions(-) create mode 100644 zig/src/native/cuda_alias.zig diff --git a/tools/zig/bolt_ggml.py b/tools/zig/bolt_ggml.py index 04d5d88e28..612e9440cb 100644 --- a/tools/zig/bolt_ggml.py +++ b/tools/zig/bolt_ggml.py @@ -1,12 +1,4 @@ -"""Reference decode of Nemotron 3.5 Bolt's ggml-packed tensors (IQ2_XXS, Q8_0). - -Bolt's ModelOpt export stores ggml blocks in safetensors as U8 [rows, blocks, bytes]: -IQ2_XXS is 66 bytes per 256 weights (fp16 d, 32 u16 qs), Q8_0 is 34 bytes per 32 -(fp16 d, 32 int8). The IQ2_XXS tables are read from the CUDA header the kernels -include, so this file and the GPU decode the same grid. - - python3 tools/zig/bolt_ggml.py MODEL_DIR TENSOR_NAME [ROWS] -""" +"""Reference decode of Bolt's ggml tensors (IQ2_XXS, Q8_0), tables from the CUDA header: MODEL_DIR TENSOR [ROWS].""" import json import re diff --git a/tools/zig/bolt_ggml_fixture.py b/tools/zig/bolt_ggml_fixture.py index 440f0c4128..56ec1af8ca 100644 --- a/tools/zig/bolt_ggml_fixture.py +++ b/tools/zig/bolt_ggml_fixture.py @@ -1,10 +1,4 @@ -"""Fixture for `tf-cuda-test ggml`: real Bolt tensors, random rows and float64 references of the decoded weights. - -The kernels get the checkpoint's raw blocks (packed on the GPU); the references come from bolt_ggml.py's decode, -which is bit-identical to libggml's, times the same bf16 rows in float64. - - python3 tools/zig/bolt_ggml_fixture.py MODEL_DIR OUT_DIR -""" +"""Fixture for `tf-cuda-test ggml`: real Bolt tensors, random rows, float64 references (MODEL_DIR OUT_DIR).""" import json import sys diff --git a/tools/zig/bolt_image_fixture.py b/tools/zig/bolt_image_fixture.py index 93cefc13f8..5f33408adf 100644 --- a/tools/zig/bolt_image_fixture.py +++ b/tools/zig/bolt_image_fixture.py @@ -1,11 +1,4 @@ -"""Fixtures for zig/src/vision/nemotron_image.zig: torch's bicubic antialiased resize and CLIP normalization. - -A synthetic 37x23 RGB image (pixel = (7x + 13y + 29c + (x*y) % 17) % 256) resized to 48x32 (up) and to 16x16 and -24x8 (down, the antialias path), normalized as bolt_vision_reference.preprocess does; fp32 [3, H, W] each, little -endian, concatenated in that order. Also prints target_patches for a table of sizes (the Zig test's expectations). - - python3 tools/zig/bolt_image_fixture.py zig/src/vision/fixtures/nemotron_resize.bin -""" +"""Fixtures for nemotron_image.zig: torch bicubic antialiased resizes and CLIP normalization of a synthetic image.""" import sys from pathlib import Path diff --git a/tools/zig/bolt_reference.py b/tools/zig/bolt_reference.py index a169e3be8b..640d753e92 100644 --- a/tools/zig/bolt_reference.py +++ b/tools/zig/bolt_reference.py @@ -1,24 +1,4 @@ -"""Layer-by-layer fp32 PyTorch reference forward of Nemotron 3.5 Bolt's language model. - -A debugging oracle for the native engine, written from the checkpoint's modeling_nemotron_h.py -(the CUDA-kernel path, which is what vLLM runs) and its quantized export: - - * embeddings, norms, A_log/D/dt_bias/conv1d, router weight: BF16 (router bias F32) - * Mamba in_proj / out_proj: ModelOpt NVFP4 (U8 e2m1 pairs, even column in the low nibble, - an F8_E4M3 scale every 16 inputs, an F32 global `weight_scale_2`): code * scale * global - * MoE latent projections, attention q/k/v/o, lm_head: ggml Q8_0 - * routed and shared experts: ggml IQ2_XXS (decoded exactly as tools/zig/bolt_ggml.py / libggml) - -The model is streamed one layer at a time: each layer's weights are decoded, used for every -position, and dropped; MoE layers decode only the experts some position routes to. All math is -fp32. Mamba runs a sequential per-token scan; attention has no positional encoding (the model -has no RoPE). A `State` carries conv/ssm states and the attention KV so greedy generation -continues incrementally. - - python3 tools/zig/bolt_reference.py --fixture short --compare - python3 tools/zig/bolt_reference.py --fixture short --generate 16 - python3 tools/zig/bolt_reference.py --ids ids.txt --dump out/ --dump-last 8 -""" +"""Layer-by-layer fp32 PyTorch reference of Bolt's language model, streamed a layer at a time: a debugging oracle.""" import argparse import json @@ -216,8 +196,7 @@ def moe(ck, cfg, bc, li, x, routes): p = f"{PRE}layers.{li}.mixer." E, k = bc["n_routed_experts"], bc["num_experts_per_tok"] T = x.shape[0] - # router: sigmoid(x W^T) in fp32; choice by score + e_score_correction_bias (n_group = 1, so no - # group masking); weights from the un-biased scores, normalized, times routed_scaling_factor + # router: sigmoid(x W^T); choice by score + bias (n_group 1); unbiased weights, normalized, times the scaling factor assert cfg["n_group"] == 1 and cfg["topk_group"] == 1 scores = torch.sigmoid(x @ ck.vec(p + "gate.weight").T) # [T, E] choice = scores + ck.vec(p + "gate.e_score_correction_bias") @@ -247,10 +226,7 @@ def lm_head(ck, x, rows=16384): def forward(ck, ids, st=None, dump=None, dump_last=8, logits_last=None, verbose=True, splice=None): - """Run `ids` (continuing `st`); returns (logits [T or logits_last, vocab] f32, state, routes). - - `splice` = (positions LongTensor [n], rows f32 [n, hidden]) overwrites those positions' token - embeddings (the multimodal path: image-context positions take projected vision embeddings).""" + """Run `ids` (continuing `st`) -> (logits, state, routes); `splice` = (positions, rows) overwrites embeddings.""" cfg = ck.cfg st = st or State() eps = cfg["layer_norm_epsilon"] diff --git a/tools/zig/bolt_vision_reference.py b/tools/zig/bolt_vision_reference.py index bb8c70f59f..361dde882b 100644 --- a/tools/zig/bolt_vision_reference.py +++ b/tools/zig/bolt_vision_reference.py @@ -1,27 +1,4 @@ -"""fp32 PyTorch reference of Nemotron 3.5 Bolt's IMAGE path, spliced into the LM reference. - -A debugging oracle for the native engine, written from the checkpoint's -image_processing_nemotron_h_omni.py / modeling_radio.py / modeling_nemotron_h_omni.py and checked -against vLLM's nano_nemotron_vl.py + radio.py (the code that produced the fixtures): - - preprocess one dynamic-resolution tile per image (no InternVL tiling, no thumbnail): the patch - grid is ceil-ish(H/16) x ceil-ish(W/16), scaled down to <= max_num_patches (13312), - up to >= min_num_patches (1024), each side rounded to even; the RGB uint8 image is - resized with torch bicubic (align_corners=False, antialias=True) in fp32, then - (x/255 - mean)/std with the CLIP mean/std - prompt + * (hp*wp/4) + replaces the chat template's "" - tower C-RADIO ViT-H/16: patch linear (768->1280, no bias) + CPE position embedding - (128x128 table bilinearly resized to max(hp,wp)^2, cropped to hp x wp) + 10 prefix - tokens (3 cls + 7 registers); 32 pre-LN blocks (LayerNorm eps 1e-6, 16 heads x 80, - fused qkv with bias, exact-erf GELU MLP 1280->5120->1280, no layer scale); no final - norm; the 10 prefix tokens are dropped, only patch tokens are used - projector pixel shuffle v2 (2x2 -> 5120 channels ordered [dy, dx, c]), RMSNorm(5120, eps 1e-5), - Linear 5120->20480, ReLU^2, Linear 20480->4096 (no biases) - splice projector rows overwrite the token embeddings at the (id 18) positions - - python3 tools/zig/bolt_vision_reference.py --case 0 --generate 8 - python3 tools/zig/bolt_vision_reference.py --case 1 --vision-only --dump out/ -""" +"""fp32 PyTorch reference of Bolt's image path (preprocess, C-RADIO tower, projector), spliced into the LM one.""" import argparse import copy diff --git a/zig/kernels/cuda/ggml.cu b/zig/kernels/cuda/ggml.cu index eb401dfed9..10d54dd578 100644 --- a/zig/kernels/cuda/ggml.cu +++ b/zig/kernels/cuda/ggml.cu @@ -1,16 +1,4 @@ -// Bolt's ggml blocks (IQ2_XXS, Q8_0) on mma.sync: grouped experts and dense projections, decode rows and prompt rows. -// -// The weight is mma's A operand (16 output columns a warp tile), the rows are B (8 rows an n8 tile), so a lane's -// output columns are the ones it decoded and each 32-input group's scale stays in the lane. The "k" order is ours, -// the same in both operands: in a 32-input group, lane t takes inputs 8t .. 8t + 3 in step 0 and 8t + 4 .. 8t + 7 in -// step 1, so an IQ2_XXS lane decodes one whole grid entry a column. Codes become exact bf16 integers; each group -// sums into fp32 by mma, then takes its column's fp32 scale. -// -// Weights are repacked at load (tf_ggml_pack), a pure byte permutation of the checkpoint's rows, same size: -// IQ2_XXS [n/16][k/256][1056]: 8 groups of [gq 8][col gq, gq + 8][8 bytes qs], then [gq 8][2] fp16 d -// Q8_0 [n/16][k/32][544]: [lane 32][(step 0, col gq), (0, gq + 8), (1, gq), (1, gq + 8)][4 int8], then d as above -// (lane (gq, t) step s holds inputs 8t + 4s .. + 3) -// A matrix's 16-column tiles follow each other; a grouped layer's experts follow each other. +// Bolt's ggml blocks (IQ2_XXS, Q8_0) on mma.sync, weights as A and rows as B: grouped experts and dense rows. #include #include @@ -19,6 +7,7 @@ #include "experts.cuh" #include "ggml_iq2xxs_tables.cuh" +// Lane t of a 32-input group holds inputs 8t .. 8t + 3 (step 0) and 8t + 4 .. 8t + 7 (step 1), in both operands. namespace tf_ggml { __device__ __forceinline__ uint4 ld_nc(const uint8_t* p) { @@ -119,8 +108,7 @@ struct Q8 { } }; -// IQ2_XXS's grid in shared memory, 256 entries of 8 bytes: indices diverge across a warp. Signs and the bf16 -// conversion are arithmetic; tables of them cost more in bank conflicts than they save. +// IQ2_XXS's grid in shared memory (indices diverge across a warp); signs and the bf16 conversion are arithmetic. constexpr int TABLE_WORDS = 256; template @@ -197,8 +185,7 @@ __device__ __forceinline__ void epilogue_rows(const float (&acc)[NT][4], void* o constexpr int U = 2; // loads a warp has in flight, twice over -// Decode rows: a CTA per TPC (item, 16-column tile) units, KW warps each over K's loads (Q groups a load), summed in -// a fixed order. +// Decode rows: a CTA per TPC (item, 16-column tile) units, KW warps each over K's loads, summed in a fixed order. template __device__ __forceinline__ void gemv_body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint8_t* __restrict__ W, int KG, int NT16, @@ -414,8 +401,7 @@ __device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X, } // namespace tf_ggml -// The entries: format x (decode with 4 warps over K a tile, decode with a warp a tile and 4 tiles a CTA, prompt) x -// epilogue (0 fp32, 1 relu^2 bf16, 3 bf16); items null for dense rows. +// Entries: format x (decode KW 4, decode TPC 4, prompt) x epilogue (0 fp32, 1 relu^2 bf16, 3 bf16); dense: no items. #define TF_GGML_GEMV(NAME, F, EPI, KW, TPC) \ extern "C" __global__ void __launch_bounds__(KW * TPC * 32) NAME( \ const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint8_t* __restrict__ W, int KG, int NT16, \ @@ -443,7 +429,7 @@ TF_GGML_ENTRIES(q8, Q8, 3, bf16) #undef TF_GGML_ENTRIES #undef TF_GGML_GEMV -// The load repack: `count` matrices [n][k / KB][B] of checkpoint blocks into the 16-column tiles above, a byte a thread. +// Load repack, a byte a thread: `count` matrices [n][k / KB][B] of checkpoint blocks into the 16-column tiles. extern "C" __global__ void tf_ggml_pack(const uint8_t* __restrict__ src, uint8_t* __restrict__ dst, int iq2, int n, int kb, long long total) { const int B = iq2 ? 66 : 34, TB = iq2 ? 1056 : 544; diff --git a/zig/kernels/cuda/ggml_iq2xxs_tables.cuh b/zig/kernels/cuda/ggml_iq2xxs_tables.cuh index 687b457d33..f3f3b4e2b1 100644 --- a/zig/kernels/cuda/ggml_iq2xxs_tables.cuh +++ b/zig/kernels/cuda/ggml_iq2xxs_tables.cuh @@ -1,7 +1,4 @@ -// IQ2_XXS lookup tables, copied from ggml (llama.cpp ggml/src/ggml-common.h), MIT License, -// Copyright (c) 2023-2026 The ggml authors. Read by the Bolt expert kernels and by -// tools/zig/bolt_ggml.py, so the reference and the GPU decode the same bytes. -// The grid sits in global memory: kernels copy it to shared memory with coalesced loads, not constant-cache ones. +// IQ2_XXS tables from ggml-common.h, MIT License, Copyright (c) 2023-2026 The ggml authors; also read by bolt_ggml.py. #pragma once #include @@ -16,6 +13,7 @@ static __device__ __constant__ const uint8_t tf_ksigns_iq2xs[128] = { 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, }; +// In global memory: kernels copy it to shared memory with coalesced loads. static __device__ const uint64_t tf_iq2xxs_grid[256] = { 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, diff --git a/zig/kernels/cuda/nemotron_ops.cu b/zig/kernels/cuda/nemotron_ops.cu index d4304e8f7e..b55bbc712d 100644 --- a/zig/kernels/cuda/nemotron_ops.cu +++ b/zig/kernels/cuda/nemotron_ops.cu @@ -292,8 +292,7 @@ extern "C" __global__ void tf_add_row(void* __restrict__ y, int y_f32, const flo } } -// LatentMoE (Bolt): a row's routed outputs in the latent width, weighted and summed in fp32, then bf16 like the -// reference's index_add_ sum. Y [R * K, L] fp32 (pair rows in slot order), WT [R * K], Z [R, L]; one block a row. +// LatentMoE (Bolt): Z [R, L] bf16, the fp32 sum of Y [R * K, L] (slot order) times WT [R * K]; one block a row. extern "C" __global__ void tf_latent_sum(const float* __restrict__ Y, const float* __restrict__ WT, __nv_bfloat16* __restrict__ Z, int K, int L) { const int64_t r = blockIdx.x; @@ -313,8 +312,7 @@ extern "C" __global__ void tf_add_bf16(__nv_bfloat16* __restrict__ out, const __ out[i] = __float2bfloat16(__bfloat162float(a[i]) + y); } -// An image prompt's chunk: rows whose SRC names a feature row take it in place of their token's embedding. -// H [rows, D] bf16, F [image rows, D] bf16, SRC [rows] (-1: a text row); one block a row. +// Image prompt chunk: row r of H [rows, D] takes F's row SRC[r] for its token embedding (-1: text); a block a row. extern "C" __global__ void tf_splice_rows(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ F, const int* __restrict__ SRC, int D) { const int64_t r = blockIdx.x; const int s = SRC[r]; diff --git a/zig/kernels/cuda/radio_vision.cu b/zig/kernels/cuda/radio_vision.cu index dcfffded2c..f930c54eb6 100644 --- a/zig/kernels/cuda/radio_vision.cu +++ b/zig/kernels/cuda/radio_vision.cu @@ -1,6 +1,4 @@ -// Nemotron 3.5 Bolt's vision tower (C-RADIO ViT-H/16) and projector: glue around cuBLASLt's bf16 GEMMs, and attention. -// The residual stream stays fp32 (its rows grow to |x| ~1700); norms, softmax and epilogues compute in fp32, GEMM -// inputs are rounded to bf16 once. Oracle: tools/zig/bolt_vision_reference.py (fp32 PyTorch). +// Bolt's C-RADIO ViT-H/16 tower and projector around cuBLASLt bf16 GEMMs; residual, norms, softmax, epilogues fp32. #include #include @@ -23,8 +21,7 @@ __device__ __forceinline__ float block_sum(float v, float* part) { return t; } -// x [prefix + n, D] fp32: rows below `prefix` the cls/register tokens, then acc (the patch Linear) plus the CPE -// position rows, bilinear from the 128x128 table: (p0 w0 + p1 w1) h0 + (p2 w0 + p3 w1) h1 as aten's interpolate. +// x [prefix + n, D] fp32: cls/register rows, then acc (patch Linear) plus CPE rows bilinear from the 128x128 table. extern "C" __global__ void __launch_bounds__(256) tf_rv_embed(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ pos, const int* __restrict__ idx, @@ -91,9 +88,7 @@ __device__ __forceinline__ void mma16816(float* c, const uint32_t* a, uint32_t b : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); } -// Full (bidirectional) attention over all `rows` tokens, heads of 80: flash style, one warp 16 queries, 64-key -// tiles in shared memory, bf16 mma.sync with fp32 scores, an online softmax in fp32 (exp2 of pre-scaled scores) and -// fp32 output accumulators; out [rows, heads * 80] bf16 (the reference's transpose(0, 1).reshape(T, C)). +// Bidirectional flash attention, heads of 80: a warp 16 queries, 64-key tiles, bf16 mma, fp32 softmax; bf16 out. #define RV_HD 80 #define RV_BK 64 extern "C" __global__ void __launch_bounds__(128) tf_rv_attention(const __nv_bfloat16* __restrict__ q, @@ -232,8 +227,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_rv_bias_residual(const floa for (int i = threadIdx.x; i < D; i += blockDim.x) x[row * D + i] += acc[row * D + i] + bf(bias[i]); } -// Pixel shuffle v2 (2x2 patches to one row of 4D, channel (dy * 2 + dx) * D + c) of the patch rows behind `prefix`, -// then RMSNorm(4D, eps) with weight w in fp32; bf16 out. A block an output row (grid gh/2 x gw/2, row-major). +// Pixel shuffle v2 (2x2 patches to a 4D row) of the rows past `prefix`, then fp32 RMSNorm(4D) by w; a block a row. extern "C" __global__ void __launch_bounds__(256) tf_rv_shuffle_rms(const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, __nv_bfloat16* __restrict__ y, int prefix, int gw, diff --git a/zig/src/core/checkpoint.zig b/zig/src/core/checkpoint.zig index 8e62ac1575..47ff2b5011 100644 --- a/zig/src/core/checkpoint.zig +++ b/zig/src/core/checkpoint.zig @@ -59,8 +59,7 @@ pub const Checkpoint = struct { io: Io, files: std.ArrayList(st.File) = .empty, used: std.StringHashMapUnmanaged(void) = .empty, - /// Prepended to every name get, has and skip take: a wrapper's part read with the names it has alone - /// ("language_model." in Nemotron 3.5 Bolt, beside its vision tower's). + /// Prepended to the names get, has and skip take ("language_model." for Bolt's LM beside its vision tower). scope: []const u8 = "", fn scoped(self: *const Checkpoint, buf: *[320]u8, name: []const u8) ![]const u8 { diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 8bafabc226..e59b21daf2 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -87,7 +87,7 @@ pub fn promptTile(f: Format) usize { }; } -/// IQ2_XXS experts (ggml.cu): relu^2 up, down to fp32 or bf16; experts packed by Iq2xxsExperts.pack from the raw blocks. +/// IQ2_XXS experts (ggml.cu): relu^2 up, down to fp32 or bf16; packed by Iq2xxsExperts.pack from the raw blocks. pub const Iq2xxsExperts = ggml.Kernels; fn iq2Run(k: *const Iq2xxsExperts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { diff --git a/zig/src/cuda/ggml.zig b/zig/src/cuda/ggml.zig index bd7be5c7e7..e7bb967ac3 100644 --- a/zig/src/cuda/ggml.zig +++ b/zig/src/cuda/ggml.zig @@ -145,8 +145,8 @@ pub const Kernels = struct { /// The common launch: `items`, `counts` and `members` 0 for dense rows (`rows` of them), else a plan's. fn go(g: *const Kernels, s: Stream, prompt: bool, out_kind: Out, kind: Kind, x: u64, stride: usize, slots: usize, w: u64, n: usize, k: usize, p: ?grouped.Plan, out: u64, units: usize, rows: usize, skip: i32) !void { try check(kind, n, k, prompt); - const ki = @intFromEnum(kind); - const oi = @intFromEnum(out_kind); + const ki = @backingInt(kind); + const oi = @backingInt(out_kind); const nt16 = n / 16; if (units == 0) return; var a: launch_.Args = .{}; @@ -187,8 +187,7 @@ pub const Kernels = struct { return if (rows <= decode_rows) g.decode(s, w, x, ldx, rows, out, out_kind) else g.promptRows(s, w, x, ldx, rows, out, out_kind); } - /// Grouped experts on a plan of at most `items` items (16 pairs decode, 64 prompt): stacked [E] packed matrices - /// (n, k) at `w`; x rows `stride` apart, token rows when `slots` > 0 (pair / slots), else pair rows; out pair rows. + /// Grouped experts, at most `items` items: [E] packed (n, k) at `w`; x token rows if `slots` > 0, else pair rows. pub fn experts(g: *const Kernels, s: Stream, prompt: bool, kind: Kind, out_kind: Out, x: u64, stride: usize, slots: usize, w: u64, n: usize, k: usize, p: grouped.Plan, out: u64, items: usize, skip: i32) !void { return g.go(s, prompt, out_kind, kind, x, stride, slots, w, n, k, p, out, items, 0, skip); } diff --git a/zig/src/families/nemotron/config.zig b/zig/src/families/nemotron/config.zig index befe4439ac..5eb9192282 100644 --- a/zig/src/families/nemotron/config.zig +++ b/zig/src/families/nemotron/config.zig @@ -327,7 +327,7 @@ fn eos(c: *Config, e: std.json.Value) void { }; } -/// nemotron_h_puzzle's block_configs: a block type per layer, each MoE layer with its own experts and width (one top-k). +/// nemotron_h_puzzle's block_configs: a block type per layer, each MoE layer its own experts and width (one top-k). fn puzzle(c: *Config, blocks: []const std.json.Value, why: *Why) !void { if (blocks.len > max_layers) { why.set("config.json has {d} layers; the native Nemotron engine reads up to {d}", .{ blocks.len, max_layers }); diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index 0e2b8fbb9e..c93f0b21ce 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -154,7 +154,7 @@ pub const Engine = struct { if (opts.slide) return error.SlideNeedsMlx; // the learner reads MLX 4-bit projections try e.k.loadModelopt(ctx); } - // Bolt (LatentMoE, ggml blocks): no MTP head in the checkpoint, and our own glue (captured sets are Lightning's) + // Bolt (LatentMoE, ggml blocks): no MTP head in the checkpoint; our own glue (captured sets are Lightning's) const mtp = opts.mtp and e.c.latent == 0; if (e.c.latent > 0) { if (e.k.triton != null) return error.CapturedGlueIsLightnings; diff --git a/zig/src/families/nemotron/cuda_glue.zig b/zig/src/families/nemotron/cuda_glue.zig index 2bf49b46a9..4a3d63d347 100644 --- a/zig/src/families/nemotron/cuda_glue.zig +++ b/zig/src/families/nemotron/cuda_glue.zig @@ -87,10 +87,9 @@ pub const AttnSeg = struct { kc: u64, vc: u64, meta: u64, row0: usize, rows: usi const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, row0: [max_segs]c_int, n: c_int }; /// The router's K slices and the shared bytes its block stages: 16 rows and 16 experts of one slice, padded a word. -/// Wide rows (Bolt's 4096) take more slices, so a block stays within the 48 KiB launches get without an opt-in. pub fn routerShape(d: usize) struct { sk: usize, shared: u32 } { var sk = triton.divisor(d / 64, 6); - while (staged(d, sk) > 48 << 10 and (d / 64) % (sk * 2) == 0) sk *= 2; + while (staged(d, sk) > 48 << 10 and (d / 64) % (sk * 2) == 0) sk *= 2; // wide rows (Bolt's): within 48 KiB, no opt-in return .{ .sk = sk, .shared = @intCast(staged(d, sk)) }; } diff --git a/zig/src/families/nemotron/cuda_lanes.zig b/zig/src/families/nemotron/cuda_lanes.zig index baf7e98a4d..9146fca387 100644 --- a/zig/src/families/nemotron/cuda_lanes.zig +++ b/zig/src/families/nemotron/cuda_lanes.zig @@ -257,8 +257,7 @@ pub const Cuda = struct { e.decide = if (s.decide != null and e.lessons() != null) s.decide.? else lesson_route.none; try e.setSampling(s.sampling); s.cached = 0; - // an image prompt: its images through the tower, their rows spliced over the placeholders as each chunk embeds; - // never resumed or kept (every image's placeholders are the same ids) + // an image prompt: tower rows spliced over the placeholders per chunk; never resumed or kept (same ids) var media: ?vision.Media = null; if (s.images.len > 0) { const tower = self.tower orelse return error.NemotronServedWithoutVision; diff --git a/zig/src/families/nemotron/cuda_modelopt.zig b/zig/src/families/nemotron/cuda_modelopt.zig index 63d369560e..5055bbe165 100644 --- a/zig/src/families/nemotron/cuda_modelopt.zig +++ b/zig/src/families/nemotron/cuda_modelopt.zig @@ -25,7 +25,7 @@ fn ggmlBlock(k: Kind) [2]usize { } fn kindOf(t: Tensor) !Kind { - // ggml blocks (Bolt) keep each row's blocks as a third axis: [rows, blocks, block bytes]; NVFP4 packs pairs as [rows, k / 2] + // ggml blocks (Bolt): [rows, blocks, block bytes]; NVFP4 packs pairs as [rows, k / 2] if (t.dtype == .u8 and t.rank == 3) return switch (t.dim(2)) { 34 => .q8_0, 66 => .iq2xxs, @@ -179,8 +179,7 @@ fn iq2Experts(L: *Loader, ck: *core.Checkpoint, label: []const u8, count: usize, return out; } -/// A LatentMoE layer (Bolt): the router and bias over this layer's experts, fc1/fc2_latent_proj, routed IQ2_XXS -/// experts in the latent width, and the shared expert at full width as a one-expert layer of its own. +/// A LatentMoE layer (Bolt): router and bias, fc1/fc2_latent_proj, latent IQ2_XXS experts, full-width shared expert. pub fn latent(L: *Loader, ck: *core.Checkpoint, label: []const u8, pre: []const u8, c: Config, layer: usize) !W.MoE { const lm = c.layerMoe(layer); const E = lm.experts; diff --git a/zig/src/families/nemotron/cuda_reuse.zig b/zig/src/families/nemotron/cuda_reuse.zig index 55084e774f..180fa040a4 100644 --- a/zig/src/families/nemotron/cuda_reuse.zig +++ b/zig/src/families/nemotron/cuda_reuse.zig @@ -1,5 +1,4 @@ //! Nemotron CUDA prompt states for the shared cache: a prompt's key and value rows, the Mamba states, the head's rows. -//! On the device, or (a discrete card's `host`) in pinned host memory: the key and value rows, then the Mamba states. const std = @import("std"); const cuda = @import("cuda"); @@ -10,12 +9,11 @@ const Engine = @import("cuda_engine.zig").Engine; const Head = @import("cuda_mtp.zig").Head; const marks = @import("cuda_marks.zig"); -/// One kept state: its Mamba states in one device buffer, its key and value rows in another, where the engine stood. /// A host state keeps both in `host` (key and value rows first); `mamba` is then an armed row's staging until `finish`. pub const Copy = struct { mamba: cuda.DeviceBuffer, kv: cuda.DeviceBuffer, - host: ?cuda.HostBuffer = null, + host: ?cuda.HostBuffer = null, // pinned: rows, then Mamba states; `mamba` is then an armed row's staging until `finish` head: bool, pos: usize, parity: usize, @@ -102,8 +100,7 @@ fn move(t: *const Target, at: usize, copy: *const Copy, into: bool, mamba: bool) return off[0] + off[1]; } -/// Copies between the bound sequence and a host copy (`into`: live to host), key and value rows first; `mamba`: the -/// Mamba states too, else from `staging` (an armed row's) when given. Queued on the engine's stream. +/// Bound sequence <-> host copy (`into`: to host), rows first; Mamba states too if `mamba`, else from `staging`. fn hostMove(t: *const Target, at: usize, host: []u8, into: bool, mamba: bool, staging: ?u64) !void { const o = t.e.ops(); var store: [max_parts]Part = undefined; @@ -191,8 +188,7 @@ pub fn finish(t: Target) !void { for (e.armed.marks()) |m| { const copy: *Copy = @ptrCast(@alignCast(m.copy)); if (t.host()) { - // the chunk wrote this row's Mamba states into the staging buffer: rows and states to the host, then the - // staging back (the stream orders its next writer after these copies) + // the chunk wrote this row's Mamba states to staging: rows and states to the host, then staging back var h = cuda.HostBuffer.alloc(e.ctx.d, stateBytes(e.c, m.at, copy.head)) catch continue; // stays not whole hostMove(&t, m.at, h.bytes, true, false, copy.mamba.ptr) catch { h.free(); diff --git a/zig/src/families/nemotron/cuda_state.zig b/zig/src/families/nemotron/cuda_state.zig index c65e23c804..f0bdb2626c 100644 --- a/zig/src/families/nemotron/cuda_state.zig +++ b/zig/src/families/nemotron/cuda_state.zig @@ -97,19 +97,19 @@ fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { const items = kern.maxItems(@intCast(pairs), c.experts + 2, 16); const lat: usize = @max(c.latent, 1); // LatentMoE's rows in the latent width (Bolt), else unused return .{ - W * c.vocab * 2, 16, R * 4, - R * D * 2, c.vocab * 2, 16, - R * D * 2, R * D * 2, R * D * 2, - R * D * 2, R * D / 64 * 4, R * D * 2, - R * c.projDim() * 2, R * cd * 2, R * xd * 2, - R * xd * 2, R * xd / 64 * 4, R * c.qkvDim() * 2, - R * qd * 2, R * qd * 2, R * qd / 64 * 4, - W * nch * qd * 4, W * nch * c.heads * 4, W * nch * c.heads * 4, + W * c.vocab * 2, 16, R * 4, + R * D * 2, c.vocab * 2, 16, + R * D * 2, R * D * 2, R * D * 2, + R * D * 2, R * D / 64 * 4, R * D * 2, + R * c.projDim() * 2, R * cd * 2, R * xd * 2, + R * xd * 2, R * xd / 64 * 4, R * c.qkvDim() * 2, + R * qd * 2, R * qd * 2, R * qd / 64 * 4, + W * nch * qd * 4, W * nch * c.heads * 4, W * nch * c.heads * 4, glue.max_router_slices * R * c.experts * 4, R * ns * 4, R * ns * 4, - pairs * 4, items * 12, 8, - pairs * 4, (pairs + 1023) / 1024 * (c.experts + 2) * 4, pairs * c.expert_width * 2, - pairs * D * 4, shared_words * 4, R * kern.pocket_width * 4, - R * lat * 2, R * lat * 2, R * 4, + pairs * 4, items * 12, 8, + pairs * 4, (pairs + 1023) / 1024 * (c.experts + 2) * 4, pairs * c.expert_width * 2, + pairs * D * 4, shared_words * 4, R * kern.pocket_width * 4, + R * lat * 2, R * lat * 2, R * 4, }; } @@ -268,8 +268,7 @@ pub const Buffers = struct { &b.logits, &b.p_meta, &b.p_ids, &b.p_hidden, &b.p_logits, &b.p_sampled, &b.emb, &b.h[0], &b.h[1], &b.y, &b.xs, &b.delta, &b.proj, &b.p_xc, &b.sy, &b.g, &b.gxs, &b.qkv, &b.q, &b.att, &b.axs, &b.po, &b.pm, &b.pl, &b.part, &b.pick, &b.wts, &b.plan.members, &b.plan.items, &b.plan.counts, - &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, - &b.lat, &b.z, &b.zeros, + &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, &b.lat, &b.z, &b.zeros, }; } diff --git a/zig/src/families/nemotron/cuda_vision.zig b/zig/src/families/nemotron/cuda_vision.zig index 2dd6c4b625..b017d534ad 100644 --- a/zig/src/families/nemotron/cuda_vision.zig +++ b/zig/src/families/nemotron/cuda_vision.zig @@ -1,14 +1,4 @@ -//! Nemotron 3.5 Bolt's image tower on CUDA: C-RADIO ViT-H/16 (vision_model.radio_model.) and the mlp1 projector, -//! bf16 weights, cuBLASLt GEMMs with fp32 accumulation, an fp32 residual stream; kernels in radio_vision.cu. -//! -//! `encode` turns one image's patches (vision/nemotron_image.zig's `prepare`) into the bf16 [tokens, 4096] rows that -//! replace the prompt's `` placeholders: `` (19), then `` (18) x tokens, then `` (20), the -//! image directly before the question. Oracle: tools/zig/bolt_vision_reference.py. -//! -//! VRAM (one arena, made at load): weights 1.681 GB (tower 1.303 GB: blocks 1.259 GB, the 128x128 position table -//! 42 MB, patch Linear and prefix tokens 2 MB; projector 377 MB) plus scratch sized for the largest image (13,312 -//! patches) 0.702 GB: 2.383 GB in all, and the caller's output rows (3,328 x 4096 bf16 at most, 27 MB). -//! The video embedder (video_embedder.weight, 3.9 MB) is checked but not loaded: video is not this module's. +//! Bolt's image tower on CUDA: C-RADIO ViT-H/16 and the mlp1 projector, bf16 weights, fp32 residual; radio_vision.cu. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -60,10 +50,9 @@ pub fn checkConfig(gpa: std.mem.Allocator, bytes: []const u8) !void { const root = parsed.value.object; const v = (root.get("vision_config") orelse return error.UnsupportedBoltVision).object; const ints = [_]struct { std.json.ObjectMap, []const u8, i64 }{ - .{ v, "hidden_size", hidden }, .{ v, "num_attention_heads", heads }, .{ v, "num_hidden_layers", depth }, - .{ v, "patch_size", patch }, .{ v, "num_cls_tokens", 3 }, .{ v, "num_registers", 7 }, - .{ root, "projector_hidden_size", projector_hidden }, .{ root, "img_context_token_id", image_token }, - .{ root, "vit_hidden_size", hidden }, + .{ v, "hidden_size", hidden }, .{ v, "num_attention_heads", heads }, .{ v, "num_hidden_layers", depth }, + .{ v, "patch_size", patch }, .{ v, "num_cls_tokens", 3 }, .{ v, "num_registers", 7 }, + .{ root, "projector_hidden_size", projector_hidden }, .{ root, "img_context_token_id", image_token }, .{ root, "vit_hidden_size", hidden }, }; for (ints) |w| if (num(w[0].get(w[1])) != w[2]) return error.UnsupportedBoltVision; const llm = (root.get("llm_config") orelse return error.UnsupportedBoltVision).object; @@ -125,7 +114,7 @@ pub const Tower = struct { mod: cuda.Module, f: Fns, lib: lt.Library, - arena: DeviceBuffer, + arena: DeviceBuffer, // weights 1.681 GB, scratch for max_patches 0.702 GB at: usize = 0, scratch: u64 = 0, embed_w: u64 = 0, @@ -139,8 +128,7 @@ pub const Tower = struct { /// Blocks `encode` runs (tests: 0 leaves the embedded rows in Stages.x); the projector runs only after all 32. run_depth: usize = depth, - /// The checkpoint's vision_model. and mlp1. tensors (bf16) from `dir` (Bolt's own directory: the language - /// model's loader skips them). Uploads synchronously; the arena is sized for `max_patches`. + /// Loads vision_model. and mlp1. (bf16) from Bolt's `dir`, synchronously, into an arena sized for `max_patches`. pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !*Tower { try check(gpa, io, dir); const t = try gpa.create(Tower); @@ -170,18 +158,18 @@ pub const Tower = struct { } } _ = try vk.expect("vision_model.summary_idxs", .i64, &.{2}); // the summary heads' cls rows: not used for images - _ = try vk.expect(model ++ "patch_generator.video_embedder.weight", .bf16, &.{ hidden, 2 * patch_width }); // video + _ = try vk.expect(model ++ "patch_generator.video_embedder.weight", .bf16, &.{ hidden, 2 * patch_width }); // video: checked, not loaded t.embed_w = try t.get(&vk, model ++ "patch_generator.embedder.weight", &.{ hidden, patch_width }); t.pos = try t.get(&vk, model ++ "patch_generator.pos_embed", &.{ 1, cpe_side * cpe_side, hidden }); t.cls = try t.get(&vk, model ++ "patch_generator.cls_token.token", &.{ prefix_tokens, hidden }); for (&t.blocks, 0..) |*b, i| { const shapes = .{ - .{ "norm1_w", "norm1.weight", .{hidden} }, .{ "norm1_b", "norm1.bias", .{hidden} }, + .{ "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.fc1.weight", .{ mlp, hidden } }, .{ "fc1_b", "mlp.fc1.bias", .{mlp} }, - .{ "fc2_w", "mlp.fc2.weight", .{ hidden, mlp } }, .{ "fc2_b", "mlp.fc2.bias", .{hidden} }, + .{ "norm2_w", "norm2.weight", .{hidden} }, .{ "norm2_b", "norm2.bias", .{hidden} }, + .{ "fc1_w", "mlp.fc1.weight", .{ mlp, hidden } }, .{ "fc1_b", "mlp.fc1.bias", .{mlp} }, + .{ "fc2_w", "mlp.fc2.weight", .{ hidden, mlp } }, .{ "fc2_b", "mlp.fc2.bias", .{hidden} }, }; inline for (shapes) |sh| { var nb: [128]u8 = undefined; @@ -243,8 +231,7 @@ pub const Tower = struct { return p; } - /// One image's patches (fp32 [gh * gw, 768], nemotron_image.prepare's) to bf16 [gh * gw / 4, 4096] rows at `out`, - /// on stream `s`. Returns once the host's inputs are copied (the stream is synchronized); `out` is then written. + /// One image's patches (fp32 [gh * gw, 768]) to bf16 [gh * gw / 4, 4096] rows at `out` on `s`; inputs copied. pub fn encode(t: *Tower, s: cuda.Stream, patches: []const f32, gh: u32, gw: u32, out: u64) !Stages { const n = @as(usize, gh) * gw; if (gh == 0 or gw == 0 or gh % 2 != 0 or gw % 2 != 0 or n > max_patches) return error.BadImageGrid; @@ -370,8 +357,7 @@ fn bf16(v: f32) u16 { return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); } -/// CPE: the 128x128 table resized (bilinear, align_corners=False) to max(gh, gw) square and cropped to gh x gw. -/// Each patch's 4 table rows (y0 x0, y0 x1, y1 x0, y1 x1) and weights (w0, w1, h0, h1), as aten computes them in fp32. +/// CPE taps: each patch's 4 rows and weights in the 128x128 table resized to max(gh, gw) and cropped, as aten (fp32). pub fn cpeTaps(gh: u32, gw: u32, idx: []i32, wts: []f32) void { const m = @max(gh, gw); for (0..gh) |i| { @@ -468,8 +454,7 @@ pub const Media = struct { } }; -/// Each image's run of image_token rows between image_start and image_end, in order, its length the image's tokens; -/// `src` gets each row's feature row (-1 for text). Returns the feature rows. +/// Each image's image_token run between image_start and image_end; `src` gets each row's feature row (-1: text). pub fn placeholders(prompt: []const u32, images: []const lanes_.Image, src: []i32) !usize { @memset(src, -1); var next: usize = 0; // feature rows handed out diff --git a/zig/src/families/nemotron/cuda_weights.zig b/zig/src/families/nemotron/cuda_weights.zig index 086fbc07a1..26a3ba6834 100644 --- a/zig/src/families/nemotron/cuda_weights.zig +++ b/zig/src/families/nemotron/cuda_weights.zig @@ -30,8 +30,7 @@ pub const Pocket = kern.Pocket; /// The shared expert's down projection rest in two halves (experts E and E + 1); its learned change reads both. pub const MoE = struct { router: u64, bias: u64, experts: Experts, rest: [2]u64 = .{ 0, 0 }, adapter: ?Adapter = null, pocket: ?Pocket = null, latent: ?Latent = null }; -/// LatentMoE's own parts (Bolt): rows go to the latent width and back around the routed experts, and the shared -/// expert reads the full-width rows on its own (`experts` then holds the routed experts alone). +/// LatentMoE's parts (Bolt): latent-width projections around the routed `experts`; `shared` reads full-width rows. pub const Latent = struct { fc1: QLinear, fc2: QLinear, shared: Experts }; pub const Block = struct { kind: Kind, norm: u64, mamba: Mamba = undefined, attn: Attention = undefined, moe: MoE = undefined }; pub const Mtp = struct { enorm: u64, hnorm: u64, eh_proj: QLinear, attn_norm: u64, attn: Attention, moe_norm: u64, moe: MoE, final_norm: u64 }; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 0bed892d5d..f77d0bedc2 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -12,6 +12,8 @@ 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 aliases = @import("cuda_alias.zig").aliases; +const reads = @import("cuda_alias.zig").reads; /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{nemotron.native}; @@ -38,17 +40,6 @@ pub const families: []const api.Family = blk: { break :blk &final; }; -/// The other config model_types a family reads (a wrapper's, as Bolt's nemotron_h_omni). -fn aliases(comptime F: type) []const []const u8 { - return if (@hasDecl(F, "also_model_types")) F.also_model_types else &.{}; -} - -fn reads(comptime F: type, model_type: []const u8) bool { - if (std.mem.eql(u8, model_type, F.model_type)) return true; - for (aliases(F)) |t| if (std.mem.eql(u8, model_type, t)) return true; - return false; -} - const gib: f64 = 1 << 30; /// The chip class gate entries name ("nvidia-sm121" for a GB10) of the device TF_CUDA_DEVICE picks; null without one. @@ -423,7 +414,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const host_snaps = @hasField(@TypeOf(loaded), "host_snaps") and loaded.host_snaps; const cache = if (!@hasField(@TypeOf(loaded), "target")) 0 else if (host_snaps) budget.hostCacheBytes(a, io, o.prompt_cache_gib) 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 {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 "", + 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{s}", .{ toGib(cache), if (host_snaps) " 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); diff --git a/zig/src/native/cuda_alias.zig b/zig/src/native/cuda_alias.zig new file mode 100644 index 0000000000..0a9e9bc71a --- /dev/null +++ b/zig/src/native/cuda_alias.zig @@ -0,0 +1,13 @@ +//! The config model_types a CUDA family reads: its own, then its wrappers' (as Bolt's nemotron_h_omni). +const std = @import("std"); + +/// The other config model_types a family reads (`also_model_types`). +pub fn aliases(comptime F: type) []const []const u8 { + return if (@hasDecl(F, "also_model_types")) F.also_model_types else &.{}; +} + +pub fn reads(comptime F: type, model_type: []const u8) bool { + if (std.mem.eql(u8, model_type, F.model_type)) return true; + for (aliases(F)) |t| if (std.mem.eql(u8, model_type, t)) return true; + return false; +} diff --git a/zig/src/native/cuda_memory.zig b/zig/src/native/cuda_memory.zig index 1d3a3b438c..6942f3e04d 100644 --- a/zig/src/native/cuda_memory.zig +++ b/zig/src/native/cuda_memory.zig @@ -102,8 +102,7 @@ pub fn cacheBytes(asked: ?f64, room: u64, streams: u64) u64 { return @min(room -| streams, want); } -/// Kept prompt states' budget in pinned host memory: --prompt-cache-gib (0: none), else 8 GiB, never past half of -/// what /proc/meminfo calls available. +/// Pinned host budget for kept prompt states: --prompt-cache-gib (0: none), else 8 GiB, at most half of MemAvailable. 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 * gib) else 0) else 8 << 30; const text = procText(a, io, "/proc/meminfo") orelse return 0; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 36e4812af1..3b5f0a1746 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -216,7 +216,7 @@ pub const Server = struct { return .{ .kind = kind, .text = text, .after = if (thinking) srv.text.tokenId(srv.markers.close) else null, .compiler = c }; } - /// The model's grammar compiler: tokenizer.json, config.json's vocab_size (text_config or llm_config first), the eos ids. + /// The grammar compiler: tokenizer.json, config.json's vocab_size (text_config or llm_config first), the eos ids. fn grammarCompiler(srv: *Server) ?*api.grammar.Compiler { srv.grammar_mutex.lockUncancelable(srv.io); defer srv.grammar_mutex.unlock(srv.io); diff --git a/zig/src/vision/nemotron_image.zig b/zig/src/vision/nemotron_image.zig index 3fa997e797..0abac37bc2 100644 --- a/zig/src/vision/nemotron_image.zig +++ b/zig/src/vision/nemotron_image.zig @@ -1,8 +1,4 @@ -//! Nemotron 3.5 Bolt's image preparation (image_processing_nemotron_h_omni.py as vLLM runs it): one dynamic-resolution -//! tile, torch's fp32 bicubic antialiased resize, CLIP normalization and C-RADIO's 16x16 patches. -//! -//! The oracle is tools/zig/bolt_vision_reference.py (`preprocess`, `target_patches`); its arithmetic is copied here -//! step for step. Decoding is the shared decoder's (qwen_image.decode); this module starts from RGB bytes. +//! Bolt's image prep as vLLM runs it: one tile, torch's fp32 bicubic antialiased resize, CLIP norm, 16x16 patches. const std = @import("std"); const qi = @import("qwen_image.zig"); @@ -114,8 +110,7 @@ fn taps(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Taps { return t; } -/// F.interpolate(mode="bicubic", align_corners=False, antialias=True) of fp32 planes [3, h, w] to [3, oh, ow]: -/// the width pass first, then the height pass (aten's separable order); a side already the size is not resampled. +/// Bicubic antialiased F.interpolate (align_corners=False) of fp32 [3, h, w] to [3, oh, ow]: width pass, then height. pub fn resizePlanes(gpa: std.mem.Allocator, src: []const f32, w: u32, h: u32, ow: u32, oh: u32) Error![]f32 { var cur = try gpa.dupe(f32, src); var cw: usize = w; @@ -230,8 +225,8 @@ fn testImage(gpa: std.mem.Allocator, w: u32, h: u32) !Image { test "the patch grid matches target_patches (tools/zig/bolt_image_fixture.py's table)" { const cases = [_][4]u32{ - .{ 2172, 724, 136, 46 }, .{ 640, 480, 40, 30 }, .{ 37, 23, 40, 28 }, .{ 4000, 3000, 132, 100 }, - .{ 16, 16, 32, 32 }, .{ 100, 5000, 8, 314 }, .{ 1024, 1024, 64, 64 }, .{ 512, 512, 32, 32 }, + .{ 2172, 724, 136, 46 }, .{ 640, 480, 40, 30 }, .{ 37, 23, 40, 28 }, .{ 4000, 3000, 132, 100 }, + .{ 16, 16, 32, 32 }, .{ 100, 5000, 8, 314 }, .{ 1024, 1024, 64, 64 }, .{ 512, 512, 32, 32 }, .{ 3000, 100, 188, 8 }, .{ 1920, 1080, 120, 68 }, .{ 800, 600, 50, 38 }, }; for (cases) |c| { diff --git a/zig/tests/bolt_vision_main.zig b/zig/tests/bolt_vision_main.zig index a9d8fb1a83..d64e3900be 100644 --- a/zig/tests/bolt_vision_main.zig +++ b/zig/tests/bolt_vision_main.zig @@ -1,10 +1,4 @@ -//! `tf-bolt-vision [reps]`: Nemotron 3.5 Bolt's image path end to end against the fp32 reference. -//! -//! The dump directory is tools/zig/bolt_vision_reference.py's (`--vision-only --dump DIR`): image_rgb.npy (the decoded -//! RGB the reference resized), meta.json, pixel_values.npy, tower_embed.npy, tower_out.npy, projector_out.npy. Host -//! preparation, the embedded rows, the last block's rows and the projector's rows are each compared (max |d|, the -//! relative Frobenius error, cosine over the whole tensor and the worst row's), then `reps` timed calls of `encode`. -//! Exits 1 when the projector rows' cosine is below 0.999 or the host pixels are off by more than 1e-4. +//! `tf-bolt-vision [reps]`: Bolt's image path against bolt_vision_reference.py's --dump. const std = @import("std"); const cuda = @import("cuda"); const nemotron = @import("nemotron");