From eb56483683a60591c77a58614f7ac8bddcbf4eb2 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 00:51:50 +0800 Subject: [PATCH 01/21] idol: Idol's denoiser as a native Metal family, with tf-idol and its reference case tool --- tools/zig/idol_case.py | 243 ++++++++++++++++++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol.metal | 329 +++++++++++++++++++++++++ zig/kernels/metal/idol_attention.metal | 149 +++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/families/idol/config.zig | 71 ++++++ zig/src/families/idol/host.zig | 147 +++++++++++ zig/src/families/idol/idol.zig | 19 ++ zig/src/families/idol/kernels.zig | 209 ++++++++++++++++ zig/src/families/idol/model.zig | 244 ++++++++++++++++++ zig/src/families/idol/reference.zig | 79 ++++++ zig/src/families/idol/weights.zig | 310 +++++++++++++++++++++++ zig/src/tensorfold.zig | 1 + zig/tests/idol.zig | 242 ++++++++++++++++++ 14 files changed, 2046 insertions(+) create mode 100644 tools/zig/idol_case.py create mode 100644 zig/kernels/metal/idol.metal create mode 100644 zig/kernels/metal/idol_attention.metal create mode 100644 zig/src/families/idol/config.zig create mode 100644 zig/src/families/idol/host.zig create mode 100644 zig/src/families/idol/idol.zig create mode 100644 zig/src/families/idol/kernels.zig create mode 100644 zig/src/families/idol/model.zig create mode 100644 zig/src/families/idol/reference.zig create mode 100644 zig/src/families/idol/weights.zig create mode 100644 zig/tests/idol.zig diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py new file mode 100644 index 0000000000..83453a8e54 --- /dev/null +++ b/tools/zig/idol_case.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""Idol reference cases for tf-idol from the publisher's idol-sdk (torch): the denoiser's schedule and its blocks.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +USAGE = """\ + case WEIGHTS SDK OUT --features F [F] noise, each step's input and both velocities, the final sample + blocks WEIGHTS SDK CASE [--step S] every block's rows at step S, and each block in float32 on its input + exact WEIGHTS SDK CASE [--steps N] both velocities in float32 on the case's first N inputs (the ruler) +WEIGHTS is a speridlabs/idol snapshot; SDK an idol-sdk checkout at the pinned commit. F holds feature1 (DINOv3 rows +[1, 4101, 1280]) and feature2 (FLUX.2 VAE rows [1, 4101, 128]), in one file or two, as the SDK's ImageEncoder returns. +""" + +SDK_COMMIT = "4702a7f00850e34b6e83a2170cc123c3e9570744" +SDK_SHA256 = { + "idol/models/dit.py": "fbe725cc478b0acd245d9404ac8d9df327ebfbf79fe51c954d94b005c5bd6af4", + "idol/models/layers.py": "2dbf24107396194636c9a2e89e7acee1bcf91b14dba1b4cde0318cb85b3f482f", + "idol/sampler.py": "532987b5599be678361c2fcf566335bb3e1b0171696f985818d10b7460bfc8b0", +} +WEIGHTS_REVISION = "2731c2291fc65a797c5cb627a5e5526fc1490902" +DENOISER_ARGS = {"q_token_length": 8192, "in_channels": 16, "cam_channels": 5, "model_channels": 1024, + "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, "num_refiner_blocks": 2, + "num_heads": 16, "mlp_ratio": 4} + + +def sdk(path: str): + """The SDK's denoiser and sampler, imported once each file used has the pinned commit's sha256.""" + root = Path(path) + for name, want in SDK_SHA256.items(): + digest = hashlib.sha256((root / name).read_bytes()).hexdigest() + if digest != want: + sys.exit(f"{root / name} has sha256 {digest}; this tool runs only idol-sdk {SDK_COMMIT[:12]}'s ({want})") + sys.path.insert(0, str(root)) + from idol.models.dit import LatentSeqMMFlowModel + from idol.sampler import sample_flow + return LatentSeqMMFlowModel, sample_flow + + +def denoiser(weights: str, sdk_dir: str, device: str): + """The SDK's denoiser with the snapshot's weights, strictly loaded, as IdolPipeline builds it.""" + import torch + from safetensors.torch import load_file + + model_cls, _ = sdk(sdk_dir) + config = json.loads((Path(weights) / "config.json").read_text()) + args = config["denoiser"]["args"] + if config["denoiser"]["name"] != "LatentSeqMMFlowModel" or args != DENOISER_ARGS: + sys.exit(f"{weights}/config.json is not speridlabs/idol {WEIGHTS_REVISION[:12]}'s denoiser") + model = model_cls(**args) + model.load_state_dict(load_file(str(Path(weights) / "denoiser.safetensors")), strict=True) + return model.eval().requires_grad_(False).to(torch.device(device)) + + +def features(paths: list[str]): + """feature1 and feature2 [1, 4101, C] float32 from one or two safetensors files.""" + from safetensors.torch import load_file + + found = {} + for path in paths: + found.update(load_file(path)) + for key, width in (("feature1", 1280), ("feature2", 128)): + if key not in found or tuple(found[key].shape) != (1, 4101, width): + sys.exit(f"the features need {key} [1, 4101, {width}]") + return found["feature1"].float(), found["feature2"].float() + + +def condition(f1, f2, device: str): + """The sampler's cond and neg_cond (both streams zeroed), as IdolPipeline passes them.""" + from tensordict import TensorDict + + cond = TensorDict({"feature1": f1, "feature2": f2}, batch_size=1, device=device) + return cond, cond.apply(lambda x: x * 0) + + +def sample(latent, camera, device: str): + """The sampler's TensorDict for one latent [8192, 16] and camera [5].""" + from tensordict import TensorDict + + return TensorDict({"latent": latent[None].to(device), "camera": camera.view(1, 1, 5).to(device)}, batch_size=1, + device=device) + + +def flush(device: str) -> None: + import torch + + if device == "mps": + torch.mps.synchronize() + + +def case(args) -> int: + import torch + from safetensors.torch import save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + model = denoiser(args.weights, args.sdk, args.device) + _, sample_flow = sdk(args.sdk) + f1, f2 = features(args.features) + cond, neg = condition(f1, f2, args.device) + gen = torch.Generator("cpu").manual_seed(args.seed) + latent = torch.randn((8192, 16), generator=gen, dtype=torch.float32) + camera = torch.randn((5,), generator=gen, dtype=torch.float32) + calls = [] + forward = model.forward + + def recorded(x, t, c): + mark = time.perf_counter() + v = forward(x, t, c) + flush(args.device) + calls.append({"latent": x["latent"][0].float().cpu(), "camera": x["camera"][0, 0].float().cpu(), + "t": float(t[0]), "negative": bool(c["feature1"].abs().max() == 0), + "v": v["latent"][0].float().cpu(), "v_camera": v["camera"][0, 0].float().cpu(), + "s": time.perf_counter() - mark}) + return v + + model.forward = recorded + with torch.no_grad(): + final = sample_flow(model, sample(latent, camera, args.device), cond=cond, neg_cond=neg, steps=args.steps, + rescale_t=args.rescale_t, cfg_strength=args.cfg, verbose=False, batch_cfg=False) + model.forward = forward + pos, negs = calls[0::2], calls[1::2] + if args.cfg <= 0 or any(c["negative"] for c in pos) or not all(c["negative"] for c in negs): + sys.exit("the sampler's calls are not (conditional, negative) pairs") + stack = lambda rows, key: torch.stack([r[key] for r in rows]).contiguous() + tensors = {"feature1": f1[0].contiguous(), "feature2": f2[0].contiguous(), "latent": latent, "camera": camera, + "timesteps": torch.tensor([c["t"] for c in pos], dtype=torch.float32), + "inputs": stack(pos, "latent"), "inputs_camera": stack(pos, "camera"), + "velocity": stack(pos, "v"), "velocity_camera": stack(pos, "v_camera"), + "velocity_negative": stack(negs, "v"), "velocity_negative_camera": stack(negs, "v_camera"), + "sample": final["latent"][0].float().cpu().contiguous(), + "sample_camera": final["camera"][0, 0].float().cpu().contiguous()} + save_file(tensors, str(out / "case.safetensors")) + meta = {"seed": args.seed, "steps": args.steps, "cfg": args.cfg, "rescale_t": args.rescale_t, + "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, "weights": WEIGHTS_REVISION, + "torch_forward_s": [round(c["s"], 3) for c in calls]} + (out / "reference.json").write_text(json.dumps(meta, indent=1) + "\n") + times = sorted(meta["torch_forward_s"]) + print(json.dumps({k: v for k, v in meta.items() if k != "torch_forward_s"} | {"median_forward_s": times[len(times) // 2]})) + return 0 + + +def blocks(args) -> int: + """Each block's rows in the conditional forward at one step, and each block again in float32 on the same input.""" + import copy + + import torch + from safetensors.torch import load_file, save_file + + case_dir = Path(args.case_dir) + data = load_file(str(case_dir / "case.safetensors")) + model = denoiser(args.weights, args.sdk, args.device) + cond, _ = condition(data["feature1"][None], data["feature2"][None], args.device) + seen = {} + + def keep(name): + def hook(_module, _inputs, output): + seen[name] = output.detach().contiguous() + return hook + + stages = [("noise_refiner", model.noise_refiner), ("context_refiner", model.context_refiner), ("block", model.blocks)] + hooks = [m.register_forward_hook(keep(f"{name}.{i}")) for name, layers in stages for i, m in enumerate(layers)] + x = sample(data["inputs"][args.step], data["inputs_camera"][args.step], args.device) + t = data["timesteps"][args.step : args.step + 1].to(args.device) + with torch.no_grad(): + v = model(x, t, cond) + for h in hooks: + h.remove() + # the ruler: each joint block in float32 on the fp16 rows the SDK fed it, with an unrounded modulation + t_emb = model.t_embedder(t) + t_mod = model.adaLN_modulation(t_emb).float() + exact = {} + for i in range(1, len(model.blocks)): + block = copy.deepcopy(model.blocks[i]).float() + h = seen[f"block.{i - 1}"].float() + exact[f"block.{i}"] = block(h, mod=t_mod, rotary_emb=model.repo_layers[i](h))[0].cpu().contiguous() + del block + tensors = {k: x[0].cpu().contiguous() for k, x in seen.items()} + tensors["velocity"] = v["latent"][0].float().cpu().contiguous() + tensors["velocity_camera"] = v["camera"][0, 0].float().cpu().contiguous() + save_file(tensors, str(case_dir / f"blocks{args.step}.safetensors")) + save_file(exact, str(case_dir / f"blocks{args.step}_f32.safetensors")) + drift = (tensors["velocity"] - data["velocity"][args.step]).abs().max().item() + print(json.dumps({"step": args.step, "tensors": len(tensors), "velocity_max_abs_vs_case": drift})) + return 0 + + +def exact(args) -> int: + """Both branches in float32 throughout on the case's own inputs: the fp16 reference's distance from exact.""" + import torch + from safetensors.torch import load_file, save_file + + case_dir = Path(args.case_dir) + data = load_file(str(case_dir / "case.safetensors")) + model = denoiser(args.weights, args.sdk, args.device).float() + model.dtype = torch.float32 + cond, neg = condition(data["feature1"][None], data["feature2"][None], args.device) + keys = ("velocity", "velocity_camera", "velocity_negative", "velocity_negative_camera") + rows = {k: [] for k in keys} + with torch.no_grad(): + for s in range(min(args.steps, data["inputs"].shape[0])): + x = sample(data["inputs"][s], data["inputs_camera"][s], args.device) + t = data["timesteps"][s : s + 1].to(args.device) + for prefix, c in (("velocity", cond), ("velocity_negative", neg)): + v = model(x, t, c) + rows[prefix].append(v["latent"][0].float().cpu()) + rows[prefix + "_camera"].append(v["camera"][0, 0].float().cpu()) + save_file({k + "_f32": torch.stack(v).contiguous() for k, v in rows.items()}, str(case_dir / "exact.safetensors")) + print(json.dumps({"steps": len(rows["velocity"])})) + return 0 + + +def main(argv: list[str]) -> int: + ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = ap.add_subparsers(dest="cmd", required=True) + for name in ("case", "blocks", "exact"): + p = sub.add_parser(name) + p.add_argument("weights") + p.add_argument("sdk") + p.add_argument("out_dir" if name == "case" else "case_dir") + p.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + if name == "case": + p.add_argument("--features", nargs="+", required=True) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--steps", type=int, default=50) + p.add_argument("--cfg", type=float, default=3.0) + p.add_argument("--rescale-t", type=float, default=3.0) + elif name == "blocks": + p.add_argument("--step", type=int, default=0) + else: + p.add_argument("--steps", type=int, default=50) + args = ap.parse_args(argv) + return {"case": case, "blocks": blocks, "exact": exact}[args.cmd](args) + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 5bac04a9b0..f08c980d2f 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -317,6 +317,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-qwen-vision", .path = "zig/tests/qwen_vision.zig", .about = "The Qwen-VL tower on picture files: patches, grids and merged rows written for a float32 reference" }, .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" }, .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" }, + .{ .name = "tf-idol", .path = "zig/tests/idol.zig", .about = "Idol's denoiser on a reference case: each step's two branches compared, then the guided schedule timed" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/idol.metal b/zig/kernels/metal/idol.metal new file mode 100644 index 0000000000..1a8048d1b2 --- /dev/null +++ b/zig/kernels/metal/idol.metal @@ -0,0 +1,329 @@ +// Idol's denoiser steps after gemm_nax.metal: fp16 rows rounded where idol-sdk's torch code rounds them. + +inline float id_h(float x) { return float(half(x)); } + +// One threadgroup's sum of every thread's `v` (at most 1,024 threads), returned to all of them. +inline float id_sum(float v, threadgroup float* part, uint n, uint lane, uint sg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float s = lane < (n + 31) / 32 ? part[lane] : 0.0f; + s = simd_sum(s); + if (lane == 0) part[32] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float out = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return out; +} + +// torch's tanh GELU of an fp16 value, computed in fp32. +inline float id_gelu(float x) { + const float inner = 0.7978845608028654f * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + metal::precise::tanh(inner)); +} + +// ---- y [M, N] = fp16(x w^T + b) (EPI 1: then GELU, rounded again): tfp::gemm_nax's block, fp32 sums ---- +template +inline void id_linear(const device half* A, const device half* B, const device half* bias, device half* D, + const device int* P, uint3 tg, uint sg, uint lane) { + constexpr int BM = 64, BN = 128, BK = 256, WN = 4; + const int tx = int(tg.x), ty = int(tg.y); + if (tx >= P[6] || ty >= P[7]) return; + const int lda = P[3], ldb = P[4], ldd = P[5]; + const int row = ty * BM + 32 * int(sg / WN), col = tx * BN + 32 * int(sg % WN); + A += long(row) * lda; + B += long(col) * ldb; + D += long(row) * ldd + col; + const int rows = AM ? 32 : min(32, P[0] - row), cols = AN ? 32 : min(32, P[1] - col); + const bool full = (AM || rows == 32) && (AN || cols == 32); + const short2 home = tfp::frag_home(ushort(lane)); + tfp::frag acc[2][2]; + if (full) { + tfp::gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } else { + tfp::gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } + TF_UNROLL + for (short j = 0; j < 2; j++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + const int c = col + 16 * j + home.x + TF_COL(e); + const float b = c < P[1] ? float(bias[c]) : 0.0f; + TF_UNROLL + for (short i = 0; i < 2; i++) { + const float v = id_h(acc[i][j][e] + b); + acc[i][j][e] = EPI == 1 ? id_gelu(v) : v; + } + } + } + if (full) { + tfp::gemm_store(acc, D, ldd, rows, cols, home); + } else if (rows > 0 && cols > 0) { + tfp::gemm_store(acc, D, ldd, rows, cols, home); + } +} + +#define ID_LINEAR(NAME, EPI, AM, AN) \ + [[max_total_threads_per_threadgroup(256)]] [[kernel]] void NAME( \ + const device half* A [[buffer(0)]], const device half* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device half* D [[buffer(3)]], const device half* bias [[buffer(4)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + id_linear(A, B, bias, D, P, tg, sg, lane); \ + } +ID_LINEAR(id_linear_00, 0, false, false) +ID_LINEAR(id_linear_01, 0, false, true) +ID_LINEAR(id_linear_10, 0, true, false) +ID_LINEAR(id_linear_11, 0, true, true) +ID_LINEAR(id_linear_gelu_00, 1, false, false) +ID_LINEAR(id_linear_gelu_01, 1, false, true) +ID_LINEAR(id_linear_gelu_10, 1, true, false) +ID_LINEAR(id_linear_gelu_11, 1, true, true) + +struct IdGemm32 { + uint M, N, K, ldy; +}; + +// y [M, N] (rows ldy apart) = x [M, K] w^T, fp32 FMA on 8x8 simdgroup matrices; K % 32, N % 64 (or N < 64, % 8) = 0. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void id_gemm32( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], device float* y [[buffer(2)]], + constant IdGemm32& p [[buffer(3)]], uint2 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + constexpr uint BK = 32, LD = BK + 4; + threadgroup float ta[64 * LD]; + threadgroup float tb[64 * LD]; + threadgroup float stage[4][64]; + const uint row0 = tg.y * 64, col0 = tg.x * 64; + const uint r0 = (sg / 2) * 32, c0 = (sg % 2) * 32; + const uint cols = min(64u, p.N - col0); + simdgroup_float8x8 acc[4][4]; + for (uint i = 0; i < 4; i++) + for (uint j = 0; j < 4; j++) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (uint k = 0; k < p.K; k += BK) { + // 64 x 32 of x and of w: 2,048 values each, 16 a thread as float4s; rows past M or N read as zeros + for (uint e = t; e < 64 * BK / 4; e += 128) { + const uint r = e / (BK / 4), c = (e % (BK / 4)) * 4; + const float4 va = row0 + r < p.M ? *(const device float4*)(x + ulong(row0 + r) * p.K + k + c) : float4(0.0f); + const float4 vb = r < cols ? *(const device float4*)(w + ulong(col0 + r) * p.K + k + c) : float4(0.0f); + *(threadgroup float4*)(ta + r * LD + c) = va; + *(threadgroup float4*)(tb + r * LD + c) = vb; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 a[4], b[4]; + for (uint i = 0; i < 4; i++) simdgroup_load(a[i], ta + (r0 + 8 * i) * LD + kk, LD); + for (uint j = 0; j < 4; j++) simdgroup_load(b[j], tb + (c0 + 8 * j) * LD + kk, LD, ulong2(0, 0), true); + for (uint i = 0; i < 4; i++) + for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i = 0; i < 4; i++) { + for (uint j = 0; j < 4; j++) { + const uint r = row0 + r0 + 8 * i, c = c0 + 8 * j; + if (r >= p.M || c >= cols) continue; + if (r + 8 <= p.M) { + simdgroup_store(acc[i][j], y + ulong(r) * p.ldy + col0 + c, p.ldy); + } else { + simdgroup_store(acc[i][j], stage[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = lane; e < 64; e += 32) + if (r + e / 8 < p.M) y[ulong(r + e / 8) * p.ldy + col0 + c + e % 8] = stage[sg][e]; + simdgroup_barrier(mem_flags::mem_threadgroup); + } + } + } +} + +// c [M, N] fp32 = a [M, K] b^T over fp16 operands, tfp::gemm_splitk_nax's one partition (fp32 sums). +#define ID_SPLIT(NAME, AM) \ + [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ + const device half* A [[buffer(0)]], const device half* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device float* C [[buffer(3)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::gemm_splitk_nax(A, B, C, P, tg, sg, lane); \ + } +ID_SPLIT(id_split_0, false) +ID_SPLIT(id_split_1, true) + +struct IdNorm { + uint width; + float eps; + uint mode; // 0 plain, 1 affine (fp32 w, b), 2 modulated (fp16 1 + scale, shift) +}; + +// LayerNorm32 of an fp16 row in fp32, rounded to fp16; then affine in fp32, or modulated in fp16. A threadgroup a row. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_norm( + const device half* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device half* s1 [[buffer(3)]], const device half* sh [[buffer(4)]], device half* y [[buffer(5)]], + constant IdNorm& p [[buffer(6)]], uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint n [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device half* xr = x + ulong(row) * p.width; + device half* yr = y + ulong(row) * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = id_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(id_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + float v = (float(xr[i]) - mean) * inv; + if (p.mode == 1) v = v * w[i] + b[i]; + v = id_h(v); + if (p.mode == 2) v = id_h(id_h(v * float(s1[i])) + float(sh[i])); + yr[i] = half(v); + } +} + +// Affine LayerNorm32 of an fp16 row as fp16 [hi | hi | lo] of 3 * width: times [w_hi | w_lo | w_hi], near-fp32 sums. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_norm_split( + const device half* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device half* y [[buffer(3)]], constant IdNorm& p [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device half* xr = x + ulong(row) * p.width; + device half* yr = y + ulong(row) * 3 * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = id_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(id_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + const float v = (float(xr[i]) - mean) * inv * w[i] + b[i]; + const half hi = half(v); + yr[i] = hi; + yr[p.width + i] = hi; + yr[2 * p.width + i] = half(v - float(hi)); + } +} + +// LayerNorm of an fp16 row in fp32 kept in fp32: affine (w, b) when given, else times (1 + scale) plus shift (mode 2). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_norm32( + const device half* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device float* y [[buffer(3)]], constant IdNorm& p [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device half* xr = x + ulong(row) * p.width; + device float* yr = y + ulong(row) * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = id_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(id_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + const float v = (float(xr[i]) - mean) * inv; + yr[i] = p.mode == 2 ? v * (1.0f + w[i]) + b[i] : v * w[i] + b[i]; + } +} + +struct IdRepo { + uint rows; + uint stride; // fp32 values from one row's gate sums to the next; content sums follow gate's 128 + float inv; // the maps' power-of-two scale undone, exactly +}; + +// RePo: d = final_map(silu(gate) * content), 3 per head; angles clamp_mul(d, freqs) * pi as cos, sin [rows, 16, 32]. +[[max_total_threads_per_threadgroup(64)]] [[kernel]] void id_repo( + const device float* gc [[buffer(0)]], const device float* wf [[buffer(1)]], const device float* f [[buffer(2)]], + device float* cs [[buffer(3)]], device float* sn [[buffer(4)]], constant IdRepo& p [[buffer(5)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]]) { + threadgroup float u[128]; + threadgroup float d[48]; + const device float* g = gc + ulong(row) * p.stride; + for (uint i = t; i < 128; i += 64) { + const float a = g[i] * p.inv; + u[i] = (a / (1.0f + metal::precise::exp(-a))) * (g[128 + i] * p.inv); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < 48) { + float s = 0.0f; + for (uint k = 0; k < 128; k++) s = fma(u[k], wf[t * 128 + k], s); + d[t] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // freqs: 10 x, 10 y, 12 z; each pair of the 32 takes its axis's offset + for (uint i = t; i < 512; i += 64) { + const uint h = i / 32, j = i % 32; + const uint axis = j < 10 ? 0 : (j < 20 ? 1 : 2); + const float x = d[h * 3 + axis], fr = f[j]; + const float ft = metal::precise::tanh(fr); + const float ang = (x * ft + x * (fr - ft)) * 3.14159265358979323846f; + cs[ulong(row) * 512 + i] = metal::precise::cos(ang); + sn[ulong(row) * 512 + i] = metal::precise::sin(ang); + } +} + +struct IdHeads { + uint rows; + float scale; // sqrt(head width): MultiHeadRMSNorm's +}; + +// q, k of qkv rows: pairs turned by the row and head's cos, sin; F.normalize * gamma * scale; out [16, rows, 64]. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_qk( + const device half* qkv [[buffer(0)]], const device float* gq [[buffer(1)]], const device float* gk [[buffer(2)]], + const device float* cs [[buffer(3)]], const device float* sn [[buffer(4)]], device half* Q [[buffer(5)]], + device half* K [[buffer(6)]], device half* V [[buffer(7)]], constant IdHeads& p [[buffer(8)]], + uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + const uint r = tg.x, h = tg.y * 8 + sg; + if (h >= 16) return; + const ulong from = ulong(r) * 3072 + h * 64 + lane * 2; + const ulong at = ulong(r) * 512 + h * 32 + lane; + const float c = cs[at], s = sn[at]; + const float q0 = float(qkv[from]), q1 = float(qkv[from + 1]); + const float k0 = float(qkv[from + 1024]), k1 = float(qkv[from + 1025]); + const float rq0 = id_h(q0 * c - q1 * s), rq1 = id_h(q0 * s + q1 * c); + const float rk0 = id_h(k0 * c - k1 * s), rk1 = id_h(k0 * s + k1 * c); + const float qn = max(metal::precise::sqrt(simd_sum(rq0 * rq0 + rq1 * rq1)), 1e-12f); + const float kn = max(metal::precise::sqrt(simd_sum(rk0 * rk0 + rk1 * rk1)), 1e-12f); + const uint g = h * 64 + lane * 2; + const ulong to = (ulong(h) * p.rows + r) * 64 + lane * 2; + Q[to] = half(rq0 / qn * gq[g] * p.scale); + Q[to + 1] = half(rq1 / qn * gq[g + 1] * p.scale); + K[to] = half(rk0 / kn * gk[g] * p.scale); + K[to + 1] = half(rk1 / kn * gk[g + 1] * p.scale); + V[to] = qkv[from + 2048]; + V[to + 1] = qkv[from + 2049]; +} + +// x = fp16(x + fp16(g * y)) by column; g is a modulation gate, or ones for an ungated residual. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_gate_add( + device half* x [[buffer(0)]], const device half* y [[buffer(1)]], const device half* g [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + x[i] = half(float(x[i]) + id_h(float(g[at.x]) * float(y[i]))); +} + +// x = fp16(fp16(z w^T + b) + pe): the latent's input projection (16 inputs, fp32) and its fp16 position embedding. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_embed( + const device float* z [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device half* pe [[buffer(3)]], device half* x [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + const uint c = at.x, r = at.y; + float s = 0.0f; + for (uint k = 0; k < 16; k++) s = fma(z[r * 16 + k], w[c * 16 + k], s); + const ulong i = ulong(r) * 1024 + c; + x[i] = half(float(half(s + b[c])) + float(pe[i])); +} + +// y = fp16(a + ba + c + bc) by column: the two condition projections' fp32 sums and biases (each torch Linear). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_cond( + const device float* a [[buffer(0)]], const device float* ba [[buffer(1)]], const device float* c [[buffer(2)]], + const device float* bc [[buffer(3)]], device half* y [[buffer(4)]], constant uint& width [[buffer(5)]], + uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + y[i] = half((a[i] + ba[at.x]) + (c[i] + bc[at.x])); +} + +// y = x + b by column in fp32 (an fp32 Linear's bias), in place. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void id_bias32( + device float* y [[buffer(0)]], const device float* b [[buffer(1)]], constant uint& width [[buffer(2)]], + uint2 at [[thread_position_in_grid]]) { + y[ulong(at.y) * width + at.x] += b[at.x]; +} diff --git a/zig/kernels/metal/idol_attention.metal b/zig/kernels/metal/idol_attention.metal new file mode 100644 index 0000000000..f676630fc5 --- /dev/null +++ b/zig/kernels/metal/idol_attention.metal @@ -0,0 +1,149 @@ +// Idol's attention: tfp::attention_nax's online softmax at head width 64, every query row to every key row. + +constant constexpr float kIdMasked = -3.402823466e+38f; +constant constexpr int kIdAnyCols = 1 << 30; + +// Fold one fragment's rows into r[2] (rows home.y, home.y + 8): pairs, then lanes 1 and 8 apart, then into r. +template +inline void id_fold(thread const tfp::frag& f, thread float (&r)[2]) { + TF_UNROLL + for (short h = 0; h < 2; h++) { + const short b = 4 * h; + float t = MAX ? max(max(f[b], f[b + 1]), max(f[b + 2], f[b + 3])) : (f[b] + f[b + 1]) + (f[b + 2] + f[b + 3]); + const float u = simd_shuffle_xor(t, ushort(1)); + t = MAX ? max(t, u) : t + u; + const float w = simd_shuffle_xor(t, ushort(8)); + t = MAX ? max(t, w) : t + w; + r[h] = MAX ? max(r[h], t) : r[h] + t; + } +} + +// One 64-query block of one head over 32-key blocks; P and F as Kernels.attention sets them (kernels.zig). +template +inline void id_attention(const device T* Q, const device T* K, const device T* V, device T* O, const device int* P, + const device float* F, uint3 tg, uint sg, uint lane) { + constexpr int BQ = 64, BK = 32, D = 64; + const int h = int(tg.y), lq = P[1], lk = P[2]; + Q += (long(h) * lq + long(tg.x) * BQ) * D; + K += long(h) * lk * D; + V += long(h) * lk * D; + const int os = P[0] * D; + O += long(tg.x) * BQ * os + h * D; + const float scale2 = F[0] * 1.44269504089f; + const short tm = 16 * short(sg); + Q += tm * D; + const short2 home = tfp::frag_home(ushort(lane)); + + tfp::frag acc[4]; + TF_UNROLL + for (short i = 0; i < 4; i++) acc[i] = tfp::frag(0); + float top[2] = {kIdMasked, kIdMasked}, total[2] = {0.0f, 0.0f}; + const bool last_q = int(tg.x) == P[4]; + const short q_rows = short(P[6] - tm), k_rows = short(P[7]); + + for (int kb = 0; kb < P[3]; kb++) { + const bool last_k = kb == P[5]; + tfp::frag s[2] = {tfp::frag(0), tfp::frag(0)}; + TF_UNROLL + for (short d = 0; d < 4; d++) { + tfp::frag q, k0, k1; + if (!AQ && last_q) { + tfp::frag_get_in(q, Q, D, 0, 16 * d, home, q_rows, kIdAnyCols); + } else { + tfp::frag_get(q, Q, D, 0, 16 * d, home); + } + if (!AK && last_k) { + tfp::frag_get_t_in(k0, K, D, 0, 16 * d, home, k_rows, kIdAnyCols); + tfp::frag_get_t_in(k1, K, D, 16, 16 * d, home, k_rows, kIdAnyCols); + } else { + tfp::frag_get_t(k0, K, D, 0, 16 * d, home); + tfp::frag_get_t(k1, K, D, 16, 16 * d, home); + } + tfp::mma_16x32(s[0], s[1], q, k0, k1); + } + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] *= scale2; + } + if (!AK && last_k) { + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kIdMasked; + } + } + float top_new[2] = {top[0], top[1]}; + id_fold(s[0], top_new); + id_fold(s[1], top_new); + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = fast::exp2(s[f][e] - top_new[e >> 2]); + } + float scale_old[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) { + scale_old[r] = fast::exp2(top[r] - top_new[r]); + top[r] = top_new[r]; + total[r] = total[r] * scale_old[r]; + } + id_fold(s[0], total); + id_fold(s[1], total); + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * scale_old[e >> 2]; + } + simdgroup_barrier(mem_flags::mem_none); + TF_UNROLL + for (short d = 0; d < 4; d += 2) { + TF_UNROLL + for (short k = 0; k < 2; k++) { + tfp::frag v0, v1; + if (!AK && last_k) { + tfp::frag_get_in(v0, V, D, 16 * k, 16 * d, home, k_rows, kIdAnyCols); + tfp::frag_get_in(v1, V, D, 16 * k, 16 * d + 16, home, k_rows, kIdAnyCols); + } else { + tfp::frag_get(v0, V, D, 16 * k, 16 * d, home); + tfp::frag_get(v1, V, D, 16 * k, 16 * d + 16, home); + } + tfp::mma_16x32(acc[d], acc[d + 1], s[k], v0, v1); + } + } + K += BK * D; + V += BK * D; + } + + threadgroup_barrier(mem_flags::mem_none); + float inv[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) inv[r] = 1.f / total[r]; + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * inv[e >> 2]; + } + O += tm * os; + if (!AQ && last_q) { + if (q_rows <= 0) return; + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put_in(acc[i], O, os, 0, 16 * i, home, q_rows, kIdAnyCols); + } else { + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put(acc[i], O, os, 0, 16 * i, home); + } +} + +#define ID_ATTENTION(NAME, T, AQ, AK) \ + [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ + const device T* Q [[buffer(0)]], const device T* K [[buffer(1)]], const device T* V [[buffer(2)]], \ + const device int32_t* P [[buffer(3)]], const device float* F [[buffer(4)]], device T* O [[buffer(5)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + id_attention(Q, K, V, O, P, F, tg, sg, lane); \ + } +ID_ATTENTION(id_attention_00, half, false, false) +ID_ATTENTION(id_attention_01, half, false, true) +ID_ATTENTION(id_attention_10, half, true, false) +ID_ATTENTION(id_attention_11, half, true, true) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index f5d5861ad8..ef3221a88b 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -133,6 +133,8 @@ pub const core_qwen_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile( pub const qwen_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("qwen_image.metal"); /// MOSS-Image's transformer: the prefill attention template, then its own steps (projections: core_qwen_vision). pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal"); +/// Idol's denoiser: the prefill GEMM template, then its own steps and head-64 attention. +pub const idol = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol.metal") ++ @embedFile("idol_attention.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/families/idol/config.zig b/zig/src/families/idol/config.zig new file mode 100644 index 0000000000..d399755886 --- /dev/null +++ b/zig/src/families/idol/config.zig @@ -0,0 +1,71 @@ +//! Idol's denoiser geometry, and the check that a snapshot's config.json names that denoiser. +const std = @import("std"); + +pub const hidden = 1024; +pub const heads = 16; +pub const head_dim = hidden / heads; +pub const ffn = 4 * hidden; +pub const layers = 24; +/// The noise and condition refiners' blocks, each. +pub const refiners = 2; +pub const latent_rows = 8192; +pub const latent_channels = 16; +pub const camera_channels = 5; +/// DINOv3 ViT-H+/16 rows at 1024 x 1024: 64 x 64 patches, the class token and four registers. +pub const cond_rows = 4101; +pub const cond_channels = 1280; +pub const cond2_channels = 128; +/// Latent, condition and camera rows through the joint blocks. +pub const joint_rows = latent_rows + cond_rows + 1; +/// RePo's hidden width (model_channels / 8) and its frequencies per head: 10, 10 and 12 for x, y and z. +pub const repo_hidden = hidden / 8; +pub const repo_pairs = head_dim / 2; +pub const time_dim = 256; +/// Every transformer block's LayerNorm32, fixed in idol-sdk's layers.py, not a config setting. +pub const block_eps: f32 = 1e-6; +/// RePo's LayerNorm32 and the output's F.layer_norm take torch's default. +pub const default_eps: f32 = 1e-5; +/// Modulated blocks: the noise refiner's, then the joint blocks. +pub const modulated = refiners + layers; + +/// Accepts the snapshot config whose denoiser this family runs (LatentSeqMMFlowModel at its released size). +pub fn check(gpa: std.mem.Allocator, text: []const u8) !void { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadConfig; + const den = doc.value.object.get("denoiser") orelse return error.BadConfig; + if (den != .object) return error.BadConfig; + const name = den.object.get("name") orelse return error.BadConfig; + if (name != .string or !std.mem.eql(u8, name.string, "LatentSeqMMFlowModel")) return error.NotIdol; + const args = den.object.get("args") orelse return error.BadConfig; + if (args != .object) return error.BadConfig; + const want = .{ .{ "q_token_length", latent_rows }, .{ "in_channels", latent_channels }, .{ "cam_channels", camera_channels }, .{ "model_channels", hidden }, .{ "cond_channels", cond_channels }, .{ "cond2_channels", cond2_channels }, .{ "num_blocks", layers }, .{ "num_refiner_blocks", refiners }, .{ "num_heads", heads }, .{ "mlp_ratio", 4 } }; + inline for (want) |kv| { + const v = args.object.get(kv[0]) orelse return error.BadConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.warn("idol: denoiser {s} is not {d}; this family runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } +} + +const sample = + \\{"denoiser": {"name": "LatentSeqMMFlowModel", "args": {"q_token_length": 8192, "in_channels": 16, + \\ "cam_channels": 5, "model_channels": 1024, "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, + \\ "num_refiner_blocks": 2, "num_heads": 16, "mlp_ratio": 4}}, "image_size": 1024, "max_voxel_level": 8} +; + +test "the released denoiser config passes; another depth, width or class does not" { + const a = std.testing.allocator; + try check(a, sample); + const deeper = try std.mem.replaceOwned(u8, a, sample, "\"num_blocks\": 24", "\"num_blocks\": 12"); + defer a.free(deeper); + try std.testing.expectError(error.UnsupportedGeometry, check(a, deeper)); + const missing = try std.mem.replaceOwned(u8, a, sample, "\"mlp_ratio\": 4", "\"mlp\": 4"); + defer a.free(missing); + try std.testing.expectError(error.BadConfig, check(a, missing)); + const other = try std.mem.replaceOwned(u8, a, sample, "LatentSeqMMFlowModel", "LatentFlowModel"); + defer a.free(other); + try std.testing.expectError(error.NotIdol, check(a, other)); + try std.testing.expectError(error.BadConfig, check(a, "{\"decoder\": {}}")); +} diff --git a/zig/src/families/idol/host.zig b/zig/src/families/idol/host.zig new file mode 100644 index 0000000000..33dd008649 --- /dev/null +++ b/zig/src/families/idol/host.zig @@ -0,0 +1,147 @@ +//! The host's part of an Idol denoiser pass: the schedule, the time modulation, the camera token and its velocity. +const std = @import("std"); +const cfg = @import("config.zig"); +const Host = @import("weights.zig").Host; + +const H = cfg.hidden; + +pub fn half(x: f32) f32 { + return @as(f16, @floatCast(x)); +} + +/// y[o] = sum over k of x[k] w[o, k] + b[o] in fp32, w stored (outputs, inputs). +pub fn project(w: []const f32, b: []const f32, x: []const f32, y: []f32) void { + std.debug.assert(w.len == x.len * y.len and b.len == y.len); + const V = @Vector(8, f32); + for (y, 0..) |*o, row| { + const ws = w[row * x.len ..][0..x.len]; + var acc: V = @splat(0); + var k: usize = 0; + while (k + 8 <= x.len) : (k += 8) acc += @as(V, ws[k..][0..8].*) * @as(V, x[k..][0..8].*); + var sum = @reduce(.Add, acc); + while (k < x.len) : (k += 1) sum += ws[k] * x[k]; + o.* = sum + b[row]; + } +} + +fn silu(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} + +fn gelu(x: f32) f32 { + return 0.5 * x * (1.0 + std.math.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))); +} + +/// The sampler's times: t = s u / (1 + (s - 1) u) for u from 1 to 0 in `steps` (steps + 1 values), in float64. +pub fn schedule(out: []f64, rescale: f64) void { + const steps = out.len - 1; + for (out, 0..) |*t, i| { + const u = 1.0 - @as(f64, @floatFromInt(i)) / @as(f64, @floatFromInt(steps)); + t.* = rescale * u / (1.0 + (rescale - 1.0) * u); + } +} + +/// The modulated blocks' six fp16 vectors each (shift, 1 + scale, gate for attention, then for the MLP). +pub const vectors = 6; +pub const Mods = [cfg.modulated][vectors][H]f16; + +/// t_embedder's output [1024] at model time `t` (1000 times the flow time), in fp32. +pub fn timeEmbedding(w: Host, t: f32, out: *[H]f32) void { + const halfdim = cfg.time_dim / 2; + var freq: [cfg.time_dim]f32 = undefined; + const lg: f32 = @floatCast(-@log(@as(f64, 10000))); + for (0..halfdim) |i| { + const a = t * @exp(lg * @as(f32, @floatFromInt(i)) / @as(f32, halfdim)); + freq[i] = @cos(a); + freq[halfdim + i] = @sin(a); + } + var mid: [H]f32 = undefined; + project(w.t1, w.t1_b, &freq, &mid); + for (&mid) |*v| v.* = silu(v.*); + project(w.t2, w.t2_b, &mid, out); +} + +/// Every modulated block's vectors and the output's fp32 shift and scale at model time `t`. +pub fn modulation(w: Host, t: f32, mods: *Mods, head: *[2][H]f32) void { + var emb: [H]f32 = undefined; + timeEmbedding(w, t, &emb); + var act: [H]f32 = undefined; + for (&act, emb) |*o, v| o.* = silu(v); + var tmod: [6 * H]f32 = undefined; + project(w.ada, w.ada_b, &act, &tmod); + for (&tmod) |*v| v.* = half(v.*); + for (mods, w.shifts) |*m, shift| { + for (0..6 * H) |i| { + const v = half(tmod[i] + half(shift[i])); + const part = i / H; + // the scales (parts 1 and 4) enter as fp16(1 + scale) + m[part][i % H] = @floatCast(if (part == 1 or part == 4) half(1.0 + v) else v); + } + } + for (0..2) |p| for (0..H) |i| { + head[p][i] = w.head[p * H + i] + emb[i]; + }; +} + +/// The camera token: cam_refiner's MLP of the five camera values, in fp32, then fp16. +pub fn cameraToken(w: Host, camera: *const [cfg.camera_channels]f32, out: []f16) void { + var mid: [H]f32 = undefined; + project(w.cam1, w.cam1_b, camera, &mid); + for (&mid) |*v| v.* = gelu(v.*); + var tok: [H]f32 = undefined; + project(w.cam2, w.cam2_b, &mid, &tok); + for (out, tok) |*o, v| o.* = @floatCast(v); +} + +/// The camera velocity from the last joint row: F.layer_norm (eps 1e-5), modulated by `head`, then cam_out_layer. +pub fn cameraVelocity(w: Host, row: []const f16, head: *const [2][H]f32, out: *[cfg.camera_channels]f32) void { + var x: [H]f32 = undefined; + var mean: f32 = 0; + for (&x, row) |*o, v| { + o.* = v; + mean += o.*; + } + mean /= H; + var q: f32 = 0; + for (x) |v| q += (v - mean) * (v - mean); + const inv = 1.0 / @sqrt(q / H + cfg.default_eps); + // shift is the table's first half, scale its second + for (&x, 0..) |*v, i| v.* = ((v.* - mean) * inv) * (1.0 + head[1][i]) + head[0][i]; + project(w.cam_out, w.cam_out_b, &x, out); +} + +/// The latents' fixed position embedding [8192, 1024] in fp16: sin then cos of each axis at 2^k frequencies. +pub fn positions(pos: []const f32, out: []f16) void { + std.debug.assert(pos.len == cfg.latent_rows * 3 and out.len == cfg.latent_rows * H); + const fd = H / 3 / 2; + const max_res = 16; + var freqs: [fd]f32 = undefined; + const rest = fd - max_res; + for (&freqs, 0..) |*f, i| { + const e: f32 = if (i < max_res) @floatFromInt(i) else @as(f32, @floatFromInt(i - max_res)) / @as(f32, rest) * max_res; + f.* = std.math.pow(f32, 2, e); + } + for (0..cfg.latent_rows) |r| { + const o = out[r * H ..][0..H]; + // outer(x.reshape(-1), freqs) then cat(sin, cos) per value: each axis's fd sines, then its fd cosines + for (0..3) |a| for (0..fd) |j| { + const ang = pos[r * 3 + a] * freqs[j] * 2 * std.math.pi; + o[a * 2 * fd + j] = @floatCast(@sin(ang)); + o[a * 2 * fd + fd + j] = @floatCast(@cos(ang)); + }; + for (o[6 * fd ..]) |*v| v.* = 0; + } +} + +test "the schedule runs from 1 to 0, denser at high noise for rescale 3" { + var t: [51]f64 = undefined; + schedule(&t, 3.0); + try std.testing.expectEqual(@as(f64, 1), t[0]); + try std.testing.expectEqual(@as(f64, 0), t[50]); + try std.testing.expectApproxEqAbs(@as(f64, 0.75), t[25], 1e-12); +} + +test "fp16 rounding is to nearest, ties to even" { + try std.testing.expectEqual(@as(f32, 1.0), half(1.0 + 1.0 / 4096.0)); + try std.testing.expectEqual(@as(f32, 1.0 + 1.0 / 512.0), half(1.0 + 3.0 / 2048.0)); +} diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig new file mode 100644 index 0000000000..a2aff60944 --- /dev/null +++ b/zig/src/families/idol/idol.zig @@ -0,0 +1,19 @@ +//! Idol image to 3D: its denoiser on Metal (a denoiser beside the lane engine, not a token family). +pub const config = @import("config.zig"); +pub const host = @import("host.zig"); +pub const weights = @import("weights.zig"); +pub const kernels = @import("kernels.zig"); +pub const model = @import("model.zig"); +pub const reference = @import("reference.zig"); +pub const Model = model.Model; +pub const Condition = model.Condition; +pub const Timing = model.Timing; + +test { + _ = config; + _ = host; + _ = weights; + _ = kernels; + _ = model; + _ = reference; +} diff --git a/zig/src/families/idol/kernels.zig b/zig/src/families/idol/kernels.zig new file mode 100644 index 0000000000..2efe840cdc --- /dev/null +++ b/zig/src/families/idol/kernels.zig @@ -0,0 +1,209 @@ +//! Idol's Metal pipelines: its fp16 projections and steps, fp32 projections and head-64 attention. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const cfg = @import("config.zig"); +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// Each pipeline by its function. +pub const names = [_][]const u8{ + "id_linear_00", + "id_linear_01", + "id_linear_10", + "id_linear_11", + "id_linear_gelu_00", + "id_linear_gelu_01", + "id_linear_gelu_10", + "id_linear_gelu_11", + "id_norm", + "id_norm32", + "id_repo", + "id_qk", + "id_gate_add", + "id_embed", + "id_cond", + "id_bias32", + "id_gemm32", + "id_split_0", + "id_split_1", + "id_norm_split", + "id_attention_00", + "id_attention_01", + "id_attention_10", + "id_attention_11", +}; +const linear_first = 0; +const attention_first = blk: { + for (names, 0..) |n, i| if (std.mem.eql(u8, n, "id_attention_00")) break :blk i; + unreachable; +}; + +pub const Norm = enum(u32) { plain = 0, affine = 1, modulated = 2 }; + +pub const Kernels = struct { + own: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Kernels { + var k: Kernels = undefined; + const own = try frags.source(device, gpa, sources.idol); + defer gpa.free(own); + k.own = try mtl.Library.fromSource(device, own, mtl.CompileOptions.mlx()); + errdefer k.own.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.own, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.own.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } + + /// y [M, N] (rows `ldd` apart) = x [M, K] w^T in fp32 FMA, w as stored [N, K]; K % 32, N % 64 (or N < 64, % 8) = 0. + pub fn linear32(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: usize, N: usize, K: usize, ldd: usize) void { + std.debug.assert(K % 32 == 0 and (N % 64 == 0 or (N < 64 and N % 8 == 0))); + const params = [4]u32{ @intCast(M), @intCast(N), @intCast(K), @intCast(ldd) }; + bind(e, k.get("id_gemm32"), &.{ x, w, y }, params); + e.dispatchGroups(mtl.Size.of((N + 63) / 64, (M + 63) / 64, 1), mtl.Size.of(128, 1, 1)); + } + + /// c [M, 256] fp32 = x3 [M, 3072] w3^T over fp16: a split LayerNorm's rows and a RePo layer's split maps. + pub fn split(k: *const Kernels, e: mtl.ComputeEncoder, x3: Ref, w3: Ref, c: Ref, M: usize) void { + const N = 2 * cfg.repo_hidden; + const K = 3 * cfg.hidden; + const tn: usize = N / 64; + const tm = (M + 63) / 64; + const swz: u6 = if (tm <= 3) 0 else 1; + e.setPipeline(if (M % 64 == 0) k.get("id_split_1") else k.get("id_split_0")); + e.setBuffer(x3.buf, x3.off, 0); + e.setBuffer(w3.buf, w3.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 1, M * N, K, swz, 0, 0 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(c.buf, c.off, 3); + const tiles = (tn << swz) * ((tm + (@as(usize, 1) << swz) - 1) >> swz); + e.dispatchThreads(mtl.Size.of(tiles * 32, 2, 2), mtl.Size.of(32, 2, 2)); + } + + /// x3 rows [hi | hi | lo] of the affine LayerNorm32 of x's rows, for `split`. + pub fn normSplit(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, x3: Ref, rows: usize, eps: f32) void { + const params = extern struct { width: u32, eps: f32, mode: u32 }{ .width = cfg.hidden, .eps = eps, .mode = 1 }; + bind(e, k.get("id_norm_split"), &.{ x, w, b, x3 }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [M, N] = fp16(x w^T + b) in fp16 with fp32 sums, then GELU (rounded again) when `gelu`. + pub fn linear(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, M: usize, N: usize, K: usize, gelu: bool) void { + const tn = (N + 127) / 128; + const tm = (M + 63) / 64; + const form = @as(usize, @intFromBool(M % 64 == 0)) * 2 + @intFromBool(N % 128 == 0); + e.setPipeline(k.p[linear_first + @as(usize, @intFromBool(gelu)) * 4 + form]); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 0, K / 256 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(y.buf, y.off, 3); + e.setBuffer(b.buf, b.off, 4); + e.dispatchGroups(mtl.Size.of(tn, tm, 1), mtl.Size.of(256, 1, 1)); + } + + /// o rows [Lq, heads * 64] = softmax(q k^T / 8) v for q [heads, Lq, 64] and k, v [heads, Lk, 64]. + pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, lq: usize, lk: usize) void { + const form = @as(usize, @intFromBool(lq % 64 == 0)) * 2 + @intFromBool(lk % 32 == 0); + const nq = (lq + 63) / 64; + const nk = (lk + 31) / 32; + e.setPipeline(k.p[attention_first + form]); + for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + // heads, q rows, k rows, k blocks, last q block, last k block, rows in the last q block and in the last k block + var params: [8]i32 = .{ cfg.heads, @intCast(lq), @intCast(lk), @intCast(nk), @intCast(nq - 1), @intCast(nk - 1), @intCast(lq - (nq - 1) * 64), @intCast(lk - (nk - 1) * 32) }; + e.setBytes(std.mem.asBytes(¶ms), 3); + var f: [4]f32 = .{ 1.0 / @sqrt(@as(f32, cfg.head_dim)), 0, 0, 0 }; + e.setBytes(std.mem.asBytes(&f), 4); + e.setBuffer(o.buf, o.off, 5); + e.dispatchGroups(mtl.Size.of(nq, cfg.heads, 1), mtl.Size.of(32, 4, 1)); + } + + /// y = LayerNorm32 of x's rows in fp32 rounded to fp16, then affine (w, b fp32) or modulated (s1, sh fp16). + pub fn norm(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, s1: Ref, sh: Ref, y: Ref, rows: usize, eps: f32, mode: Norm) void { + const params = extern struct { width: u32, eps: f32, mode: u32 }{ .width = cfg.hidden, .eps = eps, .mode = @backingInt(mode) }; + bind(e, k.get("id_norm"), &.{ x, w, b, s1, sh, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// y fp32 = LayerNorm of x's rows times w plus b (affine), or times (1 + w) plus b (modulated). + pub fn norm32(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, rows: usize, eps: f32, mode: Norm) void { + const params = extern struct { width: u32, eps: f32, mode: u32 }{ .width = cfg.hidden, .eps = eps, .mode = @backingInt(mode) }; + bind(e, k.get("id_norm32"), &.{ x, w, b, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// cos, sin [rows, 16, 32] of RePo from its gate and content sums (fp32 rows of 256) and final map. + pub fn repo(k: *const Kernels, e: mtl.ComputeEncoder, gc: Ref, inv: f32, final: Ref, freqs: Ref, cos: Ref, sin: Ref, rows: usize) void { + const params = extern struct { rows: u32, stride: u32, inv: f32 }{ .rows = @intCast(rows), .stride = 2 * cfg.repo_hidden, .inv = inv }; + bind(e, k.get("id_repo"), &.{ gc, final, freqs, cos, sin }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(64, 1, 1)); + } + + /// q, k, v heads [16, rows, 64] from packed qkv rows: rotated by cos, sin, then normed by gq, gk. + pub fn heads(k: *const Kernels, e: mtl.ComputeEncoder, qkv: Ref, gq: Ref, gk: Ref, cos: Ref, sin: Ref, q: Ref, kk: Ref, v: Ref, rows: usize) void { + const params = extern struct { rows: u32, scale: f32 }{ .rows = @intCast(rows), .scale = @sqrt(@as(f32, cfg.head_dim)) }; + bind(e, k.get("id_qk"), &.{ qkv, gq, gk, cos, sin, q, kk, v }, params); + e.dispatchGroups(mtl.Size.of(rows, 2, 1), mtl.Size.of(256, 1, 1)); + } + + /// x [rows, 1024] += y gated by the fp16 vector g, in place. + pub fn gateAdd(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, y: Ref, g: Ref, rows: usize) void { + bind(e, k.get("id_gate_add"), &.{ x, y, g }, @as(u32, cfg.hidden)); + e.dispatchThreads(mtl.Size.of(cfg.hidden, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// x [8192, 1024] = fp16(fp16(z w^T + b) + pe). + pub fn embed(k: *const Kernels, e: mtl.ComputeEncoder, z: Ref, w: Ref, b: Ref, pe: Ref, x: Ref) void { + e.setPipeline(k.get("id_embed")); + for ([_]Ref{ z, w, b, pe, x }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.dispatchThreads(mtl.Size.of(cfg.hidden, cfg.latent_rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [rows, 1024] = fp16((a + ba) + (c + bc)). + pub fn cond(k: *const Kernels, e: mtl.ComputeEncoder, a: Ref, ba: Ref, c: Ref, bc: Ref, y: Ref, rows: usize) void { + bind(e, k.get("id_cond"), &.{ a, ba, c, bc, y }, @as(u32, cfg.hidden)); + e.dispatchThreads(mtl.Size.of(cfg.hidden, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [rows, width] += b by column, fp32. + pub fn bias32(k: *const Kernels, e: mtl.ComputeEncoder, y: Ref, b: Ref, rows: usize, width: usize) void { + bind(e, k.get("id_bias32"), &.{ y, b }, @as(u32, @intCast(width))); + e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(@min(width, 256), 1, 1)); + } +}; + +/// One kernel's buffers in order, then its parameters at the next index. +pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); +} + +test "every Idol pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + if (!device.tensorUnits()) return error.SkipZigTest; + var k = try Kernels.load(std.testing.allocator, device); + k.deinit(); +} diff --git a/zig/src/families/idol/model.zig b/zig/src/families/idol/model.zig new file mode 100644 index 0000000000..27c6e12a13 --- /dev/null +++ b/zig/src/families/idol/model.zig @@ -0,0 +1,244 @@ +//! Idol's denoiser on Metal: the image condition refined once, then latent, condition and camera rows each step. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const host = @import("host.zig"); +const wts = @import("weights.zig"); +const ks = @import("kernels.zig"); +const Ref = ks.Ref; + +const H = cfg.hidden; +const F = cfg.ffn; +const L = cfg.joint_rows; +const R = cfg.latent_rows; +const C = cfg.cond_rows; + +pub const Timing = struct { wall: f64, gpu: f64, host: f64 }; + +/// A refined image condition: the condition refiner's rows [4101, 1024] fp16. +pub const Condition = struct { + rows: mtl.Buffer, + + pub fn deinit(c: *Condition) void { + c.rows.deinit(); + } +}; + +/// The scratch buffer's parts in bytes, in the order of Model's views. +const parts = [_]usize{ + L * H * 2, L * H * 2, L * 3 * H * 2, L * H * 2, L * H * 2, L * H * 2, L * H * 2, L * H * 2, + L * F * 2, L * H * 4, L * 2 * cfg.repo_hidden * 4, L * cfg.heads * cfg.repo_pairs * 4, L * cfg.heads * cfg.repo_pairs * 4, R * cfg.latent_channels * 4, R * cfg.latent_channels * 4, cfg.modulated * host.vectors * H * 2, + 2 * H * 4, C * cfg.cond_channels * 4, C * cfg.cond2_channels * 4, H * 2, R * H * 2, +}; + +pub fn scratchBytes() usize { + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n, 256); + return total; +} + +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + scratch: mtl.Buffer, + // fp16 rows: the stream, a norm's output, packed qkv, heads [16, rows, 64], attention rows, a projection's output + x: Ref, + xn: Ref, + qkv: Ref, + q: Ref, + k_: Ref, + v: Ref, + att: Ref, + y: Ref, + /// the MLP's GELU rows + u: Ref, + /// fp32 rows: RePo's norm (and the condition's two projections), its gate | content sums, cos and sin + ln: Ref, + gc: Ref, + cos: Ref, + sin: Ref, + lat: Ref, + vel: Ref, + mods: Ref, + head: Ref, + f1: Ref, + f2: Ref, + ones: Ref, + pe: Ref, + /// The model time `mods` and `head` hold, if any. + time: ?f32 = null, + + /// The denoiser of snapshot folder `dir`. + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !*Model { + // qualified on the M5's tensor units alone: the projections and attention take another path before them + if (!device.tensorUnits()) { + std.log.err("idol: {s} has no tensor units; this family runs on the M5's alone, the only chip it was checked on", .{device.name()}); + return error.TensorUnitsRequired; + } + const t = try gpa.create(Model); + errdefer gpa.destroy(t); + t.* = undefined; + t.gpa = gpa; + t.device = device; + t.time = null; + t.queue = try device.queue(); + errdefer t.queue.deinit(); + t.k = try ks.Kernels.load(gpa, device); + errdefer t.k.deinit(); + t.w = try wts.load(gpa, io, device, dir); + errdefer t.w.deinit(); + t.scratch = try device.buffer(scratchBytes(), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &t.x, &t.xn, &t.qkv, &t.q, &t.k_, &t.v, &t.att, &t.y, &t.u, &t.ln, &t.gc, &t.cos, &t.sin, &t.lat, &t.vel, &t.mods, &t.head, &t.f1, &t.f2, &t.ones, &t.pe }; + for (views, parts) |view, n| { + view.* = .{ .buf = t.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + @memset(halves(t.ones, H), 1.0); + host.positions(t.w.host.pos, halves(t.pe, R * H)); + return t; + } + + pub fn deinit(t: *Model) void { + t.scratch.deinit(); + t.w.deinit(); + t.k.deinit(); + t.queue.deinit(); + t.gpa.destroy(t); + } + + fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn halves(r: Ref, n: usize) []f16 { + return @as([*]f16, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + /// The modulated block `i`'s vector `part` (0 shift, 1 one plus scale, 2 gate; 3 to 5 the MLP's). + fn vec(t: *const Model, i: usize, part: usize) Ref { + return t.mods.at(((i * host.vectors) + part) * H * 2); + } + + /// The image condition: DINOv3 rows [4101, 1280] and FLUX.2 VAE rows [4101, 128] (zeros: the negative branch). + pub fn condition(t: *Model, feature1: []const f32, feature2: []const f32) !Condition { + if (feature1.len != C * cfg.cond_channels or feature2.len != C * cfg.cond2_channels) return error.ConditionRows; + var c: Condition = .{ .rows = try t.device.buffer(C * H * 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked) }; + errdefer c.rows.deinit(); + @memcpy(floats(t.f1, feature1.len), feature1); + @memcpy(floats(t.f2, feature2.len), feature2); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + const a = t.ln; + const b = t.ln.at(C * H * 4); + t.k.linear32(e, t.f1, t.w.cond, a, C, H, cfg.cond_channels, H); + t.k.linear32(e, t.f2, t.w.cond2, b, C, H, cfg.cond2_channels, H); + t.k.cond(e, a, t.w.cond_b, b, t.w.cond2_b, t.x, C); + for (t.w.context, t.w.context_repo) |blk, r| t.block(e, blk, r, C, null); + _ = try finish(cb, e, started, 0); + @memcpy(halves(.{ .buf = c.rows }, C * H), halves(t.x, C * H)); + return c; + } + + /// The velocity of latent [8192, 16] and camera [5] at model time `time` (1000 times t) under condition `c`. + pub fn forward(t: *Model, c: *const Condition, latent: []const f32, camera: *const [cfg.camera_channels]f32, time: f32, velocity: []f32, camera_velocity: *[cfg.camera_channels]f32) !Timing { + if (latent.len != R * cfg.latent_channels or velocity.len != latent.len) return error.LatentRows; + const started = mtl.clock.seconds(); + if (t.time != time) { + t.time = null; + host.modulation(t.w.host, time, @ptrCast(@alignCast(t.mods.addr())), @ptrCast(@alignCast(t.head.addr()))); + t.time = time; + } + @memcpy(floats(t.lat, latent.len), latent); + @memcpy(halves(t.x.at(R * H * 2), C * H), halves(.{ .buf = c.rows }, C * H)); + host.cameraToken(t.w.host, camera, halves(t.x.at((L - 1) * H * 2), H)); + const spent = mtl.clock.seconds() - started; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.k.embed(e, t.lat, t.w.input, t.w.input_b, t.pe, t.x); + for (t.w.noise, t.w.noise_repo, 0..) |blk, r, i| t.block(e, blk, r, R, i); + for (t.w.blocks, t.w.repo, 0..) |blk, r, i| t.block(e, blk, r, L, cfg.refiners + i); + const head: Ref = t.head; + t.k.norm32(e, t.x, head.at(H * 4), head, t.ln, R, cfg.default_eps, .modulated); + t.k.linear32(e, t.ln, t.w.out, t.vel, R, cfg.latent_channels, H, cfg.latent_channels); + t.k.bias32(e, t.vel, t.w.out_b, R, cfg.latent_channels); + const timing = try finish(cb, e, started, spent); + @memcpy(velocity, floats(t.vel, velocity.len)); + host.cameraVelocity(t.w.host, halves(t.x.at((L - 1) * H * 2), H), @ptrCast(@alignCast(t.head.addr())), camera_velocity); + return timing; + } + + /// Modulated block `index` alone at model time `time` on `rows` (the noise refiner's 8192, else 12294) given rows. + pub fn blockAt(t: *Model, index: usize, rows: []const f16, time: f32) !void { + const n: usize = if (index < cfg.refiners) R else L; + if (index >= cfg.modulated or rows.len != n * H) return error.BlockRows; + if (t.time != time) { + t.time = null; + host.modulation(t.w.host, time, @ptrCast(@alignCast(t.mods.addr())), @ptrCast(@alignCast(t.head.addr()))); + t.time = time; + } + @memcpy(halves(t.x, rows.len), rows); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + if (index < cfg.refiners) t.block(e, t.w.noise[index], t.w.noise_repo[index], n, index) else t.block(e, t.w.blocks[index - cfg.refiners], t.w.repo[index - cfg.refiners], n, index); + _ = try finish(cb, e, mtl.clock.seconds(), 0); + } + + /// The stream's first `rows` rows as the last pass left them, fp16. + pub fn stream(t: *const Model, rows: usize) []const f16 { + return halves(t.x, rows * H); + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder, started: f64, spent: f64) !Timing { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return .{ .wall = mtl.clock.seconds() - started, .gpu = cb.gpuSeconds(), .host = spent }; + } + + /// One block over the stream's first `rows` rows: modulated block `mod`, or the condition refiner's (null). + fn block(t: *const Model, e: mtl.ComputeEncoder, b: wts.Block, r: wts.Repo, rows: usize, mod: ?usize) void { + const k = &t.k; + // RePo: angles from the block's input rows + k.normSplit(e, t.x, r.norm, r.norm_b, t.qkv, rows, cfg.default_eps); + k.split(e, t.qkv, r.maps, t.gc, rows); + k.repo(e, t.gc, r.inv, r.final, r.freqs, t.cos, t.sin, rows); + if (mod) |m| { + k.norm(e, t.x, t.ln, t.ln, t.vec(m, 1), t.vec(m, 0), t.xn, rows, cfg.block_eps, .modulated); + } else { + k.norm(e, t.x, b.norm1, b.norm1_b, t.ones, t.ones, t.xn, rows, cfg.block_eps, .affine); + } + k.linear(e, t.xn, b.qkv, b.qkv_b, t.qkv, rows, 3 * H, H, false); + k.heads(e, t.qkv, b.q_norm, b.k_norm, t.cos, t.sin, t.q, t.k_, t.v, rows); + k.attention(e, t.q, t.k_, t.v, t.att, rows, rows); + k.linear(e, t.att, b.out, b.out_b, t.y, rows, H, H, false); + k.gateAdd(e, t.x, t.y, if (mod) |m| t.vec(m, 2) else t.ones, rows); + if (mod) |m| { + k.norm(e, t.x, t.ln, t.ln, t.vec(m, 4), t.vec(m, 3), t.xn, rows, cfg.block_eps, .modulated); + } else { + k.norm(e, t.x, b.norm2, b.norm2_b, t.ones, t.ones, t.xn, rows, cfg.block_eps, .affine); + } + k.linear(e, t.xn, b.up, b.up_b, t.u, rows, F, H, true); + k.linear(e, t.u, b.down, b.down_b, t.y, rows, H, F, false); + k.gateAdd(e, t.x, t.y, if (mod) |m| t.vec(m, 5) else t.ones, rows); + } +}; + +test "scratch for the joint rows fits a GPU buffer, each part 256-aligned" { + try std.testing.expect(scratchBytes() < 1 << 31); + try std.testing.expectEqual(@as(usize, 12294), L); +} diff --git a/zig/src/families/idol/reference.zig b/zig/src/families/idol/reference.zig new file mode 100644 index 0000000000..549b92fa99 --- /dev/null +++ b/zig/src/families/idol/reference.zig @@ -0,0 +1,79 @@ +//! Comparing a velocity with idol-sdk's: cosine, largest difference and worst row, and the gate against fp32. +const std = @import("std"); + +pub const Match = struct { max_abs: f64, cosine: f64, worst_row: f64 }; + +/// Largest absolute difference, cosine, and the largest relative L2 error of one row of `width`. +pub fn match(ours: []const f32, theirs: []const f32, width: usize) Match { + std.debug.assert(ours.len == theirs.len and ours.len % width == 0); + var dot: f64 = 0; + var aa: f64 = 0; + var bb: f64 = 0; + var most: f64 = 0; + var worst: f64 = 0; + var row: usize = 0; + while (row < ours.len) : (row += width) { + var rd: f64 = 0; + var rb: f64 = 0; + for (ours[row..][0..width], theirs[row..][0..width]) |a, b| { + const d = @as(f64, a) - b; + dot += @as(f64, a) * b; + aa += @as(f64, a) * a; + bb += @as(f64, b) * b; + rd += d * d; + rb += @as(f64, b) * b; + most = @max(most, @abs(d)); + } + worst = @max(worst, @sqrt(rd / @max(rb, 1e-30))); + } + return .{ .max_abs = most, .cosine = dot / @sqrt(@max(aa * bb, 1e-300)), .worst_row = worst }; +} + +/// How far ours may sit from the fp32 ruler: this many times the torch fp16 reference's own distance. +pub const factor = 2.0; + +/// Against the fp32 ruler: worst row and 1 - cosine within `factor` times the fp16 reference's own; non-finite fails. +pub fn within(ours: Match, reference: Match) bool { + if (!std.math.isFinite(ours.cosine) or !std.math.isFinite(ours.worst_row)) return false; + return ours.worst_row <= factor * reference.worst_row and (1 - ours.cosine) <= factor * (1 - reference.cosine); +} + +/// The cosine half of `within`: for a velocity, whose worst row the gate takes over the whole run instead. +pub fn withinCosine(ours: Match, reference: Match) bool { + return std.math.isFinite(ours.cosine) and (1 - ours.cosine) <= factor * (1 - reference.cosine); +} + +test "a velocity matches itself exactly" { + const v = [_]f32{ 1, -2, 3, 0.5, 0.25, -1 }; + const m = match(&v, &v, 3); + try std.testing.expectEqual(@as(f64, 0), m.max_abs); + try std.testing.expectEqual(@as(f64, 0), m.worst_row); + try std.testing.expectApproxEqAbs(@as(f64, 1), m.cosine, 1e-15); +} + +test "a planted 5% error in one row fails the gate though everything else repeats exactly" { + const a = std.testing.allocator; + const n = 64 * 16; + const ruler = try a.alloc(f32, n); + defer a.free(ruler); + const fp16 = try a.alloc(f32, n); + defer a.free(fp16); + const ours = try a.alloc(f32, n); + defer a.free(ours); + var rng = std.Random.DefaultPrng.init(7); + for (ruler, fp16, ours) |*r, *f, *o| { + r.* = rng.random().floatNorm(f32); + f.* = r.* * (1 + 1e-3 * rng.random().floatNorm(f32)); + o.* = r.* * (1 + 1e-3 * rng.random().floatNorm(f32)); + } + const reference = match(fp16, ruler, 16); + try std.testing.expect(within(match(ours, ruler, 16), reference)); + for (ours[5 * 16 ..][0..16]) |*o| o.* *= 1.05; + try std.testing.expect(!within(match(ours, ruler, 16), reference)); +} + +test "non-finite values fail the gate" { + const ruler = [_]f32{ 1, 2, 3, 4 }; + const bad = [_]f32{ 1, std.math.nan(f32), 3, 4 }; + try std.testing.expect(!within(match(&bad, &ruler, 2), match(&ruler, &ruler, 2))); +} diff --git a/zig/src/families/idol/weights.zig b/zig/src/families/idol/weights.zig new file mode 100644 index 0000000000..27d872168c --- /dev/null +++ b/zig/src/families/idol/weights.zig @@ -0,0 +1,310 @@ +//! Idol's denoiser tensors (fp32 as stored): the blocks' linears to fp16 at load as idol-sdk does, the rest fp32. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const cfg = @import("config.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +const H = cfg.hidden; + +/// A transformer block: fp16 linears and biases, fp32 head RMSNorm gammas [16, 64]; norms only in the refiner. +pub const Block = struct { qkv: Ref, qkv_b: Ref, out: Ref, out_b: Ref, up: Ref, up_b: Ref, down: Ref, down_b: Ref, q_norm: Ref, k_norm: Ref, norm1: Ref = undefined, norm1_b: Ref = undefined, norm2: Ref = undefined, norm2_b: Ref = undefined }; + +/// A RePo layer: fp32 norm; gate | content maps over `inv` (a power of two) as fp16 [hi | lo | hi]; final map, freqs. +pub const Repo = struct { norm: Ref, norm_b: Ref, maps: Ref, inv: f32, final: Ref, freqs: Ref }; + +const repos = 2 * cfg.refiners + cfg.layers; +const map_bytes = 2 * cfg.repo_hidden * 3 * H * 2; + +/// The power of two that puts a map's largest magnitude in [2^13, 2^14), so its fp16 halves keep their bits. +fn mapScale(gate: []const f32, content: []const f32) f32 { + var most: f32 = 0; + for (gate) |v| most = @max(most, @abs(v)); + for (content) |v| most = @max(most, @abs(v)); + if (most == 0 or !std.math.isFinite(most)) return 1; + return std.math.ldexp(@as(f32, 1), 13 - std.math.ilogb(most)); +} + +/// Row `n` of a map as [hi | lo | hi] fp16 of `scale` times its values. +fn split(src: []const f32, scale: f32, dst: []f16) void { + for (src, 0..) |v, i| { + const x = v * scale; + const hi: f16 = @floatCast(x); + dst[i] = hi; + dst[H + i] = @floatCast(x - @as(f32, hi)); + dst[2 * H + i] = hi; + } +} + +/// The fp32 vectors and small matrices the host reads each timestep. +pub const Host = struct { + t1: []const f32, + t1_b: []const f32, + t2: []const f32, + t2_b: []const f32, + ada: []const f32, + ada_b: []const f32, + cam1: []const f32, + cam1_b: []const f32, + cam2: []const f32, + cam2_b: []const f32, + cam_out: []const f32, + cam_out_b: []const f32, + /// Each modulated block's own shift table [6144]: the noise refiner's, then the joint blocks'. + shifts: [cfg.modulated][]const f32, + /// The output's shift table [2, 1024]. + head: []const f32, + pos: []const f32, +}; + +pub const Weights = struct { + ck: ck.Checkpoint, + noise: [cfg.refiners]Block, + context: [cfg.refiners]Block, + blocks: [cfg.layers]Block, + noise_repo: [cfg.refiners]Repo, + context_repo: [cfg.refiners]Repo, + repo: [cfg.layers]Repo, + input: Ref, + input_b: Ref, + cond: Ref, + cond_b: Ref, + cond2: Ref, + cond2_b: Ref, + out: Ref, + out_b: Ref, + host: Host, + + pub fn deinit(w: *Weights) void { + w.ck.deinit(); + } +}; + +const linears = [_][]const u8{ "attn.qkv", "attn.out", "mlp.mlp.0", "mlp.mlp.2" }; + +/// The blocks' linear weights stay in the file: load reads each once into its fp16 copy. +fn later(name: []const u8, rank: usize) bool { + if (rank != 2) return false; + const owned = std.mem.startsWith(u8, name, "blocks.") or std.mem.startsWith(u8, name, "noise_refiner.") or std.mem.startsWith(u8, name, "context_refiner."); + return owned and std.mem.endsWith(u8, name, ".weight"); +} + +fn check(t: ck.Tensor, name: []const u8, shape: []const usize) !void { + if (t.dtype != .f32 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("idol: {s} is {t} {any}; the denoiser reads f32 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } +} + +fn take(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) !Ref { + const t = try c.get(name); + try check(t, name, shape); + if (t.offset % 16 != 0) return error.UnalignedTensor; + return .{ .buf = t.buffer, .off = t.offset }; +} + +fn floats(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) ![]const f32 { + const t = try c.get(name); + try check(t, name, shape); + return t.host(f32); +} + +/// fp16 values written into an owned buffer, bump-allocated 256-aligned. +const Halves = struct { + buf: mtl.Buffer, + at: usize = 0, + stage: []f32, + + fn put(h: *Halves, src: []const f32) Ref { + const dst = @as([*]f16, @ptrCast(@alignCast(h.buf.contents() + h.at)))[0..src.len]; + for (dst, src) |*o, v| o.* = @floatCast(v); + const r: Ref = .{ .buf = h.buf, .off = h.at }; + h.at += std.mem.alignForward(usize, src.len * 2, ck.alignment); + return r; + } + + /// A deferred fp32 weight read in staged pieces, rounded to fp16. + fn read(h: *Halves, c: *const ck.Checkpoint, name: []const u8, shape: [2]usize) !Ref { + const d = c.deferred.get(name) orelse return error.MissingTensor; + if (d.dtype != .f32 or d.rank != 2 or d.shape[0] != shape[0] or d.shape[1] != shape[1]) { + std.log.err("idol: {s} is {t} {any}; the denoiser reads f32 {any}", .{ name, d.dtype, d.shape[0..d.rank], shape }); + return error.UnexpectedTensor; + } + const n = shape[0] * shape[1]; + const r: Ref = .{ .buf = h.buf, .off = h.at }; + const dst = @as([*]f16, @ptrCast(@alignCast(h.buf.contents() + h.at)))[0..n]; + var done: usize = 0; + while (done < n) { + const part = @min(h.stage.len, n - done); + try c.read(name, done * 4, std.mem.sliceAsBytes(h.stage[0..part])); + for (dst[done..][0..part], h.stage[0..part]) |*o, v| o.* = @floatCast(v); + done += part; + } + h.at += std.mem.alignForward(usize, n * 2, ck.alignment); + return r; + } +}; + +const per_block = blk: { + var n: usize = 0; + for ([_]usize{ 3 * H * H, 3 * H, H * H, H, cfg.ffn * H, cfg.ffn, H * cfg.ffn, H }) |v| n += std.mem.alignForward(usize, v * 2, ck.alignment); + break :blk n; +}; + +fn block(c: *const ck.Checkpoint, h: *Halves, prefix: []const u8, i: usize, affine: bool) !Block { + var name: [96]u8 = undefined; + const P = struct { + fn n(buf: []u8, p: []const u8, j: usize, s: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.{s}", .{ p, j, s }); + } + }; + var b: Block = undefined; + b.qkv = try h.read(c, try P.n(&name, prefix, i, "attn.qkv.weight"), .{ 3 * H, H }); + b.qkv_b = h.put(try floats(c, try P.n(&name, prefix, i, "attn.qkv.bias"), &.{3 * H})); + b.out = try h.read(c, try P.n(&name, prefix, i, "attn.out.weight"), .{ H, H }); + b.out_b = h.put(try floats(c, try P.n(&name, prefix, i, "attn.out.bias"), &.{H})); + b.up = try h.read(c, try P.n(&name, prefix, i, "mlp.mlp.0.weight"), .{ cfg.ffn, H }); + b.up_b = h.put(try floats(c, try P.n(&name, prefix, i, "mlp.mlp.0.bias"), &.{cfg.ffn})); + b.down = try h.read(c, try P.n(&name, prefix, i, "mlp.mlp.2.weight"), .{ H, cfg.ffn }); + b.down_b = h.put(try floats(c, try P.n(&name, prefix, i, "mlp.mlp.2.bias"), &.{H})); + b.q_norm = try take(c, try P.n(&name, prefix, i, "attn.q_norm.gamma"), &.{ cfg.heads, cfg.head_dim }); + b.k_norm = try take(c, try P.n(&name, prefix, i, "attn.k_norm.gamma"), &.{ cfg.heads, cfg.head_dim }); + if (affine) { + b.norm1 = try take(c, try P.n(&name, prefix, i, "norm1.weight"), &.{H}); + b.norm1_b = try take(c, try P.n(&name, prefix, i, "norm1.bias"), &.{H}); + b.norm2 = try take(c, try P.n(&name, prefix, i, "norm2.weight"), &.{H}); + b.norm2_b = try take(c, try P.n(&name, prefix, i, "norm2.bias"), &.{H}); + } + return b; +} + +/// A RePo layer; its two maps into `pack` at `slot`, and its three frequency vectors as 32 values after every map. +fn repo(c: *const ck.Checkpoint, pack: mtl.Buffer, slot: usize, prefix: []const u8, i: usize) !Repo { + var name: [96]u8 = undefined; + var r: Repo = undefined; + r.norm = try take(c, try std.fmt.bufPrint(&name, "{s}.{d}.norm.weight", .{ prefix, i }), &.{H}); + r.norm_b = try take(c, try std.fmt.bufPrint(&name, "{s}.{d}.norm.bias", .{ prefix, i }), &.{H}); + const gate = try floats(c, try std.fmt.bufPrint(&name, "{s}.{d}.gate_map.weight", .{ prefix, i }), &.{ cfg.repo_hidden, H }); + const content = try floats(c, try std.fmt.bufPrint(&name, "{s}.{d}.content_map.weight", .{ prefix, i }), &.{ cfg.repo_hidden, H }); + const scale = mapScale(gate, content); + const maps = @as([*]f16, @ptrCast(@alignCast(pack.contents() + slot * map_bytes)))[0 .. 2 * cfg.repo_hidden * 3 * H]; + for (0..2 * cfg.repo_hidden) |n| { + const src = if (n < cfg.repo_hidden) gate[n * H ..][0..H] else content[(n - cfg.repo_hidden) * H ..][0..H]; + split(src, scale, maps[n * 3 * H ..][0 .. 3 * H]); + } + r.maps = .{ .buf = pack, .off = slot * map_bytes }; + r.inv = 1 / scale; + r.final = try take(c, try std.fmt.bufPrint(&name, "{s}.{d}.final_map.weight", .{ prefix, i }), &.{ 3 * cfg.heads, cfg.repo_hidden }); + const freqs = repos * map_bytes; + const dst = @as([*]f32, @ptrCast(@alignCast(pack.contents() + freqs)))[slot * cfg.repo_pairs ..][0..cfg.repo_pairs]; + var at: usize = 0; + inline for (.{ .{ "freqs_0", 10 }, .{ "freqs_1", 10 }, .{ "freqs_2", 12 } }) |f| { + const src = try floats(c, try std.fmt.bufPrint(&name, "{s}.{d}.{s}", .{ prefix, i, f[0] }), &.{f[1]}); + @memcpy(dst[at..][0..f[1]], src); + at += f[1]; + } + r.freqs = .{ .buf = pack, .off = freqs + slot * cfg.repo_pairs * 4 }; + return r; +} + +/// The snapshot folder `dir`: config.json checked, then denoiser.safetensors's every tensor, each read once. +pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "config.json" }), arena, .limited(1 << 20)); + try cfg.check(gpa, config); + var w: Weights = undefined; + w.ck = ck.Checkpoint.init(gpa); + errdefer w.ck.deinit(); + const path = try std.fmt.allocPrintSentinel(arena, "{s}/denoiser.safetensors", .{dir}, 0); + try w.ck.addFileSplit(device, path, "", null, later); + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const blocks = 2 * cfg.refiners + cfg.layers; + const fp16 = try device.buffer(blocks * per_block, opts); + w.ck.own(fp16, blocks * per_block) catch |err| { + fp16.deinit(); + return err; + }; + const pack_bytes = repos * (map_bytes + cfg.repo_pairs * 4); + const pack = try device.buffer(pack_bytes, opts); + w.ck.own(pack, pack_bytes) catch |err| { + pack.deinit(); + return err; + }; + var h: Halves = .{ .buf = fp16, .stage = try arena.alloc(f32, 4 << 20) }; + for (&w.noise, &w.context, &w.noise_repo, &w.context_repo, 0..) |*n, *x, *nr, *xr, i| { + n.* = try block(&w.ck, &h, "noise_refiner", i, false); + x.* = try block(&w.ck, &h, "context_refiner", i, true); + nr.* = try repo(&w.ck, pack, i, "noise_repo_layers", i); + xr.* = try repo(&w.ck, pack, cfg.refiners + i, "context_repo_layers", i); + } + for (&w.blocks, &w.repo, 0..) |*b, *r, i| { + b.* = try block(&w.ck, &h, "blocks", i, false); + r.* = try repo(&w.ck, pack, 2 * cfg.refiners + i, "repo_layers", i); + } + w.ck.closeFiles(); + w.input = try take(&w.ck, "input_layer.weight", &.{ H, cfg.latent_channels }); + w.input_b = try take(&w.ck, "input_layer.bias", &.{H}); + w.cond = try take(&w.ck, "cond_embedder.weight", &.{ H, cfg.cond_channels }); + w.cond_b = try take(&w.ck, "cond_embedder.bias", &.{H}); + w.cond2 = try take(&w.ck, "cond_embedder2.weight", &.{ H, cfg.cond2_channels }); + w.cond2_b = try take(&w.ck, "cond_embedder2.bias", &.{H}); + w.out = try take(&w.ck, "out_layer.weight", &.{ cfg.latent_channels, H }); + w.out_b = try take(&w.ck, "out_layer.bias", &.{cfg.latent_channels}); + var hw: Host = undefined; + hw.t1 = try floats(&w.ck, "t_embedder.mlp.0.weight", &.{ H, cfg.time_dim }); + hw.t1_b = try floats(&w.ck, "t_embedder.mlp.0.bias", &.{H}); + hw.t2 = try floats(&w.ck, "t_embedder.mlp.2.weight", &.{ H, H }); + hw.t2_b = try floats(&w.ck, "t_embedder.mlp.2.bias", &.{H}); + hw.ada = try floats(&w.ck, "adaLN_modulation.1.weight", &.{ 6 * H, H }); + hw.ada_b = try floats(&w.ck, "adaLN_modulation.1.bias", &.{6 * H}); + hw.cam1 = try floats(&w.ck, "cam_refiner.mlp.0.weight", &.{ H, cfg.camera_channels }); + hw.cam1_b = try floats(&w.ck, "cam_refiner.mlp.0.bias", &.{H}); + hw.cam2 = try floats(&w.ck, "cam_refiner.mlp.2.weight", &.{ H, H }); + hw.cam2_b = try floats(&w.ck, "cam_refiner.mlp.2.bias", &.{H}); + hw.cam_out = try floats(&w.ck, "cam_out_layer.weight", &.{ cfg.camera_channels, H }); + hw.cam_out_b = try floats(&w.ck, "cam_out_layer.bias", &.{cfg.camera_channels}); + var name: [64]u8 = undefined; + for (&hw.shifts, 0..) |*s, i| { + const n = if (i < cfg.refiners) try std.fmt.bufPrint(&name, "noise_refiner.{d}.shift_table", .{i}) else try std.fmt.bufPrint(&name, "blocks.{d}.shift_table", .{i - cfg.refiners}); + s.* = try floats(&w.ck, n, &.{ 1, 6 * H }); + } + hw.head = try floats(&w.ck, "shift_table", &.{ 1, 2, H }); + hw.pos = try floats(&w.ck, "pos_pe", &.{ 1, cfg.latent_rows, 3 }); + w.host = hw; + // a block's 4 linears, 4 biases and 2 gammas, the context refiner's 4 norms, the shift tables; 8 a RePo layer + const all = 2 * cfg.refiners + cfg.layers; + const read = 10 * all + 4 * cfg.refiners + cfg.modulated + 8 * all + 8 + 12 + 2; + const total = w.ck.tensors.count() + w.ck.deferred.count(); + if (total != read) { + std.log.err("idol: denoiser.safetensors has {d} tensors; this family reads {d}", .{ total, read }); + return error.UnexpectedTensors; + } + return w; +} + +test "a map's split halves sum back to its scaled values within fp32's rounding" { + var src: [H]f32 = undefined; + var rng = std.Random.DefaultPrng.init(3); + for (&src) |*v| v.* = 0.05 * rng.random().floatNorm(f32); + const scale = mapScale(&src, &src); + try std.testing.expect(scale >= 1); + var dst: [3 * H]f16 = undefined; + split(&src, scale, &dst); + for (src, 0..) |v, i| { + const back = (@as(f32, dst[i]) + @as(f32, dst[H + i])) / scale; + try std.testing.expect(@abs(back - v) <= @abs(v) * 1e-6 + 1e-12); + try std.testing.expectEqual(dst[i], dst[2 * H + i]); + } +} + +test "the blocks' linear weights, and only they, stay in the file for their fp16 copies" { + try std.testing.expect(later("blocks.3.attn.qkv.weight", 2)); + try std.testing.expect(later("context_refiner.1.mlp.mlp.2.weight", 2)); + try std.testing.expect(!later("blocks.3.attn.qkv.bias", 1)); + try std.testing.expect(!later("blocks.3.attn.q_norm.gamma", 2)); + try std.testing.expect(!later("context_refiner.0.norm1.weight", 1)); + try std.testing.expect(!later("repo_layers.0.gate_map.weight", 2)); + try std.testing.expect(!later("cond_embedder.weight", 2)); +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 23655fbcbe..89eb6e678c 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -13,6 +13,7 @@ pub const qwen_vl_tower = @import("core/qwen_vl/metal.zig"); pub const qwen_image = @import("families/qwen_image/qwen_image.zig"); /// MOSS-Image's transformer on Metal: a denoiser driven by tf-moss-image, outside the lane engine. pub const moss_image = @import("families/moss_image/moss_image.zig"); +pub const idol = @import("families/idol/idol.zig"); pub const checkpoint = @import("core/checkpoint_metal.zig"); pub const shard_edit = @import("core/shard_edit.zig"); pub const lora_merge = @import("core/lora_merge_dir.zig"); diff --git a/zig/tests/idol.zig b/zig/tests/idol.zig new file mode 100644 index 0000000000..4359d38e45 --- /dev/null +++ b/zig/tests/idol.zig @@ -0,0 +1,242 @@ +//! tf-idol: Idol's denoiser on a case from tools/zig/idol_case.py, compared with idol-sdk's, gated and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const id = tf.idol; +const cfg = id.config; +const ref = id.reference; + +const usage = + \\tf-idol SNAPSHOT_DIR CASE_DIR [blocks STEP] + \\SNAPSHOT_DIR is a speridlabs/idol snapshot (config.json, denoiser.safetensors). CASE_DIR (tools/zig/idol_case.py) + \\holds case.safetensors: both image conditions, the starting noise and the SDK's run (each step's input and both + \\branches' velocities), and optionally exact.safetensors (the same steps in float32). Each step runs both branches on + \\the SDK's input; with exact.safetensors each 1 - cosine against float32 must be within twice the SDK's own, and + \\so must the worst row over the run. Then the whole guided schedule runs from the noise, timed; CASE_DIR gets tf-idol.json and + \\sample.f32 (8192 x 16, then the 5 camera values). Exit 1 on a failed gate or a forward that does not repeat. + \\With `blocks STEP`, runs each joint block alone on the SDK's input rows at that step (blocksSTEP.safetensors) + \\and gates it against the same block in float32 (blocksSTEP_f32.safetensors), worst row and 1 - cosine. +; + +/// The card's sampling defaults. +const guidance: f32 = 3.0; +const rescale: f64 = 3.0; + +fn take(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) ![]f32 { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + if (t.dtype != .f32 or t.numel() != n) { + std.debug.print("{s} is not {d} float32 values\n", .{ name, n }); + return error.BadCaseTensor; + } + const out = try a.alloc(f32, n); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn json(list: *std.ArrayList(u8), a: std.mem.Allocator, key: []const u8, m: ref.Match) !void { + try print(list, a, ", \"{s}\": {{\"max_abs\": {d:.6}, \"cosine\": {d:.9}, \"worst_row\": {d:.6}}}", .{ key, m.max_abs, m.cosine, m.worst_row }); +} + +const Case = struct { + steps: usize, + timesteps: []f32, + inputs: []f32, + inputs_camera: []f32, + velocity: [2][]f32, + velocity_camera: [2][]f32, + exact: ?[2][]f32, +}; + +const N = cfg.latent_rows * cfg.latent_channels; +const labels = [2][]const u8{ "conditional", "negative" }; + +/// Each step's two branches on the SDK's input against its velocities; returns false when a gate fails. +fn teacher(a: std.mem.Allocator, c: Case, model: *id.Model, conds: *[2]id.Condition, out: *std.ArrayList(u8)) !bool { + var pass = true; + var worst = [2]f64{ 0, 0 }; + const v = try a.alloc(f32, N); + var vc: [cfg.camera_channels]f32 = undefined; + try print(out, a, ", \"teacher_forced\": [", .{}); + for (0..c.steps) |s| { + try print(out, a, "{s}{{\"step\": {d}, \"time\": {d:.4}", .{ if (s == 0) "" else ", ", s + 1, c.timesteps[s] }); + const cam: *const [cfg.camera_channels]f32 = c.inputs_camera[s * cfg.camera_channels ..][0..cfg.camera_channels]; + for (0..2) |b| { + _ = try model.forward(&conds[b], c.inputs[s * N ..][0..N], cam, c.timesteps[s], v, &vc); + const theirs = c.velocity[b][s * N ..][0..N]; + const m = ref.match(v, theirs, cfg.latent_channels); + const mc = ref.match(&vc, c.velocity_camera[b][s * cfg.camera_channels ..][0..cfg.camera_channels], cfg.camera_channels); + std.debug.print("step {d} {s}: against the SDK max |diff| {d:.5}, cosine {d:.8}, worst row {d:.5}; camera cosine {d:.7}\n", .{ s + 1, labels[b], m.max_abs, m.cosine, m.worst_row, mc.cosine }); + try json(out, a, try std.fmt.allocPrint(a, "{s}_sdk", .{labels[b]}), m); + try json(out, a, try std.fmt.allocPrint(a, "{s}_camera_sdk", .{labels[b]}), mc); + if (c.exact) |ex| if (ex[b].len >= (s + 1) * N) { + const want = ex[b][s * N ..][0..N]; + const ours = ref.match(v, want, cfg.latent_channels); + const sdk = ref.match(theirs, want, cfg.latent_channels); + const ok = ref.withinCosine(ours, sdk); + pass = pass and ok; + worst[0] = @max(worst[0], ours.worst_row); + worst[1] = @max(worst[1], sdk.worst_row); + std.debug.print(" against float32: worst row {d:.5}, cosine {d:.8}; the SDK's own {d:.5}, {d:.8}: {s}\n", .{ ours.worst_row, ours.cosine, sdk.worst_row, sdk.cosine, if (ok) "pass" else "FAIL" }); + try json(out, a, try std.fmt.allocPrint(a, "{s}_fp32", .{labels[b]}), ours); + try json(out, a, try std.fmt.allocPrint(a, "{s}_sdk_vs_fp32", .{labels[b]}), sdk); + }; + } + try out.append(a, '}'); + } + if (c.exact != null) { + const ok = worst[0] <= ref.factor * worst[1]; + std.debug.print("worst row over the run against float32: {d:.5}; the SDK's own {d:.5}: {s}\n", .{ worst[0], worst[1], if (ok) "pass" else "FAIL" }); + pass = pass and ok; + } + try print(out, a, "], \"worst_row_fp32\": {d:.6}, \"sdk_worst_row_fp32\": {d:.6}", .{ worst[0], worst[1] }); + return pass; +} + +/// The same step twice with the other branch between: identical bytes or not. +fn repeats(a: std.mem.Allocator, c: Case, model: *id.Model, conds: *[2]id.Condition) !bool { + const first = try a.alloc(f32, N); + const again = try a.alloc(f32, N); + var vc: [cfg.camera_channels]f32 = undefined; + const cam: *const [cfg.camera_channels]f32 = c.inputs_camera[0..cfg.camera_channels]; + _ = try model.forward(&conds[0], c.inputs[0..N], cam, c.timesteps[0], first, &vc); + _ = try model.forward(&conds[1], c.inputs[0..N], cam, c.timesteps[@min(1, c.steps - 1)], again, &vc); + _ = try model.forward(&conds[0], c.inputs[0..N], cam, c.timesteps[0], again, &vc); + return std.mem.eql(u8, std.mem.sliceAsBytes(first), std.mem.sliceAsBytes(again)); +} + +/// The case tensor `name`, fp16 or fp32 rows, as f32. +fn rows(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]f32 { + const t = f.get(name) orelse return error.MissingTensor; + const out = try a.alloc(f32, t.numel()); + switch (t.dtype) { + .f16 => for (out, @as([*]align(1) const f16, @ptrCast(t.bytes.ptr))[0..out.len]) |*o, v| { + o.* = v; + }, + .f32 => @memcpy(std.mem.sliceAsBytes(out), t.bytes), + else => return error.BadCaseTensor, + } + return out; +} + +/// Each joint block alone on the SDK's own input rows at step `step`, against that block in float32; false on a fail. +fn blockGates(a: std.mem.Allocator, io: std.Io, dir: []const u8, step: usize, c: Case, model: *id.Model) !bool { + var f = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/blocks{d}.safetensors", .{ dir, step })); + defer f.close(io); + var g = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/blocks{d}_f32.safetensors", .{ dir, step })); + defer g.close(io); + const H = cfg.hidden; + var pass = true; + var name: [32]u8 = undefined; + for (1..cfg.layers) |i| { + const input = try rows(a, &f, try std.fmt.bufPrint(&name, "block.{d}", .{i - 1})); + const in16 = try a.alloc(f16, input.len); + for (in16, input) |*o, v| o.* = @floatCast(v); + try model.blockAt(cfg.refiners + i, in16, c.timesteps[step]); + const ours = try a.alloc(f32, input.len); + for (ours, model.stream(cfg.joint_rows)) |*o, v| o.* = v; + const label = try std.fmt.bufPrint(&name, "block.{d}", .{i}); + const want = try rows(a, &g, label); + const m = ref.match(ours, want, H); + const r = ref.match(try rows(a, &f, label), want, H); + const ok = ref.within(m, r); + pass = pass and ok; + std.debug.print("step {d} block {d} against float32: worst row {d:.5}, cosine {d:.9}; the SDK's own {d:.5}, {d:.9}: {s}\n", .{ step + 1, i, m.worst_row, m.cosine, r.worst_row, r.cosine, if (ok) "pass" else "FAIL" }); + } + return pass; +} + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 3 and args.len != 5) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/case.safetensors", .{args[2]})); + defer case.close(io); + const steps = (case.get("timesteps") orelse return error.MissingTensor).numel(); + const S = steps; + const cam_n = cfg.camera_channels; + var c: Case = .{ .steps = S, .timesteps = try take(a, &case, "timesteps", S), .inputs = try take(a, &case, "inputs", S * N), .inputs_camera = try take(a, &case, "inputs_camera", S * cam_n), .velocity = .{ try take(a, &case, "velocity", S * N), try take(a, &case, "velocity_negative", S * N) }, .velocity_camera = .{ try take(a, &case, "velocity_camera", S * cam_n), try take(a, &case, "velocity_negative_camera", S * cam_n) }, .exact = null }; + const exact_path = try std.fmt.allocPrint(a, "{s}/exact.safetensors", .{args[2]}); + var exact_file: ?tf.safetensors.File = tf.safetensors.File.open(a, io, exact_path) catch null; + defer if (exact_file) |*f| f.close(io); + if (exact_file) |*f| { + const n = (f.get("velocity_f32") orelse return error.MissingTensor).numel(); + c.exact = .{ try take(a, f, "velocity_f32", n), try take(a, f, "velocity_negative_f32", n) }; + } + const f1 = try take(a, &case, "feature1", cfg.cond_rows * cfg.cond_channels); + const f2 = try take(a, &case, "feature2", cfg.cond_rows * cfg.cond2_channels); + const noise = try take(a, &case, "latent", N); + const noise_cam = try take(a, &case, "camera", cam_n); + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const model = try id.Model.load(std.heap.page_allocator, io, device, args[1]); + defer model.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: {d} steps; loaded in {d:.1} s\n", .{ device.name(), S, load_s }); + + var out: std.ArrayList(u8) = .empty; + const conditioning = mtl.clock.seconds(); + var conds: [2]id.Condition = undefined; + conds[0] = try model.condition(f1, f2); + defer conds[0].deinit(); + @memset(f1, 0); + @memset(f2, 0); + conds[1] = try model.condition(f1, f2); + defer conds[1].deinit(); + const condition_s = mtl.clock.seconds() - conditioning; + if (args.len == 5) return if (try blockGates(a, io, args[2], try std.fmt.parseInt(usize, args[4], 10), c, model)) 0 else 1; + try print(&out, a, "{{\"device\": \"{s}\", \"steps\": {d}, \"load_s\": {d:.2}, \"condition_s\": {d:.4}", .{ device.name(), S, load_s, condition_s }); + const pass = try teacher(a, c, model, &conds, &out); + const repeat = try repeats(a, c, model, &conds); + std.debug.print("the same step twice, another between: {s}\n", .{if (repeat) "identical bytes" else "DIFFERENT bytes"}); + try print(&out, a, ", \"gates_pass\": {}, \"repeat_identical\": {}, \"run\": [", .{ pass and c.exact != null, repeat }); + + const ts = try a.alloc(f64, S + 1); + id.host.schedule(ts, rescale); + const x = try a.alloc(f32, N + cam_n); + @memcpy(x[0..N], noise); + @memcpy(x[N..], noise_cam); + const vc = try a.alloc(f32, N); + const vu = try a.alloc(f32, N); + var cc: [cam_n]f32 = undefined; + var cu: [cam_n]f32 = undefined; + var total: f64 = 0; + for (0..S) |s| { + const time: f32 = @floatCast(1000 * ts[s]); + const cam: *const [cam_n]f32 = x[N..][0..cam_n]; + const p = try model.forward(&conds[0], x[0..N], cam, time, vc, &cc); + const q = try model.forward(&conds[1], x[0..N], cam, time, vu, &cu); + const by: f32 = @floatCast(ts[s] - ts[s + 1]); + for (x[0..N], vc, vu) |*o, pc, pu| o.* -= ((1 + guidance) * pc - guidance * pu) * by; + for (x[N..], cc, cu) |*o, pc, pu| o.* -= ((1 + guidance) * pc - guidance * pu) * by; + const wall = p.wall + q.wall; + total += wall; + std.debug.print("step {d}/{d}: {d:.3} s ({d:.3} s on the GPU)\n", .{ s + 1, S, wall, p.gpu + q.gpu }); + try print(&out, a, "{s}{{\"wall_s\": {d:.4}, \"gpu_s\": {d:.4}}}", .{ if (s == 0) "" else ", ", wall, p.gpu + q.gpu }); + } + if (case.get("sample") != null) { + const sample = try take(a, &case, "sample", N); + const m = ref.match(x[0..N], sample, cfg.latent_channels); + std.debug.print("final sample against the SDK's: cosine {d:.6}, worst row {d:.4}\n", .{ m.cosine, m.worst_row }); + try print(&out, a, "]", .{}); + try json(&out, a, "sample_sdk", m); + try print(&out, a, ", \"denoise_s\": {d:.3}}}\n", .{total}); + } else try print(&out, a, "], \"denoise_s\": {d:.3}}}\n", .{total}); + std.debug.print("denoise: {d:.2} s for {d} steps ({d} forwards)\n", .{ total, S, 2 * S }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/sample.f32", .{args[2]}), .data = std.mem.sliceAsBytes(x) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol.json", .{args[2]}), .data = out.items }); + return if (repeat and (c.exact == null or pass)) 0 else 1; +} From 6b1b074e29a5719b2ab997ee58082e5e5948b09a Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:39:20 +0800 Subject: [PATCH 02/21] contract: Idol's denoiser kernels through gates 1-5 Every pipeline of idol.metal and idol_attention.metal is registered: id_linear's eight forms (plain and GELU), id_split's two, id_gemm32, id_norm (plain, affine, modulated), id_norm32 (affine, modulated), id_norm_split, id_repo, id_qk, id_gate_add, id_embed, id_cond, id_bias32 and the attention's four forms. Each output is checked at idol-sdk's fp16 points as the range its fp32 steps allow; the NAX projections against the tensor unit's 16-deep step bound, the attention against attention_ref's formula with P at the tensor unit's fp32. Inputs at the large served shapes repeat at periods prime to every tile, so float64 stays cheap and every output is checked. 19 controls: rounding points dropped or added (gate 3), fp16 sums in id_norm32, id_repo, id_qk, id_embed and id_gemm32, and the NAX fp32 path in id_gemm32's place (gate 4), and the attention's tail keys unmasked. Gate 5 holds the 24 launches, each reserved by its source. Known finding: id_qk gives a head holding a NaN +-inf in its other outputs where F.normalize gives NaN, as Metal's max drops the NaN norm that clamp_min keeps. --- docs/kernel-contract.md | 6 + zig/src/core/contract.zig | 1 + zig/src/core/contract/idol_records.zig | 34 ++ zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_idol.zig | 190 +++++++++ zig/tests/contract/metal_idol_attention.zig | 233 +++++++++++ zig/tests/contract/metal_idol_common.zig | 245 +++++++++++ zig/tests/contract/metal_idol_gemm.zig | 430 +++++++++++++++++++ zig/tests/contract/metal_idol_norms.zig | 251 +++++++++++ zig/tests/contract/metal_idol_rows.zig | 442 ++++++++++++++++++++ zig/tests/contract/metal_idol_steps.zig | 426 +++++++++++++++++++ zig/tests/contract/metal_limits.zig | 3 + zig/tests/contract/metal_main.zig | 2 + 13 files changed, 2264 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/idol_records.zig create mode 100644 zig/tests/contract/metal_idol.zig create mode 100644 zig/tests/contract/metal_idol_attention.zig create mode 100644 zig/tests/contract/metal_idol_common.zig create mode 100644 zig/tests/contract/metal_idol_gemm.zig create mode 100644 zig/tests/contract/metal_idol_norms.zig create mode 100644 zig/tests/contract/metal_idol_rows.zig create mode 100644 zig/tests/contract/metal_idol_steps.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b690ec8f9e..472881c3b5 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -117,6 +117,8 @@ MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 an `tf_glm_scores` scores a NaN dot as 0 where float64 keeps the NaN: its relu is `fmaxf(x, 0)`, as `tl.maximum` is in the Python family's Triton glue it was ported from. Nemotron's Metal `tf_xup_w3` gives 0 for a NaN dot where float64 keeps the NaN: its relu2 is `metal::max(bfloat(acc), 0.0f)`; the other up pipelines give its bytes on the same inputs. +Idol's `id_qk` gives a head holding a NaN +-inf in its other outputs where torch's `F.normalize` gives NaN: its +norm is `max(precise::sqrt(ss), 1e-12f)`, and Metal's `max` drops the NaN that `clamp_min` keeps. HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is @@ -185,6 +187,10 @@ CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) unscaled tile's outputs rounded and scaled on the host. HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out). +Metal also checks Idol's denoiser (`metal_idol*.zig`): its fp16 projections and GELU form, the split-K fp32 rows, +`id_gemm32` (true fp32 sums: the NAX fp32 path in its place is a control gate 4 must catch), the LayerNorms, RePo's +angles, the heads, the elementwise steps and the attention's four forms, each at idol-sdk's fp16 points (19 controls). + RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. CUDA also checks GLM-5.3-Flash's KV page kernels on a page table out of order: the latent write and the pooled keys diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 6828c844cc..4153341694 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -25,6 +25,7 @@ test { _ = @import("contract/glm_rows_records.zig"); _ = @import("contract/nemotron_experts_records.zig"); _ = @import("contract/moss_image_records.zig"); + _ = @import("contract/idol_records.zig"); _ = @import("contract/metal_pair_records.zig"); _ = spec; _ = gates; diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig new file mode 100644 index 0000000000..ae39549164 --- /dev/null +++ b/zig/src/core/contract/idol_records.zig @@ -0,0 +1,34 @@ +//! Idol's denoiser launches (idol.metal and idol_attention.metal after gemm_nax.metal), each source-reserved. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "idol id_linear_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_01", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_10", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_11", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_01", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_10", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_linear_gelu_11", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_split_0", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_split_1", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_gemm32", .threads = 128, .tg_memory = 19456, .reserved = 128 }, + .{ .name = "idol id_norm", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol id_norm32", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol id_norm_split", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol id_repo", .threads = 64, .tg_memory = 704, .reserved = 64 }, + .{ .name = "idol id_qk", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_gate_add", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_embed", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_cond", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_bias32", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol id_attention_00", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_attention_01", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_attention_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol id_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, +}; + +test "the Idol launches fit every Apple family, one record a pipeline" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); + try std.testing.expectEqual(@as(usize, 24), records.len); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index d48c6541bc..f55d146bf7 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// The Metal collector keeps measured slots and the restored vision declarations. pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ - @import("moss_image_records.zig").records ++ @import("minimax_records.zig").records; + @import("moss_image_records.zig").records ++ @import("minimax_records.zig").records ++ @import("idol_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/tests/contract/metal_idol.zig b/zig/tests/contract/metal_idol.zig new file mode 100644 index 0000000000..50d1284ae4 --- /dev/null +++ b/zig/tests/contract/metal_idol.zig @@ -0,0 +1,190 @@ +//! Idol's Metal kernels under the contract, each at idol-sdk's fp16 points; controls drop or add one rounding point. +const std = @import("std"); +const mtl = @import("metal"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const gemm = @import("metal_idol_gemm.zig"); +const norms = @import("metal_idol_norms.zig"); +const steps = @import("metal_idol_steps.zig"); +const elems = @import("metal_idol_rows.zig"); +const attn = @import("metal_idol_attention.zig"); + +const g = id.g; +const ks = id.ks; +const Control = id.Control; +const Lib = id.Lib; + +/// The NAX fp32 GEMM (the tensor unit's fp32 path, not fp32-accurate on M5) in id_gemm32's place. +const nax_gemm32 = + \\[[max_total_threads_per_threadgroup(128)]] [[kernel]] void id_gemm32( + \\ const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], device float* y [[buffer(2)]], + \\ constant IdGemm32& p [[buffer(3)]], uint2 tg [[threadgroup_position_in_grid]], + \\ uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + \\ const int row = int(tg.y) * 64 + 32 * int(sg / 2), col = int(tg.x) * 64 + 32 * int(sg % 2); + \\ const int rows = min(32, int(p.M) - row), cols = min(32, int(p.N) - col); + \\ const short2 home = tfp::frag_home(ushort(lane)); + \\ tfp::frag acc[2][2]; + \\ tfp::gemm_block(acc, x + long(row) * p.K, w + long(col) * p.K, int(p.K), int(p.K), + \\ int(p.K), int(p.K) / 32, rows, cols, home); + \\ if (rows > 0 && cols > 0) tfp::gemm_store(acc, y + long(row) * p.ldy + col, int(p.ldy), rows, cols, home); + \\} +; + +const mma_line = "for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]);"; +const mma_half = "for (uint j = 0; j < 4; j++) { simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]); " ++ + "acc[i][j].thread_elements()[0] = id_h(acc[i][j].thread_elements()[0]); " ++ + "acc[i][j].thread_elements()[1] = id_h(acc[i][j].thread_elements()[1]); }"; + +const controls = [_]Control{ + .{ .kernel = "id_linear_gelu_01", .what = "drops fp16(x w + b) before the GELU", .edits = &.{.{ .from = "const float v = id_h(acc[i][j][e] + b);", .to = "const float v = (acc[i][j][e] + b);" }} }, + .{ .kernel = "id_linear_01", .what = "adds fp16(x w) before the bias", .edits = &.{.{ .from = "const float v = id_h(acc[i][j][e] + b);", .to = "const float v = id_h(id_h(acc[i][j][e]) + b);" }} }, + .{ .kernel = "id_norm modulated", .what = "drops fp16(y s1)", .edits = &.{.{ .from = "v = id_h(id_h(v * float(s1[i])) + float(sh[i]));", .to = "v = id_h((v * float(s1[i])) + float(sh[i]));" }} }, + .{ .kernel = "id_norm affine", .what = "adds fp16(y) before the affine", .edits = &.{.{ .from = "if (p.mode == 1) v = v * w[i] + b[i];", .to = "if (p.mode == 1) v = id_h(v) * w[i] + b[i];" }} }, + .{ .kernel = "id_norm32 modulated", .what = "adds fp16(y) before the modulation", .edits = &.{.{ .from = "yr[i] = p.mode == 2 ? v * (1.0f + w[i])", .to = "yr[i] = p.mode == 2 ? id_h(v) * (1.0f + w[i])" }} }, + .{ .kernel = "id_norm32 modulated", .what = "sums the simdgroups in fp16", .gate = .sums, .edits = &.{.{ .from = " v = simd_sum(v);\n if (lane == 0) part[sg] = v;", .to = " v = float(simd_sum(half(v)));\n if (lane == 0) part[sg] = v;" }} }, + .{ .kernel = "id_norm_split", .what = "rounds the row to fp16 before its split", .edits = &.{.{ .from = "const float v = (float(xr[i]) - mean) * inv * w[i] + b[i];", .to = "const float v = id_h((float(xr[i]) - mean) * inv * w[i] + b[i]);" }} }, + .{ .kernel = "id_repo", .what = "adds fp16(u)", .edits = &.{.{ .from = "u[i] = (a / (1.0f + metal::precise::exp(-a))) * (g[128 + i] * p.inv);", .to = "u[i] = id_h((a / (1.0f + metal::precise::exp(-a))) * (g[128 + i] * p.inv));" }} }, + .{ .kernel = "id_repo", .what = "sums d in fp16", .gate = .sums, .edits = &.{.{ .from = "s = fma(u[k], wf[t * 128 + k], s);", .to = "s = id_h(fma(u[k], wf[t * 128 + k], s));" }} }, + .{ .kernel = "id_qk", .what = "drops fp16(q0 c - q1 s)", .edits = &.{.{ .from = "const float rq0 = id_h(q0 * c - q1 * s), rq1", .to = "const float rq0 = (q0 * c - q1 * s), rq1" }} }, + .{ .kernel = "id_qk", .what = "sums q's squares in fp16", .gate = .sums, .edits = &.{.{ .from = "simd_sum(rq0 * rq0 + rq1 * rq1)", .to = "float(simd_sum(half(rq0 * rq0 + rq1 * rq1)))" }} }, + .{ .kernel = "id_gate_add", .what = "drops fp16(g y)", .edits = &.{.{ .from = "half(float(x[i]) + id_h(float(g[at.x]) * float(y[i])));", .to = "half(float(x[i]) + (float(g[at.x]) * float(y[i])));" }} }, + .{ .kernel = "id_embed", .what = "drops fp16(z w + b)", .edits = &.{.{ .from = "x[i] = half(float(half(s + b[c])) + float(pe[i]));", .to = "x[i] = half((s + b[c]) + float(pe[i]));" }} }, + .{ .kernel = "id_embed", .what = "sums z w in fp16", .gate = .sums, .edits = &.{.{ .from = "s = fma(z[r * 16 + k], w[c * 16 + k], s);", .to = "s = id_h(fma(z[r * 16 + k], w[c * 16 + k], s));" }} }, + .{ .kernel = "id_cond", .what = "adds fp16(a + ba)", .edits = &.{.{ .from = "y[i] = half((a[i] + ba[at.x]) + (c[i] + bc[at.x]));", .to = "y[i] = half(id_h(a[i] + ba[at.x]) + (c[i] + bc[at.x]));" }} }, + .{ .kernel = "id_bias32", .what = "adds fp16(b)", .edits = &.{.{ .from = "y[ulong(at.y) * width + at.x] += b[at.x];", .to = "y[ulong(at.y) * width + at.x] += id_h(b[at.x]);" }} }, + .{ .kernel = "id_gemm32", .what = "sums in fp16", .gate = .sums, .edits = &.{.{ .from = mma_line, .to = mma_half }} }, + .{ .kernel = "id_gemm32", .what = "takes the NAX fp32 path", .gate = .sums, .edits = &.{.{ .from = "void id_gemm32(", .to = "void id_gemm32_simd(" }}, .append = nax_gemm32 }, + .{ .kernel = "id_attention_00", .what = "keys past Lk unmasked", .edits = &.{.{ .from = "s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kIdMasked;", .to = "s[f][e] = s[f][e];" }} }, +}; + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + libs: std.ArrayList(*Lib) = .empty, + names: std.ArrayList([]u8) = .empty, + made: std.ArrayList(Made) = .empty, + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: *Lib, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: id.Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + for ([_]gemm.Kind{ .linear, .gelu, .split }) |kind| for (0..4) |f| { + const form: u2 = @intCast(f); + if (kind == .split and form & 1 == 0) continue; + const fun = gemm.pipeline(kind, form); + const name = (try w.name(fun)) orelse continue; + const t = try s.make(gemm.Gemm16, .{ .a = s.a, .rig = rig, .lib = lib, .kind = kind, .form = form, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + }; + if (try w.name("id_gemm32")) |name| { + const t = try s.make(gemm.Gemm32, .{ .a = s.a, .rig = rig, .lib = lib, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + const norm_kinds = [_]struct { kind: norms.Kind, mode: ks.Norm, fun: []const u8 }{ + .{ .kind = .norm, .mode = .plain, .fun = "id_norm plain" }, + .{ .kind = .norm, .mode = .affine, .fun = "id_norm affine" }, + .{ .kind = .norm, .mode = .modulated, .fun = "id_norm modulated" }, + .{ .kind = .norm32, .mode = .affine, .fun = "id_norm32 affine" }, + .{ .kind = .norm32, .mode = .modulated, .fun = "id_norm32 modulated" }, + .{ .kind = .split, .mode = .affine, .fun = "id_norm_split" }, + }; + for (norm_kinds) |nk| { + const name = (try w.name(nk.fun)) orelse continue; + const t = try s.make(norms.Norm, .{ .a = s.a, .rig = rig, .lib = lib, .kind = nk.kind, .mode = nk.mode, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("id_repo")) |name| { + const t = try s.make(steps.Repo, .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("id_qk")) |name| { + const t = try s.make(steps.Qk, .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + inline for (.{ .{ elems.GateAdd, "id_gate_add" }, .{ elems.Embed, "id_embed" }, .{ elems.Cond, "id_cond" }, .{ elems.Bias32, "id_bias32" } }) |e| { + if (try w.name(e[1])) |name| { + const t = try s.make(e[0], .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + for (attn.forms, 0..) |f, i| { + const name = (try w.name(f.name)) orelse continue; + const t = try s.make(attn.Attention, .{ .a = s.a, .rig = rig, .lib = lib, .form = i, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + if (!device.tensorUnits()) { + std.debug.print("UNSUPPORTED metal idol: the family runs on tensor units (M5) alone; its kernels unchecked here\n", .{}); + return s; + } + for ([_]?Control{null} ++ controls) |x| { + const lib = try a.create(Lib); + lib.* = .{ .a = a, .device = device, .ctl = x }; + s.libs.append(a, lib) catch |err| { + a.destroy(lib); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.libs.items) |l| { + l.deinit(); + a.destroy(l); + } + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_attention.zig b/zig/tests/contract/metal_idol_attention.zig new file mode 100644 index 0000000000..fc75fac747 --- /dev/null +++ b/zig/tests/contract/metal_idol_attention.zig @@ -0,0 +1,233 @@ +//! id_attention under the contract: softmax(q k^T / 8) v at head width 64, against attention_ref's float64 formula. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const D = id.D; +const kref = id.kref; + +const heads = id.cfg.heads; + +/// Query rows repeat at this period (prime to the 64-row blocks): float64 costs 89 queries a head over every key. +const query_period = 89; + +/// Each form's name, served lengths and edges by (Lq % 64 == 0)(Lk % 32 == 0); Lq is n, Lk is k. +pub const forms = [4]struct { name: []const u8, served: []const c.Shape, tiles: [3]c.Tile }{ + .{ + .name = "id_attention_00", + .served = &.{ .{ .rows = heads, .n = 12294, .k = 12294 }, .{ .rows = heads, .n = 4101, .k = 4101 } }, + .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 32 } }, + }, + // no served call has Lq != Lk: the cross-length forms are checked at their tails + .{ + .name = "id_attention_01", + .served = &.{ .{ .rows = heads, .n = 4101, .k = 4128 }, .{ .rows = heads, .n = 100, .k = 64 } }, + .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 64, .step = 32, .min = 32 } }, + }, + .{ + .name = "id_attention_10", + .served = &.{ .{ .rows = heads, .n = 4096, .k = 4101 }, .{ .rows = heads, .n = 128, .k = 70 } }, + .tiles = .{ .{}, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 32 } }, + }, + .{ + .name = "id_attention_11", + .served = &.{.{ .rows = heads, .n = 8192, .k = 8192 }}, + .tiles = .{ .{}, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 64, .step = 32, .min = 32 } }, + }, +}; + +pub fn form(lq: usize, lk: usize) usize { + return @as(usize, @intFromBool(lq % 64 == 0)) * 2 + @intFromBool(lk % 32 == 0); +} + +/// The scores' 64-deep sums, then the keys' sums. +fn attentionDepth(s: c.Shape) usize { + return D + s.k; +} + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + form: usize, + name: []const u8, + lq: usize = 0, + lk: usize = 0, + stress: bool = false, + q: []f16 = &.{}, + k: []f16 = &.{}, + v: []f16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: the launch's 16 heads (fixed; a gate's fewer rows run all 16); n: query rows Lq; k: key rows Lk. + pub fn spec(name: []const u8, f: usize) c.Spec { + return .{ .name = name, .accum = .f32, .depth = attentionDepth, .ops = 8, .out = .f16, .tiles = forms[f].tiles, .served = forms[f].served }; + } + + pub fn deinit(t: *Attention) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Attention) void { + t.a.free(t.q); + t.a.free(t.k); + t.a.free(t.v); + t.q = &.{}; + t.k = &.{}; + t.v = &.{}; + } + pub fn kernel(t: *Attention, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.k == 0 or form(shape.n, shape.k) != t.form) return false; + t.rig.reset(); + t.free(); + t.lq = shape.n; + t.lk = shape.k; + t.stress = fill == .stress; + const L = t.lk; + t.q = try t.a.alloc(f16, heads * t.lq * D); + t.k = try t.a.alloc(f16, heads * L * D); + t.v = try t.a.alloc(f16, heads * L * D); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores (q = k = 0), each column's values integers whose partial sums reach ~2^23 and cancel to 4 + @memset(t.q, 0); + @memset(t.k, 0); + const half = L / 2; + for (0..heads) |h| for (0..L) |j| { + const at = if (j < half) j else j - half; + const positive: f32 = if (at % 2 == 0) 127 else 4064; + const residual: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + const v: f32 = if (j < half) positive else if (j < 2 * half) residual - positive else 0; + @memset(t.v[(h * L + j) * D ..][0..D], @floatCast(v)); + }; + } else { + id.randomF16(rng, t.q, -2, 2); + id.randomF16(rng, t.k, -2, 2); + id.randomF16(rng, t.v, -1, 1); + } + if (fill == .nonfinite) { + // head 0: values NaN, +Inf, -Inf, and both infinities in one column; head 1: a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < L) { + t.v[p[0] * D + p[1]] = v; + }; + t.q[(t.lq + 0) * D + 5] = id.nan16; + if (L > 7) t.k[(L + 7) * D + 3] = id.inf16; + } + // every head's queries repeat at the period, so float64 computes each distinct one once + const p = @min(t.lq, query_period); + for (0..heads) |h| for (p..t.lq) |i| { + @memcpy(t.q[(h * t.lq + i) * D ..][0..D], t.q[(h * t.lq + i % p) * D ..][0..D]); + }; + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.q)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.k)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.v)); + t.out = try t.rig.output(heads * t.lq * D * 2); + return true; + } + + pub fn run(t: *Attention) !void { + try t.lib.need(forms[t.form].name); + const p = t.rig.begin(); + t.lib.k.attention(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.lq, t.lk); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return heads * t.lq * D; + } + + /// attention_ref.head's formula, one pass over the keys a distinct query; O rows [Lq, heads * 64]. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.lk; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + const kf = try t.a.alloc(f64, L * D); + defer t.a.free(kf); + const vf = try t.a.alloc(f64, L * D); + defer t.a.free(vf); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + var qf: [D]f64 = undefined; + var out: [D]f64 = undefined; + var tol: [D]f64 = undefined; + const p = @min(t.lq, query_period); + for (0..heads) |h| { + for (kf, vf, t.k[h * L * D ..][0 .. L * D], t.v[h * L * D ..][0 .. L * D]) |*kd, *vd, kb, vb| { + kd.* = kb; + vd.* = vb; + } + for (0..p) |i| { + var zero = true; + for (&qf, t.q[(h * t.lq + i) * D ..][0..D]) |*x, b| { + x.* = b; + zero = zero and x.* == 0; + } + var m: f64 = -std.math.inf(f64); + var nan = false; + for (s, 0..) |*sj, j| { + var dot: f64 = 0; + if (!zero) { + const Vec = @Vector(8, f64); + var acc: Vec = @splat(0); + var e: usize = 0; + while (e < D) : (e += 8) acc += @as(Vec, qf[e..][0..8].*) * @as(Vec, kf[j * D + e ..][0..8].*); + dot = @reduce(.Add, acc); + } + sj.* = dot * scale; + nan = nan or std.math.isNan(sj.*); + m = @max(m, sj.*); + } + if (nan) m = std.math.nan(f64); + var l: f64 = 0; + @memset(&out, 0); + @memset(&tol, 0); + for (s, 0..) |sj, j| { + const pj = @exp(sj - m); + l += pj; + for (&out, &tol, vf[j * D ..][0..D]) |*o, *tl, vv| { + o.* += pj * vv; + tl.* += pj * @abs(vv); + } + } + var r = i; + while (r < t.lq) : (r += p) { + const at = (r * heads + h) * D; + for (out, tol, 0..) |o, tl, e| { + want[at + e] = o / l; + // P enters P x V at the tensor unit's fp32 (~2^-10 relative); O rounds to fp16 + mag[at + e] = tl / l * 0x1p-9 + @abs(o / l) * 0x1p-10 + 1e-6; + } + } + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, tolerance: []const f64, allow: []f64) !void { + for (want, tolerance, allow) |w, tl, *l| l.* = if (t.stress) c.boundExact(spec(t.name, t.form), w) else if (std.math.isFinite(w)) tl else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.lk; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..heads) |h| for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = t.v[(h * L + j) * D + e]; + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..t.lq) |i| out[(i * heads + h) * D + e] = o; + }; + } +}; diff --git a/zig/tests/contract/metal_idol_common.zig b/zig/tests/contract/metal_idol_common.zig new file mode 100644 index 0000000000..222b6b22b4 --- /dev/null +++ b/zig/tests/contract/metal_idol_common.zig @@ -0,0 +1,245 @@ +//! Idol's kernel tests' shared parts: float64 values at fp16 points, their ranges, inputs and the family's library. +const std = @import("std"); +const mtl = @import("metal"); +pub const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +pub const c = contract.spec; +pub const g = contract.gates; +pub const rows = contract.rows; +pub const ks = tf.idol.kernels; +pub const cfg = tf.idol.config; +pub const H = cfg.hidden; +pub const D = cfg.head_dim; + +pub fn noSums(_: c.Shape) usize { + return 0; +} + +/// A value rounded once to fp16 (the kernels' half(fp32) of a value float64 holds exactly, or float64's own). +pub fn hf(v: f64) f64 { + return @as(f16, @floatCast(v)); +} + +pub fn kref(r: rig_.Ref) ks.Ref { + return .{ .buf = r.buf, .off = r.off }; +} + +/// Every value the kernel may hold at a step: its fp32 steps widen both ends, its fp16 points round them. +pub const Range = struct { + lo: f64, + hi: f64, + pub fn of(v: f64) Range { + return .{ .lo = v, .hi = v }; + } + /// An fp32 step good to `ulps` units in the last place (a correctly rounded one: 0.5). + pub fn fp32(a: Range, ulps: f64) Range { + const u = ulps * 0x1p-23; + return .{ .lo = a.lo - u * @abs(a.lo) - 0x1p-149, .hi = a.hi + u * @abs(a.hi) + 0x1p-149 }; + } + pub fn fp16(a: Range) Range { + return .{ .lo = hf(a.lo), .hi = hf(a.hi) }; + } + pub fn add(a: Range, b: Range) Range { + return .{ .lo = a.lo + b.lo, .hi = a.hi + b.hi }; + } + pub fn neg(a: Range) Range { + return .{ .lo = -a.hi, .hi = -a.lo }; + } + pub fn sub(a: Range, b: Range) Range { + return a.add(b.neg()); + } + pub fn widen(a: Range, e: f64) Range { + return .{ .lo = a.lo - e, .hi = a.hi + e }; + } + pub fn mul(a: Range, b: Range) Range { + const p = [4]f64{ a.lo * b.lo, a.lo * b.hi, a.hi * b.lo, a.hi * b.hi }; + return .{ .lo = @min(@min(p[0], p[1]), @min(p[2], p[3])), .hi = @max(@max(p[0], p[1]), @max(p[2], p[3])) }; + } + /// An increasing function at both ends. + pub fn up(a: Range, comptime f: fn (f64) f64) Range { + return .{ .lo = f(a.lo), .hi = f(a.hi) }; + } + /// 1 / sqrt(a), a positive. + pub fn rsqrt(a: Range) Range { + return .{ .lo = 1 / @sqrt(a.hi), .hi = 1 / @sqrt(a.lo) }; + } + /// 1 / a, a positive. + pub fn recip(a: Range) Range { + return .{ .lo = 1 / a.hi, .hi = 1 / a.lo }; + } + /// a^2 over the range. + pub fn square(a: Range) Range { + const m = @max(a.lo * a.lo, a.hi * a.hi); + return .{ .lo = if (a.lo <= 0 and a.hi >= 0) 0 else @min(a.lo * a.lo, a.hi * a.hi), .hi = m }; + } +}; + +pub fn tanh(x: f64) f64 { + return std.math.tanh(x); +} + +pub fn sqrt(x: f64) f64 { + return @sqrt(x); +} + +// Metal's error bounds with the engine's options (fast math on, the kernels' own `precise::` calls). +pub const div_ulps = 2.5; +pub const rsqrt_ulps = 2; +pub const sqrt_ulps = 3; +pub const exp_ulps = 4; +pub const tanh_ulps = 5; +pub const trig_ulps = 4; + +/// f over [lo, hi] (cos or sin): its ends, and 1 or -1 where an extreme (a multiple of pi past `at`) falls inside. +fn trigOver(lo: f64, hi: f64, comptime f: fn (f64) f64, at: f64) Range { + const pi = std.math.pi; + if (!(hi - lo < 2 * pi)) return .{ .lo = -1, .hi = 1 }; + var out = Range{ .lo = @min(f(lo), f(hi)), .hi = @max(f(lo), f(hi)) }; + var k = @ceil((lo - at) / pi); + while (k * pi + at <= hi) : (k += 1) { + if (@mod(k, 2) == 0) out.hi = 1 else out.lo = -1; + } + return out; +} + +fn cos(x: f64) f64 { + return @cos(x); +} + +fn sin(x: f64) f64 { + return @sin(x); +} + +pub fn cosOver(lo: f64, hi: f64) Range { + return trigOver(lo, hi, cos, 0); +} + +pub fn sinOver(lo: f64, hi: f64) Range { + return trigOver(lo, hi, sin, std.math.pi / 2.0); +} + +/// Each output's float64 value at the kernel's rounding points, and its range, kept for the bounds and intervals. +pub const Ends = struct { + a: std.mem.Allocator, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + + pub fn deinit(e: *Ends) void { + e.a.free(e.lo); + e.a.free(e.hi); + e.lo = &.{}; + e.hi = &.{}; + } + + pub fn size(e: *Ends, n: usize) !void { + e.deinit(); + e.lo = try e.a.alloc(f64, n); + e.hi = try e.a.alloc(f64, n); + } + + pub fn put(e: *Ends, want: []f64, mag: []f64, i: usize, w: f64, r: Range) void { + want[i] = w; + mag[i] = if (std.math.isFinite(w)) @abs(w) else 0; + e.lo[i] = r.lo; + e.hi[i] = r.hi; + } + + /// The allowance about the float64 value: its range's farther end (a non-finite value is checked by class). + pub fn bounds(e: *const Ends, want: []const f64, allow: []f64) void { + for (allow, want, e.lo, e.hi) |*l, w, lo, hi| l.* = if (std.math.isFinite(w)) @max(@abs(w - lo), @abs(hi - w)) else 0; + } + + pub fn intervals(e: *const Ends, low: []f64, high: []f64) void { + @memcpy(low, e.lo); + @memcpy(high, e.hi); + } +}; + +pub fn randomF16(rng: std.Random, out: []f16, lo: f32, hi: f32) void { + for (out) |*v| v.* = @floatCast(lo + rng.float(f32) * (hi - lo)); +} + +pub fn randomF32(rng: std.Random, out: []f32, lo: f32, hi: f32) void { + for (out) |*v| v.* = lo + rng.float(f32) * (hi - lo); +} + +/// An integer in [lo, hi]. +pub fn int(rng: std.Random, lo: i32, hi: i32) f32 { + return @floatFromInt(lo + @as(i32, @intCast(rng.uintLessThan(u32, @intCast(hi - lo + 1))))); +} + +/// +1 or -1. +pub fn sign(rng: std.Random) f32 { + return if (rng.boolean()) 1 else -1; +} + +pub const nan16: f16 = std.math.nan(f16); +pub const inf16: f16 = std.math.inf(f16); +pub const nan32: f32 = std.math.nan(f32); +pub const inf32: f32 = std.math.inf(f32); + +/// [count, width] rows of which row r is `distinct`'s row r % (its rows): a period prime to every tile. +pub fn periodic(comptime T: type, a: std.mem.Allocator, distinct: []const T, width: usize, count: usize) ![]T { + const p = distinct.len / width; + const out = try a.alloc(T, count * width); + for (0..count) |r| @memcpy(out[r * width ..][0..width], distinct[(r % p) * width ..][0..width]); + return out; +} + +/// One text change a control makes to the family's source. +pub const Edit = struct { from: []const u8, to: []const u8 }; + +/// A control: the kernel it breaks, the gate that must go red, its source edits and any text it appends. +pub const Control = struct { kernel: []const u8, what: []const u8, gate: g.Gate = .shapes, edits: []const Edit = &.{}, append: []const u8 = "" }; + +/// The family's source as the engine compiles it (or a control's copy), its pipelines made as tests ask for them. +pub const Lib = struct { + a: std.mem.Allocator, + device: mtl.Device, + ctl: ?Control, + k: ks.Kernels = undefined, + built: bool = false, + made: [ks.names.len]bool = @splat(false), + + fn text(l: *Lib) ![]u8 { + var out = try tf.frags.source(l.device, l.a, sources.idol); + const x = l.ctl orelse return out; + for (x.edits) |e| { + errdefer l.a.free(out); + if (std.mem.count(u8, out, e.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, e.from }); + return error.ControlTextNotFoundOnce; + } + const next = try std.mem.replaceOwned(u8, l.a, out, e.from, e.to); + l.a.free(out); + out = next; + } + if (x.append.len == 0) return out; + defer l.a.free(out); + return std.mem.concat(l.a, u8, &.{ out, "\n", x.append }); + } + + /// The pipeline `name` made on this library (compiled on first use). + pub fn need(l: *Lib, name: []const u8) !void { + if (!l.built) { + const src = try l.text(); + defer l.a.free(src); + l.k.own = try mtl.Library.fromSource(l.device, src, mtl.CompileOptions.mlx()); + l.built = true; + } + for (ks.names, 0..) |n, i| if (std.mem.eql(u8, n, name)) { + if (!l.made[i]) l.k.p[i] = try mtl.Pipeline.init(l.device, l.k.own, n, false); + l.made[i] = true; + return; + }; + return error.NoSuchPipeline; + } + + pub fn deinit(l: *Lib) void { + for (l.k.p, l.made) |p, m| if (m) p.deinit(); + if (l.built) l.k.own.deinit(); + } +}; diff --git a/zig/tests/contract/metal_idol_gemm.zig b/zig/tests/contract/metal_idol_gemm.zig new file mode 100644 index 0000000000..7cb64720f4 --- /dev/null +++ b/zig/tests/contract/metal_idol_gemm.zig @@ -0,0 +1,430 @@ +//! id_linear (plain and GELU), id_split and id_gemm32 under the contract: products against float64, sums proven fp32. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const H = id.H; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +/// Rows of x and of w (with the bias) repeat at these periods, prime to every tile: 97 x 131 dots a shape. +const row_period = 97; +const col_period = 131; + +pub const Kind = enum { linear, gelu, split }; + +/// An output's float64 value and its range. +const Out = struct { w: f64, r: Range }; + +/// The M5 tensor unit's fp32 sums: 36 roundings of 2^-24 of |acc| + sum |a b| a 16-deep step (attention_range's). +fn naxDepth(s: c.Shape) usize { + return s.k / 16 * contract.attention_range.mma_step + 8; +} + +/// float64's dot of an fp16 row pair and the tensor unit's allowance on its fp32 sums, step by 16-deep step. +fn naxDot(x: []const f16, w: []const f16) [3]f64 { + var s: f64 = 0; + var spread: f64 = 0; + var mag: f64 = 0; + var i: usize = 0; + while (i < x.len) : (i += 16) { + var t: f64 = 0; + var abs: f64 = 0; + for (x[i..@min(i + 16, x.len)], w[i..@min(i + 16, x.len)]) |a, b| { + const p = @as(f64, a) * @as(f64, b); + t += p; + abs += @abs(p); + } + spread += @abs(s) + abs; + s += t; + mag += abs; + } + return .{ s, spread * contract.attention_range.mma_step * 0x1p-24, mag }; +} + +const gelu_c: f64 = @as(f32, 0.7978845608028654); +const gelu_k: f64 = @as(f32, 0.044715); + +fn inner(x: f64) f64 { + return gelu_c * (x + gelu_k * x * x * x); +} + +/// torch's tanh GELU in float64 at the kernel's fp32 constants. +fn gelu(x: f64) f64 { + return 0.5 * x * (1 + std.math.tanh(inner(x))); +} + +/// The GELU's fp32 steps over a range of fp16 inputs: its inner polynomial rises, tanh rises, the rest as products. +fn geluRange(v: Range) Range { + const t = (Range{ .lo = inner(v.lo), .hi = inner(v.hi) }).fp32(3).up(id.tanh).fp32(id.tanh_ulps); + return (Range{ .lo = 0.5 * v.lo, .hi = 0.5 * v.hi }).mul(t.add(Range.of(1)).fp32(0.5)).fp32(0.5); +} + +const Served = struct { served: []const c.Shape, tiles: [3]c.Tile }; + +/// Forms by (M % 64 == 0)(N % 128 == 0); rows: 64-row tiles in whole-row forms, a last tile's rows first in the rest. +fn envelope(kind: Kind, form: u2) Served { + const row_tail: c.Tile = .{ .size = 64 }; + const row_whole: c.Tile = .{ .size = 2, .step = 1, .min = 1 }; + const col_tail: c.Tile = .{ .size = 128 }; + const col_whole: c.Tile = .{ .size = 256, .step = 128, .min = 128 }; + const k_tile: c.Tile = .{ .size = 256, .step = 32, .min = 32 }; + const rt = if (form & 2 != 0) row_whole else row_tail; + if (kind == .split) return .{ + .served = if (form & 2 != 0) &.{.{ .rows = 8192 / 64, .n = 256, .k = 3 * H }} else &.{ .{ .rows = 6, .n = 256, .k = 3 * H }, .{ .rows = 12294, .n = 256, .k = 3 * H }, .{ .rows = 4101, .n = 256, .k = 3 * H } }, + .tiles = .{ rt, .{}, .{} }, + }; + const ct = if (form & 1 != 0) col_whole else col_tail; + const tiles: [3]c.Tile = .{ rt, ct, k_tile }; + if (kind == .gelu) return .{ + .tiles = tiles, + .served = switch (form) { + 3 => &.{.{ .rows = 8192 / 64, .n = 4 * H, .k = H }}, + 1 => &.{ .{ .rows = 6, .n = 4 * H, .k = H }, .{ .rows = 12294, .n = 4 * H, .k = H }, .{ .rows = 4101, .n = 4 * H, .k = H } }, + // no served call takes a width short of 128 columns: these forms are checked at one + 2 => &.{.{ .rows = 4096 / 64, .n = 1000, .k = H }}, + 0 => &.{ .{ .rows = 5, .n = 1000, .k = H }, .{ .rows = 4101, .n = 1000, .k = H } }, + }, + }; + return .{ .tiles = tiles, .served = switch (form) { + 3 => &.{ .{ .rows = 8192 / 64, .n = 3 * H, .k = H }, .{ .rows = 8192 / 64, .n = H, .k = H }, .{ .rows = 8192 / 64, .n = H, .k = 4 * H } }, + 1 => &.{ .{ .rows = 6, .n = 3 * H, .k = H }, .{ .rows = 12294, .n = 3 * H, .k = H }, .{ .rows = 12294, .n = H, .k = H }, .{ .rows = 12294, .n = H, .k = 4 * H }, .{ .rows = 4101, .n = 3 * H, .k = H }, .{ .rows = 4101, .n = H, .k = H }, .{ .rows = 4101, .n = H, .k = 4 * H } }, + 2 => &.{.{ .rows = 4096 / 64, .n = 1000, .k = H }}, + 0 => &.{ .{ .rows = 5, .n = 1000, .k = H }, .{ .rows = 4101, .n = 1000, .k = H } }, + } }; +} + +/// The pipeline a kind and form run. +pub fn pipeline(kind: Kind, form: u2) []const u8 { + const names = [3][4][]const u8{ + .{ "id_linear_00", "id_linear_01", "id_linear_10", "id_linear_11" }, + .{ "id_linear_gelu_00", "id_linear_gelu_01", "id_linear_gelu_10", "id_linear_gelu_11" }, + .{ "id_split_0", "id_split_0", "id_split_1", "id_split_1" }, + }; + return names[@backingInt(kind)][form]; +} + +pub const Gemm16 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + kind: Kind, + /// split's forms are 1 and 3: its N, 256, is whole 128-column tiles. + form: u2, + name: []const u8, + m: usize = 0, + n: usize = 0, + k: usize = 0, + stress: bool = false, + x: []f16 = &.{}, + w: []f16 = &.{}, + b: []f16 = &.{}, + out: rig_.Ref = undefined, + ins: [3]rig_.Ref = undefined, + ends: Ends, + + pub fn spec(name: []const u8, kind: Kind, form: u2) c.Spec { + const e = envelope(kind, form); + return .{ + .name = name, + .accum = .f32, + .depth = naxDepth, + .ops = 1, + .out = if (kind == .split) .f32 else .f16, + .tiles = e.tiles, + .served = e.served, + }; + } + + pub fn deinit(t: *Gemm16) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Gemm16) void { + t.a.free(t.x); + t.a.free(t.w); + t.a.free(t.b); + t.x = &.{}; + t.w = &.{}; + t.b = &.{}; + } + pub fn kernel(t: *Gemm16, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.kind, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Gemm16); + + pub fn prepare(t: *Gemm16, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const m = if (t.form & 2 != 0) shape.rows * 64 else shape.rows; + if (m == 0 or shape.n == 0 or shape.k == 0 or shape.k % 2 != 0) return false; + if (t.kind == .split and (shape.n != 256 or shape.k != 3 * H)) return false; + const form = @as(u2, @intFromBool(m % 64 == 0)) * 2 + @intFromBool(shape.n % 128 == 0); + if (form != t.form) return false; + t.rig.reset(); + t.free(); + t.m = m; + t.n = shape.n; + t.k = shape.k; + t.stress = fill == .stress; + const K = shape.k; + const px = @min(m, row_period); + const pw = @min(shape.n, col_period); + t.x = try t.a.alloc(f16, px * K); + t.w = try t.a.alloc(f16, pw * K); + t.b = try t.a.alloc(f16, pw); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // integers in [hi / 2, hi), hi^2 K <= 2^24: fp32 sums exact in any order, a 32nd of a half past fp16's + const hi: i32 = @intFromFloat(@sqrt(@as(f64, 0x1p24) / @as(f64, @floatFromInt(K)))); + for (0..px) |r| for (0..K / 2) |i| { + const v = id.int(rng, @divTrunc(hi, 2), hi - 1); + t.x[r * K + i] = @floatCast(v); + // x's halves differ by 1 in their first 8: the sums cancel to a few hundred + t.x[r * K + K / 2 + i] = @floatCast(v - @as(f32, if (i < 8) 1 else 0)); + }; + for (0..pw) |j| for (0..K / 2) |i| { + const v = id.int(rng, @divTrunc(hi, 2), hi - 1); + t.w[j * K + i] = @floatCast(v); + t.w[j * K + K / 2 + i] = @floatCast(-v); + }; + for (t.b) |*v| v.* = @floatCast(id.int(rng, -4, 4)); + } else { + id.randomF16(rng, t.x, -2, 2); + id.randomF16(rng, t.w, -0.25, 0.25); + id.randomF16(rng, t.b, -1, 1); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, K / 3 }, .{ 1, K / 2 }, .{ 2, K - 1 }, .{ 3, 0 }, .{ 3, K - 1 } }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < px) { + t.x[p[0] * K + p[1]] = v; + }; + if (pw > 2) t.b[2] = id.inf16; + } + const xs = try id.periodic(f16, t.a, t.x, K, m); + defer t.a.free(xs); + const ws = try id.periodic(f16, t.a, t.w, K, shape.n); + defer t.a.free(ws); + const bs = try id.periodic(f16, t.a, t.b, 1, shape.n); + defer t.a.free(bs); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(xs)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(ws)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(bs)); + t.out = try t.rig.output(m * shape.n * @as(usize, if (t.kind == .split) 4 else 2)); + return true; + } + + pub fn run(t: *Gemm16) !void { + try t.lib.need(pipeline(t.kind, t.form)); + const p = t.rig.begin(); + const k = &t.lib.k; + switch (t.kind) { + .split => k.split(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.m), + else => k.linear(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.m, t.n, t.k, t.kind == .gelu), + } + try t.rig.finish(p); + } + + pub fn count(t: *Gemm16) usize { + return t.m * t.n; + } + + /// The output of a sum `s` (range `sr`) of column j: its bias, fp16 rounding and GELU; split's fp32 sum as is. + fn finish(t: *const Gemm16, s: f64, sr: Range, j: usize) Out { + if (t.kind == .split) return .{ .w = s, .r = sr.fp32(0.5) }; + const b: f64 = t.b[j]; + const v = hf(s + b); + const vr = sr.add(Range.of(b)).fp32(0.5).fp16(); + if (t.kind == .linear) return .{ .w = v, .r = vr }; + return .{ .w = hf(gelu(v)), .r = geluRange(vr).fp16() }; + } + + /// The distinct outputs' values and ranges by `value(r, j)`, then every output from its row's and column's. + fn table(t: *Gemm16, ends: *Ends, want: []f64, mag: []f64, ctx: anytype, comptime value: anytype) !void { + const px = t.x.len / t.k; + const pw = t.w.len / t.k; + const tw = try t.a.alloc(f64, 3 * px * pw); + defer t.a.free(tw); + for (0..px) |r| for (0..pw) |j| { + const o = value(ctx, t, r, j); + tw[3 * (r * pw + j) ..][0..3].* = .{ o.w, o.r.lo, o.r.hi }; + }; + for (0..t.m) |r| for (0..t.n) |j| { + const e = tw[3 * ((r % px) * pw + j % pw) ..][0..3]; + ends.put(want, mag, r * t.n + j, e[0], .{ .lo = e[1], .hi = e[2] }); + }; + } + + pub fn reference(t: *Gemm16, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const f = struct { + fn value(_: void, s: *const Gemm16, r: usize, j: usize) Out { + const d = naxDot(s.x[r * s.k ..][0..s.k], s.w[j * s.k ..][0..s.k]); + return s.finish(d[0], if (s.stress) Range.of(d[0]) else Range.of(d[0]).widen(d[1]), j); + } + }; + try t.table(&t.ends, want, mag, {}, f.value); + } + pub fn bounds(t: *Gemm16, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Gemm16, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Gemm16, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + const terms = try t.a.alloc(f64, t.k); + defer t.a.free(terms); + const Ctx = struct { acc: c.Accum, order: g.Order, terms: []f64 }; + const f = struct { + fn value(x: Ctx, s: *const Gemm16, r: usize, j: usize) Out { + for (x.terms, s.x[r * s.k ..][0..s.k], s.w[j * s.k ..][0..s.k]) |*v, a, b| v.* = @as(f64, a) * @as(f64, b); + const sum = id.rows.sumIn(x.acc, x.order, x.terms); + return s.finish(sum, Range.of(sum), j); + } + }; + try t.table(&scratch, out, mag, Ctx{ .acc = acc, .order = order, .terms = terms }, f.value); + } +}; + +// --- id_gemm32: fp32 x [M, K] times w [N, K]^T, sums that must be fp32's (the NAX fp32 path's are not) ----------- + +/// Each 8x8 step's eight FMAs in one chain over K. +fn gemm32Depth(s: c.Shape) usize { + return s.k + 8; +} + +pub const Gemm32 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = gemm32Depth, + .out = .f32, + .tiles = .{ .{ .size = 64 }, .{ .size = 64, .step = 8, .min = 8 }, .{ .size = 32, .step = 32, .min = 32 } }, + .served = &.{ .{ .rows = 4101, .n = H, .k = 1280 }, .{ .rows = 4101, .n = H, .k = 128 }, .{ .rows = 8192, .n = 16, .k = H } }, + }; + } + + pub fn deinit(t: *Gemm32) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Gemm32) void { + t.a.free(t.x); + t.a.free(t.w); + t.x = &.{}; + t.w = &.{}; + } + pub fn kernel(t: *Gemm32, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Gemm32); + + pub fn prepare(t: *Gemm32, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const n_ok = shape.n % 64 == 0 or (shape.n < 64 and shape.n % 8 == 0); + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 32 != 0 or !n_ok) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const K = shape.k; + const px = @min(shape.rows, row_period); + const pw = @min(shape.n, col_period); + t.x = try t.a.alloc(f32, px * K); + t.w = try t.a.alloc(f32, pw * K); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // 13-bit integers against +-1, in turn on either side: exact fp32 products and sums below 2^24 + for (0..px) |r| for (0..K) |i| { + t.x[r * K + i] = id.sign(rng) * if (i % 2 == 0) id.int(rng, 4097, 8191) else 1; + }; + for (0..pw) |j| for (0..K) |i| { + t.w[j * K + i] = id.sign(rng) * if (i % 2 == 1) id.int(rng, 4097, 8191) else 1; + }; + } else { + id.randomF32(rng, t.x, -2, 2); + id.randomF32(rng, t.w, -0.25, 0.25); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, K / 3 }, .{ 1, K / 2 }, .{ 2, K - 1 }, .{ 3, 0 }, .{ 3, K - 1 } }; + for (at, [_]f32{ id.nan32, id.inf32, -id.inf32, id.inf32, -id.inf32 }) |p, v| if (p[0] < px) { + t.x[p[0] * K + p[1]] = v; + }; + } + const xs = try id.periodic(f32, t.a, t.x, K, shape.rows); + defer t.a.free(xs); + const ws = try id.periodic(f32, t.a, t.w, K, shape.n); + defer t.a.free(ws); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(xs)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(ws)); + t.out = try t.rig.output(shape.rows * shape.n * 4); + return true; + } + + pub fn run(t: *Gemm32) !void { + try t.lib.need("id_gemm32"); + const p = t.rig.begin(); + t.lib.k.linear32(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.shape.rows, t.shape.n, t.shape.k, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Gemm32) usize { + return t.shape.rows * t.shape.n; + } + + fn each(t: *Gemm32, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) !void { + const K = t.shape.k; + const px = t.x.len / K; + const pw = t.w.len / K; + const tw = try t.a.alloc(f64, 2 * px * pw); + defer t.a.free(tw); + const terms = try t.a.alloc(f64, K); + defer t.a.free(terms); + for (0..px) |r| for (0..pw) |j| { + var s: f64 = 0; + var m: f64 = 0; + for (terms, t.x[r * K ..][0..K], t.w[j * K ..][0..K]) |*v, a, b| { + v.* = @as(f64, a) * @as(f64, b); + s += v.*; + m += @abs(v.*); + } + tw[2 * (r * pw + j) ..][0..2].* = .{ if (acc) |x| id.rows.sumIn(x, order, terms) else s, m }; + }; + for (0..t.shape.rows) |r| for (0..t.shape.n) |j| { + const e = tw[2 * ((r % px) * pw + j % pw) ..][0..2]; + want[r * t.shape.n + j] = e[0]; + mag[r * t.shape.n + j] = e[1]; + }; + } + + pub fn reference(t: *Gemm32, want: []f64, mag: []f64) !void { + try t.each(want, mag, null, .sequential); + } + + pub fn simulate(t: *Gemm32, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + try t.each(out, mag, acc, order); + } +}; diff --git a/zig/tests/contract/metal_idol_norms.zig b/zig/tests/contract/metal_idol_norms.zig new file mode 100644 index 0000000000..c9ea523659 --- /dev/null +++ b/zig/tests/contract/metal_idol_norms.zig @@ -0,0 +1,251 @@ +//! id_norm, id_norm32 and id_norm_split under the contract: LayerNorm32 in fp32, at idol-sdk's fp16 points. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const ks = id.ks; +const cfg = id.cfg; +const H = id.H; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +pub const Kind = enum { norm, norm32, split }; + +const norm_why = "an fp16 output follows the sums as var^-1/2: a bf16 sum's error is a few output ulps"; + +/// A thread's 4 values of the 1,024 in turn, the simdgroup tree, then the tree over the threadgroup's 8 simdgroups. +const sum_depth = 1024 / 256 + 5 + 5; + +fn normDepth(_: c.Shape) usize { + return sum_depth + 3; +} + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + kind: Kind, + mode: ks.Norm, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + x: []f16 = &.{}, + /// w then b, fp32 + wb: []f32 = &.{}, + /// the modulation's 1 + scale, then its shift, fp16 + mod: []f16 = &.{}, + ins: [5]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n and k: the width, 1,024 (the launcher's). + pub fn spec(name: []const u8, kind: Kind, mode: ks.Norm) c.Spec { + const all: []const c.Shape = &.{ .{ .rows = 12294, .n = H, .k = H }, .{ .rows = 8192, .n = H, .k = H }, .{ .rows = 4101, .n = H, .k = H } }; + return .{ + .name = name, + .accum = .f32, + .depth = normDepth, + .ops = 2 * (id.div_ulps + id.rsqrt_ulps) + 3, + .out = if (kind == .norm32) .f32 else .f16, + .served = switch (kind) { + .norm => switch (mode) { + .modulated => all[0..2], + .affine => all[2..], + // no block calls the plain form: checked at every stream's length + .plain => all, + }, + // the output head's modulated LayerNorm; the affine form at the same rows + .norm32 => all[1..2], + .split => all, + }, + .sums = if (kind == .norm) .{ .bf16 = 0, .why = norm_why } else .{}, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + t.a.free(t.x); + t.a.free(t.wb); + t.a.free(t.mod); + t.x = &.{}; + t.wb = &.{}; + t.mod = &.{}; + } + pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.kind, t.mode), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Norm); + + fn eps(t: *const Norm) f32 { + return if (t.kind == .norm) cfg.block_eps else cfg.default_eps; + } + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H or shape.k != H) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.stress = fill == .stress; + t.x = try t.a.alloc(f16, shape.rows * H); + t.wb = try t.a.alloc(f32, 2 * H); + t.mod = try t.a.alloc(f16, 2 * H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // integers mirrored in each row's halves: a mean of exactly 0 and squares summing to ~2^22, exact in fp32 + for (0..shape.rows) |r| for (0..H / 2) |i| { + const v = id.int(rng, -64, 64); + t.x[r * H + i] = @floatCast(v); + t.x[r * H + H / 2 + i] = @floatCast(-v); + }; + } else id.randomF16(rng, t.x, -3, 3); + id.randomF32(rng, t.wb[0..H], 0.25, 2); + id.randomF32(rng, t.wb[H..], -1, 1); + id.randomF16(rng, t.mod[0..H], 0.5, 1.5); + id.randomF16(rng, t.mod[H..], -1, 1); + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, H / 3 }, .{ 1, H / 2 }, .{ 2, H - 1 }, .{ 3, 0 }, .{ 3, H - 1 } }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < shape.rows) { + t.x[p[0] * H + p[1]] = v; + }; + t.wb[5] = id.nan32; + t.wb[H + 6] = id.inf32; + t.mod[7] = -id.inf16; + t.mod[H + 8] = id.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wb[0..H])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.wb[H..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.mod[0..H])); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.mod[H..])); + t.out = try t.rig.output(t.count() * @as(usize, if (t.kind == .norm32) 4 else 2)); + return true; + } + + pub fn run(t: *Norm) !void { + const k = &t.lib.k; + try t.lib.need(switch (t.kind) { + .norm => "id_norm", + .norm32 => "id_norm32", + .split => "id_norm_split", + }); + const p = t.rig.begin(); + const in = t.ins; + switch (t.kind) { + .norm => k.norm(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(in[3]), kref(in[4]), kref(t.out), t.shape.rows, t.eps(), t.mode), + .norm32 => k.norm32(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(t.out), t.shape.rows, t.eps(), t.mode), + .split => k.normSplit(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(t.out), t.shape.rows, t.eps()), + } + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * H * @as(usize, if (t.kind == .split) 3 else 1); + } + + /// Row r's outputs from the kernel's mean (range `mr`) and sum of squared deviations (range `qr`). + fn row(t: *Norm, ends: *Ends, r: usize, mean: f64, mr: Range, q: f64, qr: Range, want: []f64, mag: []f64) void { + const n: f64 = H; + const e: f64 = t.eps(); + const inv = 1 / @sqrt(q / n + e); + const ir = (Range{ .lo = qr.lo / n, .hi = qr.hi / n }).add(Range.of(e)).fp32(0.5).rsqrt().fp32(id.rsqrt_ulps); + for (0..H) |i| { + const x: f64 = t.x[r * H + i]; + const v = (x - mean) * inv; + const vr = Range.of(x).sub(mr).fp32(0.5).mul(ir).fp32(0.5); + const w: f64 = t.wb[i]; + const b: f64 = t.wb[H + i]; + const s1: f64 = t.mod[i]; + const sh: f64 = t.mod[H + i]; + const at = r * H + i; + switch (t.kind) { + .norm => switch (t.mode) { + .plain => ends.put(want, mag, at, hf(v), vr.fp16()), + .affine => ends.put(want, mag, at, hf(v * w + b), vr.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5).fp16()), + .modulated => { + const y = vr.fp16().mul(Range.of(s1)).fp32(0.5).fp16().add(Range.of(sh)).fp32(0.5).fp16(); + ends.put(want, mag, at, hf(hf(hf(v) * s1) + sh), y); + }, + }, + .norm32 => switch (t.mode) { + .modulated => ends.put(want, mag, at, v * (1 + w) + b, vr.mul(Range.of(1 + w).fp32(0.5)).fp32(0.5).add(Range.of(b)).fp32(0.5)), + else => ends.put(want, mag, at, v * w + b, vr.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5)), + }, + .split => { + const y = v * w + b; + const yr = vr.mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5); + const hi = hf(y); + const hr = yr.fp16(); + // hi + lo: lo = fp16(y - hi), y - hi exact in fp32 for whichever fp16 hi the range rounds to + const lr = (Range{ .lo = yr.lo - hr.hi, .hi = yr.hi - hr.lo }).fp16(); + const base = r * 3 * H + i; + ends.put(want, mag, base, hi, hr); + ends.put(want, mag, base + H, hi, hr); + ends.put(want, mag, base + 2 * H, hf(y - hi), lr); + }, + } + } + } + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const n: f64 = H; + const gs = c.gamma(sum_depth, c.unit(.f32)); + const gq = c.gamma(sum_depth + 3, c.unit(.f32)); + for (0..t.shape.rows) |r| { + const xs = t.x[r * H ..][0..H]; + var s: f64 = 0; + var abs: f64 = 0; + for (xs) |v| { + s += v; + abs += @abs(@as(f64, v)); + } + const mean = s / n; + var q: f64 = 0; + for (xs) |v| q += (v - mean) * (v - mean); + if (t.stress) { + // a mean of 0 exactly, the deviations the integers and their squares' sums exact in fp32 + t.row(&t.ends, r, mean, Range.of(mean), q, Range.of(q), want, mag); + continue; + } + const mr = Range{ .lo = (s - gs * abs) / n, .hi = (s + gs * abs) / n }; + // the squares about the kernel's mean exceed float64's by n (m' - mean)^2, then fp32 rounds them + const dm = @max(@abs(mr.lo - mean), @abs(mr.hi - mean)); + const qr = Range{ .lo = q * (1 - gq), .hi = (q + n * dm * dm) * (1 + gq) }; + t.row(&t.ends, r, mean, mr, q, qr, want, mag); + } + } + pub fn bounds(t: *Norm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Norm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Norm, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var terms: [H]f64 = undefined; + for (0..t.shape.rows) |r| { + const xs = t.x[r * H ..][0..H]; + for (&terms, xs) |*v, x| v.* = x; + const mean = contract.rows.sumIn(acc, order, &terms) / H; + for (&terms, xs) |*v, x| v.* = (x - mean) * (x - mean); + const q = contract.rows.sumIn(acc, order, &terms); + t.row(&scratch, r, mean, Range.of(mean), q, Range.of(q), out, mag); + } + } +}; diff --git a/zig/tests/contract/metal_idol_rows.zig b/zig/tests/contract/metal_idol_rows.zig new file mode 100644 index 0000000000..63fa5b1043 --- /dev/null +++ b/zig/tests/contract/metal_idol_rows.zig @@ -0,0 +1,442 @@ +//! id_gate_add, id_embed, id_cond and id_bias32 under the contract: elementwise steps at idol-sdk's fp16 points. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const cfg = id.cfg; +const H = id.H; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +/// The steps' common parts: the outputs' ends, and its bounds and intervals from them. +fn Common(comptime T: type) type { + return struct { + pub fn bounds(t: *T, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *T, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + }; +} + +// --- id_gate_add: x = fp16(x + fp16(g y)), in place ------------------------------------------------------------ + +pub const GateAdd = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + old: []f16 = &.{}, + y: []f16 = &.{}, + gv: []f16 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the stream's width, 1,024 (the launcher's). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = id.noSums, + .ops = 1, + .mids = 1, + .out = .f16, + .in_place = true, + .served = &.{ .{ .rows = 12294, .n = H, .k = 1 }, .{ .rows = 8192, .n = H, .k = 1 }, .{ .rows = 4101, .n = H, .k = 1 } }, + }; + } + + pub fn deinit(t: *GateAdd) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *GateAdd) void { + t.a.free(t.old); + t.a.free(t.y); + t.a.free(t.gv); + t.old = &.{}; + t.y = &.{}; + t.gv = &.{}; + } + pub fn kernel(t: *GateAdd, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = blk: { + var hooks = g.hooks(GateAdd); + hooks.overlap = overlap; + break :blk hooks; + }; + /// The stream is its own output. + fn overlap(_: *anyopaque) bool { + return true; + } + pub const bounds = Common(GateAdd).bounds; + pub const intervals = Common(GateAdd).intervals; + + pub fn prepare(t: *GateAdd, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + const n = shape.rows * H; + t.old = try t.a.alloc(f16, n); + t.y = try t.a.alloc(f16, n); + t.gv = try t.a.alloc(f16, H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF16(rng, t.old, -8, 8); + id.randomF16(rng, t.y, -4, 4); + id.randomF16(rng, t.gv, -2, 2); + if (fill == .nonfinite) { + for ([_]f16{ id.nan16, id.inf16, -id.inf16 }, 0..) |v, i| { + t.old[(i * 13) % n] = v; + t.y[(i * 17 + 1) % n] = v; + } + t.gv[5] = id.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.y)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gv)); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *GateAdd) !void { + try t.lib.need("id_gate_add"); + // in place: the stream's rows go back before each independent run + @memcpy(t.rig.bytes(t.out, t.count() * 2), std.mem.sliceAsBytes(t.old)); + const p = t.rig.begin(); + t.lib.k.gateAdd(p.e, kref(t.out), kref(t.ins[0]), kref(t.ins[1]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *GateAdd) usize { + return t.rows * H; + } + + pub fn reference(t: *GateAdd, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.old, t.y, 0..) |xb, yb, i| { + const x: f64 = xb; + const y: f64 = yb; + const gv: f64 = t.gv[i % H]; + const r = Range.of(gv * y).fp32(0.5).fp16().add(Range.of(x)).fp32(0.5).fp16(); + t.ends.put(want, mag, i, hf(x + hf(gv * y)), r); + } + } +}; + +// --- id_embed: x = fp16(fp16(z w^T + b) + pe), z [8192, 16] fp32 -------------------------------------------------- + +/// The 16 inputs' FMA chain. +fn embedDepth(_: c.Shape) usize { + return cfg.latent_channels + 1; +} + +pub const Embed = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + stress: bool = false, + z: []f32 = &.{}, + /// w [1024, 16] then b [1024] + wb: []f32 = &.{}, + pe: []f16 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const R = cfg.latent_rows; + const K = cfg.latent_channels; + + /// Rows: the latent's 8,192, which the launcher fixes (a gate's fewer rows run all of them); k: its 16 channels. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = embedDepth, .ops = 2, .mids = 1, .out = .f16, .served = &.{.{ .rows = R, .n = H, .k = K }} }; + } + + pub fn deinit(t: *Embed) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Embed) void { + t.a.free(t.z); + t.a.free(t.wb); + t.a.free(t.pe); + t.z = &.{}; + t.wb = &.{}; + t.pe = &.{}; + } + pub fn kernel(t: *Embed, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Embed); + pub const bounds = Common(Embed).bounds; + pub const intervals = Common(Embed).intervals; + + pub fn prepare(t: *Embed, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H or shape.k != K) return false; + t.rig.reset(); + t.free(); + t.stress = fill == .stress; + t.z = try t.a.alloc(f32, R * K); + t.wb = try t.a.alloc(f32, H * K + H); + t.pe = try t.a.alloc(f16, R * H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // products to 2^16 whose halves cancel but for the first: partial sums near 2^19, exact in fp32 + for (0..R) |r| for (0..K / 2) |k| { + t.z[r * K + k] = id.int(rng, 1, 255); + t.z[r * K + K / 2 + k] = t.z[r * K + k] + @as(f32, if (k == 0) 1 else 0); + }; + for (0..H) |o| for (0..K / 2) |k| { + t.wb[o * K + k] = id.sign(rng) * id.int(rng, 1, 255); + t.wb[o * K + K / 2 + k] = -t.wb[o * K + k]; + }; + for (t.wb[H * K ..]) |*v| v.* = id.int(rng, -8, 8); + for (t.pe) |*v| v.* = @floatCast(id.int(rng, -8, 8)); + } else { + id.randomF32(rng, t.z, -3, 3); + id.randomF32(rng, t.wb, -0.5, 0.5); + id.randomF16(rng, t.pe, -1, 1); + } + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| t.z[i * K + 3 * i] = v; + t.wb[H * K + 9] = id.nan32; + t.pe[5 * H + 11] = id.inf16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.z)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wb[0 .. H * K])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.wb[H * K ..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.pe)); + t.out = try t.rig.output(R * H * 2); + return true; + } + + pub fn run(t: *Embed) !void { + try t.lib.need("id_embed"); + const p = t.rig.begin(); + t.lib.k.embed(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.out)); + try t.rig.finish(p); + } + + pub fn count(_: *Embed) usize { + return R * H; + } + + fn each(t: *Embed, ends: *Ends, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) void { + const gd = c.gamma(@floatFromInt(embedDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var terms: [K]f64 = undefined; + for (0..R) |r| for (0..H) |o| { + var s: f64 = 0; + var abs: f64 = 0; + for (&terms, t.z[r * K ..][0..K], t.wb[o * K ..][0..K]) |*v, z, w| { + v.* = @as(f64, z) * @as(f64, w); + s += v.*; + abs += @abs(v.*); + } + var sr = Range.of(s).widen(gd * abs); + if (acc) |x| { + s = contract.rows.sumIn(x, order, &terms); + sr = Range.of(s); + } else if (t.stress) sr = Range.of(s); + const b: f64 = t.wb[H * K + o]; + const pe: f64 = t.pe[r * H + o]; + const y = sr.add(Range.of(b)).fp32(0.5).fp16().add(Range.of(pe)).fp32(0.5).fp16(); + ends.put(want, mag, r * H + o, hf(hf(s + b) + pe), y); + }; + } + + pub fn reference(t: *Embed, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + t.each(&t.ends, want, mag, null, .sequential); + } + + pub fn simulate(t: *Embed, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + t.each(&scratch, out, mag, acc, order); + } +}; + +// --- id_cond: y = fp16((a + ba) + (c + bc)), the two condition projections' fp32 sums and biases ----------------- + +pub const Cond = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + /// a, c [rows, 1024], then ba, bc [1024] + v: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: the condition's 4,101; n: its width, 1,024 (the launcher's). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .ops = 3, .out = .f16, .served = &.{.{ .rows = cfg.cond_rows, .n = H, .k = 1 }} }; + } + + pub fn deinit(t: *Cond) void { + t.a.free(t.v); + t.v = &.{}; + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Cond, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Cond); + pub const bounds = Common(Cond).bounds; + pub const intervals = Common(Cond).intervals; + + pub fn prepare(t: *Cond, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H) return false; + t.rig.reset(); + t.a.free(t.v); + t.rows = shape.rows; + const n = shape.rows * H; + t.v = try t.a.alloc(f32, 2 * n + 2 * H); + var prng = std.Random.DefaultPrng.init(seed); + id.randomF32(prng.random(), t.v, -8, 8); + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| { + t.v[(i * 13) % n] = v; + t.v[n + (i * 17 + 1) % n] = v; + } + t.v[2 * n + 5] = id.inf32; + t.v[2 * n + H + 6] = id.nan32; + } + for (0..4) |i| { + const part = if (i < 2) t.v[i * n ..][0..n] else t.v[2 * n + (i - 2) * H ..][0..H]; + t.ins[i] = try t.rig.input(std.mem.sliceAsBytes(part)); + } + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *Cond) !void { + try t.lib.need("id_cond"); + const p = t.rig.begin(); + t.lib.k.cond(p.e, kref(t.ins[0]), kref(t.ins[2]), kref(t.ins[1]), kref(t.ins[3]), kref(t.out), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Cond) usize { + return t.rows * H; + } + + pub fn reference(t: *Cond, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const n = t.count(); + for (0..n) |i| { + const av: f64 = t.v[i]; + const cv: f64 = t.v[n + i]; + const ba: f64 = t.v[2 * n + i % H]; + const bc: f64 = t.v[2 * n + H + i % H]; + const r = Range.of(av + ba).fp32(0.5).add(Range.of(cv + bc).fp32(0.5)).fp32(0.5).fp16(); + t.ends.put(want, mag, i, hf((av + ba) + (cv + bc)), r); + } + } +}; + +// --- id_bias32: y += b by column, fp32, in place -------------------------------------------------------------- + +pub const Bias32 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + shape: c.Shape = undefined, + old: []f32 = &.{}, + b: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: the latent's; n: the velocity's 16 channels. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .ops = 1, .out = .f32, .in_place = true, .served = &.{.{ .rows = cfg.latent_rows, .n = cfg.latent_channels, .k = 1 }} }; + } + + pub fn deinit(t: *Bias32) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Bias32) void { + t.a.free(t.old); + t.a.free(t.b); + t.old = &.{}; + t.b = &.{}; + } + pub fn kernel(t: *Bias32, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = blk: { + var hooks = g.hooks(Bias32); + hooks.overlap = overlap; + break :blk hooks; + }; + /// The sums are their own output. + fn overlap(_: *anyopaque) bool { + return true; + } + pub const bounds = Common(Bias32).bounds; + pub const intervals = Common(Bias32).intervals; + + pub fn prepare(t: *Bias32, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const n = shape.rows * shape.n; + t.old = try t.a.alloc(f32, n); + t.b = try t.a.alloc(f32, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF32(rng, t.old, -16, 16); + id.randomF32(rng, t.b, -2, 2); + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| t.old[(i * 13) % n] = v; + if (shape.n > 6) t.b[5] = id.inf32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.out = try t.rig.output(n * 4); + return true; + } + + pub fn run(t: *Bias32) !void { + try t.lib.need("id_bias32"); + @memcpy(t.rig.bytes(t.out, t.count() * 4), std.mem.sliceAsBytes(t.old)); + const p = t.rig.begin(); + t.lib.k.bias32(p.e, kref(t.out), kref(t.ins[0]), t.shape.rows, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Bias32) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *Bias32, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.old, 0..) |y, i| { + const b: f64 = t.b[i % t.shape.n]; + t.ends.put(want, mag, i, y + b, Range.of(y + b).fp32(0.5)); + } + } +}; diff --git a/zig/tests/contract/metal_idol_steps.zig b/zig/tests/contract/metal_idol_steps.zig new file mode 100644 index 0000000000..617729535b --- /dev/null +++ b/zig/tests/contract/metal_idol_steps.zig @@ -0,0 +1,426 @@ +//! id_repo and id_qk under the contract: RePo's angles in fp32, and the heads rotated, normed and rounded to fp16. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const H = id.H; +const D = id.D; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; + +const streams: []const c.Shape = &.{ .{ .rows = 12294, .n = 512, .k = 128 }, .{ .rows = 8192, .n = 512, .k = 128 }, .{ .rows = 4101, .n = 512, .k = 128 } }; + +// --- id_repo: u = silu(g) c, d = u final^T (48), angles x f pi as cos and sin [rows, 16, 32] -------------------- + +const pi32: f64 = @as(f32, std.math.pi); + +/// d's FMA chain over 128. +fn repoDepth(_: c.Shape) usize { + return 128 + 1; +} + +/// A pair's axis: 10 x, 10 y, then 12 z. +fn axis(j: usize) usize { + return if (j < 10) 0 else if (j < 20) 1 else 2; +} + +pub const Repo = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + stress: bool = false, + /// gate sums then content sums, 128 each a row + gc: []f32 = &.{}, + /// the final map [48, 128] + wf: []f32 = &.{}, + freqs: []f32 = &.{}, + inv: f32 = 0.5, + ins: [3]rig_.Ref = undefined, + outs: [2]rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: 16 heads x 32 pairs; k: RePo's hidden width. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = repoDepth, .ops = id.exp_ulps + id.div_ulps + 2, .out = .f32, .served = streams }; + } + + pub fn deinit(t: *Repo) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Repo) void { + t.a.free(t.gc); + t.a.free(t.wf); + t.a.free(t.freqs); + t.gc = &.{}; + t.wf = &.{}; + t.freqs = &.{}; + } + pub fn kernel(t: *Repo, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Repo); + + pub fn prepare(t: *Repo, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != 512 or shape.k != 128) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + t.stress = fill == .stress; + t.gc = try t.a.alloc(f32, shape.rows * 256); + t.wf = try t.a.alloc(f32, 48 * 128); + t.freqs = try t.a.alloc(f32, 32); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // gates of 20 to 40 (sigmoid 1 in fp32), integer contents: u's halves differ by a gate twice, d small + for (0..shape.rows) |r| for (0..64) |i| { + const gr = t.gc[r * 256 ..][0..256]; + gr[i] = id.int(rng, 20, 40) / t.inv; + gr[i + 64] = gr[i]; + gr[128 + i] = id.sign(rng) * id.int(rng, 1, 8) / t.inv; + gr[128 + i + 64] = gr[128 + i] + @as(f32, if (i < 2) 1 / t.inv else 0); + }; + for (0..48) |o| for (0..64) |i| { + t.wf[o * 128 + i] = id.sign(rng) * id.int(rng, 1, 32); + t.wf[o * 128 + 64 + i] = -t.wf[o * 128 + i]; + }; + id.randomF32(rng, t.freqs, 0.001, 0.01); + } else { + id.randomF32(rng, t.gc, -4, 4); + id.randomF32(rng, t.wf, -0.2, 0.2); + id.randomF32(rng, t.freqs, 0.05, 1.5); + } + if (fill == .nonfinite) { + const at = [_]usize{ 0 * 256 + 5, 1 * 256 + 7, 2 * 256 + 128 + 3 }; + for (at, [_]f32{ id.nan32, id.inf32, -id.inf32 }) |p, v| if (p < t.gc.len) { + t.gc[p] = v; + }; + t.freqs[4] = id.nan32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.gc)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wf)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.freqs)); + for (&t.outs) |*o| o.* = try t.rig.output(shape.rows * 512 * 4); + return true; + } + + pub fn run(t: *Repo) !void { + try t.lib.need("id_repo"); + const p = t.rig.begin(); + t.lib.k.repo(p.e, kref(t.ins[0]), t.inv, kref(t.ins[1]), kref(t.ins[2]), kref(t.outs[0]), kref(t.outs[1]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Repo) usize { + return 2 * t.rows * 512; + } + + /// Row r's u values in float64 and their ranges through the kernel's exp, add, division and product. + fn us(t: *const Repo, r: usize, u: *[128]f64, ur: *[128]Range) void { + const inv: f64 = t.inv; + for (u, ur, 0..) |*v, *vr, i| { + const a = @as(f64, t.gc[r * 256 + i]) * inv; + const cc = @as(f64, t.gc[r * 256 + 128 + i]) * inv; + v.* = a / (1 + @exp(-a)) * cc; + // a >= 20: 1 + exp(-a) is 1 in fp32 and a / 1 is a, so u = a c exactly + if (t.stress) { + v.* = a * cc; + vr.* = Range.of(v.*); + continue; + } + const den = Range.of(@exp(-a)).fp32(id.exp_ulps).add(Range.of(1)).fp32(0.5); + vr.* = Range.of(a).mul(den.recip()).fp32(id.div_ulps).mul(Range.of(cc)).fp32(0.5); + } + } + + /// Row r's outputs from its 48 offsets d (ranges dr). + fn row(t: *Repo, ends: *Ends, r: usize, d: *const [48]f64, dr: *const [48]Range, want: []f64, mag: []f64) void { + for (0..512) |i| { + const h = i / 32; + const j = i % 32; + const x = d[h * 3 + axis(j)]; + const xr = dr[h * 3 + axis(j)]; + const f: f64 = t.freqs[j]; + const th = std.math.tanh(f); + const ang = (x * th + x * (f - th)) * pi32; + // tanh's error cancels between the two products: fp32 rounds each product, the sum and the scale + const xm = @max(@abs(xr.lo), @abs(xr.hi)); + const e = 2 * 0x1p-24 * pi32 * xm * (@abs(th) + 2 * @abs(f - th) + 2 * @abs(f)) + 0x1p-140; + const lo = @min(xr.lo * f, xr.hi * f) * pi32 - e; + const hi = @max(xr.lo * f, xr.hi * f) * pi32 + e; + ends.put(want, mag, r * 512 + i, @cos(ang), id.cosOver(lo, hi).fp32(id.trig_ulps)); + ends.put(want, mag, (t.rows + r) * 512 + i, @sin(ang), id.sinOver(lo, hi).fp32(id.trig_ulps)); + } + } + + pub fn reference(t: *Repo, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const gd = c.gamma(@floatFromInt(repoDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var u: [128]f64 = undefined; + var ur: [128]Range = undefined; + var d: [48]f64 = undefined; + var dr: [48]Range = undefined; + for (0..t.rows) |r| { + t.us(r, &u, &ur); + for (&d, &dr, 0..) |*dv, *dv_r, o| { + var s: f64 = 0; + var rr = Range.of(0); + var abs: f64 = 0; + for (u, ur, t.wf[o * 128 ..][0..128]) |uv, uvr, wv| { + s += uv * wv; + rr = rr.add(uvr.mul(Range.of(wv))); + abs += @max(@abs(uvr.lo), @abs(uvr.hi)) * @abs(@as(f64, wv)); + } + dv.* = s; + dv_r.* = if (t.stress) rr else rr.widen(gd * abs); + } + t.row(&t.ends, r, &d, &dr, want, mag); + } + } + pub fn bounds(t: *Repo, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Repo, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Repo, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var u: [128]f64 = undefined; + var ur: [128]Range = undefined; + var d: [48]f64 = undefined; + var dr: [48]Range = undefined; + var terms: [128]f64 = undefined; + for (0..t.rows) |r| { + t.us(r, &u, &ur); + for (&d, &dr, 0..) |*dv, *dv_r, o| { + for (&terms, u, t.wf[o * 128 ..][0..128]) |*v, uv, wv| v.* = uv * wv; + dv.* = contract.rows.sumIn(acc, order, &terms); + dv_r.* = Range.of(dv.*); + } + t.row(&scratch, r, &d, &dr, out, mag); + } + } +}; + +// --- id_qk: (q, k) pairs turned by cos and sin and rounded, F.normalize, times gamma and 8, rounded; v copied ----- + +const qk_why = "an fp16 output follows the head's sum of squares as its square root: a bf16 sum is off a few ulps"; + +const qk_known = "a NaN in a head gives its other outputs +-inf where F.normalize gives NaN: max(sqrt(ss), 1e-12f) drops it"; + +/// A lane's two squares, then the simdgroup tree. +fn qkDepth(_: c.Shape) usize { + return 1 + 5; +} + +pub const Qk = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *id.Lib, + name: []const u8, + rows: usize = 0, + stress: bool = false, + qkv: []f16 = &.{}, + /// gq then gk [1024] each + gam: []f32 = &.{}, + /// cos then sin [rows, 512] each + rot: []f32 = &.{}, + ins: [5]rig_.Ref = undefined, + outs: [3]rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: heads; k: the head width. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = qkDepth, + .ops = id.sqrt_ulps + id.div_ulps + 2, + .out = .f16, + .served = &.{ .{ .rows = 12294, .n = 16, .k = D }, .{ .rows = 8192, .n = 16, .k = D }, .{ .rows = 4101, .n = 16, .k = D } }, + .sums = .{ .bf16 = 0, .why = qk_why }, + .known = .{ .gate = .nonfinite, .finding = qk_known }, + }; + } + + pub fn deinit(t: *Qk) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Qk) void { + t.a.free(t.qkv); + t.a.free(t.gam); + t.a.free(t.rot); + t.qkv = &.{}; + t.gam = &.{}; + t.rot = &.{}; + } + pub fn kernel(t: *Qk, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Qk); + + pub fn prepare(t: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != 16 or shape.k != D) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + t.stress = fill == .stress; + const n = shape.rows; + t.qkv = try t.a.alloc(f16, n * 3 * H); + t.gam = try t.a.alloc(f32, 2 * H); + t.rot = try t.a.alloc(f32, 2 * n * 512); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF16(rng, t.qkv, -4, 4); + id.randomF32(rng, t.gam, 0.25, 2); + for (0..n * 512) |i| { + if (t.stress) { + // quarter turns: exact rotations of integer heads, whose squares sum to ~2^20 exactly + const turns = [4][2]f32{ .{ 1, 0 }, .{ 0, 1 }, .{ -1, 0 }, .{ 0, -1 } }; + const turn = turns[rng.uintLessThan(usize, 4)]; + t.rot[i] = turn[0]; + t.rot[n * 512 + i] = turn[1]; + continue; + } + const angle = rng.float(f32) * 2 * std.math.pi; + t.rot[i] = @cos(angle); + t.rot[n * 512 + i] = @sin(angle); + } + if (t.stress) for (0..n) |r| for (0..2 * H) |i| { + t.qkv[r * 3 * H + i] = @floatCast(id.int(rng, -128, 128)); + }; + if (fill == .nonfinite) { + // q NaN, k +inf, v -inf, in turn; then a NaN cos + const at = [_]usize{ 0 * 3 * H + 0 * D + 5, 1 * 3 * H + H + 2 * D + 9, 2 * 3 * H + 2 * H + 3 * D + 1 }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16 }) |p, v| if (p < t.qkv.len) { + t.qkv[p] = v; + }; + if (n > 3) t.rot[3 * 512 + 4 * 32 + 6] = id.nan32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gam[0..H])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.gam[H..])); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.rot[0 .. n * 512])); + t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.rot[n * 512 ..])); + for (&t.outs) |*o| o.* = try t.rig.output(n * H * 2); + return true; + } + + pub fn run(t: *Qk) !void { + try t.lib.need("id_qk"); + const p = t.rig.begin(); + const in = t.ins; + t.lib.k.heads(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(in[3]), kref(in[4]), kref(t.outs[0]), kref(t.outs[1]), kref(t.outs[2]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Qk) usize { + return 3 * t.rows * H; + } + + /// Head h of row r of q (0) or k (1): its pairs turned and rounded, as float64 values and ranges. + fn turned(t: *const Qk, which: usize, r: usize, h: usize, v: *[D]f64, vr: *[D]Range) void { + const x = t.qkv[r * 3 * H + which * H + h * D ..][0..D]; + for (0..D / 2) |p| { + const cs: f64 = t.rot[r * 512 + h * 32 + p]; + const sn: f64 = t.rot[(t.rows + r) * 512 + h * 32 + p]; + const a: f64 = x[2 * p]; + const b: f64 = x[2 * p + 1]; + v[2 * p] = hf(a * cs - b * sn); + v[2 * p + 1] = hf(a * sn + b * cs); + const ac = Range.of(a * cs).fp32(0.5); + const bs = Range.of(b * sn).fp32(0.5); + const as_ = Range.of(a * sn).fp32(0.5); + const bc = Range.of(b * cs).fp32(0.5); + vr[2 * p] = ac.sub(bs).fp32(0.5).fp16(); + vr[2 * p + 1] = as_.add(bc).fp32(0.5).fp16(); + } + } + + /// The head's outputs from its sum of squares `ss` (range `ssr`). + fn head(t: *Qk, ends: *Ends, which: usize, r: usize, h: usize, v: *const [D]f64, vr: *const [D]Range, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const e12: f64 = @as(f32, 1e-12); + // torch's clamp_min keeps a NaN norm (@max would drop it) + const norm = if (std.math.isNan(ss)) ss else @max(@sqrt(ss), e12); + const nr0 = ssr.up(id.sqrt).fp32(id.sqrt_ulps); + const nr = Range{ .lo = @max(nr0.lo, e12), .hi = @max(nr0.hi, e12) }; + const gam = t.gam[which * H + h * D ..][0..D]; + const base = which * t.rows * H + (h * t.rows + r) * D; + for (0..D) |i| { + const gv: f64 = gam[i]; + const w = hf(v[i] / norm * gv * 8); + const rr = vr[i].mul(nr.recip()).fp32(id.div_ulps).mul(Range.of(gv)).fp32(0.5).mul(Range.of(8)).fp32(0.5).fp16(); + ends.put(want, mag, base + i, w, rr); + } + } + + /// v's heads, copied bit for bit. + fn values(t: *Qk, ends: *Ends, want: []f64, mag: []f64) void { + for (0..t.rows) |r| for (0..16) |h| for (0..D) |i| { + const f: f64 = t.qkv[r * 3 * H + 2 * H + h * D + i]; + ends.put(want, mag, 2 * t.rows * H + (h * t.rows + r) * D + i, f, Range.of(f)); + }; + } + + pub fn reference(t: *Qk, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const gd = c.gamma(@floatFromInt(qkDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var v: [D]f64 = undefined; + var vr: [D]Range = undefined; + for (0..2) |which| for (0..t.rows) |r| for (0..16) |h| { + t.turned(which, r, h, &v, &vr); + var ss: f64 = 0; + var lo: f64 = 0; + var hi: f64 = 0; + for (v, vr) |x, xr| { + ss += x * x; + const q = xr.square(); + lo += q.lo; + hi += q.hi; + } + const ssr = if (t.stress) Range.of(ss) else Range{ .lo = lo * (1 - gd), .hi = hi * (1 + gd) }; + t.head(&t.ends, which, r, h, &v, &vr, ss, ssr, want, mag); + }; + t.values(&t.ends, want, mag); + } + pub fn bounds(t: *Qk, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Qk, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + var v: [D]f64 = undefined; + var vr: [D]Range = undefined; + var terms: [D]f64 = undefined; + for (0..2) |which| for (0..t.rows) |r| for (0..16) |h| { + t.turned(which, r, h, &v, &vr); + for (&terms, v) |*s, x| s.* = x * x; + const ss = contract.rows.sumIn(acc, order, &terms); + t.head(&scratch, which, r, h, &v, &vr, ss, Range.of(ss), out, mag); + }; + t.values(&scratch, out, mag); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 796ed74d6b..ece8073ecc 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -102,6 +102,9 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol ")) { + text = sources.idol; + frags = true; } else if (std.mem.startsWith(u8, d.name, "core deltanet ")) { text = sources.core_deltanet; } else if (std.mem.startsWith(u8, d.name, "qwen27 ")) { diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index b32e2baa9e..9799191182 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -427,6 +427,8 @@ pub fn main(init: std.process.Init) !u8 { defer native_gdn.destroy(); const moss_image = try @import("metal_moss_image.zig").Suite.create(a, device, queue, &kernels, o.only); defer moss_image.destroy(); + const idol = try @import("metal_idol.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { From e99604cca4ac2615c04dfd70fb19c44df80b0464 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:48:17 +0800 Subject: [PATCH 03/21] idol: id_qk keeps a NaN norm as F.normalize does; the recipe and changelog line --- CHANGELOG.md | 2 + docs/kernel-contract.md | 2 - docs/recipes/idol.md | 129 ++++++++++++++++++++++++ zig/kernels/metal/idol.metal | 6 +- zig/tests/contract/metal_idol_steps.zig | 3 - 5 files changed, 135 insertions(+), 7 deletions(-) create mode 100644 docs/recipes/idol.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f676abf2a..cad527edba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (fixes #1757). Thanks to @jayleaton. + - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. - Qwen3.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 472881c3b5..b2b465117a 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -117,8 +117,6 @@ MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 an `tf_glm_scores` scores a NaN dot as 0 where float64 keeps the NaN: its relu is `fmaxf(x, 0)`, as `tl.maximum` is in the Python family's Triton glue it was ported from. Nemotron's Metal `tf_xup_w3` gives 0 for a NaN dot where float64 keeps the NaN: its relu2 is `metal::max(bfloat(acc), 0.0f)`; the other up pipelines give its bytes on the same inputs. -Idol's `id_qk` gives a head holding a NaN +-inf in its other outputs where torch's `F.normalize` gives NaN: its -norm is `max(precise::sqrt(ss), 1e-12f)`, and Metal's `max` drops the NaN that `clamp_min` keeps. HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md new file mode 100644 index 0000000000..27995b8120 --- /dev/null +++ b/docs/recipes/idol.md @@ -0,0 +1,129 @@ +# Idol's denoiser on Metal + +Idol (speridlabs) turns one image into a 3D Gaussian splat. A 0.37B flow transformer denoises 8,192 latent rows of 16 +channels, conditioned on the image's DINOv3 and FLUX.2 VAE rows; two decoders then grow an octree from that latent and +emit the Gaussians. `zig/src/families/idol` runs the denoiser on the Metal runtime from the released checkpoint: + +- once an image: the condition embedding (DINOv3's 1,280 channels and the VAE's 128, each projected to 1,024 and + summed) and the two condition-refiner blocks (no modulation, affine LayerNorms); +- at each step: the two noise-refiner blocks over the latent rows, then the 24 joint blocks over 12,294 rows (8,192 + latent, 4,101 condition, one camera token), then the output head and the camera velocity. + +Every block is 16 heads of 64 at width 1,024 with a tanh-GELU MLP of 4,096. Its rotary is RePo: a learned map from +each row's LayerNorm to per-head 3D positions, angles from 10, 10 and 12 frequencies per axis. + +The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_cond - 3 v_uncond` (two passes a step: +with the condition, and with zeros). + +The image encoders, the decoders and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's +SDK in torch, a development reference. It is a denoiser, not a token family: it does not register with the lane +engine, the server or `families.detect`. + +It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them. + +## Checkpoint + +`speridlabs/idol` at `2731c2291fc65a797c5cb627a5e5526fc1490902`, Apache 2.0, ungated. `config.json` must name +`LatentSeqMMFlowModel` at the released geometry, and `denoiser.safetensors` must hold its 560 tensors; anything else is +refused. The reference is idol-sdk at `4702a7f00850e34b6e83a2170cc123c3e9570744` (Apache 2.0). + +## Precision + +The SDK loads the checkpoint's fp32 weights and converts the blocks' linear layers to fp16; everything else stays fp32. +This family keeps that split: + +- fp16 block projections with fp32 sums; the bias joins the sums before one fp16 rounding, GELU after it. +- LayerNorms in fp32, rounded to fp16; the modulation's shift and scale in fp16 ops, as torch's fp16 tensors run them. +- RePo in fp32: its LayerNorm, both 1,024 to 128 maps, the 128 to 48 map and the angles. +- Attention: fp16 q, k and v, fp32 scores and sums. q and k are rotated, then L2-normalized and scaled by the + checkpoint's gamma times 8. +- The condition embedding and the output projection in fp32; the time MLP, modulation and camera MLPs on the host. + +### The tensor units' fp32 matmul is not fp32 + +On the M5, a GEMM whose inputs are fp32 on the tensor units (`tfp::gemm_nax`, the Qwen-VL tower's +`qwv_gemm_f32`, MLX's fp32 matmul) rounds its inputs to about fp16: its error against float64 is about 7.5e-4 of the +result, against about 1e-7 for an fp32 FMA. RePo's maps at that precision move block 8's attention out of its gate. + +`id_norm_split` therefore writes each RePo input row as fp16 `[hi | hi | lo]` (`lo = x - hi`), and the map's weights are +stored `[hi | lo | hi]`, scaled by a power of two so `lo` stays normal. One fp16 split-K GEMM with K = 3,072 then sums +`hi*hi + hi*lo + lo*hi` in fp32; `id_repo` undoes the scale. With it every block passes its gate. `id_gemm32`, +an fp32 FMA GEMM on simdgroup matrices, is exact but takes 0.47 s a forward if it runs RePo; it serves the condition +embedding (once an image) and the output projection. + +## Kernel contract + +`zig/tests/contract/metal_idol*.zig` registers every kernel of `idol.metal` and `idol_attention.metal`: + +- `id_linear`'s eight forms (plain and GELU, each tile tail) and the split-K halves `id_split_0` and `id_split_1`; +- `id_norm` (plain, affine, modulated), `id_norm32` (affine, modulated) and `id_norm_split`; +- `id_repo`, `id_qk`, `id_gate_add`, `id_embed`, `id_cond`, `id_bias32` and `id_gemm32`; +- the attention's four forms. + +Each output must lie in the range its fp32 steps allow at the SDK's fp16 rounding points. Nineteen controls drop or +add one rounding point or sum in fp16, and each must go red at its gate. One runs the tensor units' fp32 GEMM in +`id_gemm32`'s place: gate 4 misses it by 13 million times the exact-sum bound. `id_qk` keeps a NaN norm as +`F.normalize` does. Forms no engine call reaches (`id_linear_00` and `_10`, attention `01` and `10`, plain `id_norm`, +affine `id_norm32`) are checked at made-up shapes. The tensor-unit kernels do not build for M1 or M2 offline, as +MOSS-Image's do not; the rest build at their reserved thread counts. + +`tf-kernel-contract --only "metal idol"` on the M5 Ultra: + +``` +SUMMARY 27 kernels: 27 pass gates 1-4, 0 fail, 0 known findings; 19 controls: 19 red as planted, 0 missed; 0 limit problems +``` + +## Run a case + +`idol_case.py` imports the SDK's `dit.py`, `layers.py` and `sampler.py` only when each file's sha256 is the pinned +commit's, and the snapshot only when it is the pinned revision. + +```sh +hf download speridlabs/idol --revision 2731c2291fc65a797c5cb627a5e5526fc1490902 --local-dir WEIGHTS +git clone https://github.com/speridlabs/idol-sdk SDK && git -C SDK checkout 4702a7f00850e34b6e83a2170cc123c3e9570744 +python3 tools/zig/idol_case.py case WEIGHTS SDK case/ --features dino.safetensors vae.safetensors +python3 tools/zig/idol_case.py exact WEIGHTS SDK case/ # the same steps with an fp32 denoiser +python3 tools/zig/idol_case.py blocks WEIGHTS SDK case/ --step 0 # each block's input and output, fp16 and fp32 +zig build tf-idol -Dcpu=apple_m1 +zig-out/bin/tf-idol WEIGHTS case/ +zig-out/bin/tf-idol WEIGHTS case/ blocks 0 +``` + +The features are the SDK's own `feature1` [1, 4101, 1280] and `feature2` [1, 4101, 128] for an image. The case holds the +seeded noise (torch's CPU generator), the schedule, the SDK's input, both velocities and camera velocities at every +step, and its final sample. `exact` adds fp32 velocities on the same inputs; `blocks` adds every block's input and +output at one step, and each joint block run alone in fp32 on the SDK's fp16 input. + +`tf-idol` gates: + +- **Each step**, both branches on the SDK's own input: cosine against the fp32 velocity within 2x the SDK's own + distance (`1 - cos`). Worst row over the run within 2x the SDK's worst row. Fifty fp16 steps make single rows chaotic, + so the per-row gate runs per block instead: +- **Each block** alone on the SDK's fp16 input, against that block in fp32: worst row and cosine within 2x the SDK's. +- **Repeat**: the same step twice, another between, byte for byte. + +It then runs the whole guided schedule from the noise, times it, and writes `sample.f32`. `reference.zig`'s tests plant +a 5% error in one row and require the gate to fail it. + +## Measured + +Setup: + +- Mac Studio M5 Ultra 256 GB, macOS 27; this commit's `tf-idol`. +- The SDK on torch 2.13.0 (MPS, fp16 blocks, `ATTN_BACKEND=sdpa`). +- The SDK's avocado example: its own DINOv3 and VAE features, seed 42, 50 steps, guidance 3. + +Results: + +- **Velocity**, both branches of all 50 steps on the SDK's own input (100 of 100 pass): + - against the SDK: cosine 0.999997 or better; + - against fp32: our worst row over the run 0.0456, the SDK's 0.0398; per step our `1 - cos` is the SDK's within + the 2x bound at every step, and the smaller at 40 of the 100. +- **Blocks**, at steps 1 and 41, each of the 23 joint blocks after the first on the SDK's input: all pass, worst rows + 0.0007 to 0.0017 against fp32, where the SDK's are 0.0007 to 0.0025. +- **Sample:** our 50 guided steps from the same noise end at cosine 0.99999 to the SDK's sample. +- **Repeat:** identical bytes. +- **Speed**, one session: 0.273 s a forward (27.3 s for 50 guided steps), 3.4x the SDK's 0.939 s (the median of its + 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. +- **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the + denoiser's fp16 weights in shared Metal buffers. diff --git a/zig/kernels/metal/idol.metal b/zig/kernels/metal/idol.metal index 1a8048d1b2..f687b50e5a 100644 --- a/zig/kernels/metal/idol.metal +++ b/zig/kernels/metal/idol.metal @@ -281,8 +281,10 @@ struct IdHeads { const float k0 = float(qkv[from + 1024]), k1 = float(qkv[from + 1025]); const float rq0 = id_h(q0 * c - q1 * s), rq1 = id_h(q0 * s + q1 * c); const float rk0 = id_h(k0 * c - k1 * s), rk1 = id_h(k0 * s + k1 * c); - const float qn = max(metal::precise::sqrt(simd_sum(rq0 * rq0 + rq1 * rq1)), 1e-12f); - const float kn = max(metal::precise::sqrt(simd_sum(rk0 * rk0 + rk1 * rk1)), 1e-12f); + // F.normalize's clamp_min keeps a NaN norm, where Metal's max would drop it + const float sq = metal::precise::sqrt(simd_sum(rq0 * rq0 + rq1 * rq1)); + const float sk = metal::precise::sqrt(simd_sum(rk0 * rk0 + rk1 * rk1)); + const float qn = isnan(sq) ? sq : max(sq, 1e-12f), kn = isnan(sk) ? sk : max(sk, 1e-12f); const uint g = h * 64 + lane * 2; const ulong to = (ulong(h) * p.rows + r) * 64 + lane * 2; Q[to] = half(rq0 / qn * gq[g] * p.scale); diff --git a/zig/tests/contract/metal_idol_steps.zig b/zig/tests/contract/metal_idol_steps.zig index 617729535b..0202b284b2 100644 --- a/zig/tests/contract/metal_idol_steps.zig +++ b/zig/tests/contract/metal_idol_steps.zig @@ -220,8 +220,6 @@ pub const Repo = struct { const qk_why = "an fp16 output follows the head's sum of squares as its square root: a bf16 sum is off a few ulps"; -const qk_known = "a NaN in a head gives its other outputs +-inf where F.normalize gives NaN: max(sqrt(ss), 1e-12f) drops it"; - /// A lane's two squares, then the simdgroup tree. fn qkDepth(_: c.Shape) usize { return 1 + 5; @@ -253,7 +251,6 @@ pub const Qk = struct { .out = .f16, .served = &.{ .{ .rows = 12294, .n = 16, .k = D }, .{ .rows = 8192, .n = 16, .k = D }, .{ .rows = 4101, .n = 16, .k = D } }, .sums = .{ .bf16 = 0, .why = qk_why }, - .known = .{ .gate = .nonfinite, .finding = qk_known }, }; } From 3114e80268272a005bc925f2bedd53d332b658fe Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:54:43 +0800 Subject: [PATCH 04/21] changelog: #1766 for Idol's denoiser --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cad527edba..38fd52543d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (fixes #1757). Thanks to @jayleaton. +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. From 427c7465eea4dc617afb871114586b1de95ac464 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:09:00 +0800 Subject: [PATCH 05/21] idol: attention rounds P to fp16 before P V, so its denominator sums the values the units multiply --- CHANGELOG.md | 2 +- docs/recipes/idol.md | 15 ++++++++------- zig/kernels/metal/idol_attention.metal | 3 ++- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 38fd52543d..a04a4c04eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.273 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 27995b8120..007f76fe6b 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -36,7 +36,8 @@ This family keeps that split: - LayerNorms in fp32, rounded to fp16; the modulation's shift and scale in fp16 ops, as torch's fp16 tensors run them. - RePo in fp32: its LayerNorm, both 1,024 to 128 maps, the 128 to 48 map and the angles. - Attention: fp16 q, k and v, fp32 scores and sums. q and k are rotated, then L2-normalized and scaled by the - checkpoint's gamma times 8. + checkpoint's gamma times 8. The tensor units take an fp32 P short, which left outputs 3.5e-4 low on average, so + `id_attention` rounds P to fp16 itself and sums the softmax's denominator over those same values. - The condition embedding and the output projection in fp32; the time MLP, modulation and camera MLPs on the host. ### The tensor units' fp32 matmul is not fp32 @@ -116,14 +117,14 @@ Setup: Results: - **Velocity**, both branches of all 50 steps on the SDK's own input (100 of 100 pass): - - against the SDK: cosine 0.999997 or better; - - against fp32: our worst row over the run 0.0456, the SDK's 0.0398; per step our `1 - cos` is the SDK's within - the 2x bound at every step, and the smaller at 40 of the 100. + - against the SDK: cosine 0.999998 or better; + - against fp32: our worst row over the run 0.0474, the SDK's 0.0398; per step our `1 - cos` is the SDK's within + the 2x bound at every step, and the smaller at 75 of the 100. - **Blocks**, at steps 1 and 41, each of the 23 joint blocks after the first on the SDK's input: all pass, worst rows - 0.0007 to 0.0017 against fp32, where the SDK's are 0.0007 to 0.0025. -- **Sample:** our 50 guided steps from the same noise end at cosine 0.99999 to the SDK's sample. + 0.0008 to 0.0017 against fp32, where the SDK's are 0.0007 to 0.0025. +- **Sample:** our 50 guided steps from the same noise end at cosine 0.999993 to the SDK's sample. - **Repeat:** identical bytes. -- **Speed**, one session: 0.273 s a forward (27.3 s for 50 guided steps), 3.4x the SDK's 0.939 s (the median of its +- **Speed**, one session: 0.275 s a forward (27.5 s for 50 guided steps), 3.4x the SDK's 0.939 s (the median of its 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. - **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the denoiser's fp16 weights in shared Metal buffers. diff --git a/zig/kernels/metal/idol_attention.metal b/zig/kernels/metal/idol_attention.metal index f676630fc5..8cc6919b33 100644 --- a/zig/kernels/metal/idol_attention.metal +++ b/zig/kernels/metal/idol_attention.metal @@ -79,7 +79,8 @@ inline void id_attention(const device T* Q, const device T* K, const device T* V TF_UNROLL for (short f = 0; f < 2; f++) { TF_UNROLL - for (short e = 0; e < 8; e++) s[f][e] = fast::exp2(s[f][e] - top_new[e >> 2]); + // the units take an fp32 P short: round it to fp16 here, so the denominator sums the values P V multiplies + for (short e = 0; e < 8; e++) s[f][e] = float(half(fast::exp2(s[f][e] - top_new[e >> 2]))); } float scale_old[2]; TF_UNROLL From 94a0392497f8dad77c0fd237e1880fcf007c7b70 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:03:17 +0800 Subject: [PATCH 06/21] idol: the recipe names the tensor units' fp32 operand truncation (#1772) --- CHANGELOG.md | 2 +- docs/recipes/idol.md | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a04a4c04eb..2abaa9af8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM is about fp16-accurate. The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 007f76fe6b..5b506262a9 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -43,8 +43,9 @@ This family keeps that split: ### The tensor units' fp32 matmul is not fp32 On the M5, a GEMM whose inputs are fp32 on the tensor units (`tfp::gemm_nax`, the Qwen-VL tower's -`qwv_gemm_f32`, MLX's fp32 matmul) rounds its inputs to about fp16: its error against float64 is about 7.5e-4 of the -result, against about 1e-7 for an fp32 FMA. RePo's maps at that precision move block 8's attention out of its gate. +`qwv_gemm_f32`, MLX's fp32 matmul) truncates each input to 10 mantissa bits, toward zero (#1772): its error against +float64 is about 7.5e-4 of the result, against about 1e-7 for an fp32 FMA. RePo's maps at that precision move +block 8's attention out of its gate. `id_norm_split` therefore writes each RePo input row as fp16 `[hi | hi | lo]` (`lo = x - hi`), and the map's weights are stored `[hi | lo | hi]`, scaled by a power of two so `lo` stays normal. One fp16 split-K GEMM with K = 3,072 then sums From 5d5e9fa86db874457603dfa9d619f9828250d38e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:52:06 +0800 Subject: [PATCH 07/21] idol: the FLUX.2 VAE image encoder natively on Metal, in fp32, with tf-idol-vae and its case Builds the denoiser's feature2 [4101, 128] from the preprocessed picture as idol-sdk's ImageEncoder does: the encoder, quant_conv, the posterior's sample with the SDK's noise as an input, the 2 x 2 patches, the batch norm's statistics and five zero rows. fp32 throughout on the SDK's bf16 weights (FLUX.2-klein-4B, or FLUX.2-dev rounded): Winograd 3x3s, implicit-GEMM downsamples and 1x1s, two-pass GroupNorm, single-head attention. idol_case.py vae records the SDK's MPS rows, its noise and a float64 ruler; tf-idol-vae gates worst row and 1 - cosine within twice the SDK's distance from float64 (ours 5.1e-5 and 5.0e-11 against 2.7e-2 and 5.1e-5), and repeat bytes. All eight kernels register with the contract: 13 forms pass gates 1-4, 12 controls red. --- CHANGELOG.md | 2 + docs/kernel-contract.md | 2 + docs/recipes/idol.md | 99 +++- tools/zig/idol_case.py | 148 ++++- zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol_vae.metal | 446 +++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/core/contract/idol_records.zig | 12 +- zig/src/families/idol/idol.zig | 8 +- zig/src/families/idol/vae.zig | 233 ++++++++ zig/src/families/idol/vae_kernels.zig | 154 ++++++ zig/src/families/idol/vae_weights.zig | 308 +++++++++++ zig/tests/contract/metal_idol_vae.zig | 176 ++++++ zig/tests/contract/metal_idol_vae_rows.zig | 437 +++++++++++++++ zig/tests/contract/metal_idol_vae_steps.zig | 578 ++++++++++++++++++++ zig/tests/contract/metal_idol_vae_wino.zig | 261 +++++++++ zig/tests/contract/metal_limits.zig | 2 + zig/tests/contract/metal_main.zig | 2 + zig/tests/idol_vae.zig | 132 +++++ 19 files changed, 2998 insertions(+), 5 deletions(-) create mode 100644 zig/kernels/metal/idol_vae.metal create mode 100644 zig/src/families/idol/vae.zig create mode 100644 zig/src/families/idol/vae_kernels.zig create mode 100644 zig/src/families/idol/vae_weights.zig create mode 100644 zig/tests/contract/metal_idol_vae.zig create mode 100644 zig/tests/contract/metal_idol_vae_rows.zig create mode 100644 zig/tests/contract/metal_idol_vae_steps.zig create mode 100644 zig/tests/contract/metal_idol_vae_wino.zig create mode 100644 zig/tests/idol_vae.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 2abaa9af8a..810a7ef6f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (after #1766, for #1757). Thanks to @jayleaton. + - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. - Qwen3.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b2b465117a..1d920643c7 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -188,6 +188,8 @@ HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row Metal also checks Idol's denoiser (`metal_idol*.zig`): its fp16 projections and GELU form, the split-K fp32 rows, `id_gemm32` (true fp32 sums: the NAX fp32 path in its place is a control gate 4 must catch), the LayerNorms, RePo's angles, the heads, the elementwise steps and the attention's four forms, each at idol-sdk's fp16 points (19 controls). +Its image encoder (`metal_idol_vae*.zig`, FLUX.2's VAE encoder in fp32) registers its convolutions (direct, implicit +GEMM and Winograd), GroupNorm's passes, the 512-wide attention and the condition rows against float64 (12 controls). RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 5b506262a9..a449405e51 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -15,7 +15,8 @@ each row's LayerNorm to per-head 3D positions, angles from 10, 10 and 12 frequen The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_cond - 3 v_uncond` (two passes a step: with the condition, and with zeros). -The image encoders, the decoders and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's +The FLUX.2 VAE image encoder runs natively too ([below](#vae-encoder)); DINOv3, the decoders and the PLY writer are +not ported yet. `tools/zig/idol_case.py` runs the publisher's SDK in torch, a development reference. It is a denoiser, not a token family: it does not register with the lane engine, the server or `families.detect`. @@ -129,3 +130,99 @@ Results: 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. - **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the denoiser's fp16 weights in shared Metal buffers. + +## VAE encoder + +`zig/src/families/idol/vae*.zig` (with `idol_vae.metal`) builds the denoiser's `feature2` [4101, 128] from the 1,024 x +1,024 picture `IdolPipeline.preprocess` returns, the way the SDK's `ImageEncoder._flux2_tokens` does: + +- FLUX.2's VAE encoder: `conv_in`, four `DownEncoderBlock2D` stages of two ResNet blocks (128, 256, 512, 512 wide; + GroupNorm 32 and SiLU) with a stride-2 3x3 downsample padded right and below after the first three, the mid block + (ResNet, single-head attention of 512 over the 16,384 positions, ResNet), `conv_norm_out`, SiLU, `conv_out` and + `quant_conv`, on `2 x - 1` of the picture; +- the posterior's **sample**, not its mean: `mean + exp(logvar / 2) * noise` with logvar clamped to [-30, 20]. Some + channels' standard deviations reach 1, so the noise matters. The SDK draws it from torch's MPS generator in bf16 after + `torch.manual_seed(seed)`, which `encode` cannot reproduce; it takes the noise [128, 128, 32] as an input, and the + case records the SDK's; +- each 2 x 2 patch of the 32 latent channels as channel `c * 4 + dy * 2 + dx` (128 channels on a 64 x 64 grid), less + the batch norm's `running_mean` over `sqrt(running_var + 1e-4)`; +- five zero rows, then the 4,096 patches in row order: 4,101 rows. + +### Weights and precision + +`black-forest-labs/FLUX.2-klein-4B` `vae/` at `e7b7dc27f91deacad38e78976d1f2b499d76a294` (Apache 2.0, ungated, bf16). +The SDK loads the VAE with `torch_dtype=torch.bfloat16`. Klein's file is FLUX.2-dev's fp32 VAE rounded to bf16, so both +give the same weights. The loader takes either: it rounds an fp32 file to bf16 as torch does. It reads the 106 encoder +tensors, `quant_conv` and the two batch-norm statistics, and refuses any other geometry. + +The SDK then runs the whole encoder in bf16. This port computes in fp32 on those bf16 weights: 3x3 convolutions by +Winograd F(2x2, 3x3) with `G g G^T` made in float64 and rounded once, the stride-2 downsamples and 1x1s as implicit GEMMs, +GroupNorm statistics in two passes (the mean, then the squared deviations), fp32 attention, `precise::exp`. None of it +uses the tensor units, whose fp32 GEMM is about fp16-accurate (above). + +### Gates + +```sh +hf download black-forest-labs/FLUX.2-klein-4B --revision e7b7dc27f91deacad38e78976d1f2b499d76a294 --include 'vae/*' --local-dir KLEIN +python3 tools/zig/idol_case.py vae KLEIN/vae SDK picture.png case/ --stages +zig build tf-idol-vae -Dcpu=apple_m1 +zig-out/bin/tf-idol-vae KLEIN/vae case/ [features.safetensors] +``` + +`idol_case.py vae` imports `idol/conditioning.py` and diffusers 0.42.1's VAE files only when their sha256 values are +the pinned ones, and the weights only when their sha256 is klein's or FLUX.2-dev's. It runs the SDK's own +`_flux2_tokens` on MPS, seeded as `IdolPipeline` seeds it, and records the noise it draws. Then it runs the same steps +in float64 on the CPU with that noise: the ruler, from the exact picture rather than its bf16 rounding. `--stages` adds +float64's output of each down block, the mid block and `quant_conv`. + +`tf-idol-vae` gates our rows against float64. Worst row (the largest relative L2 error of a row) and `1 - cosine` +must both be within twice the SDK's own distance, and the padding rows exactly zero. Five more encodes must repeat the +first byte for byte. `vae.zig`'s test plants a 2% error in one row, then a non-zero padding value, and requires the +gate to fail each. + +### Kernel contract + +`metal_idol_vae*.zig` registers all eight kernels in 13 forms, each against float64: `iv_conv_in`, `iv_conv` (1x1, +1x1 plus residual, stride-2 3x3), `iv_wino` (plain, plus residual), `iv_norm` (SiLU, plain), `iv_group_part` (sums, +squared deviations), `iv_group_finish`, `iv_attention` and `iv_tokens`. + +There are twelve controls: + +- inputs staged in half (`iv_conv_in`, the 1x1, Winograd's patch); +- symmetric padding in the stride-2 3x3; +- a dropped term of Winograd's output transform; +- a half SiLU; +- half sums (Winograd, both GroupNorm passes), caught at gate 4; +- attention without its rescale; +- swapped patch axes; +- a NaN logvar lost to Metal's `clamp`. `torch.clamp` keeps NaN, so `iv_tokens` does too. + +Winograd's float64 reference is its own transform, so a non-finite input reaches its whole 2 x 2 tile. The pictures +and 1x1 row counts are smaller than the engine's, and rows only tile. These kernels use no tensor ops; their records +reserve their thread counts in the source. + +``` +SUMMARY 13 kernels: 13 pass gates 1-4, 0 fail, 0 known findings; 12 controls: 12 red as planted, 0 missed; 0 limit problems +``` + +### Measured + +On the M5 Ultra, with the avocado example the SDK preprocessed, seed 42: + +| against float64 | worst row | 1 - cosine | +| --- | --- | --- | +| ours | 5.1e-5 | 5.0e-11 | +| the SDK (bf16, MPS) | 2.73e-2 | 5.05e-5 | + +- **Gate:** passes, by about 500 times on worst row. +- **Against the SDK:** cosine 0.999949, worst row 0.0275. The case's SDK rows equal the SDK's own recorded + `feature2` for this picture bit for bit, so its noise is the pipeline's. +- **Stages:** each down block, the mid block and the moments sit within 4e-5 (worst pixel) of float64. +- **Repeat:** identical bytes. +- **Speed:** 0.185 s an encode (median of five, 0.185 s of it on the GPU), where the SDK takes 0.095 s on torch MPS + (median of five, after a first call at 0.19 s). The SDK's bf16 runs on the tensor units; this fp32 path runs at 9 + to 14 TFLOP/s, and the attention, which reads K and V once every 8 queries, takes about 40 ms of it. Load: 0.3 to + 0.4 s. +- **Memory:** 3.28 GB peak footprint (`/usr/bin/time -l`, the whole `tf-idol-vae` run with its case and stage + copies), 2 GB of it four 512 MB picture tensors for the 1,024 x 1,024 x 128 stage. +- **FLUX.2-dev's fp32 file** in klein's place gives the same rows bit for bit, as rounding it to bf16 should. diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py index 83453a8e54..3bb1d378cd 100644 --- a/tools/zig/idol_case.py +++ b/tools/zig/idol_case.py @@ -14,8 +14,11 @@ case WEIGHTS SDK OUT --features F [F] noise, each step's input and both velocities, the final sample blocks WEIGHTS SDK CASE [--step S] every block's rows at step S, and each block in float32 on its input exact WEIGHTS SDK CASE [--steps N] both velocities in float32 on the case's first N inputs (the ruler) + vae VAE SDK IMAGE OUT [--stages] the FLUX.2 VAE rows of a preprocessed image: the SDK's, then float64's WEIGHTS is a speridlabs/idol snapshot; SDK an idol-sdk checkout at the pinned commit. F holds feature1 (DINOv3 rows [1, 4101, 1280]) and feature2 (FLUX.2 VAE rows [1, 4101, 128]), in one file or two, as the SDK's ImageEncoder returns. +VAE is the vae/ folder of black-forest-labs/FLUX.2-klein-4B at the pinned revision (FLUX.2-dev's gives the same bf16 +weights); IMAGE the 1024 x 1024 RGB picture IdolPipeline.preprocess returns. """ SDK_COMMIT = "4702a7f00850e34b6e83a2170cc123c3e9570744" @@ -25,6 +28,21 @@ "idol/sampler.py": "532987b5599be678361c2fcf566335bb3e1b0171696f985818d10b7460bfc8b0", } WEIGHTS_REVISION = "2731c2291fc65a797c5cb627a5e5526fc1490902" +CONDITIONING_SHA256 = "8c9cc50d327df3a5269cda392ef46cf4ed9c1e603f42f46050beda83f355f5ab" +DIFFUSERS_VERSION = "0.42.1" +DIFFUSERS_SHA256 = { + "models/autoencoders/autoencoder_kl_flux2.py": "1afe08d76a93e994a453bd9036b57d173516211c2e6b9794ea64f485bea57f57", + "models/autoencoders/vae.py": "59b51fdf2d951da45d697569271e0c99adf4a959ec0f91885f3d9314e925a12c", + "models/unets/unet_2d_blocks.py": "545d158c9d98fb8971c501303117fd3e32d27ae28506b89005c65548404f7437", + "models/resnet.py": "77ad799719fa51bf0fc703f0486915b791c9b69673983abeb3e32c6edcc76650", + "models/downsampling.py": "bce174aa8007b2c40787331e7f46b31eef1eab7785f16b4e54733553b3d14724", + "models/attention_processor.py": "59871fd4e66fc14a26f0e0464706ab24680e13dcd93161c55123678a5e671772", + "utils/torch_utils.py": "bce0a1a4488557e8f80dc02de5782c2b510e91ba2e07f7ae8d26168cee9afc2a", +} +# FLUX.2-klein-4B's vae/ (Apache 2.0, ungated) at its pinned revision, and FLUX.2-dev's: the same weights once in bf16 +VAE_REVISION = "e7b7dc27f91deacad38e78976d1f2b499d76a294" +VAE_SHA256 = {"ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04": "FLUX.2-klein-4B", + "d64f3a68e1cc4f9f4e29b6e0da38a0204fe9a49f2d4053f0ec1fa1ca02f9c4b5": "FLUX.2-dev"} DENOISER_ARGS = {"q_token_length": 8192, "in_channels": 16, "cam_channels": 5, "model_channels": 1024, "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, "num_refiner_blocks": 2, "num_heads": 16, "mlp_ratio": 4} @@ -216,9 +234,137 @@ def exact(args) -> int: return 0 +def flux2_vae(vae_dir: str, sdk_dir: str): + """The SDK's ImageEncoder and diffusers' FLUX.2 VAE (bf16, as the SDK loads it), each file used checked first.""" + import diffusers + import torch + + def pinned(path: Path, want: str, what: str) -> None: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != want: + sys.exit(f"{path} has sha256 {digest}; this tool runs only {what}'s ({want})") + + pinned(Path(sdk_dir) / "idol/conditioning.py", CONDITIONING_SHA256, f"idol-sdk {SDK_COMMIT[:12]}") + if diffusers.__version__ != DIFFUSERS_VERSION: + sys.exit(f"diffusers is {diffusers.__version__}; this tool runs only {DIFFUSERS_VERSION}") + root = Path(diffusers.__file__).parent + for name, want in DIFFUSERS_SHA256.items(): + pinned(root / name, want, f"diffusers {DIFFUSERS_VERSION}") + weights = Path(vae_dir) / "diffusion_pytorch_model.safetensors" + digest = hashlib.sha256(weights.read_bytes()).hexdigest() + if digest not in VAE_SHA256: + sys.exit(f"{weights} is neither FLUX.2-klein-4B {VAE_REVISION[:12]}'s vae/ nor FLUX.2-dev's") + sys.path.insert(0, str(Path(sdk_dir))) + from diffusers.models import AutoencoderKLFlux2 + from idol.conditioning import ImageEncoder + + vae = AutoencoderKLFlux2.from_pretrained(vae_dir, torch_dtype=torch.bfloat16) + return ImageEncoder, vae.eval().requires_grad_(False), VAE_SHA256[digest] + + +def tokens64(vae, pixels, noise): + """ImageEncoder._flux2_tokens's steps in float64 with the SDK's noise: the ruler, and each stage when asked.""" + import torch + + latents = vae.encode(pixels * 2 - 1).latent_dist + latents = latents.mean + latents.std * noise + batch, channels, height, width = latents.shape + latents = latents.view(batch, channels, height // 2, 2, width // 2, 2).permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch, channels * 4, height // 2, width // 2) + mean = vae.bn.running_mean.view(1, -1, 1, 1) + std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps) + rows = ((latents - mean) / std).flatten(2).transpose(1, 2) + return torch.cat([torch.zeros(batch, 5, rows.shape[2], dtype=rows.dtype), rows], dim=1) + + +def distance(ours, ruler) -> dict: + """1 - cosine and the worst row's relative L2 error against the ruler, the rows of 128 as tf-idol-vae reads them.""" + a, b = ours.double().reshape(-1, 128), ruler.double().reshape(-1, 128) + cos = float((a * b).sum() / (a.norm() * b.norm())) + rows = ((a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)).max() + return {"one_minus_cos": 1 - cos, "worst_row": float(rows), "max_abs": float((a - b).abs().max())} + + +def vae_case(args) -> int: + """The SDK's feature2 for one image on MPS (seeded as IdolPipeline seeds it), its noise, and the float64 ruler.""" + import types + + import diffusers.models.autoencoders.vae as vae_module + import numpy as np + import torch + from PIL import Image + from safetensors.torch import save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + image_encoder, vae, source = flux2_vae(args.vae, args.sdk) + picture = Image.open(args.image).convert("RGB") + if picture.size != (1024, 1024): + sys.exit(f"{args.image} is {picture.size}; the encoders see the 1024 x 1024 picture preprocess returns") + pixels = torch.from_numpy(np.array(picture)).permute(2, 0, 1).float().div(255.0)[None] + encoder = types.SimpleNamespace(flux2_vae=vae.to(args.device)) + drawn = [] + randn = vae_module.randn_tensor + + def recorded(*a, **k): + drawn.append(randn(*a, **k)) + return drawn[-1] + + vae_module.randn_tensor = recorded + times = [] + with torch.no_grad(): + for i in range(args.repeats + 1): + torch.manual_seed(args.seed) + mark = time.perf_counter() + rows = image_encoder._flux2_tokens(encoder, pixels.to(args.device)) + flush(args.device) + times.append(time.perf_counter() - mark) + if i == 0: + sdk_rows, noise = rows[0].cpu(), drawn[0].float().cpu() + vae_module.randn_tensor = randn + if any(not torch.equal(d.float().cpu(), noise) for d in drawn): + sys.exit("the seeded noise differs between encodes") + stages = {} + hooks = [] + if args.stages: + named = [(f"down.{i}", b) for i, b in enumerate(vae.encoder.down_blocks)] + [("mid", vae.encoder.mid_block)] + named.append(("moments", vae.quant_conv)) + for name, module in named: + def keep(_m, _i, output, name=name): + stages[name] = output[0].permute(1, 2, 0).float().contiguous() + hooks.append(module.register_forward_hook(keep)) + vae64 = vae.to("cpu").double() + with torch.no_grad(): + mark = time.perf_counter() + ruler = tokens64(vae64, pixels.double(), noise.double())[0] + ruler_s = time.perf_counter() - mark + for h in hooks: + h.remove() + tensors = {"image": pixels[0].permute(1, 2, 0).contiguous(), "noise": noise[0].permute(1, 2, 0).contiguous(), + "sdk": sdk_rows.float().contiguous(), "ruler": ruler.float().contiguous()} + save_file(tensors, str(out / "vae.safetensors")) + if stages: + save_file(stages, str(out / "vae_stages.safetensors")) + sdk = distance(sdk_rows, ruler) + meta = {"seed": args.seed, "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, + "diffusers": DIFFUSERS_VERSION, "vae": source, "vae_revision": VAE_REVISION, + "sdk_vs_float64": sdk, "float64_s": round(ruler_s, 2), "encode_s": [round(t, 4) for t in times]} + (out / "vae.json").write_text(json.dumps(meta, indent=1) + "\n") + warm = sorted(times[1:]) or times + print(json.dumps(meta | {"median_encode_s": warm[len(warm) // 2]})) + return 0 + + def main(argv: list[str]) -> int: ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("vae") + for name in ("vae", "sdk", "image", "out_dir"): + p.add_argument(name) + p.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--repeats", type=int, default=5, help="timed encodes after the first") + p.add_argument("--stages", action="store_true", help="float64's down blocks, mid block and moments too") for name in ("case", "blocks", "exact"): p = sub.add_parser(name) p.add_argument("weights") @@ -236,7 +382,7 @@ def main(argv: list[str]) -> int: else: p.add_argument("--steps", type=int, default=50) args = ap.parse_args(argv) - return {"case": case, "blocks": blocks, "exact": exact}[args.cmd](args) + return {"case": case, "blocks": blocks, "exact": exact, "vae": vae_case}[args.cmd](args) if __name__ == "__main__": diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index f08c980d2f..b2ed5b4efa 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -318,6 +318,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" }, .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" }, .{ .name = "tf-idol", .path = "zig/tests/idol.zig", .about = "Idol's denoiser on a reference case: each step's two branches compared, then the guided schedule timed" }, + .{ .name = "tf-idol-vae", .path = "zig/tests/idol_vae.zig", .about = "Idol's image encoder on a reference case: its condition rows gated against float64 beside the SDK's, then timed" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/idol_vae.metal b/zig/kernels/metal/idol_vae.metal new file mode 100644 index 0000000000..de517018b1 --- /dev/null +++ b/zig/kernels/metal/idol_vae.metal @@ -0,0 +1,446 @@ +// Idol's image encoder (FLUX.2's VAE encoder) in fp32, each picture tensor [rows, columns, channels]. + +#include +using namespace metal; + +#define TF_UNROLL_IV _Pragma("clang loop unroll(full)") + +struct IvConvIn { + uint rows; + uint cols; +}; + +// y [pixels, 128] = conv3x3(2 x - 1) + b of the picture x [pixels, 3] in [0, 1], zero padded; w [27, 128] by tap. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void iv_conv_in( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device float* y [[buffer(3)]], constant IvConvIn& p [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + if (at.y >= p.rows * p.cols) return; + const int py = int(at.y / p.cols); + const int px = int(at.y % p.cols); + float s = 0.0f; + TF_UNROLL_IV + for (int ky = 0; ky < 3; ky++) { + TF_UNROLL_IV + for (int kx = 0; kx < 3; kx++) { + const int iy = py + ky - 1; + const int ix = px + kx - 1; + if (iy < 0 || ix < 0 || iy >= int(p.rows) || ix >= int(p.cols)) continue; + const device float* v = x + ulong(uint(iy) * p.cols + uint(ix)) * 3; + TF_UNROLL_IV + for (int c = 0; c < 3; c++) s = fma(v[c] * 2.0f - 1.0f, w[((ky * 3 + kx) * 3 + c) * 128 + at.x], s); + } + } + y[ulong(at.y) * 128 + at.x] = s + b[at.x]; +} + +struct IvConv { + uint rows; + uint cols; + uint in_rows; + uint in_cols; + uint cin; + uint cout; + uint ldw; + uint taps; + uint add; +}; + +// y = conv(x) + b (+ r): a 1x1 (taps 1) or a 3x3 of stride 2 padded right and below (taps 9), w [taps * cin, ldw]. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void iv_conv( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device float* r [[buffer(3)]], device float* y [[buffer(4)]], constant IvConv& p [[buffer(5)]], + uint2 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint BK = 32; + constexpr uint lda = BK + 4; + // an implicit GEMM, x [in_rows, in_cols, cin] read a tap at a time: a 64 x 64 tile a threadgroup, fp32 sums + threadgroup float As[64 * lda]; + threadgroup float Cs[64 * 68]; + const uint pixels = p.rows * p.cols; + const uint m0 = tg.y * 64; + const uint n0 = tg.x * 64; + const uint arow = t / 2; + const uint acol = (t % 2) * 16; + const uint pix = m0 + arow; + const uint oy = pix / p.cols; + const uint ox = pix % p.cols; + const uint stride = p.taps == 9 ? 2 : 1; + const uint sr = (sg / 2) * 32; + const uint sc = (sg % 2) * 32; + simdgroup_float8x8 acc[4][4]; + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) acc[i][j] = simdgroup_float8x8(0.0f); + } + const uint K = p.taps * p.cin; + for (uint k0 = 0; k0 < K; k0 += BK) { + const uint tap = k0 / p.cin; + const uint iy = oy * stride + tap / 3; + const uint ix = ox * stride + tap % 3; + threadgroup float4* dst = (threadgroup float4*)(As + arow * lda + acol); + if (pix < pixels && iy < p.in_rows && ix < p.in_cols) { + const device float4* src = (const device float4*)(x + ulong(iy * p.in_cols + ix) * p.cin + k0 % p.cin + acol); + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) dst[i] = src[i]; + } else { + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) dst[i] = float4(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const device float* wk = w + ulong(k0) * p.ldw + n0 + sc; + TF_UNROLL_IV + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 am[4]; + simdgroup_float8x8 bm[4]; + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) simdgroup_load(am[i], As + (sr + i * 8) * lda + kk, lda); + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_load(bm[j], wk + ulong(kk) * p.ldw + j * 8, p.ldw); + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(acc[i][j], am[i], bm[j], acc[i][j]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + TF_UNROLL_IV + for (uint i = 0; i < 4; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_store(acc[i][j], Cs + (sr + i * 8) * 68 + sc + j * 8, 68); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = t; e < 64 * 64; e += 128) { + const uint row = m0 + e / 64; + const uint col = n0 + e % 64; + if (row >= pixels || col >= p.cout) continue; + const ulong at = ulong(row) * p.cout + col; + float v = Cs[(e / 64) * 68 + e % 64] + b[col]; + if (p.add) v = v + r[at]; + y[at] = v; + } +} + +// The patch rows (or columns) a, b and b's sign that B^T's row i combines: (0, 2, -), (1, 2, +), (2, 1, -), (1, 3, -). +inline uint iv_wa(uint i) { return i == 2 ? 2 : (i == 0 ? 0 : 1); } +inline uint iv_wb(uint i) { return i == 0 ? 2 : (i == 3 ? 3 : (i == 1 ? 2 : 1)); } +inline float iv_ws(uint i) { return i == 1 ? 1.0f : -1.0f; } + +// y = conv3x3(x) + b (+ r) by Winograd F(2x2, 3x3), u = G g G^T [16, cin, ldw]; 16 tiles by 64 channels a group. +[[max_total_threads_per_threadgroup(512)]] [[kernel]] void iv_wino( + const device float* x [[buffer(0)]], const device float* u [[buffer(1)]], const device float* b [[buffer(2)]], + const device float* r [[buffer(3)]], device float* y [[buffer(4)]], constant IvConv& p [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint NF = 8; + constexpr uint BK = 16; + constexpr uint TILES = 16; + constexpr uint PR = 6; + constexpr uint PC = 18; + // the input patch [6, 18, 16]; the epilogue's [16 points, 16 tiles, 32 channels] after + threadgroup float buf[16 * 16 * 32]; + const uint n0 = tg.x * NF * 8; + const int y0 = int(tg.y) * 4; + const int x0 = int(tg.z) * 16; + // this lane's two elements of every 8 x 8 fragment: row fm, columns fn and fn + 1 + const uint fm = ((lane / 4) & 4) + (lane / 2) % 4; + const uint fn = ((lane / 4) & 2) * 2 + (lane % 2) * 2; + const uint pi = sg / 4; + const uint pj = sg % 4; + simdgroup_float8x8 acc[2][NF]; + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) { + TF_UNROLL_IV + for (uint j = 0; j < NF; j++) acc[i][j] = simdgroup_float8x8(0.0f); + } + for (uint c0 = 0; c0 < p.cin; c0 += BK) { + for (uint e = t; e < PR * PC * (BK / 4); e += 512) { + const uint px = e / (BK / 4); + const uint q = e % (BK / 4); + const int iy = y0 - 1 + int(px / PC); + const int ix = x0 - 1 + int(px % PC); + float4 v = float4(0.0f); + if (iy >= 0 && ix >= 0 && iy < int(p.rows) && ix < int(p.cols)) + v = *(const device float4*)(x + ulong(uint(iy) * p.cols + uint(ix)) * p.cin + c0 + q * 4); + *(threadgroup float4*)(buf + px * BK + q * 4) = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + TF_UNROLL_IV + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 vm[2]; + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) { + // B^T d B at this point, for tile i * 8 + fm and channels fn, fn + 1 + const uint tile = i * 8 + fm; + const threadgroup float* d = buf + (((tile / 8) * 2) * PC + (tile % 8) * 2) * BK + kk + fn; + TF_UNROLL_IV + for (uint c = 0; c < 2; c++) { + const threadgroup float* q = d + c; + const float wa = q[(iv_wa(pi) * PC + iv_wa(pj)) * BK] + iv_ws(pi) * q[(iv_wb(pi) * PC + iv_wa(pj)) * BK]; + const float wb = q[(iv_wa(pi) * PC + iv_wb(pj)) * BK] + iv_ws(pi) * q[(iv_wb(pi) * PC + iv_wb(pj)) * BK]; + vm[i].thread_elements()[c] = wa + iv_ws(pj) * wb; + } + } + const device float* up_ = u + (ulong(sg) * p.cin + c0 + kk) * p.ldw + n0; + TF_UNROLL_IV + for (uint j = 0; j < NF; j++) { + simdgroup_float8x8 um; + simdgroup_load(um, up_ + j * 8, p.ldw); + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) simdgroup_multiply_accumulate(acc[i][j], vm[i], um, acc[i][j]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + // the output transform A^T M A, 32 channels at a time + for (uint h = 0; h < NF / 4; h++) { + TF_UNROLL_IV + for (uint i = 0; i < 2; i++) { + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) simdgroup_store(acc[i][h * 4 + j], buf + (sg * TILES + i * 8) * 32 + j * 8, 32); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + { + // one tile and channel a thread: its 16 points, then the four outputs A^T M A, rows then columns + const uint co = t % 32; + const uint tile = t / 32; + const uint col = n0 + h * 32 + co; + const threadgroup float* m = buf + tile * 32 + co; + float s[4][2]; + TF_UNROLL_IV + for (uint j = 0; j < 4; j++) { + const float m0 = m[(0 * 4 + j) * TILES * 32]; + const float m1 = m[(1 * 4 + j) * TILES * 32]; + const float m2 = m[(2 * 4 + j) * TILES * 32]; + const float m3 = m[(3 * 4 + j) * TILES * 32]; + s[j][0] = (m0 + m1) + m2; + s[j][1] = (m1 - m2) - m3; + } + const int ty = y0 + int(tile / 8) * 2; + const int tx = x0 + int(tile % 8) * 2; + TF_UNROLL_IV + for (uint rr = 0; rr < 2; rr++) { + const float out[2] = {(s[0][rr] + s[1][rr]) + s[2][rr], (s[1][rr] - s[2][rr]) - s[3][rr]}; + TF_UNROLL_IV + for (uint cc = 0; cc < 2; cc++) { + const int oy = ty + int(rr); + const int ox = tx + int(cc); + if (col >= p.cout || oy >= int(p.rows) || ox >= int(p.cols)) continue; + const ulong at = ulong(uint(oy) * p.cols + uint(ox)) * p.cout + col; + float v = out[cc] + b[col]; + if (p.add) v = v + r[at]; + y[at] = v; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +struct IvNorm { + uint pixels; + uint channels; + uint silu; +}; + +// y = GroupNorm(x) by stats (mean, rstd per group) with g, beta, then SiLU when silu; x, y [pixels, channels]. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void iv_norm( + const device float* x [[buffer(0)]], const device float* stats [[buffer(1)]], const device float* g [[buffer(2)]], + const device float* beta [[buffer(3)]], device float* y [[buffer(4)]], constant IvNorm& p [[buffer(5)]], + uint2 at [[thread_position_in_grid]]) { + if (at.y >= p.pixels || at.x >= p.channels) return; + const uint grp = at.x / (p.channels / 32); + const ulong i = ulong(at.y) * p.channels + at.x; + float v = (x[i] - stats[grp]) * stats[32 + grp]; + v = v * g[at.x] + beta[at.x]; + y[i] = p.silu ? v / (1.0f + precise::exp(-v)) : v; +} + +struct IvGroups { + uint pixels; + uint channels; + uint chunk; + uint chunks; + uint centered; + float eps; +}; + +// One chunk of pixels' sums of x a group, or of (x - mean)^2 when centered, into part [32, chunks]: a thread a channel. +[[max_total_threads_per_threadgroup(512)]] [[kernel]] void iv_group_part( + const device float* x [[buffer(0)]], const device float* stats [[buffer(1)]], device float* part [[buffer(2)]], + constant IvGroups& p [[buffer(3)]], uint tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]]) { + threadgroup float sums[512]; + const uint size = p.channels / 32; + const uint first = tg * p.chunk; + const uint last = min(first + p.chunk, p.pixels); + const float mean = p.centered ? stats[t / size] : 0.0f; + float s = 0.0f; + for (uint px = first; px < last; px++) { + const float v = x[ulong(px) * p.channels + t]; + const float d = v - mean; + s += p.centered ? d * d : v; + } + sums[t] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < 32) { + float all = 0.0f; + for (uint c = 0; c < size; c++) all += sums[t * size + c]; + part[t * p.chunks + tg] = all; + } +} + +// stats from part's chunk sums, a simdgroup a group: mean[32] when not centered, else rstd[32] = rsqrt(var + eps). +[[max_total_threads_per_threadgroup(1024)]] [[kernel]] void iv_group_finish( + const device float* part [[buffer(0)]], device float* stats [[buffer(1)]], constant IvGroups& p [[buffer(2)]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + float s = 0.0f; + for (uint i = lane; i < p.chunks; i += 32) s += part[sg * p.chunks + i]; + s = simd_sum(s); + if (lane != 0) return; + const float n = float(p.pixels) * float(p.channels / 32); + if (p.centered) + stats[32 + sg] = metal::precise::rsqrt(s / n + p.eps); + else + stats[sg] = s / n; +} + +struct IvAttention { + uint rows; + float scale; +}; + +// o [rows, 512] = softmax(q k^T * scale) v of qkv rows [rows, 1536] (both padded to 32 rows): 8 queries a group. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void iv_attention( + const device float* qkv [[buffer(0)]], device float* o [[buffer(1)]], constant IvAttention& p [[buffer(2)]], + uint tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint D = 512; + constexpr uint Q = D / 4; + constexpr uint ld = 3 * D; + constexpr uint BK = 32; + threadgroup float Sp[4][8 * BK]; + threadgroup float Ps[8 * BK]; + threadgroup float As[8]; + threadgroup uint rescale; + // the rows of this lane's two elements of every 8 x 8 fragment + const uint fm = ((lane / 4) & 4) + (lane / 2) % 4; + const uint q0 = tg * 8; + const uint d0 = sg * Q; + simdgroup_float8x8 qm[Q / 8]; + simdgroup_float8x8 acc[Q / 8]; + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + simdgroup_load(qm[d], qkv + ulong(q0) * ld + d0 + d * 8, ld); + acc[d] = simdgroup_float8x8(0.0f); + } + const uint row = t / 16; + const uint col = (t % 16) * 2; + float best = -INFINITY; + float total = 0.0f; + for (uint k0 = 0; k0 < p.rows; k0 += BK) { + const device float* kp = qkv + ulong(k0) * ld + D + d0; + TF_UNROLL_IV + for (uint j = 0; j < BK / 8; j++) { + simdgroup_float8x8 s = simdgroup_float8x8(0.0f); + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + simdgroup_float8x8 km; + simdgroup_load(km, kp + ulong(j * 8) * ld + d * 8, ld, ulong2(0, 0), true); + simdgroup_multiply_accumulate(s, qm[d], km, s); + } + simdgroup_store(s, Sp[sg] + j * 8, BK); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float v[2]; + float most = best; + TF_UNROLL_IV + for (uint c = 0; c < 2; c++) { + const uint i = row * BK + col + c; + const float sum = (Sp[0][i] + Sp[1][i]) + (Sp[2][i] + Sp[3][i]); + v[c] = k0 + col + c < p.rows ? sum * p.scale : -INFINITY; + most = max(most, v[c]); + } + TF_UNROLL_IV + for (ushort m = 1; m < 16; m <<= 1) most = max(most, simd_shuffle_xor(most, m)); + const float alpha = precise::exp(best - most); + float add = 0.0f; + TF_UNROLL_IV + for (uint c = 0; c < 2; c++) { + const float e = precise::exp(v[c] - most); + Ps[row * BK + col + c] = e; + add += e; + } + TF_UNROLL_IV + for (ushort m = 1; m < 16; m <<= 1) add += simd_shuffle_xor(add, m); + total = total * alpha + add; + best = most; + if (t % 16 == 0) As[row] = alpha; + if (t == 0) rescale = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t % 16 == 0 && alpha != 1.0f) rescale = 1; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (rescale) { + const float f = As[fm]; + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + acc[d].thread_elements()[0] *= f; + acc[d].thread_elements()[1] *= f; + } + } + const device float* vp = qkv + ulong(k0) * ld + 2 * D + d0; + TF_UNROLL_IV + for (uint j = 0; j < BK / 8; j++) { + simdgroup_float8x8 pm; + simdgroup_load(pm, Ps + j * 8, BK); + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + simdgroup_float8x8 vm; + simdgroup_load(vm, vp + ulong(j * 8) * ld + d * 8, ld); + simdgroup_multiply_accumulate(acc[d], pm, vm, acc[d]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t % 16 == 0) As[row] = 1.0f / total; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float inv = As[fm]; + TF_UNROLL_IV + for (uint d = 0; d < Q / 8; d++) { + acc[d].thread_elements()[0] *= inv; + acc[d].thread_elements()[1] *= inv; + simdgroup_store(acc[d], o + ulong(q0) * D + d0 + d * 8, D); + } +} + +struct IvTokens { + uint rows; + uint cols; + uint pad; +}; + +// Rows [pad + patches, 128]: pad zeros, then each 2 x 2 patch of z = mean + std * noise as (z - m) / s. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void iv_tokens( + const device float* moments [[buffer(0)]], const device float* noise [[buffer(1)]], + const device float* m [[buffer(2)]], const device float* s [[buffer(3)]], device float* out [[buffer(4)]], + constant IvTokens& p [[buffer(5)]], uint2 at [[thread_position_in_grid]]) { + const uint w2 = p.cols / 2; + if (at.x >= 128 || at.y >= p.pad + (p.rows / 2) * w2) return; + const ulong o = ulong(at.y) * 128 + at.x; + if (at.y < p.pad) { + out[o] = 0.0f; + return; + } + // moments [rows, cols, 64] as mean | logvar, noise [rows, cols, 32]; a patch's channel c * 4 + dy * 2 + dx + const uint patch = at.y - p.pad; + const uint c = at.x / 4; + const uint py = (patch / w2) * 2 + (at.x % 4) / 2; + const uint px = (patch % w2) * 2 + at.x % 2; + const uint pixel = py * p.cols + px; + // torch.clamp keeps a NaN, where Metal's clamp (fmin, fmax) would drop it + const float raw = moments[ulong(pixel) * 64 + 32 + c]; + const float logvar = isnan(raw) ? raw : clamp(raw, -30.0f, 20.0f); + const float z = moments[ulong(pixel) * 64 + c] + precise::exp(0.5f * logvar) * noise[ulong(pixel) * 32 + c]; + out[o] = (z - m[at.x]) / s[at.x]; +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index ef3221a88b..51aee74e04 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -135,6 +135,8 @@ pub const qwen_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile(" pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal"); /// Idol's denoiser: the prefill GEMM template, then its own steps and head-64 attention. pub const idol = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol.metal") ++ @embedFile("idol_attention.metal"); +/// Idol's image encoder (FLUX.2's VAE encoder) in fp32. +pub const idol_vae = @embedFile("idol_vae.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig index ae39549164..1644357cef 100644 --- a/zig/src/core/contract/idol_records.zig +++ b/zig/src/core/contract/idol_records.zig @@ -1,4 +1,4 @@ -//! Idol's denoiser launches (idol.metal and idol_attention.metal after gemm_nax.metal), each source-reserved. +//! Idol's launches: the denoiser's (idol.metal, idol_attention.metal after gemm_nax.metal), the image encoder's. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; pub const records = [_]M{ @@ -26,9 +26,17 @@ pub const records = [_]M{ .{ .name = "idol id_attention_01", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "idol id_attention_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "idol id_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol vae iv_conv_in", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol vae iv_conv", .threads = 128, .tg_memory = 26624, .reserved = 128 }, + .{ .name = "idol vae iv_wino", .threads = 512, .tg_memory = 32768, .reserved = 512 }, + .{ .name = "idol vae iv_norm", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol vae iv_group_part", .threads = 512, .tg_memory = 2048, .reserved = 512 }, + .{ .name = "idol vae iv_group_finish", .threads = 1024, .tg_memory = 0, .reserved = 1024 }, + .{ .name = "idol vae iv_attention", .threads = 128, .tg_memory = 5168, .reserved = 128 }, + .{ .name = "idol vae iv_tokens", .threads = 256, .tg_memory = 0, .reserved = 256 }, }; test "the Idol launches fit every Apple family, one record a pipeline" { try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); - try std.testing.expectEqual(@as(usize, 24), records.len); + try std.testing.expectEqual(@as(usize, 32), records.len); } diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig index a2aff60944..70fa00e6e6 100644 --- a/zig/src/families/idol/idol.zig +++ b/zig/src/families/idol/idol.zig @@ -1,10 +1,13 @@ -//! Idol image to 3D: its denoiser on Metal (a denoiser beside the lane engine, not a token family). +//! Idol image to 3D: its denoiser and image encoder on Metal (beside the lane engine, not a token family). pub const config = @import("config.zig"); pub const host = @import("host.zig"); pub const weights = @import("weights.zig"); pub const kernels = @import("kernels.zig"); pub const model = @import("model.zig"); pub const reference = @import("reference.zig"); +pub const vae = @import("vae.zig"); +pub const vae_kernels = @import("vae_kernels.zig"); +pub const vae_weights = @import("vae_weights.zig"); pub const Model = model.Model; pub const Condition = model.Condition; pub const Timing = model.Timing; @@ -16,4 +19,7 @@ test { _ = kernels; _ = model; _ = reference; + _ = vae; + _ = vae_kernels; + _ = vae_weights; } diff --git a/zig/src/families/idol/vae.zig b/zig/src/families/idol/vae.zig new file mode 100644 index 0000000000..890cda28e4 --- /dev/null +++ b/zig/src/families/idol/vae.zig @@ -0,0 +1,233 @@ +//! Idol's image encoder on Metal: FLUX.2's VAE encoder in fp32, a 1024 x 1024 picture to the condition's 4,101 rows. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("vae_kernels.zig"); +const wts = @import("vae_weights.zig"); +const ref = @import("reference.zig"); +const Ref = ks.Ref; + +/// The picture the encoder sees, its latent grid, and the condition rows: 5 zero rows, then a row a 2 x 2 patch. +pub const size = 1024; +pub const latent = size / 8; +pub const pad = 5; +pub const rows = pad + (latent / 2) * (latent / 2); +pub const channels = 4 * wts.latent_channels; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +/// Where `encode` copies each stage's output when asked, [rows, columns, channels] fp32 (null: not wanted). +pub const Stages = struct { down: [4]?[]f32 = @splat(null), mid: ?[]f32 = null, moments: ?[]f32 = null }; + +/// The widest picture tensor: the first stage's 1024 x 1024 x 128. +const most = size * size * wts.widths[0]; + +pub const Encoder = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + scratch: mtl.Buffer, + image: Ref, + noise: Ref, + t: [4]Ref, + moments: Ref, + out: Ref, + stats: Ref, + part: Ref, + + /// The vae/ folder `dir` (FLUX.2-klein-4B's, or FLUX.2-dev's rounded to bf16 as the SDK loads it). + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !*Encoder { + const v = try gpa.create(Encoder); + errdefer gpa.destroy(v); + v.* = undefined; + v.gpa = gpa; + v.device = device; + v.queue = try device.queue(); + errdefer v.queue.deinit(); + v.k = try ks.Kernels.load(device); + errdefer v.k.deinit(); + v.w = try wts.load(gpa, io, device, dir); + errdefer v.w.deinit(); + const L = latent * latent; + const parts = [_]usize{ size * size * 3, L * wts.latent_channels, most, most, most, most, L * 2 * wts.latent_channels, rows * channels, 2 * ks.groups, ks.groups * ks.chunks(size * size) }; + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n * 4, 256); + v.scratch = try device.buffer(total, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &v.image, &v.noise, &v.t[0], &v.t[1], &v.t[2], &v.t[3], &v.moments, &v.out, &v.stats, &v.part }; + for (views, parts) |view, n| { + view.* = .{ .buf = v.scratch, .off = at }; + at += std.mem.alignForward(usize, n * 4, 256); + } + return v; + } + + pub fn deinit(v: *Encoder) void { + v.scratch.deinit(); + v.w.deinit(); + v.k.deinit(); + v.queue.deinit(); + v.gpa.destroy(v); + } + + /// The condition rows [4101, 128] of `image` [1024, 1024, 3] in [0, 1] with the posterior's `noise` [128, 128, 32]. + pub fn encode(v: *Encoder, image: []const f32, noise: []const f32, out: []f32, stages: Stages) !Timing { + if (image.len != size * size * 3) return error.PictureSize; + if (noise.len != latent * latent * wts.latent_channels or out.len != rows * channels) return error.LatentRows; + @memcpy(floats(v.image, image.len), image); + @memcpy(floats(v.noise, noise.len), noise); + const started = mtl.clock.seconds(); + var gpu: f64 = 0; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var cb = v.queue.commandBuffer(); + var e = cb.compute(.serial); + var R: usize = size; + ks.convIn(e, v.k.get("iv_conv_in"), v.image, v.w.conv_in.w, v.w.conv_in.b, v.t[0], R, R); + var x: usize = 0; + for (v.w.stages, 0..) |s, i| { + for (s.resnets) |r| x = v.resnet(e, r, x, R); + if (s.down) |d| { + const y = (x + 1) % v.t.len; + ks.conv(e, v.k.get("iv_conv"), v.t[x], d.w, d.b, null, v.t[y], .{ .rows = R / 2, .cols = R / 2, .in_rows = R, .in_cols = R, .cin = d.cin, .cout = d.cout, .ldw = d.ldw, .taps = 9 }); + x = y; + R /= 2; + } + if (stages.down[i]) |dst| { + gpu += try finish(cb, e); + @memcpy(dst, floats(v.t[x], dst.len)); + cb = v.queue.commandBuffer(); + e = cb.compute(.serial); + } + } + x = v.resnet(e, v.w.mid[0], x, R); + x = v.attention(e, x, R * R); + x = v.resnet(e, v.w.mid[1], x, R); + if (stages.mid) |dst| { + gpu += try finish(cb, e); + @memcpy(dst, floats(v.t[x], dst.len)); + cb = v.queue.commandBuffer(); + e = cb.compute(.serial); + } + const n = (x + 1) % v.t.len; + const o = (x + 2) % v.t.len; + v.norm(e, v.t[x], v.t[n], R * R, wts.widths[3], v.w.norm_out, true); + v.conv(e, v.w.conv_out, v.t[n], v.t[o], R, null); + v.conv(e, v.w.quant, v.t[o], v.moments, R, null); + ks.tokens(e, v.k.get("iv_tokens"), v.moments, v.noise, v.w.bn_mean, v.w.bn_std, v.out, R, R, pad); + gpu += try finish(cb, e); + if (stages.moments) |dst| @memcpy(dst, floats(v.moments, dst.len)); + @memcpy(out, floats(v.out, out.len)); + return .{ .wall = mtl.clock.seconds() - started, .gpu = gpu }; + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !f64 { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol vae: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return cb.gpuSeconds(); + } + + /// One residual block on t[x] at R x R, its output in another of t; returns that one's index. + fn resnet(v: *const Encoder, e: mtl.ComputeEncoder, r: wts.Resnet, x: usize, R: usize) usize { + const n = (x + 1) % v.t.len; + const h = (x + 2) % v.t.len; + const s = (x + 3) % v.t.len; + v.norm(e, v.t[x], v.t[n], R * R, r.conv1.cin, r.norm1, true); + v.conv(e, r.conv1, v.t[n], v.t[h], R, null); + v.norm(e, v.t[h], v.t[n], R * R, r.conv1.cout, r.norm2, true); + const skip = if (r.shortcut) |sc| blk: { + v.conv(e, sc, v.t[x], v.t[s], R, null); + break :blk v.t[s]; + } else v.t[x]; + v.conv(e, r.conv2, v.t[n], v.t[h], R, skip); + return h; + } + + /// The mid block's self-attention on t[x] (`n` positions): GroupNorm, q | k | v, one head of 512, out + t[x]. + fn attention(v: *const Encoder, e: mtl.ComputeEncoder, x: usize, n: usize) usize { + const xn = (x + 1) % v.t.len; + const qkv = v.t[(x + 2) % v.t.len]; + const att = v.t[(x + 3) % v.t.len]; + const W = ks.width; + v.norm(e, v.t[x], v.t[xn], n, W, v.w.attn_norm, false); + ks.conv(e, v.k.get("iv_conv"), v.t[xn], v.w.qkv.w, v.w.qkv.b, null, qkv, .{ .rows = n, .cols = 1, .in_rows = n, .in_cols = 1, .cin = W, .cout = 3 * W, .ldw = 3 * W }); + ks.attend(e, v.k.get("iv_attention"), qkv, att, n); + ks.conv(e, v.k.get("iv_conv"), att, v.w.attn_out.w, v.w.attn_out.b, v.t[x], v.t[xn], .{ .rows = n, .cols = 1, .in_rows = n, .in_cols = 1, .cin = W, .cout = W, .ldw = W }); + return xn; + } + + /// y = GroupNorm(x) with `nm` (then SiLU when `silu`), x and y [pixels, channels], statistics in two passes. + fn norm(v: *const Encoder, e: mtl.ComputeEncoder, x: Ref, y: Ref, n: usize, ch: usize, nm: wts.Norm, silu: bool) void { + for ([_]bool{ false, true }) |centered| { + ks.groupPart(e, v.k.get("iv_group_part"), x, v.stats, v.part, n, ch, centered); + ks.groupFinish(e, v.k.get("iv_group_finish"), v.part, v.stats, n, ch, centered); + } + ks.normed(e, v.k.get("iv_norm"), x, v.stats, nm.g, nm.b, y, n, ch, silu); + } + + /// y = conv(x) + b (+ r) at R x R, a same-size 3x3 (Winograd) or a 1x1. + fn conv(v: *const Encoder, e: mtl.ComputeEncoder, c: wts.Conv, x: Ref, y: Ref, R: usize, r: ?Ref) void { + const s: ks.ConvShape = .{ .rows = R, .cols = R, .in_rows = R, .in_cols = R, .cin = c.cin, .cout = c.cout, .ldw = c.ldw }; + switch (c.kind) { + .wino => ks.winograd(e, v.k.get("iv_wino"), x, c.w, c.b, r, y, s), + .one => ks.conv(e, v.k.get("iv_conv"), x, c.w, c.b, r, y, s), + .strided, .input => unreachable, + } + } +}; + +fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; +} + +/// The encoder's gate against the float64 ruler: worst row and 1 - cosine within twice the SDK's own, rows of 128. +pub const Verdict = struct { ours: ref.Match, sdk: ref.Match, pass: bool }; + +pub fn gate(ours: []const f32, sdk: []const f32, ruler: []const f32) Verdict { + const o = ref.match(ours[pad * channels ..], ruler[pad * channels ..], channels); + const s = ref.match(sdk[pad * channels ..], ruler[pad * channels ..], channels); + const zeros = for (ours[0 .. pad * channels]) |z| { + if (z != 0) break false; + } else true; + return .{ .ours = o, .sdk = s, .pass = zeros and ref.within(o, s) }; +} + +test "the encoder's scratch fits a GPU buffer and the rows are the denoiser's condition rows" { + try std.testing.expectEqual(@as(usize, 4101), rows); + try std.testing.expectEqual(@as(usize, 128), channels); + try std.testing.expect(4 * most * 4 < 1 << 33); +} + +test "a planted mismatch in one condition row fails the encoder's gate, as does a padding row not zero" { + const a = std.testing.allocator; + const n = rows * channels; + const ruler = try a.alloc(f32, n); + defer a.free(ruler); + const sdk = try a.alloc(f32, n); + defer a.free(sdk); + const ours = try a.alloc(f32, n); + defer a.free(ours); + var rng = std.Random.DefaultPrng.init(11); + for (ruler, sdk, ours, 0..) |*r, *s, *o, i| { + const z = i < pad * channels; + r.* = if (z) 0 else rng.random().floatNorm(f32); + s.* = r.* * (1 + 4e-3 * rng.random().floatNorm(f32)); + o.* = r.* * (1 + 1e-6 * rng.random().floatNorm(f32)); + } + try std.testing.expect(gate(ours, sdk, ruler).pass); + const row = 2000; + for (ours[row * channels ..][0..channels]) |*o| o.* *= 1.02; + const v = gate(ours, sdk, ruler); + try std.testing.expect(!v.pass); + try std.testing.expect(v.ours.worst_row > 0.019); + for (ours[row * channels ..][0..channels], ruler[row * channels ..][0..channels]) |*o, r| o.* = r; + try std.testing.expect(gate(ours, sdk, ruler).pass); + ours[3 * channels + 7] = 1e-3; + try std.testing.expect(!gate(ours, sdk, ruler).pass); +} diff --git a/zig/src/families/idol/vae_kernels.zig b/zig/src/families/idol/vae_kernels.zig new file mode 100644 index 0000000000..091005fee8 --- /dev/null +++ b/zig/src/families/idol/vae_kernels.zig @@ -0,0 +1,154 @@ +//! Idol's image encoder pipelines (idol_vae.metal, fp32): convolutions, GroupNorm, attention and the condition rows. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// Each pipeline by its function. +pub const names = [_][]const u8{ "iv_conv_in", "iv_conv", "iv_wino", "iv_norm", "iv_group_part", "iv_group_finish", "iv_attention", "iv_tokens" }; + +/// GroupNorm's groups and epsilon, the pixels a statistics part sums, and the mid block's attention width. +pub const groups = 32; +pub const eps: f32 = 1e-6; +pub const chunk = 256; +pub const width = 512; + +pub const Kernels = struct { + lib: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(device: mtl.Device) !Kernels { + return loadSource(device, sources.idol_vae); + } + + /// The pipelines of `text` (the served source, or a contract control's copy of it). + pub fn loadSource(device: mtl.Device, text: []const u8) !Kernels { + var k: Kernels = undefined; + k.lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + errdefer k.lib.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.lib, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.lib.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } +}; + +/// One kernel's buffers in order, then its parameters at the next index. +pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); +} + +/// iv_conv_in: y [rows * cols, 128] = conv3x3(2 x - 1) + b of the picture x [rows * cols, 3], w [27, 128]. +pub fn convIn(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, y: Ref, rows: usize, cols: usize) void { + bind(e, p, &.{ x, w, b, y }, [2]u32{ @intCast(rows), @intCast(cols) }); + e.dispatchThreads(mtl.Size.of(128, rows * cols, 1), mtl.Size.of(128, 2, 1)); +} + +/// A convolution's size: the output picture, its input picture, channels, the weights' row length and taps (1 or 9). +pub const ConvShape = struct { rows: usize, cols: usize, in_rows: usize, in_cols: usize, cin: usize, cout: usize, ldw: usize, taps: usize = 1 }; + +fn convParams(s: ConvShape, add: bool) extern struct { rows: u32, cols: u32, in_rows: u32, in_cols: u32, cin: u32, cout: u32, ldw: u32, taps: u32, add: u32 } { + return .{ .rows = @intCast(s.rows), .cols = @intCast(s.cols), .in_rows = @intCast(s.in_rows), .in_cols = @intCast(s.in_cols), .cin = @intCast(s.cin), .cout = @intCast(s.cout), .ldw = @intCast(s.ldw), .taps = @intCast(s.taps), .add = @intFromBool(add) }; +} + +/// iv_conv: y = conv(x) + b (+ r when given), a 1x1 or a stride-2 3x3; cin % 32 = 0, ldw % 64 = 0. +pub fn conv(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + std.debug.assert(s.cin % 32 == 0 and s.ldw % 64 == 0 and (s.taps == 1 or s.taps == 9)); + bind(e, p, &.{ x, w, b, r orelse y, y }, convParams(s, r != null)); + e.dispatchGroups(mtl.Size.of(s.ldw / 64, (s.rows * s.cols + 63) / 64, 1), mtl.Size.of(128, 1, 1)); +} + +/// iv_wino: y = conv3x3(x) + b (+ r when given) by Winograd F(2x2, 3x3), same size; cin % 16 = 0, ldw % 64 = 0. +pub fn winograd(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, u: Ref, b: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + std.debug.assert(s.cin % 16 == 0 and s.ldw % 64 == 0); + bind(e, p, &.{ x, u, b, r orelse y, y }, convParams(s, r != null)); + e.dispatchGroups(mtl.Size.of(s.ldw / 64, (s.rows + 3) / 4, (s.cols + 15) / 16), mtl.Size.of(512, 1, 1)); +} + +/// The Winograd F(2x2, 3x3) form of one 3x3 kernel g: G g G^T in float64, each of its 16 values rounded once. +pub fn transform(g: *const [9]f32) [16]f32 { + const G = [4][3]f64{ .{ 1, 0, 0 }, .{ 0.5, 0.5, 0.5 }, .{ 0.5, -0.5, 0.5 }, .{ 0, 0, 1 } }; + var out: [16]f32 = undefined; + for (0..4) |i| for (0..4) |j| { + var v: f64 = 0; + for (0..3) |a| for (0..3) |b| { + v += G[i][a] * @as(f64, g[a * 3 + b]) * G[j][b]; + }; + out[i * 4 + j] = @floatCast(v); + }; + return out; +} + +/// iv_norm: y = GroupNorm(x) by stats (32 means, then 32 rstd) with g and beta, then SiLU when `silu`. +pub fn normed(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, stats: Ref, g: Ref, beta: Ref, y: Ref, n: usize, channels: usize, silu: bool) void { + const params = extern struct { pixels: u32, channels: u32, silu: u32 }{ .pixels = @intCast(n), .channels = @intCast(channels), .silu = @intFromBool(silu) }; + bind(e, p, &.{ x, stats, g, beta, y }, params); + e.dispatchThreads(mtl.Size.of(channels, n, 1), mtl.Size.of(32, 8, 1)); +} + +/// GroupNorm's parts over `n` pixels: one sum a chunk of `chunk` pixels and a group. +pub fn chunks(n: usize) usize { + return (n + chunk - 1) / chunk; +} + +fn groupParams(n: usize, channels: usize, centered: bool) extern struct { pixels: u32, channels: u32, chunk: u32, chunks: u32, centered: u32, eps: f32 } { + return .{ .pixels = @intCast(n), .channels = @intCast(channels), .chunk = chunk, .chunks = @intCast(chunks(n)), .centered = @intFromBool(centered), .eps = eps }; +} + +/// iv_group_part: part [32, chunks] = each chunk's sum of a group's values, or (centered) of (x - mean)^2. +pub fn groupPart(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, stats: Ref, part: Ref, n: usize, channels: usize, centered: bool) void { + std.debug.assert(channels % groups == 0 and channels <= 512); + bind(e, p, &.{ x, stats, part }, groupParams(n, channels, centered)); + e.dispatchGroups(mtl.Size.of(chunks(n), 1, 1), mtl.Size.of(channels, 1, 1)); +} + +/// iv_group_finish: stats' 32 means from part, or (centered) its 32 rstd. +pub fn groupFinish(e: mtl.ComputeEncoder, p: mtl.Pipeline, part: Ref, stats: Ref, n: usize, channels: usize, centered: bool) void { + bind(e, p, &.{ part, stats }, groupParams(n, channels, centered)); + e.dispatchGroups(mtl.Size.of(1, 1, 1), mtl.Size.of(32 * groups, 1, 1)); +} + +/// iv_attention: o [rows, 512] = softmax(q k^T / sqrt(512)) v of qkv rows [rows, 1536], both padded to 32 rows. +pub fn attend(e: mtl.ComputeEncoder, p: mtl.Pipeline, qkv: Ref, o: Ref, rows: usize) void { + const params = extern struct { rows: u32, scale: f32 }{ .rows = @intCast(rows), .scale = 1.0 / @sqrt(@as(f32, width)) }; + bind(e, p, &.{ qkv, o }, params); + e.dispatchGroups(mtl.Size.of((rows + 7) / 8, 1, 1), mtl.Size.of(128, 1, 1)); +} + +/// iv_tokens: out [pad + rows * cols / 4, 128] from moments [rows, cols, 64], noise [rows, cols, 32], m and s [128]. +pub fn tokens(e: mtl.ComputeEncoder, p: mtl.Pipeline, moments: Ref, noise: Ref, m: Ref, s: Ref, out: Ref, rows: usize, cols: usize, pad: usize) void { + bind(e, p, &.{ moments, noise, m, s, out }, [3]u32{ @intCast(rows), @intCast(cols), @intCast(pad) }); + e.dispatchThreads(mtl.Size.of(128, pad + rows * cols / 4, 1), mtl.Size.of(128, 2, 1)); +} + +test "every image encoder pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var k = try Kernels.load(device); + k.deinit(); +} + +test "the Winograd transform of a delta kernel is G's outer products" { + var g: [9]f32 = @splat(0); + g[4] = 1; + const u = transform(&g); + const col = [4]f32{ 0, 0.5, -0.5, 0 }; + for (0..4) |i| for (0..4) |j| try std.testing.expectEqual(col[i] * col[j], u[i * 4 + j]); +} diff --git a/zig/src/families/idol/vae_weights.zig b/zig/src/families/idol/vae_weights.zig new file mode 100644 index 0000000000..817794aa39 --- /dev/null +++ b/zig/src/families/idol/vae_weights.zig @@ -0,0 +1,308 @@ +//! Idol's image encoder weights: FLUX.2's VAE encoder, quant_conv and batch norm statistics, as the SDK holds them. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const ks = @import("vae_kernels.zig"); +const Ref = ks.Ref; + +/// The encoder's widths by stage, and the rest of the released geometry. +pub const widths = [4]usize{ 128, 256, 512, 512 }; +pub const latent_channels = 32; +pub const resnets = 2; +pub const batch_norm_eps: f32 = 1e-4; + +/// Accepts vae/config.json at FLUX.2's released encoder geometry, else an error. +pub fn check(gpa: std.mem.Allocator, text: []const u8) !void { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadConfig; + const o = doc.value.object; + const name = o.get("_class_name") orelse return error.BadConfig; + if (name != .string or !std.mem.eql(u8, name.string, "AutoencoderKLFlux2")) return error.NotFlux2Vae; + const want = .{ .{ "in_channels", 3 }, .{ "latent_channels", latent_channels }, .{ "layers_per_block", resnets }, .{ "norm_num_groups", ks.groups } }; + inline for (want) |kv| { + const v = o.get(kv[0]) orelse return error.BadConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.warn("idol vae: {s} is not {d}; the encoder runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } + const lists = .{ .{ "block_out_channels", [_]i64{ 128, 256, 512, 512 } }, .{ "patch_size", [_]i64{ 2, 2 } } }; + inline for (lists) |kv| { + const v = o.get(kv[0]) orelse return error.BadConfig; + if (v != .array or v.array.items.len != kv[1].len) return error.UnsupportedGeometry; + for (v.array.items, kv[1]) |item, n| if (item != .integer or item.integer != n) return error.UnsupportedGeometry; + } + inline for (.{ "mid_block_add_attention", "use_quant_conv" }) |f| { + const v = o.get(f) orelse return error.BadConfig; + if (v != .bool or !v.bool) return error.UnsupportedGeometry; + } + const blocks = o.get("down_block_types") orelse return error.BadConfig; + if (blocks != .array or blocks.array.items.len != widths.len) return error.UnsupportedGeometry; + for (blocks.array.items) |b| if (b != .string or !std.mem.eql(u8, b.string, "DownEncoderBlock2D")) return error.UnsupportedGeometry; + const act = o.get("act_fn") orelse return error.BadConfig; + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.UnsupportedGeometry; + const bn = o.get("batch_norm_eps") orelse return error.BadConfig; + const bn_eps: f64 = switch (bn) { + .float => |f| f, + .integer => |n| @floatFromInt(n), + else => return error.BadConfig, + }; + if (@as(f32, @floatCast(bn_eps)) != batch_norm_eps) return error.UnsupportedGeometry; +} + +/// A convolution's weights (1x1 [cin, ldw], stride-2 3x3 [9 * cin, ldw], 3x3 G g G^T [16, cin, ldw]) and bias. +pub const Conv = struct { w: Ref, b: Ref, cin: usize, cout: usize, ldw: usize, kind: Kind }; +pub const Kind = enum { one, wino, strided, input }; +pub const Norm = struct { g: Ref, b: Ref }; +pub const Resnet = struct { norm1: Norm, conv1: Conv, norm2: Norm, conv2: Conv, shortcut: ?Conv }; +pub const Stage = struct { resnets: [resnets]Resnet, down: ?Conv }; + +pub const Weights = struct { + buf: mtl.Buffer, + conv_in: Conv, + stages: [4]Stage, + mid: [2]Resnet, + attn_norm: Norm, + qkv: Conv, + attn_out: Conv, + norm_out: Norm, + conv_out: Conv, + quant: Conv, + /// The batch norm's running mean and sqrt(running_var + eps), each of 128. + bn_mean: Ref, + bn_std: Ref, + + pub fn deinit(w: *Weights) void { + w.buf.deinit(); + } +}; + +/// The bf16 value an fp32 weight rounds to (to nearest, ties to even), as torch's cast to bf16 makes it. +pub fn bf16(v: f32) f32 { + if (std.math.isNan(v)) return v; + const bits: u32 = @bitCast(v); + const rounded = bits +% (0x7fff + ((bits >> 16) & 1)); + return @bitCast(rounded & 0xffff0000); +} + +/// The folder `dir` (config.json, diffusion_pytorch_model.safetensors in bf16 or f32): every encoder tensor, packed. +pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "config.json" }), arena, .limited(1 << 20)); + try check(gpa, config); + var c = ck.Checkpoint.init(gpa); + defer c.deinit(); + try c.addFile(device, try std.fmt.allocPrintSentinel(arena, "{s}/diffusion_pytorch_model.safetensors", .{dir}, 0), ""); + var w: Weights = undefined; + var p: Packer = .{ .c = &c, .arena = arena }; + try layout(&w, &p); + if (p.read != expected(&c)) { + std.log.err("idol vae: the checkpoint has {d} encoder tensors; the encoder reads {d}", .{ expected(&c), p.read }); + return error.UnexpectedTensors; + } + w.buf = try device.buffer(p.at * 4, mtl.ResourceOptions.shared); + errdefer w.buf.deinit(); + p.out = @as([*]f32, @ptrCast(@alignCast(w.buf.contents())))[0..p.at]; + p.at = 0; + p.read = 0; + p.buf = w.buf; + try layout(&w, &p); + return w; +} + +/// Walks every tensor in the order it is packed: sizes it on the first pass, copies it on the second. +fn layout(w: *Weights, p: *Packer) !void { + var name: [96]u8 = undefined; + w.conv_in = try p.conv("encoder.conv_in", 3, widths[0], .input); + var prev = widths[0]; + for (&w.stages, widths, 0..) |*s, out, i| { + for (&s.resnets, 0..) |*r, j| { + r.* = try p.resnet(try std.fmt.bufPrint(&name, "encoder.down_blocks.{d}.resnets.{d}", .{ i, j }), if (j == 0) prev else out, out); + } + s.down = if (i + 1 < widths.len) try p.conv(try std.fmt.bufPrint(&name, "encoder.down_blocks.{d}.downsamplers.0.conv", .{i}), out, out, .strided) else null; + prev = out; + } + const W = widths[3]; + for (&w.mid, 0..) |*r, i| r.* = try p.resnet(try std.fmt.bufPrint(&name, "encoder.mid_block.resnets.{d}", .{i}), W, W); + const at = "encoder.mid_block.attentions.0"; + w.attn_norm = try p.norm(at ++ ".group_norm", W); + w.qkv = try p.linears(&.{ at ++ ".to_q", at ++ ".to_k", at ++ ".to_v" }, W); + w.attn_out = try p.linears(&.{at ++ ".to_out.0"}, W); + w.norm_out = try p.norm("encoder.conv_norm_out", W); + w.conv_out = try p.conv("encoder.conv_out", W, 2 * latent_channels, .wino); + w.quant = try p.conv("quant_conv", 2 * latent_channels, 2 * latent_channels, .one); + const n = 4 * latent_channels; + const mean = try p.vector("bn.running_mean", n); + const var_ = try p.vector("bn.running_var", n); + w.bn_mean = p.alloc(n); + w.bn_std = p.alloc(n); + if (p.out) |out| { + // the SDK's sqrt(var + eps) in bf16, here in float64 rounded once + for (out[w.bn_mean.off / 4 ..][0..n], out[w.bn_std.off / 4 ..][0..n], mean, var_) |*mo, *so, mv, v| { + mo.* = mv; + so.* = @floatCast(@sqrt(@as(f64, v) + 1e-4)); + } + } +} + +/// The checkpoint's encoder-side tensors: the encoder's, quant_conv's and the batch norm's two statistics. +fn expected(c: *const ck.Checkpoint) usize { + var n: usize = 0; + var it = c.tensors.keyIterator(); + while (it.next()) |k| { + if (std.mem.startsWith(u8, k.*, "encoder.") or std.mem.startsWith(u8, k.*, "quant_conv.") or + std.mem.eql(u8, k.*, "bn.running_mean") or std.mem.eql(u8, k.*, "bn.running_var")) n += 1; + } + return n; +} + +/// Lays the encoder's tensors out in one fp32 buffer: sizes on the first pass (out null), copies on the second. +const Packer = struct { + c: *const ck.Checkpoint, + arena: std.mem.Allocator, + buf: mtl.Buffer = undefined, + out: ?[]f32 = null, + at: usize = 0, + read: usize = 0, + + fn alloc(p: *Packer, n: usize) Ref { + const r: Ref = .{ .buf = p.buf, .off = p.at * 4 }; + p.at = std.mem.alignForward(usize, p.at + n, 64); + return r; + } + + /// Tensor `name` of `shape` as the bf16 values the SDK computes with (a bf16 file's own, an f32 file's rounded). + fn tensor(p: *Packer, name: []const u8, shape: []const usize) ![]const f32 { + const t = try p.c.get(name); + if ((t.dtype != .f32 and t.dtype != .bf16) or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("idol vae: {s} is {t} {any}; the encoder reads bf16 or f32 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } + p.read += 1; + if (p.out == null) return &.{}; + const n = t.count(); + const v = try p.arena.alloc(f32, n); + switch (t.dtype) { + .bf16 => for (v, t.host(u16)) |*o, h| { + o.* = @bitCast(@as(u32, h) << 16); + }, + else => for (v, t.host(f32)) |*o, f| { + o.* = bf16(f); + }, + } + return v; + } + + fn vector(p: *Packer, name: []const u8, n: usize) ![]const f32 { + return p.tensor(name, &.{n}); + } + + /// A copied vector `name` of n. + fn copy(p: *Packer, name: []const u8, n: usize) !Ref { + const v = try p.vector(name, n); + const r = p.alloc(n); + if (p.out) |out| @memcpy(out[r.off / 4 ..][0..n], v); + return r; + } + + /// Convolution `name` (.weight [cout, cin, k, k], .bias) in `kind`'s layout, its columns padded to ldw. + fn conv(p: *Packer, name: []const u8, cin: usize, cout: usize, kind: Kind) !Conv { + var buf: [96]u8 = undefined; + const taps: usize = if (kind == .one) 1 else 9; + const k: usize = if (kind == .one) 1 else 3; + const ldw = if (kind == .input) cout else std.mem.alignForward(usize, cout, 64); + const rows: usize = switch (kind) { + .one => cin, + .wino => 16 * cin, + .strided, .input => 9 * cin, + }; + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ cout, cin, k, k }); + const r = p.alloc(rows * ldw); + if (p.out) |out| { + const dst = out[r.off / 4 ..][0 .. rows * ldw]; + @memset(dst, 0); + for (0..cout) |o| for (0..cin) |i| { + const g = w[(o * cin + i) * taps ..][0..taps]; + switch (kind) { + .one => dst[i * ldw + o] = g[0], + .wino => for (ks.transform(g[0..9]), 0..) |v, pt| { + dst[(pt * cin + i) * ldw + o] = v; + }, + .strided, .input => for (g, 0..) |v, tap| { + dst[(tap * cin + i) * ldw + o] = v; + }, + } + }; + } + const b = try p.copy(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), cout); + return .{ .w = r, .b = b, .cin = cin, .cout = cout, .ldw = ldw, .kind = kind }; + } + + /// Linear layers `names` (.weight [width, width], .bias) stacked as one 1x1 convolution. + fn linears(p: *Packer, names_: []const []const u8, width: usize) !Conv { + var buf: [96]u8 = undefined; + const n = names_.len * width; + const r = p.alloc(width * n); + for (names_, 0..) |name, part| { + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ width, width }); + if (p.out) |out| { + const dst = out[r.off / 4 ..][0 .. width * n]; + for (0..width) |o| for (0..width) |i| { + dst[i * n + part * width + o] = w[o * width + i]; + }; + } + } + const b = p.alloc(n); + for (names_, 0..) |name, i| { + const v = try p.vector(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), width); + if (p.out) |out| @memcpy(out[b.off / 4 + i * width ..][0..width], v); + } + return .{ .w = r, .b = b, .cin = width, .cout = n, .ldw = n, .kind = .one }; + } + + fn norm(p: *Packer, name: []const u8, width: usize) !Norm { + var buf: [96]u8 = undefined; + return .{ .g = try p.copy(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), width), .b = try p.copy(try std.fmt.bufPrint(&buf, "{s}.bias", .{name}), width) }; + } + + fn resnet(p: *Packer, name: []const u8, cin: usize, cout: usize) !Resnet { + var buf: [96]u8 = undefined; + return .{ + .norm1 = try p.norm(try std.fmt.bufPrint(&buf, "{s}.norm1", .{name}), cin), + .conv1 = try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv1", .{name}), cin, cout, .wino), + .norm2 = try p.norm(try std.fmt.bufPrint(&buf, "{s}.norm2", .{name}), cout), + .conv2 = try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv2", .{name}), cout, cout, .wino), + .shortcut = if (cin != cout) try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv_shortcut", .{name}), cin, cout, .one) else null, + }; + } +}; + +test "only FLUX.2's released encoder geometry passes" { + const a = std.testing.allocator; + const sample = + \\{"_class_name": "AutoencoderKLFlux2", "act_fn": "silu", "batch_norm_eps": 0.0001, "block_out_channels": [128, 256, 512, 512], + \\ "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"], "in_channels": 3, + \\ "latent_channels": 32, "layers_per_block": 2, "mid_block_add_attention": true, "norm_num_groups": 32, + \\ "patch_size": [2, 2], "use_quant_conv": true} + ; + try check(a, sample); + const narrow = try std.mem.replaceOwned(u8, a, sample, "[128, 256, 512, 512]", "[96, 192, 384, 384]"); + defer a.free(narrow); + try std.testing.expectError(error.UnsupportedGeometry, check(a, narrow)); + const no_quant = try std.mem.replaceOwned(u8, a, sample, "\"use_quant_conv\": true", "\"use_quant_conv\": false"); + defer a.free(no_quant); + try std.testing.expectError(error.UnsupportedGeometry, check(a, no_quant)); + const other = try std.mem.replaceOwned(u8, a, sample, "AutoencoderKLFlux2", "AutoencoderKL"); + defer a.free(other); + try std.testing.expectError(error.NotFlux2Vae, check(a, other)); +} + +test "an f32 weight rounds to bf16 as torch's cast does: to nearest, ties to even" { + try std.testing.expectEqual(@as(f32, 1.0), bf16(1.0 + 0x1p-9)); + try std.testing.expectEqual(@as(f32, 1.0 + 0x1p-7), bf16(1.0 + 0x1p-7 + 0x1p-9)); + try std.testing.expectEqual(@as(f32, 1.0 + 0x1p-6), bf16(1.0 + 0x1p-7 + 0x1p-8)); + try std.testing.expectEqual(@as(f32, 1.0 + 0x1p-7), bf16(1.0 + 0x1p-7 + 0x1p-8 - 0x1p-20)); +} diff --git a/zig/tests/contract/metal_idol_vae.zig b/zig/tests/contract/metal_idol_vae.zig new file mode 100644 index 0000000000..5ec14ca3b2 --- /dev/null +++ b/zig/tests/contract/metal_idol_vae.zig @@ -0,0 +1,176 @@ +//! Idol's image encoder kernels (idol_vae.metal, fp32) under the contract; each control breaks one gate. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); +const wn = @import("metal_idol_vae_wino.zig"); +const rw = @import("metal_idol_vae_rows.zig"); + +const g = id.g; +const Ends = id.Ends; + +/// A control: the kernel it breaks, the gate that must catch it, and the one text change that breaks it. +const Control = struct { kernel: []const u8, gate: g.Gate, what: []const u8, from: []const u8, to: []const u8 }; + +const wino_mma = "for (uint i = 0; i < 2; i++) simdgroup_multiply_accumulate(acc[i][j], vm[i], um, acc[i][j]);"; +const wino_mma_half = "for (uint i = 0; i < 2; i++) { simdgroup_multiply_accumulate(acc[i][j], vm[i], um, acc[i][j]); " ++ + "acc[i][j].thread_elements()[0] = float(half(acc[i][j].thread_elements()[0])); " ++ + "acc[i][j].thread_elements()[1] = float(half(acc[i][j].thread_elements()[1])); }"; +const patch_axes = "const uint py = (patch / w2) * 2 + (at.x % 4) / 2;\n const uint px = (patch % w2) * 2 + at.x % 2;"; +const patch_swapped = "const uint py = (patch / w2) * 2 + at.x % 2;\n const uint px = (patch % w2) * 2 + (at.x % 4) / 2;"; + +const controls = [_]Control{ + .{ .kernel = "iv_conv_in", .gate = .shapes, .what = "stages 2 x - 1 in half", .from = "s = fma(v[c] * 2.0f - 1.0f,", .to = "s = fma(float(half(v[c] * 2.0f - 1.0f))," }, + .{ .kernel = "iv_conv 1x1", .gate = .shapes, .what = "stages its inputs in half", .from = "for (uint i = 0; i < 4; i++) dst[i] = src[i];", .to = "for (uint i = 0; i < 4; i++) dst[i] = float4(half4(src[i]));" }, + .{ .kernel = "iv_conv 3x3 stride 2", .gate = .shapes, .what = "pads above as padding=1 would", .from = "const uint iy = oy * stride + tap / 3;", .to = "const uint iy = oy * stride + tap / 3 - (p.taps == 9 ? 1 : 0);" }, + .{ .kernel = "iv_wino 3x3", .gate = .shapes, .what = "stages its patch in half", .from = "*(threadgroup float4*)(buf + px * BK + q * 4) = v;", .to = "*(threadgroup float4*)(buf + px * BK + q * 4) = float4(half4(v));" }, + .{ .kernel = "iv_wino 3x3 + residual", .gate = .shapes, .what = "drops a term of the output transform", .from = "s[j][1] = (m1 - m2) - m3;", .to = "s[j][1] = (m1 - m2);" }, + .{ .kernel = "iv_wino 3x3", .gate = .sums, .what = "sums in half", .from = wino_mma, .to = wino_mma_half }, + .{ .kernel = "iv_norm silu", .gate = .shapes, .what = "its SiLU in half", .from = "y[i] = p.silu ? v / (1.0f + precise::exp(-v)) : v;", .to = "y[i] = p.silu ? float(half(v) / (1.0h + exp(-half(v)))) : v;" }, + .{ .kernel = "iv_group_part", .gate = .sums, .what = "sums in half", .from = "s += p.centered ? d * d : v;", .to = "s = float(half(s) + half(p.centered ? d * d : v));" }, + .{ .kernel = "iv_group_finish", .gate = .sums, .what = "sums in half", .from = "for (uint i = lane; i < p.chunks; i += 32) s += part[sg * p.chunks + i];", .to = "for (uint i = lane; i < p.chunks; i += 32) s = float(half(s) + half(part[sg * p.chunks + i]));" }, + .{ .kernel = "iv_attention", .gate = .shapes, .what = "keeps stale maxima (no rescale)", .from = " if (rescale) {", .to = " if (false) {" }, + .{ .kernel = "iv_tokens", .gate = .shapes, .what = "swaps the patch's axes", .from = patch_axes, .to = patch_swapped }, + .{ .kernel = "iv_tokens", .gate = .nonfinite, .what = "clamps a NaN logvar (Metal's clamp)", .from = "const float logvar = isnan(raw) ? raw : clamp(raw, -30.0f, 20.0f);", .to = "const float logvar = clamp(raw, -30.0f, 20.0f);" }, +}; + +/// The encoder's source as the engine compiles it, with `ctl`'s change made where it occurs, exactly once. +fn compile(a: std.mem.Allocator, device: mtl.Device, ctl: ?Control) !mtl.Library { + const text = if (ctl) |x| blk: { + if (std.mem.count(u8, sources.idol_vae, x.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, x.from }); + return error.ControlTextNotFoundOnce; + } + break :blk try std.mem.replaceOwned(u8, a, sources.idol_vae, x.from, x.to); + } else try a.dupe(u8, sources.idol_vae); + defer a.free(text); + return mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); +} + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + pipes: std.ArrayList(mtl.Pipeline) = .empty, + names: std.ArrayList([]u8) = .empty, + libs: std.ArrayList(mtl.Library) = .empty, + made: std.ArrayList(Made) = .empty, + + fn pipeline(s: *Suite, device: mtl.Device, lib: mtl.Library, name: []const u8) !mtl.Pipeline { + const p = try mtl.Pipeline.init(device, lib, name, false); + errdefer p.deinit(); + try s.pipes.append(s.a, p); + return p; + } + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: mtl.Library, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol vae {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + if (try w.name("iv_conv_in")) |name| { + const t = try s.make(st.ConvIn, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_conv_in"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for (st.conv_forms) |f| { + const name = (try w.name(f.label)) orelse continue; + const t = try s.make(st.Conv, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_conv"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for (wn.forms) |f| { + const name = (try w.name(f.label)) orelse continue; + const t = try s.make(wn.Wino, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_wino"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ true, false }) |silu| { + const name = (try w.name(if (silu) "iv_norm silu" else "iv_norm")) orelse continue; + const t = try s.make(st.Norm, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_norm"), .silu = silu, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ false, true }) |centered| { + const name = (try w.name(if (centered) "iv_group_part centered" else "iv_group_part")) orelse continue; + const t = try s.make(rw.GroupPart, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_group_part"), .centered = centered, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("iv_group_finish")) |name| { + const t = try s.make(rw.GroupFinish, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_group_finish"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("iv_attention")) |name| { + const t = try s.make(rw.Attention, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_attention"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("iv_tokens")) |name| { + const t = try s.make(st.Tokens, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "iv_tokens"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + for ([_]?Control{null} ++ controls) |x| { + const lib = try compile(a, device, x); + s.libs.append(a, lib) catch |err| { + lib.deinit(); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.pipes.items) |p| p.deinit(); + s.pipes.deinit(a); + for (s.libs.items) |l| l.deinit(); + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_vae_rows.zig b/zig/tests/contract/metal_idol_vae_rows.zig new file mode 100644 index 0000000000..10824594ab --- /dev/null +++ b/zig/tests/contract/metal_idol_vae_rows.zig @@ -0,0 +1,437 @@ +//! iv_group_part, iv_group_finish and iv_attention under the contract: fp32 sums against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const mi = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); + +const c = mi.c; +const g = mi.g; +const Range = mi.Range; +const Ends = mi.Ends; +const dec = st.ks; +const ref = st.ref; +const uniform = st.uniform; +const plant = st.plant; +const nan = st.nan; +const inf = st.inf; + +// --- iv_group_part: GroupNorm's chunk sums, of the values or of their squared deviations --------------------------- + +pub const GroupPart = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + centered: bool, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + stats: [64]f32 = undefined, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// A thread's chunk of pixels, then its group's channels (up to 16) summed in turn. + fn depth(s: c.Shape) usize { + return @min(s.k, dec.chunk) + s.n / 32 + 1; + } + + /// Rows unused (1); n: channels (32 groups); k: pixels, summed in chunks. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{}, .{ .size = 128, .step = 32, .min = 32 }, .{ .size = dec.chunk } }, + .served = &.{ .{ .rows = 1, .n = 512, .k = 16384 }, .{ .rows = 1, .n = 512, .k = 16500 }, .{ .rows = 1, .n = 512, .k = 65536 }, .{ .rows = 1, .n = 256, .k = 262144 }, .{ .rows = 1, .n = 128, .k = 1048576 } }, + }; + } + + pub fn deinit(t: *GroupPart) void { + t.a.free(t.x); + t.rig.deinit(); + } + pub fn kernel(t: *GroupPart, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(GroupPart); + + pub fn prepare(t: *GroupPart, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.n % 32 != 0 or shape.k == 0) return false; + t.rig.reset(); + t.a.free(t.x); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.k * shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // integers: plain sums reach ~2^21 and cancel; deviations within 18 keep squares' sums under 2^24 + for (0..32) |grp| t.stats[grp] = @floatFromInt(rng.intRangeAtMost(i32, -4, 4)); + for (0..shape.k) |p| for (0..shape.n) |ch| { + const grp = ch / (shape.n / 32); + const in = p % dec.chunk; + const half = dec.chunk / 2; + const v: f32 = if (t.centered) + t.stats[grp] + @as(f32, @floatFromInt(rng.intRangeAtMost(i32, -18, 18))) + else if (in < half) @as(f32, if (in % 2 == 0) 127 else 511) else -@as(f32, if ((in - half) % 2 == 0) 127 else 511) + @as(f32, if (in - half < 4) 1 else 0); + t.x[p * shape.n + ch] = v; + }; + } else { + uniform(rng, t.x, -3, 3); + uniform(rng, t.stats[0..32], -0.5, 0.5); + if (fill == .nonfinite) plant(t.x, shape.n); + } + uniform(rng, t.stats[32..], 0.5, 2); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(&t.stats)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *GroupPart) !void { + const p = t.rig.begin(); + dec.groupPart(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.out), t.shape.k, t.shape.n, t.centered); + try t.rig.finish(p); + } + + pub fn count(t: *GroupPart) usize { + return 32 * dec.chunks(t.shape.k); + } + + /// Output (group, chunk)'s terms: each pixel's channels of the group, or their squared deviations. + fn terms(t: *const GroupPart, grp: usize, ck: usize, out: []f64) usize { + const size = t.shape.n / 32; + var n: usize = 0; + for (ck * dec.chunk..@min((ck + 1) * dec.chunk, t.shape.k)) |p| for (0..size) |i| { + const v: f64 = t.x[p * t.shape.n + grp * size + i]; + const d = v - t.stats[grp]; + out[n] = if (t.centered) d * d else v; + n += 1; + }; + return n; + } + + pub fn reference(t: *GroupPart, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, dec.chunk * (t.shape.n / 32)); + defer t.a.free(buf); + const chunks = dec.chunks(t.shape.k); + for (0..32) |grp| for (0..chunks) |ck| { + const n = t.terms(grp, ck, buf); + var s: f64 = 0; + var m: f64 = 0; + for (buf[0..n]) |v| { + s += v; + m += @abs(v); + } + want[grp * chunks + ck] = s; + mag[grp * chunks + ck] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *GroupPart, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, dec.chunk * (t.shape.n / 32)); + defer t.a.free(buf); + const chunks = dec.chunks(t.shape.k); + for (0..32) |grp| for (0..chunks) |ck| { + const n = t.terms(grp, ck, buf); + out[grp * chunks + ck] = contract.rows.sumIn(acc, order, buf[0..n]); + }; + } +}; + +// --- iv_group_finish: the 32 means, then the 32 rstd, from the chunk sums ------------------------------------------- + +pub const GroupFinish = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + part: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Each lane's chunks / 32 sums in turn, then the simdgroup tree. + fn depth(s: c.Shape) usize { + return (s.k + 31) / 32 + 5; + } + + /// Rows unused (1); n: the 32 groups; k: the chunks. The group size is the widest stage's, 512 / 32. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{}, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 1, .n = 32, .k = 64 }, .{ .rows = 1, .n = 32, .k = 65 }, .{ .rows = 1, .n = 32, .k = 256 }, .{ .rows = 1, .n = 32, .k = 1024 }, .{ .rows = 1, .n = 32, .k = 4096 } }, + }; + } + + pub fn deinit(t: *GroupFinish) void { + t.a.free(t.part); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *GroupFinish, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(GroupFinish); + + fn pixels(t: *const GroupFinish) usize { + return t.shape.k * dec.chunk; + } + + pub fn prepare(t: *GroupFinish, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != 32 or shape.k == 0) return false; + t.rig.reset(); + t.a.free(t.part); + t.shape = shape; + t.part = try t.a.alloc(f32, 32 * shape.k); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // integer sums near 2^16 (past fp16's range), the second half of chunks their negatives plus 64 + for (0..32) |grp| for (0..shape.k) |ck| { + const half = shape.k / 2; + const v: f32 = if (ck < half) 65533 + @as(f32, @floatFromInt(ck % 7)) else if (ck < 2 * half) -65533 - @as(f32, @floatFromInt((ck - half) % 7)) + 64 else 64; + t.part[grp * shape.k + ck] = v; + }; + } else { + uniform(rng, t.part, 0, 5000); + if (fill == .nonfinite) plant(t.part, shape.k); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.part)); + t.out = try t.rig.output(64 * 4); + return true; + } + + pub fn run(t: *GroupFinish) !void { + const p = t.rig.begin(); + for ([_]bool{ false, true }) |centered| dec.groupFinish(p.e, t.pipe, ref(t.ins[0]), ref(t.out), t.pixels(), 512, centered); + try t.rig.finish(p); + } + + pub fn count(_: *GroupFinish) usize { + return 64; + } + + fn outputs(t: *GroupFinish, sums: [32]f64, mags: [32]f64, want: []f64, mag: []f64, ends: *Ends) void { + const n: f64 = @floatFromInt(t.pixels() * 16); + const nf: f64 = @as(f32, @floatCast(n)); + const eps: f64 = dec.eps; + for (0..32) |grp| { + const s = sums[grp]; + const e = c.gamma(@floatFromInt(depth(t.shape)), c.unit(.f32)) * mags[grp]; + const sr = Range{ .lo = s - e, .hi = s + e }; + const mean = (Range{ .lo = sr.lo / nf, .hi = sr.hi / nf }).fp32(mi.div_ulps); + ends.put(want, mag, grp, s / n, mean); + const pos = mean.add(Range.of(eps)).fp32(0.5); + const inv = if (pos.lo > 0) pos.rsqrt().fp32(mi.rsqrt_ulps) else Range.of(1 / @sqrt(s / n + eps)); + ends.put(want, mag, 32 + grp, 1 / @sqrt(s / n + eps), inv); + } + } + + pub fn reference(t: *GroupFinish, want: []f64, mag: []f64) !void { + try t.ends.size(64); + var sums: [32]f64 = @splat(0); + var mags: [32]f64 = @splat(0); + for (0..32) |grp| for (t.part[grp * t.shape.k ..][0..t.shape.k]) |v| { + sums[grp] += v; + mags[grp] += @abs(@as(f64, v)); + }; + t.outputs(sums, mags, want, mag, &t.ends); + } + pub fn bounds(t: *GroupFinish, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *GroupFinish, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *GroupFinish, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try t.a.alloc(f64, t.shape.k); + defer t.a.free(terms); + const n: f64 = @floatFromInt(t.pixels() * 16); + for (0..32) |grp| { + for (terms, t.part[grp * t.shape.k ..][0..t.shape.k]) |*v, p| v.* = p; + const s = contract.rows.sumIn(acc, order, terms); + out[grp] = s / n; + out[32 + grp] = 1 / @sqrt(s / n + dec.eps); + } + } +}; + +// --- iv_attention: softmax(q k^T / sqrt(512)) v, one head of 512, fp32 throughout ---------------------------------- + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + qkv: []f32 = &.{}, + tol: []f64 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + const D = 512; + const ld = 3 * D; + + /// The scores' 512-deep sums, the keys' sums and one rescale a block of 32 keys. + fn depth(s: c.Shape) usize { + return D + s.k + s.k / 32 + 8; + } + + /// Rows: the queries compared (the first `rows` of all k); n: the head width; k: the positions L. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = mi.div_ulps + 2, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 24, .n = D, .k = 16384 }, .{ .rows = 24, .n = D, .k = 16500 }, .{ .rows = 24, .n = D, .k = 16128 }, .{ .rows = 40, .n = D, .k = 40 } }, + }; + } + + pub fn deinit(t: *Attention) void { + t.a.free(t.qkv); + t.a.free(t.tol); + t.rig.deinit(); + } + pub fn kernel(t: *Attention, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != D or shape.k == 0 or shape.rows == 0 or shape.rows > shape.k) return false; + t.rig.reset(); + t.a.free(t.qkv); + t.shape = shape; + t.stress = fill == .stress; + const L = shape.k; + const padded = std.mem.alignForward(usize, L, 32); + t.qkv = try t.a.alloc(f32, padded * ld); + @memset(t.qkv, 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores (q = k = 0); each column's values integers whose partial sums reach ~2^23 and cancel to 4 + const half = L / 2; + for (0..L) |j| { + const at = if (j < half) j else j - half; + const positive: f32 = if (at % 2 == 0) 127 else 2041; + const residual: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + const v: f32 = if (j < half) positive else if (j < 2 * half) residual - positive else 0; + @memset(t.qkv[j * ld + 2 * D ..][0..D], v); + } + } else { + for (0..L) |j| uniform(rng, t.qkv[j * ld ..][0..ld], -1, 1); + if (fill == .nonfinite) { + // values NaN, +Inf, -Inf, and both infinities in one column; a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]f32{ nan, inf, -inf, inf, -inf }) |p, v| if (p[0] < L) { + t.qkv[p[0] * ld + 2 * D + p[1]] = v; + }; + if (L > 6) t.qkv[1 * ld + 5] = nan; + if (L > 7) t.qkv[7 * ld + D + 3] = inf; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.out = try t.rig.output(padded * D * 4); + return true; + } + + pub fn run(t: *Attention) !void { + const p = t.rig.begin(); + dec.attend(p.e, t.pipe, ref(t.ins[0]), ref(t.out), t.shape.k); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return t.shape.rows * D; + } + + /// Each compared query's output in float64, and its bound: the scores' error moved through exp, then the sums'. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.shape.k; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + const u = c.unit(.f32); + t.a.free(t.tol); + t.tol = try t.a.alloc(f64, t.count()); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + const sabs = try t.a.alloc(f64, L); + defer t.a.free(sabs); + var out: [D]f64 = undefined; + var spread: [D]f64 = undefined; + for (0..t.shape.rows) |i| { + const q = t.qkv[i * ld ..][0..D]; + var m: f64 = -std.math.inf(f64); + var nan_row = false; + var worst: f64 = 0; + for (s, sabs, 0..) |*sj, *aj, j| { + const k = t.qkv[j * ld + D ..][0..D]; + var dot: f64 = 0; + var mag_: f64 = 0; + for (q, k) |a, b| { + dot += @as(f64, a) * b; + mag_ += @abs(@as(f64, a) * b); + } + sj.* = dot * scale; + aj.* = mag_ * scale; + nan_row = nan_row or std.math.isNan(sj.*); + m = @max(m, sj.*); + worst = @max(worst, aj.*); + } + if (nan_row) m = std.math.nan(f64); + // a score's error: its sums and the scale, then the subtraction of the maximum + const es = c.gamma(D + 4, u) * worst + 2 * u * (@abs(m) + worst); + const rel = 2 * es + (mi.exp_ulps + 1) * 0x1p-24; + var l: f64 = 0; + @memset(&out, 0); + @memset(&spread, 0); + for (s, 0..) |sj, j| { + const p = @exp(sj - m); + l += p; + for (&out, &spread, t.qkv[j * ld + 2 * D ..][0..D]) |*o, *sp, v| { + o.* += p * v; + sp.* += p * @abs(@as(f64, v)); + } + } + const sums = c.gamma(@floatFromInt(depth(t.shape)), u); + for (out, spread, 0..) |o, sp, e| { + const at = i * D + e; + want[at] = o / l; + mag[at] = sp / l; + t.tol[at] = (2 * rel + 2 * sums) * sp / l; + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, _: []const f64, allow: []f64) !void { + for (allow, want, t.tol) |*l, w, tl| l.* = if (t.stress) c.boundExact(spec(t.name), w) else if (std.math.isFinite(w)) tl + c.boundExact(spec(t.name), w) else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.shape.k; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = t.qkv[j * ld + 2 * D + e]; + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..t.shape.rows) |i| out[i * D + e] = o; + } + } +}; diff --git a/zig/tests/contract/metal_idol_vae_steps.zig b/zig/tests/contract/metal_idol_vae_steps.zig new file mode 100644 index 0000000000..5fb097959f --- /dev/null +++ b/zig/tests/contract/metal_idol_vae_steps.zig @@ -0,0 +1,578 @@ +//! iv_conv_in, iv_conv, iv_norm and iv_tokens under the contract: Idol's image encoder's fp32 steps against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); + +const c = id.c; +const g = id.g; +const Range = id.Range; +const Ends = id.Ends; +pub const ks = tf.idol.vae_kernels; + +pub fn ref(r: rig_.Ref) ks.Ref { + return .{ .buf = r.buf, .off = r.off }; +} + +pub fn uniform(rng: std.Random, out: []f32, lo: f32, hi: f32) void { + for (out) |*v| v.* = lo + rng.float(f32) * (hi - lo); +} + +pub const nan = std.math.nan(f32); +pub const inf = std.math.inf(f32); + +/// NaN, +Inf, -Inf at three places, and +Inf with -Inf at two places `apart` further on (one row holding both). +pub fn plant(x: []f32, apart: usize) void { + for ([_]usize{ 0, 1, 2 }, [_]f32{ nan, inf, -inf }) |at, v| x[(at * 7) % x.len] = v; + if (x.len > 3 * apart + 1) { + x[3 * apart] = inf; + x[3 * apart + 1] = -inf; + } +} + +// --- iv_conv_in: conv3x3(2 x - 1) + b of the picture, 3 channels to 128 ---------------------------------------------- + +pub const ConvIn = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + b: []f32 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// The 27 taps' products, 2 x - 1's rounding and the bias. + fn depth(_: c.Shape) usize { + return 27 + 2; + } + + /// Rows: the picture's rows; n: the 128 output channels; k: its columns (the encoder's picture is 1024 x 1024). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 2 }, .{}, .{ .size = 2 } }, + .served = &.{ .{ .rows = 12, .n = 128, .k = 1024 }, .{ .rows = 1024, .n = 128, .k = 12 }, .{ .rows = 37, .n = 128, .k = 41 } }, + }; + } + + pub fn deinit(t: *ConvIn) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *ConvIn) void { + inline for (.{ &t.x, &t.w, &t.b }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *ConvIn, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(ConvIn); + + pub fn prepare(t: *ConvIn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0 or shape.n != 128) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.k * 3); + t.w = try t.a.alloc(f32, 27 * 128); + t.b = try t.a.alloc(f32, 128); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // 2 x - 1 an odd integer (127 or 2041), channel 1 two more, channel 2 zero; weights 509, -509, 1 cancel + for (0..shape.rows * shape.k) |p| { + const v: f32 = if (p % 2 == 0) 2041 else 127; + t.x[p * 3] = (v + 1) / 2; + t.x[p * 3 + 1] = (v + 3) / 2; + t.x[p * 3 + 2] = 0.5; + } + for (0..27) |r| @memset(t.w[r * 128 ..][0..128], switch (r % 3) { + 0 => 509, + 1 => -509, + else => 1, + }); + @memset(t.b, 0); + } else { + uniform(rng, t.x, 0, 1); + uniform(rng, t.w, -0.5, 0.5); + uniform(rng, t.b, -1, 1); + if (fill == .nonfinite) plant(t.x, shape.k * 3); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *ConvIn) !void { + const p = t.rig.begin(); + ks.convIn(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.out), t.shape.rows, t.shape.k); + try t.rig.finish(p); + } + + pub fn count(t: *ConvIn) usize { + return t.shape.rows * t.shape.k * 128; + } + + /// Output (pixel, o)'s 27 products in the kernel's order (row, column, channel), zero past the picture's edge. + fn terms(t: *const ConvIn, px: usize, o: usize, out: *[27]f64) void { + const py: isize = @intCast(px / t.shape.k); + const pxx: isize = @intCast(px % t.shape.k); + for (0..3) |ky| for (0..3) |kx| for (0..3) |ch| { + const iy = py + @as(isize, @intCast(ky)) - 1; + const ix = pxx + @as(isize, @intCast(kx)) - 1; + const at = (ky * 3 + kx) * 3 + ch; + out[at] = if (iy < 0 or ix < 0 or iy >= t.shape.rows or ix >= t.shape.k) 0 else blk: { + const v: f64 = t.x[(@as(usize, @intCast(iy)) * t.shape.k + @as(usize, @intCast(ix))) * 3 + ch]; + break :blk (2 * v - 1) * t.w[at * 128 + o]; + }; + }; + } + + pub fn reference(t: *ConvIn, want: []f64, mag: []f64) !void { + var buf: [27]f64 = undefined; + for (0..t.shape.rows * t.shape.k) |px| for (0..128) |o| { + t.terms(px, o, &buf); + var s: f64 = 0; + var m: f64 = 0; + for (buf) |v| { + s += v; + m += @abs(v); + } + want[px * 128 + o] = s + t.b[o]; + m += @abs(@as(f64, t.b[o])); + mag[px * 128 + o] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *ConvIn, acc: c.Accum, order: g.Order, out: []f64) !void { + var buf: [27]f64 = undefined; + for (0..t.shape.rows * t.shape.k) |px| for (0..128) |o| { + t.terms(px, o, &buf); + out[px * 128 + o] = contract.rows.sumIn(acc, order, &buf) + t.b[o]; + }; + } +}; + +// --- iv_conv: a 1x1 or stride-2 3x3 convolution as an implicit GEMM, fp32 sums, + bias (+ a residual) --------------- + +/// A form the encoder launches: its taps (1, or 9 at stride 2), a residual, and its served shapes. +pub const Form = struct { label: []const u8, taps: usize = 1, add: bool = false, served: []const c.Shape }; + +/// A 1x1's rows are pixels (fewer than the engine's: rows only tile); a 3x3's are output rows of 9 columns. +pub const conv_forms = [_]Form{ + .{ .label = "iv_conv 1x1", .served = &.{ .{ .rows = 1100, .n = 1536, .k = 512 }, .{ .rows = 1100, .n = 256, .k = 128 }, .{ .rows = 1100, .n = 512, .k = 256 }, .{ .rows = 1100, .n = 64, .k = 64 } } }, + .{ .label = "iv_conv 1x1 + residual", .add = true, .served = &.{.{ .rows = 1100, .n = 512, .k = 512 }} }, + .{ .label = "iv_conv 3x3 stride 2", .taps = 9, .served = &.{ .{ .rows = 8, .n = 128, .k = 128 }, .{ .rows = 8, .n = 256, .k = 256 }, .{ .rows = 8, .n = 512, .k = 512 } } }, +}; + +const out_cols = 9; + +pub const Conv = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: Form, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + b: []f32 = &.{}, + r: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Each output sums taps * cin products, then the bias and the residual. + fn depth1(s: c.Shape) usize { + return s.k + 2; + } + fn depth9(s: c.Shape) usize { + return 9 * s.k + 2; + } + + /// Rows: pixels (1x1) or output rows (3x3); n: output channels; k: input channels. + pub fn spec(name: []const u8, f: Form) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = if (f.taps == 9) depth9 else depth1, + .out = .f32, + .tiles = .{ if (f.taps == 9) .{ .size = 8 } else .{ .size = 64 }, .{ .size = 64 }, .{ .size = 32, .step = 32, .min = 32 } }, + .served = f.served, + }; + } + + pub fn deinit(t: *Conv) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Conv) void { + inline for (.{ &t.x, &t.w, &t.b, &t.r }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Conv, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Conv); + + fn ldw(t: *const Conv) usize { + return std.mem.alignForward(usize, t.shape.n, 64); + } + fn cols(t: *const Conv) usize { + return if (t.form.taps == 9) out_cols else 1; + } + fn pixels(t: *const Conv) usize { + return t.shape.rows * t.cols(); + } + fn inRows(t: *const Conv) usize { + return if (t.form.taps == 9) 2 * t.shape.rows else t.shape.rows; + } + fn inCols(t: *const Conv) usize { + return if (t.form.taps == 9) 2 * out_cols else 1; + } + + pub fn prepare(t: *Conv, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 32 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const L = t.ldw(); + const K = t.form.taps * shape.k; + const in_px = t.inRows() * t.inCols(); + t.x = try t.a.alloc(f32, in_px * shape.k); + t.w = try t.a.alloc(f32, K * L); + t.b = try t.a.alloc(f32, shape.n); + t.r = try t.a.alloc(f32, if (t.form.add) t.pixels() * shape.n else 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + @memset(t.w, 0); + if (fill == .stress) { + // the channels' second half mirroring the first, weights 1 then -1: partial sums reach ~2^18 and cancel + const half = shape.k / 2; + for (0..in_px) |p| for (0..shape.k) |ci| { + const at = if (ci < half) ci else ci - half; + const v: f32 = if (at % 2 == 0) 127 else 2041; + t.x[p * shape.k + ci] = if (ci < half) v else v + @as(f32, if (at < 4) 1 else 0); + }; + for (0..K) |k| for (0..shape.n) |o| { + t.w[k * L + o] = if (k % shape.k < half) 1 else -1; + }; + @memset(t.b, 0); + for (t.r) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + } else { + uniform(rng, t.x, -2, 2); + for (0..K) |k| uniform(rng, t.w[k * L ..][0..shape.n], -0.1, 0.1); + uniform(rng, t.b, -1, 1); + uniform(rng, t.r, -3, 3); + if (fill == .nonfinite) plant(t.x, shape.k); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + if (t.form.add) t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.r)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Conv) !void { + const p = t.rig.begin(); + const s: ks.ConvShape = .{ .rows = t.shape.rows, .cols = t.cols(), .in_rows = t.inRows(), .in_cols = t.inCols(), .cin = t.shape.k, .cout = t.shape.n, .ldw = t.ldw(), .taps = t.form.taps }; + ks.conv(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), if (t.form.add) ref(t.ins[3]) else null, ref(t.out), s); + try t.rig.finish(p); + } + + pub fn count(t: *Conv) usize { + return t.pixels() * t.shape.n; + } + + /// Output (p, o)'s products in the kernel's order, tap then input channel; a tap past the right or lower edge is 0. + fn terms(t: *const Conv, p: usize, o: usize, out: []f64) void { + const K = t.shape.k; + const oy = p / t.cols(); + const ox = p % t.cols(); + const stride: usize = if (t.form.taps == 9) 2 else 1; + for (0..t.form.taps) |tap| { + const iy = oy * stride + tap / 3; + const ix = ox * stride + tap % 3; + for (0..K) |ci| { + const inside = iy < t.inRows() and ix < t.inCols(); + out[tap * K + ci] = if (inside) @as(f64, t.x[(iy * t.inCols() + ix) * K + ci]) * t.w[(tap * K + ci) * t.ldw() + o] else 0; + } + } + } + + pub fn reference(t: *Conv, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, t.form.taps * t.shape.k); + defer t.a.free(buf); + for (0..t.pixels()) |p| for (0..t.shape.n) |o| { + t.terms(p, o, buf); + var sum: f64 = 0; + var m: f64 = 0; + for (buf) |v| { + sum += v; + m += @abs(v); + } + const i = p * t.shape.n + o; + const r: f64 = if (t.form.add) t.r[i] else 0; + want[i] = sum + t.b[o] + r; + m += @abs(@as(f64, t.b[o])) + @abs(r); + mag[i] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *Conv, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, t.form.taps * t.shape.k); + defer t.a.free(buf); + for (0..t.pixels()) |p| for (0..t.shape.n) |o| { + t.terms(p, o, buf); + const i = p * t.shape.n + o; + out[i] = contract.rows.sumIn(acc, order, buf) + t.b[o] + (if (t.form.add) t.r[i] else 0); + }; + } +}; + +// --- iv_norm: GroupNorm by given statistics, then SiLU or not ------------------------------------------------------- + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + silu: bool, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + stats: [64]f32 = undefined, + gamma: []f32 = &.{}, + beta: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: pixels; n: channels (32 groups of n / 32). + pub fn spec(name: []const u8, silu: bool) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = id.noSums, + .ops = 4 + id.exp_ulps + id.div_ulps + 1, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{ .size = 128, .step = 32, .min = 32 }, .{} }, + .served = if (silu) &.{ .{ .rows = 4100, .n = 128, .k = 1 }, .{ .rows = 4100, .n = 256, .k = 1 }, .{ .rows = 4100, .n = 512, .k = 1 } } else &.{.{ .rows = 4100, .n = 512, .k = 1 }}, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + inline for (.{ &t.x, &t.gamma, &t.beta }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.silu), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Norm); + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 32 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.n); + t.gamma = try t.a.alloc(f32, shape.n); + t.beta = try t.a.alloc(f32, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // values stay within +-60 after the norm, short of fp32 exp's overflow in the SiLU + uniform(rng, t.x, -4, 4); + uniform(rng, t.stats[0..32], -1, 1); + uniform(rng, t.stats[32..], 0.2, 3); + uniform(rng, t.gamma, -2, 2); + uniform(rng, t.beta, -1, 1); + if (fill == .nonfinite) plant(t.x, shape.n); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(&t.stats)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.gamma)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.beta)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Norm) !void { + const p = t.rig.begin(); + ks.normed(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.ins[3]), ref(t.out), t.shape.rows, t.shape.n, t.silu); + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * t.shape.n; + } + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const size = t.shape.n / 32; + for (t.x, 0..) |xv, i| { + const ch = i % t.shape.n; + const grp = ch / size; + const x: f64 = xv; + const m: f64 = t.stats[grp]; + const r: f64 = t.stats[32 + grp]; + const gm: f64 = t.gamma[ch]; + const b: f64 = t.beta[ch]; + const v = (x - m) * r * gm + b; + var vr = Range.of(x - m).fp32(0.5).mul(Range.of(r)).fp32(0.5).mul(Range.of(gm)).fp32(0.5).add(Range.of(b)).fp32(0.5); + var w = v; + if (t.silu) { + w = v / (1 + @exp(-v)); + const e = (Range{ .lo = @exp(-vr.hi), .hi = @exp(-vr.lo) }).fp32(id.exp_ulps); + vr = vr.mul(e.add(Range.of(1)).fp32(0.5).recip()).fp32(id.div_ulps); + } + t.ends.put(want, mag, i, w, vr); + } + } + pub fn bounds(t: *Norm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Norm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; + +// --- iv_tokens: the posterior's sample, patchified and normalized by the batch norm's statistics -------------------- + +pub const Tokens = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + moments: []f32 = &.{}, + noise: []f32 = &.{}, + m: [128]f32 = undefined, + s: [128]f32 = undefined, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const pad = 5; + + /// Rows: the latent grid's rows; n: the 128 channels of a row; k: the grid's columns. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = id.noSums, + .ops = id.exp_ulps + id.div_ulps + 3, + .out = .f32, + .tiles = .{ .{ .size = 2, .step = 2, .min = 2 }, .{}, .{ .size = 2, .step = 2, .min = 2 } }, + .served = &.{ .{ .rows = 128, .n = 128, .k = 128 }, .{ .rows = 6, .n = 128, .k = 10 } }, + }; + } + + pub fn deinit(t: *Tokens) void { + t.a.free(t.moments); + t.a.free(t.noise); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Tokens, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Tokens); + + fn latents(t: *const Tokens) usize { + return t.shape.rows * t.shape.k; + } + + pub fn prepare(t: *Tokens, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != 128 or shape.rows == 0 or shape.k == 0 or shape.rows % 2 != 0 or shape.k % 2 != 0) return false; + t.rig.reset(); + t.a.free(t.moments); + t.a.free(t.noise); + t.shape = shape; + const px = t.latents(); + t.moments = try t.a.alloc(f32, px * 64); + t.noise = try t.a.alloc(f32, px * 32); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..px) |p| { + uniform(rng, t.moments[p * 64 ..][0..32], -3, 3); + // logvars past both of the clamp's ends (-30, 20) as well as inside them + uniform(rng, t.moments[p * 64 + 32 ..][0..32], -40, 30); + } + uniform(rng, t.noise, -3, 3); + uniform(rng, &t.m, -1, 1); + uniform(rng, &t.s, 0.5, 4); + if (fill == .nonfinite) { + plant(t.moments, 64); + plant(t.noise, 32); + t.moments[5 * 64 + 40] = nan; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.moments)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.noise)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(&t.m)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(&t.s)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Tokens) !void { + const p = t.rig.begin(); + ks.tokens(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.ins[3]), ref(t.out), t.shape.rows, t.shape.k, pad); + try t.rig.finish(p); + } + + pub fn count(t: *Tokens) usize { + return (pad + t.latents() / 4) * 128; + } + + /// torch.clamp(x, -30, 20): NaN stays NaN. + fn clamp(x: f64) f64 { + return if (std.math.isNan(x)) x else std.math.clamp(x, -30, 20); + } + + pub fn reference(t: *Tokens, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (0..pad * 128) |i| t.ends.put(want, mag, i, 0, Range.of(0)); + const w2 = t.shape.k / 2; + for (0..t.latents() / 4) |patch| for (0..128) |k| { + const ch = k / 4; + const py = (patch / w2) * 2 + (k % 4) / 2; + const px = (patch % w2) * 2 + k % 2; + const p = py * t.shape.k + px; + const mean: f64 = t.moments[p * 64 + ch]; + const lv = clamp(t.moments[p * 64 + 32 + ch]); + const n: f64 = t.noise[p * 32 + ch]; + const sd = @exp(0.5 * lv); + const z = mean + sd * n; + const v = (z - t.m[k]) / t.s[k]; + const er = (Range{ .lo = @exp(0.5 * lv), .hi = @exp(0.5 * lv) }).fp32(id.exp_ulps); + const zr = er.mul(Range.of(n)).fp32(0.5).add(Range.of(mean)).fp32(0.5); + const vr = zr.sub(Range.of(t.m[k])).fp32(0.5).mul(Range.of(1 / @as(f64, t.s[k]))).fp32(id.div_ulps); + t.ends.put(want, mag, (pad + patch) * 128 + k, v, vr); + }; + } + pub fn bounds(t: *Tokens, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Tokens, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; diff --git a/zig/tests/contract/metal_idol_vae_wino.zig b/zig/tests/contract/metal_idol_vae_wino.zig new file mode 100644 index 0000000000..9e55840dae --- /dev/null +++ b/zig/tests/contract/metal_idol_vae_wino.zig @@ -0,0 +1,261 @@ +//! iv_wino under the contract: Winograd F(2x2, 3x3) against the float64 convolution, in Winograd's own float64 form. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); + +const c = id.c; +const g = id.g; +const ks = st.ks; +const ref = st.ref; +const uniform = st.uniform; +const plant = st.plant; + +/// B^T, G and A^T of F(2x2, 3x3). +const BT = [4][4]f64{ .{ 1, 0, -1, 0 }, .{ 0, 1, 1, 0 }, .{ 0, -1, 1, 0 }, .{ 0, 1, 0, -1 } }; +const G = [4][3]f64{ .{ 1, 0, 0 }, .{ 0.5, 0.5, 0.5 }, .{ 0.5, -0.5, 0.5 }, .{ 0, 0, 1 } }; +const AT = [2][4]f64{ .{ 1, 1, 1, 0 }, .{ 0, 1, -1, -1 } }; + +/// A 3x3 form the encoder launches: with a residual or not. +pub const Form = struct { label: []const u8, add: bool = false, served: []const c.Shape }; + +/// The pictures here are 8 x 18 (the encoder's up to 1024 x 1024): the sums' depth is the input channels alone. +pub const forms = [_]Form{ + .{ .label = "iv_wino 3x3", .served = &.{ .{ .rows = 8, .n = 128, .k = 128 }, .{ .rows = 8, .n = 256, .k = 128 }, .{ .rows = 8, .n = 256, .k = 256 }, .{ .rows = 8, .n = 512, .k = 256 }, .{ .rows = 8, .n = 512, .k = 512 }, .{ .rows = 8, .n = 64, .k = 512 } } }, + .{ .label = "iv_wino 3x3 + residual", .add = true, .served = &.{ .{ .rows = 8, .n = 128, .k = 128 }, .{ .rows = 8, .n = 256, .k = 256 }, .{ .rows = 8, .n = 512, .k = 512 } } }, +}; + +const cols = 18; + +pub const Wino = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: Form, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + /// the weights as released [cout, cin, 9], and their transform as the decoder packs it [16, cin, ldw] + w: []f32 = &.{}, + u: []f32 = &.{}, + b: []f32 = &.{}, + r: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// The transforms' roundings (2 in B^T d B, 1 in G g G^T), cin products, A^T M A's 4, the bias and the residual. + fn depth(s: c.Shape) usize { + return s.k + 3 + 4 + 2; + } + + /// Rows: the output picture's rows (18 columns); n: output channels; k: input channels. + pub fn spec(name: []const u8, f: Form) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 4, .step = 2, .min = 2 }, .{ .size = 64 }, .{ .size = 16, .step = 16, .min = 16 } }, + .served = f.served, + }; + } + + pub fn deinit(t: *Wino) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Wino) void { + inline for (.{ &t.x, &t.w, &t.u, &t.b, &t.r }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + } + pub fn kernel(t: *Wino, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Wino); + + fn ldw(t: *const Wino) usize { + return std.mem.alignForward(usize, t.shape.n, 64); + } + fn pixels(t: *const Wino) usize { + return t.shape.rows * cols; + } + + pub fn prepare(t: *Wino, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows % 2 != 0 or shape.n == 0 or shape.k == 0 or shape.k % 16 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const K = shape.k; + const L = t.ldw(); + const in_px = t.pixels(); + t.x = try t.a.alloc(f32, in_px * K); + t.w = try t.a.alloc(f32, shape.n * K * 9); + t.u = try t.a.alloc(f32, 16 * K * L); + t.b = try t.a.alloc(f32, shape.n); + t.r = try t.a.alloc(f32, if (t.form.add) t.pixels() * shape.n else 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // equal pixels, mirrored channel halves, weights 1 then -1: quarters throughout, sums cancel + const half = K / 2; + for (0..in_px) |p| for (0..K) |ci| { + const pos = if (ci < half) ci else ci - half; + const v: f32 = if (pos % 2 == 0) 127 else 1021; + t.x[p * K + ci] = if (ci < half) v else v + @as(f32, if (pos < 4) 1 else 0); + }; + for (0..shape.n) |o| for (0..K) |ci| @memset(t.w[(o * K + ci) * 9 ..][0..9], if (ci < half) 1 else -1); + @memset(t.b, 0); + for (t.r) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + } else { + uniform(rng, t.x, -2, 2); + uniform(rng, t.w, -0.1, 0.1); + uniform(rng, t.b, -1, 1); + uniform(rng, t.r, -3, 3); + if (fill == .nonfinite) plant(t.x, K); + } + @memset(t.u, 0); + for (0..shape.n) |o| for (0..K) |ci| { + for (ks.transform(t.w[(o * K + ci) * 9 ..][0..9]), 0..) |v, pt| t.u[(pt * K + ci) * L + o] = v; + }; + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.u)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + if (t.form.add) t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.r)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Wino) !void { + const p = t.rig.begin(); + const s: ks.ConvShape = .{ .rows = t.shape.rows, .cols = cols, .in_rows = t.shape.rows, .in_cols = cols, .cin = t.shape.k, .cout = t.shape.n, .ldw = t.ldw() }; + ks.winograd(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), if (t.form.add) ref(t.ins[3]) else null, ref(t.out), s); + try t.rig.finish(p); + } + + pub fn count(t: *Wino) usize { + return t.pixels() * t.shape.n; + } + + /// The input at picture position (y, x), zero in the padding. + fn at(t: *const Wino, y: isize, x: isize, ci: usize) f64 { + if (y < 0 or x < 0 or y >= t.shape.rows or x >= cols) return 0; + const uy: usize = @intCast(y); + const ux: usize = @intCast(x); + return t.x[(uy * cols + ux) * t.shape.k + ci]; + } + + /// B^T d B of tile (ty, tx)'s patch for channel ci, and the same of |d| (the magnitude its roundings act on). + fn tileV(t: *const Wino, ty: usize, tx: usize, ci: usize, v: *[16]f64, va: *[16]f64) void { + var d: [4][4]f64 = undefined; + for (0..4) |i| for (0..4) |j| { + d[i][j] = t.at(@as(isize, @intCast(ty * 2 + i)) - 1, @as(isize, @intCast(tx * 2 + j)) - 1, ci); + }; + for (0..4) |i| for (0..4) |j| { + var s: f64 = 0; + var m: f64 = 0; + for (0..4) |a| for (0..4) |b| { + const k = BT[i][a] * BT[j][b]; + if (k == 0) continue; + s += k * d[a][b]; + m += @abs(k * d[a][b]); + }; + v[i * 4 + j] = s; + va[i * 4 + j] = m; + }; + } + + /// G g G^T exactly, and the same of |g|. + fn exactU(t: *const Wino, o: usize, ci: usize, u: *[16]f64, ua: *[16]f64) void { + const w = t.w[(o * t.shape.k + ci) * 9 ..][0..9]; + for (0..4) |i| for (0..4) |j| { + var s: f64 = 0; + var m: f64 = 0; + for (0..3) |a| for (0..3) |b| { + const k = G[i][a] * G[j][b]; + s += k * w[a * 3 + b]; + m += @abs(k * w[a * 3 + b]); + }; + u[i * 4 + j] = s; + ua[i * 4 + j] = m; + }; + } + + /// Each output as A^T M A, M each point's products V U summed by `sum` (U exact or as packed); `mag` its bound's. + fn each(t: *const Wino, packed_u: bool, sum: anytype, out: []f64, mag: ?[]f64) !void { + const K = t.shape.k; + const vs = try t.a.alloc([16]f64, K); + defer t.a.free(vs); + const vas = try t.a.alloc([16]f64, K); + defer t.a.free(vas); + const terms = try t.a.alloc(f64, K); + defer t.a.free(terms); + const us = try t.a.alloc([16]f64, K); + defer t.a.free(us); + const uas = try t.a.alloc([16]f64, K); + defer t.a.free(uas); + for (0..t.shape.rows / 2) |ty| for (0..cols / 2) |tx| { + for (0..K) |ci| t.tileV(ty, tx, ci, &vs[ci], &vas[ci]); + for (0..t.shape.n) |o| { + for (0..K) |ci| t.exactU(o, ci, &us[ci], &uas[ci]); + var m: [16]f64 = undefined; + var ma: [16]f64 = undefined; + for (0..16) |pt| { + var mm: f64 = 0; + for (0..K) |ci| { + const uv = if (packed_u) @as(f64, t.u[(pt * K + ci) * t.ldw() + o]) else us[ci][pt]; + terms[ci] = vs[ci][pt] * uv; + mm += vas[ci][pt] * uas[ci][pt]; + } + m[pt] = sum.of(terms); + ma[pt] = mm; + } + for (0..2) |rr| for (0..2) |cc| { + var y: f64 = 0; + var ya: f64 = 0; + for (0..4) |i| for (0..4) |j| { + const k = AT[rr][i] * AT[cc][j]; + if (k == 0) continue; + y += k * m[i * 4 + j]; + ya += @abs(k) * ma[i * 4 + j]; + }; + const px = (ty * 2 + rr) * cols + tx * 2 + cc; + const idx = px * t.shape.n + o; + const res: f64 = if (t.form.add) t.r[idx] else 0; + out[idx] = y + t.b[o] + res; + if (mag) |mg| { + const total = ya + @abs(@as(f64, t.b[o])) + @abs(res); + mg[idx] = if (std.math.isFinite(total)) total else 0; + } + }; + } + }; + } + + const Exact = struct { + fn of(_: Exact, terms: []const f64) f64 { + var s: f64 = 0; + for (terms) |v| s += v; + return s; + } + }; + const Replay = struct { + acc: c.Accum, + order: g.Order, + fn of(r: Replay, terms: []const f64) f64 { + return contract.rows.sumIn(r.acc, r.order, terms); + } + }; + + pub fn reference(t: *Wino, want: []f64, mag: []f64) !void { + try t.each(false, Exact{}, want, mag); + } + + pub fn simulate(t: *Wino, acc: c.Accum, order: g.Order, out: []f64) !void { + try t.each(true, Replay{ .acc = acc, .order = order }, out, null); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index ece8073ecc..ff3088ccc0 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -102,6 +102,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol vae ")) { + text = sources.idol_vae; } else if (std.mem.startsWith(u8, d.name, "idol ")) { text = sources.idol; frags = true; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 9799191182..4de4d69394 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -429,6 +429,8 @@ pub fn main(init: std.process.Init) !u8 { defer moss_image.destroy(); const idol = try @import("metal_idol.zig").Suite.create(a, device, queue, &kernels, o.only); defer idol.destroy(); + const idol_vae = try @import("metal_idol_vae.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol_vae.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { diff --git a/zig/tests/idol_vae.zig b/zig/tests/idol_vae.zig new file mode 100644 index 0000000000..3a0e38b5d5 --- /dev/null +++ b/zig/tests/idol_vae.zig @@ -0,0 +1,132 @@ +//! tf-idol-vae: Idol's image encoder on a case from tools/zig/idol_case.py's vae, gated against float64 and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const vae = tf.idol.vae; +const ref = tf.idol.reference; + +const usage = + \\tf-idol-vae VAE_DIR CASE_DIR [FEATURES] + \\VAE_DIR is FLUX.2-klein-4B's vae/ folder (config.json, diffusion_pytorch_model.safetensors). CASE_DIR (idol_case.py vae) + \\holds vae.safetensors: the picture [1024, 1024, 3] in [0, 1], the SDK's posterior noise [128, 128, 32], the SDK's + \\condition rows (`sdk`, [4101, 128]) and float64's on the same picture and noise (`ruler`). Our rows' worst row and + \\1 - cosine against float64 must be within twice the SDK's own, and two more encodes must repeat them byte for byte. + \\With vae_stages.safetensors, each stage is compared with float64's too. FEATURES, a safetensors file holding + \\`feature2` [1, 4101, 128] (the SDK's own run), is compared as well. CASE_DIR gets tf-idol-vae.json and rows.f32. +; + +/// Encodes timed after the gated one. +const timed = 5; + +fn take(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) ![]f32 { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + if (t.dtype != .f32 or t.numel() != n) { + std.debug.print("{s} is not {d} float32 values\n", .{ name, n }); + return error.BadCaseTensor; + } + const out = try a.alloc(f32, n); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn json(list: *std.ArrayList(u8), a: std.mem.Allocator, key: []const u8, m: ref.Match) !void { + try print(list, a, ", \"{s}\": {{\"max_abs\": {d:.6}, \"cosine\": {d:.12}, \"worst_row\": {d:.8}}}", .{ key, m.max_abs, m.cosine, m.worst_row }); +} + +const stage_names = [_][]const u8{ "down.0", "down.1", "down.2", "down.3", "mid", "moments" }; +const stage_sizes = [_]usize{ 512 * 512 * 128, 256 * 256 * 256, 128 * 128 * 512, 128 * 128 * 512, 128 * 128 * 512, 128 * 128 * 64 }; +const stage_widths = [_]usize{ 128, 256, 512, 512, 512, 64 }; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 3 and args.len != 4) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/vae.safetensors", .{args[2]})); + defer case.close(io); + const n = vae.rows * vae.channels; + const image = try take(a, &case, "image", vae.size * vae.size * 3); + const noise = try take(a, &case, "noise", vae.latent * vae.latent * 32); + const sdk = try take(a, &case, "sdk", n); + const ruler = try take(a, &case, "ruler", n); + var stages_file: ?tf.safetensors.File = tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/vae_stages.safetensors", .{args[2]})) catch null; + defer if (stages_file) |*f| f.close(io); + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const enc = try vae.Encoder.load(std.heap.page_allocator, io, device, args[1]); + defer enc.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: encoder loaded in {d:.2} s\n", .{ device.name(), load_s }); + var out: std.ArrayList(u8) = .empty; + try print(&out, a, "{{\"device\": \"{s}\", \"load_s\": {d:.3}", .{ device.name(), load_s }); + + if (stages_file) |*f| { + var stages: vae.Stages = .{}; + var mine: [stage_names.len][]f32 = undefined; + for (&mine, stage_sizes) |*m, s| m.* = try a.alloc(f32, s); + for (0..4) |i| stages.down[i] = mine[i]; + stages.mid = mine[4]; + stages.moments = mine[5]; + const scratch = try a.alloc(f32, n); + _ = try enc.encode(image, noise, scratch, stages); + try print(&out, a, ", \"stages\": {{\"_\": 0", .{}); + for (stage_names, mine, stage_widths) |name, m, w| { + const want = try take(a, f, name, m.len); + const r = ref.match(m, want, w); + std.debug.print("stage {s} against float64: worst pixel {e:.3}, 1 - cosine {e:.3}, max |diff| {e:.3}\n", .{ name, r.worst_row, 1 - r.cosine, r.max_abs }); + try json(&out, a, name, r); + } + try out.append(a, '}'); + } + + const rows = try a.alloc(f32, n); + const first = try enc.encode(image, noise, rows, .{}); + const v = vae.gate(rows, sdk, ruler); + std.debug.print("against float64: ours worst row {e:.3}, 1 - cosine {e:.3}; the SDK's {e:.3}, {e:.3}: {s}\n", .{ v.ours.worst_row, 1 - v.ours.cosine, v.sdk.worst_row, 1 - v.sdk.cosine, if (v.pass) "pass" else "FAIL" }); + const vs_sdk = ref.match(rows, sdk, vae.channels); + std.debug.print("against the SDK: worst row {d:.4}, cosine {d:.8}, max |diff| {d:.4}\n", .{ vs_sdk.worst_row, vs_sdk.cosine, vs_sdk.max_abs }); + try json(&out, a, "ours_vs_float64", v.ours); + try json(&out, a, "sdk_vs_float64", v.sdk); + try json(&out, a, "ours_vs_sdk", vs_sdk); + if (args.len == 4) { + var feats = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}", .{args[3]})); + defer feats.close(io); + const theirs = try take(a, &feats, "feature2", n); + const m = ref.match(rows, theirs, vae.channels); + const same = std.mem.eql(u8, std.mem.sliceAsBytes(theirs), std.mem.sliceAsBytes(sdk)); + std.debug.print("against {s}: worst row {d:.4}, cosine {d:.8} (the case's SDK rows {s} it)\n", .{ args[3], m.worst_row, m.cosine, if (same) "equal" else "differ from" }); + try json(&out, a, "ours_vs_features", m); + try print(&out, a, ", \"case_sdk_equals_features\": {}", .{same}); + } + + const again = try a.alloc(f32, n); + var repeat = true; + var times: [timed]f64 = undefined; + var gpus: [timed]f64 = undefined; + for (×, &gpus) |*t, *gp| { + const tm = try enc.encode(image, noise, again, .{}); + t.* = tm.wall; + gp.* = tm.gpu; + repeat = repeat and std.mem.eql(u8, std.mem.sliceAsBytes(rows), std.mem.sliceAsBytes(again)); + } + std.mem.sort(f64, ×, {}, std.sort.asc(f64)); + std.mem.sort(f64, &gpus, {}, std.sort.asc(f64)); + std.debug.print("{d} more encodes: {s}; first {d:.3} s, median {d:.4} s ({d:.4} s on the GPU)\n", .{ timed, if (repeat) "identical bytes" else "DIFFERENT bytes", first.wall, times[timed / 2], gpus[timed / 2] }); + try print(&out, a, ", \"gate_pass\": {}, \"repeat_identical\": {}, \"first_s\": {d:.4}, \"median_s\": {d:.4}, \"median_gpu_s\": {d:.4}}}\n", .{ v.pass, repeat, first.wall, times[timed / 2], gpus[timed / 2] }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/rows.f32", .{args[2]}), .data = std.mem.sliceAsBytes(rows) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-vae.json", .{args[2]}), .data = out.items }); + return if (v.pass and repeat) 0 else 1; +} From c72a497c12c70bcd814fd16d75645ddbe5202423 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 02:02:13 +0800 Subject: [PATCH 08/21] idol: native Metal octree and Gaussian decoders, tf-idol-decode The octree decoder (4 blocks) and Gaussian decoder (16 blocks) run natively on Metal from the released checkpoint. Host-side pieces: systematic resampling with a pluggable RNG (native PRNG, or the SDK's recorded draws), compaction, leaves and jitter, activations, the PLY writer and the camera JSON. tf-idol-decode decodes a latent and gates it against an idol_case.py `decode` case. New kernels live in idol_decoder.metal: the point embedding, K/V and head packing, an fp32 attention for <= 8 query rows, and a copy of the idol attention whose P.V uses an fp16 hi/lo split. This removes the -2^-12 per-pass bias that fp32 P picks up on the tensor unit. Gates are measured against fp32, allowing twice the worst of: the SDK's own distance, the fp16 floor, and seeded nudged SDK reruns. Further gates: per-block teacher-forced checks, counts identical given the SDK's probabilities and draws, activated fields, the PLY's columns, the camera, and byte-identical repeats. The comparators are covered by plant-a-mismatch unit tests. idol_case.py gains `decode`, which pins decoder.py, gaussian.py and camera.py by sha256 and records levels, draws, counts, fields, the PLY and timings. --- tools/zig/idol_case.py | 292 ++++++++++++++- zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol_decoder.metal | 306 ++++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/families/idol/decode.zig | 144 ++++++++ zig/src/families/idol/decode_reference.zig | 171 +++++++++ zig/src/families/idol/decoder.zig | 298 ++++++++++++++++ zig/src/families/idol/decoder_config.zig | 147 ++++++++ zig/src/families/idol/decoder_kernels.zig | 122 +++++++ zig/src/families/idol/decoder_weights.zig | 240 +++++++++++++ zig/src/families/idol/gaussian.zig | 206 +++++++++++ zig/src/families/idol/idol.zig | 19 +- zig/src/families/idol/octree.zig | 256 +++++++++++++ zig/tests/idol_decode.zig | 395 +++++++++++++++++++++ 14 files changed, 2592 insertions(+), 7 deletions(-) create mode 100644 zig/kernels/metal/idol_decoder.metal create mode 100644 zig/src/families/idol/decode.zig create mode 100644 zig/src/families/idol/decode_reference.zig create mode 100644 zig/src/families/idol/decoder.zig create mode 100644 zig/src/families/idol/decoder_config.zig create mode 100644 zig/src/families/idol/decoder_kernels.zig create mode 100644 zig/src/families/idol/decoder_weights.zig create mode 100644 zig/src/families/idol/gaussian.zig create mode 100644 zig/src/families/idol/octree.zig create mode 100644 zig/tests/idol_decode.zig diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py index 83453a8e54..18bcc8b6bb 100644 --- a/tools/zig/idol_case.py +++ b/tools/zig/idol_case.py @@ -14,6 +14,10 @@ case WEIGHTS SDK OUT --features F [F] noise, each step's input and both velocities, the final sample blocks WEIGHTS SDK CASE [--step S] every block's rows at step S, and each block in float32 on its input exact WEIGHTS SDK CASE [--steps N] both velocities in float32 on the case's first N inputs (the ruler) + decode WEIGHTS SDK CASE OUT [--gaussians N] [--ruler] [--time-only] + the SDK's decode of CASE's final sample: each octree level's parents, + probabilities, draws and counts, the anchors, the Gaussian decoder's raw + rows, the activated Gaussians, the PLY and the camera (OUT/decode.*) WEIGHTS is a speridlabs/idol snapshot; SDK an idol-sdk checkout at the pinned commit. F holds feature1 (DINOv3 rows [1, 4101, 1280]) and feature2 (FLUX.2 VAE rows [1, 4101, 128]), in one file or two, as the SDK's ImageEncoder returns. """ @@ -24,20 +28,30 @@ "idol/models/layers.py": "2dbf24107396194636c9a2e89e7acee1bcf91b14dba1b4cde0318cb85b3f482f", "idol/sampler.py": "532987b5599be678361c2fcf566335bb3e1b0171696f985818d10b7460bfc8b0", } +DECODE_SHA256 = { + "idol/models/decoder.py": "ffaf304f0924360782f902a126efac72084ff907b23db8721b4937f0c3175bbd", + "idol/gaussian.py": "f181ae2eaf5e44fee3055bd032d57b5545ad2e38e4d2315ae3c21a2fb4f85217", + "idol/camera.py": "5576043593454e492403b821d1e355433e71f515a2d4406a51c53c30fc097935", +} WEIGHTS_REVISION = "2731c2291fc65a797c5cb627a5e5526fc1490902" DENOISER_ARGS = {"q_token_length": 8192, "in_channels": 16, "cam_channels": 5, "model_channels": 1024, "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, "num_refiner_blocks": 2, "num_heads": 16, "mlp_ratio": 4} -def sdk(path: str): - """The SDK's denoiser and sampler, imported once each file used has the pinned commit's sha256.""" - root = Path(path) - for name, want in SDK_SHA256.items(): +def pin(root: Path, files: dict) -> None: + """Exits unless each file has the pinned commit's sha256, then puts the checkout on the import path.""" + for name, want in files.items(): digest = hashlib.sha256((root / name).read_bytes()).hexdigest() if digest != want: sys.exit(f"{root / name} has sha256 {digest}; this tool runs only idol-sdk {SDK_COMMIT[:12]}'s ({want})") - sys.path.insert(0, str(root)) + if str(root) not in sys.path: + sys.path.insert(0, str(root)) + + +def sdk(path: str): + """The SDK's denoiser and sampler, imported once each file used has the pinned commit's sha256.""" + pin(Path(path), SDK_SHA256) from idol.models.dit import LatentSeqMMFlowModel from idol.sampler import sample_flow return LatentSeqMMFlowModel, sample_flow @@ -216,9 +230,275 @@ def exact(args) -> int: return 0 +DECODER_ARGS = {"model_channels": 1024, "cond_channels": 16, "num_blocks": 4, "num_heads": 16, "mlp_ratio": 4} +GS_ARGS = {"in_channels": 3, "model_channels": 1024, "cond_channels": 16, "num_blocks": 16, "num_heads": 16, + "mlp_ratio": 4} + + +def decoders(weights: str, sdk_dir: str, device: str): + """The SDK's octree and Gaussian decoders with the snapshot's weights, strictly loaded, as IdolPipeline does.""" + import torch + from safetensors.torch import load_file + + pin(Path(sdk_dir), SDK_SHA256 | DECODE_SHA256) + from idol.models.decoder import GaussianFixedlenDecoder, OctreeProbabilityFixedlenDecoder + config = json.loads((Path(weights) / "config.json").read_text()) + d, g = config["decoder"], config["decoder_gs"] + gs_args = {k: v for k, v in g["args"].items() if k != "representation_config"} + if (d["name"], d["args"], g["name"], gs_args, config["max_voxel_level"]) != ( + "OctreeProbabilityFixedlenDecoder", DECODER_ARGS, "GaussianFixedlenDecoder", GS_ARGS, 8): + sys.exit(f"{weights}/config.json is not speridlabs/idol {WEIGHTS_REVISION[:12]}'s decoders") + built = [] + for cls, entry, name in ((OctreeProbabilityFixedlenDecoder, d, "decoder"), (GaussianFixedlenDecoder, g, "decoder_gs")): + model = cls(**entry["args"]) + model.load_state_dict(load_file(str(Path(weights) / f"{name}.safetensors")), strict=True) + built.append(model.eval().requires_grad_(False).to(torch.device(device))) + return built[0], built[1], config + + +def replay_counts(original, probs, counts, draws): + """systematic_sample again on the CPU, its torch.rand calls answered from the recorded per-row draws.""" + import torch + + c = counts.reshape(-1) + unique_n, inv = c.unique(sorted=False, return_inverse=True) + queue = [draws[(inv == i).nonzero(as_tuple=False).squeeze(1)].view(-1, 1) for i, n in enumerate(unique_n.tolist()) + if n != 0] + real = torch.rand + torch.rand = lambda *a, **k: queue.pop(0) + try: + return original(probs[None].cpu(), c[None].cpu())[0] + finally: + torch.rand = real + + +NUDGES = (1, 2, 3) + + +def nudged(model, seed: int): + """A pre-hook on the first block: 1% of its fp16 input values one ulp off, seeded (the SDK's own sensitivity).""" + import torch + + def hook(_module, args): + x = args[0] + g = torch.Generator().manual_seed(seed) + h = x.cpu() + pick = torch.rand(h.shape, generator=g) < 0.01 + up = torch.rand(h.shape, generator=g) < 0.5 + moved = torch.nextafter(h, torch.where(up, torch.full_like(h, float("inf")), torch.full_like(h, -float("inf")))) + return (torch.where(pick, moved, h).to(x.device),) + tuple(args[1:]) + + return model.blocks[0].register_forward_pre_hook(hook) + + +def sensitivity(octree, gs, xs, anchors, latent, device) -> dict: + """Each decoder again in fp16 on the SDK's inputs with nudged first-block inputs, once a seed in NUDGES.""" + import torch + + found = {} + z = latent[None].to(device) + with torch.no_grad(): + for seed in NUDGES: + for model in (octree, gs): + hook = nudged(model, seed) + if model is octree: + for lv, (x, res) in enumerate(xs, start=1): + lvl = torch.full((1,), res, dtype=torch.long, device=device) + found[f"probs16n{seed}.{lv}"] = torch.softmax(octree(x[None].to(device), lvl, z), -1)[0].float().cpu() + else: + found[f"raw16n{seed}"] = gs(anchors.to(device), z)[0].float().cpu() + hook.remove() + return found + + +def rulers(octree, gs, xs, anchors, latent, device, dev) -> dict: + """Both decoders in float32, again with the last rows rounded to fp16 once, and levels 1 to 3 block by block.""" + # rounded once: the floor no fp16 stream gets under; block by block: the SDK's rows, each block in float32 + import copy + + import torch + + found = {} + z = latent[None] + o32 = copy.deepcopy(octree).to(dev).float() + o32.dtype = torch.float32 + g32 = copy.deepcopy(gs).to(dev).float() + g32.dtype = torch.float32 + round16 = lambda _m, _i, o: o.half().float() + with torch.no_grad(): + for model, key in ((o32, "probs32"), (g32, "raw32")): + for suffix in ("", "h"): + hook = model.blocks[-1].register_forward_hook(round16) if suffix else None + if model is o32: + for lv, (x, res) in enumerate(xs, start=1): + logits = o32(x[None].to(dev), torch.full((1,), res, dtype=torch.long, device=dev), z.to(dev)) + found[f"{key}{suffix}.{lv}"] = torch.softmax(logits, dim=-1)[0].float().cpu() + else: + found[f"{key}{suffix}"] = g32(anchors.float().to(dev), z.to(dev))[0].float().cpu() + if hook: + hook.remove() + for lv, (x, res) in enumerate(xs[:3], start=1): + seen = [] + hooks = [b.register_forward_hook(lambda _m, i, o: seen.append((i[0], o))) for b in octree.blocks] + lvl = torch.full((1,), res, dtype=torch.long, device=device) + octree(x[None].to(device), lvl, z.to(device)) + for h in hooks: + h.remove() + emb = o32.adaLN_modulation(o32.l_embedder(lvl.to(dev).float())) + found[f"stem.{lv}"] = seen[0][0][0].float().cpu() + for i, (inp, o) in enumerate(seen): + found[f"block{i}.{lv}"] = o[0].float().cpu() + found[f"block{i}_f32.{lv}"] = o32.blocks[i](inp.float().to(dev), emb, z.to(dev))[0].float().cpu() + return found + + +def decode(args) -> int: + """The SDK's decode of a case's final sample, timed clean (cold, then warm), then recorded on a third run.""" + import torch + from safetensors.torch import load_file, save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + data = load_file(str(Path(args.case_dir) / "case.safetensors")) + latent, camera = data["sample"].float().contiguous(), data["sample_camera"].float().contiguous() + octree, gs, config = decoders(args.weights, args.sdk, args.device) + import idol.models.decoder as dec + from idol.camera import decode_camera + + points = max(1, args.gaussians // gs.rep_config["num_gaussians"]) + level = config["max_voxel_level"] + z = latent[None].to(args.device) + + def run(): + torch.manual_seed(args.seed) + flush(args.device) + marks = [time.perf_counter()] + anchors = dec.OctreeProbabilityFixedlenDecoder.sample(octree, z, num_points=points, level=level) + flush(args.device) + marks.append(time.perf_counter()) + raw = gs(anchors, z) + flush(args.device) + marks.append(time.perf_counter()) + g = gs.to_representation(anchors, raw)[0] + flush(args.device) + marks.append(time.perf_counter()) + return anchors, raw, g, [round(b - a, 4) for a, b in zip(marks, marks[1:])] + + timings = {} + with torch.no_grad(): + for name in ("cold", "warm"): + clean, _, g, spent = run() + timings[name] = {"octree_s": spent[0], "gaussian_decoder_s": spent[1], "activation_s": spent[2], + "decode_s": round(sum(spent), 4)} + mark = time.perf_counter() + g.save_ply(str(out / "timed.ply")) + timings["save_ply_s"] = round(time.perf_counter() - mark, 4) + (out / "timed.ply").unlink() + if args.device == "mps": + timings["mps_driver_bytes_after"] = torch.mps.driver_allocated_memory() + meta = {"seed": args.seed, "gaussians": points * gs.rep_config["num_gaussians"], "anchors": points, + "level": level, "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, + "weights": WEIGHTS_REVISION, "sdk_time": timings, "camera": decode_camera(camera)} + if args.time_only: + (out / "decode.json").write_text(json.dumps(meta, indent=1) + "\n") + print(json.dumps(meta)) + return 0 + + levels, xs, jitter = [], [], [] + forward, original, real_like = octree.forward, dec.systematic_sample, torch.rand_like + + def recorded_forward(x, lvl, c): + xs.append((x[0].float().cpu(), int(lvl[0]))) + return forward(x, lvl, c) + + def recorded_systematic(probs, counts): + draws, real = [], torch.rand + + def rand(*a, **k): + v = real(*a, **k) + draws.append(v.float().cpu().reshape(-1)) + return v + + torch.rand = rand + try: + result = original(probs, counts) + finally: + torch.rand = real + c = counts.reshape(-1) + unique_n, inv = c.unique(sorted=False, return_inverse=True) + per_row, k = torch.full((c.numel(),), float("nan")), 0 + for index, n in enumerate(unique_n.tolist()): + if n == 0: + continue + rows = (inv == index).nonzero(as_tuple=False).squeeze(1).cpu() + if k >= len(draws) or draws[k].numel() != rows.numel(): + sys.exit("systematic_sample's draws do not follow its count groups") + per_row[rows] = draws[k] + k += 1 + if k != len(draws): + sys.exit("systematic_sample drew more often than it has count groups") + levels.append({"probs": probs[0].float().cpu(), "counts_in": c.cpu(), "draws": per_row, + "counts": result[0].cpu()}) + return result + + def recorded_like(*a, **k): + v = real_like(*a, **k) + jitter.append(v.float().cpu()) + return v + + octree.forward, dec.systematic_sample, torch.rand_like = recorded_forward, recorded_systematic, recorded_like + try: + with torch.no_grad(): + anchors, raw, g, _ = run() + mods = [octree.adaLN_modulation(octree.l_embedder(torch.full((1,), float(1 << lv), device=args.device))) + .half()[0].float().cpu() for lv in range(1, level + 1)] + finally: + octree.forward, dec.systematic_sample, torch.rand_like = forward, original, real_like + g.save_ply(str(out / "reference.ply")) + offsets = torch.tensor([[i, j, k] for k in (0, 1) for j in (0, 1) for i in (0, 1)], dtype=torch.long) + tensors = {"latent": latent, "camera": camera, "freqs": octree.pos_embedder.freqs.float().cpu(), + "freqs_gs": gs.pos_embedder.freqs.float().cpu()} + mismatch = [] + for lv, ((x, res), rec) in enumerate(zip(xs, levels), start=1): + if res != 1 << lv: + sys.exit(f"octree level {lv} ran at resolution {res}") + parents = torch.round(x * (1 << (lv - 1)) - 0.5).long() + again = replay_counts(original, rec["probs"], rec["counts_in"], rec["draws"]) + mismatch.append(int((again != rec["counts"]).any(dim=1).sum())) + tensors |= {f"parents.{lv}": parents.int(), f"probs.{lv}": rec["probs"], f"draws.{lv}": rec["draws"], + f"counts_in.{lv}": rec["counts_in"].int(), f"counts.{lv}": rec["counts"].int(), + f"mods.{lv}": mods[lv - 1]} + last = levels[-1]["counts"].reshape(-1) + children = (tensors[f"parents.{level}"].long()[:, None, :] * 2 + offsets[None]).reshape(-1, 3) + tensors |= {"leaves": children[last > 0].int(), "leaf_counts": last[last > 0].int(), + "jitter": jitter[0].reshape(-1, 3), "anchors": anchors[0].float().cpu(), + "raw": raw[0].float().cpu(), "xyz": g.xyz.float().cpu(), "features_dc": g.features_dc.reshape(-1, 3).cpu(), + "scaling": g.scaling.cpu(), "rotation": g.rotation.cpu(), "opacity": g.opacity.cpu()} + if args.ruler: + tensors |= rulers(octree, gs, xs, anchors, latent, args.device, torch.device(args.ruler_device)) + tensors |= sensitivity(octree, gs, xs, anchors, latent, args.device) + save_file({k: v.contiguous() for k, v in tensors.items()}, str(out / "decode.safetensors")) + same = bool(torch.equal(clean.cpu(), anchors.cpu())) + meta |= {"parents": [int(x.shape[0]) for x, _ in xs], "leaves": int(tensors["leaves"].shape[0]), + "cpu_replay_count_mismatch_rows": mismatch, "recorded_anchors_equal_timed": same, + "ruler_device": args.ruler_device if args.ruler else None} + (out / "decode.json").write_text(json.dumps(meta, indent=1) + "\n") + print(json.dumps(meta)) + return 0 + + def main(argv: list[str]) -> int: ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) + d = sub.add_parser("decode") + for name in ("weights", "sdk", "case_dir", "out_dir"): + d.add_argument(name) + d.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + d.add_argument("--gaussians", type=int, default=262144) + d.add_argument("--seed", type=int, default=42) + d.add_argument("--ruler", action="store_true", help="both decoders again in float32 on the SDK's own inputs") + d.add_argument("--ruler-device", default="cpu", choices=("mps", "cpu", "cuda")) + d.add_argument("--time-only", action="store_true", help="the timed decode alone, nothing recorded") for name in ("case", "blocks", "exact"): p = sub.add_parser(name) p.add_argument("weights") @@ -236,7 +516,7 @@ def main(argv: list[str]) -> int: else: p.add_argument("--steps", type=int, default=50) args = ap.parse_args(argv) - return {"case": case, "blocks": blocks, "exact": exact}[args.cmd](args) + return {"case": case, "blocks": blocks, "exact": exact, "decode": decode}[args.cmd](args) if __name__ == "__main__": diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index f08c980d2f..b31097f619 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -318,6 +318,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" }, .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" }, .{ .name = "tf-idol", .path = "zig/tests/idol.zig", .about = "Idol's denoiser on a reference case: each step's two branches compared, then the guided schedule timed" }, + .{ .name = "tf-idol-decode", .path = "zig/tests/idol_decode.zig", .about = "Idol's octree and Gaussian decoders on a reference case: each level, the rows and the PLY gated, then timed" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/idol_decoder.metal b/zig/kernels/metal/idol_decoder.metal new file mode 100644 index 0000000000..77f6ae884a --- /dev/null +++ b/zig/kernels/metal/idol_decoder.metal @@ -0,0 +1,306 @@ +// Idol's octree and Gaussian decoders: their steps beside the denoiser's (idol.metal), rounded where torch rounds. + +#include +#include "../nax.h" +using namespace metal; + +// h [rows, 1024] fp32 = in_proj(p) + pe(p): per axis sin then cos of p[a] * freqs[j] * pi (170 each), then 4 zeros. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dec_embed( + const device float* p [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device float* freqs [[buffer(3)]], device float* h [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + const uint c = at.x, r = at.y; + const float x0 = p[r * 3], x1 = p[r * 3 + 1], x2 = p[r * 3 + 2]; + const float s = fma(x2, w[c * 3 + 2], fma(x1, w[c * 3 + 1], x0 * w[c * 3])) + b[c]; + float pe = 0.0f; + if (c < 1020) { + const uint a = c / 340, j = c % 340; + const float x = a == 0 ? x0 : (a == 1 ? x1 : x2); + const float ang = (x * freqs[j % 170]) * 3.14159265358979323846f; + pe = j < 170 ? metal::precise::sin(ang) : metal::precise::cos(ang); + } + h[ulong(r) * 1024 + c] = s + pe; +} + +// y16 = fp16(y + b) by column: an fp32 Linear's bias, then the cast to the blocks' fp16. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dec_half( + const device float* y [[buffer(0)]], const device float* b [[buffer(1)]], device half* o [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + o[i] = half(y[i] + b[at.x]); +} + +// kv [rows, 2048] = fp16(ctx w^T + b): to_kv of the fp16 latent [rows, 16], fp32 sums, fp16 weights and bias. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dec_kv( + const device half* ctx [[buffer(0)]], const device half* w [[buffer(1)]], const device half* b [[buffer(2)]], + device half* kv [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const uint c = at.x, r = at.y; + float s = 0.0f; + for (uint k = 0; k < 16; k++) s = fma(float(ctx[r * 16 + k]), float(w[c * 16 + k]), s); + kv[ulong(r) * 2048 + c] = half(s + float(b[c])); +} + +struct DecHeads { + uint rows; + uint stride; // fp16 values from one source row to the next + uint offset; // the first column of this tensor's 16 heads in a source row + uint norm; // 1: MultiHeadRMSNorm (F.normalize times gamma times 8), 0: a plain copy (values) +}; + +// One tensor's heads [16, rows, 64] from packed fp16 rows, RMS-normed per head in fp32 as MultiHeadRMSNorm does. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dec_heads( + const device half* src [[buffer(0)]], const device float* gamma [[buffer(1)]], device half* dst [[buffer(2)]], + constant DecHeads& p [[buffer(3)]], uint2 tg [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + const uint r = tg.x, h = tg.y * 8 + sg; + if (h >= 16) return; + const ulong from = ulong(r) * p.stride + p.offset + h * 64 + lane * 2; + const ulong to = (ulong(h) * p.rows + r) * 64 + lane * 2; + const float a0 = float(src[from]), a1 = float(src[from + 1]); + if (p.norm == 0) { + dst[to] = half(a0); + dst[to + 1] = half(a1); + return; + } + const float n = max(metal::precise::sqrt(simd_sum(a0 * a0 + a1 * a1)), 1e-12f); + const uint g = h * 64 + lane * 2; + dst[to] = half(a0 / n * gamma[g] * 8.0f); + dst[to + 1] = half(a1 / n * gamma[g + 1] * 8.0f); +} + +struct DecFew { + uint lq; + uint lk; // at most 32 * 256 keys +}; + +// o [lq, 1024] = softmax(q k^T / 8) v in fp32 for lq <= 8 (torch's MPS short-query kernel); 256 threads a row, head. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dec_attention_few( + const device half* Q [[buffer(0)]], const device half* K [[buffer(1)]], const device half* V [[buffer(2)]], + device half* O [[buffer(3)]], constant DecFew& p [[buffer(4)]], uint2 tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float q[64]; + threadgroup float part[8][64]; + threadgroup float red[9]; + const uint r = tg.x, h = tg.y; + if (t < 64) q[t] = float(Q[(ulong(h) * p.lq + r) * 64 + t]) * 0.125f; + threadgroup_barrier(mem_flags::mem_threadgroup); + const device half* kh = K + ulong(h) * p.lk * 64; + const device half* vh = V + ulong(h) * p.lk * 64; + float s[32]; + float top = -INFINITY; + for (uint i = 0; i < 32; i++) { + const uint j = t + i * 256; + s[i] = -INFINITY; + if (j >= p.lk) continue; + const device half4* kr = (const device half4*)(kh + ulong(j) * 64); + float d = 0.0f; + for (uint c = 0; c < 16; c++) { + const float4 kv = float4(kr[c]); + d = fma(q[4 * c], kv.x, fma(q[4 * c + 1], kv.y, fma(q[4 * c + 2], kv.z, fma(q[4 * c + 3], kv.w, d)))); + } + s[i] = d; + top = max(top, d); + } + top = simd_max(top); + if (lane == 0) red[sg] = top; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t == 0) { + float m = red[0]; + for (uint i = 1; i < 8; i++) m = max(m, red[i]); + red[8] = m; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + top = red[8]; + float o[64]; + for (uint c = 0; c < 64; c++) o[c] = 0.0f; + float sum = 0.0f; + for (uint i = 0; i < 32; i++) { + const uint j = t + i * 256; + if (j >= p.lk) continue; + const float e = metal::precise::exp(s[i] - top); + sum += e; + const device half4* vr = (const device half4*)(vh + ulong(j) * 64); + for (uint c = 0; c < 16; c++) { + const float4 vv = float4(vr[c]); + o[4 * c] = fma(e, vv.x, o[4 * c]); + o[4 * c + 1] = fma(e, vv.y, o[4 * c + 1]); + o[4 * c + 2] = fma(e, vv.z, o[4 * c + 2]); + o[4 * c + 3] = fma(e, vv.w, o[4 * c + 3]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + sum = simd_sum(sum); + if (lane == 0) red[sg] = sum; + for (uint c = 0; c < 64; c++) { + const float v = simd_sum(o[c]); + if (lane == 0) part[sg][c] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < 64) { + float total = 0.0f, acc = 0.0f; + for (uint i = 0; i < 8; i++) { + total += red[i]; + acc += part[i][t]; + } + O[ulong(r) * 1024 + h * 64 + t] = half(acc / total); + } +} + +// idol_attention.metal's, P V as fp16 P hi + lo: the unit takes fp32 P short (outputs low by ~2^-12, compounding). +constant constexpr float kDecMasked = -3.402823466e+38f; +constant constexpr int kDecAnyCols = 1 << 30; + +// Fold one fragment's rows into r[2] (rows home.y, home.y + 8): pairs, then lanes 1 and 8 apart, then into r. +template +inline void dec_fold(thread const tfp::frag& f, thread float (&r)[2]) { + TF_UNROLL + for (short h = 0; h < 2; h++) { + const short b = 4 * h; + float t = MAX ? max(max(f[b], f[b + 1]), max(f[b + 2], f[b + 3])) : (f[b] + f[b + 1]) + (f[b + 2] + f[b + 3]); + const float u = simd_shuffle_xor(t, ushort(1)); + t = MAX ? max(t, u) : t + u; + const float w = simd_shuffle_xor(t, ushort(8)); + t = MAX ? max(t, w) : t + w; + r[h] = MAX ? max(r[h], t) : r[h] + t; + } +} + +// One 64-query block of one head over 32-key blocks; P and F as Kernels.attention sets them (kernels.zig). +template +inline void dec_attention(const device T* Q, const device T* K, const device T* V, device T* O, const device int* P, + const device float* F, uint3 tg, uint sg, uint lane) { + constexpr int BQ = 64, BK = 32, D = 64; + const int h = int(tg.y), lq = P[1], lk = P[2]; + Q += (long(h) * lq + long(tg.x) * BQ) * D; + K += long(h) * lk * D; + V += long(h) * lk * D; + const int os = P[0] * D; + O += long(tg.x) * BQ * os + h * D; + const float scale2 = F[0] * 1.44269504089f; + const short tm = 16 * short(sg); + Q += tm * D; + const short2 home = tfp::frag_home(ushort(lane)); + + tfp::frag acc[4]; + TF_UNROLL + for (short i = 0; i < 4; i++) acc[i] = tfp::frag(0); + float top[2] = {kDecMasked, kDecMasked}, total[2] = {0.0f, 0.0f}; + const bool last_q = int(tg.x) == P[4]; + const short q_rows = short(P[6] - tm), k_rows = short(P[7]); + + for (int kb = 0; kb < P[3]; kb++) { + const bool last_k = kb == P[5]; + tfp::frag s[2] = {tfp::frag(0), tfp::frag(0)}; + TF_UNROLL + for (short d = 0; d < 4; d++) { + tfp::frag q, k0, k1; + if (!AQ && last_q) { + tfp::frag_get_in(q, Q, D, 0, 16 * d, home, q_rows, kDecAnyCols); + } else { + tfp::frag_get(q, Q, D, 0, 16 * d, home); + } + if (!AK && last_k) { + tfp::frag_get_t_in(k0, K, D, 0, 16 * d, home, k_rows, kDecAnyCols); + tfp::frag_get_t_in(k1, K, D, 16, 16 * d, home, k_rows, kDecAnyCols); + } else { + tfp::frag_get_t(k0, K, D, 0, 16 * d, home); + tfp::frag_get_t(k1, K, D, 16, 16 * d, home); + } + tfp::mma_16x32(s[0], s[1], q, k0, k1); + } + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] *= scale2; + } + if (!AK && last_k) { + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kDecMasked; + } + } + float top_new[2] = {top[0], top[1]}; + dec_fold(s[0], top_new); + dec_fold(s[1], top_new); + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = fast::exp2(s[f][e] - top_new[e >> 2]); + } + float scale_old[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) { + scale_old[r] = fast::exp2(top[r] - top_new[r]); + top[r] = top_new[r]; + total[r] = total[r] * scale_old[r]; + } + dec_fold(s[0], total); + dec_fold(s[1], total); + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * scale_old[e >> 2]; + } + tfp::frag ph[2], pl[2]; + TF_UNROLL + for (short k = 0; k < 2; k++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + ph[k][e] = half(s[k][e]); + pl[k][e] = half(s[k][e] - float(ph[k][e])); + } + } + simdgroup_barrier(mem_flags::mem_none); + TF_UNROLL + for (short d = 0; d < 4; d += 2) { + TF_UNROLL + for (short k = 0; k < 2; k++) { + tfp::frag v0, v1; + if (!AK && last_k) { + tfp::frag_get_in(v0, V, D, 16 * k, 16 * d, home, k_rows, kDecAnyCols); + tfp::frag_get_in(v1, V, D, 16 * k, 16 * d + 16, home, k_rows, kDecAnyCols); + } else { + tfp::frag_get(v0, V, D, 16 * k, 16 * d, home); + tfp::frag_get(v1, V, D, 16 * k, 16 * d + 16, home); + } + tfp::mma_16x32(acc[d], acc[d + 1], ph[k], v0, v1); + tfp::mma_16x32(acc[d], acc[d + 1], pl[k], v0, v1); + } + } + K += BK * D; + V += BK * D; + } + + threadgroup_barrier(mem_flags::mem_none); + float inv[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) inv[r] = 1.f / total[r]; + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * inv[e >> 2]; + } + O += tm * os; + if (!AQ && last_q) { + if (q_rows <= 0) return; + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put_in(acc[i], O, os, 0, 16 * i, home, q_rows, kDecAnyCols); + } else { + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put(acc[i], O, os, 0, 16 * i, home); + } +} + +#define DEC_ATTENTION(NAME, T, AQ, AK) \ + [[max_total_threads_per_threadgroup(128)]] [[kernel]] void NAME( \ + const device T* Q [[buffer(0)]], const device T* K [[buffer(1)]], const device T* V [[buffer(2)]], \ + const device int32_t* P [[buffer(3)]], const device float* F [[buffer(4)]], device T* O [[buffer(5)]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ + uint3 tg [[threadgroup_position_in_grid]]) { \ + dec_attention(Q, K, V, O, P, F, tg, sg, lane); \ + } +DEC_ATTENTION(dec_attention_00, half, false, false) +DEC_ATTENTION(dec_attention_01, half, false, true) +DEC_ATTENTION(dec_attention_10, half, true, false) +DEC_ATTENTION(dec_attention_11, half, true, true) diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index ef3221a88b..beee70dd9f 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -135,6 +135,8 @@ pub const qwen_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile(" pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal"); /// Idol's denoiser: the prefill GEMM template, then its own steps and head-64 attention. pub const idol = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol.metal") ++ @embedFile("idol_attention.metal"); +/// Idol's decoders' own steps and attention (families/idol/decoder_kernels.zig); projections, norms: the denoiser's. +pub const idol_decoder = @embedFile("idol_decoder.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/families/idol/decode.zig b/zig/src/families/idol/decode.zig new file mode 100644 index 0000000000..e80dacb608 --- /dev/null +++ b/zig/src/families/idol/decode.zig @@ -0,0 +1,144 @@ +//! IdolPipeline.decode on Metal: the octree's 8 levels to anchors, the Gaussian decoder, then the Gaussians. +const std = @import("std"); +const mtl = @import("metal"); +const dc = @import("decoder_config.zig"); +const octree = @import("octree.zig"); +const gs = @import("gaussian.zig"); +const Decoder = @import("decoder.zig").Decoder; + +pub const Timing = struct { octree: f64 = 0, octree_gpu: f64 = 0, gaussian: f64 = 0, gaussian_gpu: f64 = 0, activation: f64 = 0 }; + +/// Each level's parents and their counts, kept when a caller asks (for comparing a run's topology). +pub const Trace = struct { + parents: [dc.max_level][]const [3]u32 = @splat(&.{}), + counts: [dc.max_level][]const u32 = @splat(&.{}), + probs: [dc.max_level][]const f32 = @splat(&.{}), + + pub fn deinit(t: *Trace, gpa: std.mem.Allocator) void { + for (t.parents, t.counts, t.probs) |p, c, q| { + gpa.free(p); + gpa.free(c); + gpa.free(q); + } + t.* = .{}; + } +}; + +pub const Result = struct { + anchors: []f32, + raw: []f32, + gaussians: gs.Gaussians, + time: Timing, + + pub fn deinit(r: *Result, gpa: std.mem.Allocator) void { + gpa.free(r.anchors); + gpa.free(r.raw); + r.gaussians.free(gpa); + } +}; + +/// The anchors for `points` points: each level's probabilities on Metal, its split on the host. +pub fn anchors(gpa: std.mem.Allocator, t: *Decoder, points: usize, draws: *octree.Draws, trace: ?*Trace, time: *Timing) ![]f32 { + var level: octree.Level = .{ .coords = try gpa.dupe([3]u32, &.{.{ 0, 0, 0 }}), .counts = try gpa.dupe(u32, &.{@intCast(points)}) }; + defer { + gpa.free(level.coords); + gpa.free(level.counts); + } + for (1..dc.max_level + 1) |l| { + const rows = level.coords.len; + const centres = try gpa.alloc(f32, rows * 3); + defer gpa.free(centres); + octree.centres(level.coords, @intCast(l), centres); + const probs = try gpa.alloc(f32, rows * dc.children); + defer gpa.free(probs); + const tm = try t.octreeLogits(centres, l, probs); + time.octree += tm.wall; + time.octree_gpu += tm.gpu; + const started = mtl.clock.seconds(); + octree.softmax(probs); + const next = try octree.descend(gpa, level, probs, @intCast(l), draws, null); + time.octree += mtl.clock.seconds() - started; + if (trace) |tr| { + tr.parents[l - 1] = try gpa.dupe([3]u32, level.coords); + tr.counts[l - 1] = try gpa.dupe(u32, level.counts); + tr.probs[l - 1] = try gpa.dupe(f32, probs); + } + gpa.free(level.coords); + gpa.free(level.counts); + level = next; + } + const started = mtl.clock.seconds(); + const out = try gpa.alloc(f32, points * 3); + errdefer gpa.free(out); + try octree.anchors(level, dc.max_level, draws, out); + time.octree += mtl.clock.seconds() - started; + return out; +} + +/// The Gaussians of `latent` [8192, 16]: `gaussians` (32,768 to 1,048,576, rounded down to whole anchors). +pub fn decode(gpa: std.mem.Allocator, t: *Decoder, latent: []const f32, gaussians: usize, draws: *octree.Draws, trace: ?*Trace) !Result { + const points = try dc.anchors(gaussians); + if (points > t.cap) return error.AnchorCount; + try t.setLatent(latent); + var time: Timing = .{}; + const pts = try anchors(gpa, t, points, draws, trace, &time); + errdefer gpa.free(pts); + const raw = try gpa.alloc(f32, points * dc.raw_width); + errdefer gpa.free(raw); + const tg = try t.gaussianRows(pts, raw); + time.gaussian = tg.wall; + time.gaussian_gpu = tg.gpu; + const started = mtl.clock.seconds(); + const g = try gs.Gaussians.alloc(gpa, points * dc.per_anchor); + activateParallel(t.params(), raw, pts, g); + time.activation = mtl.clock.seconds() - started; + return .{ .anchors = pts, .raw = raw, .gaussians = g, .time = time }; +} + +/// gaussian.activate split by anchors over up to 8 threads; each Gaussian's values depend on its own row alone. +pub fn activateParallel(p: gs.Params, raw: []const f32, pts: []const f32, g: gs.Gaussians) void { + const n = pts.len / 3; + const threads = @min(8, @max(1, n / 512)); + var pool: [8]?std.Thread = @splat(null); + const per = (n + threads - 1) / threads; + for (0..threads) |i| { + const lo = @min(n, i * per); + const hi = @min(n, lo + per); + if (lo == hi) continue; + const part: gs.Gaussians = .{ + .xyz = g.xyz[lo * 3 * dc.per_anchor .. hi * 3 * dc.per_anchor], + .features_dc = g.features_dc[lo * 3 * dc.per_anchor .. hi * 3 * dc.per_anchor], + .scaling = g.scaling[lo * 3 * dc.per_anchor .. hi * 3 * dc.per_anchor], + .rotation = g.rotation[lo * 4 * dc.per_anchor .. hi * 4 * dc.per_anchor], + .opacity = g.opacity[lo * dc.per_anchor .. hi * dc.per_anchor], + }; + const args = .{ p, raw[lo * dc.raw_width .. hi * dc.raw_width], pts[lo * 3 .. hi * 3], part }; + pool[i] = std.Thread.spawn(.{}, gs.activate, args) catch blk: { + @call(.auto, gs.activate, args); + break :blk null; + }; + } + for (pool) |th| if (th) |h| h.join(); +} + +test "the parallel activation equals the serial one byte for byte" { + const a = std.testing.allocator; + const n = 1500; + const raw = try a.alloc(f32, n * dc.raw_width); + defer a.free(raw); + const pts = try a.alloc(f32, n * 3); + defer a.free(pts); + var rng = std.Random.DefaultPrng.init(5); + for (raw) |*v| v.* = rng.random().floatNorm(f32); + for (pts) |*v| v.* = rng.random().float(f32); + var pert: [dc.per_anchor * 3]f32 = undefined; + for (&pert) |*v| v.* = rng.random().floatNorm(f32); + const p: gs.Params = .{ .rep = dc.released, .base_offset_scale = -1, .perturbation = &pert }; + const one = try gs.Gaussians.alloc(a, n * dc.per_anchor); + defer one.free(a); + const many = try gs.Gaussians.alloc(a, n * dc.per_anchor); + defer many.free(a); + gs.activate(p, raw, pts, one); + activateParallel(p, raw, pts, many); + try std.testing.expectEqualSlices(f32, one.xyz.ptr[0 .. n * dc.per_anchor * 14], many.xyz.ptr[0 .. n * dc.per_anchor * 14]); +} diff --git a/zig/src/families/idol/decode_reference.zig b/zig/src/families/idol/decode_reference.zig new file mode 100644 index 0000000000..5e6fc68abd --- /dev/null +++ b/zig/src/families/idol/decode_reference.zig @@ -0,0 +1,171 @@ +//! Gating a decode against idol-sdk's: worst rows, count and topology differences, activated fields, PLY columns. +const std = @import("std"); +const gs = @import("gaussian.zig"); +const ref = @import("reference.zig"); + +/// The row of `width` with the largest relative L2 error, and that error. +pub const Worst = struct { row: usize, err: f64 }; + +pub fn worst(ours: []const f32, theirs: []const f32, width: usize) Worst { + std.debug.assert(ours.len == theirs.len and ours.len % width == 0); + var w: Worst = .{ .row = 0, .err = 0 }; + for (0..ours.len / width) |r| { + var d: f64 = 0; + var b: f64 = 0; + for (ours[r * width ..][0..width], theirs[r * width ..][0..width]) |x, y| { + d += (@as(f64, x) - y) * (@as(f64, x) - y); + b += @as(f64, y) * y; + } + const e = @sqrt(d / @max(b, 1e-30)); + if (!(e <= w.err)) w = .{ .row = r, .err = e }; + } + return w; +} + +/// The distance a gate allows twice of: the worst of the SDK's own, the fp16 floor's and the SDK's nudged reruns'. +pub fn reference(distances: []const ref.Match) ref.Match { + var r = distances[0]; + for (distances[1..]) |d| r = .{ .max_abs = @max(r.max_abs, d.max_abs), .cosine = @min(r.cosine, d.cosine), .worst_row = @max(r.worst_row, d.worst_row) }; + return r; +} + +/// Rows of `width` counts that differ, and the first. +pub const Counts = struct { rows: usize, first: ?usize }; + +pub fn counts(ours: []const u32, theirs: []const u32, width: usize) Counts { + if (ours.len != theirs.len) return .{ .rows = @max(ours.len, theirs.len) / width, .first = 0 }; + var c: Counts = .{ .rows = 0, .first = null }; + for (0..ours.len / width) |r| { + if (std.mem.eql(u32, ours[r * width ..][0..width], theirs[r * width ..][0..width])) continue; + c.rows += 1; + if (c.first == null) c.first = r; + } + return c; +} + +/// Element-wise: largest absolute difference and where, and how many exceed |d| <= rel |theirs| + abs. +pub const Field = struct { max_abs: f64, at: usize, bad: usize, exact: usize }; + +pub fn field(ours: []const f32, theirs: []const f32, rel: f64, abs: f64) Field { + std.debug.assert(ours.len == theirs.len); + var f: Field = .{ .max_abs = 0, .at = 0, .bad = 0, .exact = 0 }; + for (ours, theirs, 0..) |x, y, i| { + const d = @abs(@as(f64, x) - y); + if (@as(u32, @bitCast(x)) == @as(u32, @bitCast(y))) f.exact += 1; + if (!(d <= f.max_abs)) f = .{ .max_abs = d, .at = i, .bad = f.bad, .exact = f.exact }; + if (!(d <= rel * @abs(@as(f64, y)) + abs)) f.bad += 1; + } + return f; +} + +/// A PLY file against the SDK's: headers, vertex counts, and each property column's largest difference. +pub const Ply = struct { + header_equal: bool, + rows_equal: bool, + rows: usize, + max_abs: [gs.properties.len]f64, + worst_row: [gs.properties.len]usize, + differing: [gs.properties.len]usize, +}; + +fn body(bytes: []const u8) ?[]const u8 { + const end = "end_header\n"; + const at = std.mem.indexOf(u8, bytes, end) orelse return null; + return bytes[at + end.len ..]; +} + +pub fn ply(ours: []const u8, theirs: []const u8) Ply { + var p: Ply = .{ .header_equal = false, .rows_equal = false, .rows = 0, .max_abs = @splat(0), .worst_row = @splat(0), .differing = @splat(0) }; + const a = body(ours) orelse return p; + const b = body(theirs) orelse return p; + p.header_equal = std.mem.eql(u8, ours[0 .. ours.len - a.len], theirs[0 .. theirs.len - b.len]); + p.rows_equal = a.len == b.len and a.len % gs.row_bytes == 0; + if (!p.rows_equal) return p; + p.rows = a.len / gs.row_bytes; + for (0..p.rows) |r| for (0..gs.properties.len) |c| { + const at = r * gs.row_bytes + c * 4; + const x: f32 = @bitCast(std.mem.readInt(u32, a[at..][0..4], .little)); + const y: f32 = @bitCast(std.mem.readInt(u32, b[at..][0..4], .little)); + if (@as(u32, @bitCast(x)) == @as(u32, @bitCast(y))) continue; + p.differing[c] += 1; + const d = @abs(@as(f64, x) - y); + if (!(d <= p.max_abs[c])) { + p.max_abs[c] = d; + p.worst_row[c] = r; + } + }; + return p; +} + +/// The PLY gate: same header and rows; every column bit-equal but opacity and scales (device logs), within `tol`. +pub fn plyPass(p: Ply, tol: f64) bool { + if (!p.header_equal or !p.rows_equal) return false; + for (0..gs.properties.len) |c| { + const transcendental = c >= 9 and c <= 12; + if (transcendental) { + if (!(p.max_abs[c] <= tol)) return false; + } else if (p.differing[c] != 0) return false; + } + return true; +} + +test "the reference takes the worst distance of each kind; a planted error past twice it still fails" { + const lucky: ref.Match = .{ .max_abs = 1e-5, .cosine = 1 - 1e-9, .worst_row = 1e-4 }; + const nudged: ref.Match = .{ .max_abs = 2e-5, .cosine = 1 - 4e-9, .worst_row = 2e-4 }; + const floor: ref.Match = .{ .max_abs = 0, .cosine = 1 - 1e-10, .worst_row = 1e-5 }; + const r = reference(&.{ lucky, floor, nudged }); + try std.testing.expectEqual(nudged.worst_row, r.worst_row); + try std.testing.expectEqual(nudged.cosine, r.cosine); + try std.testing.expectEqual(lucky, reference(&.{lucky})); + try std.testing.expect(ref.within(.{ .max_abs = 0, .cosine = 1 - 7e-9, .worst_row = 3.9e-4 }, r)); + try std.testing.expect(!ref.within(.{ .max_abs = 0, .cosine = 1 - 7e-9, .worst_row = 4.1e-4 }, r)); + try std.testing.expect(!ref.within(.{ .max_abs = 0, .cosine = 1 - 9e-9, .worst_row = 1e-4 }, r)); +} + +test "a planted count difference is found at its row" { + const a = [_]u32{ 1, 0, 0, 0, 2, 0, 0, 0 }; + var b = a; + try std.testing.expectEqual(Counts{ .rows = 0, .first = null }, counts(&a, &b, 4)); + b[5] = 1; + b[4] = 1; + try std.testing.expectEqual(Counts{ .rows = 1, .first = 1 }, counts(&a, &b, 4)); + try std.testing.expectEqual(@as(usize, 2), counts(&a, b[0..4], 4).rows); +} + +test "the worst row and a field's planted error are found where they were planted" { + var x = [_]f32{ 1, 2, 3, 4, 5, 6 }; + const y = x; + x[4] = 5.5; + try std.testing.expectEqual(@as(usize, 2), worst(&x, &y, 2).row); + const f = field(&x, &y, 1e-5, 0); + try std.testing.expectEqual(@as(usize, 4), f.at); + try std.testing.expectEqual(@as(usize, 1), f.bad); + try std.testing.expectEqual(@as(usize, 5), f.exact); + x[1] = std.math.nan(f32); + try std.testing.expectEqual(@as(usize, 2), field(&x, &y, 1e-5, 0).bad); +} + +test "a PLY passes against itself; a planted header, row or quaternion bit fails" { + const a = std.testing.allocator; + var xyz = [_]f32{ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6 }; + var f = [_]f32{ 1, 2, 3, 4, 5, 6 }; + var s = [_]f32{ 0.01, 0.02, 0.03, 0.04, 0.05, 0.06 }; + var q = [_]f32{ 1, 0, 0, 0, 0.5, 0.5, 0.5, 0.5 }; + var o = [_]f32{ 0.2, 0.9 }; + const g: gs.Gaussians = .{ .xyz = &xyz, .features_dc = &f, .scaling = &s, .rotation = &q, .opacity = &o }; + const ours = try gs.ply(a, g); + defer a.free(ours); + try std.testing.expect(plyPass(ply(ours, ours), 1e-5)); + const theirs = try a.dupe(u8, ours); + defer a.free(theirs); + // the last row's rot_3, one bit off + const at = theirs.len - 4; + std.mem.writeInt(u32, theirs[at..][0..4], std.mem.readInt(u32, theirs[at..][0..4], .little) ^ 1, .little); + const p = ply(ours, theirs); + try std.testing.expect(!plyPass(p, 1e-5)); + try std.testing.expectEqual(@as(usize, 1), p.differing[16]); + try std.testing.expectEqual(@as(usize, 1), p.worst_row[16]); + try std.testing.expect(!plyPass(ply(ours, theirs[0 .. theirs.len - gs.row_bytes]), 1e-5)); + theirs[0] = 'P'; + try std.testing.expect(!plyPass(ply(ours, theirs), 1e-5)); +} diff --git a/zig/src/families/idol/decoder.zig b/zig/src/families/idol/decoder.zig new file mode 100644 index 0000000000..db519e4ebe --- /dev/null +++ b/zig/src/families/idol/decoder.zig @@ -0,0 +1,298 @@ +//! Idol's two decoders on Metal: the octree decoder's child probabilities and the Gaussian decoder's anchor rows. +const std = @import("std"); +const mtl = @import("metal"); +const dc = @import("decoder_config.zig"); +const ks = @import("kernels.zig"); +const dk = @import("decoder_kernels.zig"); +const dw = @import("decoder_weights.zig"); +const octree = @import("octree.zig"); +const gaussian = @import("gaussian.zig"); +const Ref = ks.Ref; + +const H = dc.hidden; +const F = dc.ffn; +const C = dc.context_rows; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +/// The scratch's parts for `cap` rows in bytes, in the order of Decoder's views. +fn parts(cap: usize) [17]usize { + const kv_rows = @max(cap, C); + return .{ + cap * H * 2, cap * H * 2, cap * 3 * H * 2, kv_rows * H * 2, kv_rows * H * 2, kv_rows * H * 2, + cap * H * 2, cap * H * 2, cap * F * 2, cap * H * 4, cap * H * 4, cap * dc.raw_width * 4, + cap * 3 * 4, C * dc.context_channels * 2, C * 2 * H * 2, dc.max_level * 6 * H * 2, 2 * H * 4 + 256 * 4, + }; +} + +/// Scratch bytes for `cap` rows (anchors), each part 256-aligned. +pub fn scratchBytes(cap: usize) usize { + var total: usize = 0; + for (parts(cap)) |n| total += std.mem.alignForward(usize, n, 256); + return total; +} + +pub const Decoder = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + d: dk.Kernels, + w: dw.Weights, + scratch: mtl.Buffer, + cap: usize, + // fp16 rows: the stream, a norm's output, packed qkv, heads [16, rows, 64], attention rows, a projection's output + x: Ref, + xn: Ref, + qkv: Ref, + q: Ref, + k_: Ref, + v: Ref, + att: Ref, + y: Ref, + /// the MLP's GELU rows + u: Ref, + /// fp32 rows: the stem's sums, then input_layer's (and the head's LayerNorm); the head's outputs + e32: Ref, + f32b: Ref, + out: Ref, + pts: Ref, + /// the latent [8192, 16] in fp16, and one block's to_kv rows [8192, 2048] + ctx: Ref, + kv: Ref, + /// each level's modulation [8][6][1024] fp16 + mods: Ref, + /// fp32 zeros [1024], ones fp16 [1024], then the point embedding's frequencies [170] + consts: Ref, + levels_ready: bool = false, + + /// The decoders of snapshot folder `dir`, with scratch for up to `max_anchors` anchors. + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8, max_anchors: usize) !*Decoder { + if (!device.tensorUnits()) { + std.log.err("idol: {s} has no tensor units; the decoders run on the M5's alone, the only chip checked", .{device.name()}); + return error.TensorUnitsRequired; + } + if (max_anchors == 0 or max_anchors > dc.max_anchors) return error.AnchorCount; + const t = try gpa.create(Decoder); + errdefer gpa.destroy(t); + t.* = undefined; + t.gpa = gpa; + t.device = device; + t.cap = max_anchors; + t.levels_ready = false; + t.queue = try device.queue(); + errdefer t.queue.deinit(); + t.k = try ks.Kernels.load(gpa, device); + errdefer t.k.deinit(); + t.d = try dk.Kernels.load(gpa, device); + errdefer t.d.deinit(); + t.w = try dw.load(gpa, io, device, dir); + errdefer t.w.deinit(); + t.scratch = try device.buffer(scratchBytes(max_anchors), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &t.x, &t.xn, &t.qkv, &t.q, &t.k_, &t.v, &t.att, &t.y, &t.u, &t.e32, &t.f32b, &t.out, &t.pts, &t.ctx, &t.kv, &t.mods, &t.consts }; + for (views, parts(max_anchors)) |view, n| { + view.* = .{ .buf = t.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + @memset(floats(t.consts, H), 0); + @memset(halves(t.consts.at(H * 4), H), 1.0); + octree.pointFreqs(floats(t.freqs(), dc.pe_freqs)[0..dc.pe_freqs]); + return t; + } + + pub fn deinit(t: *Decoder) void { + t.scratch.deinit(); + t.w.deinit(); + t.d.deinit(); + t.k.deinit(); + t.queue.deinit(); + t.gpa.destroy(t); + } + + fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn halves(r: Ref, n: usize) []f16 { + return @as([*]f16, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn zeros(t: *const Decoder) Ref { + return t.consts; + } + + fn ones(t: *const Decoder) Ref { + return t.consts.at(H * 4); + } + + fn freqs(t: *const Decoder) Ref { + return t.consts.at(2 * H * 4); + } + + /// Level `level`'s (1 to 8) modulation part `part` (0 shift, 1 one plus scale, 2 gate; 3 to 5 the MLP's). + fn mod(t: *const Decoder, level: usize, part: usize) Ref { + return t.mods.at(((level - 1) * 6 + part) * H * 2); + } + + /// The point embedding's frequencies as this decoder computes them. + pub fn pointFreqs(t: *const Decoder) []const f32 { + return floats(t.freqs(), dc.pe_freqs); + } + + /// The activation's fixed inputs from decoder_gs.safetensors and config.json. + pub fn params(t: *const Decoder) gaussian.Params { + return .{ .rep = t.w.rep, .base_offset_scale = t.w.base_offset_scale, .perturbation = t.w.perturbation }; + } + + /// Level `level`'s modulation as the blocks read it. + pub fn levelMods(t: *const Decoder, level: usize) []const f16 { + return halves(t.mod(level, 0), 6 * H); + } + + /// The latent both decoders attend to: [8192, 16] fp32, cast to fp16 as the SDK's `cond.type(fp16)`. + pub fn setLatent(t: *Decoder, latent: []const f32) !void { + if (latent.len != C * dc.context_channels) return error.LatentRows; + for (halves(t.ctx, latent.len), latent) |*o, v| o.* = @floatCast(v); + if (!t.levels_ready) { + for (1..dc.max_level + 1) |l| octree.levelMods(t.w.level, @intCast(l), @ptrCast(@alignCast(t.mod(l, 0).addr()))); + t.levels_ready = true; + } + } + + /// The octree decoder's logits [rows, 8] for parent centres [rows, 3] at `level` (1 to 8). + pub fn octreeLogits(t: *Decoder, centres: []const f32, level: usize, logits: []f32) !Timing { + const rows = centres.len / 3; + if (!t.levels_ready) return error.NoLatent; + if (rows == 0 or rows > t.cap or centres.len != rows * 3 or logits.len != rows * dc.children or level < 1 or level > dc.max_level) return error.OctreeRows; + @memcpy(floats(t.pts, centres.len), centres); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.stem(e, t.w.octree_stem, rows); + for (t.w.octree) |b| t.octreeBlock(e, b, level, rows); + t.head(e, t.w.octree_stem, rows, dc.children); + const timing = try finish(cb, e, started); + @memcpy(logits, floats(t.out, logits.len)); + return timing; + } + + /// The octree stream [rows, 1024] fp16 after its stem (`block` null), or after block `block` alone run on `input`. + pub fn octreeRows(t: *Decoder, centres: []const f32, input: []const f16, level: usize, block: ?usize) ![]const f16 { + const rows = centres.len / 3; + if (!t.levels_ready) return error.NoLatent; + if (rows == 0 or rows > t.cap or input.len != rows * H or level < 1 or level > dc.max_level) return error.OctreeRows; + @memcpy(floats(t.pts, centres.len), centres); + @memcpy(halves(t.x, input.len), input); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + if (block) |i| t.octreeBlock(e, t.w.octree[i], level, rows) else t.stem(e, t.w.octree_stem, rows); + _ = try finish(cb, e, mtl.clock.seconds()); + return halves(t.x, rows * H); + } + + /// One octree block: cross attention to the latent, then the MLP, each modulated by the level. + fn octreeBlock(t: *const Decoder, e: mtl.ComputeEncoder, b: dw.OctreeBlock, level: usize, rows: usize) void { + t.k.norm(e, t.x, t.zeros(), t.zeros(), t.mod(level, 1), t.mod(level, 0), t.xn, rows, dc.block_eps, .modulated); + t.cross(e, b.cross, rows); + t.k.gateAdd(e, t.x, t.y, t.mod(level, 2), rows); + t.k.norm(e, t.x, t.zeros(), t.zeros(), t.mod(level, 4), t.mod(level, 3), t.xn, rows, dc.block_eps, .modulated); + t.mlp(e, b.mlp, rows); + t.k.gateAdd(e, t.x, t.y, t.mod(level, 5), rows); + } + + /// The Gaussian decoder's raw rows [anchors, 480] fp32 for anchors [anchors, 3]. + pub fn gaussianRows(t: *Decoder, anchors: []const f32, raw: []f32) !Timing { + const rows = anchors.len / 3; + if (!t.levels_ready) return error.NoLatent; + if (rows == 0 or rows > t.cap or anchors.len != rows * 3 or raw.len != rows * dc.raw_width) return error.AnchorRows; + @memcpy(floats(t.pts, anchors.len), anchors); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.stem(e, t.w.gaussian_stem, rows); + for (t.w.gaussian) |b| t.gaussianBlock(e, b, rows); + t.head(e, t.w.gaussian_stem, rows, dc.raw_width); + const timing = try finish(cb, e, started); + @memcpy(raw, floats(t.out, raw.len)); + return timing; + } + + /// One Gaussian decoder block: self attention over the anchors, cross attention to the latent, then the MLP. + fn gaussianBlock(t: *const Decoder, e: mtl.ComputeEncoder, b: dw.GaussianBlock, rows: usize) void { + const k = &t.k; + t.selfAttention(e, b.self_attn, rows); + k.norm(e, t.x, b.norm2, b.norm2_b, t.ones(), t.ones(), t.xn, rows, dc.block_eps, .affine); + t.cross(e, b.cross, rows); + k.gateAdd(e, t.x, t.y, t.ones(), rows); + k.norm(e, t.x, t.zeros(), t.zeros(), t.ones(), t.ones(), t.xn, rows, dc.block_eps, .plain); + t.mlp(e, b.mlp, rows); + k.gateAdd(e, t.x, t.y, t.ones(), rows); + } + + /// x += to_out(attention(q_rms, k_rms, v of to_qkv(LayerNorm(x)))) over the anchors. + fn selfAttention(t: *const Decoder, e: mtl.ComputeEncoder, sa: dw.SelfAttention, rows: usize) void { + const k = &t.k; + k.norm(e, t.x, t.zeros(), t.zeros(), t.ones(), t.ones(), t.xn, rows, dc.block_eps, .plain); + k.linear(e, t.xn, sa.qkv, sa.qkv_b, t.qkv, rows, 3 * H, H, false); + t.d.heads(e, t.qkv, sa.gq, t.q, rows, 3 * H, 0); + t.d.heads(e, t.qkv, sa.gk, t.k_, rows, 3 * H, H); + t.d.heads(e, t.qkv, null, t.v, rows, 3 * H, 2 * H); + t.d.attention(e, t.q, t.k_, t.v, t.att, rows, rows); + k.linear(e, t.att, sa.out, sa.out_b, t.y, rows, H, H, false); + k.gateAdd(e, t.x, t.y, t.ones(), rows); + } + + /// x [rows, 1024] fp16 = fp16(input_layer(in_proj(p) + pe(p))), all fp32 before the cast. + fn stem(t: *const Decoder, e: mtl.ComputeEncoder, s: dw.Stem, rows: usize) void { + t.d.embed(e, t.pts, s.in_w, s.in_b, t.freqs(), t.e32, rows); + dk.gemm32(&t.k, e, t.e32, s.input, t.f32b, rows, H, H); + t.d.toHalf(e, t.f32b, s.input_b, t.x, rows, H); + } + + /// out [rows, n] fp32 = out_proj(F.layer_norm(x)) with the stream cast to fp32. + fn head(t: *const Decoder, e: mtl.ComputeEncoder, s: dw.Stem, rows: usize, n: usize) void { + t.k.norm32(e, t.x, t.zeros(), t.zeros(), t.f32b, rows, dc.out_eps, .modulated); + dk.gemm32(&t.k, e, t.f32b, s.out, t.out, rows, n, H); + t.k.bias32(e, t.out, s.out_b, rows, n); + } + + /// y = to_out(attention(q_rms(to_q(xn)), k_rms(to_kv(latent)), v)) over the latent's 8192 rows; fp32 for few rows. + fn cross(t: *const Decoder, e: mtl.ComputeEncoder, c: dw.Cross, rows: usize) void { + t.k.linear(e, t.xn, c.q, c.q_b, t.att, rows, H, H, false); + t.d.heads(e, t.att, c.gq, t.q, rows, H, 0); + t.d.kv(e, t.ctx, c.kv, c.kv_b, t.kv); + t.d.heads(e, t.kv, c.gk, t.k_, C, 2 * H, 0); + t.d.heads(e, t.kv, null, t.v, C, 2 * H, H); + if (rows <= dk.few_rows) t.d.attentionFew(e, t.q, t.k_, t.v, t.att, rows, C) else t.d.attention(e, t.q, t.k_, t.v, t.att, rows, C); + t.k.linear(e, t.att, c.out, c.out_b, t.y, rows, H, H, false); + } + + /// y = down(GELU(up(xn))). + fn mlp(t: *const Decoder, e: mtl.ComputeEncoder, m: dw.Mlp, rows: usize) void { + t.k.linear(e, t.xn, m.up, m.up_b, t.u, rows, F, H, true); + t.k.linear(e, t.u, m.down, m.down_b, t.y, rows, H, F, false); + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder, started: f64) !Timing { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol: a decoder pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return .{ .wall = mtl.clock.seconds() - started, .gpu = cb.gpuSeconds() }; + } +}; + +test "decoder scratch at a million Gaussians fits one GPU buffer" { + try std.testing.expect(scratchBytes(dc.max_anchors) < 1 << 32); + try std.testing.expect(scratchBytes(1) > C * 2 * H * 2); +} diff --git a/zig/src/families/idol/decoder_config.zig b/zig/src/families/idol/decoder_config.zig new file mode 100644 index 0000000000..cf2c035e4d --- /dev/null +++ b/zig/src/families/idol/decoder_config.zig @@ -0,0 +1,147 @@ +//! Idol's two decoders' geometry, and their config.json check: the octree decoder and the Gaussian decoder. +const std = @import("std"); + +pub const hidden = 1024; +pub const heads = 16; +pub const ffn = 4 * hidden; +pub const octree_blocks = 4; +pub const gs_blocks = 16; +/// The latent both decoders attend to: the denoiser's final sample [8192, 16]. +pub const context_rows = 8192; +pub const context_channels = 16; +/// An octree cell's children: the decoder's 8 logits a parent. +pub const children = 8; +pub const max_level = 8; +pub const per_anchor = 32; +/// The Gaussian decoder's outputs an anchor: offset 3, colour 3, scale 3, rotation 4, opacity 1, offset scale 1. +pub const raw_width = per_anchor * 15; +/// PcdAbsolutePositionEmbedderv2's frequencies an axis: channels // 3 // 2. +pub const pe_freqs = hidden / 3 / 2; +/// SinusoidalEmbedder's frequency vector (the level embedding). +pub const level_freqs = 256; +pub const level_period: f32 = 1024; +pub const block_eps: f32 = 1e-6; +pub const out_eps: f32 = 1e-5; +/// IdolPipeline.decode's Gaussian counts: 32,768 to 1,048,576, a multiple of 32. +pub const min_gaussians = 32768; +pub const max_gaussians = 1 << 20; +pub const max_anchors = max_gaussians / per_anchor; + +/// representation_config: the activations' learning-rate multipliers, offsets and biases, as python's doubles. +pub const Representation = struct { + lr_xyz: f64, + lr_features_dc: f64, + lr_scaling: f64, + lr_rotation: f64, + lr_opacity: f64, + perturbe_size: f64, + filter: f64, + scaling_bias: f64, + opacity_bias: f64, +}; + +/// The released representation_config, for tests. +pub const released: Representation = .{ .lr_xyz = 1, .lr_features_dc = 1, .lr_scaling = 1, .lr_rotation = 0.1, .lr_opacity = 1, .perturbe_size = 1.5, .filter = 0.0009, .scaling_bias = 0.004, .opacity_bias = 0.1 }; + +fn number(v: std.json.Value) !f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadConfig, + }; +} + +fn field(o: std.json.ObjectMap, key: []const u8) !std.json.Value { + return o.get(key) orelse error.BadConfig; +} + +fn object(v: std.json.Value) !std.json.ObjectMap { + return if (v == .object) v.object else error.BadConfig; +} + +fn args(root: std.json.ObjectMap, key: []const u8, class: []const u8) !std.json.ObjectMap { + const d = try object(try field(root, key)); + const name = try field(d, "name"); + if (name != .string or !std.mem.eql(u8, name.string, class)) return error.NotIdol; + return object(try field(d, "args")); +} + +fn expect(o: std.json.ObjectMap, comptime want: anytype) !void { + inline for (want) |kv| { + const v = try field(o, kv[0]); + if (v != .integer or v.integer != kv[1]) { + std.log.warn("idol: decoder {s} is not {d}; this family runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } +} + +/// The snapshot config's decoders checked against the released geometry; returns their representation_config. +pub fn check(gpa: std.mem.Allocator, text: []const u8) !Representation { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + const root = try object(doc.value); + const oct = try args(root, "decoder", "OctreeProbabilityFixedlenDecoder"); + try expect(oct, .{ .{ "model_channels", hidden }, .{ "cond_channels", context_channels }, .{ "num_blocks", octree_blocks }, .{ "num_heads", heads }, .{ "mlp_ratio", 4 } }); + const gs = try args(root, "decoder_gs", "GaussianFixedlenDecoder"); + try expect(gs, .{ .{ "in_channels", 3 }, .{ "model_channels", hidden }, .{ "cond_channels", context_channels }, .{ "num_blocks", gs_blocks }, .{ "num_heads", heads }, .{ "mlp_ratio", 4 } }); + try expect(root, .{.{ "max_voxel_level", max_level }}); + const rc = try object(try field(gs, "representation_config")); + try expect(rc, .{.{ "num_gaussians", per_anchor }}); + const lr = try object(try field(rc, "lr")); + const f = struct { + fn f(o: std.json.ObjectMap, k: []const u8) !f64 { + return number(try field(o, k)); + } + }.f; + return .{ + .lr_xyz = try f(lr, "_xyz"), + .lr_features_dc = try f(lr, "_features_dc"), + .lr_scaling = try f(lr, "_scaling"), + .lr_rotation = try f(lr, "_rotation"), + .lr_opacity = try f(lr, "_opacity"), + .perturbe_size = try f(rc, "perturbe_size"), + .filter = try f(rc, "3d_filter_kernel_size"), + .scaling_bias = try f(rc, "scaling_bias"), + .opacity_bias = try f(rc, "opacity_bias"), + }; +} + +/// Anchors for `gaussians` as IdolPipeline.decode counts them: max(1, gaussians // 32), within the card's range. +pub fn anchors(gaussians: usize) !usize { + if (gaussians < min_gaussians or gaussians > max_gaussians) return error.GaussianCount; + return gaussians / per_anchor; +} + +const sample = + \\{"decoder": {"name": "OctreeProbabilityFixedlenDecoder", "args": {"model_channels": 1024, "cond_channels": 16, + \\ "num_blocks": 4, "num_heads": 16, "mlp_ratio": 4}}, "decoder_gs": {"name": "GaussianFixedlenDecoder", "args": + \\ {"in_channels": 3, "model_channels": 1024, "cond_channels": 16, "num_blocks": 16, "num_heads": 16, "mlp_ratio": 4, + \\ "representation_config": {"num_gaussians": 32, "lr": {"_xyz": 1.0, "_features_dc": 1.0, "_opacity": 1.0, + \\ "_scaling": 1.0, "_rotation": 0.1}, "perturbe_size": 1.5, "3d_filter_kernel_size": 0.0009, + \\ "scaling_bias": 0.004, "opacity_bias": 0.1}}}, "image_size": 1024, "max_voxel_level": 8} +; + +test "the released decoders pass with their representation_config; another depth or class does not" { + const a = std.testing.allocator; + try std.testing.expectEqual(released, try check(a, sample)); + const deeper = try std.mem.replaceOwned(u8, a, sample, "\"num_blocks\": 16", "\"num_blocks\": 12"); + defer a.free(deeper); + try std.testing.expectError(error.UnsupportedGeometry, check(a, deeper)); + const level = try std.mem.replaceOwned(u8, a, sample, "\"max_voxel_level\": 8", "\"max_voxel_level\": 9"); + defer a.free(level); + try std.testing.expectError(error.UnsupportedGeometry, check(a, level)); + const other = try std.mem.replaceOwned(u8, a, sample, "GaussianFixedlenDecoder", "GaussianDecoder"); + defer a.free(other); + try std.testing.expectError(error.NotIdol, check(a, other)); + const missing = try std.mem.replaceOwned(u8, a, sample, "\"opacity_bias\"", "\"opacity\""); + defer a.free(missing); + try std.testing.expectError(error.BadConfig, check(a, missing)); +} + +test "anchors: the card's range, rounded down to whole anchors" { + try std.testing.expectEqual(@as(usize, 8192), try anchors(262144)); + try std.testing.expectEqual(@as(usize, 1024), try anchors(32768 + 31)); + try std.testing.expectError(error.GaussianCount, anchors(1000)); + try std.testing.expectError(error.GaussianCount, anchors((1 << 20) + 32)); +} diff --git a/zig/src/families/idol/decoder_kernels.zig b/zig/src/families/idol/decoder_kernels.zig new file mode 100644 index 0000000000..3c393e3cfd --- /dev/null +++ b/zig/src/families/idol/decoder_kernels.zig @@ -0,0 +1,122 @@ +//! The decoders' own Metal pipelines (idol_decoder.metal); their projections, norms and attention are the denoiser's. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const ks = @import("kernels.zig"); +const dc = @import("decoder_config.zig"); +const Ref = ks.Ref; +const bind = ks.bind; + +pub const names = [_][]const u8{ "dec_embed", "dec_half", "dec_kv", "dec_heads", "dec_attention_few", "dec_attention_00", "dec_attention_01", "dec_attention_10", "dec_attention_11" }; +const attention_first = 5; + +/// Query rows the fp32 attention takes: torch on MPS runs its fp32 vector kernel up to 8, its fp16 graph past them. +pub const few_rows = 8; + +comptime { + std.debug.assert(std.mem.eql(u8, names[attention_first], "dec_attention_00")); +} + +pub const Kernels = struct { + own: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Kernels { + var k: Kernels = undefined; + const own = try frags.source(device, gpa, sources.idol_decoder); + defer gpa.free(own); + k.own = try mtl.Library.fromSource(device, own, mtl.CompileOptions.mlx()); + errdefer k.own.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.own, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.own.deinit(); + } + + fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } + + /// h [rows, 1024] fp32 = in_proj(p) + the points' sinusoidal embedding, p [rows, 3] fp32. + pub fn embed(k: *const Kernels, e: mtl.ComputeEncoder, p: Ref, w: Ref, b: Ref, freqs: Ref, h: Ref, rows: usize) void { + buffers(e, k.get("dec_embed"), &.{ p, w, b, freqs, h }); + e.dispatchThreads(mtl.Size.of(dc.hidden, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// o [rows, width] fp16 = y + b by column. + pub fn toHalf(k: *const Kernels, e: mtl.ComputeEncoder, y: Ref, b: Ref, o: Ref, rows: usize, width: usize) void { + bind(e, k.get("dec_half"), &.{ y, b, o }, @as(u32, @intCast(width))); + e.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// kv [8192, 2048] fp16 = to_kv of the fp16 latent rows. + pub fn kv(k: *const Kernels, e: mtl.ComputeEncoder, ctx: Ref, w: Ref, b: Ref, out: Ref) void { + buffers(e, k.get("dec_kv"), &.{ ctx, w, b, out }); + e.dispatchThreads(mtl.Size.of(2 * dc.hidden, dc.context_rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// dst [16, rows, 64] from columns `offset` on of `src` rows `stride` apart, RMS-normed by `gamma` when given. + pub fn heads(k: *const Kernels, e: mtl.ComputeEncoder, src: Ref, gamma: ?Ref, dst: Ref, rows: usize, stride: usize, offset: usize) void { + const P = extern struct { rows: u32, stride: u32, offset: u32, norm: u32 }; + const params: P = .{ .rows = @intCast(rows), .stride = @intCast(stride), .offset = @intCast(offset), .norm = @intFromBool(gamma != null) }; + bind(e, k.get("dec_heads"), &.{ src, gamma orelse src, dst }, params); + e.dispatchGroups(mtl.Size.of(rows, 2, 1), mtl.Size.of(256, 1, 1)); + } + + /// o rows [lq, 1024] = softmax(q k^T / 8) v for q [16, lq, 64], k, v [16, lk, 64]; P V in two fp16 halves of P. + pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, lq: usize, lk: usize) void { + const form = @as(usize, @intFromBool(lq % 64 == 0)) * 2 + @intFromBool(lk % 32 == 0); + const nq = (lq + 63) / 64; + const nk = (lk + 31) / 32; + e.setPipeline(k.p[attention_first + form]); + for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + // heads, q rows, k rows, k blocks, last q block, last k block, rows in the last q block and in the last k block + var params: [8]i32 = .{ dc.heads, @intCast(lq), @intCast(lk), @intCast(nk), @intCast(nq - 1), @intCast(nk - 1), @intCast(lq - (nq - 1) * 64), @intCast(lk - (nk - 1) * 32) }; + e.setBytes(std.mem.asBytes(¶ms), 3); + var f: [4]f32 = .{ 0.125, 0, 0, 0 }; + e.setBytes(std.mem.asBytes(&f), 4); + e.setBuffer(o.buf, o.off, 5); + e.dispatchGroups(mtl.Size.of(nq, dc.heads, 1), mtl.Size.of(32, 4, 1)); + } + + /// o rows [lq, 1024] = softmax(q k^T / 8) v in fp32 for lq <= 8 query rows and up to 8,192 keys. + pub fn attentionFew(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, lq: usize, lk: usize) void { + std.debug.assert(lq <= few_rows and lk <= 32 * 256); + bind(e, k.get("dec_attention_few"), &.{ q, kb, vb, o }, [2]u32{ @intCast(lq), @intCast(lk) }); + e.dispatchGroups(mtl.Size.of(lq, dc.heads, 1), mtl.Size.of(256, 1, 1)); + } +}; + +/// One kernel's buffers in order, no parameters. +fn buffers(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); +} + +/// y [M, N] (rows N apart) = x [M, K] w^T in fp32 FMA by the denoiser's id_gemm32; any N % 8 = 0 (its last tile). +pub fn gemm32(k: *const ks.Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: usize, N: usize, K: usize) void { + std.debug.assert(K % 32 == 0 and N % 8 == 0); + const params = [4]u32{ @intCast(M), @intCast(N), @intCast(K), @intCast(N) }; + bind(e, k.get("id_gemm32"), &.{ x, w, y }, params); + e.dispatchGroups(mtl.Size.of((N + 63) / 64, (M + 63) / 64, 1), mtl.Size.of(128, 1, 1)); +} + +test "every decoder pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + if (!device.tensorUnits()) return error.SkipZigTest; + var k = try Kernels.load(std.testing.allocator, device); + k.deinit(); +} diff --git a/zig/src/families/idol/decoder_weights.zig b/zig/src/families/idol/decoder_weights.zig new file mode 100644 index 0000000000..800758dbb1 --- /dev/null +++ b/zig/src/families/idol/decoder_weights.zig @@ -0,0 +1,240 @@ +//! Idol's decoders' tensors (fp32 as stored): the blocks' linears to fp16 at load as idol-sdk does, the rest fp32. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const dc = @import("decoder_config.zig"); +const octree = @import("octree.zig"); +const Ref = @import("kernels.zig").Ref; + +const H = dc.hidden; +const F = dc.ffn; + +/// Cross attention to the latent: fp16 to_q, to_kv, to_out and biases, fp32 RMSNorm gammas [16, 64]. +pub const Cross = struct { q: Ref, q_b: Ref, kv: Ref, kv_b: Ref, out: Ref, out_b: Ref, gq: Ref, gk: Ref }; +/// Self attention over the anchors: fp16 to_qkv, to_out and biases, fp32 gammas. +pub const SelfAttention = struct { qkv: Ref, qkv_b: Ref, out: Ref, out_b: Ref, gq: Ref, gk: Ref }; +pub const Mlp = struct { up: Ref, up_b: Ref, down: Ref, down_b: Ref }; +pub const OctreeBlock = struct { cross: Cross, mlp: Mlp }; +/// A Gaussian decoder block; norm2 is its one affine LayerNorm32 (fp32). +pub const GaussianBlock = struct { self_attn: SelfAttention, cross: Cross, norm2: Ref, norm2_b: Ref, mlp: Mlp }; +/// The fp32 stem and head: in_proj [1024, 3], input_layer [1024, 1024], out_proj [N, 1024], and biases. +pub const Stem = struct { in_w: Ref, in_b: Ref, input: Ref, input_b: Ref, out: Ref, out_b: Ref }; + +pub const Weights = struct { + ck: ck.Checkpoint, + octree: [dc.octree_blocks]OctreeBlock, + gaussian: [dc.gs_blocks]GaussianBlock, + octree_stem: Stem, + gaussian_stem: Stem, + level: octree.LevelWeights, + base_offset_scale: f32, + /// points_offset_perturbation [32, 3]. + perturbation: []const f32, + rep: dc.Representation, + + pub fn deinit(w: *Weights) void { + w.ck.deinit(); + } +}; + +/// The blocks' linear weights stay in their files: load reads each once into its fp16 copy. +fn later(name: []const u8, rank: usize) bool { + if (rank != 2) return false; + const owned = std.mem.startsWith(u8, name, "oct.blocks.") or std.mem.startsWith(u8, name, "gs.blocks."); + return owned and std.mem.endsWith(u8, name, ".weight"); +} + +fn check(name: []const u8, dtype: ck.DType, got: []const usize, shape: []const usize) !void { + if (dtype != .f32 or !std.mem.eql(usize, got, shape)) { + std.log.err("idol: {s} is {t} {any}; the decoders read f32 {any}", .{ name, dtype, got, shape }); + return error.UnexpectedTensor; + } +} + +fn take(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) !Ref { + const t = try c.get(name); + try check(name, t.dtype, t.shape[0..t.rank], shape); + if (t.offset % 16 != 0) return error.UnalignedTensor; + return .{ .buf = t.buffer, .off = t.offset }; +} + +fn floats(c: *const ck.Checkpoint, name: []const u8, shape: []const usize) ![]const f32 { + const t = try c.get(name); + try check(name, t.dtype, t.shape[0..t.rank], shape); + return t.host(f32); +} + +/// fp16 copies bump-allocated 256-aligned in one owned buffer. +const Halves = struct { + c: *const ck.Checkpoint, + buf: mtl.Buffer, + at: usize = 0, + stage: []f32, + name: [96]u8 = undefined, + + fn bump(h: *Halves, n: usize) Ref { + const r: Ref = .{ .buf = h.buf, .off = h.at }; + h.at += std.mem.alignForward(usize, n * 2, ck.alignment); + return r; + } + + fn dst(h: *Halves, r: Ref, n: usize) []f16 { + return @as([*]f16, @ptrCast(@alignCast(h.buf.contents() + r.off)))[0..n]; + } + + /// The bias `prefix.leaf.bias` [n] as fp16. + fn bias(h: *Halves, prefix: []const u8, leaf: []const u8, n: usize) !Ref { + const name = try std.fmt.bufPrint(&h.name, "{s}.{s}.bias", .{ prefix, leaf }); + const src = try floats(h.c, name, &.{n}); + const r = h.bump(n); + for (h.dst(r, n), src) |*o, v| o.* = @floatCast(v); + return r; + } + + /// The deferred weight `prefix.leaf.weight` [rows, cols], read in staged pieces and rounded to fp16. + fn weight(h: *Halves, prefix: []const u8, leaf: []const u8, rows: usize, cols: usize) !Ref { + const name = try std.fmt.bufPrint(&h.name, "{s}.{s}.weight", .{ prefix, leaf }); + const d = h.c.deferred.get(name) orelse return error.MissingTensor; + try check(name, d.dtype, d.shape[0..d.rank], &.{ rows, cols }); + const n = rows * cols; + const r = h.bump(n); + const out = h.dst(r, n); + var done: usize = 0; + while (done < n) { + const part = @min(h.stage.len, n - done); + try h.c.read(name, done * 4, std.mem.sliceAsBytes(h.stage[0..part])); + for (out[done..][0..part], h.stage[0..part]) |*o, v| o.* = @floatCast(v); + done += part; + } + return r; + } +}; + +fn sized(comptime parts: []const usize) usize { + var n: usize = 0; + for (parts) |v| n += std.mem.alignForward(usize, v * 2, ck.alignment); + return n; +} + +const cross_halves = [_]usize{ H * H, H, 2 * H * dc.context_channels, 2 * H, H * H, H }; +const mlp_halves = [_]usize{ F * H, F, H * F, H }; +const self_halves = [_]usize{ 3 * H * H, 3 * H, H * H, H }; +const octree_bytes = sized(&(cross_halves ++ mlp_halves)); +const gaussian_bytes = sized(&(cross_halves ++ mlp_halves ++ self_halves)); +const fp16_bytes = dc.octree_blocks * octree_bytes + dc.gs_blocks * gaussian_bytes; + +fn gamma(c: *const ck.Checkpoint, buf: *[96]u8, prefix: []const u8, leaf: []const u8) !Ref { + return take(c, try std.fmt.bufPrint(buf, "{s}.{s}.gamma", .{ prefix, leaf }), &.{ dc.heads, H / dc.heads }); +} + +fn cross(h: *Halves, block: []const u8) !Cross { + var p: [96]u8 = undefined; + const pre = try std.fmt.bufPrint(&p, "{s}.cross_attn", .{block}); + var g: [96]u8 = undefined; + return .{ + .q = try h.weight(pre, "to_q", H, H), + .q_b = try h.bias(pre, "to_q", H), + .kv = try h.weight(pre, "to_kv", 2 * H, dc.context_channels), + .kv_b = try h.bias(pre, "to_kv", 2 * H), + .out = try h.weight(pre, "to_out", H, H), + .out_b = try h.bias(pre, "to_out", H), + .gq = try gamma(h.c, &g, pre, "q_rms_norm"), + .gk = try gamma(h.c, &g, pre, "k_rms_norm"), + }; +} + +fn mlp(h: *Halves, block: []const u8) !Mlp { + var p: [96]u8 = undefined; + const pre = try std.fmt.bufPrint(&p, "{s}.mlp.mlp", .{block}); + return .{ .up = try h.weight(pre, "0", F, H), .up_b = try h.bias(pre, "0", F), .down = try h.weight(pre, "2", H, F), .down_b = try h.bias(pre, "2", H) }; +} + +fn stem(c: *const ck.Checkpoint, prefix: []const u8, out: usize) !Stem { + var n: [96]u8 = undefined; + return .{ + .in_w = try take(c, try std.fmt.bufPrint(&n, "{s}.in_proj.weight", .{prefix}), &.{ H, 3 }), + .in_b = try take(c, try std.fmt.bufPrint(&n, "{s}.in_proj.bias", .{prefix}), &.{H}), + .input = try take(c, try std.fmt.bufPrint(&n, "{s}.input_layer.weight", .{prefix}), &.{ H, H }), + .input_b = try take(c, try std.fmt.bufPrint(&n, "{s}.input_layer.bias", .{prefix}), &.{H}), + .out = try take(c, try std.fmt.bufPrint(&n, "{s}.out_proj.weight", .{prefix}), &.{ out, H }), + .out_b = try take(c, try std.fmt.bufPrint(&n, "{s}.out_proj.bias", .{prefix}), &.{out}), + }; +} + +/// The snapshot folder `dir`: config.json's decoders checked, then both decoders' safetensors, each tensor once. +pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "config.json" }), arena, .limited(1 << 20)); + var w: Weights = undefined; + w.rep = try dc.check(gpa, config); + w.ck = ck.Checkpoint.init(gpa); + errdefer w.ck.deinit(); + for ([_][2][]const u8{ .{ "decoder", "oct." }, .{ "decoder_gs", "gs." } }) |f| { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/{s}.safetensors", .{ dir, f[0] }, 0); + try w.ck.addFileSplit(device, path, f[1], null, later); + } + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const fp16 = try device.buffer(fp16_bytes, opts); + w.ck.own(fp16, fp16_bytes) catch |err| { + fp16.deinit(); + return err; + }; + var h: Halves = .{ .c = &w.ck, .buf = fp16, .stage = try arena.alloc(f32, 4 << 20) }; + var name: [96]u8 = undefined; + for (&w.octree, 0..) |*b, i| { + const pre = try std.fmt.bufPrint(&name, "oct.blocks.{d}", .{i}); + b.cross = try cross(&h, pre); + b.mlp = try mlp(&h, pre); + } + var g: [96]u8 = undefined; + for (&w.gaussian, 0..) |*b, i| { + const pre = try std.fmt.bufPrint(&name, "gs.blocks.{d}", .{i}); + var sp: [96]u8 = undefined; + const sa = try std.fmt.bufPrint(&sp, "{s}.self_attn", .{pre}); + b.self_attn = .{ + .qkv = try h.weight(sa, "to_qkv", 3 * H, H), + .qkv_b = try h.bias(sa, "to_qkv", 3 * H), + .out = try h.weight(sa, "to_out", H, H), + .out_b = try h.bias(sa, "to_out", H), + .gq = try gamma(&w.ck, &g, sa, "q_rms_norm"), + .gk = try gamma(&w.ck, &g, sa, "k_rms_norm"), + }; + b.cross = try cross(&h, pre); + b.norm2 = try take(&w.ck, try std.fmt.bufPrint(&g, "{s}.norm2.weight", .{pre}), &.{H}); + b.norm2_b = try take(&w.ck, try std.fmt.bufPrint(&g, "{s}.norm2.bias", .{pre}), &.{H}); + b.mlp = try mlp(&h, pre); + } + std.debug.assert(h.at == fp16_bytes); + w.ck.closeFiles(); + w.octree_stem = try stem(&w.ck, "oct", dc.children); + w.gaussian_stem = try stem(&w.ck, "gs", dc.raw_width); + w.level = .{ + .l1 = try floats(&w.ck, "oct.l_embedder.mlp.0.weight", &.{ H, dc.level_freqs }), + .l1_b = try floats(&w.ck, "oct.l_embedder.mlp.0.bias", &.{H}), + .l2 = try floats(&w.ck, "oct.l_embedder.mlp.2.weight", &.{ H, H }), + .l2_b = try floats(&w.ck, "oct.l_embedder.mlp.2.bias", &.{H}), + .ada = try floats(&w.ck, "oct.adaLN_modulation.1.weight", &.{ 6 * H, H }), + .ada_b = try floats(&w.ck, "oct.adaLN_modulation.1.bias", &.{6 * H}), + }; + w.base_offset_scale = (try floats(&w.ck, "gs.base_offset_scale", &.{}))[0]; + w.perturbation = try floats(&w.ck, "gs.points_offset_perturbation", &.{ dc.per_anchor, 3 }); + // the octree's 12 a block and 12 more; the Gaussian decoder's 20 a block and 8 more + const read = 12 * dc.octree_blocks + 12 + 20 * dc.gs_blocks + 8; + const total = w.ck.tensors.count() + w.ck.deferred.count(); + if (total != read) { + std.log.err("idol: the decoders' files have {d} tensors; this family reads {d}", .{ total, read }); + return error.UnexpectedTensors; + } + return w; +} + +test "the blocks' linear weights, and only they, stay in their files for their fp16 copies" { + try std.testing.expect(later("oct.blocks.3.cross_attn.to_kv.weight", 2)); + try std.testing.expect(later("gs.blocks.15.self_attn.to_qkv.weight", 2)); + try std.testing.expect(!later("gs.blocks.0.cross_attn.q_rms_norm.gamma", 2)); + try std.testing.expect(!later("gs.blocks.0.norm2.weight", 1)); + try std.testing.expect(!later("gs.input_layer.weight", 2)); + try std.testing.expect(!later("oct.adaLN_modulation.1.weight", 2)); +} diff --git a/zig/src/families/idol/gaussian.zig b/zig/src/families/idol/gaussian.zig new file mode 100644 index 0000000000..0b2acdb953 --- /dev/null +++ b/zig/src/families/idol/gaussian.zig @@ -0,0 +1,206 @@ +//! Idol's Gaussians from the Gaussian decoder's rows (idol-sdk's gaussian.py), the 3DGS PLY, and the camera. +const std = @import("std"); +const dc = @import("decoder_config.zig"); + +const G = dc.per_anchor; + +/// Activated Gaussians, z-up inside [-0.5, 0.5]^3: xyz [n, 3], colour [n, 3], scale [n, 3], quaternion wxyz [n, 4]. +pub const Gaussians = struct { + xyz: []f32, + features_dc: []f32, + scaling: []f32, + rotation: []f32, + opacity: []f32, + + pub fn alloc(gpa: std.mem.Allocator, n: usize) !Gaussians { + const all = try gpa.alloc(f32, n * 14); + return .{ .xyz = all[0 .. 3 * n], .features_dc = all[3 * n .. 6 * n], .scaling = all[6 * n .. 9 * n], .rotation = all[9 * n .. 13 * n], .opacity = all[13 * n .. 14 * n] }; + } + + pub fn free(g: Gaussians, gpa: std.mem.Allocator) void { + gpa.free(g.xyz.ptr[0 .. g.opacity.len * 14]); + } + + pub fn count(g: Gaussians) usize { + return g.opacity.len; + } +}; + +/// torch's F.softplus (beta 1, threshold 20). +pub fn softplus(x: f32) f32 { + return if (x > 20) x else std.math.log1p(@exp(x)); +} + +fn sigmoid(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// inverse_softplus(torch.tensor(scaling_bias)) and inverse_sigmoid(torch.tensor(opacity_bias)), in fp32. +pub fn offsets(rep: dc.Representation) [2]f32 { + const s: f32 = @floatCast(rep.scaling_bias); + const o: f32 = @floatCast(rep.opacity_bias); + return .{ s + @log(-std.math.expm1(-s)), @log(o / (1 - o)) }; +} + +/// The activation's fixed inputs: representation_config, base_offset_scale and points_offset_perturbation [32, 3]. +pub const Params = struct { rep: dc.Representation, base_offset_scale: f32, perturbation: []const f32 }; + +/// to_representation then Gaussian(): raw rows [anchors, 480] at anchors [anchors, 3] into `out` (anchors x 32). +pub fn activate(p: Params, raw: []const f32, anchors: []const f32, out: Gaussians) void { + const n = anchors.len / 3; + std.debug.assert(raw.len == n * dc.raw_width and out.count() == n * G and p.perturbation.len == G * 3); + const off = offsets(p.rep); + const r = p.rep; + const lr = [5]f32{ @floatCast(r.lr_xyz), @floatCast(r.lr_features_dc), @floatCast(r.lr_scaling), @floatCast(r.lr_rotation), @floatCast(r.lr_opacity) }; + const size: f32 = @floatCast(r.perturbe_size); + const filter2: f32 = @floatCast(r.filter * r.filter); + for (0..n) |a| { + const line = raw[a * dc.raw_width ..][0..dc.raw_width]; + for (0..G) |g| { + const i = a * G + g; + const scale = softplus(line[448 + g] + p.base_offset_scale); + for (0..3) |ax| { + const o = line[g * 3 + ax] * lr[0] + p.perturbation[g * 3 + ax]; + const offset = ((std.math.tanh(o) * 0.5) * size) * scale; + out.xyz[i * 3 + ax] = (offset + anchors[a * 3 + ax]) * 1.0 + -0.5; + out.features_dc[i * 3 + ax] = line[96 + g * 3 + ax] * lr[1]; + const s = softplus(line[192 + g * 3 + ax] * lr[2] + off[0]); + out.scaling[i * 3 + ax] = @sqrt(s * s + filter2); + } + var q: [4]f32 = undefined; + var sq: f32 = 0; + for (&q, 0..) |*v, c| { + v.* = line[288 + g * 4 + c] * lr[3] + @as(f32, if (c == 0) 1 else 0); + sq += v.* * v.*; + } + const norm = @max(@sqrt(sq), 1e-12); + for (q, 0..) |v, c| out.rotation[i * 4 + c] = v / norm; + out.opacity[i] = sigmoid(line[416 + g] * lr[4] + off[1]); + } + } +} + +pub const properties = [_][]const u8{ "x", "y", "z", "nx", "ny", "nz", "f_dc_0", "f_dc_1", "f_dc_2", "opacity", "scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3" }; +pub const row_bytes = properties.len * 4; + +/// The header plyfile writes for `n` vertices of the 17 float properties. +pub fn header(buf: []u8, n: usize) ![]const u8 { + var w: std.Io.Writer = .fixed(buf); + try w.print("ply\nformat binary_little_endian 1.0\nelement vertex {d}\n", .{n}); + for (properties) |name| try w.print("property float {s}\n", .{name}); + try w.writeAll("end_header\n"); + return w.buffered(); +} + +/// One PLY row as save_ply writes it: (x, -z, y), clamped logit opacity and log scale, the quaternion turned in f64. +pub fn row(g: Gaussians, i: usize, out: *[properties.len]f32) void { + const eps: f32 = 1e-6; + const hi: f32 = @floatCast(@as(f64, 1) - 1e-6); + const xyz = g.xyz[i * 3 ..][0..3]; + out[0] = xyz[0]; + out[1] = -xyz[2]; + out[2] = xyz[1]; + out[3] = 0; + out[4] = 0; + out[5] = 0; + @memcpy(out[6..9], g.features_dc[i * 3 ..][0..3]); + const o = std.math.clamp(g.opacity[i], eps, hi); + out[9] = @log(o / (1 - o)); + for (0..3) |a| out[10 + a] = @log(@max(g.scaling[i * 3 + a], eps)); + const aw: f64 = @sqrt(@as(f64, 0.5)); + const ax = aw; + const ay: f64 = 0; + const az: f64 = 0; + const b = g.rotation[i * 4 ..][0..4]; + const bw: f64 = b[0]; + const bx: f64 = b[1]; + const by: f64 = b[2]; + const bz: f64 = b[3]; + out[13] = @floatCast(aw * bw - ax * bx - ay * by - az * bz); + out[14] = @floatCast(aw * bx + ax * bw + ay * bz - az * by); + out[15] = @floatCast(aw * by - ax * bz + ay * bw + az * bx); + out[16] = @floatCast(aw * bz + ax * by - ay * bx + az * bw); +} + +/// The whole PLY file in memory: header then every row, little-endian float32. +pub fn ply(gpa: std.mem.Allocator, g: Gaussians) ![]u8 { + var head: [1024]u8 = undefined; + const h = try header(&head, g.count()); + const bytes = try gpa.alloc(u8, h.len + g.count() * row_bytes); + @memcpy(bytes[0..h.len], h); + for (0..g.count()) |i| { + var r: [properties.len]f32 = undefined; + row(g, i, &r); + for (r, 0..) |v, c| std.mem.writeInt(u32, bytes[h.len + i * row_bytes + c * 4 ..][0..4], @bitCast(v), .little); + } + return bytes; +} + +/// decode_camera: yaw and pitch (radians), radius, vertical field of view (degrees) from the camera latent [5]. +pub const Camera = struct { yaw: f64, pitch: f64, radius: f64, fov: f64 }; + +pub fn camera(c: [5]f32) Camera { + const norm = @sqrt(c[0] * c[0] + c[1] * c[1] + c[2] * c[2]); + var d = [3]f64{ 0, 0, 1 }; + if (norm > 1e-6) for (&d, 0..) |*v, i| { + v.* = c[i] / norm; + }; + const inv: f64 = c[3]; + const hw: f64 = c[4]; + const deg = 180.0 / std.math.pi; + const radius: f64 = if (inv < 1e-3) 1000.0 else 1.0 / inv; + const fov = if (inv < 1e-3) deg * 2 * std.math.atan(hw / radius) else deg * 2 * std.math.atan(hw * inv); + return .{ .yaw = std.math.atan2(d[0], d[1]), .pitch = std.math.asin(std.math.clamp(d[2], -1.0, 1.0)), .radius = radius, .fov = fov }; +} + +test "the activation offsets are torch's fp32 values for the released biases" { + const o = offsets(dc.released); + try std.testing.expectEqual(@as(f32, -5.519460201263428), o[0]); + try std.testing.expectEqual(@as(f32, -2.1972246170043945), o[1]); +} + +test "zero rows give unit quaternions, the bias opacity and the scale floor" { + const a = std.testing.allocator; + var raw: [dc.raw_width]f32 = @splat(0); + const pert: [G * 3]f32 = @splat(0); + const g = try Gaussians.alloc(a, G); + defer g.free(a); + activate(.{ .rep = dc.released, .base_offset_scale = 0, .perturbation = &pert }, &raw, &.{ 0.5, 0.5, 0.5 }, g); + try std.testing.expectApproxEqAbs(@as(f32, 0.1), g.opacity[5], 1e-7); + try std.testing.expectEqual(@as(f32, 1), g.rotation[5 * 4]); + try std.testing.expectEqual(@as(f32, 0), g.xyz[7]); + try std.testing.expect(@abs(g.scaling[0] - @sqrt(0.004 * 0.004 + 0.0009 * 0.0009)) < 1e-6); + raw[288] = std.math.nan(f32); + activate(.{ .rep = dc.released, .base_offset_scale = 0, .perturbation = &pert }, &raw, &.{ 0.5, 0.5, 0.5 }, g); + try std.testing.expect(std.math.isNan(g.rotation[0])); +} + +test "a PLY row: y-up position, the quaternion turned about x, clamped logit and log" { + var xyz = [3]f32{ 0.1, 0.2, 0.3 }; + var f = [3]f32{ 1, 2, 3 }; + var s = [3]f32{ 0, 1, 2 }; + var q = [4]f32{ 1, 0, 0, 0 }; + var o = [1]f32{1}; + const g: Gaussians = .{ .xyz = &xyz, .features_dc = &f, .scaling = &s, .rotation = &q, .opacity = &o }; + var r: [properties.len]f32 = undefined; + row(g, 0, &r); + try std.testing.expectEqualSlices(f32, &.{ 0.1, -0.3, 0.2, 0, 0, 0, 1, 2, 3 }, r[0..9]); + try std.testing.expectApproxEqAbs(@as(f32, 13.802318), r[9], 1e-5); + try std.testing.expectApproxEqAbs(@log(@as(f32, 1e-6)), r[10], 1e-6); + try std.testing.expectEqual(@as(f32, @floatCast(@sqrt(@as(f64, 0.5)))), r[13]); + try std.testing.expectEqual(r[13], r[14]); + var head: [1024]u8 = undefined; + const h = try header(&head, 262144); + try std.testing.expect(std.mem.startsWith(u8, h, "ply\nformat binary_little_endian 1.0\nelement vertex 262144\nproperty float x\n")); + try std.testing.expect(std.mem.endsWith(u8, h, "property float rot_3\nend_header\n")); +} + +test "the camera: orthographic below an inverse distance of 1e-3, else perspective" { + const p = camera(.{ 0, 1, 0, 0.5, 0.25 }); + try std.testing.expectApproxEqAbs(@as(f64, 0), p.yaw, 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 2), p.radius, 1e-12); + try std.testing.expectApproxEqAbs(2 * std.math.atan(@as(f64, 0.125)) * 180 / std.math.pi, p.fov, 1e-9); + const o = camera(.{ 0, 0, 0, 0, 1 }); + try std.testing.expectEqual(@as(f64, 1000), o.radius); + try std.testing.expectApproxEqAbs(std.math.pi / 2.0, o.pitch, 1e-12); +} diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig index a2aff60944..5aee4da8cb 100644 --- a/zig/src/families/idol/idol.zig +++ b/zig/src/families/idol/idol.zig @@ -1,10 +1,19 @@ -//! Idol image to 3D: its denoiser on Metal (a denoiser beside the lane engine, not a token family). +//! Idol image to 3D: its denoiser and decoders on Metal (beside the lane engine, not a token family). pub const config = @import("config.zig"); pub const host = @import("host.zig"); pub const weights = @import("weights.zig"); pub const kernels = @import("kernels.zig"); pub const model = @import("model.zig"); pub const reference = @import("reference.zig"); +pub const decoder_config = @import("decoder_config.zig"); +pub const decoder_kernels = @import("decoder_kernels.zig"); +pub const decoder_weights = @import("decoder_weights.zig"); +pub const decoder = @import("decoder.zig"); +pub const octree = @import("octree.zig"); +pub const gaussian = @import("gaussian.zig"); +pub const decode = @import("decode.zig"); +pub const decode_reference = @import("decode_reference.zig"); +pub const Decoder = decoder.Decoder; pub const Model = model.Model; pub const Condition = model.Condition; pub const Timing = model.Timing; @@ -16,4 +25,12 @@ test { _ = kernels; _ = model; _ = reference; + _ = decoder_config; + _ = decoder_kernels; + _ = decoder_weights; + _ = decoder; + _ = octree; + _ = gaussian; + _ = decode; + _ = decode_reference; } diff --git a/zig/src/families/idol/octree.zig b/zig/src/families/idol/octree.zig new file mode 100644 index 0000000000..f473137a89 --- /dev/null +++ b/zig/src/families/idol/octree.zig @@ -0,0 +1,256 @@ +//! The octree decoder's host part: the level modulation, systematic resampling, compaction and the leaves' jitter. +const std = @import("std"); +const dc = @import("decoder_config.zig"); +const project = @import("host.zig").project; + +const H = dc.hidden; +const P = dc.children; + +fn half(x: f32) f32 { + return @as(f16, @floatCast(x)); +} + +fn silu(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} + +/// torch.linspace(start, end, n) in fp32 as its CPU kernel computes it: fma from the nearer end. +pub fn linspace(out: []f32, start: f32, end: f32) void { + const n = out.len; + const step = (end - start) / @as(f32, @floatFromInt(n - 1)); + for (out, 0..) |*o, i| { + o.* = if (i < n / 2) @mulAdd(f32, step, @floatFromInt(i), start) else @mulAdd(f32, -step, @floatFromInt(n - 1 - i), end); + } +} + +/// Where torch's CPU pow (built at init, not saved: the model's constants) rounds 2^linspace the other way, in ulps. +const torch_pow = [_]struct { usize, i32 }{ .{ 43, -1 }, .{ 72, -1 }, .{ 95, -1 }, .{ 113, 1 }, .{ 138, -1 }, .{ 159, 1 } }; + +/// PcdAbsolutePositionEmbedderv2's frequencies torch.pow(2.0, linspace(0, 10, 170)): correctly rounded but six. +pub fn pointFreqs(out: *[dc.pe_freqs]f32) void { + linspace(out, 0, 10); + for (out) |*f| f.* = @floatCast(@exp2(@as(f64, f.*))); + for (torch_pow) |e| out[e[0]] = @bitCast(@as(i32, @bitCast(out[e[0]])) + e[1]); +} + +/// The level modulation's parts in block order: shift, 1 + scale, gate for attention, then for the MLP; fp16. +pub const Mods = [6][H]f16; + +/// The fp32 tensors the level embedding reads: l_embedder's MLP, then adaLN_modulation's linear. +pub const LevelWeights = struct { l1: []const f32, l1_b: []const f32, l2: []const f32, l2_b: []const f32, ada: []const f32, ada_b: []const f32 }; + +/// adaLN_modulation(l_embedder(2^level)) as fp16, the two scales as fp16(1 + scale), as the blocks read them. +pub fn levelMods(w: LevelWeights, level: u5, out: *Mods) void { + const hf = dc.level_freqs / 2; + const t: f32 = @floatFromInt(@as(u32, 1) << level); + const lg: f32 = @floatCast(-@log(@as(f64, dc.level_period))); + var freq: [dc.level_freqs]f32 = undefined; + for (0..hf) |i| { + const e: f32 = lg * @as(f32, @floatFromInt(i)) / @as(f32, hf); + const a = ((t * @as(f32, @floatCast(@exp(@as(f64, e))))) * 2) * std.math.pi; + freq[i] = @cos(a); + freq[hf + i] = @sin(a); + } + var mid: [H]f32 = undefined; + project(w.l1, w.l1_b, &freq, &mid); + for (&mid) |*v| v.* = silu(v.*); + var emb: [H]f32 = undefined; + project(w.l2, w.l2_b, &mid, &emb); + for (&emb) |*v| v.* = silu(v.*); + var mod: [6 * H]f32 = undefined; + project(w.ada, w.ada_b, &emb, &mod); + for (mod, 0..) |v, i| { + const part = i / H; + out[part][i % H] = @floatCast(if (part == 1 or part == 4) half(1.0 + half(v)) else v); + } +} + +/// The parents' cell centres in [0, 1]: (coords + 0.5) / 2^(level - 1), fp32 [rows, 3]. +pub fn centres(coords: []const [3]u32, level: u5, out: []f32) void { + const res: f32 = @floatFromInt(@as(u32, 1) << (level - 1)); + for (coords, 0..) |c, r| for (0..3) |a| { + out[r * 3 + a] = (@as(f32, @floatFromInt(c[a])) + 0.5) / res; + }; +} + +/// softmax over each row of 8 logits, in place. +pub fn softmax(rows: []f32) void { + var r: usize = 0; + while (r < rows.len) : (r += P) { + const x = rows[r..][0..P]; + var top = x[0]; + for (x[1..]) |v| top = @max(top, v); + var sum: f32 = 0; + for (x) |*v| { + v.* = @exp(v.* - top); + sum += v.*; + } + for (x) |*v| v.* /= sum; + } +} + +/// Where systematic_sample's uniform draws come from: replayed from a reference run, or a generator. +pub const Draws = union(enum) { + native: std.Random.DefaultPrng, + replay: *const Replay, + + /// The draw for the row of parent `coord` at `level`, `row` in the level's order. + pub fn draw(d: *Draws, level: u5, coord: [3]u32, row: usize) f32 { + return switch (d.*) { + .native => |*g| g.random().float(f32), + .replay => |r| r.draw(level, coord, row), + }; + } + + /// The leaves' jitter value `i` (point i / 3, axis i % 3). + pub fn jitter(d: *Draws, i: usize) f32 { + return switch (d.*) { + .native => |*g| g.random().float(f32), + .replay => |r| if (i < r.jitter.len) r.jitter[i] else std.math.nan(f32), + }; + } +}; + +/// A reference run's draws: each level's parents and their draws (NaN where torch drew none), and the jitter. +pub const Replay = struct { + parents: [dc.max_level][]const [3]u32, + draws: [dc.max_level][]const f32, + jitter: []const f32, + + /// The reference's draw for the parent at `coord`: by row when the rows agree, else by searching its parents. + pub fn draw(r: *const Replay, level: u5, coord: [3]u32, row: usize) f32 { + const ps = r.parents[level - 1]; + if (row < ps.len and std.mem.eql(u32, &ps[row], &coord)) return r.draws[level - 1][row]; + for (ps, 0..) |p, i| if (std.mem.eql(u32, &p, &coord)) return r.draws[level - 1][i]; + return std.math.nan(f32); + } +}; + +/// systematic_sample for one parent: `n` points among 8 children by probabilities `p` and offset draw `u`. +pub fn split(p: *const [P]f32, n: u32, u: f32, out: *[P]u32) void { + out.* = @splat(0); + if (n == 0) return; + var q: [P]f32 = undefined; + var sum: f32 = 0; + for (&q, p) |*o, v| { + o.* = @max(v, 0); + sum += o.*; + } + const by = @max(sum, 1); + for (&q) |*v| v.* = if (sum == 0) 1.0 / @as(f32, P) else v.* / by; + // cdf clamped to fp32's 1 - 1e-12, which is 1 + var cdf: [P]f32 = undefined; + var acc: f32 = 0; + for (&cdf, q) |*c, v| { + acc += v; + c.* = @min(acc, 1.0); + } + const nf: f32 = @floatFromInt(n); + const first = u / nf; + for (0..n) |j| { + const us = @min(first + @as(f32, @floatFromInt(j)) / nf, 1.0); + var idx: usize = 0; + while (idx < P and cdf[idx] < us) idx += 1; + out[@min(idx, P - 1)] += 1; + } +} + +/// One octree level: the parents' counts split by `probs` [rows, 8]; the children that got points, in order. +pub const Level = struct { + coords: []const [3]u32, + counts: []const u32, +}; + +/// The 8 children of every parent, split by probs and draws, kept where they got points (parent-major, i + 2j + 4k). +pub fn descend(gpa: std.mem.Allocator, parents: Level, probs: []const f32, level: u5, draws: *Draws, splits: ?[]u32) !Level { + std.debug.assert(probs.len == parents.coords.len * P); + var coords: std.ArrayList([3]u32) = .empty; + errdefer coords.deinit(gpa); + var counts: std.ArrayList(u32) = .empty; + errdefer counts.deinit(gpa); + for (parents.coords, parents.counts, 0..) |c, n, r| { + var got: [P]u32 = undefined; + const u = if (n == 0) 0 else draws.draw(level, c, r); + split(probs[r * P ..][0..P], n, u, &got); + if (splits) |s| @memcpy(s[r * P ..][0..P], &got); + for (got, 0..) |m, i| { + if (m == 0) continue; + const o = [3]u32{ @intCast(i & 1), @intCast((i >> 1) & 1), @intCast(i >> 2) }; + try coords.append(gpa, .{ c[0] * 2 + o[0], c[1] * 2 + o[1], c[2] * 2 + o[2] }); + try counts.append(gpa, m); + } + } + return .{ .coords = try coords.toOwnedSlice(gpa), .counts = try counts.toOwnedSlice(gpa) }; +} + +/// The anchors: each leaf repeated by its count, (coords + jitter) / 2^level, fp32 [points, 3]. +pub fn anchors(leaves: Level, level: u5, draws: *Draws, out: []f32) !void { + var at: usize = 0; + const res: f32 = @floatFromInt(@as(u32, 1) << level); + for (leaves.coords, leaves.counts) |c, n| for (0..n) |_| { + if (at + 3 > out.len) return error.AnchorCount; + for (0..3) |a| { + const j = draws.jitter(at + a); + if (std.math.isNan(j)) return error.MissingDraw; + out[at + a] = (@as(f32, @floatFromInt(c[a])) + j) / res; + } + at += 3; + }; + if (at != out.len) return error.AnchorCount; +} + +test "linspace takes each value from its nearer end with an fma, as torch's CPU kernel" { + var v: [7]f32 = undefined; + linspace(&v, 0, 1); + try std.testing.expectEqual(@as(f32, 0), v[0]); + try std.testing.expectEqual(@as(f32, 1), v[6]); + var f: [dc.pe_freqs]f32 = undefined; + pointFreqs(&f); + try std.testing.expectEqual(@as(f32, 1), f[0]); + try std.testing.expectEqual(@as(f32, 1024), f[dc.pe_freqs - 1]); + try std.testing.expectEqual(@as(f32, 0.05917159840464592), blk: { + var l: [dc.pe_freqs]f32 = undefined; + linspace(&l, 0, 10); + break :blk l[1]; + }); +} + +test "systematic resampling keeps the count and follows the cdf" { + var out: [P]u32 = undefined; + const p = [P]f32{ 0.5, 0, 0.25, 0, 0, 0, 0, 0.25 }; + split(&p, 8, 0.5, &out); + try std.testing.expectEqualSlices(u32, &.{ 4, 0, 2, 0, 0, 0, 0, 2 }, &out); + // one point: the draw alone picks the child + split(&p, 1, 0.6, &out); + try std.testing.expectEqualSlices(u32, &.{ 0, 0, 1, 0, 0, 0, 0, 0 }, &out); + // a zero row falls back to uniform, and sums under 1 are not rescaled (torch divides by max(sum, 1)) + split(&@as([P]f32, @splat(0)), 8, 0.5, &out); + try std.testing.expectEqualSlices(u32, &@as([P]u32, @splat(1)), &out); + // searchsorted left: a draw landing on a cdf step goes to the lower child, so u = 0 doubles the first + split(&@as([P]f32, @splat(0)), 8, 0.0, &out); + try std.testing.expectEqualSlices(u32, &.{ 2, 1, 1, 1, 1, 1, 1, 0 }, &out); + const low = [P]f32{ 0.25, 0.25, 0, 0, 0, 0, 0, 0 }; + split(&low, 4, 0.5, &out); + try std.testing.expectEqualSlices(u32, &.{ 1, 1, 0, 0, 0, 0, 0, 2 }, &out); +} + +test "descend packs children parent-major in i + 2j + 4k order; anchors land inside their leaves" { + const a = std.testing.allocator; + const parents: Level = .{ .coords = &.{ .{ 0, 0, 0 }, .{ 1, 0, 1 } }, .counts = &.{ 2, 1 } }; + var probs: [2 * P]f32 = @splat(0); + probs[3] = 1; // child (1, 1, 0) of the first parent + probs[P + 4] = 1; // child (0, 0, 1) of the second + var d: Draws = .{ .native = std.Random.DefaultPrng.init(1) }; + const next = try descend(a, parents, &probs, 2, &d, null); + defer a.free(next.coords); + defer a.free(next.counts); + try std.testing.expectEqualSlices(u32, &.{ 2, 1 }, next.counts); + try std.testing.expectEqual([3]u32{ 1, 1, 0 }, next.coords[0]); + try std.testing.expectEqual([3]u32{ 2, 0, 3 }, next.coords[1]); + var pts: [9]f32 = undefined; + try anchors(next, 2, &d, &pts); + for (0..3) |p| for (0..3) |ax| { + const c: f32 = @floatFromInt(next.coords[if (p < 2) 0 else 1][ax]); + try std.testing.expect(pts[p * 3 + ax] >= c / 4 and pts[p * 3 + ax] < (c + 1) / 4); + }; +} diff --git a/zig/tests/idol_decode.zig b/zig/tests/idol_decode.zig new file mode 100644 index 0000000000..15a68b5f6d --- /dev/null +++ b/zig/tests/idol_decode.zig @@ -0,0 +1,395 @@ +//! tf-idol-decode: Idol's octree and Gaussian decoders on a case from tools/zig/idol_case.py decode, gated and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const id = tf.idol; +const dc = id.decoder_config; +const oc = id.octree; +const gs = id.gaussian; +const dr = id.decode_reference; +const ref = id.reference; + +const usage = + \\tf-idol-decode SNAPSHOT_DIR CASE_DIR + \\tf-idol-decode SNAPSHOT_DIR CASE_DIR time GAUSSIANS [OUT.ply] + \\SNAPSHOT_DIR is a speridlabs/idol snapshot (config.json, decoder.safetensors, decoder_gs.safetensors). CASE_DIR + \\holds decode.safetensors, decode.json and reference.ply from `idol_case.py decode` (--ruler: float32 rulers). + \\Gates, each on the SDK's own inputs: every level's counts from the SDK's probabilities and draws identical; its + \\probabilities against float32 within twice the worst of the SDK's own distance, the fp16 floor's (float32, its + \\last rows rounded once) and the SDK's reruns with 1% of a first block's inputs an ulp off; levels 1 to 3 block + \\by block, rows and updates, within twice the SDK's; the Gaussian decoder's rows as the levels; fields; the PLY's + \\header, rows and columns; the camera; byte-identical repeats. A free run on the SDK's draws reports where its + \\topology leaves the SDK's. Writes CASE_DIR/tf-idol-decode.json, .ply and camera.json; exit 1 on a failed gate. + \\With `time`, decodes case.safetensors's sample at GAUSSIANS with the native generator, twice, timed. +; + +fn tensor(f: *const tf.safetensors.File, name: []const u8) !tf.safetensors.Tensor { + return f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; +} + +fn floats(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]f32 { + const t = try tensor(f, name); + if (t.dtype != .f32) return error.BadCaseTensor; + const out = try a.alloc(f32, t.numel()); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn ints(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]u32 { + const t = try tensor(f, name); + if (t.dtype != .i32) return error.BadCaseTensor; + const out = try a.alloc(u32, t.numel()); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn coords(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]const [3]u32 { + const v = try ints(a, f, name); + return @as([*]const [3]u32, @ptrCast(v.ptr))[0 .. v.len / 3]; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn verdict(ok: bool) []const u8 { + return if (ok) "pass" else "FAIL"; +} + +/// The SDK's run of the case: levels, draws and jitter. +const Case = struct { + f: *const tf.safetensors.File, + parents: [dc.max_level][]const [3]u32, + draws: [dc.max_level][]f32, + replay: oc.Replay, +}; + +/// Each level teacher-forced on the SDK's parents: probabilities, the level modulation, and counts. +fn levels(a: std.mem.Allocator, d: *id.Decoder, c: *const Case, out: *std.ArrayList(u8)) !bool { + var pass = true; + try print(out, a, ", \"levels\": [", .{}); + for (1..dc.max_level + 1) |l| { + var name: [32]u8 = undefined; + const parents = c.parents[l - 1]; + const rows = parents.len; + const mods = try floats(a, c.f, try std.fmt.bufPrint(&name, "mods.{d}", .{l})); + // the case holds adaLN's fp16 output; the blocks read the scales (parts 1 and 4) as fp16(1 + scale) + var mod_diff: usize = 0; + for (d.levelMods(l), mods, 0..) |x, y, i| { + const part = i / dc.hidden; + const want: f16 = if (part == 1 or part == 4) @floatCast(1 + y) else @floatCast(y); + mod_diff += @intFromBool(x != want); + } + const centres = try a.alloc(f32, rows * 3); + oc.centres(parents, @intCast(l), centres); + const ours = try a.alloc(f32, rows * 8); + _ = try d.octreeLogits(centres, l, ours); + oc.softmax(ours); + const sdk = try floats(a, c.f, try std.fmt.bufPrint(&name, "probs.{d}", .{l})); + const m = ref.match(ours, sdk, 8); + const w = dr.worst(ours, sdk, 8); + std.debug.print("level {d}: {d} parents; probabilities against the SDK: max |diff| {e:.3}, cosine {d:.9}, worst row {d:.5} (row {d})\n", .{ l, rows, m.max_abs, m.cosine, w.err, w.row }); + try print(out, a, "{s}{{\"level\": {d}, \"parents\": {d}, \"sdk\": {{\"max_abs\": {e:.4}, \"cosine\": {d:.10}, \"worst_row\": {d:.6}, \"worst_at\": {d}}}", .{ if (l == 1) "" else ", ", l, rows, m.max_abs, m.cosine, w.err, w.row }); + const ruler = try std.fmt.bufPrint(&name, "probs32.{d}", .{l}); + if (c.f.get(ruler) != null) { + const exact = try floats(a, c.f, ruler); + const floor = try floats(a, c.f, try std.fmt.bufPrint(&name, "probs32h.{d}", .{l})); + pass = (try rulerGate(a, out, ours, sdk, exact, floor, try nudges(a, c.f, "probs16n", try std.fmt.allocPrint(a, ".{d}", .{l})), 8)) and pass; + } + const counts_in = try ints(a, c.f, try std.fmt.bufPrint(&name, "counts_in.{d}", .{l})); + const want = try ints(a, c.f, try std.fmt.bufPrint(&name, "counts.{d}", .{l})); + const level: oc.Level = .{ .coords = parents, .counts = counts_in }; + const theirs = try a.alloc(u32, rows * 8); + const mine = try a.alloc(u32, rows * 8); + var replay: oc.Draws = .{ .replay = &c.replay }; + _ = try oc.descend(a, level, sdk, @intCast(l), &replay, theirs); + _ = try oc.descend(a, level, ours, @intCast(l), &replay, mine); + const same = dr.counts(theirs, want, 8); + const drift = dr.counts(mine, want, 8); + pass = pass and same.rows == 0; + std.debug.print(" counts from the SDK's probabilities and draws: {d} rows differ: {s}; from ours: {d} rows differ", .{ same.rows, verdict(same.rows == 0), drift.rows }); + if (drift.first) |r| std.debug.print(" (first at parent {d}, cell {any}: ours {any}, the SDK's {any})", .{ r, parents[r], mine[r * 8 ..][0..8], want[r * 8 ..][0..8] }); + std.debug.print("; level modulation: {d} of 6144 fp16 values differ from the SDK's\n", .{mod_diff}); + try print(out, a, ", \"counts_sdk_probs_rows_differ\": {d}, \"counts_our_probs_rows_differ\": {d}, \"first_differing_parent\": {d}, \"mods_differ\": {d}}}", .{ same.rows, drift.rows, if (drift.first) |r| @as(i64, @intCast(r)) else -1, mod_diff }); + } + try print(out, a, "]", .{}); + return pass; +} + +/// The full decode on the SDK's draws: where its topology first leaves the SDK's, and whether the anchors match. +fn freeRun(a: std.mem.Allocator, d: *id.Decoder, c: *const Case, latent: []const f32, n: usize, out: *std.ArrayList(u8)) !id.decode.Result { + var trace: id.decode.Trace = .{}; + var draws: oc.Draws = .{ .replay = &c.replay }; + const r = try id.decode.decode(a, d, latent, n, &draws, &trace); + var first: ?usize = null; + try print(out, a, ", \"free_run\": {{\"parents_equal\": [", .{}); + for (1..dc.max_level + 1) |l| { + const ours = trace.parents[l - 1]; + const theirs = c.parents[l - 1]; + const equal = ours.len == theirs.len and std.mem.eql(u8, std.mem.sliceAsBytes(ours), std.mem.sliceAsBytes(theirs)); + if (!equal and first == null) first = l; + try print(out, a, "{s}{}", .{ if (l == 1) "" else ", ", equal }); + if (!equal) std.debug.print("free run: level {d} has {d} parents, the SDK's {d}\n", .{ l, ours.len, theirs.len }); + } + const sdk_anchors = try floats(a, c.f, "anchors"); + const equal = std.mem.eql(u8, std.mem.sliceAsBytes(r.anchors), std.mem.sliceAsBytes(sdk_anchors)); + std.debug.print("free run on the SDK's draws: topology {s}; anchors {s} the SDK's\n", .{ if (first) |l| try std.fmt.allocPrint(a, "first differs at level {d}", .{l}) else "identical at every level", if (equal) "bit-equal to" else "differ from" }); + try print(out, a, "], \"first_divergent_level\": {d}, \"anchors_equal\": {}", .{ if (first) |l| @as(i64, @intCast(l)) else -1, equal }); + try print(out, a, ", \"octree_s\": {d:.4}, \"gaussian_s\": {d:.4}, \"activation_s\": {d:.4}}}", .{ r.time.octree, r.time.gaussian, r.time.activation }); + return r; +} + +/// Ours against float32 within twice the worst of the SDK's own distance, the fp16 floor's and its nudged reruns'. +fn rulerGate(a: std.mem.Allocator, out: *std.ArrayList(u8), ours: []const f32, sdk: []const f32, exact: []const f32, floor: []const f32, nudged: []const []const f32, width: usize) !bool { + const m = ref.match(ours, exact, width); + const r = ref.match(sdk, exact, width); + const f = ref.match(floor, exact, width); + var all: [8]ref.Match = undefined; + all[0] = r; + all[1] = f; + var worst_nudge: ref.Match = .{ .max_abs = 0, .cosine = 1, .worst_row = 0 }; + for (nudged, 0..) |n, i| { + all[2 + i] = ref.match(n, exact, width); + worst_nudge = dr.reference(&.{ worst_nudge, all[2 + i] }); + } + const w = dr.worst(ours, exact, width); + const ok = ref.within(m, dr.reference(all[0 .. 2 + nudged.len])); + std.debug.print(" against float32: worst row {d:.5} (row {d}), 1 - cosine {e:.3}; the SDK's {d:.5}, {e:.3}; nudged, worst {d:.5}, {e:.3}; floor {d:.5}, {e:.3}: {s}\n", .{ m.worst_row, w.row, 1 - m.cosine, r.worst_row, 1 - r.cosine, worst_nudge.worst_row, 1 - worst_nudge.cosine, f.worst_row, 1 - f.cosine, verdict(ok) }); + try print(out, a, ", \"fp32\": {{\"worst_row\": {d:.6}, \"worst_at\": {d}, \"cosine\": {d:.10}, \"sdk_worst_row\": {d:.6}, \"sdk_cosine\": {d:.10}, \"nudged_worst_row\": {d:.6}, \"nudged_cosine\": {d:.10}, \"floor_worst_row\": {d:.6}, \"floor_cosine\": {d:.10}, \"pass\": {}}}", .{ m.worst_row, w.row, m.cosine, r.worst_row, r.cosine, worst_nudge.worst_row, worst_nudge.cosine, f.worst_row, f.cosine, ok }); + return ok; +} + +/// The SDK's nudged reruns the case holds: `prefix`, the seed, then `suffix`. +fn nudges(a: std.mem.Allocator, f: *const tf.safetensors.File, prefix: []const u8, suffix: []const u8) ![]const []const f32 { + var list: std.ArrayList([]const f32) = .empty; + for (1..4) |seed| { + const key = try std.fmt.allocPrint(a, "{s}{d}{s}", .{ prefix, seed, suffix }); + if (f.get(key) != null) try list.append(a, try floats(a, f, key)); + } + return list.items; +} + +fn fieldGate(a: std.mem.Allocator, out: *std.ArrayList(u8), name: []const u8, ours: []const f32, theirs: []const f32) !bool { + const f = dr.field(ours, theirs, 1e-5, 1e-6); + std.debug.print("activated {s}: max |diff| {e:.3} at {d}, {d} of {d} bit-equal, {d} past 1e-5 rel + 1e-6: {s}\n", .{ name, f.max_abs, f.at, f.exact, ours.len, f.bad, verdict(f.bad == 0) }); + try print(out, a, ", \"{s}\": {{\"max_abs\": {e:.4}, \"at\": {d}, \"exact\": {d}, \"bad\": {d}}}", .{ name, f.max_abs, f.at, f.exact, f.bad }); + return f.bad == 0; +} + +/// The Gaussian decoder, the activation and the PLY, each on the SDK's own inputs. +fn gaussians(a: std.mem.Allocator, io: std.Io, d: *id.Decoder, c: *const Case, dir: []const u8, out: *std.ArrayList(u8)) !bool { + var pass = true; + const anchors = try floats(a, c.f, "anchors"); + const sdk = try floats(a, c.f, "raw"); + const raw = try a.alloc(f32, sdk.len); + _ = try d.gaussianRows(anchors, raw); + const m = ref.match(raw, sdk, dc.raw_width); + const w = dr.worst(raw, sdk, dc.raw_width); + std.debug.print("Gaussian decoder on the SDK's {d} anchors against the SDK: max |diff| {e:.3}, cosine {d:.9}, worst row {d:.5} (anchor {d})\n", .{ anchors.len / 3, m.max_abs, m.cosine, w.err, w.row }); + try print(out, a, ", \"gaussian_decoder\": {{\"sdk\": {{\"max_abs\": {e:.4}, \"cosine\": {d:.10}, \"worst_row\": {d:.6}, \"worst_at\": {d}}}", .{ m.max_abs, m.cosine, w.err, w.row }); + if (c.f.get("raw32") != null) { + const nudged = try nudges(a, c.f, "raw16n", ""); + pass = (try rulerGate(a, out, raw, sdk, try floats(a, c.f, "raw32"), try floats(a, c.f, "raw32h"), nudged, dc.raw_width)) and pass; + } + try print(out, a, "}}, \"activated\": {{", .{}); + const n = anchors.len / 3 * dc.per_anchor; + const g = try gs.Gaussians.alloc(a, n); + id.decode.activateParallel(d.params(), sdk, anchors, g); + const theirs: gs.Gaussians = .{ .xyz = try floats(a, c.f, "xyz"), .features_dc = try floats(a, c.f, "features_dc"), .scaling = try floats(a, c.f, "scaling"), .rotation = try floats(a, c.f, "rotation"), .opacity = try floats(a, c.f, "opacity") }; + try print(out, a, "\"rows\": {d}", .{n}); + inline for (.{ "xyz", "features_dc", "scaling", "rotation", "opacity" }) |name| { + pass = (try fieldGate(a, out, name, @field(g, name), @field(theirs, name))) and pass; + } + const reference = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.allocPrint(a, "{s}/reference.ply", .{dir}), a, .limited(1 << 31)); + const mine = try gs.ply(a, theirs); + const p = dr.ply(mine, reference); + const ok = dr.plyPass(p, 1e-5); + pass = pass and ok; + std.debug.print("PLY of the SDK's Gaussians against its reference.ply: header {s}, {d} rows {s}; columns differing:", .{ if (p.header_equal) "equal" else "DIFFERENT", p.rows, if (p.rows_equal) "equal" else "DIFFERENT" }); + for (gs.properties, p.differing, p.max_abs) |name, k, mx| if (k > 0) std.debug.print(" {s} {d} (max {e:.2})", .{ name, k, mx }); + std.debug.print(": {s}\n", .{verdict(ok)}); + try print(out, a, "}}, \"ply\": {{\"header_equal\": {}, \"rows_equal\": {}, \"differing\": {any}, \"max_abs\": [", .{ p.header_equal, p.rows_equal, p.differing }); + for (p.max_abs, 0..) |mx, i| try print(out, a, "{s}{e:.3}", .{ if (i == 0) "" else ", ", mx }); + try print(out, a, "], \"pass\": {}}}", .{ok}); + return pass; +} + +fn camera(a: std.mem.Allocator, io: std.Io, c: *const Case, dir: []const u8, out: *std.ArrayList(u8)) !bool { + const v = try floats(a, c.f, "camera"); + const cam = gs.camera(v[0..5].*); + const text = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.allocPrint(a, "{s}/decode.json", .{dir}), a, .limited(1 << 20)); + const doc = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + const sdk = doc.value.object.get("camera") orelse return error.BadCase; + var most: f64 = 0; + inline for (.{ "yaw", "pitch", "radius", "fov" }) |k| { + const s = sdk.object.get(k) orelse return error.BadCase; + const want: f64 = switch (s) { + .float => |x| x, + .integer => |x| @floatFromInt(x), + else => return error.BadCase, + }; + most = @max(most, @abs(@field(cam, k) - want)); + } + const json = try std.fmt.allocPrint(a, "{{\"yaw\": {d}, \"pitch\": {d}, \"radius\": {d}, \"fov\": {d}}}\n", .{ cam.yaw, cam.pitch, cam.radius, cam.fov }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/camera.json", .{dir}), .data = json }); + const ok = most <= 1e-6; + std.debug.print("camera {s}against the SDK's: max |diff| {e:.3}: {s}\n", .{ json, most, verdict(ok) }); + try print(out, a, ", \"camera_max_abs\": {e:.3}", .{most}); + return ok; +} + +/// The whole decode twice for each generator: identical PLY bytes, or not. +fn repeats(a: std.mem.Allocator, d: *id.Decoder, c: *const Case, latent: []const f32, n: usize, first: []const u8) !bool { + var draws: oc.Draws = .{ .replay = &c.replay }; + var again = try id.decode.decode(a, d, latent, n, &draws, null); + const same_replay = std.mem.eql(u8, first, try gs.ply(a, again.gaussians)); + var bytes: [2][]u8 = undefined; + for (&bytes) |*b| { + var native: oc.Draws = .{ .native = std.Random.DefaultPrng.init(42) }; + again = try id.decode.decode(a, d, latent, n, &native, null); + b.* = try gs.ply(a, again.gaussians); + } + const same_native = std.mem.eql(u8, bytes[0], bytes[1]); + std.debug.print("the decode again on the SDK's draws: {s}; twice on the native generator: {s}\n", .{ if (same_replay) "identical bytes" else "DIFFERENT bytes", if (same_native) "identical bytes" else "DIFFERENT bytes" }); + return same_replay and same_native; +} + +fn timed(a: std.mem.Allocator, io: std.Io, device: mtl.Device, args: []const [:0]const u8) !u8 { + const n = try std.fmt.parseInt(usize, args[4], 10); + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/case.safetensors", .{args[2]})); + defer case.close(io); + const latent = try floats(a, &case, "sample"); + const loading = mtl.clock.seconds(); + const d = try id.Decoder.load(std.heap.page_allocator, io, device, args[1], try dc.anchors(n)); + defer d.deinit(); + const load_s = mtl.clock.seconds() - loading; + var line: std.ArrayList(u8) = .empty; + try print(&line, a, "{{\"device\": \"{s}\", \"gaussians\": {d}, \"load_s\": {d:.2}", .{ device.name(), n, load_s }); + for ([_][]const u8{ "cold", "warm" }) |label| { + var draws: oc.Draws = .{ .native = std.Random.DefaultPrng.init(42) }; + const started = mtl.clock.seconds(); + var r = try id.decode.decode(a, d, latent, n, &draws, null); + const decoded = mtl.clock.seconds() - started; + const writing = mtl.clock.seconds(); + const bytes = try gs.ply(a, r.gaussians); + if (args.len == 6) try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = args[5], .data = bytes }); + const ply_s = mtl.clock.seconds() - writing; + const t = r.time; + std.debug.print("{s}: decode {d:.3} s (octree {d:.3}, {d:.3} on the GPU; Gaussian decoder {d:.3}, {d:.3} on the GPU; activation {d:.3}); PLY {d:.3} s\n", .{ label, decoded, t.octree, t.octree_gpu, t.gaussian, t.gaussian_gpu, t.activation, ply_s }); + try print(&line, a, ", \"{s}\": {{\"decode_s\": {d:.4}, \"octree_s\": {d:.4}, \"octree_gpu_s\": {d:.4}, \"gaussian_s\": {d:.4}, \"gaussian_gpu_s\": {d:.4}, \"activation_s\": {d:.4}, \"ply_s\": {d:.4}}}", .{ label, decoded, t.octree, t.octree_gpu, t.gaussian, t.gaussian_gpu, t.activation, ply_s }); + r.deinit(a); + a.free(bytes); + } + try print(&line, a, ", \"metal_allocated_bytes\": {d}}}\n", .{device.allocated()}); + std.debug.print("{s}", .{line.items}); + return 0; +} + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + const time_mode = args.len >= 5 and std.mem.eql(u8, args[3], "time"); + if (!(args.len == 3 or (time_mode and args.len <= 6))) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + if (time_mode) return timed(std.heap.page_allocator, io, device, args); + const dir = args[2]; + var file = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/decode.safetensors", .{dir})); + defer file.close(io); + var c: Case = .{ .f = &file, .parents = undefined, .draws = undefined, .replay = undefined }; + var name: [32]u8 = undefined; + for (0..dc.max_level) |i| { + c.parents[i] = try coords(a, &file, try std.fmt.bufPrint(&name, "parents.{d}", .{i + 1})); + c.draws[i] = try floats(a, &file, try std.fmt.bufPrint(&name, "draws.{d}", .{i + 1})); + } + c.replay = .{ .parents = c.parents, .draws = undefined, .jitter = try floats(a, &file, "jitter") }; + for (&c.replay.draws, c.draws) |*o, v| o.* = v; + const latent = try floats(a, &file, "latent"); + const n = (try tensor(&file, "anchors")).numel() / 3 * dc.per_anchor; + const d = try id.Decoder.load(std.heap.page_allocator, io, device, args[1], try dc.anchors(n)); + defer d.deinit(); + try d.setLatent(latent); + var out: std.ArrayList(u8) = .empty; + try print(&out, a, "{{\"device\": \"{s}\", \"gaussians\": {d}", .{ device.name(), n }); + const sdk_freqs = try floats(a, &file, "freqs"); + var freq_diff: usize = 0; + for (d.pointFreqs(), sdk_freqs) |x, y| freq_diff += @intFromBool(x != y); + std.debug.print("point embedding frequencies: {d} of {d} differ from torch's by its CPU pow's rounding\n", .{ freq_diff, sdk_freqs.len }); + try print(&out, a, ", \"freqs_differ\": {d}", .{freq_diff}); + var pass = try levels(a, d, &c, &out); + if (file.get("probs32.1") != null) pass = (try blockGates(a, d, &c, &out)) and pass; + var run = try freeRun(a, d, &c, latent, n, &out); + const ours = try gs.ply(a, run.gaussians); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-decode.ply", .{dir}), .data = ours }); + const reference = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.allocPrint(a, "{s}/reference.ply", .{dir}), a, .limited(1 << 31)); + const whole = dr.ply(ours, reference); + const shape = whole.header_equal and whole.rows_equal; + pass = pass and shape; + std.debug.print("our own PLY: header {s} the SDK's, rows {s}: {s}\n", .{ if (whole.header_equal) "equal to" else "DIFFERENT from", if (whole.rows_equal) "equal" else "DIFFERENT", verdict(shape) }); + try print(&out, a, ", \"own_ply_header_equal\": {}, \"own_ply_rows_equal\": {}", .{ whole.header_equal, whole.rows_equal }); + pass = (try gaussians(a, io, d, &c, dir, &out)) and pass; + pass = (try camera(a, io, &c, dir, &out)) and pass; + const repeat = try repeats(a, d, &c, latent, n, ours); + run.deinit(a); + pass = pass and repeat; + try print(&out, a, ", \"repeat_identical\": {}, \"gates_pass\": {}}}\n", .{ repeat, pass }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-decode.json", .{dir}), .data = out.items }); + std.debug.print("gates: {s}\n", .{if (pass) "pass" else "FAIL"}); + return if (pass) 0 else 1; +} + +/// Levels 1 to 3 block by block, each block alone on the SDK's own input rows against that block in float32. +fn blockGates(a: std.mem.Allocator, d: *id.Decoder, c: *const Case, out: *std.ArrayList(u8)) !bool { + var pass = true; + var name: [32]u8 = undefined; + try print(out, a, ", \"blocks\": [", .{}); + for (1..4) |l| { + const parents = c.parents[l - 1]; + const centres = try a.alloc(f32, parents.len * 3); + oc.centres(parents, @intCast(l), centres); + var input = try floats(a, c.f, try std.fmt.bufPrint(&name, "stem.{d}", .{l})); + const zero = try a.alloc(f16, input.len); + @memset(zero, 0); + const stem = try a.alloc(f32, input.len); + for (stem, try d.octreeRows(centres, zero, l, null)) |*o, v| o.* = v; + const ms = ref.match(stem, input, dc.hidden); + std.debug.print("level {d} stem against the SDK's: max |diff| {e:.3}, worst row {d:.6}\n", .{ l, ms.max_abs, ms.worst_row }); + try print(out, a, "{s}{{\"level\": {d}, \"stem_max_abs\": {e:.3}, \"stem_worst_row\": {d:.6}, \"each\": [", .{ if (l == 1) "" else ", ", l, ms.max_abs, ms.worst_row }); + for (0..dc.octree_blocks) |b| { + const in16 = try a.alloc(f16, input.len); + for (in16, input) |*o, v| o.* = @floatCast(v); + const ours = try a.alloc(f32, input.len); + for (ours, try d.octreeRows(centres, in16, l, b)) |*o, v| o.* = v; + const sdk = try floats(a, c.f, try std.fmt.bufPrint(&name, "block{d}.{d}", .{ b, l })); + const exact = try floats(a, c.f, try std.fmt.bufPrint(&name, "block{d}_f32.{d}", .{ b, l })); + const m = ref.match(ours, exact, dc.hidden); + const r = ref.match(sdk, exact, dc.hidden); + // the residual update (output less the fp16 input) too: the output's own rounding hides it in the rows + const deltas = try a.alloc(f32, 3 * input.len); + for (0..input.len) |i| { + const base: f32 = in16[i]; + deltas[i] = ours[i] - base; + deltas[input.len + i] = sdk[i] - base; + deltas[2 * input.len + i] = exact[i] - base; + } + const dm = ref.match(deltas[0..input.len], deltas[2 * input.len ..], dc.hidden); + const dsdk = ref.match(deltas[input.len .. 2 * input.len], deltas[2 * input.len ..], dc.hidden); + const ok = ref.within(m, r) and ref.within(dm, dsdk); + pass = pass and ok; + std.debug.print(" block {d} against float32: worst row {d:.6}, 1 - cosine {e:.3}; the SDK's own {d:.6}, {e:.3}; its update {d:.5}, {e:.3}; the SDK's {d:.5}, {e:.3}: {s}\n", .{ b, m.worst_row, 1 - m.cosine, r.worst_row, 1 - r.cosine, dm.worst_row, 1 - dm.cosine, dsdk.worst_row, 1 - dsdk.cosine, verdict(ok) }); + try print(out, a, "{s}{{\"worst_row\": {d:.6}, \"cosine\": {d:.10}, \"sdk_worst_row\": {d:.6}, \"sdk_cosine\": {d:.10}, \"update_worst_row\": {d:.6}, \"update_cosine\": {d:.10}, \"sdk_update_worst_row\": {d:.6}, \"sdk_update_cosine\": {d:.10}, \"pass\": {}}}", .{ if (b == 0) "" else ", ", m.worst_row, m.cosine, r.worst_row, r.cosine, dm.worst_row, dm.cosine, dsdk.worst_row, dsdk.cosine, ok }); + input = sdk; + } + try print(out, a, "]}}", .{}); + } + try print(out, a, "]", .{}); + return pass; +} From 5302c79884162c5fde65c9022d22207360381e8e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:06:05 +0800 Subject: [PATCH 09/21] idol: register the decoders' nine Metal kernels with the kernel contract metal_idol_decoder*.zig registers dec_embed, dec_half, dec_kv, dec_heads, dec_attention_few and dec_attention_{00,01,10,11} as "metal idol decoder ...", against float64 at idol-sdk's fp32 and fp16 points. Fifteen controls each go red at their gate: an added or dropped rounding point, fp16 sums, P's lo half dropped (caught on heads built so O is P's fp16 remainder alone), keys past Lk unmasked, an input written, and the two NaN rules below undone. Gate-5 records are in idol_records.zig; the common Lib is now generic over a source and its launchers. Gate 1 found two real rule breaks, fixed with finite outputs unchanged (the 262k and 32k decode cases give byte-identical PLYs and gate lines): dec_heads clamped its norm by max, dropping a NaN; and the tiled attention gave NaN for an infinite value, because an fp16 lo of 0 or of the opposite sign met the infinity. The lo half now meets the values clamped to fp16's finite range (about 2% at 262k, 7% at 1M). docs/recipes/idol.md gains a Decoders part with the precision, gates, run-a-case commands, measurements and the contract SUMMARY. --- CHANGELOG.md | 2 + docs/kernel-contract.md | 4 + docs/recipes/idol.md | 164 ++++- zig/kernels/metal/idol_decoder.metal | 14 +- zig/src/core/contract/idol_records.zig | 20 +- zig/tests/contract/metal_idol_common.zig | 95 +-- zig/tests/contract/metal_idol_decoder.zig | 147 +++++ .../contract/metal_idol_decoder_attention.zig | 296 +++++++++ .../contract/metal_idol_decoder_steps.zig | 562 ++++++++++++++++++ zig/tests/contract/metal_limits.zig | 3 + zig/tests/contract/metal_main.zig | 2 + 11 files changed, 1257 insertions(+), 52 deletions(-) create mode 100644 zig/tests/contract/metal_idol_decoder.zig create mode 100644 zig/tests/contract/metal_idol_decoder_attention.zig create mode 100644 zig/tests/contract/metal_idol_decoder_steps.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 2abaa9af8a..40ae4002d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol's octree and Gaussian decoders now run natively on M5 Metal, from the denoiser's latent to the 3DGS PLY and camera. On an M5 Ultra, 262,144 Gaussians decode in 0.343 s against the SDK's 0.739 s, and 1,048,576 in 1.963 s against 4.653 s, at a 2.43 GB peak footprint against 5.64 GB. Given the SDK's probabilities and draws, every octree level's counts match, and the PLY's header and rows equal the SDK's. Both decoders sit as close to an fp32 run as the SDK's own; `tf-idol-decode` runs those gates from a `tools/zig/idol_case.py decode` case. The decoders' nine kernels are registered with the kernel contract (after #1766, for #1757). Thanks to @jayleaton. + - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. - Qwen3.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b2b465117a..5e71dae356 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -188,6 +188,10 @@ HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row Metal also checks Idol's denoiser (`metal_idol*.zig`): its fp16 projections and GELU form, the split-K fp32 rows, `id_gemm32` (true fp32 sums: the NAX fp32 path in its place is a control gate 4 must catch), the LayerNorms, RePo's angles, the heads, the elementwise steps and the attention's four forms, each at idol-sdk's fp16 points (19 controls). +Its decoders' own kernels (`metal_idol_decoder*.zig`) are checked too: the point embedding, the stem's fp16 cast, +`to_kv`, the heads' RMS norm, the attention's four forms (P V as fp16 P hi + lo) and its fp32 short form. Each is checked +at idol-sdk's fp32 and fp16 points, with 15 controls. One drops P's lo half and must go red on heads built so that O is +P's fp16 remainder alone. RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 5b506262a9..dea9f414cf 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -1,8 +1,9 @@ -# Idol's denoiser on Metal +# Idol on Metal Idol (speridlabs) turns one image into a 3D Gaussian splat. A 0.37B flow transformer denoises 8,192 latent rows of 16 channels, conditioned on the image's DINOv3 and FLUX.2 VAE rows; two decoders then grow an octree from that latent and -emit the Gaussians. `zig/src/families/idol` runs the denoiser on the Metal runtime from the released checkpoint: +emit the Gaussians. `zig/src/families/idol` runs the denoiser and both decoders on the Metal runtime from the released +checkpoint ([Decoders](#decoders) below). The denoiser: - once an image: the condition embedding (DINOv3's 1,280 channels and the VAE's 128, each projected to 1,024 and summed) and the two condition-refiner blocks (no modulation, affine LayerNorms); @@ -15,8 +16,8 @@ each row's LayerNorm to per-head 3D positions, angles from 10, 10 and 12 frequen The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_cond - 3 v_uncond` (two passes a step: with the condition, and with zeros). -The image encoders, the decoders and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's -SDK in torch, a development reference. It is a denoiser, not a token family: it does not register with the lane +The image encoders are not ported yet. `tools/zig/idol_case.py` runs the publisher's SDK in torch, a development +reference. It is a denoiser, not a token family: it does not register with the lane engine, the server or `families.detect`. It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them. @@ -129,3 +130,158 @@ Results: 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. - **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the denoiser's fp16 weights in shared Metal buffers. + +## Decoders + +`IdolPipeline.decode` turns the denoiser's final latent [8192, 16] into Gaussians. Both decoders run on Metal from +`decoder.safetensors` and `decoder_gs.safetensors`; the steps between them run on the host +(`zig/src/families/idol/decode.zig`): + +- **Octree decoder** (4 blocks): eight levels, each from the occupied parent cells' centres to 8 child logits. Each + block is cross attention from the cells to the latent, then the MLP, both modulated by the level's embedding + (`adaLN_modulation(l_embedder(2^level))`, computed once on the host). The host takes the softmax, splits each + parent's points among its children by `systematic_sample`, and keeps the children that got points. After level 8 + each leaf becomes its count of anchors, jittered inside the cell. +- **Gaussian decoder** (16 blocks): self attention over the anchors, cross attention to the latent, then the MLP. + Each anchor gets 480 outputs: 32 Gaussians of offset, colour, scale, rotation, opacity and offset scale. +- **Host**: `gaussian.py`'s activations, the 3DGS PLY and the camera JSON (`gaussian.zig`). + +A request asks for 32,768 to 1,048,576 Gaussians, 32 an anchor. `Decoder.load` sizes its scratch for the largest it +will serve. + +### Precision + +The SDK converts the blocks' linear layers to fp16. The stem, the head and the latent's rows before the blocks stay +fp32. This port keeps that split: + +- Stem in fp32: `in_proj` of the 3 coordinates as one FMA chain, plus the point embedding (`dec_embed`: sin and cos + of torch's fp32 angle `fp32(fp32(x f) pi)`, the frequencies `torch.pow`'s on the CPU, bit for bit). Then + `input_layer` by the denoiser's `id_gemm32` and its bias, rounded once to fp16 (`dec_half`). +- Blocks in fp16, through the denoiser's kernels: `id_linear` (fp32 sums, bias, one fp16 rounding, GELU after it), + `id_norm` and `id_gate_add`. +- Cross attention: `dec_kv` computes `to_kv` of the fp16 latent as 16-deep fp32 FMA sums, then the bias, rounded to + fp16. `dec_heads` applies MultiHeadRMSNorm in fp32 (`F.normalize`, times gamma times 8) and rounds to fp16. +- Attention: + - Scores and sums are fp32. P meets V on the tensor units as two fp16 halves, `hi = fp16(P)` and + `lo = fp16(P - hi)`, so P keeps about 22 bits. A single fp32 P would lose bits on the units: each pass would come + out about 2^-12 low, and the blocks would compound it. + - At 8 query rows or fewer (the octree's first levels), torch on MPS runs its fp32 vector kernel, and + `dec_attention_few` runs fp32 throughout as well. +- Head: the denoiser's `id_norm32` LayerNorm on the stream cast to fp32, `id_gemm32` and `id_bias32`. +- Host: softmax and resampling in fp32, as `systematic_sample` runs them. + +### Gates + +`tf-idol-decode` gates a case from `idol_case.py decode` on the SDK's own inputs: + +- **Counts.** Every octree level's split, run on the SDK's probabilities and draws, must give the SDK's counts row for + row. +- **Probabilities and raw rows.** Each level's probabilities and the Gaussian decoder's rows are compared against the + same decoder in fp32 on the same inputs. They must stay within twice the worst of three distances: + - the SDK's own distance to fp32; + - the fp16 floor: fp32 with only its last rows rounded once; + - the SDK rerun with 1% of its first block's fp16 inputs nudged one ulp, three seeds. + + Including the nudged reruns matters. They show how far the SDK itself moves under a change no fp16 stream can + avoid. At level 8, one ulp on 1% of the inputs moves the SDK's worst row from 0.0415 to 0.0488, more than the gap + between this port and the SDK. +- **Blocks.** Levels 1 to 3, block by block on the SDK's fp16 rows: each block's output and its residual update, + within twice the SDK's own distance to fp32. +- **Output.** The activated fields to 1e-5 relative plus 1e-6, the PLY's header, rows and columns against the SDK's + `reference.ply`, and the camera. +- **Repeat.** A second decode on the SDK's draws, and two on the native generator, must give identical bytes. + +The free run on the SDK's draws is reported, not gated. Its topology leaves the SDK's at level 6 of the 262k case (5 +of the 32k case), and that is expected: + +- `systematic_sample` is discontinuous. A stratified draw lands on one side of a cumulative-probability edge or the + other, so a probability 1e-5 off can move one point to a sibling. +- Fed our own probabilities, a few rows a level differ by one point: 1 row at level 1, 20 rows at level 8. +- Once a child's count reaches or leaves zero, the set of occupied cells changes, and every later level follows. +- The counts gate isolates this: given the SDK's probabilities, every count matches. + +### Run a case + +`idol_case.py decode` decodes a denoiser case's final sample with the SDK. It imports `decoder.py`, `gaussian.py` and +`camera.py` only at their pinned sha256. It records each level's parents, probabilities, draws and counts, the anchors, +the raw rows, the activated Gaussians, the PLY and the camera. `--ruler` adds the fp32 runs, the fp16 floor, the +nudged reruns and levels 1 to 3 block by block. + +```sh +python3 tools/zig/idol_case.py decode WEIGHTS SDK case/ decode262k/ --gaussians 262144 --ruler +zig build tf-idol-decode -Dcpu=apple_m1 +zig-out/bin/tf-idol-decode WEIGHTS decode262k/ # the gates; writes .json, .ply, camera.json +zig-out/bin/tf-idol-decode WEIGHTS case/ time 262144 out.ply # the native generator, cold then warm +``` + +### Measured + +Setup: + +- The same Mac and SDK as above. +- The avocado case's final sample, seed 42. +- 262,144 Gaussians unless noted. + +Results: + +- **Counts:** identical at all 8 levels, given the SDK's probabilities and draws. +- **Octree probabilities**, worst row against fp32: level 1 0.00019 (the SDK 0.00009), level 8 0.0453 (the SDK + 0.0415, nudged 0.0488). +- **Gaussian decoder**, worst row against fp32: 0.133, `1 - cos` 4.5e-6. The SDK's are 0.211 and 8.1e-6; its nudged + reruns reach 0.195. At 32,768 Gaussians this port's worst row is 0.0213 and the SDK's 0.0072, within the bound + (nudged 0.0305). +- **PLY:** header and all rows equal to the SDK's. Opacity and scale columns are within 1.9e-6; the other columns are + bit-equal. +- **Repeat:** identical bytes. +- **Speed and memory**, warm, `tf-idol-decode ... time` against the SDK's warm decode. The footprint is the peak + process footprint (`/usr/bin/time -l`). + + | Gaussians | ours | SDK | our peak footprint | SDK's | + | ---: | ---: | ---: | ---: | ---: | + | 32,768 | 0.071 s | 0.170 s | 1.06 GB | 2.97 GB | + | 262,144 | 0.343 s | 0.739 s | 1.33 GB | — | + | 1,048,576 | 1.963 s | 4.653 s | 2.43 GB | 5.64 GB | + + At 1M, the Gaussian decoder's self attention over 32,768 anchors takes most of the time. + +### Kernel contract + +`zig/tests/contract/metal_idol_decoder*.zig` registers the 9 pipelines of `idol_decoder.metal`, named +`metal idol decoder ...`, so `--only "metal idol"` runs them with the denoiser's: + +- `dec_embed`, `dec_half`, `dec_kv` and `dec_heads`; +- the attention's four tiled forms and `dec_attention_few`. + +Form `10` is never served, because every call has Lq = Lk or Lk = 8,192; it is checked at made-up shapes. + +Gate 3 compares each step against float64 at the SDK's fp32 and fp16 points, as a range over its fp32 steps. +Attention's allowance is: + +- 2^-17 of `sum P |v| / l`, for the fp32 scores, exp2 and sums; +- each P's lo-half floor (2^-25, or P where P is smaller) times |v|; +- O's own fp16 rounding. + +Half the heads are built so O is P's fp16 remainder alone. On those heads, a P V without P's lo half, or with P rounded +to fp16, gives 0. Fifteen controls must each go red at their gate: + +- **Shapes:** add an fp16 rounding point (`pe`, `y` before the bias, `ctx w` before the bias, `x / n`), drop P's lo + half, round the short form's P to fp16, leave keys past Lk unmasked. +- **Sums:** sum in fp16 (`in_proj`, `to_kv`, the heads' squares, P V in the short and the tiled forms). +- **Repeat:** `dec_half` writing its input. +- **Nonfinite:** the two NaN and infinity rules below, undone. + +Gate 1 found two rules the first version broke, and both kernels are fixed: + +- `dec_heads` clamped its norm by `max`, which drops a NaN; it now keeps it, as `clamp_min` does. +- The tiled attention gave NaN for an infinite value: an fp16 `lo` of 0, or of the opposite sign, times infinity. The + lo half now meets the values clamped to fp16's finite range, while the hi half still carries the infinity. +- Finite outputs are unchanged. The 262k and 32k cases give byte-identical PLYs and gate lines before and after. +- The clamp costs about 2% at 262k and 6 to 7% at 32k and 1M. + +The tensor-unit kernels do not build for M1 or M2 offline; the rest build at their reserved thread counts. + +`tf-kernel-contract --only "metal idol"` on the M5 Ultra, the denoiser's 27 kernels and 19 controls with these: + +``` +SUMMARY 36 kernels: 36 pass gates 1-4, 0 fail, 0 known findings; 34 controls: 34 red as planted, 0 missed; 0 limit problems +``` diff --git a/zig/kernels/metal/idol_decoder.metal b/zig/kernels/metal/idol_decoder.metal index 77f6ae884a..ab351220e0 100644 --- a/zig/kernels/metal/idol_decoder.metal +++ b/zig/kernels/metal/idol_decoder.metal @@ -61,7 +61,9 @@ struct DecHeads { dst[to + 1] = half(a1); return; } - const float n = max(metal::precise::sqrt(simd_sum(a0 * a0 + a1 * a1)), 1e-12f); + // torch's clamp_min keeps a NaN norm, which max would drop + const float len = metal::precise::sqrt(simd_sum(a0 * a0 + a1 * a1)); + const float n = isnan(len) ? len : max(len, 1e-12f); const uint g = h * 64 + lane * 2; dst[to] = half(a0 / n * gamma[g] * 8.0f); dst[to + 1] = half(a1 / n * gamma[g + 1] * 8.0f); @@ -165,6 +167,13 @@ inline void dec_fold(thread const tfp::frag& f, thread float (&r)[2]) { } } +// v with each infinity at fp16's largest finite value of its sign (P's lo half meets these values). +inline void dec_finite(thread tfp::frag& v) { + thread half4* q = (thread half4*)&v; + q[0] = clamp(q[0], half4(-65504.0h), half4(65504.0h)); + q[1] = clamp(q[1], half4(-65504.0h), half4(65504.0h)); +} + // One 64-query block of one head over 32-key blocks; P and F as Kernels.attention sets them (kernels.zig). template inline void dec_attention(const device T* Q, const device T* K, const device T* V, device T* O, const device int* P, @@ -265,6 +274,9 @@ inline void dec_attention(const device T* Q, const device T* K, const device T* tfp::frag_get(v1, V, D, 16 * k, 16 * d + 16, home); } tfp::mma_16x32(acc[d], acc[d + 1], ph[k], v0, v1); + // the lo half meets the values held finite: a 0 or negative lo times an infinite value would make P V NaN + dec_finite(v0); + dec_finite(v1); tfp::mma_16x32(acc[d], acc[d + 1], pl[k], v0, v1); } } diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig index ae39549164..22a12df75a 100644 --- a/zig/src/core/contract/idol_records.zig +++ b/zig/src/core/contract/idol_records.zig @@ -1,7 +1,9 @@ -//! Idol's denoiser launches (idol.metal and idol_attention.metal after gemm_nax.metal), each source-reserved. +//! Idol's launches, each source-reserved: the denoiser's (idol.metal, idol_attention.metal), then idol_decoder.metal's. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; -pub const records = [_]M{ +pub const records = denoiser ++ decoder; + +const denoiser = [_]M{ .{ .name = "idol id_linear_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "idol id_linear_01", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "idol id_linear_10", .threads = 256, .tg_memory = 0, .reserved = 256 }, @@ -28,7 +30,19 @@ pub const records = [_]M{ .{ .name = "idol id_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, }; +const decoder = [_]M{ + .{ .name = "idol decoder dec_embed", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol decoder dec_half", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol decoder dec_kv", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol decoder dec_heads", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol decoder dec_attention_few", .threads = 256, .tg_memory = 2352, .reserved = 256 }, + .{ .name = "idol decoder dec_attention_00", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol decoder dec_attention_01", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol decoder dec_attention_10", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "idol decoder dec_attention_11", .threads = 128, .tg_memory = 0, .reserved = 128 }, +}; + test "the Idol launches fit every Apple family, one record a pipeline" { try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); - try std.testing.expectEqual(@as(usize, 24), records.len); + try std.testing.expectEqual(@as(usize, 24 + 9), records.len); } diff --git a/zig/tests/contract/metal_idol_common.zig b/zig/tests/contract/metal_idol_common.zig index 222b6b22b4..fd13ab7177 100644 --- a/zig/tests/contract/metal_idol_common.zig +++ b/zig/tests/contract/metal_idol_common.zig @@ -195,51 +195,58 @@ pub const Edit = struct { from: []const u8, to: []const u8 }; /// A control: the kernel it breaks, the gate that must go red, its source edits and any text it appends. pub const Control = struct { kernel: []const u8, what: []const u8, gate: g.Gate = .shapes, edits: []const Edit = &.{}, append: []const u8 = "" }; -/// The family's source as the engine compiles it (or a control's copy), its pipelines made as tests ask for them. -pub const Lib = struct { - a: std.mem.Allocator, - device: mtl.Device, - ctl: ?Control, - k: ks.Kernels = undefined, - built: bool = false, - made: [ks.names.len]bool = @splat(false), - - fn text(l: *Lib) ![]u8 { - var out = try tf.frags.source(l.device, l.a, sources.idol); - const x = l.ctl orelse return out; - for (x.edits) |e| { - errdefer l.a.free(out); - if (std.mem.count(u8, out, e.from) != 1) { - std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, e.from }); - return error.ControlTextNotFoundOnce; +/// The denoiser's library: idol.metal after gemm_nax.metal, through kernels.zig's launchers. +pub const Lib = LibOf(ks.Kernels, &ks.names, sources.idol); + +/// A source as the engine compiles it (or a control's copy), its pipelines made as tests ask for them; K's launchers. +pub fn LibOf(comptime K: type, comptime names: []const []const u8, comptime source: []const u8) type { + return struct { + a: std.mem.Allocator, + device: mtl.Device, + ctl: ?Control, + k: K = undefined, + built: bool = false, + made: [names.len]bool = @splat(false), + + const This = @This(); + + fn text(l: *This) ![]u8 { + var out = try tf.frags.source(l.device, l.a, source); + const x = l.ctl orelse return out; + for (x.edits) |e| { + errdefer l.a.free(out); + if (std.mem.count(u8, out, e.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, e.from }); + return error.ControlTextNotFoundOnce; + } + const next = try std.mem.replaceOwned(u8, l.a, out, e.from, e.to); + l.a.free(out); + out = next; } - const next = try std.mem.replaceOwned(u8, l.a, out, e.from, e.to); - l.a.free(out); - out = next; + if (x.append.len == 0) return out; + defer l.a.free(out); + return std.mem.concat(l.a, u8, &.{ out, "\n", x.append }); } - if (x.append.len == 0) return out; - defer l.a.free(out); - return std.mem.concat(l.a, u8, &.{ out, "\n", x.append }); - } - - /// The pipeline `name` made on this library (compiled on first use). - pub fn need(l: *Lib, name: []const u8) !void { - if (!l.built) { - const src = try l.text(); - defer l.a.free(src); - l.k.own = try mtl.Library.fromSource(l.device, src, mtl.CompileOptions.mlx()); - l.built = true; + + /// The pipeline `name` made on this library (compiled on first use). + pub fn need(l: *This, name: []const u8) !void { + if (!l.built) { + const src = try l.text(); + defer l.a.free(src); + l.k.own = try mtl.Library.fromSource(l.device, src, mtl.CompileOptions.mlx()); + l.built = true; + } + for (names, 0..) |n, i| if (std.mem.eql(u8, n, name)) { + if (!l.made[i]) l.k.p[i] = try mtl.Pipeline.init(l.device, l.k.own, n, false); + l.made[i] = true; + return; + }; + return error.NoSuchPipeline; } - for (ks.names, 0..) |n, i| if (std.mem.eql(u8, n, name)) { - if (!l.made[i]) l.k.p[i] = try mtl.Pipeline.init(l.device, l.k.own, n, false); - l.made[i] = true; - return; - }; - return error.NoSuchPipeline; - } - pub fn deinit(l: *Lib) void { - for (l.k.p, l.made) |p, m| if (m) p.deinit(); - if (l.built) l.k.own.deinit(); - } -}; + pub fn deinit(l: *This) void { + for (l.k.p, l.made) |p, m| if (m) p.deinit(); + if (l.built) l.k.own.deinit(); + } + }; +} diff --git a/zig/tests/contract/metal_idol_decoder.zig b/zig/tests/contract/metal_idol_decoder.zig new file mode 100644 index 0000000000..2370e20ad2 --- /dev/null +++ b/zig/tests/contract/metal_idol_decoder.zig @@ -0,0 +1,147 @@ +//! idol_decoder.metal's kernels under the contract at idol-sdk's points; controls break one rounding, sum or rule each. +const std = @import("std"); +const mtl = @import("metal"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const steps = @import("metal_idol_decoder_steps.zig"); +const attn = @import("metal_idol_decoder_attention.zig"); + +const g = id.g; +const Control = id.Control; +const Lib = steps.Lib; + +const pv_lo = "tfp::mma_16x32(acc[d], acc[d + 1], pl[k], v0, v1);"; +const pv_half = pv_lo ++ " for (short e = 0; e < 8; e++) { acc[d][e] = float(half(acc[d][e])); " ++ + "acc[d + 1][e] = float(half(acc[d + 1][e])); }"; +const embed_sum = "const float s = fma(x2, w[c * 3 + 2], fma(x1, w[c * 3 + 1], x0 * w[c * 3])) + b[c];"; +const embed_half = "const float s = float(fma(half(x2), half(w[c * 3 + 2]), fma(half(x1), half(w[c * 3 + 1]), " ++ + "half(x0) * half(w[c * 3])))) + b[c];"; +const kv_sum = "s = fma(float(ctx[r * 16 + k]), float(w[c * 16 + k]), s);"; +const to_half = "o[i] = half(y[i] + b[at.x]);"; +/// dec_attention_few's four P V chains, each step rounded to fp16. +const few_half = blk: { + var out: [4]id.Edit = undefined; + for (&out, [_][]const u8{ "", " + 1", " + 2", " + 3" }, "xyzw") |*e, at, x| { + const slot = "o[4 * c" ++ at ++ "]"; + const step = "fma(e, vv." ++ [_]u8{x} ++ ", " ++ slot ++ ")"; + e.* = .{ .from = slot ++ " = " ++ step ++ ";", .to = slot ++ " = float(half(" ++ step ++ "));" }; + } + break :blk out; +}; + +const controls = [_]Control{ + .{ .kernel = "dec_embed", .what = "adds fp16(pe)", .edits = &.{.{ .from = "h[ulong(r) * 1024 + c] = s + pe;", .to = "h[ulong(r) * 1024 + c] = s + float(half(pe));" }} }, + .{ .kernel = "dec_embed", .what = "sums in_proj in fp16", .gate = .sums, .edits = &.{.{ .from = embed_sum, .to = embed_half }} }, + .{ .kernel = "dec_half", .what = "adds fp16(y) before the bias", .edits = &.{.{ .from = to_half, .to = "o[i] = half(float(half(y[i])) + b[at.x]);" }} }, + .{ .kernel = "dec_half", .what = "writes its input", .gate = .repeat, .edits = &.{.{ .from = to_half, .to = to_half ++ " ((device float*)y)[i] = 0.0f;" }} }, + .{ .kernel = "dec_kv", .what = "adds fp16(ctx w) before the bias", .edits = &.{.{ .from = "half(s + float(b[c]));", .to = "half(float(half(s)) + float(b[c]));" }} }, + .{ .kernel = "dec_kv", .what = "sums in fp16", .gate = .sums, .edits = &.{.{ .from = kv_sum, .to = "s = float(half(fma(float(ctx[r * 16 + k]), float(w[c * 16 + k]), s)));" }} }, + .{ .kernel = "dec_heads", .what = "clamps the norm by max, which drops a NaN", .gate = .nonfinite, .edits = &.{.{ .from = "const float n = isnan(len) ? len : max(len, 1e-12f);", .to = "const float n = max(len, 1e-12f);" }} }, + .{ .kernel = "dec_heads", .what = "adds fp16(x / n)", .edits = &.{.{ .from = "dst[to] = half(a0 / n * gamma[g] * 8.0f);", .to = "dst[to] = half(float(half(a0 / n)) * gamma[g] * 8.0f);" }} }, + .{ .kernel = "dec_heads", .what = "sums the squares in fp16", .gate = .sums, .edits = &.{.{ .from = "simd_sum(a0 * a0 + a1 * a1)", .to = "float(simd_sum(half(a0 * a0 + a1 * a1)))" }} }, + .{ .kernel = "dec_attention_few", .what = "rounds P to fp16", .edits = &.{.{ .from = "const float e = metal::precise::exp(s[i] - top);", .to = "const float e = float(half(metal::precise::exp(s[i] - top)));" }} }, + .{ .kernel = "dec_attention_few", .what = "sums P V in fp16", .gate = .sums, .edits = &few_half }, + .{ .kernel = "dec_attention_10", .what = "P's lo half meets infinite values", .gate = .nonfinite, .edits = &.{.{ .from = "dec_finite(v0);", .to = "" }} }, + .{ .kernel = "dec_attention_00", .what = "keys past Lk unmasked", .edits = &.{.{ .from = "s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kDecMasked;", .to = "s[f][e] = s[f][e];" }} }, + .{ .kernel = "dec_attention_01", .what = "sums P V in fp16", .gate = .sums, .edits = &.{.{ .from = pv_lo, .to = pv_half }} }, + .{ .kernel = "dec_attention_11", .what = "drops P's lo half", .edits = &.{.{ .from = pv_lo, .to = "" }} }, +}; + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + libs: std.ArrayList(*Lib) = .empty, + names: std.ArrayList([]u8) = .empty, + made: std.ArrayList(Made) = .empty, + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: *Lib, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: id.Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol decoder {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + inline for (.{ .{ steps.Embed, "dec_embed" }, .{ steps.Half, "dec_half" }, .{ steps.Kv, "dec_kv" }, .{ steps.Heads, "dec_heads" } }) |e| { + if (try w.name(e[1])) |name| { + const t = try s.make(e[0], .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + for (attn.forms, 0..) |f, i| { + const name = (try w.name(f.name)) orelse continue; + const t = try s.make(attn.Attention, .{ .a = s.a, .rig = rig, .lib = lib, .form = i, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + if (!device.tensorUnits()) { + std.debug.print("UNSUPPORTED metal idol decoder: the family runs on tensor units (M5) alone; its kernels unchecked here\n", .{}); + return s; + } + for ([_]?Control{null} ++ controls) |x| { + const lib = try a.create(Lib); + lib.* = .{ .a = a, .device = device, .ctl = x }; + s.libs.append(a, lib) catch |err| { + a.destroy(lib); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.libs.items) |l| { + l.deinit(); + a.destroy(l); + } + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_decoder_attention.zig b/zig/tests/contract/metal_idol_decoder_attention.zig new file mode 100644 index 0000000000..6b24301161 --- /dev/null +++ b/zig/tests/contract/metal_idol_decoder_attention.zig @@ -0,0 +1,296 @@ +//! The decoders' attention under the contract: softmax(q k^T / 8) v, P V as fp16 P hi + lo, or fp32 for 8 queries. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const steps = @import("metal_idol_decoder_steps.zig"); + +const c = id.c; +const g = id.g; +const D = id.D; +const kref = id.kref; +const dk = id.tf.idol.decoder_kernels; +const dc = id.tf.idol.decoder_config; + +const heads = dc.heads; +const C = dc.context_rows; + +/// Query rows repeat at this period (prime to the 64-row blocks): float64 costs 89 queries a head over every key. +const query_period = 89; + +/// Heads from here on make O one P's fp16 remainder alone (see `exact`): P without its lo half gives 0 there. +const exact_heads = 8; +const p_dot: f16 = -1; +const p_hi: f16 = @floatCast(@exp(-1.0 / 8.0)); + +/// Tiled forms by (Lq % 64 == 0)(Lk % 32 == 0), then the short form; rows: the 16 heads, n: Lq, k: Lk. +pub const forms = [5]struct { name: []const u8, served: []const c.Shape, tiles: [3]c.Tile }{ + .{ + .name = "dec_attention_00", + // self attention over 1,025 and 8,193 anchors (32,800 and 262,176 Gaussians) + .served = &.{ .{ .rows = heads, .n = 1025, .k = 1025 }, .{ .rows = heads, .n = 8193, .k = 8193 } }, + .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 32 } }, + }, + .{ + .name = "dec_attention_01", + // octree parents to the latent, the fewest past the short form, and 1,056 anchors' self attention + .served = &.{ .{ .rows = heads, .n = 1000, .k = C }, .{ .rows = heads, .n = 9, .k = C }, .{ .rows = heads, .n = 1056, .k = 1056 } }, + .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 64, .step = 32, .min = 32 } }, + }, + // Lq = Lk or Lk = 8,192 in every served call, so this form is never served: checked at its tails + .{ + .name = "dec_attention_10", + .served = &.{ .{ .rows = heads, .n = 4096, .k = 4101 }, .{ .rows = heads, .n = 128, .k = 70 } }, + .tiles = .{ .{}, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 32 } }, + }, + .{ + .name = "dec_attention_11", + // 32k Gaussians' anchors, parents to the latent, then 262k's and 1M's anchors + .served = &.{ + .{ .rows = heads, .n = 1024, .k = 1024 }, .{ .rows = heads, .n = 4096, .k = C }, + .{ .rows = heads, .n = 8192, .k = 8192 }, .{ .rows = heads, .n = 32768, .k = 32768 }, + }, + .tiles = .{ .{}, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 64, .step = 32, .min = 32 } }, + }, + .{ + .name = "dec_attention_few", + // the first octree levels' 1 to 8 parents against the latent + .served = &.{ .{ .rows = heads, .n = dk.few_rows, .k = C }, .{ .rows = heads, .n = 1, .k = C } }, + .tiles = .{ .{}, .{ .size = dk.few_rows }, .{ .size = 256 } }, + }, +}; +pub const few = 4; + +pub fn form(lq: usize, lk: usize) usize { + return @as(usize, @intFromBool(lq % 64 == 0)) * 2 + @intFromBool(lk % 32 == 0); +} + +/// The scores' 64-deep sums, then the keys' sums. +fn attentionDepth(s: c.Shape) usize { + return D + s.k; +} + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *steps.Lib, + form: usize, + name: []const u8, + lq: usize = 0, + lk: usize = 0, + stress: bool = false, + q: []f16 = &.{}, + k: []f16 = &.{}, + v: []f16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn spec(name: []const u8, f: usize) c.Spec { + return .{ .name = name, .accum = .f32, .depth = attentionDepth, .ops = 8, .out = .f16, .tiles = forms[f].tiles, .served = forms[f].served }; + } + + pub fn deinit(t: *Attention) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Attention) void { + t.a.free(t.q); + t.a.free(t.k); + t.a.free(t.v); + t.q = &.{}; + t.k = &.{}; + t.v = &.{}; + } + pub fn kernel(t: *Attention, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + fn takes(t: *const Attention, lq: usize, lk: usize) bool { + if (lq == 0 or lk == 0) return false; + if (t.form == few) return lq <= dk.few_rows and lk <= 32 * 256; + return form(lq, lk) == t.form; + } + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!t.takes(shape.n, shape.k)) return false; + t.rig.reset(); + t.free(); + t.lq = shape.n; + t.lk = shape.k; + t.stress = fill == .stress; + const L = t.lk; + t.q = try t.a.alloc(f16, heads * t.lq * D); + t.k = try t.a.alloc(f16, heads * L * D); + t.v = try t.a.alloc(f16, heads * L * D); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // q = k = 0; odd fp16 integers, partial sums under 2^24 that cancel to 4; a thread's fp16 chain loses bits + @memset(t.q, 0); + @memset(t.k, 0); + const half = L / 2; + const big: f32 = if (half / 2 * (127 + 2047) < 1 << 24) 2047 else 1023; + for (0..heads) |h| for (0..L) |j| { + const at = if (j < half) j else j - half; + const positive: f32 = if (at % 2 == 0) 127 else big; + const residual: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + const v: f32 = if (j < half) positive else if (j < 2 * half) residual - positive else 0; + @memset(t.v[(h * L + j) * D ..][0..D], @floatCast(v)); + }; + } else { + id.randomF16(rng, t.q, -2, 2); + id.randomF16(rng, t.k, -2, 2); + id.randomF16(rng, t.v, -1, 1); + if (L >= 2) for (exact_heads..heads) |h| t.exact(rng, h); + } + if (fill == .nonfinite) { + // head 0: values NaN, +Inf, -Inf, and both infinities in one column; head 1: a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]f16{ id.nan16, id.inf16, -id.inf16, id.inf16, -id.inf16 }) |p, v| if (p[0] < L) { + t.v[p[0] * D + p[1]] = v; + }; + t.q[(t.lq + 0) * D + 5] = id.nan16; + if (L > 7) t.k[(L + 7) * D + 3] = id.inf16; + } + // every head's queries repeat at the period, so float64 computes each distinct one once + const p = @min(t.lq, query_period); + for (0..heads) |h| for (p..t.lq) |i| { + @memcpy(t.q[(h * t.lq + i) * D ..][0..D], t.q[(h * t.lq + i % p) * D ..][0..D]); + }; + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.q)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.k)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.v)); + t.out = try t.rig.output(t.lq * heads * D * 2); + return true; + } + + /// Head h: scores 0, -1/8 (last key) and -256 (P = 1, e^-1/8, 0), values -2^m fp16(P) and 2^m: O = 2^m lo / l. + fn exact(t: *Attention, rng: std.Random, h: usize) void { + const L = t.lk; + for (0..t.lq) |i| { + const q = t.q[(h * t.lq + i) * D ..][0..D]; + @memset(q, 0); + q[0] = 1; + } + for (0..L) |j| { + const k = t.k[(h * L + j) * D ..][0..D]; + @memset(k, 0); + k[0] = if (j == 0) 0 else if (j == L - 1) p_dot else -2048; + } + for (0..D) |e| { + const m: f32 = @floatFromInt(@as(u32, 1) << @intCast(rng.uintLessThan(u32, 11))); + const s = id.sign(rng) * m; + t.v[(h * L) * D + e] = @floatCast(-s * @as(f32, p_hi)); + t.v[(h * L + L - 1) * D + e] = @floatCast(s); + } + } + + pub fn run(t: *Attention) !void { + try t.lib.need(forms[t.form].name); + const p = t.rig.begin(); + const x = &t.lib.k; + if (t.form == few) { + x.attentionFew(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.lq, t.lk); + } else x.attention(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.lq, t.lk); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return heads * t.lq * D; + } + + /// Float64 O [Lq, 16 * 64]; allowance 2^-17 sum P |v| / l + sum min(P, 2^-25) |v| / l + O's rounding. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.lk; + const scale: f64 = 0.125; + const kf = try t.a.alloc(f64, L * D); + defer t.a.free(kf); + const vf = try t.a.alloc(f64, L * D); + defer t.a.free(vf); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + var qf: [D]f64 = undefined; + var out: [D]f64 = undefined; + var tol: [D]f64 = undefined; + var floor: [D]f64 = undefined; + const p = @min(t.lq, query_period); + for (0..heads) |h| { + for (kf, vf, t.k[h * L * D ..][0 .. L * D], t.v[h * L * D ..][0 .. L * D]) |*kd, *vd, kb, vb| { + kd.* = kb; + vd.* = vb; + } + for (0..p) |i| { + var zero = true; + for (&qf, t.q[(h * t.lq + i) * D ..][0..D]) |*x, b| { + x.* = b; + zero = zero and x.* == 0; + } + var m: f64 = -std.math.inf(f64); + var nan = false; + for (s, 0..) |*sj, j| { + var dot: f64 = 0; + if (!zero) { + const Vec = @Vector(8, f64); + var acc: Vec = @splat(0); + var e: usize = 0; + while (e < D) : (e += 8) acc += @as(Vec, qf[e..][0..8].*) * @as(Vec, kf[j * D + e ..][0..8].*); + dot = @reduce(.Add, acc); + } + sj.* = dot * scale; + nan = nan or std.math.isNan(sj.*); + m = @max(m, sj.*); + } + if (nan) m = std.math.nan(f64); + var l: f64 = 0; + @memset(&out, 0); + @memset(&tol, 0); + @memset(&floor, 0); + for (s, 0..) |sj, j| { + const pj = @exp(sj - m); + const least = @min(pj, 0x1p-25); + l += pj; + for (&out, &tol, &floor, vf[j * D ..][0..D]) |*o, *tl, *fl, vv| { + o.* += pj * vv; + tl.* += pj * @abs(vv); + fl.* += least * @abs(vv); + } + } + var r = i; + while (r < t.lq) : (r += p) { + const at = (r * heads + h) * D; + for (out, tol, floor, 0..) |o, tl, fl, e| { + want[at + e] = o / l; + mag[at + e] = tl / l * 0x1p-17 + fl / l + @abs(o / l) * 0x1p-10 + 0x1p-24; + } + } + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, tolerance: []const f64, allow: []f64) !void { + for (want, tolerance, allow) |w, tl, *l| l.* = if (t.stress) c.boundExact(spec(t.name, t.form), w) else if (std.math.isFinite(w)) tl else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.lk; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..heads) |h| for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = t.v[(h * L + j) * D + e]; + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..t.lq) |i| out[(i * heads + h) * D + e] = o; + }; + } +}; + +test "the exact heads' P has a fp16 remainder far from both zero and a tie" { + const p = @exp(@as(f64, p_dot) / 8); + const ulp = 0x1p-11; + const lo = (p - @as(f64, p_hi)) / ulp; + try std.testing.expect(@abs(lo) > 0.25 and @abs(lo) < 0.4); +} diff --git a/zig/tests/contract/metal_idol_decoder_steps.zig b/zig/tests/contract/metal_idol_decoder_steps.zig new file mode 100644 index 0000000000..91c90d7e3c --- /dev/null +++ b/zig/tests/contract/metal_idol_decoder_steps.zig @@ -0,0 +1,562 @@ +//! dec_embed, dec_half, dec_kv and dec_heads under the contract: the decoders' steps at idol-sdk's rounding points. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const sources = @import("kernel_sources"); + +const c = id.c; +const g = id.g; +const H = id.H; +const D = id.D; +const Range = id.Range; +const Ends = id.Ends; +const hf = id.hf; +const kref = id.kref; +const dk = id.tf.idol.decoder_kernels; +const dc = id.tf.idol.decoder_config; + +/// The decoders' library: idol_decoder.metal through decoder_kernels.zig's launchers. +pub const Lib = id.LibOf(dk.Kernels, &dk.names, sources.idol_decoder); + +/// The steps' common parts: the outputs' ends, and its bounds and intervals from them. +fn Common(comptime T: type) type { + return struct { + pub fn bounds(t: *T, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *T, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + pub fn simulate(t: *T, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + t.each(&scratch, out, mag, acc, order); + } + }; +} + +/// A value rounded once to fp32. +fn f32r(v: f64) f64 { + return @as(f32, @floatCast(v)); +} + +/// The rows the decoders call a step at: one octree parent, the fewest past the short attention, 32k, 262k and 1M. +const served_rows = [_]usize{ 1024, 1, 9, 8192, 32768 }; + +/// Those rows at `n` outputs a row and `k` inputs. +fn servedAt(comptime n: usize, comptime k: usize) [served_rows.len]c.Shape { + var out: [served_rows.len]c.Shape = undefined; + for (&out, served_rows) |*s, r| s.* = .{ .rows = r, .n = n, .k = k }; + return out; +} + +const embed_served = servedAt(H, 3); +const half_served = servedAt(H, 1); +const heads_served = servedAt(16, D); + +// --- dec_embed: h = fp32(fma(x2, w2, fma(x1, w1, x0 w0)) + b) + pe, pe = sin or cos of fp32(fp32(x f) pi) ------ + +/// The three-input FMA chain, then the bias. +fn embedDepth(_: c.Shape) usize { + return 4; +} + +pub const Embed = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *Lib, + name: []const u8, + rows: usize = 0, + stress: bool = false, + p: []f32 = &.{}, + /// w [1024, 3], b [1024], then the frequencies [170] + wbf: []f32 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const pi32: f64 = @as(f32, std.math.pi); + + /// Rows: points; n: the stream's 1,024; k: the three coordinates. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = embedDepth, .ops = 2 + id.trig_ulps, .out = .f32, .served = &embed_served }; + } + + pub fn deinit(t: *Embed) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Embed) void { + t.a.free(t.p); + t.a.free(t.wbf); + t.p = &.{}; + t.wbf = &.{}; + } + pub fn kernel(t: *Embed, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Embed); + pub const bounds = Common(Embed).bounds; + pub const intervals = Common(Embed).intervals; + pub const simulate = Common(Embed).simulate; + + pub fn prepare(t: *Embed, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H or shape.k != 3) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + t.stress = fill == .stress; + t.p = try t.a.alloc(f32, t.rows * 3); + t.wbf = try t.a.alloc(f32, H * 4 + dc.pe_freqs); + const w = t.wbf[0 .. H * 3]; + const b = t.wbf[H * 3 .. H * 4]; + const f = t.wbf[H * 4 ..][0..dc.pe_freqs]; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // x0 w0 and x1 w1 near +-2^16 (past fp16's range) cancelling to a few hundred; no embedding (frequencies 0) + for (0..t.rows) |r| { + t.p[r * 3] = id.int(rng, 200, 300); + t.p[r * 3 + 1] = t.p[r * 3]; + t.p[r * 3 + 2] = id.int(rng, 1, 8); + } + for (0..H) |o| { + w[o * 3] = id.sign(rng) * id.int(rng, 220, 300); + w[o * 3 + 1] = -w[o * 3] + id.int(rng, -1, 1); + w[o * 3 + 2] = id.int(rng, -8, 8); + } + for (b) |*v| v.* = id.int(rng, -8, 8); + @memset(f, 0); + } else { + id.randomF32(rng, t.p, 0, 1); + id.randomF32(rng, w, -1, 1); + id.randomF32(rng, b, -0.5, 0.5); + id.tf.idol.octree.pointFreqs(f); + } + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| t.p[i * 3 + i] = v; + b[9] = id.nan32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.p)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(b)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(f)); + t.out = try t.rig.output(t.rows * H * 4); + return true; + } + + pub fn run(t: *Embed) !void { + try t.lib.need("dec_embed"); + const p = t.rig.begin(); + t.lib.k.embed(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.ins[3]), kref(t.out), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Embed) usize { + return t.rows * H; + } + + fn each(t: *Embed, ends: *Ends, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) void { + const gd = c.gamma(3, c.unit(.f32)); + const w = t.wbf[0 .. H * 3]; + const b = t.wbf[H * 3 .. H * 4]; + const f = t.wbf[H * 4 ..][0..dc.pe_freqs]; + var terms: [3]f64 = undefined; + for (0..t.rows) |r| for (0..H) |o| { + const x = t.p[r * 3 ..][0..3]; + var s: f64 = 0; + var abs: f64 = 0; + for (&terms, x, w[o * 3 ..][0..3]) |*v, xv, wv| { + v.* = @as(f64, xv) * @as(f64, wv); + s += v.*; + abs += @abs(v.*); + } + var sr = if (t.stress) Range.of(s) else Range.of(s).widen(gd * abs); + if (acc) |y| { + s = contract.rows.sumIn(y, order, &terms); + sr = Range.of(s); + } + const bv: f64 = b[o]; + var pe: f64 = 0; + var per = Range.of(0); + if (o < 1020) { + const axis = o / 340; + const j = o % 340; + const xf = @as(f64, x[axis]) * @as(f64, f[j % 170]); + // torch's fp32 angle; the kernel's within an ulp of it whichever way its fast math associates + const ang = f32r(f32r(xf) * pi32); + const ar = Range.of(xf).fp32(0.5).mul(Range.of(pi32)).fp32(1); + if (j < 170) { + pe = @sin(ang); + per = id.sinOver(ar.lo, ar.hi).fp32(id.trig_ulps); + } else { + pe = @cos(ang); + per = id.cosOver(ar.lo, ar.hi).fp32(id.trig_ulps); + } + if (std.math.isNan(ang)) per = Range.of(ang); + } + const y = sr.add(Range.of(bv)).fp32(0.5).add(per).fp32(0.5); + ends.put(want, mag, r * H + o, s + bv + pe, y); + }; + } + + pub fn reference(t: *Embed, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + t.each(&t.ends, want, mag, null, .sequential); + } +}; + +// --- dec_half: o = fp16(y + b) by column ------------------------------------------------------------------------ + +pub const Half = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *Lib, + name: []const u8, + rows: usize = 0, + /// y [rows, 1024] then b [1024] + yb: []f32 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: points; n: the stream's 1,024 (the stem's only width). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .ops = 1, .out = .f16, .served = &half_served }; + } + + pub fn deinit(t: *Half) void { + t.a.free(t.yb); + t.yb = &.{}; + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Half, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Half); + pub const bounds = Common(Half).bounds; + pub const intervals = Common(Half).intervals; + + pub fn prepare(t: *Half, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != H) return false; + t.rig.reset(); + t.a.free(t.yb); + t.rows = shape.rows; + const n = t.rows * H; + t.yb = try t.a.alloc(f32, n + H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF32(rng, t.yb[0..n], -8, 8); + id.randomF32(rng, t.yb[n..], -2, 2); + if (fill == .nonfinite) { + for ([_]f32{ id.nan32, id.inf32, -id.inf32 }, 0..) |v, i| t.yb[(i * 13) % n] = v; + t.yb[n + 6] = id.inf32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.yb[0..n])); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.yb[n..])); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *Half) !void { + try t.lib.need("dec_half"); + const p = t.rig.begin(); + t.lib.k.toHalf(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.out), t.rows, H); + try t.rig.finish(p); + } + + pub fn count(t: *Half) usize { + return t.rows * H; + } + + pub fn reference(t: *Half, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const n = t.count(); + for (0..n) |i| { + const v = @as(f64, t.yb[i]) + @as(f64, t.yb[n + i % H]); + t.ends.put(want, mag, i, hf(v), Range.of(v).fp32(0.5).fp16()); + } + } +}; + +// --- dec_kv: kv = fp16(ctx w^T + b), the latent's 16 channels in fp32 FMA ----------------------------------------- + +const C = dc.context_rows; +const K = dc.context_channels; + +/// The 16 inputs' FMA chain, then the bias. +fn kvDepth(_: c.Shape) usize { + return K + 1; +} + +pub const Kv = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *Lib, + name: []const u8, + stress: bool = false, + /// ctx [8192, 16], w [2048, 16], b [2048] + v: []f16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const N = 2 * H; + + /// Rows: the latent's 8,192, which the launcher fixes (a gate's fewer rows run all of them); k: its 16 channels. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = kvDepth, .ops = 1, .out = .f16, .served = &.{.{ .rows = C, .n = N, .k = K }} }; + } + + pub fn deinit(t: *Kv) void { + t.a.free(t.v); + t.v = &.{}; + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Kv, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Kv); + pub const bounds = Common(Kv).bounds; + pub const intervals = Common(Kv).intervals; + pub const simulate = Common(Kv).simulate; + + pub fn prepare(t: *Kv, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != N or shape.k != K) return false; + t.rig.reset(); + t.a.free(t.v); + t.stress = fill == .stress; + t.v = try t.a.alloc(f16, C * K + N * K + N); + const ctx = t.v[0 .. C * K]; + const w = t.v[C * K ..][0 .. N * K]; + const b = t.v[C * K + N * K ..]; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // products to 2^16 whose halves cancel but for the first: partial sums near 2^19, exact in fp32 + for (0..C) |r| for (0..K / 2) |k| { + ctx[r * K + k] = @floatCast(id.int(rng, 1, 255)); + ctx[r * K + K / 2 + k] = @floatCast(@as(f32, ctx[r * K + k]) + @as(f32, if (k == 0) 1 else 0)); + }; + for (0..N) |o| for (0..K / 2) |k| { + w[o * K + k] = @floatCast(id.sign(rng) * id.int(rng, 1, 255)); + w[o * K + K / 2 + k] = -w[o * K + k]; + }; + for (b) |*x| x.* = @floatCast(id.int(rng, -8, 8)); + } else { + id.randomF16(rng, ctx, -3, 3); + id.randomF16(rng, w, -0.5, 0.5); + id.randomF16(rng, b, -0.5, 0.5); + } + if (fill == .nonfinite) { + for ([_]f16{ id.nan16, id.inf16, -id.inf16 }, 0..) |x, i| ctx[i * K + 3 * i] = x; + b[9] = id.nan16; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(ctx)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(b)); + t.out = try t.rig.output(C * N * 2); + return true; + } + + pub fn run(t: *Kv) !void { + try t.lib.need("dec_kv"); + const p = t.rig.begin(); + t.lib.k.kv(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out)); + try t.rig.finish(p); + } + + pub fn count(_: *Kv) usize { + return C * N; + } + + fn each(t: *Kv, ends: *Ends, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) void { + const gd = c.gamma(K, c.unit(.f32)); + const ctx = t.v[0 .. C * K]; + const w = t.v[C * K ..][0 .. N * K]; + const b = t.v[C * K + N * K ..]; + var terms: [K]f64 = undefined; + for (0..C) |r| for (0..N) |o| { + var s: f64 = 0; + var abs: f64 = 0; + for (&terms, ctx[r * K ..][0..K], w[o * K ..][0..K]) |*x, z, wv| { + x.* = @as(f64, z) * @as(f64, wv); + s += x.*; + abs += @abs(x.*); + } + var sr = if (t.stress) Range.of(s) else Range.of(s).widen(gd * abs); + if (acc) |y| { + s = contract.rows.sumIn(y, order, &terms); + sr = Range.of(s); + } + const bv: f64 = b[o]; + ends.put(want, mag, r * N + o, hf(s + bv), sr.add(Range.of(bv)).fp32(0.5).fp16()); + }; + } + + pub fn reference(t: *Kv, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + t.each(&t.ends, want, mag, null, .sequential); + } +}; + +// --- dec_heads: q and k heads F.normalize(x) times gamma and 8, rounded; v copied ----------------------------------- + +const heads_why = "an fp16 output follows the head's sum of squares as its square root: a bf16 sum is off a few ulps"; + +/// A lane's two squares, then the simdgroup tree. +fn headsDepth(_: c.Shape) usize { + return 1 + 5; +} + +pub const Heads = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *Lib, + name: []const u8, + rows: usize = 0, + stress: bool = false, + /// packed q, k, v rows [rows, 3072], as to_qkv writes them + qkv: []f16 = &.{}, + /// gq then gk [1024] each + gam: []f32 = &.{}, + ins: [3]rig_.Ref = undefined, + outs: [3]rig_.Ref = undefined, + ends: Ends, + + /// Rows: points; n: heads; k: the head width. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = headsDepth, + .ops = id.sqrt_ulps + id.div_ulps + 2, + .out = .f16, + .served = &heads_served, + .sums = .{ .bf16 = 0, .why = heads_why }, + }; + } + + pub fn deinit(t: *Heads) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Heads) void { + t.a.free(t.qkv); + t.a.free(t.gam); + t.qkv = &.{}; + t.gam = &.{}; + } + pub fn kernel(t: *Heads, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Heads); + pub const bounds = Common(Heads).bounds; + pub const intervals = Common(Heads).intervals; + + pub fn prepare(t: *Heads, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != 16 or shape.k != D) return false; + t.rig.reset(); + t.free(); + t.rows = shape.rows; + t.stress = fill == .stress; + const n = shape.rows; + t.qkv = try t.a.alloc(f16, n * 3 * H); + t.gam = try t.a.alloc(f32, 2 * H); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + id.randomF16(rng, t.qkv, -4, 4); + id.randomF32(rng, t.gam, 0.25, 2); + // integer heads, whose squares sum to ~2^20 exactly + if (t.stress) for (t.qkv) |*v| { + v.* = @floatCast(id.int(rng, -128, 128)); + }; + if (fill == .nonfinite) { + // q NaN, k +inf, v -inf, in turn; then a NaN gamma + const at = [_]usize{ 0 * 3 * H + 0 * D + 5, 1 * 3 * H + H + 2 * D + 9, 2 * 3 * H + 2 * H + 3 * D + 1 }; + for (at, [_]f16{ id.nan16, id.inf16, -id.inf16 }) |p, v| if (p < t.qkv.len) { + t.qkv[p] = v; + }; + t.gam[H + 4 * D + 6] = id.nan32; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.gam[0..H])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.gam[H..])); + for (&t.outs) |*o| o.* = try t.rig.output(n * H * 2); + return true; + } + + pub fn run(t: *Heads) !void { + try t.lib.need("dec_heads"); + const p = t.rig.begin(); + const k = &t.lib.k; + k.heads(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.outs[0]), t.rows, 3 * H, 0); + k.heads(p.e, kref(t.ins[0]), kref(t.ins[2]), kref(t.outs[1]), t.rows, 3 * H, H); + k.heads(p.e, kref(t.ins[0]), null, kref(t.outs[2]), t.rows, 3 * H, 2 * H); + try t.rig.finish(p); + } + + pub fn count(t: *Heads) usize { + return 3 * t.rows * H; + } + + /// The head's outputs from its sum of squares `ss` (range `ssr`). + fn head(t: *Heads, ends: *Ends, which: usize, r: usize, h: usize, ss: f64, ssr: Range, want: []f64, mag: []f64) void { + const e12: f64 = @as(f32, 1e-12); + // torch's clamp_min keeps a NaN norm (@max would drop it) + const norm = if (std.math.isNan(ss)) ss else @max(@sqrt(ss), e12); + const nr0 = ssr.up(id.sqrt).fp32(id.sqrt_ulps); + const nr = Range{ .lo = @max(nr0.lo, e12), .hi = @max(nr0.hi, e12) }; + const x = t.qkv[r * 3 * H + which * H + h * D ..][0..D]; + const gam = t.gam[which * H + h * D ..][0..D]; + const base = which * t.rows * H + (h * t.rows + r) * D; + for (0..D) |i| { + const v: f64 = x[i]; + const gv: f64 = gam[i]; + const rr = Range.of(v).mul(nr.recip()).fp32(id.div_ulps).mul(Range.of(gv)).fp32(0.5).mul(Range.of(8)).fp32(0.5).fp16(); + ends.put(want, mag, base + i, hf(v / norm * gv * 8), rr); + } + } + + /// v's heads, copied bit for bit. + fn values(t: *Heads, ends: *Ends, want: []f64, mag: []f64) void { + for (0..t.rows) |r| for (0..16) |h| for (0..D) |i| { + const f: f64 = t.qkv[r * 3 * H + 2 * H + h * D + i]; + ends.put(want, mag, 2 * t.rows * H + (h * t.rows + r) * D + i, f, Range.of(f)); + }; + } + + fn each(t: *Heads, ends: *Ends, want: []f64, mag: []f64, acc: ?c.Accum, order: g.Order) void { + const gd = c.gamma(@floatFromInt(headsDepth(.{ .rows = 1, .n = 1, .k = 1 })), c.unit(.f32)); + var terms: [D]f64 = undefined; + for (0..2) |which| for (0..t.rows) |r| for (0..16) |h| { + var ss: f64 = 0; + for (&terms, t.qkv[r * 3 * H + which * H + h * D ..][0..D]) |*s, x| { + s.* = @as(f64, x) * @as(f64, x); + ss += s.*; + } + var ssr = if (t.stress) Range.of(ss) else Range{ .lo = ss * (1 - gd), .hi = ss * (1 + gd) }; + if (acc) |y| { + ss = contract.rows.sumIn(y, order, &terms); + ssr = Range.of(ss); + } + t.head(ends, which, r, h, ss, ssr, want, mag); + }; + t.values(ends, want, mag); + } + + pub fn reference(t: *Heads, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + t.each(&t.ends, want, mag, null, .sequential); + } + + pub const simulate = Common(Heads).simulate; +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index ece8073ecc..a125f644bb 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -102,6 +102,9 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol decoder ")) { + text = sources.idol_decoder; + frags = true; } else if (std.mem.startsWith(u8, d.name, "idol ")) { text = sources.idol; frags = true; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 9799191182..1c7eb98291 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -429,6 +429,8 @@ pub fn main(init: std.process.Init) !u8 { defer moss_image.destroy(); const idol = try @import("metal_idol.zig").Suite.create(a, device, queue, &kernels, o.only); defer idol.destroy(); + const idol_decoder = try @import("metal_idol_decoder.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol_decoder.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { From 7824bf6dc04aadbbe7b0dd37a456affb9b21a9b0 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:08:46 +0800 Subject: [PATCH 10/21] changelog: #1773 for Idol's decoders --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 40ae4002d3..921e5727e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. -- Idol's octree and Gaussian decoders now run natively on M5 Metal, from the denoiser's latent to the 3DGS PLY and camera. On an M5 Ultra, 262,144 Gaussians decode in 0.343 s against the SDK's 0.739 s, and 1,048,576 in 1.963 s against 4.653 s, at a 2.43 GB peak footprint against 5.64 GB. Given the SDK's probabilities and draws, every octree level's counts match, and the PLY's header and rows equal the SDK's. Both decoders sit as close to an fp32 run as the SDK's own; `tf-idol-decode` runs those gates from a `tools/zig/idol_case.py decode` case. The decoders' nine kernels are registered with the kernel contract (after #1766, for #1757). Thanks to @jayleaton. +- Idol's octree and Gaussian decoders now run natively on M5 Metal, from the denoiser's latent to the 3DGS PLY and camera. On an M5 Ultra, 262,144 Gaussians decode in 0.343 s against the SDK's 0.739 s, and 1,048,576 in 1.963 s against 4.653 s, at a 2.43 GB peak footprint against 5.64 GB. Given the SDK's probabilities and draws, every octree level's counts match, and the PLY's header and rows equal the SDK's. Both decoders sit as close to an fp32 run as the SDK's own; `tf-idol-decode` runs those gates from a `tools/zig/idol_case.py decode` case. The decoders' nine kernels are registered with the kernel contract (#1773, after #1766, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. From 7747ed59dae2ac0d7559d45a12f738a0ecae65fe Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:10:09 +0800 Subject: [PATCH 11/21] idol: the VAE recipe names the tensor units' fp32 operand truncation (#1772) --- docs/recipes/idol.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index a449405e51..a3b4b734db 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -156,9 +156,9 @@ give the same weights. The loader takes either: it rounds an fp32 file to bf16 a tensors, `quant_conv` and the two batch-norm statistics, and refuses any other geometry. The SDK then runs the whole encoder in bf16. This port computes in fp32 on those bf16 weights: 3x3 convolutions by -Winograd F(2x2, 3x3) with `G g G^T` made in float64 and rounded once, the stride-2 downsamples and 1x1s as implicit GEMMs, -GroupNorm statistics in two passes (the mean, then the squared deviations), fp32 attention, `precise::exp`. None of it -uses the tensor units, whose fp32 GEMM is about fp16-accurate (above). +Winograd F(2x2, 3x3) with `G g G^T` made in float64 and rounded once, the stride-2 downsamples and 1x1s as implicit +GEMMs, GroupNorm statistics in two passes (the mean, then the squared deviations), fp32 attention, `precise::exp`. None +of it uses the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (above). ### Gates From b5a74eb1e6593652fb240d28f4816b30f3e3cabd Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:10:49 +0800 Subject: [PATCH 12/21] changelog: #1774 for Idol's VAE encoder --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 810a7ef6f2..2c5ff309ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. -- Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (after #1766, for #1757). Thanks to @jayleaton. +- Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (#1774, after #1766, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. From 4e5d028fdb6e86b764661110ca283ede95fd6541 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:06:26 +0800 Subject: [PATCH 13/21] idol: DINOv3 ViT-H+/16 image encoder on Metal, with tf-idol-dino and idol_case.py dino The denoiser's feature1 [4101, 1280] from the SDK's 1024 x 1024 image, as idol-sdk's ImageEncoder makes it: transformers' DINOv3ViTModel in bf16, its last_hidden_state, then F.layer_norm in fp32. Weights from the ungated timm/vit_huge_plus_patch16_dinov3.lvd1689m (DINOv3 License), rounded to bf16 at load. Patch conv as one GEMM, class token and four registers, 2D RoPE on the patches (transformers' fp32 periods, tables bit-equal to torch's), 32 blocks of 20 heads of 64 with LayerScale and SwiGLU, rounded where torch on MPS rounds (measured): one rounding a linear or norm, two for SiLU times up and LayerScale with the residual, three for the rotary. The attention rounds P to bf16 itself. idol_case.py dino pins idol-sdk's conditioning.py, transformers 5.22.1's DINOv3 files and the weights by sha256, and writes the SDK's feature1 on MPS, a float32 CPU ruler, every block's rows and each block in float32. With --official it proves timm's weights Meta's: 550 tensors byte for byte, feature1 bit for bit. On the avocado: feature1 worst row against float32 0.1353 (torch 0.1686), 1 - cos 5.37e-5 (5.57e-5); every block on torch's own input within the 2x gate; repeat byte-identical; 0.120 s an encode against the SDK's 0.261 s on MPS. --- tools/zig/idol_case.py | 18 ++ tools/zig/idol_dino_case.py | 256 ++++++++++++++++++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol_dino.metal | 313 +++++++++++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/families/idol/dino.zig | 17 ++ zig/src/families/idol/dino_config.zig | 46 ++++ zig/src/families/idol/dino_host.zig | 101 ++++++++ zig/src/families/idol/dino_kernels.zig | 127 +++++++++ zig/src/families/idol/dino_model.zig | 213 +++++++++++++++ zig/src/families/idol/dino_reference.zig | 68 +++++ zig/src/families/idol/dino_weights.zig | 196 ++++++++++++++ zig/src/families/idol/idol.zig | 5 +- zig/tests/idol_dino.zig | 188 ++++++++++++++ 14 files changed, 1550 insertions(+), 1 deletion(-) create mode 100644 tools/zig/idol_dino_case.py create mode 100644 zig/kernels/metal/idol_dino.metal create mode 100644 zig/src/families/idol/dino.zig create mode 100644 zig/src/families/idol/dino_config.zig create mode 100644 zig/src/families/idol/dino_host.zig create mode 100644 zig/src/families/idol/dino_kernels.zig create mode 100644 zig/src/families/idol/dino_model.zig create mode 100644 zig/src/families/idol/dino_reference.zig create mode 100644 zig/src/families/idol/dino_weights.zig create mode 100644 zig/tests/idol_dino.zig diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py index 83453a8e54..74b76887fd 100644 --- a/tools/zig/idol_case.py +++ b/tools/zig/idol_case.py @@ -14,8 +14,13 @@ case WEIGHTS SDK OUT --features F [F] noise, each step's input and both velocities, the final sample blocks WEIGHTS SDK CASE [--step S] every block's rows at step S, and each block in float32 on its input exact WEIGHTS SDK CASE [--steps N] both velocities in float32 on the case's first N inputs (the ruler) + dino DINO SDK OUT --image PNG [--official DIR] [--sdk-features F] + DINOv3 ViT-H+/16's feature1 for a 1024 x 1024 image as the SDK encodes it + (bf16), a float32 CPU ruler, every block's rows, and each block in float32 WEIGHTS is a speridlabs/idol snapshot; SDK an idol-sdk checkout at the pinned commit. F holds feature1 (DINOv3 rows [1, 4101, 1280]) and feature2 (FLUX.2 VAE rows [1, 4101, 128]), in one file or two, as the SDK's ImageEncoder returns. +DINO is timm/vit_huge_plus_patch16_dinov3.lvd1689m (ungated); --official DIR, Meta's gated +facebook/dinov3-vith16plus-pretrain-lvd1689m, checks that it is the same model, tensor by tensor and feature by feature. """ SDK_COMMIT = "4702a7f00850e34b6e83a2170cc123c3e9570744" @@ -219,6 +224,15 @@ def exact(args) -> int: def main(argv: list[str]) -> int: ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("dino") + p.add_argument("weights") + p.add_argument("sdk") + p.add_argument("out_dir") + p.add_argument("--image", required=True) + p.add_argument("--official") + p.add_argument("--sdk-features") + p.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + p.add_argument("--no-blocks", action="store_true") for name in ("case", "blocks", "exact"): p = sub.add_parser(name) p.add_argument("weights") @@ -236,6 +250,10 @@ def main(argv: list[str]) -> int: else: p.add_argument("--steps", type=int, default=50) args = ap.parse_args(argv) + if args.cmd == "dino": + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from idol_dino_case import dino + return dino(args, sdk_root=args.sdk) return {"case": case, "blocks": blocks, "exact": exact}[args.cmd](args) diff --git a/tools/zig/idol_dino_case.py b/tools/zig/idol_dino_case.py new file mode 100644 index 0000000000..62f0062340 --- /dev/null +++ b/tools/zig/idol_dino_case.py @@ -0,0 +1,256 @@ +"""idol_case.py dino: DINOv3 ViT-H+/16's feature1 as idol-sdk's ImageEncoder makes it, its ruler and its blocks.""" + +# timm's ungated weights in the transformers graph the SDK builds from Meta's gated ones; --official proves them one + +from __future__ import annotations + +import hashlib +import json +import re +import sys +import time +from pathlib import Path + +SDK_FILES = {"idol/conditioning.py": "8c9cc50d327df3a5269cda392ef46cf4ed9c1e603f42f46050beda83f355f5ab"} +TRANSFORMERS_VERSION = "5.22.1" +TRANSFORMERS_FILES = { + "models/dinov3_vit/modeling_dinov3_vit.py": "28c72d5adedc7398a7988d2568dce3b1e399f277e853836cb8bce8c9515be1e3", + "models/dinov3_vit/configuration_dinov3_vit.py": "9a13d3c9ea8020aaed7057db28a7ffe0ffd9bb094a6178261bcc45ed04e9bbfc", +} +TIMM_REVISION = "dd0addc09788111fa893d24799a44744ba022eee" +TIMM_SHA256 = "0d54131c846c91a78ec7641482f32e45b27cb4895f83653c038f305f18c7e61c" +OFFICIAL_REVISION = "c807c9eeea853df70aec4069e6f56b28ddc82acc" +OFFICIAL_SHA256 = "3e1d4d18b9bfa9f28fad8e9de6a783f1313532d3460efa4cd0b12521d81d1a4d" +# facebook/dinov3-vith16plus-pretrain-lvd1689m's config.json at OFFICIAL_REVISION, the graph the SDK builds +CONFIG = {"hidden_act": "silu", "hidden_size": 1280, "image_size": 224, "intermediate_size": 5120, "key_bias": False, + "layer_norm_eps": 1e-05, "layerscale_value": 1.0, "mlp_bias": True, "num_attention_heads": 20, + "num_channels": 3, "num_hidden_layers": 32, "num_register_tokens": 4, "patch_size": 16, + "pos_embed_jitter": None, "pos_embed_rescale": 2.0, "pos_embed_shift": None, "proj_bias": True, + "query_bias": True, "rope_theta": 100.0, "use_gated_mlp": True, "value_bias": True, + "attention_dropout": 0.0, "drop_path_rate": 0.0} + + +def sha256(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(16 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def pinned(root: Path, files: dict[str, str], what: str) -> None: + for name, want in files.items(): + digest = sha256(root / name) + if digest != want: + sys.exit(f"{root / name} has sha256 {digest}; this tool runs only {what}'s ({want})") + + +def official_name(k: str) -> str | None: + """A timm tensor's name in transformers' DINOv3ViTModel (checkpoint names, no live `model.` prefix).""" + direct = {"cls_token": "embeddings.cls_token", "reg_token": "embeddings.register_tokens", + "patch_embed.proj.weight": "embeddings.patch_embeddings.weight", + "patch_embed.proj.bias": "embeddings.patch_embeddings.bias", "norm.weight": "norm.weight", + "norm.bias": "norm.bias"} + if k in direct: + return direct[k] + m = re.fullmatch(r"blocks\.(\d+)\.(.+)", k) + if not m: + return None + i, s = m.groups() + subs = {"attn.proj.": "attention.o_proj.", "gamma_1": "layer_scale1.lambda1", "gamma_2": "layer_scale2.lambda1", + "mlp.fc1_g.": "mlp.gate_proj.", "mlp.fc1_x.": "mlp.up_proj.", "mlp.fc2.": "mlp.down_proj.", + "norm1.": "norm1.", "norm2.": "norm2."} + for a, b in subs.items(): + if s.startswith(a): + return f"layer.{i}.{b}{s[len(a):]}" + return None + + +def mapped(path: Path): + """(official name, fp32 tensor) for every timm tensor; a fused qkv weight as its q, k and v thirds.""" + from safetensors import safe_open + + with safe_open(str(path), framework="pt", device="cpu") as f: + for k in f.keys(): + x = f.get_tensor(k) + if k.endswith(".attn.qkv.weight"): + layer = k.split(".")[1] + for q, part in zip("qkv", x.chunk(3, dim=0)): + yield f"layer.{layer}.attention.{q}_proj.weight", part.contiguous() + continue + name = official_name(k) + if name is None: + sys.exit(f"{path} holds {k}, which no DINOv3ViTModel tensor takes") + yield name, x + + +def build(weights: Path, dtype): + """DINOv3ViTModel on the official config with timm's tensors; q and v biases and the mask token zero.""" + import torch + from transformers import DINOv3ViTConfig, DINOv3ViTModel + + model = DINOv3ViTModel(DINOv3ViTConfig(**CONFIG)) + model.config._attn_implementation = "sdpa" + sd = model.state_dict() + live = lambda k: k if k in sd else "model." + k + used = set() + with torch.no_grad(): + for name, x in mapped(weights / "model.safetensors"): + key = live(name) + if key not in sd or tuple(sd[key].shape) != tuple(x.shape): + sys.exit(f"timm's {name} {tuple(x.shape)} has no place in DINOv3ViTModel") + sd[key].copy_(x) + used.add(key) + rest = [k for k in sd if k not in used] + for k in rest: + if k != "embeddings.mask_token" and not k.endswith(("attention.q_proj.bias", "attention.v_proj.bias")): + sys.exit(f"DINOv3ViTModel's {k} has no timm tensor") + sd[k].zero_() + # as from_pretrained(dtype=...) does: the parameters in `dtype`, the generated RoPE inv_freq kept in float32 + inv_freq = model.rope_embeddings.inv_freq.clone() + model = model.to(dtype) + model.rope_embeddings.inv_freq = inv_freq + return model.eval().requires_grad_(False) + + +def identity(weights: Path, official: Path) -> dict: + """Every timm tensor against its official one, bitwise; the official q and v biases must be zeros.""" + import torch + from safetensors import safe_open + + digest = sha256(official / "model.safetensors") + if digest != OFFICIAL_SHA256: + sys.exit(f"{official}/model.safetensors has sha256 {digest}, not {OFFICIAL_REVISION[:12]}'s") + equal, compared, used = True, 0, set() + with safe_open(str(official / "model.safetensors"), framework="pt", device="cpu") as f: + names = set(f.keys()) + for name, x in mapped(weights / "model.safetensors"): + y = f.get_tensor(name) + same = x.dtype == y.dtype and x.shape == y.shape and torch.equal(x.view(torch.int32), y.view(torch.int32)) + equal &= same + compared += 1 + used.add(name) + zero = all(not f.get_tensor(n).any() for n in names if n.endswith(("q_proj.bias", "v_proj.bias"))) + left = sorted(names - used) + allowed = all(n == "embeddings.mask_token" or n.endswith(("q_proj.bias", "v_proj.bias")) for n in left) + return {"tensors_compared": compared, "bitwise_equal": bool(equal), "official_q_v_biases_zero": bool(zero), + "official_only": len(left), "official_only_allowed": bool(allowed), "official_sha256": digest} + + +def encoder(model, device: str): + """The SDK's ImageEncoder holding `model` as its dinov3, without loading the FLUX.2 VAE it never touches here.""" + import torch + from idol.conditioning import ImageEncoder + + enc = object.__new__(ImageEncoder) + enc.device = torch.device(device) + enc.dinov3 = model.to(enc.device) + return enc + + +def flush(device: str) -> None: + import torch + + if device == "mps": + torch.mps.synchronize() + + +def dino(args, sdk_root: str) -> int: + import numpy as np + import torch + import torch.nn.functional as F + import transformers + from PIL import Image + from safetensors.torch import load_file, save_file + + pinned(Path(sdk_root), SDK_FILES, "idol-sdk 4702a7f00850's") + if transformers.__version__ != TRANSFORMERS_VERSION: + sys.exit(f"transformers {transformers.__version__}: this tool runs only {TRANSFORMERS_VERSION}") + pinned(Path(transformers.__file__).parent, TRANSFORMERS_FILES, f"transformers {TRANSFORMERS_VERSION}") + sys.path.insert(0, sdk_root) + weights, out = Path(args.weights), Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + digest = sha256(weights / "model.safetensors") + if digest != TIMM_SHA256: + sys.exit(f"{weights}/model.safetensors has sha256 {digest}; not timm's {TIMM_REVISION[:12]}") + image = Image.open(args.image).convert("RGB") + if image.size != (1024, 1024): + sys.exit(f"{args.image} is {image.size}; the encoder takes the SDK's 1024 x 1024 preprocessed image") + # IdolPipeline.__call__'s pixels: [1, 3, 1024, 1024] in [0, 1] + pixels = torch.from_numpy(np.array(image)).permute(2, 0, 1).float().div(255.0)[None] + if args.device == "mps": + torch.mps.set_per_process_memory_fraction(0.2) + + ruler = build(weights, torch.float32) + model = build(weights, torch.bfloat16) + enc = encoder(model, args.device) + seen, rope = {}, {} + + def keep(i): + def hook(_m, a, kw): + seen[f"block.{i}"] = a[0][0].detach().cpu().contiguous() + if i == 0: + cos, sin = kw["position_embeddings"] + rope["cos"], rope["sin"] = cos.cpu().contiguous(), sin.cpu().contiguous() + return hook + + layers = enc.dinov3.model.layer + hooks = [m.register_forward_pre_hook(keep(i), with_kwargs=True) for i, m in enumerate(layers)] + hooks.append(enc.dinov3.norm.register_forward_pre_hook(lambda _m, a: seen.__setitem__("block.32", a[0][0].cpu()))) + hooks.append(enc.dinov3.norm.register_forward_hook(lambda _m, _a, o: seen.__setitem__("hidden", o[0].cpu()))) + times = [] + with torch.no_grad(): + for _ in range(2): + mark = time.perf_counter() + feature1 = enc._dinov3_tokens(pixels.to(args.device)) + flush(args.device) + times.append(time.perf_counter() - mark) + for h in hooks: + h.remove() + feature1 = feature1[0].cpu().contiguous() + enc.dinov3.to("cpu") + meta = {"sdk_encode_s": [round(t, 4) for t in times], "device": args.device, "torch": torch.__version__, + "transformers": transformers.__version__, "timm_revision": TIMM_REVISION, "timm_sha256": digest, + "image_sha256": sha256(Path(args.image)), "attention": enc.dinov3.config._attn_implementation} + # the ruler: the same graph in float32 on the CPU (the M5's fp32 GEMM on its tensor units is ~fp16-accurate) + mark = time.perf_counter() + f32 = encoder(ruler, "cpu") + feature1_f32 = f32._dinov3_tokens(pixels)[0].contiguous() + meta["ruler_s"] = round(time.perf_counter() - mark, 2) + tensors = {"pixels": pixels[0].contiguous(), "feature1": feature1, "feature1_f32": feature1_f32, + "hidden": seen["hidden"].contiguous(), "cos": rope["cos"], "sin": rope["sin"]} + exact = {} + if not args.no_blocks: + # each block in float32 on the bf16 rows torch fed it, with float32 cos and sin + norm = ((pixels - torch.tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1)) + / torch.tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1)) + exact["embed"] = ruler.embeddings(norm)[0].contiguous() + pos = ruler.rope_embeddings(norm) + for i, layer in enumerate(ruler.model.layer): + exact[f"block.{i + 1}"] = layer(seen[f"block.{i}"][None].float(), position_embeddings=pos)[0].contiguous() + last = ruler.norm(seen["block.32"][None].float()) + exact["feature1"] = F.layer_norm(last, last.shape[-1:])[0].contiguous() + save_file({k: v for k, v in seen.items() if k.startswith("block.")}, str(out / "dino_blocks.safetensors")) + save_file(exact, str(out / "dino_blocks_f32.safetensors")) + save_file(tensors, str(out / "dino.safetensors")) + d = (feature1.double() - feature1_f32.double()) + meta["sdk_vs_ruler"] = {"max_abs": float(d.abs().max()), "cosine": float(F.cosine_similarity( + feature1.double().flatten(), feature1_f32.double().flatten(), dim=0)), + "worst_row": float((d.norm(dim=-1) / feature1_f32.double().norm(dim=-1)).max())} + if args.sdk_features: + theirs = load_file(args.sdk_features)["feature1"][0] + meta["sdk_features"] = {"sha256": sha256(Path(args.sdk_features)), "bitwise_equal": bool(torch.equal(theirs, feature1)), + "max_abs": float((theirs - feature1).abs().max())} + if args.official: + meta["identity"] = identity(weights, Path(args.official)) + from transformers import AutoModel + + off = AutoModel.from_pretrained(args.official, dtype=torch.bfloat16).eval().requires_grad_(False) + with torch.no_grad(): + theirs = encoder(off, args.device)._dinov3_tokens(pixels.to(args.device))[0].cpu() + meta["identity"]["feature1_bitwise_equal"] = bool(torch.equal(theirs, feature1)) + meta["identity"]["feature1_max_abs"] = float((theirs - feature1).abs().max()) + (out / "dino.json").write_text(json.dumps(meta, indent=1) + "\n") + print(json.dumps(meta)) + ok = meta.get("identity", {}).get("feature1_bitwise_equal", True) and meta.get("identity", {}).get("bitwise_equal", True) + return 0 if ok else 1 diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index f08c980d2f..5b8ca05f11 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -318,6 +318,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" }, .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" }, .{ .name = "tf-idol", .path = "zig/tests/idol.zig", .about = "Idol's denoiser on a reference case: each step's two branches compared, then the guided schedule timed" }, + .{ .name = "tf-idol-dino", .path = "zig/tests/idol_dino.zig", .about = "Idol's DINOv3 image encoder on a reference case: feature1 and each block gated against torch's, then timed" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/kernels/metal/idol_dino.metal b/zig/kernels/metal/idol_dino.metal new file mode 100644 index 0000000000..600515ad0d --- /dev/null +++ b/zig/kernels/metal/idol_dino.metal @@ -0,0 +1,313 @@ +// Idol's DINOv3 ViT-H+/16 encoder after gemm_nax.metal: bf16 rows rounded where transformers' torch code rounds them. + +inline float dn_b(float x) { return float(bfloat(x)); } + +// One threadgroup's sum of every thread's `v` (at most 1,024 threads), returned to all of them. +inline float dn_sum(float v, threadgroup float* part, uint n, uint lane, uint sg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float s = lane < (n + 31) / 32 ? part[lane] : 0.0f; + s = simd_sum(s); + if (lane == 0) part[32] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float out = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return out; +} + +// ---- y [M, N] (rows ldd apart) = bf16(x w^T + b): tfp::gemm_nax's block over bf16, fp32 sums, one rounding ---- +template +inline void dn_linear(const device bfloat* A, const device bfloat* B, const device bfloat* bias, device bfloat* D, + const device int* P, uint3 tg, uint sg, uint lane) { + constexpr int BM = 64, BN = 128, BK = 256, WN = 4; + const int tx = int(tg.x), ty = int(tg.y); + if (tx >= P[6] || ty >= P[7]) return; + const int lda = P[3], ldb = P[4], ldd = P[5]; + const int row = ty * BM + 32 * int(sg / WN), col = tx * BN + 32 * int(sg % WN); + A += long(row) * lda; + B += long(col) * ldb; + D += long(row) * ldd + col; + const int rows = AM ? 32 : min(32, P[0] - row), cols = 32; + const short2 home = tfp::frag_home(ushort(lane)); + tfp::frag acc[2][2]; + if (AM || rows == 32) { + tfp::gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } else { + tfp::gemm_block(acc, A, B, lda, ldb, P[2], P[9], rows, cols, home); + } + TF_UNROLL + for (short j = 0; j < 2; j++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + const float b = float(bias[col + 16 * j + home.x + TF_COL(e)]); + TF_UNROLL + for (short i = 0; i < 2; i++) acc[i][j][e] = dn_b(acc[i][j][e] + b); + } + } + if (AM || rows == 32) { + tfp::gemm_store(acc, D, ldd, rows, cols, home); + } else if (rows > 0) { + tfp::gemm_store(acc, D, ldd, rows, cols, home); + } +} + +// N % 128 = 0 always (1,280, 3,840, 10,240); _1 when M % 64 = 0 (the 4,096 patches), _0 for the 4,101 rows. +#define DN_LINEAR(NAME, AM) \ + [[max_total_threads_per_threadgroup(256)]] [[kernel]] void NAME( \ + const device bfloat* A [[buffer(0)]], const device bfloat* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device bfloat* D [[buffer(3)]], const device bfloat* bias [[buffer(4)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + dn_linear(A, B, bias, D, P, tg, sg, lane); \ + } +DN_LINEAR(dn_linear_0, false) +DN_LINEAR(dn_linear_1, true) + +struct DnNorm { + uint width; + float eps; +}; + +// nn.LayerNorm of a bf16 row: mean and variance in fp32, then bf16(n w + b) with bf16 w, b. A threadgroup a row. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dn_norm( + const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], const device bfloat* b [[buffer(2)]], + device bfloat* y [[buffer(3)]], constant DnNorm& p [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device bfloat* xr = x + ulong(row) * p.width; + device bfloat* yr = y + ulong(row) * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = dn_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(dn_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) yr[i] = bfloat((float(xr[i]) - mean) * inv * float(w[i]) + float(b[i])); +} + +// The encoder's last norm (bf16, as dn_norm) into h, then ImageEncoder's F.layer_norm of it in fp32 (no affine) into y. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dn_final( + const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], const device bfloat* b [[buffer(2)]], + device bfloat* h [[buffer(3)]], device float* y [[buffer(4)]], constant DnNorm& p [[buffer(5)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint n [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + threadgroup float v[1280]; + const device bfloat* xr = x + ulong(row) * p.width; + float a = 0.0f; + for (uint i = t; i < p.width; i += n) a += float(xr[i]); + const float mean = dn_sum(a, part, n, lane, sg) / float(p.width); + float q = 0.0f; + for (uint i = t; i < p.width; i += n) q += (float(xr[i]) - mean) * (float(xr[i]) - mean); + const float inv = metal::precise::rsqrt(dn_sum(q, part, n, lane, sg) / float(p.width) + p.eps); + float a2 = 0.0f; + for (uint i = t; i < p.width; i += n) { + const bfloat o = bfloat((float(xr[i]) - mean) * inv * float(w[i]) + float(b[i])); + h[ulong(row) * p.width + i] = o; + v[i] = float(o); + a2 += float(o); + } + const float mean2 = dn_sum(a2, part, n, lane, sg) / float(p.width); + float q2 = 0.0f; + for (uint i = t; i < p.width; i += n) q2 += (v[i] - mean2) * (v[i] - mean2); + const float inv2 = metal::precise::rsqrt(dn_sum(q2, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) y[ulong(row) * p.width + i] = (v[i] - mean2) * inv2; +} + +struct DnHeads { + uint rows; + uint prefix; // the class token and registers: no rotary +}; + +// q, k, v [20, rows, 64] of qkv rows; past the prefix, q and k as bf16(bf16(x cos) + bf16(rotate_half(x) sin)). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dn_heads( + const device bfloat* qkv [[buffer(0)]], const device bfloat* cs [[buffer(1)]], const device bfloat* sn [[buffer(2)]], + device bfloat* Q [[buffer(3)]], device bfloat* K [[buffer(4)]], device bfloat* V [[buffer(5)]], + constant DnHeads& p [[buffer(6)]], uint2 at [[thread_position_in_grid]]) { + const uint r = at.y, h = at.x / 32, j = at.x % 32; + if (h >= 20) return; + const ulong from = ulong(r) * 3840 + h * 64 + j; + const ulong to = (ulong(h) * p.rows + r) * 64 + j; + V[to] = qkv[from + 2560]; + V[to + 32] = qkv[from + 2592]; + const float q0 = float(qkv[from]), q1 = float(qkv[from + 32]); + const float k0 = float(qkv[from + 1280]), k1 = float(qkv[from + 1312]); + if (r < p.prefix) { + Q[to] = bfloat(q0); + Q[to + 32] = bfloat(q1); + K[to] = bfloat(k0); + K[to + 32] = bfloat(k1); + return; + } + const ulong t = ulong(r - p.prefix) * 64 + j; + const float c0 = float(cs[t]), c1 = float(cs[t + 32]), s0 = float(sn[t]), s1 = float(sn[t + 32]); + Q[to] = bfloat(dn_b(q0 * c0) + dn_b(-q1 * s0)); + Q[to + 32] = bfloat(dn_b(q1 * c1) + dn_b(q0 * s1)); + K[to] = bfloat(dn_b(k0 * c0) + dn_b(-k1 * s0)); + K[to + 32] = bfloat(dn_b(k1 * c1) + dn_b(k0 * s1)); +} + +// x = bf16(x + bf16(y g)) by column: LayerScale (g, bf16) then the residual, in place. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dn_scale_add( + device bfloat* x [[buffer(0)]], const device bfloat* y [[buffer(1)]], const device bfloat* g [[buffer(2)]], + constant uint& width [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + x[i] = bfloat(float(x[i]) + dn_b(float(g[at.x]) * float(y[i]))); +} + +// h [rows, width] = bf16(bf16(silu(gate)) up) from gu rows [gate | up] of 2 * width: the gated MLP's middle. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void dn_swiglu( + const device bfloat* gu [[buffer(0)]], device bfloat* h [[buffer(1)]], constant uint& width [[buffer(2)]], + uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * 2 * width + at.x; + const float g = float(gu[i]); + h[ulong(at.y) * width + at.x] = bfloat(dn_b(g / (1.0f + metal::precise::exp(-g))) * float(gu[i + width])); +} + +// ---- Attention: tfp::attention_nax's online softmax at head width 64 over bf16 q, k, v, every row to every row ---- + +constant constexpr float kDnMasked = -3.402823466e+38f; +constant constexpr int kDnAnyCols = 1 << 30; + +// Fold one fragment's rows into r[2] (rows home.y, home.y + 8): pairs, then lanes 1 and 8 apart, then into r. +template +inline void dn_fold(thread const tfp::frag& f, thread float (&r)[2]) { + TF_UNROLL + for (short h = 0; h < 2; h++) { + const short b = 4 * h; + float t = MAX ? max(max(f[b], f[b + 1]), max(f[b + 2], f[b + 3])) : (f[b] + f[b + 1]) + (f[b + 2] + f[b + 3]); + const float u = simd_shuffle_xor(t, ushort(1)); + t = MAX ? max(t, u) : t + u; + const float w = simd_shuffle_xor(t, ushort(8)); + t = MAX ? max(t, w) : t + w; + r[h] = MAX ? max(r[h], t) : r[h] + t; + } +} + +// One 64-query block of one head over 32-key blocks, both with tails; P and F as Kernels.attention sets them. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void dn_attention( + const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], const device bfloat* V [[buffer(2)]], + const device int32_t* P [[buffer(3)]], const device float* F [[buffer(4)]], device bfloat* O [[buffer(5)]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + constexpr int BQ = 64, BK = 32, D = 64; + const int h = int(tg.y), lq = P[1], lk = P[2]; + Q += (long(h) * lq + long(tg.x) * BQ) * D; + K += long(h) * lk * D; + V += long(h) * lk * D; + const int os = P[0] * D; + O += long(tg.x) * BQ * os + h * D; + const float scale2 = F[0] * 1.44269504089f; + const short tm = 16 * short(sg); + Q += tm * D; + const short2 home = tfp::frag_home(ushort(lane)); + + tfp::frag acc[4]; + TF_UNROLL + for (short i = 0; i < 4; i++) acc[i] = tfp::frag(0); + float top[2] = {kDnMasked, kDnMasked}, total[2] = {0.0f, 0.0f}; + const bool last_q = int(tg.x) == P[4]; + const short q_rows = short(P[6] - tm), k_rows = short(P[7]); + + for (int kb = 0; kb < P[3]; kb++) { + const bool last_k = kb == P[5]; + tfp::frag s[2] = {tfp::frag(0), tfp::frag(0)}; + TF_UNROLL + for (short d = 0; d < 4; d++) { + tfp::frag q, k0, k1; + if (last_q) { + tfp::frag_get_in(q, Q, D, 0, 16 * d, home, q_rows, kDnAnyCols); + } else { + tfp::frag_get(q, Q, D, 0, 16 * d, home); + } + if (last_k) { + tfp::frag_get_t_in(k0, K, D, 0, 16 * d, home, k_rows, kDnAnyCols); + tfp::frag_get_t_in(k1, K, D, 16, 16 * d, home, k_rows, kDnAnyCols); + } else { + tfp::frag_get_t(k0, K, D, 0, 16 * d, home); + tfp::frag_get_t(k1, K, D, 16, 16 * d, home); + } + tfp::mma_16x32(s[0], s[1], q, k0, k1); + } + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] *= scale2; + } + if (last_k) { + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kDnMasked; + } + } + float top_new[2] = {top[0], top[1]}; + dn_fold(s[0], top_new); + dn_fold(s[1], top_new); + tfp::frag pb[2]; + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + // P enters P V as bf16, V's type: the denominator sums those same rounded values + for (short e = 0; e < 8; e++) { + pb[f][e] = bfloat(fast::exp2(s[f][e] - top_new[e >> 2])); + s[f][e] = float(pb[f][e]); + } + } + float scale_old[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) { + scale_old[r] = fast::exp2(top[r] - top_new[r]); + top[r] = top_new[r]; + total[r] = total[r] * scale_old[r]; + } + dn_fold(s[0], total); + dn_fold(s[1], total); + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * scale_old[e >> 2]; + } + simdgroup_barrier(mem_flags::mem_none); + TF_UNROLL + for (short d = 0; d < 4; d += 2) { + TF_UNROLL + for (short k = 0; k < 2; k++) { + tfp::frag v0, v1; + if (last_k) { + tfp::frag_get_in(v0, V, D, 16 * k, 16 * d, home, k_rows, kDnAnyCols); + tfp::frag_get_in(v1, V, D, 16 * k, 16 * d + 16, home, k_rows, kDnAnyCols); + } else { + tfp::frag_get(v0, V, D, 16 * k, 16 * d, home); + tfp::frag_get(v1, V, D, 16 * k, 16 * d + 16, home); + } + tfp::mma_16x32(acc[d], acc[d + 1], pb[k], v0, v1); + } + } + K += BK * D; + V += BK * D; + } + + threadgroup_barrier(mem_flags::mem_none); + float inv[2]; + TF_UNROLL + for (short r = 0; r < 2; r++) inv[r] = 1.f / total[r]; + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * inv[e >> 2]; + } + O += tm * os; + if (last_q) { + if (q_rows <= 0) return; + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put_in(acc[i], O, os, 0, 16 * i, home, q_rows, kDnAnyCols); + } else { + TF_UNROLL + for (short i = 0; i < 4; i++) tfp::frag_put(acc[i], O, os, 0, 16 * i, home); + } +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index ef3221a88b..43b3ceb71b 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -135,6 +135,8 @@ pub const qwen_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile(" pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal"); /// Idol's denoiser: the prefill GEMM template, then its own steps and head-64 attention. pub const idol = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol.metal") ++ @embedFile("idol_attention.metal"); +/// Idol's DINOv3 ViT-H+/16 image encoder: the prefill GEMM template, then its own bf16 steps and attention. +pub const idol_dino = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol_dino.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/families/idol/dino.zig b/zig/src/families/idol/dino.zig new file mode 100644 index 0000000000..ec297ede10 --- /dev/null +++ b/zig/src/families/idol/dino.zig @@ -0,0 +1,17 @@ +//! Idol's DINOv3 ViT-H+/16 image encoder on Metal: an image to the denoiser's feature1 [4101, 1280]. +pub const config = @import("dino_config.zig"); +pub const host = @import("dino_host.zig"); +pub const weights = @import("dino_weights.zig"); +pub const kernels = @import("dino_kernels.zig"); +pub const model = @import("dino_model.zig"); +pub const reference = @import("dino_reference.zig"); +pub const Encoder = model.Encoder; + +test { + _ = config; + _ = host; + _ = weights; + _ = kernels; + _ = model; + _ = reference; +} diff --git a/zig/src/families/idol/dino_config.zig b/zig/src/families/idol/dino_config.zig new file mode 100644 index 0000000000..99f45a1e99 --- /dev/null +++ b/zig/src/families/idol/dino_config.zig @@ -0,0 +1,46 @@ +//! DINOv3 ViT-H+/16's geometry at Idol's 1024 x 1024 input, and the check that a timm snapshot names that model. +const std = @import("std"); + +pub const width = 1280; +pub const heads = 20; +pub const head_dim = width / heads; +/// SwiGLU's hidden width: gate and up projections of 5,120 each. +pub const ffn = 5120; +pub const layers = 32; +pub const patch = 16; +pub const image = 1024; +pub const grid = image / patch; +pub const patches = grid * grid; +/// The class token and four registers, ahead of the patches: no rotary. +pub const prefix = 5; +pub const rows = prefix + patches; +/// Every nn.LayerNorm of the model and ImageEncoder's last F.layer_norm. +pub const eps: f32 = 1e-5; +/// RoPE's base: periods 100^(4i/64) for i < 16, on each axis. +pub const rope_theta: f64 = 100; +pub const rope_freqs = head_dim / 4; +/// One patch's pixels, channel then row then column: the conv weight [1280, 3, 16, 16] flattened. +pub const patch_k = 3 * patch * patch; +/// timm's tensors: six outside the blocks (patch conv and bias, class token, registers, last norm), fifteen a block. +pub const tensors = 6 + 15 * layers; + +/// Accepts timm's vit_huge_plus_patch16_dinov3 snapshot config (1,280 features); anything else is refused. +pub fn check(gpa: std.mem.Allocator, text: []const u8) !void { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadConfig; + const arch = doc.value.object.get("architecture") orelse return error.BadConfig; + if (arch != .string or !std.mem.eql(u8, arch.string, "vit_huge_plus_patch16_dinov3")) return error.NotDinov3; + const features = doc.value.object.get("num_features") orelse return error.BadConfig; + if (features != .integer or features.integer != width) return error.UnsupportedGeometry; +} + +test "timm's DINOv3 H+/16 config passes; another architecture or width does not" { + const a = std.testing.allocator; + try check(a, "{\"architecture\": \"vit_huge_plus_patch16_dinov3\", \"num_classes\": 0, \"num_features\": 1280}"); + try std.testing.expectError(error.NotDinov3, check(a, "{\"architecture\": \"vit_large_patch16_dinov3\", \"num_features\": 1024}")); + try std.testing.expectError(error.UnsupportedGeometry, check(a, "{\"architecture\": \"vit_huge_plus_patch16_dinov3\", \"num_features\": 1024}")); + try std.testing.expectError(error.BadConfig, check(a, "[]")); + try std.testing.expectEqual(@as(usize, 4101), rows); + try std.testing.expectEqual(@as(usize, 486), tensors); +} diff --git a/zig/src/families/idol/dino_host.zig b/zig/src/families/idol/dino_host.zig new file mode 100644 index 0000000000..d4b4f08dfc --- /dev/null +++ b/zig/src/families/idol/dino_host.zig @@ -0,0 +1,101 @@ +//! The host's part of an Idol DINOv3 encode: the RoPE tables and the normalized patches, rounded as torch rounds them. +const std = @import("std"); +const cfg = @import("dino_config.zig"); + +/// fp32 to bf16 bits, ties to even, as torch's .to(bfloat16) rounds (NaN stays NaN). +pub fn bits(v: f32) u16 { + const u: u32 = @bitCast(v); + if (std.math.isNan(v)) return @truncate((u >> 16) | 0x40); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// The fp32 value of bf16 bits. +pub fn value(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// ImageEncoder's ImageNet mean and std, fp32 as torch.tensor makes them. +const mean = [3]f32{ 0.485, 0.456, 0.406 }; +const std_ = [3]f32{ 0.229, 0.224, 0.225 }; + +/// transformers' inv_freq: 1 / 100 ** arange(0, 1, 4 / 64) in fp32 (generated, not the checkpoint's bf16 periods). +pub fn invFreq() [cfg.rope_freqs]f32 { + var out: [cfg.rope_freqs]f32 = undefined; + for (&out, 0..) |*o, i| { + const e = @as(f64, @floatFromInt(i)) * 4.0 / cfg.head_dim; + const p: f32 = @floatCast(std.math.pow(f64, cfg.rope_theta, e)); + o.* = 1.0 / p; + } + return out; +} + +/// cos and sin [4096, 64] bf16 bits for the patches: angles 2 pi y f, 2 pi x f (fp32, as torch's ops run), tiled twice. +pub fn rope(cos: []u16, sin: []u16) void { + std.debug.assert(cos.len == cfg.patches * cfg.head_dim and sin.len == cos.len); + const inv = invFreq(); + const two_pi: f32 = @floatCast(2.0 * std.math.pi); + for (0..cfg.grid) |h| for (0..cfg.grid) |w| { + const p = h * cfg.grid + w; + // patch centers in [-1, 1]: 2 ((i + 0.5) / 64) - 1, exact in fp32 + const at = [2]f32{ 2.0 * ((@as(f32, @floatFromInt(h)) + 0.5) / cfg.grid) - 1.0, 2.0 * ((@as(f32, @floatFromInt(w)) + 0.5) / cfg.grid) - 1.0 }; + for (0..cfg.head_dim) |j| { + const axis = (j % 32) / cfg.rope_freqs; + const a: f32 = (two_pi * at[axis]) * inv[j % cfg.rope_freqs]; + const x: f64 = a; + cos[p * cfg.head_dim + j] = bits(@floatCast(@cos(x))); + sin[p * cfg.head_dim + j] = bits(@floatCast(@sin(x))); + } + }; +} + +/// The patch embedding's input [4096, 768] bf16 bits from pixels [3, 1024, 1024] in [0, 1]: (x - mean) / std in fp32. +pub fn patchRows(pixels: []const f32, out: []u16) void { + std.debug.assert(pixels.len == 3 * cfg.image * cfg.image and out.len == cfg.patches * cfg.patch_k); + const P = cfg.patch; + for (0..cfg.grid) |ph| for (0..cfg.grid) |pw| { + const row = out[(ph * cfg.grid + pw) * cfg.patch_k ..][0..cfg.patch_k]; + for (0..3) |c| for (0..P) |kh| { + const src = pixels[(c * cfg.image + ph * P + kh) * cfg.image + pw * P ..][0..P]; + for (src, 0..) |v, kw| row[(c * P + kh) * P + kw] = bits((v - mean[c]) / std_[c]); + }; + }; +} + +test "the RoPE tables: the first patch's angles, the axes' halves and the tiling" { + const a = std.testing.allocator; + const cos = try a.alloc(u16, cfg.patches * cfg.head_dim); + defer a.free(cos); + const sin = try a.alloc(u16, cfg.patches * cfg.head_dim); + defer a.free(sin); + rope(cos, sin); + const inv = invFreq(); + try std.testing.expectEqual(@as(f32, 1), inv[0]); + try std.testing.expectApproxEqRel(@as(f32, 1.0 / 74.989421), inv[15], 1e-6); + // patch (0, 0) sits at (-63/64, -63/64): its y and x halves agree, and the second 32 repeat the first + const y0: f32 = @floatCast(@cos(@as(f64, (@as(f32, @floatCast(2.0 * std.math.pi)) * (-63.0 / 64.0)) * inv[0]))); + try std.testing.expectEqual(bits(y0), cos[0]); + for (0..32) |j| try std.testing.expectEqual(cos[j], cos[32 + j]); + for (0..16) |j| try std.testing.expectEqual(sin[j], sin[16 + j]); + // patch (0, 1): x moves, y does not + try std.testing.expectEqual(cos[0], cos[cfg.head_dim]); + try std.testing.expect(cos[16] != cos[cfg.head_dim + 16]); +} + +test "bf16 rounding is to nearest, ties to even" { + try std.testing.expectEqual(@as(u16, 0x3f80), bits(1.0 + 0x1p-8)); + try std.testing.expectEqual(@as(u16, 0x3f82), bits(1.0 + 3 * 0x1p-8)); + try std.testing.expectEqual(@as(f32, -2.5), value(bits(-2.5))); +} + +test "a patch row takes its 16 x 16 pixels channel by channel, normalized" { + const a = std.testing.allocator; + const px = try a.alloc(f32, 3 * cfg.image * cfg.image); + defer a.free(px); + for (px, 0..) |*v, i| v.* = @as(f32, @floatFromInt(i % 251)) / 255.0; + const out = try a.alloc(u16, cfg.patches * cfg.patch_k); + defer a.free(out); + patchRows(px, out); + // patch (2, 5), channel 1, row 3, column 7 + const src = px[(1 * cfg.image + 2 * 16 + 3) * cfg.image + 5 * 16 + 7]; + try std.testing.expectEqual(bits((src - 0.456) / 0.224), out[(2 * cfg.grid + 5) * cfg.patch_k + (1 * 16 + 3) * 16 + 7]); +} diff --git a/zig/src/families/idol/dino_kernels.zig b/zig/src/families/idol/dino_kernels.zig new file mode 100644 index 0000000000..5567f709a1 --- /dev/null +++ b/zig/src/families/idol/dino_kernels.zig @@ -0,0 +1,127 @@ +//! Idol's DINOv3 pipelines (idol_dino.metal): bf16 projections, LayerNorms, RoPE heads, attention, the steps between. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const cfg = @import("dino_config.zig"); +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// Each pipeline by its function. +pub const names = [_][]const u8{ "dn_linear_0", "dn_linear_1", "dn_norm", "dn_final", "dn_heads", "dn_scale_add", "dn_swiglu", "dn_attention" }; + +pub const Kernels = struct { + own: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device) !Kernels { + return loadSource(gpa, device, sources.idol_dino); + } + + /// The pipelines from `text` (the contract's controls pass an edited source). + pub fn loadSource(gpa: std.mem.Allocator, device: mtl.Device, text: []const u8) !Kernels { + var k: Kernels = undefined; + const own = try frags.source(device, gpa, text); + defer gpa.free(own); + k.own = try mtl.Library.fromSource(device, own, mtl.CompileOptions.mlx()); + errdefer k.own.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.own, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.own.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } + + /// y [M, N] (rows `ldd` apart) = bf16(x [M, K] w^T + b), w as stored [N, K]; N % 128 = 0, K % 32 = 0. + pub fn linear(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, M: usize, N: usize, K: usize, ldd: usize) void { + std.debug.assert(N % 128 == 0 and K % 32 == 0); + const tn = N / 128; + const tm = (M + 63) / 64; + e.setPipeline(if (M % 64 == 0) k.get("dn_linear_1") else k.get("dn_linear_0")); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, ldd, tn, tm, 0, K / 256 }; + for (vals, 0..) |v, i| p[i] = @intCast(v); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(y.buf, y.off, 3); + e.setBuffer(b.buf, b.off, 4); + e.dispatchGroups(mtl.Size.of(tn, tm, 1), mtl.Size.of(256, 1, 1)); + } + + /// y = nn.LayerNorm of x's bf16 rows (fp32 statistics), bf16(n w + b) with bf16 w, b. + pub fn norm(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, rows: usize) void { + const params = extern struct { width: u32, eps: f32 }{ .width = cfg.width, .eps = cfg.eps }; + bind(e, k.get("dn_norm"), &.{ x, w, b, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// h = the last norm of x's rows (bf16), y = F.layer_norm of h in fp32, no affine: ImageEncoder's feature1. + pub fn final(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, h: Ref, y: Ref, rows: usize) void { + const params = extern struct { width: u32, eps: f32 }{ .width = cfg.width, .eps = cfg.eps }; + bind(e, k.get("dn_final"), &.{ x, w, b, h, y }, params); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// q, k, v heads [20, rows, 64] from qkv rows [rows, 3840]; past the prefix, q and k turned by cos, sin. + pub fn heads(k: *const Kernels, e: mtl.ComputeEncoder, qkv: Ref, cos: Ref, sin: Ref, q: Ref, kk: Ref, v: Ref, rows: usize) void { + const params = extern struct { rows: u32, prefix: u32 }{ .rows = @intCast(rows), .prefix = cfg.prefix }; + bind(e, k.get("dn_heads"), &.{ qkv, cos, sin, q, kk, v }, params); + e.dispatchThreads(mtl.Size.of(cfg.heads * 32, rows, 1), mtl.Size.of(128, 1, 1)); + } + + /// x [rows, 1280] = bf16(x + bf16(y g)), in place: LayerScale then the residual. + pub fn scaleAdd(k: *const Kernels, e: mtl.ComputeEncoder, x: Ref, y: Ref, g: Ref, rows: usize) void { + bind(e, k.get("dn_scale_add"), &.{ x, y, g }, @as(u32, cfg.width)); + e.dispatchThreads(mtl.Size.of(cfg.width, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// h [rows, 5120] = bf16(bf16(silu(gate)) up) from gu rows [gate | up]. + pub fn swiglu(k: *const Kernels, e: mtl.ComputeEncoder, gu: Ref, h: Ref, rows: usize) void { + bind(e, k.get("dn_swiglu"), &.{ gu, h }, @as(u32, cfg.ffn)); + e.dispatchThreads(mtl.Size.of(cfg.ffn, rows, 1), mtl.Size.of(256, 1, 1)); + } + + /// o rows [lq, heads * 64] = softmax(q k^T / 8) v for q [heads, lq, 64] and k, v [heads, lk, 64], bf16. + pub fn attention(k: *const Kernels, e: mtl.ComputeEncoder, q: Ref, kb: Ref, vb: Ref, o: Ref, heads_: usize, lq: usize, lk: usize) void { + const nq = (lq + 63) / 64; + const nk = (lk + 31) / 32; + e.setPipeline(k.get("dn_attention")); + for ([_]Ref{ q, kb, vb }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + // heads, q rows, k rows, k blocks, last q block, last k block, rows in the last q block and in the last k block + var params: [8]i32 = .{ @intCast(heads_), @intCast(lq), @intCast(lk), @intCast(nk), @intCast(nq - 1), @intCast(nk - 1), @intCast(lq - (nq - 1) * 64), @intCast(lk - (nk - 1) * 32) }; + e.setBytes(std.mem.asBytes(¶ms), 3); + var f: [4]f32 = .{ 1.0 / @sqrt(@as(f32, cfg.head_dim)), 0, 0, 0 }; + e.setBytes(std.mem.asBytes(&f), 4); + e.setBuffer(o.buf, o.off, 5); + e.dispatchGroups(mtl.Size.of(nq, heads_, 1), mtl.Size.of(32, 4, 1)); + } +}; + +/// One kernel's buffers in order, then its parameters at the next index. +pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); +} + +test "every Idol DINOv3 pipeline compiles and loads on this GPU" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + if (!device.tensorUnits()) return error.SkipZigTest; + var k = try Kernels.load(std.testing.allocator, device); + k.deinit(); +} diff --git a/zig/src/families/idol/dino_model.zig b/zig/src/families/idol/dino_model.zig new file mode 100644 index 0000000000..1c779d98f0 --- /dev/null +++ b/zig/src/families/idol/dino_model.zig @@ -0,0 +1,213 @@ +//! Idol's DINOv3 ViT-H+/16 on Metal: a 1024 x 1024 image to feature1 [4101, 1280] as idol-sdk's ImageEncoder makes it. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("dino_config.zig"); +const host = @import("dino_host.zig"); +const wts = @import("dino_weights.zig"); +const ks = @import("dino_kernels.zig"); +const Ref = ks.Ref; + +const W = cfg.width; +const R = cfg.rows; + +/// One encode's seconds: wall clock, on the GPU, and the host's part before it. +pub const Timing = struct { wall: f64, gpu: f64, host: f64 }; + +/// The scratch buffer's parts in bytes, in the order of Encoder's views. +const parts = [_]usize{ + R * W * 2, R * W * 2, R * 3 * W * 2, R * W * 2, R * W * 2, R * W * 2, R * W * 2, R * W * 2, + R * 2 * cfg.ffn * 2, R * cfg.ffn * 2, cfg.patches * cfg.patch_k * 2, cfg.patches * cfg.head_dim * 2, cfg.patches * cfg.head_dim * 2, R * W * 2, R * W * 4, +}; + +pub fn scratchBytes() usize { + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n, 256); + return total; +} + +/// transformers' DINOv3ViTModel in bf16, its last_hidden_state, then ImageEncoder's F.layer_norm in fp32. +pub const Encoder = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + scratch: mtl.Buffer, + // bf16 rows: the stream, a norm's output, packed qkv, heads [20, rows, 64], attention rows, a projection's output + x: Ref, + xn: Ref, + qkv: Ref, + q: Ref, + k_: Ref, + v: Ref, + att: Ref, + y: Ref, + /// the MLP's gate | up rows, then its SwiGLU rows + gu: Ref, + h: Ref, + /// the patch embedding's input [4096, 768], and RoPE's cos and sin [4096, 64], bf16 + pix: Ref, + cos: Ref, + sin: Ref, + /// last_hidden_state (bf16) and feature1 (fp32) + hidden: Ref, + out: Ref, + + /// The encoder of timm snapshot folder `dir` (config.json, model.safetensors). + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !*Encoder { + // qualified on the M5's tensor units alone, as the denoiser is + if (!device.tensorUnits()) { + std.log.err("idol dino: {s} has no tensor units; this encoder runs on the M5's alone, the only chip it was checked on", .{device.name()}); + return error.TensorUnitsRequired; + } + const t = try gpa.create(Encoder); + errdefer gpa.destroy(t); + t.* = undefined; + t.gpa = gpa; + t.device = device; + t.queue = try device.queue(); + errdefer t.queue.deinit(); + t.k = try ks.Kernels.load(gpa, device); + errdefer t.k.deinit(); + t.w = try wts.load(gpa, io, device, dir); + errdefer t.w.deinit(); + t.scratch = try device.buffer(scratchBytes(), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var at: usize = 0; + const views = [_]*Ref{ &t.x, &t.xn, &t.qkv, &t.q, &t.k_, &t.v, &t.att, &t.y, &t.gu, &t.h, &t.pix, &t.cos, &t.sin, &t.hidden, &t.out }; + for (views, parts) |view, n| { + view.* = .{ .buf = t.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + host.rope(halves(t.cos, cfg.patches * cfg.head_dim), halves(t.sin, cfg.patches * cfg.head_dim)); + return t; + } + + pub fn deinit(t: *Encoder) void { + t.scratch.deinit(); + t.w.deinit(); + t.k.deinit(); + t.queue.deinit(); + t.gpa.destroy(t); + } + + fn halves(r: Ref, n: usize) []u16 { + return @as([*]u16, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; + } + + /// The patches of `pixels` [3, 1024, 1024] in [0, 1] and the prefix rows, ready for `embed`. + fn stage(t: *Encoder, pixels: []const f32) !void { + if (pixels.len != 3 * cfg.image * cfg.image) return error.ImageSize; + host.patchRows(pixels, halves(t.pix, cfg.patches * cfg.patch_k)); + } + + /// feature1 [4101, 1280] fp32 of an image's `pixels` [3, 1024, 1024] in [0, 1] (IdolPipeline's tensor). + pub fn encode(t: *Encoder, pixels: []const f32, feature1: []f32) !Timing { + if (feature1.len != R * W) return error.FeatureRows; + const started = mtl.clock.seconds(); + try t.stage(pixels); + const spent = mtl.clock.seconds() - started; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.embed(e); + for (t.w.blocks) |b| t.block(e, b); + t.k.final(e, t.x, t.w.norm, t.w.norm_b, t.hidden, t.out, R); + const timing = try finish(cb, e, started, spent); + @memcpy(feature1, floats(t.out, R * W)); + return timing; + } + + /// The embedding alone: the stream's rows [4101, 1280] bf16 for `pixels` (the input of block 0). + pub fn embedOnly(t: *Encoder, pixels: []const f32) !void { + try t.stage(pixels); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.embed(e); + _ = try finish(cb, e, mtl.clock.seconds(), 0); + } + + /// Block `index` alone on bf16 `rows` [4101, 1280]; `stream` then holds its output. + pub fn blockAt(t: *Encoder, index: usize, rows: []const u16) !void { + if (index >= cfg.layers or rows.len != R * W) return error.BlockRows; + @memcpy(halves(t.x, R * W), rows); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.block(e, t.w.blocks[index]); + _ = try finish(cb, e, mtl.clock.seconds(), 0); + } + + /// The last norm and F.layer_norm alone on the last block's bf16 `rows`, into `feature1`. + pub fn finalAt(t: *Encoder, rows: []const u16, feature1: []f32) !void { + if (rows.len != R * W or feature1.len != R * W) return error.BlockRows; + @memcpy(halves(t.x, R * W), rows); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + t.k.final(e, t.x, t.w.norm, t.w.norm_b, t.hidden, t.out, R); + _ = try finish(cb, e, mtl.clock.seconds(), 0); + @memcpy(feature1, floats(t.out, R * W)); + } + + /// The stream's rows as the last pass left them, bf16 bits. + pub fn stream(t: *const Encoder) []const u16 { + return halves(t.x, R * W); + } + + /// last_hidden_state as the last encode left it, bf16 bits. + pub fn lastHidden(t: *const Encoder) []const u16 { + return halves(t.hidden, R * W); + } + + /// The RoPE tables load computed, bf16 bits [4096, 64] each. + pub fn ropeTables(t: *const Encoder) [2][]const u16 { + return .{ halves(t.cos, cfg.patches * cfg.head_dim), halves(t.sin, cfg.patches * cfg.head_dim) }; + } + + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder, started: f64, spent: f64) !Timing { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol dino: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return .{ .wall = mtl.clock.seconds() - started, .gpu = cb.gpuSeconds(), .host = spent }; + } + + /// The class token and registers, then the 16 x 16 conv over the patches as one GEMM into rows 5 to 4100. + fn embed(t: *Encoder, e: mtl.ComputeEncoder) void { + @memcpy(halves(t.x, cfg.prefix * W), halves(t.w.prefix, cfg.prefix * W)); + t.k.linear(e, t.pix, t.w.patch, t.w.patch_b, t.x.at(cfg.prefix * W * 2), cfg.patches, W, cfg.patch_k, W); + } + + /// One block: LayerNorm, attention (RoPE on the patches), LayerScale; LayerNorm, SwiGLU, LayerScale. + fn block(t: *const Encoder, e: mtl.ComputeEncoder, b: wts.Block) void { + const k = &t.k; + k.norm(e, t.x, b.norm1, b.norm1_b, t.xn, R); + k.linear(e, t.xn, b.qkv, t.w.zeros, t.qkv, R, 3 * W, W, 3 * W); + k.heads(e, t.qkv, t.cos, t.sin, t.q, t.k_, t.v, R); + k.attention(e, t.q, t.k_, t.v, t.att, cfg.heads, R, R); + k.linear(e, t.att, b.proj, b.proj_b, t.y, R, W, W, W); + k.scaleAdd(e, t.x, t.y, b.gamma1, R); + k.norm(e, t.x, b.norm2, b.norm2_b, t.xn, R); + k.linear(e, t.xn, b.fc1, b.fc1_b, t.gu, R, 2 * cfg.ffn, W, 2 * cfg.ffn); + k.swiglu(e, t.gu, t.h, R); + k.linear(e, t.h, b.fc2, b.fc2_b, t.y, R, W, cfg.ffn, W); + k.scaleAdd(e, t.x, t.y, b.gamma2, R); + } +}; + +test "the encoder's scratch fits one GPU buffer, each part 256-aligned" { + try std.testing.expect(scratchBytes() < 1 << 31); + try std.testing.expectEqual(@as(usize, 4101), R); +} diff --git a/zig/src/families/idol/dino_reference.zig b/zig/src/families/idol/dino_reference.zig new file mode 100644 index 0000000000..284814be1f --- /dev/null +++ b/zig/src/families/idol/dino_reference.zig @@ -0,0 +1,68 @@ +//! The DINOv3 encoder's gate: ours against the float32 ruler within twice torch bf16's distance, cosine and worst row. +const std = @import("std"); +const ref = @import("reference.zig"); +const cfg = @import("dino_config.zig"); + +pub const Match = ref.Match; +pub const match = ref.match; + +/// Ours and torch's against the ruler, and whether ours passes. +pub const Gate = struct { ours: Match, torch: Match, pass: bool }; + +/// Ours and torch's rows [n, 1280] each against the ruler's: worst row and 1 - cosine within 2x torch's. +pub fn gate(ours: []const f32, torch: []const f32, ruler: []const f32) Gate { + const o = ref.match(ours, ruler, cfg.width); + const t = ref.match(torch, ruler, cfg.width); + return .{ .ours = o, .torch = t, .pass = ref.within(o, t) }; +} + +/// The row holding the largest relative L2 error against `theirs`, and that error. +pub fn worstRow(ours: []const f32, theirs: []const f32) struct { row: usize, err: f64 } { + var worst: f64 = 0; + var at: usize = 0; + var row: usize = 0; + while (row * cfg.width < ours.len) : (row += 1) { + var rd: f64 = 0; + var rb: f64 = 0; + for (ours[row * cfg.width ..][0..cfg.width], theirs[row * cfg.width ..][0..cfg.width]) |a, b| { + rd += (@as(f64, a) - b) * (@as(f64, a) - b); + rb += @as(f64, b) * b; + } + const e = @sqrt(rd / @max(rb, 1e-30)); + if (e > worst) { + worst = e; + at = row; + } + } + return .{ .row = at, .err = worst }; +} + +test "a planted 2% error in one of 64 feature rows fails the gate, though its cosine alone would pass" { + const a = std.testing.allocator; + const n = 64 * cfg.width; + const ruler = try a.alloc(f32, n); + defer a.free(ruler); + const torch = try a.alloc(f32, n); + defer a.free(torch); + const ours = try a.alloc(f32, n); + defer a.free(ours); + var rng = std.Random.DefaultPrng.init(11); + for (ruler, torch, ours) |*r, *t, *o| { + r.* = rng.random().floatNorm(f32); + t.* = r.* * (1 + 4e-3 * rng.random().floatNorm(f32)); + o.* = r.* * (1 + 3e-3 * rng.random().floatNorm(f32)); + } + try std.testing.expect(gate(ours, torch, ruler).pass); + for (ours[37 * cfg.width ..][0..cfg.width], ruler[37 * cfg.width ..][0..cfg.width]) |*o, r| o.* += 0.02 * r; + const g = gate(ours, torch, ruler); + try std.testing.expect(!g.pass); + try std.testing.expect(ref.withinCosine(g.ours, g.torch)); + try std.testing.expectEqual(@as(usize, 37), worstRow(ours, ruler).row); +} + +test "a NaN anywhere in ours fails the gate" { + var ruler: [2 * cfg.width]f32 = @splat(1); + var ours = ruler; + ours[cfg.width + 3] = std.math.nan(f32); + try std.testing.expect(!gate(&ours, &ruler, &ruler).pass); +} diff --git a/zig/src/families/idol/dino_weights.zig b/zig/src/families/idol/dino_weights.zig new file mode 100644 index 0000000000..2fd7d173ef --- /dev/null +++ b/zig/src/families/idol/dino_weights.zig @@ -0,0 +1,196 @@ +//! DINOv3 ViT-H+/16's tensors from timm's fp32 safetensors, each rounded to bf16 at load as the SDK's dtype does. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +const cfg = @import("dino_config.zig"); +const host = @import("dino_host.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +const W = cfg.width; + +/// A block's bf16 tensors: the fused qkv [3840, 1280], gate | up [10240, 1280] as one, LayerScale's two gammas. +pub const Block = struct { qkv: Ref, proj: Ref, proj_b: Ref, fc1: Ref, fc1_b: Ref, fc2: Ref, fc2_b: Ref, norm1: Ref, norm1_b: Ref, norm2: Ref, norm2_b: Ref, gamma1: Ref, gamma2: Ref }; + +pub const Weights = struct { + ck: ck.Checkpoint, + blocks: [cfg.layers]Block, + /// The conv weight [1280, 768] and bias; the class token then the four registers [5, 1280]; the last norm. + patch: Ref, + patch_b: Ref, + prefix: Ref, + norm: Ref, + norm_b: Ref, + /// bf16 zeros [3840]: the qkv projection's bias (q and v's are zero, k has none). + zeros: Ref, + + pub fn deinit(w: *Weights) void { + w.ck.deinit(); + } +}; + +fn always(_: []const u8, _: usize) bool { + return true; +} + +/// One tensor's place in the bf16 buffer: its name, the values it holds, and where they go. +const Part = struct { name: []const u8, shape: [2]usize, at: usize }; + +fn bytes(n: usize) usize { + return std.mem.alignForward(usize, n * 2, ck.alignment); +} + +const per_block = blk: { + var n: usize = 0; + for ([_]usize{ 3 * W * W, W * W, W, 2 * cfg.ffn * W, 2 * cfg.ffn, W * cfg.ffn, W, W, W, W, W, W, W }) |v| n += bytes(v); + break :blk n; +}; +const globals = bytes(W * cfg.patch_k) + bytes(W) + bytes(cfg.prefix * W) + 2 * bytes(W) + bytes(3 * W); +pub const total_bytes = globals + cfg.layers * per_block; + +/// Reads deferred fp32 tensors into bf16 at their places, through its own staging buffer. +const Reader = struct { + c: *const ck.Checkpoint, + base: [*]u8, + stage: []f32, + + fn put(r: *Reader, p: Part) !void { + const d = r.c.deferred.get(p.name) orelse { + std.log.err("idol dino: model.safetensors has no {s}", .{p.name}); + return error.MissingTensor; + }; + const n = p.shape[0] * p.shape[1]; + if (d.dtype != .f32 or d.bytes != n * 4) { + std.log.err("idol dino: {s} is {t} {any}; the encoder reads f32 of {d} values", .{ p.name, d.dtype, d.shape[0..d.rank], n }); + return error.UnexpectedTensor; + } + var count: usize = 1; + for (d.shape[0..d.rank]) |s| count *= s; + if (count != n or (d.rank >= 2 and d.shape[0] != p.shape[0])) return error.UnexpectedTensor; + const dst = @as([*]u16, @ptrCast(@alignCast(r.base + p.at)))[0..n]; + var done: usize = 0; + while (done < n) { + const part = @min(r.stage.len, n - done); + try r.c.read(p.name, done * 4, std.mem.sliceAsBytes(r.stage[0..part])); + for (dst[done..][0..part], r.stage[0..part]) |*o, v| o.* = host.bits(v); + done += part; + } + } +}; + +/// Block `i` placed at `base` in `buf`: its Refs, and its 15 tensors' parts (gate's rows then up's rows make fc1). +fn blockAt(names: *[15][64]u8, i: usize, buf: mtl.Buffer, base: usize, parts: *[15]Part) !Block { + var at = base; + var b: Block = undefined; + const fields = [_]struct { *Ref, usize }{ + .{ &b.qkv, 3 * W * W }, .{ &b.proj, W * W }, .{ &b.proj_b, W }, .{ &b.fc1, 2 * cfg.ffn * W }, .{ &b.fc1_b, 2 * cfg.ffn }, + .{ &b.fc2, W * cfg.ffn }, .{ &b.fc2_b, W }, .{ &b.norm1, W }, .{ &b.norm1_b, W }, .{ &b.norm2, W }, + .{ &b.norm2_b, W }, .{ &b.gamma1, W }, .{ &b.gamma2, W }, + }; + for (fields) |f| { + f[0].* = .{ .buf = buf, .off = at }; + at += bytes(f[1]); + } + std.debug.assert(at - base == per_block); + const list = [_]struct { []const u8, usize, usize, usize }{ + .{ "attn.qkv.weight", 3 * W, W, b.qkv.off }, .{ "attn.proj.weight", W, W, b.proj.off }, + .{ "attn.proj.bias", W, 1, b.proj_b.off }, .{ "mlp.fc1_g.weight", cfg.ffn, W, b.fc1.off }, + .{ "mlp.fc1_x.weight", cfg.ffn, W, b.fc1.off + cfg.ffn * W * 2 }, .{ "mlp.fc1_g.bias", cfg.ffn, 1, b.fc1_b.off }, + .{ "mlp.fc1_x.bias", cfg.ffn, 1, b.fc1_b.off + cfg.ffn * 2 }, .{ "mlp.fc2.weight", W, cfg.ffn, b.fc2.off }, + .{ "mlp.fc2.bias", W, 1, b.fc2_b.off }, .{ "norm1.weight", W, 1, b.norm1.off }, + .{ "norm1.bias", W, 1, b.norm1_b.off }, .{ "norm2.weight", W, 1, b.norm2.off }, + .{ "norm2.bias", W, 1, b.norm2_b.off }, .{ "gamma_1", W, 1, b.gamma1.off }, + .{ "gamma_2", W, 1, b.gamma2.off }, + }; + for (list, parts, names) |e, *p, *n| p.* = .{ .name = try std.fmt.bufPrint(n, "blocks.{d}.{s}", .{ i, e[0] }), .shape = .{ e[1], e[2] }, .at = e[3] }; + return b; +} + +/// Blocks t, t + n, t + 2n, ... read on one thread; the first error is kept. +const Job = struct { + c: *const ck.Checkpoint, + base: [*]u8, + parts: *const [cfg.layers][15]Part, + stage: []f32, + failed: ?anyerror = null, + + fn run(j: *Job, t: usize, n: usize) void { + var r: Reader = .{ .c = j.c, .base = j.base, .stage = j.stage }; + var i = t; + while (i < cfg.layers) : (i += n) for (j.parts[i]) |p| r.put(p) catch |err| { + j.failed = err; + return; + }; + } +}; + +const threads = 8; + +/// The timm snapshot folder `dir`: config.json checked, then model.safetensors's 486 tensors, each read once into bf16. +pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(arena, &.{ dir, "config.json" }), arena, .limited(1 << 20)); + try cfg.check(gpa, config); + var w: Weights = undefined; + w.ck = ck.Checkpoint.init(gpa); + errdefer w.ck.deinit(); + const path = try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors", .{dir}, 0); + try w.ck.addFileSplit(device, path, "", null, always); + if (w.ck.deferred.count() != cfg.tensors or w.ck.tensors.count() != 0) { + std.log.err("idol dino: model.safetensors has {d} tensors; DINOv3 ViT-H+/16 has {d}", .{ w.ck.deferred.count() + w.ck.tensors.count(), cfg.tensors }); + return error.UnexpectedTensors; + } + const buf = try device.buffer(total_bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + w.ck.own(buf, total_bytes) catch |err| { + buf.deinit(); + return err; + }; + const base = buf.contents(); + var at: usize = 0; + const take = struct { + fn f(b: mtl.Buffer, pos: *usize, n: usize) Ref { + const r: Ref = .{ .buf = b, .off = pos.* }; + pos.* += bytes(n); + return r; + } + }.f; + w.patch = take(buf, &at, W * cfg.patch_k); + w.patch_b = take(buf, &at, W); + w.prefix = take(buf, &at, cfg.prefix * W); + w.norm = take(buf, &at, W); + w.norm_b = take(buf, &at, W); + w.zeros = take(buf, &at, 3 * W); + std.debug.assert(at == globals); + @memset(@as([*]u16, @ptrCast(@alignCast(base + w.zeros.off)))[0 .. 3 * W], 0); + var r: Reader = .{ .c = &w.ck, .base = base, .stage = try arena.alloc(f32, 1 << 22) }; + const head = [_]Part{ + .{ .name = "patch_embed.proj.weight", .shape = .{ W, cfg.patch_k }, .at = w.patch.off }, + .{ .name = "patch_embed.proj.bias", .shape = .{ W, 1 }, .at = w.patch_b.off }, + .{ .name = "cls_token", .shape = .{ 1, W }, .at = w.prefix.off }, + .{ .name = "reg_token", .shape = .{ 1, (cfg.prefix - 1) * W }, .at = w.prefix.off + W * 2 }, + .{ .name = "norm.weight", .shape = .{ W, 1 }, .at = w.norm.off }, + .{ .name = "norm.bias", .shape = .{ W, 1 }, .at = w.norm_b.off }, + }; + for (head) |p| try r.put(p); + const names = try arena.create([cfg.layers][15][64]u8); + const parts = try arena.create([cfg.layers][15]Part); + for (&w.blocks, 0..) |*b, i| b.* = try blockAt(&names[i], i, buf, globals + i * per_block, &parts[i]); + var jobs: [threads]Job = undefined; + var spawned: [threads]?std.Thread = @splat(null); + for (&jobs, 0..) |*j, t| { + j.* = .{ .c = &w.ck, .base = base, .parts = parts, .stage = try arena.alloc(f32, 1 << 22) }; + spawned[t] = std.Thread.spawn(.{}, Job.run, .{ j, t, threads }) catch null; + if (spawned[t] == null) j.run(t, threads); + } + for (spawned) |s| if (s) |th| th.join(); + for (jobs) |j| if (j.failed) |err| return err; + w.ck.closeFiles(); + return w; +} + +test "the bf16 layout: every part 256-aligned, the whole under one Metal buffer" { + try std.testing.expect(total_bytes < 1 << 31); + try std.testing.expectEqual(@as(usize, 0), per_block % ck.alignment); + try std.testing.expectEqual(@as(usize, 0), globals % ck.alignment); +} diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig index a2aff60944..c155ac603e 100644 --- a/zig/src/families/idol/idol.zig +++ b/zig/src/families/idol/idol.zig @@ -1,10 +1,12 @@ -//! Idol image to 3D: its denoiser on Metal (a denoiser beside the lane engine, not a token family). +//! Idol image to 3D: its denoiser and DINOv3 image encoder on Metal (beside the lane engine, not a token family). pub const config = @import("config.zig"); pub const host = @import("host.zig"); pub const weights = @import("weights.zig"); pub const kernels = @import("kernels.zig"); pub const model = @import("model.zig"); pub const reference = @import("reference.zig"); +/// Its DINOv3 ViT-H+/16 image encoder: an image to feature1. +pub const dino = @import("dino.zig"); pub const Model = model.Model; pub const Condition = model.Condition; pub const Timing = model.Timing; @@ -16,4 +18,5 @@ test { _ = kernels; _ = model; _ = reference; + _ = dino; } diff --git a/zig/tests/idol_dino.zig b/zig/tests/idol_dino.zig new file mode 100644 index 0000000000..f23859580b --- /dev/null +++ b/zig/tests/idol_dino.zig @@ -0,0 +1,188 @@ +//! tf-idol-dino: Idol's DINOv3 encoder on a case from `tools/zig/idol_case.py dino`, gated against torch and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const dn = tf.idol.dino; +const cfg = dn.config; +const ref = dn.reference; + +const usage = + \\tf-idol-dino DINO_DIR CASE_DIR [RUNS] + \\DINO_DIR is timm/vit_huge_plus_patch16_dinov3.lvd1689m (config.json, model.safetensors). CASE_DIR (idol_case.py + \\dino) holds dino.safetensors: the image's pixels, the SDK's feature1 (bf16 on MPS) and the float32 ruler's; and + \\optionally dino_blocks.safetensors (torch's rows into every block) and dino_blocks_f32.safetensors (each block in + \\float32 on them). feature1 from the image must sit as close to the ruler as the SDK's within 2x, worst row and + \\1 - cosine; so must the embedding, each block on torch's own input, and the last norm. Then the encode repeats + \\byte for byte and is timed RUNS times (default 5); CASE_DIR gets tf-idol-dino.json. Exit 1 on a failed gate. +; + +const N = cfg.rows * cfg.width; + +fn tensor(f: *const tf.safetensors.File, name: []const u8) !tf.safetensors.Tensor { + return f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; +} + +/// The case tensor `name` (f32 or bf16) as f32. +fn floats(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) ![]f32 { + const t = try tensor(f, name); + if (t.numel() != n) { + std.debug.print("{s} holds {d} values, not {d}\n", .{ name, t.numel(), n }); + return error.BadCaseTensor; + } + const out = try a.alloc(f32, n); + switch (t.dtype) { + .f32 => @memcpy(std.mem.sliceAsBytes(out), t.bytes), + .bf16 => for (out, @as([*]align(1) const u16, @ptrCast(t.bytes.ptr))[0..n]) |*o, v| { + o.* = dn.host.value(v); + }, + else => return error.BadCaseTensor, + } + return out; +} + +/// The case tensor `name` as bf16 bits. +fn bits(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) ![]u16 { + const t = try tensor(f, name); + if (t.dtype != .bf16 or t.numel() != n) { + std.debug.print("{s} is not {d} bf16 values\n", .{ name, n }); + return error.BadCaseTensor; + } + const out = try a.alloc(u16, n); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn widen(a: std.mem.Allocator, src: []const u16) ![]f32 { + const out = try a.alloc(f32, src.len); + for (out, src) |*o, v| o.* = dn.host.value(v); + return out; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn json(list: *std.ArrayList(u8), a: std.mem.Allocator, key: []const u8, m: ref.Match) !void { + try print(list, a, ", \"{s}\": {{\"max_abs\": {d:.6}, \"cosine\": {d:.9}, \"worst_row\": {d:.6}}}", .{ key, m.max_abs, m.cosine, m.worst_row }); +} + +/// One stage's gate printed and recorded; returns whether it passed. +fn stage(out: *std.ArrayList(u8), a: std.mem.Allocator, label: []const u8, ours: []const f32, torch: []const f32, ruler: []const f32) !bool { + const g = ref.gate(ours, torch, ruler); + const wo = ref.worstRow(ours, ruler); + const wt = ref.worstRow(torch, ruler); + std.debug.print("{s} against float32: worst row {d:.5} (row {d}), 1 - cos {e:.3}; torch's {d:.5} (row {d}), {e:.3}: {s}\n", .{ label, g.ours.worst_row, wo.row, 1 - g.ours.cosine, g.torch.worst_row, wt.row, 1 - g.torch.cosine, if (g.pass) "pass" else "FAIL" }); + try print(out, a, ", \"{s}\": {{\"pass\": {}, \"worst_row_index\": {d}, \"torch_worst_row_index\": {d}", .{ label, g.pass, wo.row, wt.row }); + try json(out, a, "ours_vs_ruler", g.ours); + try json(out, a, "torch_vs_ruler", g.torch); + try out.append(a, '}'); + return g.pass; +} + +/// The embedding, each block on torch's own input rows, and the last norm, each against its float32 run. +fn blocks(a: std.mem.Allocator, io: std.Io, dir: []const u8, enc: *dn.Encoder, pixels: []const f32, out: *std.ArrayList(u8)) !?bool { + var f = tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/dino_blocks.safetensors", .{dir})) catch return null; + defer f.close(io); + var g = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/dino_blocks_f32.safetensors", .{dir})); + defer g.close(io); + var pass = true; + var name: [32]u8 = undefined; + try enc.embedOnly(pixels); + const torch0 = try floats(a, &f, "block.0", N); + const embed = try widen(a, enc.stream()); + const same = std.mem.eql(u8, std.mem.sliceAsBytes(enc.stream()), std.mem.sliceAsBytes(try bits(a, &f, "block.0", N))); + std.debug.print("embedding: {s} torch's bf16 rows\n", .{if (same) "the same bytes as" else "differs from"}); + try print(out, a, ", \"embed_bytes_equal_torch\": {}", .{same}); + pass = (try stage(out, a, "embed", embed, torch0, try floats(a, &g, "embed", N))) and pass; + for (0..cfg.layers) |i| { + try enc.blockAt(i, try bits(a, &f, try std.fmt.bufPrint(&name, "block.{d}", .{i}), N)); + const label = try std.fmt.bufPrint(&name, "block.{d}", .{i + 1}); + const ours = try widen(a, enc.stream()); + pass = (try stage(out, a, label, ours, try floats(a, &f, label, N), try floats(a, &g, label, N))) and pass; + } + const last = try bits(a, &f, "block.32", N); + const ours = try a.alloc(f32, N); + try enc.finalAt(last, ours); + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/dino.safetensors", .{dir})); + defer case.close(io); + pass = (try stage(out, a, "final", ours, try floats(a, &case, "feature1", N), try floats(a, &g, "feature1", N))) and pass; + return pass; +} + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 3 and args.len != 4) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + const runs = if (args.len == 4) try std.fmt.parseInt(usize, args[3], 10) else 5; + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/dino.safetensors", .{args[2]})); + defer case.close(io); + const pixels = try floats(a, &case, "pixels", 3 * cfg.image * cfg.image); + const sdk = try floats(a, &case, "feature1", N); + const ruler = try floats(a, &case, "feature1_f32", N); + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const enc = try dn.Encoder.load(std.heap.page_allocator, io, device, args[1]); + defer enc.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: loaded in {d:.2} s\n", .{ device.name(), load_s }); + var out: std.ArrayList(u8) = .empty; + try print(&out, a, "{{\"device\": \"{s}\", \"load_s\": {d:.3}", .{ device.name(), load_s }); + + // the RoPE tables against torch's own bf16 cos and sin + const tables = enc.ropeTables(); + var equal: usize = 0; + for ([_][]const u8{ "cos", "sin" }, tables) |n, ours| { + for (ours, try bits(a, &case, n, cfg.patches * cfg.head_dim)) |x, y| equal += @intFromBool(x == y); + } + std.debug.print("RoPE tables: {d} of {d} bf16 values equal torch's\n", .{ equal, 2 * cfg.patches * cfg.head_dim }); + try print(&out, a, ", \"rope_equal\": {d}, \"rope_values\": {d}", .{ equal, 2 * cfg.patches * cfg.head_dim }); + + const ours = try a.alloc(f32, N); + const again = try a.alloc(f32, N); + const first = try enc.encode(pixels, ours); + var walls = try a.alloc(f64, runs); + var gpus = try a.alloc(f64, runs); + var repeat = true; + for (0..runs) |r| { + const t = try enc.encode(pixels, again); + walls[r] = t.wall; + gpus[r] = t.gpu; + repeat = repeat and std.mem.eql(u8, std.mem.sliceAsBytes(ours), std.mem.sliceAsBytes(again)); + } + std.mem.sort(f64, walls, {}, std.sort.asc(f64)); + std.mem.sort(f64, gpus, {}, std.sort.asc(f64)); + std.debug.print("encode: first {d:.4} s; median of {d} {d:.4} s wall, {d:.4} s on the GPU\n", .{ first.wall, runs, walls[runs / 2], gpus[runs / 2] }); + std.debug.print("{d} encodes after the first: {s}\n", .{ runs, if (repeat) "identical bytes" else "DIFFERENT bytes" }); + try print(&out, a, ", \"first_s\": {d:.4}, \"encode_s\": {d:.4}, \"gpu_s\": {d:.4}, \"host_s\": {d:.4}, \"repeat_identical\": {}", .{ first.wall, walls[runs / 2], gpus[runs / 2], first.host, repeat }); + + const vs_sdk = ref.match(ours, sdk, cfg.width); + std.debug.print("feature1 against the SDK's: max |diff| {d:.5}, cosine {d:.8}, worst row {d:.5} (row {d})\n", .{ vs_sdk.max_abs, vs_sdk.cosine, vs_sdk.worst_row, ref.worstRow(ours, sdk).row }); + try json(&out, a, "feature1_vs_sdk", vs_sdk); + var pass = try stage(&out, a, "feature1", ours, sdk, ruler); + if (case.get("hidden") != null) { + const theirs = try bits(a, &case, "hidden", N); + var same: usize = 0; + for (enc.lastHidden(), theirs) |x, y| same += @intFromBool(x == y); + std.debug.print("last_hidden_state: {d:.2}% of bf16 values equal the SDK's\n", .{100 * @as(f64, @floatFromInt(same)) / N}); + try print(&out, a, ", \"hidden_equal_fraction\": {d:.5}", .{@as(f64, @floatFromInt(same)) / N}); + } + if (try blocks(a, io, args[2], enc, pixels, &out)) |ok| { + pass = pass and ok; + try print(&out, a, ", \"blocks_pass\": {}", .{ok}); + } + try print(&out, a, ", \"gates_pass\": {}}}\n", .{pass}); + std.debug.print("gates: {s}\n", .{if (pass) "all pass" else "FAIL"}); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-dino.json", .{args[2]}), .data = out.items }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/feature1.f32", .{args[2]}), .data = std.mem.sliceAsBytes(ours) }); + return if (pass and repeat) 0 else 1; +} From 61ea7b23bdfe78df2f3e2f0dbab864fca164b5cf Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:06:26 +0800 Subject: [PATCH 14/21] contract: Idol's DINOv3 kernels through gates 1-5; the recipe and changelog line Every pipeline of idol_dino.metal is registered: dn_linear_0 and dn_linear_1, dn_norm, dn_final, dn_heads, dn_scale_add, dn_swiglu and dn_attention, each checked at transformers' bf16 points. 11 controls: a bf16 rounding added or dropped (gate 3), fp16 sums in dn_norm and dn_final and P V rounded to bf16 each key block (gate 4), the attention's tail keys unmasked. Gate 5 holds the 8 launches, each reserved by its source. SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; 11 controls: 11 red as planted, 0 missed; 0 limit problems --- CHANGELOG.md | 2 + docs/kernel-contract.md | 2 + docs/recipes/idol.md | 100 +++++- zig/src/core/contract.zig | 1 + zig/src/core/contract/idol_dino_records.zig | 18 ++ zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_idol_dino.zig | 135 ++++++++ .../contract/metal_idol_dino_attention.zig | 204 ++++++++++++ zig/tests/contract/metal_idol_dino_common.zig | 103 ++++++ zig/tests/contract/metal_idol_dino_gemm.zig | 222 +++++++++++++ zig/tests/contract/metal_idol_dino_norms.zig | 252 +++++++++++++++ zig/tests/contract/metal_idol_dino_rows.zig | 294 ++++++++++++++++++ zig/tests/contract/metal_limits.zig | 3 + zig/tests/contract/metal_main.zig | 2 + 14 files changed, 1337 insertions(+), 3 deletions(-) create mode 100644 zig/src/core/contract/idol_dino_records.zig create mode 100644 zig/tests/contract/metal_idol_dino.zig create mode 100644 zig/tests/contract/metal_idol_dino_attention.zig create mode 100644 zig/tests/contract/metal_idol_dino_common.zig create mode 100644 zig/tests/contract/metal_idol_dino_gemm.zig create mode 100644 zig/tests/contract/metal_idol_dino_norms.zig create mode 100644 zig/tests/contract/metal_idol_dino_rows.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 2abaa9af8a..12f41a7149 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. +- Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (after #1766, for #1757). Thanks to @jayleaton. + - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. - Qwen3.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b2b465117a..02bc888474 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -188,6 +188,8 @@ HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row Metal also checks Idol's denoiser (`metal_idol*.zig`): its fp16 projections and GELU form, the split-K fp32 rows, `id_gemm32` (true fp32 sums: the NAX fp32 path in its place is a control gate 4 must catch), the LayerNorms, RePo's angles, the heads, the elementwise steps and the attention's four forms, each at idol-sdk's fp16 points (19 controls). +Metal also checks Idol's DINOv3 encoder (`metal_idol_dino*.zig`): its bf16 projections, LayerNorm and final F.layer_norm, +the RoPE heads, LayerScale with the residual, SwiGLU and the attention, each at transformers' bf16 points (11 controls). RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 5b506262a9..12efeacd8e 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -15,8 +15,9 @@ each row's LayerNorm to per-head 3D positions, angles from 10, 10 and 12 frequen The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_cond - 3 v_uncond` (two passes a step: with the condition, and with zeros). -The image encoders, the decoders and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's -SDK in torch, a development reference. It is a denoiser, not a token family: it does not register with the lane +The DINOv3 image encoder runs natively too ([DINOv3 encoder](#dinov3-encoder)); the FLUX.2 VAE encoder, the decoders +and the PLY writer are not ported yet. `tools/zig/idol_case.py` runs the publisher's SDK in torch, a development +reference. It is a denoiser, not a token family: it does not register with the lane engine, the server or `families.detect`. It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them. @@ -129,3 +130,98 @@ Results: 100). Load 0.2 to 1.6 s; the condition, once an image, 0.07 s. - **Memory:** 1.84 GB peak process footprint (`/usr/bin/time -l`, the whole `tf-idol` run with its case), the denoiser's fp16 weights in shared Metal buffers. + +## DINOv3 encoder + +`dino*.zig` (kernels in `idol_dino.metal`) turns the SDK's 1024 x 1024 preprocessed image into the denoiser's +`feature1` [4101, 1280]: DINOv3 ViT-H+/16's class token, four registers and 64 x 64 patches. idol-sdk's `ImageEncoder` +(`conditioning.py`) loads `facebook/dinov3-vith16plus-pretrain-lvd1689m` through transformers' `DINOv3ViTModel` in +bf16 (`dinov3_dtype='bf16'`, the pipeline's default), normalizes the pixels with ImageNet's mean and std in fp32, and +returns `F.layer_norm(last_hidden_state.float())`: the model's own final LayerNorm (affine, bf16), then a plain one in +fp32 (eps 1e-5, no affine). + +- Patch embedding: the 16 x 16 conv as one GEMM, [4096, 768] patches (channel, row, column) against [1280, 768]. +- 32 blocks of width 1,280, 20 heads of 64; LayerNorm (eps 1e-5), attention, LayerScale; LayerNorm, SwiGLU MLP of + 5,120 (`silu(gate) * up`; gate and up run as one 10,240-column GEMM), LayerScale. +- RoPE on the 4,096 patches alone, not the five prefix rows: transformers' periods `1 / 100^(4i / 64)` for i < 16, + generated in fp32 (not the bf16 periods of Meta's checkpoint, which timm's card mentions); patch centres in + [-1, 1]; angles `2 pi y f` then `2 pi x f`, tiled twice, in `rotate_half` form; cos and sin in fp32, then bf16. The + host builds these tables at load; they equal torch's bit for bit (524,288 of 524,288 values). + +### Weights and licence + +`timm/vit_huge_plus_patch16_dinov3.lvd1689m` at `dd0addc09788111fa893d24799a44744ba022eee`, ungated: `config.json` +(`vit_huge_plus_patch16_dinov3`, 1,280 features) and `model.safetensors` (sha256 +`0d54131c846c91a78ec7641482f32e45b27cb4895f83653c038f305f18c7e61c`, 486 fp32 tensors), each rounded to bf16 at load +(ties to even, as `from_pretrained(dtype=bfloat16)` rounds). **These weights are under Meta's DINOv3 License, not +Apache 2.0**: its `LICENSE.md` ships with them, and using or redistributing them or a derivative binds you to it. + +The SDK reads Meta's gated checkpoint. `idol_case.py dino --official DIR` checks that timm's is the same model: the +names map as transformers' converter maps them (`cls_token`, `reg_token`, `patch_embed.proj`, `norm` to `embeddings.*` +and `norm.*`; `blocks.N.attn.qkv` split in q, k, v thirds; `attn.proj` to `attention.o_proj`; `gamma_1`, `gamma_2` to +`layer_scale1`, `layer_scale2`; `mlp.fc1_g`, `fc1_x`, `fc2` to `gate_proj`, `up_proj`, `down_proj`). On the avocado: + +- all 550 mapped tensors (the 32 fused qkv weights as 96 thirds) equal Meta's byte for byte; Meta's q and v biases + are all zero, and its only other unmapped tensor is the masked-input `mask_token`; +- feature1 from timm's weights in that graph equals feature1 from Meta's `from_pretrained` bit for bit on MPS, and + both equal the features the SDK pipeline recorded for the image. + +### Precision + +Where torch on MPS rounds, measured on random bf16 operands against float64 (the share of outputs each candidate +matches): a linear or the conv rounds once, `bf16(x w + b)` (99.98%; rounding before the bias, 69%); nn.LayerNorm +once (99.995%); `silu(g) * u` twice; the rotary three times (`bf16(bf16(q cos) + bf16(rotate_half(q) sin))`); and +LayerScale with the residual twice (`bf16(x + bf16(g y))`), each 100%. The kernels round at those points, with fp32 +sums and fp32 statistics. + +Attention: torch's bf16 SDPA on MPS matches no candidate scheme (fp32 or bf16 P, its denominator before or after, +bf16 scores) and sits about three times a bf16 rounding of the exact output from it. `dn_attention` keeps fp32 scores +and sums; the tensor units take P in V's type, so it rounds P to bf16 itself and sums the denominator over those +values; the output rounds to bf16. + +### Kernel contract + +`zig/tests/contract/metal_idol_dino*.zig` registers every kernel of `idol_dino.metal`: `dn_linear_0` and +`dn_linear_1` (the 4,101 rows and the 4,096 patches), `dn_norm`, `dn_final`, `dn_heads`, `dn_scale_add`, `dn_swiglu` +and `dn_attention`, each at transformers' bf16 points. Eleven controls each break one: a bf16 rounding added or +dropped (gate 3), fp16 sums in `dn_norm` and `dn_final`, P V rounded to bf16 each key block (gate 4), the attention's +tail keys unmasked. + +``` +SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; 11 controls: 11 red as planted, 0 missed; 0 limit problems +``` + +### Run it + +```sh +hf download timm/vit_huge_plus_patch16_dinov3.lvd1689m --revision dd0addc09788111fa893d24799a44744ba022eee --local-dir DINO +python3 tools/zig/idol_case.py dino DINO SDK case/ --image preprocessed.png [--official META_DIR] [--sdk-features F] +zig build tf-idol-dino -Dcpu=apple_m1 +zig-out/bin/tf-idol-dino DINO case/ +``` + +The image is the SDK's 1024 x 1024 preprocessed image (`IdolOutput.preprocessed_image`). The tool pins idol-sdk's +`conditioning.py` and transformers 5.22.1's DINOv3 files by sha256, the weights by theirs, and writes `dino.safetensors` +(the pixels, the SDK's feature1 on MPS, a float32 CPU ruler's, last_hidden_state and the RoPE tables), +`dino_blocks.safetensors` (torch's rows into each block) and `dino_blocks_f32.safetensors` (the embedding, each block +and the last norm in float32 on torch's own input rows). The ruler runs on the CPU: the M5's fp32 GEMM is not fp32. + +`tf-idol-dino` gates, each against the float32 ruler within twice torch bf16's own distance, worst row and `1 - cos`: +feature1 from the image; the embedding; each block alone on torch's input; the last norm alone. The encode must repeat +byte for byte. `dino_reference.zig`'s test plants a 2% error in one row and requires the gate to fail it. + +### Measured + +The avocado example, the same machine and torch as above (`ATTN_BACKEND=sdpa`): + +- **feature1** against float32: our worst row 0.1353, torch's 0.1686 (both row 4038, the bottom-left corner patch); + `1 - cos` 5.37e-5, torch's 5.57e-5. Against the SDK: cosine 0.999954, worst row 0.109 (row 4099). Pass. +- **Embedding**: worst row 0.0089 and `1 - cos` equal to torch's at the printed digits. Pass. +- **Blocks**, each of the 32 on torch's bf16 input: all pass; our worst row is torch's times 0.82 to 1.24, our + `1 - cos` torch's times 0.84 to 1.06. Worst rows 0.0045 to 0.0138, except block 2's row 1810 at 0.1905 (torch 0.1905). +- **Last norm**: equal to torch's at the printed digits. Pass. +- **Repeat**: identical bytes over 10 encodes. +- **Speed**: 0.120 s an encode (median of 10; 0.118 s on the GPU), the first 0.150 s, against the SDK's 0.261 s on MPS + (median of 10 warm; its first 0.385 s): 2.2x. Load 0.2 to 0.6 s. +- **Memory**: 2.28 GB peak process footprint (`/usr/bin/time -l`, `tf-idol-dino` with the case's feature rows), the bf16 + weights 1.68 GB of it. The SDK's MPS encode holds 2.24 GB of MPS driver memory. diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 4153341694..d92eb4f644 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -26,6 +26,7 @@ test { _ = @import("contract/nemotron_experts_records.zig"); _ = @import("contract/moss_image_records.zig"); _ = @import("contract/idol_records.zig"); + _ = @import("contract/idol_dino_records.zig"); _ = @import("contract/metal_pair_records.zig"); _ = spec; _ = gates; diff --git a/zig/src/core/contract/idol_dino_records.zig b/zig/src/core/contract/idol_dino_records.zig new file mode 100644 index 0000000000..e739d2ec53 --- /dev/null +++ b/zig/src/core/contract/idol_dino_records.zig @@ -0,0 +1,18 @@ +//! Idol's DINOv3 encoder launches (idol_dino.metal after gemm_nax.metal), each source-reserved. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "idol dino dn_linear_0", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol dino dn_linear_1", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol dino dn_norm", .threads = 256, .tg_memory = 144, .reserved = 256 }, + .{ .name = "idol dino dn_final", .threads = 256, .tg_memory = 5264, .reserved = 256 }, + .{ .name = "idol dino dn_heads", .threads = 128, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol dino dn_scale_add", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol dino dn_swiglu", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol dino dn_attention", .threads = 128, .tg_memory = 0, .reserved = 128 }, +}; + +test "the Idol DINOv3 launches fit every Apple family, one record a pipeline" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); + try std.testing.expectEqual(@as(usize, 8), records.len); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index f55d146bf7..cdd12c7112 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -569,7 +569,7 @@ test { } /// The Metal collector keeps measured slots and the restored vision declarations. pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ - @import("moss_image_records.zig").records ++ @import("minimax_records.zig").records ++ @import("idol_records.zig").records; + @import("moss_image_records.zig").records ++ @import("minimax_records.zig").records ++ @import("idol_records.zig").records ++ @import("idol_dino_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/tests/contract/metal_idol_dino.zig b/zig/tests/contract/metal_idol_dino.zig new file mode 100644 index 0000000000..2c7a32b3f5 --- /dev/null +++ b/zig/tests/contract/metal_idol_dino.zig @@ -0,0 +1,135 @@ +//! Idol's DINOv3 encoder kernels under the contract, each at transformers' bf16 points; controls break one point each. +const std = @import("std"); +const mtl = @import("metal"); +const rig_ = @import("rig.zig"); +const dn = @import("metal_idol_dino_common.zig"); +const gemm = @import("metal_idol_dino_gemm.zig"); +const norms = @import("metal_idol_dino_norms.zig"); +const elems = @import("metal_idol_dino_rows.zig"); +const attn = @import("metal_idol_dino_attention.zig"); + +const g = dn.g; +const Control = dn.Control; +const Lib = dn.Lib; + +const controls = [_]Control{ + .{ .kernel = "dn_linear_0", .what = "adds bf16(x w) before the bias", .edits = &.{.{ .from = "acc[i][j][e] = dn_b(acc[i][j][e] + b);", .to = "acc[i][j][e] = dn_b(dn_b(acc[i][j][e]) + b);" }} }, + .{ .kernel = "dn_linear_1", .what = "adds bf16(x w) before the bias", .edits = &.{.{ .from = "acc[i][j][e] = dn_b(acc[i][j][e] + b);", .to = "acc[i][j][e] = dn_b(dn_b(acc[i][j][e]) + b);" }} }, + .{ .kernel = "dn_norm", .what = "adds bf16(n) before the affine", .edits = &.{.{ .from = "yr[i] = bfloat((float(xr[i]) - mean) * inv * float(w[i]) + float(b[i]));", .to = "yr[i] = bfloat(dn_b((float(xr[i]) - mean) * inv) * float(w[i]) + float(b[i]));" }} }, + .{ .kernel = "dn_norm", .what = "sums the simdgroups in fp16", .gate = .sums, .edits = &.{.{ .from = " v = simd_sum(v);\n if (lane == 0) part[sg] = v;", .to = " v = float(simd_sum(half(v)));\n if (lane == 0) part[sg] = v;" }} }, + .{ .kernel = "dn_final", .what = "drops bf16(h) before F.layer_norm", .edits = &.{.{ .from = " v[i] = float(o);", .to = " v[i] = (float(xr[i]) - mean) * inv * float(w[i]) + float(b[i]);" }} }, + .{ .kernel = "dn_final", .what = "sums F.layer_norm's squares in fp16", .gate = .sums, .edits = &.{.{ .from = "q2 += (v[i] - mean2) * (v[i] - mean2);", .to = "q2 = float(half(q2 + (v[i] - mean2) * (v[i] - mean2)));" }} }, + .{ .kernel = "dn_heads", .what = "drops bf16(q cos)", .edits = &.{.{ .from = "Q[to] = bfloat(dn_b(q0 * c0) + dn_b(-q1 * s0));", .to = "Q[to] = bfloat((q0 * c0) + dn_b(-q1 * s0));" }} }, + .{ .kernel = "dn_scale_add", .what = "drops bf16(g y)", .edits = &.{.{ .from = "x[i] = bfloat(float(x[i]) + dn_b(float(g[at.x]) * float(y[i])));", .to = "x[i] = bfloat(float(x[i]) + (float(g[at.x]) * float(y[i])));" }} }, + .{ .kernel = "dn_swiglu", .what = "drops bf16(silu(g))", .edits = &.{.{ .from = "bfloat(dn_b(g / (1.0f + metal::precise::exp(-g))) * float(gu[i + width]));", .to = "bfloat((g / (1.0f + metal::precise::exp(-g))) * float(gu[i + width]));" }} }, + .{ .kernel = "dn_attention", .what = "keys past Lk unmasked", .edits = &.{.{ .from = "s[f][e] = (16 * f + home.x + TF_COL(e)) < k_rows ? s[f][e] : kDnMasked;", .to = "s[f][e] = s[f][e];" }} }, + .{ .kernel = "dn_attention", .what = "rounds the P V sums to bf16 each key block", .gate = .sums, .edits = &.{.{ .from = "for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * scale_old[e >> 2];", .to = "for (short e = 0; e < 8; e++) acc[i][e] = float(bfloat(acc[i][e] * scale_old[e >> 2]));" }} }, +}; + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + libs: std.ArrayList(*Lib) = .empty, + names: std.ArrayList([]u8) = .empty, + made: std.ArrayList(Made) = .empty, + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: *Lib, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: dn.Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol dino {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + for (gemm.forms, 0..) |fun, f| { + const name = (try w.name(fun)) orelse continue; + const t = try s.make(gemm.Linear, .{ .a = s.a, .rig = rig, .lib = lib, .form = @intCast(f), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]norms.Kind{ .norm, .final }, [_][]const u8{ "dn_norm", "dn_final" }) |kind, fun| { + const name = (try w.name(fun)) orelse continue; + const t = try s.make(norms.Norm, .{ .a = s.a, .rig = rig, .lib = lib, .kind = kind, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + inline for (.{ .{ elems.Heads, "dn_heads" }, .{ elems.ScaleAdd, "dn_scale_add" }, .{ elems.Swiglu, "dn_swiglu" } }) |e| { + if (try w.name(e[1])) |name| { + const t = try s.make(e[0], .{ .a = s.a, .rig = rig, .lib = lib, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + if (try w.name("dn_attention")) |name| { + const t = try s.make(attn.Attention, .{ .a = s.a, .rig = rig, .lib = lib, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + if (!device.tensorUnits()) { + std.debug.print("UNSUPPORTED metal idol dino: the encoder runs on tensor units (M5) alone; its kernels unchecked here\n", .{}); + return s; + } + for ([_]?Control{null} ++ controls) |x| { + const lib = try a.create(Lib); + lib.* = .{ .a = a, .device = device, .ctl = x }; + s.libs.append(a, lib) catch |err| { + a.destroy(lib); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.libs.items) |l| { + l.deinit(); + a.destroy(l); + } + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_dino_attention.zig b/zig/tests/contract/metal_idol_dino_attention.zig new file mode 100644 index 0000000000..4fdc5650fc --- /dev/null +++ b/zig/tests/contract/metal_idol_dino_attention.zig @@ -0,0 +1,204 @@ +//! dn_attention under the contract: softmax(q k^T / 8) v, bf16 heads of 64, against attention_ref's float64 formula. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const dn = @import("metal_idol_dino_common.zig"); + +const c = dn.c; +const g = dn.g; +const D = dn.D; +const kref = dn.kref; + +const heads = dn.cfg.heads; + +/// Query rows repeat at this period (prime to the 64-row blocks): float64 costs 89 queries a head over every key. +const query_period = 89; + +/// The scores' 64-deep sums, then the keys' sums. +fn attentionDepth(s: c.Shape) usize { + return D + s.k; +} + +pub const Attention = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *dn.Lib, + name: []const u8, + lq: usize = 0, + lk: usize = 0, + stress: bool = false, + q: []u16 = &.{}, + k: []u16 = &.{}, + v: []u16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: the launch's 20 heads; n: query rows Lq; k: key rows Lk. The encoder runs 4,101 against 4,101. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = attentionDepth, .ops = 8, .out = .bf16, .tiles = .{ .{}, .{ .size = 64 }, .{ .size = 32 } }, .served = &.{ .{ .rows = heads, .n = 4101, .k = 4101 }, .{ .rows = heads, .n = 100, .k = 64 } } }; + } + + pub fn deinit(t: *Attention) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Attention) void { + t.a.free(t.q); + t.a.free(t.k); + t.a.free(t.v); + t.q = &.{}; + t.k = &.{}; + t.v = &.{}; + } + pub fn kernel(t: *Attention, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Attention); + + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.k == 0) return false; + t.rig.reset(); + t.free(); + t.lq = shape.n; + t.lk = shape.k; + t.stress = fill == .stress; + const L = t.lk; + t.q = try t.a.alloc(u16, heads * t.lq * D); + t.k = try t.a.alloc(u16, heads * L * D); + t.v = try t.a.alloc(u16, heads * L * D); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // equal scores (q = k = 0); each column's values bf16 integers whose partial sums reach ~2^22, then cancel + @memset(t.q, 0); + @memset(t.k, 0); + const half = L / 2; + for (0..heads) |h| for (0..L) |j| { + const at = if (j < half) j else j - half; + const positive: f32 = if (at % 2 == 0) 127 else 4064; + // the first four 4,064s come back as -4,032: a residual of 128 over the column + const back: f32 = if (at < 8 and at % 2 == 1) 4032 else positive; + const v: f32 = if (j < half) positive else if (j < 2 * half) -back else 0; + @memset(t.v[(h * L + j) * D ..][0..D], dn.bits(v)); + }; + } else { + dn.randomBf(rng, t.q, -2, 2); + dn.randomBf(rng, t.k, -2, 2); + dn.randomBf(rng, t.v, -1, 1); + } + if (fill == .nonfinite) { + // head 0: values NaN, +Inf, -Inf, and both infinities in one column; head 1: a NaN query, an infinite key + for ([_][2]usize{ .{ 0, 0 }, .{ 1, 1 }, .{ 2, 2 }, .{ 3, 3 }, .{ 4, 3 } }, [_]u16{ dn.nan, dn.inf, dn.ninf, dn.inf, dn.ninf }) |p, v| if (p[0] < L) { + t.v[p[0] * D + p[1]] = v; + }; + t.q[(t.lq + 0) * D + 5] = dn.nan; + if (L > 7) t.k[(L + 7) * D + 3] = dn.inf; + } + // every head's queries repeat at the period, so float64 computes each distinct one once + const p = @min(t.lq, query_period); + for (0..heads) |h| for (p..t.lq) |i| { + @memcpy(t.q[(h * t.lq + i) * D ..][0..D], t.q[(h * t.lq + i % p) * D ..][0..D]); + }; + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.q)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.k)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.v)); + t.out = try t.rig.output(heads * t.lq * D * 2); + return true; + } + + pub fn run(t: *Attention) !void { + try t.lib.need("dn_attention"); + const p = t.rig.begin(); + t.lib.k.attention(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), heads, t.lq, t.lk); + try t.rig.finish(p); + } + + pub fn count(t: *Attention) usize { + return heads * t.lq * D; + } + + /// attention_ref.head's formula, one pass over the keys a distinct query; O rows [Lq, heads * 64]. + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const L = t.lk; + const scale: f64 = @as(f32, 1.0 / @sqrt(@as(f32, D))); + const kf = try t.a.alloc(f64, L * D); + defer t.a.free(kf); + const vf = try t.a.alloc(f64, L * D); + defer t.a.free(vf); + const s = try t.a.alloc(f64, L); + defer t.a.free(s); + var qf: [D]f64 = undefined; + var out: [D]f64 = undefined; + var tol: [D]f64 = undefined; + const p = @min(t.lq, query_period); + for (0..heads) |h| { + for (kf, vf, t.k[h * L * D ..][0 .. L * D], t.v[h * L * D ..][0 .. L * D]) |*kd, *vd, kb, vb| { + kd.* = dn.value(kb); + vd.* = dn.value(vb); + } + for (0..p) |i| { + var zero = true; + for (&qf, t.q[(h * t.lq + i) * D ..][0..D]) |*x, b| { + x.* = dn.value(b); + zero = zero and x.* == 0; + } + var m: f64 = -std.math.inf(f64); + var nan = false; + for (s, 0..) |*sj, j| { + var dot: f64 = 0; + if (!zero) { + const Vec = @Vector(8, f64); + var acc: Vec = @splat(0); + var e: usize = 0; + while (e < D) : (e += 8) acc += @as(Vec, qf[e..][0..8].*) * @as(Vec, kf[j * D + e ..][0..8].*); + dot = @reduce(.Add, acc); + } + sj.* = dot * scale; + nan = nan or std.math.isNan(sj.*); + m = @max(m, sj.*); + } + if (nan) m = std.math.nan(f64); + var l: f64 = 0; + @memset(&out, 0); + @memset(&tol, 0); + for (s, 0..) |sj, j| { + const pj = @exp(sj - m); + l += pj; + for (&out, &tol, vf[j * D ..][0..D]) |*o, *tl, vv| { + o.* += pj * vv; + tl.* += pj * @abs(vv); + } + } + var r = i; + while (r < t.lq) : (r += p) { + const at = (r * heads + h) * D; + for (out, tol, 0..) |o, tl, e| { + want[at + e] = o / l; + // P enters P V as bf16 (2^-8 each), so the weights move by up to 2^-7; O rounds to bf16 + mag[at + e] = tl / l * 0x1p-7 + @abs(o / l) * 0x1p-8 + 1e-6; + } + } + } + } + } + + pub fn bounds(t: *Attention, want: []const f64, tolerance: []const f64, allow: []f64) !void { + for (want, tolerance, allow) |w, tl, *l| l.* = if (t.stress) c.boundExact(spec(t.name), w) else if (std.math.isFinite(w)) tl else 0; + } + + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!t.stress) return error.AttentionStressNeedsEqualScores; + const L = t.lk; + const terms = try t.a.alloc(f64, L); + defer t.a.free(terms); + const ones = try t.a.alloc(f64, L); + defer t.a.free(ones); + @memset(ones, 1); + const l = contract.rows.sumIn(acc, order, ones); + for (0..heads) |h| for (0..D) |e| { + for (terms, 0..) |*v, j| v.* = dn.value(t.v[(h * L + j) * D + e]); + const o = contract.rows.sumIn(acc, order, terms) / l; + for (0..t.lq) |i| out[(i * heads + h) * D + e] = o; + }; + } +}; diff --git a/zig/tests/contract/metal_idol_dino_common.zig b/zig/tests/contract/metal_idol_dino_common.zig new file mode 100644 index 0000000000..7fe07c3183 --- /dev/null +++ b/zig/tests/contract/metal_idol_dino_common.zig @@ -0,0 +1,103 @@ +//! Idol's DINOv3 kernel tests' shared parts: float64 values at bf16 points, bf16 inputs and the encoder's library. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const id = @import("metal_idol_common.zig"); + +pub const tf = id.tf; +pub const c = id.c; +pub const g = id.g; +pub const rows = id.rows; +pub const Range = id.Range; +pub const Ends = id.Ends; +pub const Control = id.Control; +pub const kref = id.kref; +pub const noSums = id.noSums; +pub const periodic = id.periodic; +pub const int = id.int; +pub const sign = id.sign; +pub const ks = tf.idol.dino.kernels; +pub const cfg = tf.idol.dino.config; +pub const W = cfg.width; +pub const D = cfg.head_dim; + +/// A value rounded once to bf16 (ties to even). +pub const bf = contract.bf16_round.bf16; + +/// A range's ends each rounded to bf16. +pub fn bfr(r: Range) Range { + return .{ .lo = bf(r.lo), .hi = bf(r.hi) }; +} + +/// fp32 bits rounded to bf16 bits, and back. +pub const bits = contract.bf16_round.bits; +pub const value = contract.bf16_round.value; + +/// bf16 values uniform in [lo, hi]. +pub fn randomBf(rng: std.Random, out: []u16, lo: f32, hi: f32) void { + for (out) |*v| v.* = bits(lo + rng.float(f32) * (hi - lo)); +} + +pub const nan: u16 = 0x7fc0; +pub const inf: u16 = 0x7f80; +pub const ninf: u16 = 0xff80; + +pub const div_ulps = id.div_ulps; +pub const rsqrt_ulps = id.rsqrt_ulps; +pub const exp_ulps = id.exp_ulps; + +/// The encoder's source as the engine compiles it (or a control's copy), its pipelines made as tests ask for them. +pub const Lib = struct { + a: std.mem.Allocator, + device: mtl.Device, + ctl: ?Control, + k: ks.Kernels = undefined, + built: bool = false, + made: [ks.names.len]bool = @splat(false), + + fn text(l: *Lib) ![]u8 { + var out = try tf.frags.source(l.device, l.a, sources.idol_dino); + const x = l.ctl orelse return out; + for (x.edits) |e| { + errdefer l.a.free(out); + if (std.mem.count(u8, out, e.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, e.from }); + return error.ControlTextNotFoundOnce; + } + const next = try std.mem.replaceOwned(u8, l.a, out, e.from, e.to); + l.a.free(out); + out = next; + } + if (x.append.len == 0) return out; + defer l.a.free(out); + return std.mem.concat(l.a, u8, &.{ out, "\n", x.append }); + } + + /// The pipeline `name` made on this library (compiled on first use). + pub fn need(l: *Lib, name: []const u8) !void { + if (!l.built) { + const src = try l.text(); + defer l.a.free(src); + l.k.own = try mtl.Library.fromSource(l.device, src, mtl.CompileOptions.mlx()); + l.built = true; + } + for (ks.names, 0..) |n, i| if (std.mem.eql(u8, n, name)) { + if (!l.made[i]) l.k.p[i] = try mtl.Pipeline.init(l.device, l.k.own, n, false); + l.made[i] = true; + return; + }; + return error.NoSuchPipeline; + } + + pub fn deinit(l: *Lib) void { + for (l.k.p, l.made) |p, m| if (m) p.deinit(); + if (l.built) l.k.own.deinit(); + } +}; + +test "bf16 ranges round both ends to nearest even" { + const r = bfr(.{ .lo = 1 + 0x1p-8, .hi = 1 + 3 * 0x1p-8 }); + try std.testing.expectEqual(@as(f64, 1), r.lo); + try std.testing.expectEqual(@as(f64, 1 + 0x1p-6), r.hi); +} diff --git a/zig/tests/contract/metal_idol_dino_gemm.zig b/zig/tests/contract/metal_idol_dino_gemm.zig new file mode 100644 index 0000000000..d9b0bfd15c --- /dev/null +++ b/zig/tests/contract/metal_idol_dino_gemm.zig @@ -0,0 +1,222 @@ +//! dn_linear under the contract: bf16 projections, fp32 sums on the tensor units, the bias in them, one bf16 rounding. +const std = @import("std"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const dn = @import("metal_idol_dino_common.zig"); + +const c = dn.c; +const g = dn.g; +const W = dn.W; +const Range = dn.Range; +const Ends = dn.Ends; +const bf = dn.bf; +const kref = dn.kref; + +/// Rows of x and of w (with the bias) repeat at these periods, prime to every tile: 97 x 131 dots a shape. +const row_period = 97; +const col_period = 131; + +/// The M5 tensor unit's fp32 sums: its 16-deep steps (attention_range's), as for the denoiser's fp16 projections. +fn naxDepth(s: c.Shape) usize { + return s.k / 16 * contract.attention_range.mma_step + 8; +} + +/// float64's dot of a bf16 row pair, the tensor unit's allowance on its fp32 sums step by step, and the terms' size. +fn naxDot(x: []const u16, w: []const u16) [3]f64 { + var s: f64 = 0; + var spread: f64 = 0; + var mag: f64 = 0; + var i: usize = 0; + while (i < x.len) : (i += 16) { + var t: f64 = 0; + var abs: f64 = 0; + for (x[i..@min(i + 16, x.len)], w[i..@min(i + 16, x.len)]) |a, b| { + const p = @as(f64, dn.value(a)) * @as(f64, dn.value(b)); + t += p; + abs += @abs(p); + } + spread += @abs(s) + abs; + s += t; + mag += abs; + } + return .{ s, spread * contract.attention_range.mma_step * 0x1p-24, mag }; +} + +/// The two forms by M % 64 == 0: the 4,101 rows (a last tile's rows), the 4,096 patches (whole 64-row tiles). +pub const forms = [2][]const u8{ "dn_linear_0", "dn_linear_1" }; + +fn envelope(form: u1) struct { served: []const c.Shape, tiles: [3]c.Tile } { + const col: c.Tile = .{ .size = 256, .step = 128, .min = 128 }; + const k: c.Tile = .{ .size = 256, .step = 32, .min = 32 }; + if (form == 1) return .{ .served = &.{.{ .rows = 4096 / 64, .n = W, .k = dn.cfg.patch_k }}, .tiles = .{ .{ .size = 2, .step = 1, .min = 1 }, col, k } }; + // first the 4,101 rows' last tile alone (5 rows): the single-shape gates take the first served shape's rows + return .{ + .served = &.{ .{ .rows = 5, .n = W, .k = W }, .{ .rows = 4101, .n = 3 * W, .k = W }, .{ .rows = 4101, .n = W, .k = W }, .{ .rows = 4101, .n = 2 * dn.cfg.ffn, .k = W }, .{ .rows = 4101, .n = W, .k = dn.cfg.ffn } }, + .tiles = .{ .{ .size = 64 }, col, k }, + }; +} + +pub const Linear = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *dn.Lib, + form: u1, + name: []const u8, + m: usize = 0, + n: usize = 0, + k: usize = 0, + stress: bool = false, + x: []u16 = &.{}, + w: []u16 = &.{}, + b: []u16 = &.{}, + out: rig_.Ref = undefined, + ins: [3]rig_.Ref = undefined, + ends: Ends, + + /// Rows: M (form 1: 64-row tiles); n: N, whole 128-column tiles; k: K. + pub fn spec(name: []const u8, form: u1) c.Spec { + const e = envelope(form); + return .{ .name = name, .accum = .f32, .depth = naxDepth, .ops = 1, .out = .bf16, .tiles = e.tiles, .served = e.served }; + } + + pub fn deinit(t: *Linear) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Linear) void { + t.a.free(t.x); + t.a.free(t.w); + t.a.free(t.b); + t.x = &.{}; + t.w = &.{}; + t.b = &.{}; + } + pub fn kernel(t: *Linear, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Linear); + + pub fn prepare(t: *Linear, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const m = if (t.form == 1) shape.rows * 64 else shape.rows; + if (m == 0 or shape.n == 0 or shape.n % 128 != 0 or shape.k == 0 or shape.k % 32 != 0) return false; + if (@intFromBool(m % 64 == 0) != t.form) return false; + t.rig.reset(); + t.free(); + t.m = m; + t.n = shape.n; + t.k = shape.k; + t.stress = fill == .stress; + const K = shape.k; + const px = @min(m, row_period); + const pw = @min(shape.n, col_period); + t.x = try t.a.alloc(u16, px * K); + t.w = try t.a.alloc(u16, pw * K); + t.b = try t.a.alloc(u16, pw); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // integers in [hi / 2, hi) (below 256, exact in bf16), hi^2 K <= 2^24: fp32 sums exact in any order + const hi: i32 = @min(256, @as(i32, @intFromFloat(@sqrt(@as(f64, 0x1p24) / @as(f64, @floatFromInt(K)))))); + for (0..px) |r| for (0..K / 2) |i| { + const v = dn.int(rng, @divTrunc(hi, 2), hi - 1); + t.x[r * K + i] = dn.bits(v); + // x's halves differ by 1 in their first 8: the sums cancel to a few hundred + t.x[r * K + K / 2 + i] = dn.bits(v - @as(f32, if (i < 8) 1 else 0)); + }; + for (0..pw) |j| for (0..K / 2) |i| { + const v = dn.int(rng, @divTrunc(hi, 2), hi - 1); + t.w[j * K + i] = dn.bits(v); + t.w[j * K + K / 2 + i] = dn.bits(-v); + }; + for (t.b) |*v| v.* = dn.bits(dn.int(rng, -4, 4)); + } else { + dn.randomBf(rng, t.x, -2, 2); + dn.randomBf(rng, t.w, -0.25, 0.25); + dn.randomBf(rng, t.b, -1, 1); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, K / 3 }, .{ 1, K / 2 }, .{ 2, K - 1 }, .{ 3, 0 }, .{ 3, K - 1 } }; + for (at, [_]u16{ dn.nan, dn.inf, dn.ninf, dn.inf, dn.ninf }) |p, v| if (p[0] < px) { + t.x[p[0] * K + p[1]] = v; + }; + if (pw > 2) t.b[2] = dn.inf; + } + const xs = try dn.periodic(u16, t.a, t.x, K, m); + defer t.a.free(xs); + const ws = try dn.periodic(u16, t.a, t.w, K, shape.n); + defer t.a.free(ws); + const bs = try dn.periodic(u16, t.a, t.b, 1, shape.n); + defer t.a.free(bs); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(xs)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(ws)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(bs)); + t.out = try t.rig.output(m * shape.n * 2); + return true; + } + + pub fn run(t: *Linear) !void { + try t.lib.need(forms[t.form]); + const p = t.rig.begin(); + t.lib.k.linear(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.out), t.m, t.n, t.k, t.n); + try t.rig.finish(p); + } + + pub fn count(t: *Linear) usize { + return t.m * t.n; + } + + /// The distinct outputs' values and ranges from (sum, its range) by row and column, then every output's. + fn table(t: *Linear, ends: *Ends, want: []f64, mag: []f64, ctx: anytype, comptime sum: anytype) !void { + const px = t.x.len / t.k; + const pw = t.w.len / t.k; + const tw = try t.a.alloc(f64, 3 * px * pw); + defer t.a.free(tw); + for (0..px) |r| for (0..pw) |j| { + const s = sum(ctx, t, r, j); + const b: f64 = dn.value(t.b[j]); + const vr = dn.bfr(s[1].add(Range.of(b)).fp32(0.5)); + tw[3 * (r * pw + j) ..][0..3].* = .{ bf(s[0] + b), vr.lo, vr.hi }; + }; + for (0..t.m) |r| for (0..t.n) |j| { + const e = tw[3 * ((r % px) * pw + j % pw) ..][0..3]; + ends.put(want, mag, r * t.n + j, e[0], .{ .lo = e[1], .hi = e[2] }); + }; + } + + pub fn reference(t: *Linear, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const f = struct { + fn sum(_: void, s: *const Linear, r: usize, j: usize) struct { f64, Range } { + const d = naxDot(s.x[r * s.k ..][0..s.k], s.w[j * s.k ..][0..s.k]); + return .{ d[0], if (s.stress) Range.of(d[0]) else Range.of(d[0]).widen(d[1]) }; + } + }; + try t.table(&t.ends, want, mag, {}, f.sum); + } + pub fn bounds(t: *Linear, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Linear, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Linear, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + const terms = try t.a.alloc(f64, t.k); + defer t.a.free(terms); + const Ctx = struct { acc: c.Accum, order: g.Order, terms: []f64 }; + const f = struct { + fn sum(x: Ctx, s: *const Linear, r: usize, j: usize) struct { f64, Range } { + for (x.terms, s.x[r * s.k ..][0..s.k], s.w[j * s.k ..][0..s.k]) |*v, a, b| v.* = @as(f64, dn.value(a)) * @as(f64, dn.value(b)); + const v = dn.rows.sumIn(x.acc, x.order, x.terms); + return .{ v, Range.of(v) }; + } + }; + try t.table(&scratch, out, mag, Ctx{ .acc = acc, .order = order, .terms = terms }, f.sum); + } +}; diff --git a/zig/tests/contract/metal_idol_dino_norms.zig b/zig/tests/contract/metal_idol_dino_norms.zig new file mode 100644 index 0000000000..64d055de17 --- /dev/null +++ b/zig/tests/contract/metal_idol_dino_norms.zig @@ -0,0 +1,252 @@ +//! dn_norm and dn_final under the contract: nn.LayerNorm of bf16 rows in fp32, then (dn_final) F.layer_norm in fp32. +const std = @import("std"); +const rig_ = @import("rig.zig"); +const dn = @import("metal_idol_dino_common.zig"); + +const c = dn.c; +const g = dn.g; +const cfg = dn.cfg; +const W = dn.W; +const Range = dn.Range; +const Ends = dn.Ends; +const bf = dn.bf; +const kref = dn.kref; + +pub const Kind = enum { norm, final }; + +const norm_why = "a bf16 output follows the sums as var^-1/2: a bf16 sum's error is a few output ulps"; + +/// A thread's 5 values of the 1,280 in turn, the simdgroup tree, then the tree over the threadgroup's 8 simdgroups. +const sum_depth = W / 256 + 5 + 5; + +fn normDepth(_: c.Shape) usize { + return sum_depth + 3; +} + +/// A row's statistics from the kernel's mean (range `mr`) and squared deviations (range `qr`): 1 / sqrt(q / n + eps). +fn invRange(qr: Range) Range { + const n: f64 = W; + return (Range{ .lo = qr.lo / n, .hi = qr.hi / n }).add(Range.of(cfg.eps)).fp32(0.5).rsqrt().fp32(dn.rsqrt_ulps); +} + +/// The mean's and squares' ranges of a row of exact values `xs` under fp32 sums (or exact ones, `exact`). +fn stats(xs: []const f64, exact: bool) struct { mean: f64, q: f64, mr: Range, qr: Range } { + const n: f64 = W; + const gs = c.gamma(sum_depth, c.unit(.f32)); + const gq = c.gamma(sum_depth + 3, c.unit(.f32)); + var s: f64 = 0; + var abs: f64 = 0; + for (xs) |v| { + s += v; + abs += @abs(v); + } + const mean = s / n; + var q: f64 = 0; + for (xs) |v| q += (v - mean) * (v - mean); + if (exact) return .{ .mean = mean, .q = q, .mr = Range.of(mean), .qr = Range.of(q) }; + const mr = (Range{ .lo = (s - gs * abs) / n, .hi = (s + gs * abs) / n }).fp32(dn.div_ulps); + // the squares about the kernel's mean exceed float64's by n (m' - mean)^2, then fp32 rounds them + const dm = @max(@abs(mr.lo - mean), @abs(mr.hi - mean)); + return .{ .mean = mean, .q = q, .mr = mr, .qr = .{ .lo = q * (1 - gq), .hi = (q + n * dm * dm) * (1 + gq) } }; +} + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *dn.Lib, + kind: Kind, + name: []const u8, + shape: c.Shape = undefined, + stress: bool = false, + x: []u16 = &.{}, + /// w then b, bf16 + wb: []u16 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + /// dn_final's bf16 rows, working memory + hidden: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n and k: the width, 1,280 (the launcher's). + pub fn spec(name: []const u8, kind: Kind) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = normDepth, + .ops = @as(f64, if (kind == .final) 2 else 1) * (2 * (dn.div_ulps + dn.rsqrt_ulps) + 3), + .out = if (kind == .final) .f32 else .bf16, + .served = &.{.{ .rows = cfg.rows, .n = W, .k = W }}, + .sums = if (kind == .norm) .{ .bf16 = 0, .why = norm_why } else .{}, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.ends.deinit(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + t.a.free(t.x); + t.a.free(t.wb); + t.x = &.{}; + t.wb = &.{}; + } + pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.kind), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Norm); + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != W or shape.k != W) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.stress = fill == .stress; + t.x = try t.a.alloc(u16, shape.rows * W); + t.wb = try t.a.alloc(u16, 2 * W); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (t.stress) { + // integers mirrored in each row's halves: a mean of exactly 0 and squares summing to ~2^22, exact in fp32 + for (0..shape.rows) |r| for (0..W / 2) |i| { + const v = dn.int(rng, -64, 64); + t.x[r * W + i] = dn.bits(v); + t.x[r * W + W / 2 + i] = dn.bits(-v); + }; + } else dn.randomBf(rng, t.x, -3, 3); + dn.randomBf(rng, t.wb[0..W], 0.25, 2); + dn.randomBf(rng, t.wb[W..], -1, 1); + if (t.stress and t.kind == .final) { + // w = 0: the norm's rows are b, integers mirrored too, so the second norm's sums are exact as well + @memset(t.wb[0..W], 0); + for (0..W / 2) |i| { + const v = dn.int(rng, -64, 64); + t.wb[W + i] = dn.bits(v); + t.wb[W + W / 2 + i] = dn.bits(-v); + } + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, W / 3 }, .{ 1, W / 2 }, .{ 2, W - 1 }, .{ 3, 0 }, .{ 3, W - 1 } }; + for (at, [_]u16{ dn.nan, dn.inf, dn.ninf, dn.inf, dn.ninf }) |p, v| if (p[0] < shape.rows) { + t.x[p[0] * W + p[1]] = v; + }; + // dn_final's rows each fold every column into F.layer_norm: a non-finite w or b would reach every row + if (t.kind == .norm) { + t.wb[5] = dn.nan; + t.wb[W + 6] = dn.inf; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.wb[0..W])); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.wb[W..])); + if (t.kind == .final) { + // the bf16 last_hidden_state is working memory here: the fp32 rows check its rounding + t.hidden = try t.rig.scratch(shape.rows * W * 2); + t.out = try t.rig.output(shape.rows * W * 4); + } else t.out = try t.rig.output(shape.rows * W * 2); + return true; + } + pub fn run(t: *Norm) !void { + const k = &t.lib.k; + try t.lib.need(if (t.kind == .final) "dn_final" else "dn_norm"); + const p = t.rig.begin(); + const in = t.ins; + if (t.kind == .final) { + k.final(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(t.hidden), kref(t.out), t.shape.rows); + } else k.norm(p.e, kref(in[0]), kref(in[1]), kref(in[2]), kref(t.out), t.shape.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * W; + } + + /// Row r's outputs, its sums (both norms') taken from `sums` when replayed, else bounded as fp32's. + fn row(t: *Norm, ends: *Ends, r: usize, want: []f64, mag: []f64, replay: ?Replay) void { + var xs: [W]f64 = undefined; + for (&xs, t.x[r * W ..][0..W]) |*o, v| o.* = dn.value(v); + var st = stats(&xs, t.stress); + if (replay) |rp| { + st.mean = rp.sum(&xs) / W; + var d: [W]f64 = undefined; + for (&d, xs) |*o, x| o.* = (x - st.mean) * (x - st.mean); + st.q = rp.sum(&d); + st.mr = Range.of(st.mean); + st.qr = Range.of(st.q); + } + const inv = 1 / @sqrt(st.q / W + cfg.eps); + const ir = invRange(st.qr); + var hs: [W]f64 = undefined; + var hr: [W]Range = undefined; + for (0..W) |i| { + const w: f64 = dn.value(t.wb[i]); + const b: f64 = dn.value(t.wb[W + i]); + const vr = Range.of(xs[i]).sub(st.mr).fp32(0.5).mul(ir).fp32(0.5).mul(Range.of(w)).fp32(0.5).add(Range.of(b)).fp32(0.5); + hs[i] = bf((xs[i] - st.mean) * inv * w + b); + hr[i] = dn.bfr(vr); + if (t.kind == .norm) ends.put(want, mag, r * W + i, hs[i], hr[i]); + } + if (t.kind == .norm) return; + // F.layer_norm of the bf16 rows: each value in its range, the statistics over those ranges + var s2 = stats(&hs, t.stress); + if (replay) |rp| { + s2.mean = rp.sum(&hs) / W; + var d: [W]f64 = undefined; + for (&d, hs) |*o, x| o.* = (x - s2.mean) * (x - s2.mean); + s2.q = rp.sum(&d); + s2.mr = Range.of(s2.mean); + s2.qr = Range.of(s2.q); + } + var lo: f64 = 0; + var hi: f64 = 0; + for (hr) |h| { + lo += h.lo; + hi += h.hi; + } + const spread = Range{ .lo = s2.mr.lo + (lo - s2.mean * W) / W, .hi = s2.mr.hi + (hi - s2.mean * W) / W }; + const mr = Range{ .lo = @min(s2.mr.lo, spread.lo), .hi = @max(s2.mr.hi, spread.hi) }; + var qlo: f64 = 0; + var qhi: f64 = 0; + for (hr) |h| { + const d = h.sub(mr).square(); + qlo += d.lo; + qhi += d.hi; + } + const qr = Range{ .lo = @min(s2.qr.lo, qlo * (1 - c.gamma(sum_depth + 3, c.unit(.f32)))), .hi = @max(s2.qr.hi, qhi * (1 + c.gamma(sum_depth + 3, c.unit(.f32)))) }; + const inv2 = 1 / @sqrt(s2.q / W + cfg.eps); + const ir2 = invRange(if (t.stress and replay == null) Range.of(s2.q) else qr); + for (0..W) |i| { + const y = (hs[i] - s2.mean) * inv2; + const yr = if (t.stress and replay == null) Range.of(hs[i]).sub(Range.of(s2.mean)).mul(ir2).fp32(0.5) else hr[i].sub(mr).fp32(0.5).mul(ir2).fp32(0.5); + ends.put(want, mag, r * W + i, y, yr); + } + } + + const Replay = struct { + acc: c.Accum, + order: g.Order, + fn sum(rp: Replay, terms: []f64) f64 { + return dn.rows.sumIn(rp.acc, rp.order, terms); + } + }; + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (0..t.shape.rows) |r| t.row(&t.ends, r, want, mag, null); + } + pub fn bounds(t: *Norm, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Norm, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + + pub fn simulate(t: *Norm, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try t.a.alloc(f64, out.len); + defer t.a.free(mag); + var scratch: Ends = .{ .a = t.a }; + defer scratch.deinit(); + try scratch.size(t.count()); + for (0..t.shape.rows) |r| t.row(&scratch, r, out, mag, .{ .acc = acc, .order = order }); + } +}; diff --git a/zig/tests/contract/metal_idol_dino_rows.zig b/zig/tests/contract/metal_idol_dino_rows.zig new file mode 100644 index 0000000000..751f8509ea --- /dev/null +++ b/zig/tests/contract/metal_idol_dino_rows.zig @@ -0,0 +1,294 @@ +//! dn_heads, dn_scale_add and dn_swiglu under the contract: elementwise steps at transformers' bf16 points. +const std = @import("std"); +const rig_ = @import("rig.zig"); +const dn = @import("metal_idol_dino_common.zig"); + +const c = dn.c; +const g = dn.g; +const cfg = dn.cfg; +const W = dn.W; +const D = dn.D; +const Range = dn.Range; +const Ends = dn.Ends; +const bf = dn.bf; +const kref = dn.kref; + +/// The steps' common parts: their buffers' lifetime, and the bounds and intervals from the outputs' ends. +fn Common(comptime T: type) type { + return struct { + pub fn deinit(t: *T) void { + for (&t.v) |*v| { + t.a.free(v.*); + v.* = &.{}; + } + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *T, control: ?g.Gate) g.Kernel { + return .{ .spec = T.spec(t.name), .ctx = t, .vt = &T.vt, .control = control }; + } + pub fn bounds(t: *T, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *T, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } + fn fresh(t: *T) void { + t.rig.reset(); + for (&t.v) |*v| { + t.a.free(v.*); + v.* = &.{}; + } + } + }; +} + +/// bf16(a * b) for bf16 a and b: the product is exact in fp32, so one rounding. +fn prod(a: f64, b: f64) Range { + return dn.bfr(Range.of(a * b)); +} + +// --- dn_heads: q, k, v [20, rows, 64] from qkv rows; q and k turned past the prefix by bf16 cos, sin --------------- + +pub const Heads = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *dn.Lib, + name: []const u8, + rows: usize = 0, + /// qkv [rows, 3840], then cos and sin [rows - 5, 64] + v: [3][]u16 = .{ &.{}, &.{}, &.{} }, + ins: [3]rig_.Ref = undefined, + outs: [3]rig_.Ref = undefined, + ends: Ends, + + /// Rows: the stream's 4,101 (the prefix's 5 unturned); n: qkv's width. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = dn.noSums, .ops = 1, .mids = 2, .out = .bf16, .served = &.{.{ .rows = cfg.rows, .n = 3 * W, .k = 1 }} }; + } + const C = Common(Heads); + pub const deinit = C.deinit; + pub const kernel = C.kernel; + pub const bounds = C.bounds; + pub const intervals = C.intervals; + const vt = g.hooks(Heads); + + pub fn prepare(t: *Heads, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows <= cfg.prefix or shape.n != 3 * W) return false; + C.fresh(t); + t.rows = shape.rows; + const patches = shape.rows - cfg.prefix; + t.v[0] = try t.a.alloc(u16, shape.rows * 3 * W); + t.v[1] = try t.a.alloc(u16, patches * D); + t.v[2] = try t.a.alloc(u16, patches * D); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + dn.randomBf(rng, t.v[0], -4, 4); + dn.randomBf(rng, t.v[1], -1, 1); + dn.randomBf(rng, t.v[2], -1, 1); + if (fill == .nonfinite) { + // a q in the prefix and two past it, a v; a cos and a sin of the patches + for ([_]u16{ dn.nan, dn.inf, dn.ninf }, 0..) |v, i| { + t.v[0][((i * 2) % shape.rows) * 3 * W + i * 97] = v; + t.v[0][i * 3 * W + 2 * W + 11 * i] = v; + } + t.v[1][(1 % patches) * D + 5] = dn.inf; + t.v[2][(2 % patches) * D + 40] = dn.nan; + } + for (&t.ins, t.v) |*r, v| r.* = try t.rig.input(std.mem.sliceAsBytes(v)); + for (&t.outs) |*r| r.* = try t.rig.output(cfg.heads * shape.rows * D * 2); + return true; + } + + pub fn run(t: *Heads) !void { + try t.lib.need("dn_heads"); + const p = t.rig.begin(); + t.lib.k.heads(p.e, kref(t.ins[0]), kref(t.ins[1]), kref(t.ins[2]), kref(t.outs[0]), kref(t.outs[1]), kref(t.outs[2]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Heads) usize { + return 3 * cfg.heads * t.rows * D; + } + + pub fn reference(t: *Heads, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const per = cfg.heads * t.rows * D; + for (0..t.rows) |r| for (0..cfg.heads) |h| for (0..D) |j| { + const to = (h * t.rows + r) * D + j; + const x = t.v[0][r * 3 * W ..]; + const vv: f64 = dn.value(x[2 * W + h * D + j]); + t.ends.put(want, mag, 2 * per + to, vv, Range.of(vv)); + for (0..2) |m| { + const at = m * W + h * D; + const own: f64 = dn.value(x[at + j]); + if (r < cfg.prefix) { + t.ends.put(want, mag, m * per + to, own, Range.of(own)); + continue; + } + // rotate_half: -x[j + 32] for j < 32, x[j - 32] past it + const other: f64 = if (j < 32) -@as(f64, dn.value(x[at + j + 32])) else dn.value(x[at + j - 32]); + const tb = (r - cfg.prefix) * D + j; + const cs: f64 = dn.value(t.v[1][tb]); + const sn: f64 = dn.value(t.v[2][tb]); + const a = prod(own, cs); + const b = prod(other, sn); + t.ends.put(want, mag, m * per + to, bf(bf(own * cs) + bf(other * sn)), dn.bfr(a.add(b).fp32(0.5))); + } + }; + } +}; + +// --- dn_scale_add: x = bf16(x + bf16(g y)), in place ------------------------------------------------------------- + +pub const ScaleAdd = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *dn.Lib, + name: []const u8, + rows: usize = 0, + /// the stream's rows as they go in, y, then g [1280] + v: [3][]u16 = .{ &.{}, &.{}, &.{} }, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: positions; n: the width, 1,280 (the launcher's). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = dn.noSums, .ops = 1, .mids = 1, .out = .bf16, .in_place = true, .served = &.{.{ .rows = cfg.rows, .n = W, .k = 1 }} }; + } + const C = Common(ScaleAdd); + pub const deinit = C.deinit; + pub const kernel = C.kernel; + pub const bounds = C.bounds; + pub const intervals = C.intervals; + const vt = blk: { + var hooks = g.hooks(ScaleAdd); + hooks.overlap = overlap; + break :blk hooks; + }; + /// The stream is its own output. + fn overlap(_: *anyopaque) bool { + return true; + } + + pub fn prepare(t: *ScaleAdd, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != W) return false; + C.fresh(t); + t.rows = shape.rows; + const n = shape.rows * W; + t.v[0] = try t.a.alloc(u16, n); + t.v[1] = try t.a.alloc(u16, n); + t.v[2] = try t.a.alloc(u16, W); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + dn.randomBf(rng, t.v[0], -8, 8); + dn.randomBf(rng, t.v[1], -4, 4); + dn.randomBf(rng, t.v[2], -2, 2); + if (fill == .nonfinite) { + for ([_]u16{ dn.nan, dn.inf, dn.ninf }, 0..) |v, i| { + t.v[0][(i * 13) % n] = v; + t.v[1][(i * 17 + 1) % n] = v; + } + t.v[2][5] = dn.nan; + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.v[1])); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.v[2])); + t.out = try t.rig.output(n * 2); + return true; + } + + pub fn run(t: *ScaleAdd) !void { + try t.lib.need("dn_scale_add"); + // in place: the stream's rows go back before each independent run + @memcpy(t.rig.bytes(t.out, t.count() * 2), std.mem.sliceAsBytes(t.v[0])); + const p = t.rig.begin(); + t.lib.k.scaleAdd(p.e, kref(t.out), kref(t.ins[0]), kref(t.ins[1]), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *ScaleAdd) usize { + return t.rows * W; + } + + pub fn reference(t: *ScaleAdd, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (t.v[0], t.v[1], 0..) |xb, yb, i| { + const x: f64 = dn.value(xb); + const y: f64 = dn.value(yb); + const gv: f64 = dn.value(t.v[2][i % W]); + t.ends.put(want, mag, i, bf(x + bf(gv * y)), dn.bfr(prod(gv, y).add(Range.of(x)).fp32(0.5))); + } + } +}; + +// --- dn_swiglu: h = bf16(bf16(silu(gate)) up) from rows [gate | up] ------------------------------------------------- + +pub const Swiglu = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + lib: *dn.Lib, + name: []const u8, + rows: usize = 0, + v: [1][]u16 = .{&.{}}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const F = cfg.ffn; + + /// Rows: positions; n: the MLP's 5,120 (the launcher's). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = dn.noSums, .ops = dn.exp_ulps + dn.div_ulps + 1, .mids = 1, .out = .bf16, .served = &.{.{ .rows = cfg.rows, .n = F, .k = 1 }} }; + } + const C = Common(Swiglu); + pub const deinit = C.deinit; + pub const kernel = C.kernel; + pub const bounds = C.bounds; + pub const intervals = C.intervals; + const vt = g.hooks(Swiglu); + + pub fn prepare(t: *Swiglu, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != F) return false; + C.fresh(t); + t.rows = shape.rows; + const n = shape.rows * 2 * F; + t.v[0] = try t.a.alloc(u16, n); + var prng = std.Random.DefaultPrng.init(seed); + dn.randomBf(prng.random(), t.v[0], -8, 8); + if (fill == .nonfinite) { + for ([_]u16{ dn.nan, dn.inf, dn.ninf }, 0..) |v, i| { + t.v[0][(i * 13) % F] = v; + t.v[0][2 * F + F + i * 17] = v; + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.v[0])); + t.out = try t.rig.output(shape.rows * F * 2); + return true; + } + + pub fn run(t: *Swiglu) !void { + try t.lib.need("dn_swiglu"); + const p = t.rig.begin(); + t.lib.k.swiglu(p.e, kref(t.ins[0]), kref(t.out), t.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Swiglu) usize { + return t.rows * F; + } + + pub fn reference(t: *Swiglu, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + for (0..t.rows) |r| for (0..F) |j| { + const gv: f64 = dn.value(t.v[0][r * 2 * F + j]); + const u: f64 = dn.value(t.v[0][r * 2 * F + F + j]); + const s = gv / (1 + @exp(-gv)); + // silu's fp32 steps: exp(-g), 1 + it, then g over that at fast math's division + const den = Range.of(@exp(-gv)).fp32(dn.exp_ulps).add(Range.of(1)).fp32(0.5); + const sr = dn.bfr(Range.of(gv).mul(den.recip()).fp32(dn.div_ulps)); + const y = Range{ .lo = @min(bf(sr.lo) * u, bf(sr.hi) * u), .hi = @max(bf(sr.lo) * u, bf(sr.hi) * u) }; + t.ends.put(want, mag, r * F + j, bf(bf(s) * u), dn.bfr(y)); + }; + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index ece8073ecc..dd84e79e50 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -102,6 +102,9 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol dino ")) { + text = sources.idol_dino; + frags = true; } else if (std.mem.startsWith(u8, d.name, "idol ")) { text = sources.idol; frags = true; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 9799191182..dcad0eab2b 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -429,6 +429,8 @@ pub fn main(init: std.process.Init) !u8 { defer moss_image.destroy(); const idol = try @import("metal_idol.zig").Suite.create(a, device, queue, &kernels, o.only); defer idol.destroy(); + const idol_dino = try @import("metal_idol_dino.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol_dino.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { From 900b7312db75d3ed55e45b77550ca8dc3844da5e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:12:57 +0800 Subject: [PATCH 15/21] changelog: #1775 for Idol's DINOv3 encoder --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 12f41a7149..8ed2301c91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol (speridlabs/idol, image to 3D Gaussian splat), first step: its 0.37B denoiser runs natively on M5 Metal from the released checkpoint, the condition refiner once an image and all 50 guided steps on the GPU. On an M5 Ultra a forward takes 0.275 s against the publisher's SDK on torch MPS at 0.939 s, with a 1.84 GB peak footprint. Each step's velocities on the SDK's own input sit as close to an fp32 run as the SDK's own, and each joint block alone passes the same gate; `tf-idol` runs those gates from a case `tools/zig/idol_case.py` records. RePo's fp32 maps run as an fp16 hi/lo split on the tensor units, whose fp32 GEMM truncates its inputs to 10 mantissa bits (#1772). The image encoders and the decoders are not ported yet (#1766, fixes #1757). Thanks to @jayleaton. -- Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (after #1766, for #1757). Thanks to @jayleaton. +- Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (#1775, after #1766, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. From 25b6e6f5097f56b7ad8e81da73072741257cf4bc Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:36:26 +0800 Subject: [PATCH 16/21] idol: the vae case moves to idol_vae_case.py, as dino's did, so the merged case tool stays under 600 lines The records file's header comment is one line again. --- tools/zig/idol_case.py | 142 +----------------------- tools/zig/idol_vae_case.py | 148 +++++++++++++++++++++++++ zig/src/core/contract/idol_records.zig | 3 +- 3 files changed, 154 insertions(+), 139 deletions(-) create mode 100644 tools/zig/idol_vae_case.py diff --git a/tools/zig/idol_case.py b/tools/zig/idol_case.py index cc0ab0ede7..d295e23893 100644 --- a/tools/zig/idol_case.py +++ b/tools/zig/idol_case.py @@ -42,21 +42,6 @@ "idol/camera.py": "5576043593454e492403b821d1e355433e71f515a2d4406a51c53c30fc097935", } WEIGHTS_REVISION = "2731c2291fc65a797c5cb627a5e5526fc1490902" -CONDITIONING_SHA256 = "8c9cc50d327df3a5269cda392ef46cf4ed9c1e603f42f46050beda83f355f5ab" -DIFFUSERS_VERSION = "0.42.1" -DIFFUSERS_SHA256 = { - "models/autoencoders/autoencoder_kl_flux2.py": "1afe08d76a93e994a453bd9036b57d173516211c2e6b9794ea64f485bea57f57", - "models/autoencoders/vae.py": "59b51fdf2d951da45d697569271e0c99adf4a959ec0f91885f3d9314e925a12c", - "models/unets/unet_2d_blocks.py": "545d158c9d98fb8971c501303117fd3e32d27ae28506b89005c65548404f7437", - "models/resnet.py": "77ad799719fa51bf0fc703f0486915b791c9b69673983abeb3e32c6edcc76650", - "models/downsampling.py": "bce174aa8007b2c40787331e7f46b31eef1eab7785f16b4e54733553b3d14724", - "models/attention_processor.py": "59871fd4e66fc14a26f0e0464706ab24680e13dcd93161c55123678a5e671772", - "utils/torch_utils.py": "bce0a1a4488557e8f80dc02de5782c2b510e91ba2e07f7ae8d26168cee9afc2a", -} -# FLUX.2-klein-4B's vae/ (Apache 2.0, ungated) at its pinned revision, and FLUX.2-dev's: the same weights once in bf16 -VAE_REVISION = "e7b7dc27f91deacad38e78976d1f2b499d76a294" -VAE_SHA256 = {"ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04": "FLUX.2-klein-4B", - "d64f3a68e1cc4f9f4e29b6e0da38a0204fe9a49f2d4053f0ec1fa1ca02f9c4b5": "FLUX.2-dev"} DENOISER_ARGS = {"q_token_length": 8192, "in_channels": 16, "cam_channels": 5, "model_channels": 1024, "cond_channels": 1280, "cond2_channels": 128, "num_blocks": 24, "num_refiner_blocks": 2, "num_heads": 16, "mlp_ratio": 4} @@ -510,127 +495,6 @@ def recorded_like(*a, **k): return 0 -def flux2_vae(vae_dir: str, sdk_dir: str): - """The SDK's ImageEncoder and diffusers' FLUX.2 VAE (bf16, as the SDK loads it), each file used checked first.""" - import diffusers - import torch - - def pinned(path: Path, want: str, what: str) -> None: - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != want: - sys.exit(f"{path} has sha256 {digest}; this tool runs only {what}'s ({want})") - - pinned(Path(sdk_dir) / "idol/conditioning.py", CONDITIONING_SHA256, f"idol-sdk {SDK_COMMIT[:12]}") - if diffusers.__version__ != DIFFUSERS_VERSION: - sys.exit(f"diffusers is {diffusers.__version__}; this tool runs only {DIFFUSERS_VERSION}") - root = Path(diffusers.__file__).parent - for name, want in DIFFUSERS_SHA256.items(): - pinned(root / name, want, f"diffusers {DIFFUSERS_VERSION}") - weights = Path(vae_dir) / "diffusion_pytorch_model.safetensors" - digest = hashlib.sha256(weights.read_bytes()).hexdigest() - if digest not in VAE_SHA256: - sys.exit(f"{weights} is neither FLUX.2-klein-4B {VAE_REVISION[:12]}'s vae/ nor FLUX.2-dev's") - sys.path.insert(0, str(Path(sdk_dir))) - from diffusers.models import AutoencoderKLFlux2 - from idol.conditioning import ImageEncoder - - vae = AutoencoderKLFlux2.from_pretrained(vae_dir, torch_dtype=torch.bfloat16) - return ImageEncoder, vae.eval().requires_grad_(False), VAE_SHA256[digest] - - -def tokens64(vae, pixels, noise): - """ImageEncoder._flux2_tokens's steps in float64 with the SDK's noise: the ruler, and each stage when asked.""" - import torch - - latents = vae.encode(pixels * 2 - 1).latent_dist - latents = latents.mean + latents.std * noise - batch, channels, height, width = latents.shape - latents = latents.view(batch, channels, height // 2, 2, width // 2, 2).permute(0, 1, 3, 5, 2, 4) - latents = latents.reshape(batch, channels * 4, height // 2, width // 2) - mean = vae.bn.running_mean.view(1, -1, 1, 1) - std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps) - rows = ((latents - mean) / std).flatten(2).transpose(1, 2) - return torch.cat([torch.zeros(batch, 5, rows.shape[2], dtype=rows.dtype), rows], dim=1) - - -def distance(ours, ruler) -> dict: - """1 - cosine and the worst row's relative L2 error against the ruler, the rows of 128 as tf-idol-vae reads them.""" - a, b = ours.double().reshape(-1, 128), ruler.double().reshape(-1, 128) - cos = float((a * b).sum() / (a.norm() * b.norm())) - rows = ((a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)).max() - return {"one_minus_cos": 1 - cos, "worst_row": float(rows), "max_abs": float((a - b).abs().max())} - - -def vae_case(args) -> int: - """The SDK's feature2 for one image on MPS (seeded as IdolPipeline seeds it), its noise, and the float64 ruler.""" - import types - - import diffusers.models.autoencoders.vae as vae_module - import numpy as np - import torch - from PIL import Image - from safetensors.torch import save_file - - out = Path(args.out_dir) - out.mkdir(parents=True, exist_ok=True) - image_encoder, vae, source = flux2_vae(args.vae, args.sdk) - picture = Image.open(args.image).convert("RGB") - if picture.size != (1024, 1024): - sys.exit(f"{args.image} is {picture.size}; the encoders see the 1024 x 1024 picture preprocess returns") - pixels = torch.from_numpy(np.array(picture)).permute(2, 0, 1).float().div(255.0)[None] - encoder = types.SimpleNamespace(flux2_vae=vae.to(args.device)) - drawn = [] - randn = vae_module.randn_tensor - - def recorded(*a, **k): - drawn.append(randn(*a, **k)) - return drawn[-1] - - vae_module.randn_tensor = recorded - times = [] - with torch.no_grad(): - for i in range(args.repeats + 1): - torch.manual_seed(args.seed) - mark = time.perf_counter() - rows = image_encoder._flux2_tokens(encoder, pixels.to(args.device)) - flush(args.device) - times.append(time.perf_counter() - mark) - if i == 0: - sdk_rows, noise = rows[0].cpu(), drawn[0].float().cpu() - vae_module.randn_tensor = randn - if any(not torch.equal(d.float().cpu(), noise) for d in drawn): - sys.exit("the seeded noise differs between encodes") - stages = {} - hooks = [] - if args.stages: - named = [(f"down.{i}", b) for i, b in enumerate(vae.encoder.down_blocks)] + [("mid", vae.encoder.mid_block)] - named.append(("moments", vae.quant_conv)) - for name, module in named: - def keep(_m, _i, output, name=name): - stages[name] = output[0].permute(1, 2, 0).float().contiguous() - hooks.append(module.register_forward_hook(keep)) - vae64 = vae.to("cpu").double() - with torch.no_grad(): - mark = time.perf_counter() - ruler = tokens64(vae64, pixels.double(), noise.double())[0] - ruler_s = time.perf_counter() - mark - for h in hooks: - h.remove() - tensors = {"image": pixels[0].permute(1, 2, 0).contiguous(), "noise": noise[0].permute(1, 2, 0).contiguous(), - "sdk": sdk_rows.float().contiguous(), "ruler": ruler.float().contiguous()} - save_file(tensors, str(out / "vae.safetensors")) - if stages: - save_file(stages, str(out / "vae_stages.safetensors")) - sdk = distance(sdk_rows, ruler) - meta = {"seed": args.seed, "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, - "diffusers": DIFFUSERS_VERSION, "vae": source, "vae_revision": VAE_REVISION, - "sdk_vs_float64": sdk, "float64_s": round(ruler_s, 2), "encode_s": [round(t, 4) for t in times]} - (out / "vae.json").write_text(json.dumps(meta, indent=1) + "\n") - warm = sorted(times[1:]) or times - print(json.dumps(meta | {"median_encode_s": warm[len(warm) // 2]})) - return 0 - - def main(argv: list[str]) -> int: ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) @@ -680,7 +544,11 @@ def main(argv: list[str]) -> int: sys.path.insert(0, str(Path(__file__).resolve().parent)) from idol_dino_case import dino return dino(args, sdk_root=args.sdk) - return {"case": case, "blocks": blocks, "exact": exact, "decode": decode, "vae": vae_case}[args.cmd](args) + if args.cmd == "vae": + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from idol_vae_case import vae_case + return vae_case(args) + return {"case": case, "blocks": blocks, "exact": exact, "decode": decode}[args.cmd](args) if __name__ == "__main__": diff --git a/tools/zig/idol_vae_case.py b/tools/zig/idol_vae_case.py new file mode 100644 index 0000000000..5cfa4ef1bd --- /dev/null +++ b/tools/zig/idol_vae_case.py @@ -0,0 +1,148 @@ +"""idol_case.py vae: the SDK's FLUX.2 VAE rows for a preprocessed picture, its posterior noise and a float64 ruler.""" + +from __future__ import annotations + +import hashlib +import json +import sys +import time +from pathlib import Path + +from idol_case import SDK_COMMIT, flush + +CONDITIONING_SHA256 = "8c9cc50d327df3a5269cda392ef46cf4ed9c1e603f42f46050beda83f355f5ab" +DIFFUSERS_VERSION = "0.42.1" +DIFFUSERS_SHA256 = { + "models/autoencoders/autoencoder_kl_flux2.py": "1afe08d76a93e994a453bd9036b57d173516211c2e6b9794ea64f485bea57f57", + "models/autoencoders/vae.py": "59b51fdf2d951da45d697569271e0c99adf4a959ec0f91885f3d9314e925a12c", + "models/unets/unet_2d_blocks.py": "545d158c9d98fb8971c501303117fd3e32d27ae28506b89005c65548404f7437", + "models/resnet.py": "77ad799719fa51bf0fc703f0486915b791c9b69673983abeb3e32c6edcc76650", + "models/downsampling.py": "bce174aa8007b2c40787331e7f46b31eef1eab7785f16b4e54733553b3d14724", + "models/attention_processor.py": "59871fd4e66fc14a26f0e0464706ab24680e13dcd93161c55123678a5e671772", + "utils/torch_utils.py": "bce0a1a4488557e8f80dc02de5782c2b510e91ba2e07f7ae8d26168cee9afc2a", +} +# FLUX.2-klein-4B's vae/ (Apache 2.0, ungated) at its pinned revision, and FLUX.2-dev's: the same weights once in bf16 +VAE_REVISION = "e7b7dc27f91deacad38e78976d1f2b499d76a294" +VAE_SHA256 = {"ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04": "FLUX.2-klein-4B", + "d64f3a68e1cc4f9f4e29b6e0da38a0204fe9a49f2d4053f0ec1fa1ca02f9c4b5": "FLUX.2-dev"} + + +def flux2_vae(vae_dir: str, sdk_dir: str): + """The SDK's ImageEncoder and diffusers' FLUX.2 VAE (bf16, as the SDK loads it), each file used checked first.""" + import diffusers + import torch + + def pinned(path: Path, want: str, what: str) -> None: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != want: + sys.exit(f"{path} has sha256 {digest}; this tool runs only {what}'s ({want})") + + pinned(Path(sdk_dir) / "idol/conditioning.py", CONDITIONING_SHA256, f"idol-sdk {SDK_COMMIT[:12]}") + if diffusers.__version__ != DIFFUSERS_VERSION: + sys.exit(f"diffusers is {diffusers.__version__}; this tool runs only {DIFFUSERS_VERSION}") + root = Path(diffusers.__file__).parent + for name, want in DIFFUSERS_SHA256.items(): + pinned(root / name, want, f"diffusers {DIFFUSERS_VERSION}") + weights = Path(vae_dir) / "diffusion_pytorch_model.safetensors" + digest = hashlib.sha256(weights.read_bytes()).hexdigest() + if digest not in VAE_SHA256: + sys.exit(f"{weights} is neither FLUX.2-klein-4B {VAE_REVISION[:12]}'s vae/ nor FLUX.2-dev's") + sys.path.insert(0, str(Path(sdk_dir))) + from diffusers.models import AutoencoderKLFlux2 + from idol.conditioning import ImageEncoder + + vae = AutoencoderKLFlux2.from_pretrained(vae_dir, torch_dtype=torch.bfloat16) + return ImageEncoder, vae.eval().requires_grad_(False), VAE_SHA256[digest] + + +def tokens64(vae, pixels, noise): + """ImageEncoder._flux2_tokens's steps in float64 with the SDK's noise: the ruler, and each stage when asked.""" + import torch + + latents = vae.encode(pixels * 2 - 1).latent_dist + latents = latents.mean + latents.std * noise + batch, channels, height, width = latents.shape + latents = latents.view(batch, channels, height // 2, 2, width // 2, 2).permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch, channels * 4, height // 2, width // 2) + mean = vae.bn.running_mean.view(1, -1, 1, 1) + std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps) + rows = ((latents - mean) / std).flatten(2).transpose(1, 2) + return torch.cat([torch.zeros(batch, 5, rows.shape[2], dtype=rows.dtype), rows], dim=1) + + +def distance(ours, ruler) -> dict: + """1 - cosine and the worst row's relative L2 error against the ruler, the rows of 128 as tf-idol-vae reads them.""" + a, b = ours.double().reshape(-1, 128), ruler.double().reshape(-1, 128) + cos = float((a * b).sum() / (a.norm() * b.norm())) + rows = ((a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)).max() + return {"one_minus_cos": 1 - cos, "worst_row": float(rows), "max_abs": float((a - b).abs().max())} + + +def vae_case(args) -> int: + """The SDK's feature2 for one image on MPS (seeded as IdolPipeline seeds it), its noise, and the float64 ruler.""" + import types + + import diffusers.models.autoencoders.vae as vae_module + import numpy as np + import torch + from PIL import Image + from safetensors.torch import save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + image_encoder, vae, source = flux2_vae(args.vae, args.sdk) + picture = Image.open(args.image).convert("RGB") + if picture.size != (1024, 1024): + sys.exit(f"{args.image} is {picture.size}; the encoders see the 1024 x 1024 picture preprocess returns") + pixels = torch.from_numpy(np.array(picture)).permute(2, 0, 1).float().div(255.0)[None] + encoder = types.SimpleNamespace(flux2_vae=vae.to(args.device)) + drawn = [] + randn = vae_module.randn_tensor + + def recorded(*a, **k): + drawn.append(randn(*a, **k)) + return drawn[-1] + + vae_module.randn_tensor = recorded + times = [] + with torch.no_grad(): + for i in range(args.repeats + 1): + torch.manual_seed(args.seed) + mark = time.perf_counter() + rows = image_encoder._flux2_tokens(encoder, pixels.to(args.device)) + flush(args.device) + times.append(time.perf_counter() - mark) + if i == 0: + sdk_rows, noise = rows[0].cpu(), drawn[0].float().cpu() + vae_module.randn_tensor = randn + if any(not torch.equal(d.float().cpu(), noise) for d in drawn): + sys.exit("the seeded noise differs between encodes") + stages = {} + hooks = [] + if args.stages: + named = [(f"down.{i}", b) for i, b in enumerate(vae.encoder.down_blocks)] + [("mid", vae.encoder.mid_block)] + named.append(("moments", vae.quant_conv)) + for name, module in named: + def keep(_m, _i, output, name=name): + stages[name] = output[0].permute(1, 2, 0).float().contiguous() + hooks.append(module.register_forward_hook(keep)) + vae64 = vae.to("cpu").double() + with torch.no_grad(): + mark = time.perf_counter() + ruler = tokens64(vae64, pixels.double(), noise.double())[0] + ruler_s = time.perf_counter() - mark + for h in hooks: + h.remove() + tensors = {"image": pixels[0].permute(1, 2, 0).contiguous(), "noise": noise[0].permute(1, 2, 0).contiguous(), + "sdk": sdk_rows.float().contiguous(), "ruler": ruler.float().contiguous()} + save_file(tensors, str(out / "vae.safetensors")) + if stages: + save_file(stages, str(out / "vae_stages.safetensors")) + sdk = distance(sdk_rows, ruler) + meta = {"seed": args.seed, "device": args.device, "torch": torch.__version__, "sdk": SDK_COMMIT, + "diffusers": DIFFUSERS_VERSION, "vae": source, "vae_revision": VAE_REVISION, + "sdk_vs_float64": sdk, "float64_s": round(ruler_s, 2), "encode_s": [round(t, 4) for t in times]} + (out / "vae.json").write_text(json.dumps(meta, indent=1) + "\n") + warm = sorted(times[1:]) or times + print(json.dumps(meta | {"median_encode_s": warm[len(warm) // 2]})) + return 0 diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig index ce53fb82fb..124347e4e1 100644 --- a/zig/src/core/contract/idol_records.zig +++ b/zig/src/core/contract/idol_records.zig @@ -1,5 +1,4 @@ -//! Idol's launches, each source-reserved: the denoiser's (idol.metal, idol_attention.metal), then idol_decoder.metal's, -//! then the image encoder's (idol_vae.metal). +//! Idol's launches, each source-reserved: the denoiser's (idol.metal, idol_attention.metal), the decoders', the VAE's. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; pub const records = denoiser ++ decoder ++ vae; From 88335348084c4ae57be32bd99a62d4ab066d5e86 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:42:15 +0800 Subject: [PATCH 17/21] idol: tf-idol-generate runs one picture end to end on Metal, both encoders, the denoiser and the decoders to a PLY Each stage is timed and each model freed after it. The starting noise is torch's CPU-seeded randn, reimplemented; the VAE's posterior noise and the starting noise can be given for parity with the SDK's MPS draws. The recipe's End to end section has the stage timings, memory and PLY comparison against the SDK's own IdolPipeline run. --- CHANGELOG.md | 1 + docs/recipes/idol.md | 116 +++++++++++ zig/build/metal_targets.zig | 1 + zig/tests/idol_generate.zig | 385 ++++++++++++++++++++++++++++++++++++ 4 files changed, 503 insertions(+) create mode 100644 zig/tests/idol_generate.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c75e85421..e91429136f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol's octree and Gaussian decoders now run natively on M5 Metal, from the denoiser's latent to the 3DGS PLY and camera. On an M5 Ultra, 262,144 Gaussians decode in 0.343 s against the SDK's 0.739 s, and 1,048,576 in 1.963 s against 4.653 s, at a 2.43 GB peak footprint against 5.64 GB. Given the SDK's probabilities and draws, every octree level's counts match, and the PLY's header and rows equal the SDK's. Both decoders sit as close to an fp32 run as the SDK's own; `tf-idol-decode` runs those gates from a `tools/zig/idol_case.py decode` case. The decoders' nine kernels are registered with the kernel contract (#1773, after #1766, for #1757). Thanks to @jayleaton. - Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (#1774, after #1766, for #1757). Thanks to @jayleaton. - Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (#1775, after #1766, for #1757). Thanks to @jayleaton. +- Idol runs end to end natively on M5 Metal: `tf-idol-generate` takes one preprocessed picture through DINOv3, the FLUX.2 VAE, the denoiser's 50 guided steps and both decoders to a 3DGS PLY and camera. On an M5 Ultra it is 29.1 s against 95.4 s for the SDK's IdolPipeline on MPS (excluding its preprocess), the denoiser 27.5 s against 93.9 s, at a 2.71 GB peak footprint against 12.88 GB. Given the SDK's starting and VAE noise the PLY's points sit as close to the SDK's as sampling resolves, and the camera's angles within 0.006 rad (after #1773, #1774 and #1775, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index b749f1a508..5aa0542761 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -476,3 +476,119 @@ The avocado example, the same machine and torch as above (`ATTN_BACKEND=sdpa`): (median of 10 warm; its first 0.385 s): 2.2x. Load 0.2 to 0.6 s. - **Memory**: 2.28 GB peak process footprint (`/usr/bin/time -l`, `tf-idol-dino` with the case's feature rows), the bf16 weights 1.68 GB of it. The SDK's MPS encode holds 2.24 GB of MPS driver memory. + +## End to end + +`tf-idol-generate` (`zig/tests/idol_generate.zig`) chains the pieces above, from one preprocessed picture to a PLY: + +1. DINOv3's feature1; +2. the VAE's feature2; +3. the denoiser's condition and its 50 guided steps; +4. the octree and Gaussian decoders; +5. the PLY and the camera. + +Each model is loaded when its stage starts and freed when it ends. Background removal and cropping (BiRefNet, the SDK's +`preprocess`) are not ported, so the input is the picture `preprocess` returns, as a `pixels` or `image` tensor that +`idol_case.py dino` or `vae` writes. + +```sh +zig build tf-idol-generate -Dcpu=apple_m1 +zig-out/bin/tf-idol-generate WEIGHTS KLEIN/vae DINO case/vae.safetensors out/ # seed 42, 50 steps, cfg 3, 262,144 +zig-out/bin/tf-idol-generate WEIGHTS KLEIN/vae DINO case/vae.safetensors out/ --noise case/vae.safetensors +zig-out/bin/tf-idol-generate WEIGHTS KLEIN/vae DINO - out/ --features dino.safetensors --features vae.safetensors +``` + +Noise: + +- **The starting noise** is `torch.Generator("cpu").manual_seed(seed)`'s randn: [8192, 16], then [5], as + `idol_case.py case` draws it. The tool reimplements torch's mt19937 and its Box-Muller. The camera's 5 values are + bit-equal to torch's. 28 of the 131,072 latent values are 1 ulp off, from torch's float cos and sin. +- **The VAE posterior's noise and the octree's draws** come from native seeded generators. +- **`--noise FILE`** gives the posterior's noise (`noise`, as `idol_case.py vae` records it). If FILE also holds + `latent` and `camera`, they replace the starting noise. + +It writes `generate.ply`, `camera.json`, `start.f32`, `sample.f32` and `tf-idol-generate.json`, each stage's seconds. + +### Measured + +Setup: + +- The avocado example and the same M5 Ultra and torch 2.13. +- Seed 42, 50 steps, guidance 3, 262,144 Gaussians. +- The SDK runs its own `IdolPipeline.__call__` on MPS once (`optimize` is CUDA-only), each stage timed by wrapping + the SDK's methods with `torch.mps.synchronize()`. Its PLY is byte-identical to an earlier run. +- Ours runs twice, each in a fresh process. The weights were in the page cache both times, so cold and warm differ + only by the process. + +| stage | ours, cold | ours, warm | SDK, MPS | +| --- | ---: | ---: | ---: | +| load (all models) | 0.67 s | 0.67 s | 2.11 s (and 2.75 s of imports) | +| preprocess (BiRefNet) | not ported | not ported | 0.51 s | +| DINOv3 | 0.158 s | 0.155 s | 0.275 s | +| VAE | 0.245 s | 0.241 s | 0.119 s | +| condition | 0.072 s | 0.071 s | in the sampling | +| denoise, 50 steps x 2 | 27.49 s | 27.50 s | 93.89 s | +| octree | 0.145 s | 0.143 s | 0.612 s | +| Gaussian decoder (+ activation) | 0.222 s (+ 0.001) | 0.222 s (+ 0.001) | 0.491 s (+ 0.013) | +| PLY and camera | 0.004 s | 0.004 s | 0.018 s | +| total | 29.09 s | 29.07 s | 100.81 s (95.41 s generating, less preprocess) | + +Memory, the whole process under `/usr/bin/time -l`: + +| | peak footprint | max RSS | GPU | +| --- | ---: | ---: | --- | +| ours | 2.71 GB | 1.94 GB | 2.40 GB Metal allocated at most | +| SDK | 12.88 GB | 6.20 GB | 5.01 GB MPS allocated, 12.64 GB driver (sampled) | + +### Quality + +The PLYs cannot be equal: + +- **Starting noise:** the SDK draws it from MPS's generator after the VAE's draw, which cannot be reproduced natively. +- **VAE posterior noise:** MPS's generator too, drawn in bf16. +- **The octree's draws:** torch's multinomial and jitter in the SDK, a native generator here. One count's difference + changes the topology from that level down. +- **The denoiser:** fp16 over 50 guided steps, and our encoders are not bit-equal to the SDK's. + +The starting noise matters most. With the same features, the SDK's own sample from MPS noise and its sample from the +CPU-seeded noise have cosine 0.479. Their PLYs are as far apart as ours is from the SDK's. + +So the PLYs are compared as distributions. Statistics are per PLY: xyz in PLY units, opacity and scale activated. The +NN columns, against the SDK's pipeline PLY, are: + +- the mean symmetric nearest-neighbour distance between 4,096 points sampled from each; +- the same over all 262,144 points. + +Two independent 4,096-point samples of the SDK's own PLY lie 0.00896 apart. + +| run | xyz std | opacity mean (q50) | scale q50 | NN, sampled | NN, all points | +| --- | --- | --- | --- | ---: | ---: | +| SDK pipeline | 0.189, 0.078, 0.304 | 0.885 (0.9986) | 0.00096 | 0.00906 | 0 | +| ours, native noise | 0.176, 0.073, 0.320 | 0.423 (0.181) | 0.00113 | 0.04623 | 0.0365 | +| ours, SDK's VAE noise | 0.176, 0.074, 0.320 | 0.425 (0.186) | 0.00113 | 0.04627 | 0.0366 | +| ours, SDK's VAE and starting noise | 0.189, 0.079, 0.304 | 0.887 (0.9986) | 0.00096 | 0.00922 | 0.00166 | +| SDK, same features, CPU-seeded noise | 0.176, 0.074, 0.319 | 0.422 (0.179) | 0.00113 | 0.04617 | 0.0368 | + +The table shows: + +- **Every PLY:** 262,144 Gaussians, all finite. +- **The SDK's VAE noise alone** changes little: the starting noise still differs. +- **With both of the SDK's noises:** our sample has cosine 0.99835 to the SDK's. The PLYs are as close as the sampling + can resolve: 0.00922 against the 0.00896 baseline, 0.0017 over all points. The camera is yaw 1.6956, pitch 0.5397 + and fov 0.1094, against 1.7006, 0.5413 and 0.1096. +- **Repeat:** our cold and warm PLYs are byte-identical. + +### Parity chain + +This run gives the SDK's own recorded feature1 and feature2 (`--features`) and the CPU-seeded noise. The SDK's +reference is the same features and noise: the `idol_case.py case` sample, and its `decode` PLY. + +- **Sample:** cosine 0.999991 to the SDK's. +- **Camera:** its latent's 5 values within 0.0022 of the SDK's. +- **PLY** against the SDK's decode of its own sample: + - mean NN 0.00981 between 4,096-point samples (the baseline 0.00944); + - 0.0018 over all points; + - xyz mean 6e-5 apart. + +From the picture, our encoders move the sample to cosine 0.99437 with native VAE noise. With the SDK's noise, they +move it to 0.99953. The PLY stays 0.0021 over all points from the SDK's decode. diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index dbf26eca58..76b5f721e0 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -321,6 +321,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-idol-decode", .path = "zig/tests/idol_decode.zig", .about = "Idol's octree and Gaussian decoders on a reference case: each level, the rows and the PLY gated, then timed" }, .{ .name = "tf-idol-vae", .path = "zig/tests/idol_vae.zig", .about = "Idol's image encoder on a reference case: its condition rows gated against float64 beside the SDK's, then timed" }, .{ .name = "tf-idol-dino", .path = "zig/tests/idol_dino.zig", .about = "Idol's DINOv3 image encoder on a reference case: feature1 and each block gated against torch's, then timed" }, + .{ .name = "tf-idol-generate", .path = "zig/tests/idol_generate.zig", .about = "Idol end to end: one preprocessed picture to a 3DGS PLY and camera through both encoders, the denoiser and the decoders, timed by stage" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/tests/idol_generate.zig b/zig/tests/idol_generate.zig new file mode 100644 index 0000000000..f6db0c56ed --- /dev/null +++ b/zig/tests/idol_generate.zig @@ -0,0 +1,385 @@ +//! tf-idol-generate: Idol end to end on Metal, one picture to a 3DGS PLY: both encoders, the denoiser, the decoders. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const id = tf.idol; +const cfg = id.config; +const vae = id.vae; +const dn = id.dino; +const dc = id.decoder_config; +const gs = id.gaussian; + +const usage = + \\tf-idol-generate SNAPSHOT_DIR VAE_DIR DINO_DIR PICTURE OUT_DIR [--seed N] [--steps N] [--cfg X] [--gaussians N] + \\ [--noise FILE] [--features FILE]... + \\SNAPSHOT_DIR is a speridlabs/idol snapshot, VAE_DIR FLUX.2-klein-4B's vae/ folder, DINO_DIR + \\timm/vit_huge_plus_patch16_dinov3.lvd1689m. PICTURE is a safetensors file holding the 1,024 x 1,024 picture + \\IdolPipeline.preprocess returns, in [0, 1]: `image` [1024, 1024, 3] (idol_case.py vae) or `pixels` [3, 1024, 1024] + \\(idol_case.py dino). Defaults: seed 42, 50 steps, guidance 3, 262,144 Gaussians, as IdolPipeline's. + \\The denoiser starts from torch.Generator("cpu").manual_seed(seed)'s randn [8192, 16] then [5], as idol_case.py + \\case draws it (the SDK draws on its device's generator). The VAE posterior's noise and the octree's draws come from + \\native seeded generators. --noise FILE gives the posterior's noise instead (`noise` [128, 128, 32], as idol_case.py + \\vae records the SDK's) and, if it holds them, the starting `latent` [8192, 16] and `camera` [5]. --features FILE + \\(once or twice) gives feature1 [1, 4101, 1280] and feature2 [1, 4101, 128] and skips both encoders (PICTURE: -). + \\Writes OUT_DIR/generate.ply, camera.json, start.f32 (the starting latent and camera), sample.f32 and + \\tf-idol-generate.json (each stage's seconds). Each model is freed once its stage is done. +; + +const N = cfg.latent_rows * cfg.latent_channels; +const cam_n = cfg.camera_channels; +const f1_n = cfg.cond_rows * cfg.cond_channels; +const f2_n = cfg.cond_rows * cfg.cond2_channels; +const noise_n = vae.latent * vae.latent * id.vae_weights.latent_channels; +const side = vae.size; +const rescale: f64 = 3.0; + +/// torch's CPU generator: mt19937 seeded as CPUGeneratorImpl seeds it. +const TorchCpu = struct { + s: [624]u32, + i: usize = 624, + /// normal_distribution's cached second sample + spare: ?f64 = null, + + fn init(seed: u64) TorchCpu { + var g: TorchCpu = .{ .s = undefined }; + g.s[0] = @truncate(seed); + for (1..624) |j| g.s[j] = 1812433253 *% (g.s[j - 1] ^ (g.s[j - 1] >> 30)) +% @as(u32, @intCast(j)); + return g; + } + + fn next(g: *TorchCpu) u32 { + if (g.i >= 624) { + for (0..624) |k| { + const y = (g.s[k] & 0x80000000) | (g.s[(k + 1) % 624] & 0x7fffffff); + g.s[k] = g.s[(k + 397) % 624] ^ (y >> 1) ^ (if (y & 1 != 0) @as(u32, 0x9908b0df) else 0); + } + g.i = 0; + } + var y = g.s[g.i]; + g.i += 1; + y ^= y >> 11; + y ^= (y << 7) & 0x9d2c5680; + y ^= (y << 15) & 0xefc60000; + return y ^ (y >> 18); + } + + fn uniform(g: *TorchCpu) f32 { + return @as(f32, @floatFromInt(g.next() & 0xffffff)) * 0x1p-24; + } + + /// normal_fill_16: Box-Muller on 8 pairs in place; theta is rounded to float as torch's scalar_t is. + fn fill16(d: *[16]f32) void { + for (0..8) |j| { + const r = @sqrt(-2 * logf(1 - d[j])); + const theta: f32 = @floatCast(2.0 * std.math.pi * @as(f64, d[j + 8])); + d[j] = r * @as(f32, @floatCast(@cos(@as(f64, theta)))); + d[j + 8] = r * @as(f32, @floatCast(@sin(@as(f64, theta)))); + } + } + + /// torch.randn(out.len, generator=g) for float32: normal_fill from 16 values on, else normal_distribution. + fn randn(g: *TorchCpu, out: []f32) void { + if (out.len >= 16) { + for (out) |*v| v.* = g.uniform(); + var i: usize = 0; + while (i + 16 <= out.len) : (i += 16) fill16(out[i..][0..16]); + if (out.len % 16 != 0) { + const tail = out[out.len - 16 ..][0..16]; + for (tail) |*v| v.* = g.uniform(); + fill16(tail); + } + return; + } + for (out) |*v| v.* = @floatCast(g.normal()); + } + + fn normal(g: *TorchCpu) f64 { + if (g.spare) |s| { + g.spare = null; + return s; + } + const u_1 = g.uniform53(); + const u_2 = g.uniform53(); + const r = @sqrt(-2 * std.math.log1p(-u_2)); + const theta = 2 * std.math.pi * u_1; + g.spare = r * @sin(theta); + return r * @cos(theta); + } + + fn uniform53(g: *TorchCpu) f64 { + const hi: u64 = g.next(); + const lo: u64 = g.next(); + return @as(f64, @floatFromInt(((hi << 32) | lo) & ((1 << 53) - 1))) * 0x1p-53; + } +}; + +extern "c" fn logf(x: f32) f32; + +const Options = struct { + seed: u64 = 42, + steps: usize = 50, + guidance: f32 = 3.0, + gaussians: usize = 262144, + noise: ?[]const u8 = null, + features: [2]?[]const u8 = .{ null, null }, +}; + +fn parse(args: []const [:0]const u8) !Options { + var o: Options = .{}; + var i: usize = 6; + while (i < args.len) : (i += 2) { + if (i + 1 >= args.len) return error.Usage; + const k = args[i]; + const v = args[i + 1]; + if (std.mem.eql(u8, k, "--seed")) { + o.seed = try std.fmt.parseInt(u64, v, 10); + } else if (std.mem.eql(u8, k, "--steps")) { + o.steps = try std.fmt.parseInt(usize, v, 10); + } else if (std.mem.eql(u8, k, "--cfg")) { + o.guidance = try std.fmt.parseFloat(f32, v); + } else if (std.mem.eql(u8, k, "--gaussians")) { + o.gaussians = try std.fmt.parseInt(usize, v, 10); + } else if (std.mem.eql(u8, k, "--noise")) { + o.noise = v; + } else if (std.mem.eql(u8, k, "--features")) { + if (o.features[0] == null) o.features[0] = v else if (o.features[1] == null) o.features[1] = v else return error.Usage; + } else return error.Usage; + } + if (o.steps == 0) return error.Usage; + return o; +} + +/// Tensor `name` of `f`, float32 with `n` values, or null when absent. +fn maybe(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, n: usize) !?[]f32 { + const t = f.get(name) orelse return null; + if (t.dtype != .f32 or t.numel() != n) { + std.debug.print("{s} is not {d} float32 values\n", .{ name, n }); + return error.BadTensor; + } + const out = try a.alloc(f32, n); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +/// The picture as [3, 1024, 1024] (DINOv3's input) and [1024, 1024, 3] (the VAE's). +fn picture(a: std.mem.Allocator, io: std.Io, path: []const u8) ![2][]f32 { + var f = try tf.safetensors.File.open(a, io, path); + defer f.close(io); + const n = 3 * side * side; + const chw = try a.alloc(f32, n); + const hwc = try a.alloc(f32, n); + if (try maybe(a, &f, "pixels", n)) |p| { + @memcpy(chw, p); + for (0..side * side) |i| for (0..3) |c| { + hwc[i * 3 + c] = p[c * side * side + i]; + }; + } else if (try maybe(a, &f, "image", n)) |p| { + @memcpy(hwc, p); + for (0..side * side) |i| for (0..3) |c| { + chw[c * side * side + i] = p[i * 3 + c]; + }; + } else { + std.debug.print("{s} holds neither `pixels` [3, 1024, 1024] nor `image` [1024, 1024, 3]\n", .{path}); + return error.MissingTensor; + } + return .{ chw, hwc }; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn gb(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / 1e9; +} + +/// Seconds per stage, as printed and written. +const Times = struct { + load_dino: f64 = 0, + dino: f64 = 0, + load_vae: f64 = 0, + vae: f64 = 0, + load_denoiser: f64 = 0, + condition: f64 = 0, + denoise: f64 = 0, + load_decoders: f64 = 0, + octree: f64 = 0, + gaussian_decoder: f64 = 0, + activation: f64 = 0, + ply: f64 = 0, + total: f64 = 0, +}; + +/// feature1 and feature2 from the encoders, or from --features files. +fn encode(a: std.mem.Allocator, io: std.Io, device: mtl.Device, args: []const [:0]const u8, o: Options, t: *Times, metal_peak: *u64) ![2][]f32 { + if (o.features[0]) |first| { + var f1: ?[]f32 = null; + var f2: ?[]f32 = null; + for ([_]?[]const u8{ first, o.features[1] }) |maybe_path| if (maybe_path) |path| { + var f = try tf.safetensors.File.open(a, io, path); + defer f.close(io); + if (try maybe(a, &f, "feature1", f1_n)) |v| f1 = v; + if (try maybe(a, &f, "feature2", f2_n)) |v| f2 = v; + }; + if (f1 == null or f2 == null) { + std.debug.print("--features needs feature1 [1, 4101, 1280] and feature2 [1, 4101, 128]\n", .{}); + return error.MissingTensor; + } + std.debug.print("features: given, both encoders skipped\n", .{}); + return .{ f1.?, f2.? }; + } + const pics = try picture(a, io, args[4]); + const f1 = try a.alloc(f32, f1_n); + const f2 = try a.alloc(f32, f2_n); + { + var started = mtl.clock.seconds(); + const enc = try dn.Encoder.load(std.heap.page_allocator, io, device, args[3]); + defer enc.deinit(); + t.load_dino = mtl.clock.seconds() - started; + started = mtl.clock.seconds(); + const tm = try enc.encode(pics[0], f1); + t.dino = mtl.clock.seconds() - started; + metal_peak.* = @max(metal_peak.*, device.allocated()); + std.debug.print("DINOv3: loaded in {d:.2} s; feature1 in {d:.3} s ({d:.3} s on the GPU)\n", .{ t.load_dino, t.dino, tm.gpu }); + } + const noise = try a.alloc(f32, noise_n); + var given = false; + if (o.noise) |path| { + var f = try tf.safetensors.File.open(a, io, path); + defer f.close(io); + if (try maybe(a, &f, "noise", noise_n)) |v| { + @memcpy(noise, v); + given = true; + } + } + if (!given) { + var rng = std.Random.DefaultPrng.init(o.seed +% 0x9e3779b97f4a7c15); + for (noise) |*v| v.* = rng.random().floatNorm(f32); + } + { + var started = mtl.clock.seconds(); + const enc = try vae.Encoder.load(std.heap.page_allocator, io, device, args[2]); + defer enc.deinit(); + t.load_vae = mtl.clock.seconds() - started; + started = mtl.clock.seconds(); + const tm = try enc.encode(pics[1], noise, f2, .{}); + t.vae = mtl.clock.seconds() - started; + metal_peak.* = @max(metal_peak.*, device.allocated()); + std.debug.print("VAE: loaded in {d:.2} s; feature2 in {d:.3} s ({d:.3} s on the GPU), posterior noise {s}\n", .{ t.load_vae, t.vae, tm.gpu, if (given) "given" else "native" }); + } + return .{ f1, f2 }; +} + +/// The guided schedule from the starting noise `x` (latent then camera), in place. +fn denoise(a: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8, f: [2][]f32, x: []f32, o: Options, t: *Times, metal_peak: *u64) !void { + var started = mtl.clock.seconds(); + const model = try id.Model.load(std.heap.page_allocator, io, device, dir); + defer model.deinit(); + t.load_denoiser = mtl.clock.seconds() - started; + started = mtl.clock.seconds(); + var conds: [2]id.Condition = undefined; + conds[0] = try model.condition(f[0], f[1]); + defer conds[0].deinit(); + @memset(f[0], 0); + @memset(f[1], 0); + conds[1] = try model.condition(f[0], f[1]); + defer conds[1].deinit(); + t.condition = mtl.clock.seconds() - started; + const ts = try a.alloc(f64, o.steps + 1); + id.host.schedule(ts, rescale); + const vc = try a.alloc(f32, N); + const vu = try a.alloc(f32, N); + var cc: [cam_n]f32 = undefined; + var cu: [cam_n]f32 = undefined; + var gpu: f64 = 0; + const w = o.guidance; + started = mtl.clock.seconds(); + for (0..o.steps) |s| { + const time: f32 = @floatCast(1000 * ts[s]); + const cam: *const [cam_n]f32 = x[N..][0..cam_n]; + const p = try model.forward(&conds[0], x[0..N], cam, time, vc, &cc); + const q = try model.forward(&conds[1], x[0..N], cam, time, vu, &cu); + const by: f32 = @floatCast(ts[s] - ts[s + 1]); + for (x[0..N], vc, vu) |*v, pc, pu| v.* -= ((1 + w) * pc - w * pu) * by; + for (x[N..], cc, cu) |*v, pc, pu| v.* -= ((1 + w) * pc - w * pu) * by; + gpu += p.gpu + q.gpu; + } + t.denoise = mtl.clock.seconds() - started; + metal_peak.* = @max(metal_peak.*, device.allocated()); + std.debug.print("denoiser: loaded in {d:.2} s; condition {d:.3} s; {d} steps in {d:.2} s ({d:.2} s on the GPU)\n", .{ t.load_denoiser, t.condition, o.steps, t.denoise, gpu }); +} + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + const o = (if (args.len >= 6) parse(args) else error.Usage) catch { + std.debug.print("{s}\n", .{usage}); + return 2; + }; + const out_dir = args[5]; + try std.Io.Dir.cwd().createDirPath(io, out_dir); + const begun = mtl.clock.seconds(); + var t: Times = .{}; + var metal_peak: u64 = 0; + + // the starting noise first: it does not depend on the picture + const x = try a.alloc(f32, N + cam_n); + var gen = TorchCpu.init(o.seed); + gen.randn(x[0..N]); + gen.randn(x[N..]); + var start_given = false; + if (o.noise) |path| { + var f = try tf.safetensors.File.open(a, io, path); + defer f.close(io); + const lat = try maybe(a, &f, "latent", N); + const cam = try maybe(a, &f, "camera", cam_n); + if (lat != null and cam != null) { + @memcpy(x[0..N], lat.?); + @memcpy(x[N..], cam.?); + start_given = true; + } + } + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/start.f32", .{out_dir}), .data = std.mem.sliceAsBytes(x) }); + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + std.debug.print("{s}: seed {d}, {d} steps, guidance {d}, {d} Gaussians; starting noise {s}\n", .{ device.name(), o.seed, o.steps, o.guidance, o.gaussians, if (start_given) "given" else "torch CPU-seeded" }); + const features = try encode(a, io, device, args, o, &t, &metal_peak); + try denoise(a, io, device, args[1], features, x, o, &t, &metal_peak); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/sample.f32", .{out_dir}), .data = std.mem.sliceAsBytes(x) }); + + var started = mtl.clock.seconds(); + const d = try id.Decoder.load(std.heap.page_allocator, io, device, args[1], try dc.anchors(o.gaussians)); + t.load_decoders = mtl.clock.seconds() - started; + var draws: id.octree.Draws = .{ .native = std.Random.DefaultPrng.init(o.seed) }; + var r = try id.decode.decode(std.heap.page_allocator, d, x[0..N], o.gaussians, &draws, null); + defer r.deinit(std.heap.page_allocator); + metal_peak = @max(metal_peak, device.allocated()); + d.deinit(); + t.octree = r.time.octree; + t.gaussian_decoder = r.time.gaussian; + t.activation = r.time.activation; + std.debug.print("decoders: loaded in {d:.2} s; octree {d:.3} s, Gaussian decoder {d:.3} s, activation {d:.3} s\n", .{ t.load_decoders, t.octree, t.gaussian_decoder, t.activation }); + + started = mtl.clock.seconds(); + const bytes = try gs.ply(a, r.gaussians); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/generate.ply", .{out_dir}), .data = bytes }); + const c = gs.camera(x[N..][0..cam_n].*); + const cam_json = try std.fmt.allocPrint(a, "{{\"yaw\": {d}, \"pitch\": {d}, \"radius\": {d}, \"fov\": {d}}}\n", .{ c.yaw, c.pitch, c.radius, c.fov }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/camera.json", .{out_dir}), .data = cam_json }); + t.ply = mtl.clock.seconds() - started; + t.total = mtl.clock.seconds() - begun; + + const load = t.load_dino + t.load_vae + t.load_denoiser + t.load_decoders; + std.debug.print("PLY: {d} Gaussians, {d} bytes in {d:.3} s; camera {s}", .{ r.gaussians.count(), bytes.len, t.ply, cam_json }); + std.debug.print("total {d:.2} s: load {d:.2}, dino {d:.3}, vae {d:.3}, condition {d:.3}, denoise {d:.2}, octree {d:.3}, Gaussian decoder {d:.3} (+ activation {d:.3}), PLY {d:.3}; Metal peak {d:.2} GB\n", .{ t.total, load, t.dino, t.vae, t.condition, t.denoise, t.octree, t.gaussian_decoder, t.activation, t.ply, gb(metal_peak) }); + var out: std.ArrayList(u8) = .empty; + try print(&out, a, "{{\"device\": \"{s}\", \"seed\": {d}, \"steps\": {d}, \"cfg\": {d}, \"gaussians\": {d}, \"start_noise\": \"{s}\", \"vae_noise\": \"{s}\", \"features_given\": {}", .{ device.name(), o.seed, o.steps, o.guidance, r.gaussians.count(), if (start_given) "given" else "torch_cpu", if (o.features[0] != null) "unused" else if (o.noise != null) "given" else "native", o.features[0] != null }); + try print(&out, a, ", \"load_s\": {d:.4}", .{load}); + inline for (@typeInfo(Times).@"struct".field_names) |name| try print(&out, a, ", \"{s}_s\": {d:.4}", .{ name, @field(t, name) }); + try print(&out, a, ", \"metal_peak_allocated_gb\": {d:.3}, \"camera\": {s}}}\n", .{ gb(metal_peak), std.mem.trimEnd(u8, cam_json, "\n") }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-generate.json", .{out_dir}), .data = out.items }); + return 0; +} From 192c4e902b68ec43693712ac1038f9ec678b2e4f Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:43:48 +0800 Subject: [PATCH 18/21] changelog: #1780 for Idol end to end --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e91429136f..8b15466c84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol's octree and Gaussian decoders now run natively on M5 Metal, from the denoiser's latent to the 3DGS PLY and camera. On an M5 Ultra, 262,144 Gaussians decode in 0.343 s against the SDK's 0.739 s, and 1,048,576 in 1.963 s against 4.653 s, at a 2.43 GB peak footprint against 5.64 GB. Given the SDK's probabilities and draws, every octree level's counts match, and the PLY's header and rows equal the SDK's. Both decoders sit as close to an fp32 run as the SDK's own; `tf-idol-decode` runs those gates from a `tools/zig/idol_case.py decode` case. The decoders' nine kernels are registered with the kernel contract (#1773, after #1766, for #1757). Thanks to @jayleaton. - Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (#1774, after #1766, for #1757). Thanks to @jayleaton. - Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (#1775, after #1766, for #1757). Thanks to @jayleaton. -- Idol runs end to end natively on M5 Metal: `tf-idol-generate` takes one preprocessed picture through DINOv3, the FLUX.2 VAE, the denoiser's 50 guided steps and both decoders to a 3DGS PLY and camera. On an M5 Ultra it is 29.1 s against 95.4 s for the SDK's IdolPipeline on MPS (excluding its preprocess), the denoiser 27.5 s against 93.9 s, at a 2.71 GB peak footprint against 12.88 GB. Given the SDK's starting and VAE noise the PLY's points sit as close to the SDK's as sampling resolves, and the camera's angles within 0.006 rad (after #1773, #1774 and #1775, for #1757). Thanks to @jayleaton. +- Idol runs end to end natively on M5 Metal: `tf-idol-generate` takes one preprocessed picture through DINOv3, the FLUX.2 VAE, the denoiser's 50 guided steps and both decoders to a 3DGS PLY and camera. On an M5 Ultra it is 29.1 s against 95.4 s for the SDK's IdolPipeline on MPS (excluding its preprocess), the denoiser 27.5 s against 93.9 s, at a 2.71 GB peak footprint against 12.88 GB. Given the SDK's starting and VAE noise the PLY's points sit as close to the SDK's as sampling resolves, and the camera's angles within 0.006 rad (#1780, after #1773, #1774 and #1775, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. From ee8343a4ae9c1fc238786b52336a0f926ff15dd8 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 06:10:17 +0800 Subject: [PATCH 19/21] idol: BiRefNet background removal and preprocessing natively on Metal BiRefNet (Swin-L, ASPP-deformable decoder) in fp32 on Metal from the Idol snapshot's birefnet.safetensors, the SDK's preprocess on the host byte for byte (Pillow LANCZOS, erosion, crop, paste; torch's CPU bilinear input), the tf-idol-birefnet gate tool, its case tool tools/zig/idol_birefnet_case.py and the kernel contract for every idol_birefnet.metal kernel. --- docs/kernel-contract.md | 4 + tools/zig/idol_birefnet_case.py | 328 ++++++++++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/idol_birefnet.metal | 561 ++++++++++++++++++ zig/kernels/metal/sources.zig | 2 + zig/src/core/contract/idol_records.zig | 19 +- zig/src/families/idol/birefnet.zig | 315 ++++++++++ zig/src/families/idol/birefnet_decoder.zig | 102 ++++ zig/src/families/idol/birefnet_kernels.zig | 161 +++++ zig/src/families/idol/birefnet_reference.zig | 102 ++++ zig/src/families/idol/birefnet_weights.zig | 314 ++++++++++ zig/src/families/idol/idol.zig | 15 + zig/src/families/idol/photo.zig | 120 ++++ zig/src/families/idol/preprocess.zig | 373 ++++++++++++ zig/tests/contract/metal_idol_birefnet.zig | 202 +++++++ .../contract/metal_idol_birefnet_conv.zig | 411 +++++++++++++ .../contract/metal_idol_birefnet_norm.zig | 391 ++++++++++++ .../contract/metal_idol_birefnet_rows.zig | 541 +++++++++++++++++ .../contract/metal_idol_birefnet_window.zig | 186 ++++++ zig/tests/contract/metal_limits.zig | 2 + zig/tests/contract/metal_main.zig | 2 + zig/tests/idol_birefnet.zig | 268 +++++++++ 22 files changed, 4417 insertions(+), 3 deletions(-) create mode 100644 tools/zig/idol_birefnet_case.py create mode 100644 zig/kernels/metal/idol_birefnet.metal create mode 100644 zig/src/families/idol/birefnet.zig create mode 100644 zig/src/families/idol/birefnet_decoder.zig create mode 100644 zig/src/families/idol/birefnet_kernels.zig create mode 100644 zig/src/families/idol/birefnet_reference.zig create mode 100644 zig/src/families/idol/birefnet_weights.zig create mode 100644 zig/src/families/idol/photo.zig create mode 100644 zig/src/families/idol/preprocess.zig create mode 100644 zig/tests/contract/metal_idol_birefnet.zig create mode 100644 zig/tests/contract/metal_idol_birefnet_conv.zig create mode 100644 zig/tests/contract/metal_idol_birefnet_norm.zig create mode 100644 zig/tests/contract/metal_idol_birefnet_rows.zig create mode 100644 zig/tests/contract/metal_idol_birefnet_window.zig create mode 100644 zig/tests/idol_birefnet.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 864a69ceaf..98cef76b69 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -196,6 +196,10 @@ Its image encoder (`metal_idol_vae*.zig`, FLUX.2's VAE encoder in fp32) register GEMM and Winograd), GroupNorm's passes, the 512-wide attention and the condition rows against float64 (12 controls). Metal also checks Idol's DINOv3 encoder (`metal_idol_dino*.zig`): its bf16 projections, LayerNorm and final F.layer_norm, the RoPE heads, LayerScale with the residual, SwiGLU and the attention, each at transformers' bf16 points (11 controls). +Its background remover (`metal_idol_birefnet*.zig`, BiRefNet in fp32) registers the implicit GEMM convolution's nine +served forms (1x1, 3x3, 7x7; batch norm, ReLU, exact GELU, residual by a window map), the picture's 3x3, LayerNorm by +rows, by window map and as PatchMerging's gather, the shifted windows' attention, the align_corners resize, the patch +and tile copies, the modulated deformable columns, the pool, the gate and the head against float64 (21 controls). RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then `tf_vit_pointwise`, byte for byte. diff --git a/tools/zig/idol_birefnet_case.py b/tools/zig/idol_birefnet_case.py new file mode 100644 index 0000000000..2a2564e71d --- /dev/null +++ b/tools/zig/idol_birefnet_case.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 +"""tf-idol-birefnet's case: the SDK's preprocess of a photo, a float64 CPU ruler, and with --stages each stage's.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from idol_case import SDK_COMMIT, flush, pin # noqa: E402 + +USAGE = """\ + idol_birefnet_case.py WEIGHTS SDK IMAGE OUT [--mode synthetic|real] [--crop-margin X] [--stages] [--repeats N] +WEIGHTS is a speridlabs/idol snapshot (birefnet.safetensors, config.json); SDK an idol-sdk checkout at the pinned +commit; IMAGE a PNG, JPEG or WebP photo. Writes OUT/birefnet.safetensors (the resized photo, BiRefNet's input, logits +and alpha, torch's and float64's, and both preprocessed pictures), OUT/birefnet.json and, with --stages, +OUT/birefnet_stages.safetensors. +""" + +PREPROCESS_SHA256 = { + "idol/models/birefnet.py": "7c8e5e1021978225740bc9e0a94beef2fba1545a2ffc739235fab6aadc244068", + "idol/preprocess.py": "b6ef41020760b1d1ffbee2f82a2eb22752797d8c474f7b3eaa25f2a6413a1ffe", + "idol/pipeline.py": "091fc0fe7d8f66f5f9f63795a91f9b654e12cf879d94bf03624f39c360ddd4a1", +} +# speridlabs/idol's birefnet.safetensors: ZhengPeng7/BiRefNet e2bf8e4460fc's tensors in fp16, less 12 training-only ones +BIREFNET_SHA256 = "44bbc97b041f879975d0d6d538b8e66a25eb04d1f49aab40697544cdd97eaa8e" +PILLOW_VERSION = "13.0.0" + + +def rows(t): + """[1, C, H, W] (or [1, L, C] tokens) as float32 rows [H * W, C], the layout tf-idol-birefnet keeps.""" + t = t.detach().float().cpu() + if t.dim() == 4: + return t[0].permute(1, 2, 0).reshape(-1, t.shape[1]).contiguous() + return t[0].contiguous() + + +def distance(ours, ruler, width: int) -> dict: + """1 - cosine and the worst row's relative L2 error against the ruler, rows of `width` as tf-idol-birefnet reads.""" + a, b = ours.double().reshape(-1, width), ruler.double().reshape(-1, width) + cos = float((a * b).sum() / (a.norm() * b.norm())) + worst = (a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30) + return {"one_minus_cos": 1 - cos, "worst_row": float(worst.max()), "worst_at": int(worst.argmax()), + "max_abs": float((a - b).abs().max())} + + +def differ(a, b) -> dict: + d = (a.astype("int32") - b.astype("int32")) + d = abs(d).max(axis=-1) if d.ndim == 3 else abs(d) + return {"pixels": int((d > 0).sum()), "max": int(d.max())} + + +def birefnet(weights: str, sdk_dir: str, device: str): + """The SDK's BiRefNet with the snapshot's weights, strictly loaded as IdolPipeline loads it (fp16 on CUDA only).""" + import PIL + import torch + from safetensors.torch import load_file + + pin(Path(sdk_dir), PREPROCESS_SHA256) + if PIL.__version__ != PILLOW_VERSION: + sys.exit(f"Pillow is {PIL.__version__}; this tool runs only {PILLOW_VERSION}") + path = Path(weights) / "birefnet.safetensors" + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != BIREFNET_SHA256: + sys.exit(f"{path} has sha256 {digest}; this tool runs only speridlabs/idol's ({BIREFNET_SHA256})") + from idol.models import BiRefNet + + model = BiRefNet() + model.load_state_dict(load_file(str(path)), strict=True) + model = model.eval().requires_grad_(False).to(torch.device(device)) + if device == "cuda": + model.half() + return model + + +class Recorder: + """The background remover preprocess_image calls: the SDK's own, its input and output kept.""" + + def __init__(self, model, device): + self.model, self.device, self.calls = model, device, [] + + def remove_background(self, image): + import numpy as np + + flush(self.device) + mark = time.perf_counter() + out = self.model.remove_background(image) + flush(self.device) + # preprocess erodes the alpha in place next, so it is copied now + self.calls.append({"resized": image, "alpha": np.array(out.getchannel(3)), "seconds": time.perf_counter() - mark}) + return out + + +class Given: + """A background remover that puts a given alpha on the photo: preprocess's later steps on that alpha.""" + + def __init__(self, alpha): + self.alpha = alpha + + def remove_background(self, image): + from PIL import Image + + rgba = image.copy() + rgba.putalpha(Image.fromarray(self.alpha, mode="L")) + return rgba + + +def ruler_input(resized): + """remove_background's tensor in float64 from the resized photo: bilinear (align_corners) to 1024, ImageNet.""" + import numpy as np + import torch + import torch.nn.functional as F + + t = torch.from_numpy(np.array(resized, dtype=np.float64) / 255.0).permute(2, 0, 1)[None] + t = F.interpolate(t, size=(1024, 1024), mode="bilinear", align_corners=True) + mean = torch.tensor((0.485, 0.456, 0.406), dtype=torch.float64).view(1, 3, 1, 1) + std = torch.tensor((0.229, 0.224, 0.225), dtype=torch.float64).view(1, 3, 1, 1) + return (t - mean) / std + + +def to_alpha(logits, size): + """remove_background's alpha from logits [1, 1, 1024, 1024]: sigmoid, bilinear to (H, W), clamp * 255, truncated.""" + import torch + import torch.nn.functional as F + + a = F.interpolate(torch.sigmoid(logits), size=size, mode="bilinear", align_corners=True)[0, 0] + return (a.clamp(0, 1) * 255).to(torch.uint8).cpu().numpy() + + +def stage_hooks(model, keep): + """Hooks keeping each stage's torch input and output: both passes' layers, the encoder, squeeze and decoder.""" + hooks, seen = [], {"bb": 0} + + def on(name, fn): + def hook(_m, args, out): + fn(name, args, out) + return hook + + def layer(name, args, out): + p = seen["bb"] + x, H, W = args + keep[f"{name}.{p}.in"] = rows(x) + keep[f"{name}.{p}.out"] = rows(out[0]) + keep[f"{name}.{p}.down"] = rows(out[3]) + keep[f"{name}.{p}.hw"] = (int(H), int(W)) + if name.endswith(".3"): + seen["bb"] += 1 + + for i, lay in enumerate(model.bb.layers): + hooks.append(lay.register_forward_hook(on(f"layer.{i}", layer))) + + def block(name, args, out): + keep[f"{name}.in"], keep[f"{name}.out"] = rows(args[0]), rows(out) + + hooks.append(model.squeeze_module.register_forward_hook(on("squeeze", block))) + for n in (4, 3, 2, 1): + hooks.append(getattr(model.decoder, f"decoder_block{n}").register_forward_hook(on(f"dec{n}", block))) + + def decoder(name, args, out): + for k, v in zip(("x", "x1", "x2", "x3", "x4"), args): + keep[f"decoder.in.{k}"] = rows(v) + keep["decoder.out"] = rows(out) + + hooks.append(model.decoder.register_forward_hook(on("decoder", decoder))) + forward_enc = model._forward_enc + + def enc(x): + out = forward_enc(x) + keep["enc.in"] = rows(x) + for k, v in zip(("x1", "x2", "x3", "x4"), out): + keep[f"enc.out.{k}"] = rows(v) + return out + + model._forward_enc = enc + return hooks, forward_enc + + +def nchw(r, h, w): + """Rows [H * W, C] back to [1, C, H, W] float64.""" + return r.double().reshape(h, w, -1).permute(2, 0, 1)[None].contiguous() + + +def stage_rulers(m64, keep, gap) -> tuple[dict, dict]: + """Each stage in float64 on torch's own input; torch's distance to it per stage.""" + import torch + + out, dist = {}, {} + with torch.no_grad(): + for p in (0, 1): + for i in range(4): + key = f"layer.{i}.{p}" + H, W = keep.pop(f"{key}.hw") + x_out, _, _, x_down, _, _ = m64.bb.layers[i](keep[f"{key}.in"].double()[None], H, W) + out[f"{key}.out"], out[f"{key}.down"] = rows(x_out), rows(x_down) + side = 1024 + x = nchw(keep["enc.in"], side, side) + for k, v in zip(("x1", "x2", "x3", "x4"), m64._forward_enc(x)): + out[f"enc.out.{k}"] = rows(v) + s4 = 32 + out["squeeze.out"] = rows(m64.squeeze_module(nchw(keep["squeeze.in"], s4, s4))) + for n, s in ((4, 32), (3, 64), (2, 128), (1, 256)): + out[f"dec{n}.out"] = rows(getattr(m64.decoder, f"decoder_block{n}")(nchw(keep[f"dec{n}.in"], s, s))) + ins = [nchw(keep[f"decoder.in.{k}"], s, s) for k, s in (("x", 1024), ("x1", 256), ("x2", 128), ("x3", 64), + ("x4", 32))] + out["decoder.out"] = rows(m64.decoder(*ins)) + for k, v in out.items(): + dist[k] = distance(keep[k], v, gap(k, v)) + return out, dist + + +def width_of(key, v): + """A stage's row: a pixel's channels; the logits' row: an image row of 1,024.""" + return 1024 if key == "decoder.out" else v.shape[-1] + + +def case(args) -> int: + import copy + + import numpy as np + import torch + from PIL import Image + from safetensors.torch import save_file + + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + model = birefnet(args.weights, args.sdk, args.device) + from idol.preprocess import CROP_MARGIN, has_real_alpha, preprocess_image + + config = json.loads((Path(args.weights) / "config.json").read_text()) + size = config["image_size"] + margin = args.crop_margin if args.crop_margin is not None else CROP_MARGIN[args.mode] + photo = Image.open(args.image) + meta = {"image": Path(args.image).name, "format": photo.format, "mode": photo.mode, "size": list(photo.size), + "preprocess_mode": args.mode, "crop_margin": margin, "device": args.device, "torch": torch.__version__, + "sdk": SDK_COMMIT, "pillow": PILLOW_VERSION} + keep, hooks, forward_enc = {}, [], None + logits = [] + hooks.append(model.decoder.register_forward_hook(lambda _m, _a, o: logits.append(o.detach().float().cpu()))) + inputs = [] + # remove_background calls forward directly, so the input is the backbone's first + hooks.append(model.bb.register_forward_pre_hook(lambda _m, a: inputs.append(a[0].detach().float().cpu()))) + if args.stages: + more, forward_enc = stage_hooks(model, keep) + hooks += more + rec = Recorder(model, args.device) + with torch.no_grad(): + mark = time.perf_counter() + picture = preprocess_image(Image.open(args.image), background_remover=rec, image_size=size, + crop_margin=margin) + flush(args.device) + first_s = time.perf_counter() - mark + for h in hooks: + h.remove() + if forward_enc is not None: + model._forward_enc = forward_enc + times = [] + with torch.no_grad(): + for _ in range(args.repeats): + mark = time.perf_counter() + again = preprocess_image(Image.open(args.image), background_remover=Recorder(model, args.device), + image_size=size, crop_margin=margin) + flush(args.device) + times.append(time.perf_counter() - mark) + if not np.array_equal(np.array(again), np.array(picture)): + sys.exit("the SDK's preprocessed picture differs between runs") + tensors = {"picture": torch.from_numpy(np.array(picture))} + meta |= {"has_real_alpha": not rec.calls, "preprocess_s": round(first_s, 4), + "preprocess_repeat_s": [round(t, 4) for t in times]} + if rec.calls: + call = rec.calls[0] + resized = call["resized"] + alpha = call["alpha"] + sdk_logits = logits[0] + tensors |= {"resized": torch.from_numpy(np.array(resized)), "input": rows(inputs[0]), + "logits": sdk_logits[0, 0].contiguous(), "alpha": torch.from_numpy(alpha)} + m64 = copy.deepcopy(model).to("cpu").double() + with torch.no_grad(): + mark = time.perf_counter() + x64 = ruler_input(resized) + x1, x2, x3, x4 = m64._forward_enc(x64) + ruler = m64.decoder(x64, x1, x2, x3, m64.squeeze_module(x4)) + ruler_s = time.perf_counter() - mark + # the ruler's logits through torch's own fp32 alpha steps: float64's sigmoid would round differently + ruler_alpha = to_alpha(ruler.float().to(args.device), alpha.shape) + ruler_picture = preprocess_image(Image.open(args.image), background_remover=Given(ruler_alpha), + image_size=size, crop_margin=margin) + tensors |= {"input_ruler": rows(x64), "logits_ruler": ruler[0, 0].float().contiguous(), + "alpha_ruler": torch.from_numpy(ruler_alpha), + "picture_ruler": torch.from_numpy(np.array(ruler_picture))} + sdk_alpha_again = to_alpha(sdk_logits.to(args.device), alpha.shape) + meta |= {"resized_size": list(resized.size), "remove_background_s": round(call["seconds"], 4), + "float64_s": round(ruler_s, 2), + "input_vs_ruler_max_abs": float((rows(inputs[0]).double() - rows(x64)).abs().max()), + "logits_vs_ruler": distance(sdk_logits[0, 0], ruler[0, 0], 1024), + "alpha_vs_ruler": differ(alpha, ruler_alpha), "picture_vs_ruler": differ(np.array(picture), + np.array(ruler_picture)), + "alpha_from_recorded_logits_equal": bool(np.array_equal(sdk_alpha_again, alpha))} + if args.stages: + ruled, dist = stage_rulers(m64, keep, width_of) + meta["stages_vs_float64"] = dist + stages = {k: v for k, v in keep.items()} | {f"ruler.{k}": v for k, v in ruled.items()} + save_file({k: v.float().contiguous() for k, v in stages.items()}, str(out / "birefnet_stages.safetensors")) + else: + meta["real_alpha"] = bool(has_real_alpha(Image.open(args.image))) + save_file({k: v.contiguous() for k, v in tensors.items()}, str(out / "birefnet.safetensors")) + (out / "birefnet.json").write_text(json.dumps(meta, indent=1) + "\n") + print(json.dumps(meta)) + return 0 + + +def main(argv: list[str]) -> int: + ap = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + for name in ("weights", "sdk", "image", "out_dir"): + ap.add_argument(name) + ap.add_argument("--mode", default="synthetic", choices=("synthetic", "real")) + ap.add_argument("--crop-margin", type=float) + ap.add_argument("--device", default="mps", choices=("mps", "cpu", "cuda")) + ap.add_argument("--stages", action="store_true", help="each stage's torch rows and float64 on them") + ap.add_argument("--repeats", type=int, default=3, help="timed preprocesses after the recorded one") + return case(ap.parse_args(argv)) + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 76b5f721e0..5d581d2e6a 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -321,6 +321,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-idol-decode", .path = "zig/tests/idol_decode.zig", .about = "Idol's octree and Gaussian decoders on a reference case: each level, the rows and the PLY gated, then timed" }, .{ .name = "tf-idol-vae", .path = "zig/tests/idol_vae.zig", .about = "Idol's image encoder on a reference case: its condition rows gated against float64 beside the SDK's, then timed" }, .{ .name = "tf-idol-dino", .path = "zig/tests/idol_dino.zig", .about = "Idol's DINOv3 image encoder on a reference case: feature1 and each block gated against torch's, then timed" }, + .{ .name = "tf-idol-birefnet", .path = "zig/tests/idol_birefnet.zig", .about = "Idol's background removal and preprocessing on a reference case: BiRefNet's logits, each stage, the alpha and the picture gated, then timed" }, .{ .name = "tf-idol-generate", .path = "zig/tests/idol_generate.zig", .about = "Idol end to end: one preprocessed picture to a 3DGS PLY and camera through both encoders, the denoiser and the decoders, timed by stage" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, diff --git a/zig/kernels/metal/idol_birefnet.metal b/zig/kernels/metal/idol_birefnet.metal new file mode 100644 index 0000000000..a2c2aafba4 --- /dev/null +++ b/zig/kernels/metal/idol_birefnet.metal @@ -0,0 +1,561 @@ +// Idol's background remover (BiRefNet: Swin-L, ASPP-deformable decoder) in fp32, each tensor [rows, channels]. + +#include +using namespace metal; + +#define TF_UNROLL_BR _Pragma("clang loop unroll(full)") + +// erf in fp32 to within an ulp (musl's erff), for nn.GELU's exact form. +inline float br_erf(float x) { + const uint ix = as_type(x) & 0x7fffffffu; + const bool neg = as_type(x) >> 31; + if (ix >= 0x7f800000u) return isnan(x) ? x : (neg ? -1.0f : 1.0f); + if (ix < 0x3f580000u) { + if (ix < 0x31800000u) return 0.125f * (8.0f * x + 1.0270333290e+00f * x); + const float z = x * x; + const float r = 1.2837916613e-01f + z * (-3.2504209876e-01f + z * (-2.8481749818e-02f + z * (-5.7702702470e-03f + + z * -2.3763017452e-05f))); + const float s = 1.0f + z * (3.9791721106e-01f + z * (6.5022252500e-02f + z * (5.0813062117e-03f + + z * (1.3249473704e-04f + z * -3.9602282413e-06f)))); + return x + x * (r / s); + } + float y; + if (ix < 0x40c00000u) { + const float ax = fabs(x); + if (ix < 0x3fa00000u) { + const float s = ax - 1.0f; + const float P = -2.3621185683e-03f + s * (4.1485610604e-01f + s * (-3.7220788002e-01f + s * (3.1834661961e-01f + + s * (-1.1089469492e-01f + s * (3.5478305072e-02f + s * -2.1663755178e-03f))))); + const float Q = 1.0f + s * (1.0642088205e-01f + s * (5.4039794207e-01f + s * (7.1828655899e-02f + + s * (1.2617121637e-01f + s * (1.3637083583e-02f + s * 1.1984500103e-02f))))); + y = 8.4506291151e-01f + P / Q; + } else { + const float s = 1.0f / (ax * ax); + float R, S; + if (ix < 0x4036db6du) { + R = -9.8649440333e-03f + s * (-6.9385856390e-01f + s * (-1.0558626175e+01f + s * (-6.2375331879e+01f + + s * (-1.6239666748e+02f + s * (-1.8460508728e+02f + s * (-8.1287437439e+01f + s * -9.8143291473e+00f)))))); + S = 1.0f + s * (1.9651271820e+01f + s * (1.3765776062e+02f + s * (4.3456588745e+02f + s * (6.4538726807e+02f + + s * (4.2900814819e+02f + s * (1.0863500214e+02f + s * (6.5702495575e+00f + s * -6.0424413532e-02f))))))); + } else { + R = -9.8649431020e-03f + s * (-7.9928326607e-01f + s * (-1.7757955551e+01f + s * (-1.6063638306e+02f + + s * (-6.3756646729e+02f + s * (-1.0250950928e+03f + s * -4.8351919556e+02f))))); + S = 1.0f + s * (3.0338060379e+01f + s * (3.2579251099e+02f + s * (1.5367296143e+03f + s * (3.1998581543e+03f + + s * (2.5530502930e+03f + s * (4.7452853394e+02f + s * -2.2440952301e+01f)))))); + } + const float z = as_type(as_type(ax) & 0xffffe000u); + y = 1.0f - precise::exp(-z * z - 0.5625f) * precise::exp((z - ax) * (z + ax) + R / S) / ax; + } + } else { + y = 1.0f - 0x1p-120f; + } + return neg ? -y : y; +} + +inline float br_sigmoid(float x) { return 1.0f / (1.0f + precise::exp(-x)); } + +struct BrConv { + uint rows; + uint cols; + uint cin; + uint ldx; + uint cout; + uint ldw; + uint k; + uint ldy; + uint yoff; + uint flags; + uint row0; +}; + +// y = conv_k(x) + b; flags: 1 BN (post), 2 ReLU, 4 GELU, 8 + r, 16 rows by map; w [k * k * cin, ldw] by tap +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void br_conv( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + const device float* post [[buffer(3)]], const device float* r [[buffer(4)]], const device int* map [[buffer(5)]], + device float* y [[buffer(6)]], constant BrConv& p [[buffer(7)]], uint2 tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint BK = 32; + constexpr uint lda = BK + 4; + // an implicit GEMM, a tap at a time: a 64 x 64 tile a threadgroup, fp32 sums + threadgroup float As[64 * lda]; + threadgroup float Cs[64 * 68]; + const uint pixels = p.rows * p.cols; + const uint m0 = tg.y * 64; + const uint n0 = tg.x * 64; + const uint arow = t / 2; + const uint acol = (t % 2) * 16; + const uint pix = m0 + arow; + const int oy = int(pix / p.cols); + const int ox = int(pix % p.cols); + const int pad = int(p.k / 2); + const uint sr = (sg / 2) * 32; + const uint sc = (sg % 2) * 32; + // sums in blocks of br_block terms added into the total: one running sum over 50k terms drifts twice torch's + constexpr uint br_block = 256; + simdgroup_float8x8 acc[4][4]; + simdgroup_float8x8 part[4][4]; + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) { + TF_UNROLL_BR + for (uint j = 0; j < 4; j++) { + acc[i][j] = simdgroup_float8x8(0.0f); + part[i][j] = simdgroup_float8x8(0.0f); + } + } + const uint K = p.k * p.k * p.cin; + for (uint k0 = 0; k0 < K; k0 += BK) { + const uint tap = k0 / p.cin; + const int iy = oy + int(tap / p.k) - pad; + const int ix = ox + int(tap % p.k) - pad; + threadgroup float4* dst = (threadgroup float4*)(As + arow * lda + acol); + if (pix < pixels && iy >= 0 && ix >= 0 && iy < int(p.rows) && ix < int(p.cols)) { + const device float4* src = + (const device float4*)(x + ulong(uint(iy) * p.cols + uint(ix)) * p.ldx + k0 % p.cin + acol); + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) dst[i] = src[i]; + } else { + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) dst[i] = float4(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const device float* wk = w + ulong(k0) * p.ldw + n0 + sc; + TF_UNROLL_BR + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 am[4]; + simdgroup_float8x8 bm[4]; + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) simdgroup_load(am[i], As + (sr + i * 8) * lda + kk, lda); + TF_UNROLL_BR + for (uint j = 0; j < 4; j++) simdgroup_load(bm[j], wk + ulong(kk) * p.ldw + j * 8, p.ldw); + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) { + TF_UNROLL_BR + for (uint j = 0; j < 4; j++) simdgroup_multiply_accumulate(part[i][j], am[i], bm[j], part[i][j]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if ((k0 + BK) % br_block == 0 || k0 + BK >= K) { + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) { + TF_UNROLL_BR + for (uint j = 0; j < 4; j++) { + acc[i][j].thread_elements()[0] += part[i][j].thread_elements()[0]; + acc[i][j].thread_elements()[1] += part[i][j].thread_elements()[1]; + part[i][j] = simdgroup_float8x8(0.0f); + } + } + } + } + TF_UNROLL_BR + for (uint i = 0; i < 4; i++) { + TF_UNROLL_BR + for (uint j = 0; j < 4; j++) simdgroup_store(acc[i][j], Cs + (sr + i * 8) * 68 + sc + j * 8, 68); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = t; e < 64 * 64; e += 128) { + const uint row = m0 + e / 64; + const uint col = n0 + e % 64; + if (row >= pixels || col >= p.cout) continue; + int out = int(p.row0 + row); + if (p.flags & 16) out = map[row]; + if (out < 0) continue; + const ulong at = ulong(out) * p.ldy + p.yoff + col; + float v = Cs[(e / 64) * 68 + e % 64] + b[col]; + if (p.flags & 1) v = (v - post[col]) * post[p.ldw + col] + post[2 * p.ldw + col]; + if (p.flags & 2) v = v < 0.0f ? 0.0f : v; + if (p.flags & 4) v = (v * 0.5f) * (1.0f + br_erf(v * 0.70710677f)); + if (p.flags & 8) v = r[at] + v; + y[at] = v; + } +} + +struct BrNorm { + uint rows; + uint width; + uint mode; + uint src_w; + uint ldy; + uint yoff; + float eps; +}; + +// Row r at channel c: mode 0 row r, 1 row map[r], 2 PatchMerging's [x0 | x1 | x2 | x3] at r's 2 x 2 patch +inline float br_norm_at(const device float* x, const device int* map, constant BrNorm& p, uint r, uint c) { + if (p.mode == 0) return x[ulong(r) * p.width + c]; + if (p.mode == 1) return x[ulong(map[r]) * p.width + c]; + const uint q = p.width / 4; + const uint part = c / q; + const uint half_w = p.src_w / 2; + const uint sy = (r / half_w) * 2 + (part & 1); + const uint sx = (r % half_w) * 2 + (part >> 1); + return x[ulong(sy * p.src_w + sx) * q + c % q]; +} + +// y = LayerNorm(row) * g + beta, two fp32 passes; a map's negative rows are zeros (torch pads after the norm) +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_norm( + const device float* x [[buffer(0)]], const device int* map [[buffer(1)]], const device float* g [[buffer(2)]], + const device float* beta [[buffer(3)]], device float* y [[buffer(4)]], constant BrNorm& p [[buffer(5)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[8]; + device float* out = y + ulong(row) * p.ldy + p.yoff; + if (p.mode == 1 && map[row] < 0) { + for (uint c = t; c < p.width; c += 256) out[c] = 0.0f; + return; + } + float s = 0.0f; + for (uint c = t; c < p.width; c += 256) s += br_norm_at(x, map, p, row, c); + s = simd_sum(s); + if (lane == 0) part[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + float all = 0.0f; + TF_UNROLL_BR + for (uint i = 0; i < 8; i++) all += part[i]; + const float mean = all / float(p.width); + threadgroup_barrier(mem_flags::mem_threadgroup); + float d2 = 0.0f; + for (uint c = t; c < p.width; c += 256) { + const float d = br_norm_at(x, map, p, row, c) - mean; + d2 += d * d; + } + d2 = simd_sum(d2); + if (lane == 0) part[sg] = d2; + threadgroup_barrier(mem_flags::mem_threadgroup); + all = 0.0f; + TF_UNROLL_BR + for (uint i = 0; i < 8; i++) all += part[i]; + const float rstd = 1.0f / precise::sqrt(all / float(p.width) + p.eps); + for (uint c = t; c < p.width; c += 256) out[c] = (br_norm_at(x, map, p, row, c) - mean) * rstd * g[c] + beta[c]; +} + +struct BrWindow { + uint width; + uint heads; + uint rows; + uint cols; + uint shift; + float scale; +}; + +// Token i's region (0 to 8) in the shifted grid: rows and columns split at size - 12 and size - shift +inline uint br_region(uint i, uint wy, uint wx, constant BrWindow& p) { + const uint y = wy * 12 + i / 12; + const uint x = wx * 12 + i % 12; + const uint ry = y < p.rows - 12 ? 0u : (y < p.rows - p.shift ? 1u : 2u); + const uint rx = x < p.cols - 12 ? 0u : (x < p.cols - p.shift ? 1u : 2u); + return ry * 3 + rx; +} + +// A 12 x 12 window, one head, 48 queries: softmax(q scale k^T + bias (+ -100 across regions)) v +[[max_total_threads_per_threadgroup(192)]] [[kernel]] void br_window( + const device float* qkv [[buffer(0)]], const device float* table [[buffer(1)]], device float* o [[buffer(2)]], + constant BrWindow& p [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint N = 144; + constexpr uint QN = 48; + constexpr uint LS = 148; + threadgroup float S[QN * LS]; + const uint h = tg.y; + const uint win = tg.z; + const uint ld = 3 * p.width; + const uint base = win * N; + const uint first = tg.x * QN; + const uint qcol = h * 32; + for (uint tile = sg; tile < (QN / 8) * (N / 8); tile += 6) { + const uint ti = tile / (N / 8); + const uint tj = tile % (N / 8); + simdgroup_float8x8 s = simdgroup_float8x8(0.0f); + TF_UNROLL_BR + for (uint d = 0; d < 32; d += 8) { + simdgroup_float8x8 qm; + simdgroup_float8x8 km; + simdgroup_load(qm, qkv + ulong(base + first + ti * 8) * ld + qcol + d, ld); + qm.thread_elements()[0] *= p.scale; + qm.thread_elements()[1] *= p.scale; + simdgroup_load(km, qkv + ulong(base + tj * 8) * ld + p.width + qcol + d, ld, ulong2(0, 0), true); + simdgroup_multiply_accumulate(s, qm, km, s); + } + simdgroup_store(s, S + ti * 8 * LS + tj * 8, LS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // four threads a query row + const uint row = t / 4; + const uint sub = t % 4; + const uint qi = first + row; + const uint wy = win / (p.cols / 12); + const uint wx = win % (p.cols / 12); + threadgroup float* sr = S + row * LS; + const uint mine = p.shift ? br_region(qi, wy, wx, p) : 0u; + float most = -INFINITY; + for (uint j = sub; j < N; j += 4) { + const uint at = ((qi / 12 + 11 - j / 12) * 23 + (qi % 12 + 11 - j % 12)) * p.heads + h; + float v = sr[j] + table[at]; + if (p.shift) v = v + (br_region(j, wy, wx, p) == mine ? 0.0f : -100.0f); + sr[j] = v; + most = max(most, v); + } + most = max(most, simd_shuffle_xor(most, 1)); + most = max(most, simd_shuffle_xor(most, 2)); + float total = 0.0f; + for (uint j = sub; j < N; j += 4) { + const float e = precise::exp(sr[j] - most); + sr[j] = e; + total += e; + } + total += simd_shuffle_xor(total, 1); + total += simd_shuffle_xor(total, 2); + for (uint j = sub; j < N; j += 4) sr[j] = sr[j] / total; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint tile = sg; tile < (QN / 8) * 4; tile += 6) { + const uint ti = tile / 4; + const uint tj = tile % 4; + simdgroup_float8x8 acc = simdgroup_float8x8(0.0f); + for (uint kk = 0; kk < N; kk += 8) { + simdgroup_float8x8 pm; + simdgroup_float8x8 vm; + simdgroup_load(pm, S + ti * 8 * LS + kk, LS); + simdgroup_load(vm, qkv + ulong(base + kk) * ld + 2 * p.width + qcol + tj * 8, ld); + simdgroup_multiply_accumulate(acc, pm, vm, acc); + } + simdgroup_store(acc, o + ulong(base + first + ti * 8) * p.width + qcol + tj * 8, p.width); + } +} + +struct BrResize { + uint in_h; + uint in_w; + uint out_h; + uint out_w; + uint channels; + uint ldx; + uint xoff; + uint ldy; + uint yoff; +}; + +// y = F.interpolate(x, bilinear, align_corners=True): x [in_h * in_w, ldx] from xoff, y [out_h * out_w, ldy] from yoff. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_resize( + const device float* x [[buffer(0)]], device float* y [[buffer(1)]], constant BrResize& p [[buffer(2)]], + uint2 at [[thread_position_in_grid]]) { + if (at.x >= p.channels || at.y >= p.out_h * p.out_w) return; + const uint oy = at.y / p.out_w; + const uint ox = at.y % p.out_w; + const float sh = p.out_h > 1 ? float(p.in_h - 1) / float(p.out_h - 1) : 0.0f; + const float sw = p.out_w > 1 ? float(p.in_w - 1) / float(p.out_w - 1) : 0.0f; + const float fy = sh * float(oy); + const float fx = sw * float(ox); + const uint y0 = min(uint(fy), p.in_h - 1); + const uint x0 = min(uint(fx), p.in_w - 1); + const uint y1 = min(y0 + 1, p.in_h - 1); + const uint x1 = min(x0 + 1, p.in_w - 1); + const float ly = min(max(fy - float(y0), 0.0f), 1.0f); + const float lx = min(max(fx - float(x0), 0.0f), 1.0f); + const device float* src = x + p.xoff + at.x; + const float v00 = src[ulong(y0 * p.in_w + x0) * p.ldx]; + const float v01 = src[ulong(y0 * p.in_w + x1) * p.ldx]; + const float v10 = src[ulong(y1 * p.in_w + x0) * p.ldx]; + const float v11 = src[ulong(y1 * p.in_w + x1) * p.ldx]; + const float top = (1.0f - lx) * v00 + lx * v01; + const float bottom = (1.0f - lx) * v10 + lx * v11; + y[ulong(at.y) * p.ldy + p.yoff + at.x] = (1.0f - ly) * top + ly * bottom; +} + +struct BrTiles { + uint out_h; + uint out_w; + uint src_w; + uint channels; + uint a; + uint b; + uint mode; + uint cpad; +}; + +// Channel ch * a * b + i * b + j (0 past channels * a * b): mode 0 a x b patches, mode 1 image2patches' tiles +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_tiles( + const device float* src [[buffer(0)]], device float* dst [[buffer(1)]], constant BrTiles& p [[buffer(2)]], + uint2 at [[thread_position_in_grid]]) { + if (at.x >= p.cpad || at.y >= p.out_h * p.out_w) return; + const ulong o = ulong(at.y) * p.cpad + at.x; + const uint ab = p.a * p.b; + if (at.x >= p.channels * ab) { + dst[o] = 0.0f; + return; + } + const uint ch = at.x / ab; + const uint i = (at.x % ab) / p.b; + const uint j = at.x % p.b; + const uint oy = at.y / p.out_w; + const uint ox = at.y % p.out_w; + const uint sy = p.mode == 0 ? oy * p.a + i : i * p.out_h + oy; + const uint sx = p.mode == 0 ? ox * p.b + j : j * p.out_w + ox; + dst[o] = src[ulong(sy * p.src_w + sx) * p.channels + ch]; +} + +struct BrConv3 { + uint rows; + uint cols; + uint cout; +}; + +// y [pixels, cout] = conv3x3(x) + b of a 3-channel picture x [pixels, 3], zero padded; w [27, cout] by (tap, channel). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_conv3( + const device float* x [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device float* y [[buffer(3)]], constant BrConv3& p [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + if (at.x >= p.cout || at.y >= p.rows * p.cols) return; + const int py = int(at.y / p.cols); + const int px = int(at.y % p.cols); + float s = 0.0f; + TF_UNROLL_BR + for (int ky = 0; ky < 3; ky++) { + TF_UNROLL_BR + for (int kx = 0; kx < 3; kx++) { + const int iy = py + ky - 1; + const int ix = px + kx - 1; + if (iy < 0 || ix < 0 || iy >= int(p.rows) || ix >= int(p.cols)) continue; + const device float* v = x + ulong(uint(iy) * p.cols + uint(ix)) * 3; + TF_UNROLL_BR + for (int c = 0; c < 3; c++) s = fma(v[c], w[((ky * 3 + kx) * 3 + c) * p.cout + at.x], s); + } + } + y[ulong(at.y) * p.cout + at.x] = s + b[at.x]; +} + +struct BrDeform { + uint rows; + uint cols; + uint k; + uint first; + uint count; + uint ldo; +}; + +// torchvision's modulated deform_conv2d columns [count, k * k * 64]: 2 sigmoid(m) x the moved bilinear sample +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_deform( + const device float* x [[buffer(0)]], const device float* offs [[buffer(1)]], device float* col [[buffer(2)]], + constant BrDeform& p [[buffer(3)]], uint2 at [[thread_position_in_grid]]) { + const uint taps = p.k * p.k; + if (at.x >= taps || at.y >= p.count) return; + const uint pix = p.first + at.y; + const int oy = int(pix / p.cols); + const int ox = int(pix % p.cols); + const int pad = int(p.k / 2); + const device float* off = offs + ulong(pix) * p.ldo; + const float mask = 2.0f * br_sigmoid(off[2 * taps + at.x]); + const float h = float(oy - pad + int(at.x / p.k)) + off[2 * at.x]; + const float w = float(ox - pad + int(at.x % p.k)) + off[2 * at.x + 1]; + device float4* out = (device float4*)(col + ulong(at.y) * taps * 64 + at.x * 64); + const int H = int(p.rows); + const int W = int(p.cols); + if (!(h > -1.0f && h < float(H) && w > -1.0f && w < float(W))) { + TF_UNROLL_BR + for (uint c = 0; c < 16; c++) out[c] = mask * float4(0.0f); + return; + } + const int hl = int(floor(h)); + const int wl = int(floor(w)); + const float lh = h - float(hl); + const float lw = w - float(wl); + const float hh = 1.0f - lh; + const float hw = 1.0f - lw; + const float w1 = hh * hw; + const float w2 = hh * lw; + const float w3 = lh * hw; + const float w4 = lh * lw; + const bool top = hl >= 0; + const bool bottom = hl + 1 <= H - 1; + const bool left = wl >= 0; + const bool right = wl + 1 <= W - 1; + const device float4* r1 = (const device float4*)(x + ulong(max(hl, 0) * W + max(wl, 0)) * 64); + const device float4* r2 = (const device float4*)(x + ulong(max(hl, 0) * W + min(wl + 1, W - 1)) * 64); + const device float4* r3 = (const device float4*)(x + ulong(min(hl + 1, H - 1) * W + max(wl, 0)) * 64); + const device float4* r4 = (const device float4*)(x + ulong(min(hl + 1, H - 1) * W + min(wl + 1, W - 1)) * 64); + TF_UNROLL_BR + for (uint c = 0; c < 16; c++) { + const float4 v1 = top && left ? r1[c] : float4(0.0f); + const float4 v2 = top && right ? r2[c] : float4(0.0f); + const float4 v3 = bottom && left ? r3[c] : float4(0.0f); + const float4 v4 = bottom && right ? r4[c] : float4(0.0f); + out[c] = mask * (((w1 * v1 + w2 * v2) + w3 * v3) + w4 * v4); + } +} + +struct BrPool { + uint pixels; + uint channels; +}; + +// mean [channels] of x [pixels, channels] over its pixels: a threadgroup a channel. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_pool( + const device float* x [[buffer(0)]], device float* mean [[buffer(1)]], constant BrPool& p [[buffer(2)]], + uint ch [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[8]; + float s = 0.0f; + for (uint i = t; i < p.pixels; i += 256) s += x[ulong(i) * p.channels + ch]; + s = simd_sum(s); + if (lane == 0) part[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t != 0) return; + float all = 0.0f; + TF_UNROLL_BR + for (uint i = 0; i < 8; i++) all += part[i]; + mean[ch] = all / float(p.pixels); +} + +struct BrGate { + uint pixels; + uint channels; +}; + +// x [pixels, channels] *= sigmoid(g . w + b), g [pixels, 16] (the gradient attention's 1x1 to one channel). +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_gate( + const device float* g [[buffer(0)]], const device float* w [[buffer(1)]], const device float* b [[buffer(2)]], + device float* x [[buffer(3)]], constant BrGate& p [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + if (at.x >= p.channels || at.y >= p.pixels) return; + float s = 0.0f; + TF_UNROLL_BR + for (uint i = 0; i < 16; i++) s = fma(g[ulong(at.y) * 16 + i], w[i], s); + const ulong o = ulong(at.y) * p.channels + at.x; + x[o] = x[o] * br_sigmoid(s + b[0]); +} + +struct BrHead { + uint in_h; + uint in_w; + uint out_h; + uint out_w; + uint cp; + uint cq; +}; + +// logits = conv_out1 of [bilinear (align_corners) of a [in, cp] | q [pixels, cq]], w [cp + cq] +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void br_head( + const device float* a [[buffer(0)]], const device float* q [[buffer(1)]], const device float* w [[buffer(2)]], + const device float* b [[buffer(3)]], device float* y [[buffer(4)]], constant BrHead& p [[buffer(5)]], + uint at [[thread_position_in_grid]]) { + if (at >= p.out_h * p.out_w) return; + const uint oy = at / p.out_w; + const uint ox = at % p.out_w; + const float fy = (float(p.in_h - 1) / float(p.out_h - 1)) * float(oy); + const float fx = (float(p.in_w - 1) / float(p.out_w - 1)) * float(ox); + const uint y0 = min(uint(fy), p.in_h - 1); + const uint x0 = min(uint(fx), p.in_w - 1); + const uint y1 = min(y0 + 1, p.in_h - 1); + const uint x1 = min(x0 + 1, p.in_w - 1); + const float ly = min(max(fy - float(y0), 0.0f), 1.0f); + const float lx = min(max(fx - float(x0), 0.0f), 1.0f); + const device float* s00 = a + ulong(y0 * p.in_w + x0) * p.cp; + const device float* s01 = a + ulong(y0 * p.in_w + x1) * p.cp; + const device float* s10 = a + ulong(y1 * p.in_w + x0) * p.cp; + const device float* s11 = a + ulong(y1 * p.in_w + x1) * p.cp; + float s = 0.0f; + for (uint c = 0; c < p.cp; c++) { + const float top = (1.0f - lx) * s00[c] + lx * s01[c]; + const float bottom = (1.0f - lx) * s10[c] + lx * s11[c]; + s = fma((1.0f - ly) * top + ly * bottom, w[c], s); + } + const device float* qr = q + ulong(at) * p.cq; + for (uint c = 0; c < p.cq; c++) s = fma(qr[c], w[p.cp + c], s); + y[at] = s + b[0]; +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index bffa7ecb85..604a36ffa3 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -141,6 +141,8 @@ pub const idol_decoder = @embedFile("idol_decoder.metal"); pub const idol_vae = @embedFile("idol_vae.metal"); /// Idol's DINOv3 ViT-H+/16 image encoder: the prefill GEMM template, then its own bf16 steps and attention. pub const idol_dino = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("idol_dino.metal"); +/// Idol's background remover (BiRefNet: Swin-L and its ASPP-deformable decoder) in fp32. +pub const idol_birefnet = @embedFile("idol_birefnet.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/core/contract/idol_records.zig b/zig/src/core/contract/idol_records.zig index 124347e4e1..e7e0fcbfbe 100644 --- a/zig/src/core/contract/idol_records.zig +++ b/zig/src/core/contract/idol_records.zig @@ -1,7 +1,7 @@ -//! Idol's launches, each source-reserved: the denoiser's (idol.metal, idol_attention.metal), the decoders', the VAE's. +//! Idol's launches, each source-reserved: the denoiser's, the decoders', the VAE's and BiRefNet's. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; -pub const records = denoiser ++ decoder ++ vae; +pub const records = denoiser ++ decoder ++ vae ++ birefnet; const denoiser = [_]M{ .{ .name = "idol id_linear_00", .threads = 256, .tg_memory = 0, .reserved = 256 }, @@ -53,7 +53,20 @@ const vae = [_]M{ .{ .name = "idol vae iv_tokens", .threads = 256, .tg_memory = 0, .reserved = 256 }, }; +const birefnet = [_]M{ + .{ .name = "idol birefnet br_conv", .threads = 128, .tg_memory = 26624, .reserved = 128 }, + .{ .name = "idol birefnet br_norm", .threads = 256, .tg_memory = 32, .reserved = 256 }, + .{ .name = "idol birefnet br_window", .threads = 192, .tg_memory = 28416, .reserved = 192 }, + .{ .name = "idol birefnet br_resize", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol birefnet br_tiles", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol birefnet br_conv3", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol birefnet br_deform", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol birefnet br_pool", .threads = 256, .tg_memory = 32, .reserved = 256 }, + .{ .name = "idol birefnet br_gate", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "idol birefnet br_head", .threads = 256, .tg_memory = 0, .reserved = 256 }, +}; + test "the Idol launches fit every Apple family, one record a pipeline" { try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); - try std.testing.expectEqual(@as(usize, 24 + 9 + 8), records.len); + try std.testing.expectEqual(@as(usize, 24 + 9 + 8 + 10), records.len); } diff --git a/zig/src/families/idol/birefnet.zig b/zig/src/families/idol/birefnet.zig new file mode 100644 index 0000000000..a57bc777bf --- /dev/null +++ b/zig/src/families/idol/birefnet.zig @@ -0,0 +1,315 @@ +//! BiRefNet on Metal in fp32 (Swin-L twice, ASPP-deformable decoder): the normalized picture to its mask logits. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("birefnet_kernels.zig"); +const wts = @import("birefnet_weights.zig"); +const dec = @import("birefnet_decoder.zig"); +const pre = @import("preprocess.zig"); +const Ref = ks.Ref; + +/// The network's input side, and its two passes' patch grids (the full picture's, then the half picture's). +pub const size = 1024; +pub const grids = [2]usize{ size / 4, size / 8 }; +/// Pixels a deformable branch builds columns for at once. +pub const chunk = 8192; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +/// The padded side of a grid of `n` (a whole number of windows). +pub fn padded(n: usize) usize { + return (n + wts.window - 1) / wts.window * wts.window; +} + +/// Each scratch tensor's floats: the picture and encoder outputs persist; encoder and decoder temporaries share one. +const P = struct { + const x = size * size * 3; + const x1 = 256 * 256 * wts.channels[3]; + const x2 = 128 * 128 * wts.channels[2]; + const x3 = 64 * 64 * wts.channels[1]; + const x4 = 32 * 32 * wts.squeeze_in; + const logits = size * size; + const pooled = 64 + wts.aspp_planes; + const xh = (size / 2) * (size / 2) * 3; + const tokens = 256 * 256 * wts.embed; + const slots = 264 * 264 * wts.embed; + const big = 256 * 256 * 4 * wts.embed; + const lin = [4]usize{ 32 * 32 * dec.level_in[0], 64 * 64 * dec.level_in[1], 128 * 128 * dec.level_in[2], 256 * 256 * dec.level_in[3] }; + const out = [4]usize{ 32 * 32 * wts.channels[1], 64 * 64 * wts.channels[2], 128 * 128 * wts.channels[3], 256 * 256 * wts.channels[3] / 2 }; + const a64 = 256 * 256 * wts.inter; + const offs = 256 * 256 * 192; + const cols = chunk * 49 * wts.inter; + const cat = 256 * 256 * 5 * wts.aspp_planes; + const gdt = 256 * 256 * wts.inter; + const patches = 256 * 256 * 64; + const q48 = size * size * 48; +}; + +fn bytes(n: usize) usize { + return std.mem.alignForward(usize, n * 4, 256); +} + +pub const Net = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + k: ks.Kernels, + w: wts.Weights, + scratch: mtl.Buffer, + maps: mtl.Buffer, + /// The windows' token maps: [pass][layer][shifted], slot to token or -1. + map: [2][4][2]Ref, + x: Ref, + xs: [4]Ref, + logits: Ref, + pooled: Ref, + xh: Ref, + ta: Ref, + tb: Ref, + slots: Ref, + big: Ref, + att: Ref, + lin: [4]Ref, + out: [4]Ref, + a64: Ref, + offs: Ref, + cols: Ref, + cat: Ref, + gdt: Ref, + patches: Ref, + q48: Ref, + + /// `path` (birefnet.safetensors from a speridlabs/idol snapshot). + pub fn load(gpa: std.mem.Allocator, device: mtl.Device, path: [:0]const u8) !*Net { + const n = try gpa.create(Net); + errdefer gpa.destroy(n); + n.* = undefined; + n.gpa = gpa; + n.device = device; + n.queue = try device.queue(); + errdefer n.queue.deinit(); + n.k = try ks.Kernels.load(device); + errdefer n.k.deinit(); + n.w = try wts.load(gpa, device, path); + errdefer n.w.deinit(); + try n.buildMaps(); + errdefer n.maps.deinit(); + const persistent = [_]usize{ P.x, P.x1, P.x2, P.x3, P.x4, P.logits, P.pooled }; + const enc = [_]usize{ P.xh, P.tokens, P.tokens, P.slots, P.big, P.slots }; + const decs = P.lin ++ P.out ++ [_]usize{ P.a64, P.offs, P.cols, P.cat, P.gdt, P.patches }; + var at: usize = 0; + for (persistent) |v| at += bytes(v); + var e_end = at; + for (enc) |v| e_end += bytes(v); + var d_end = at; + for (decs) |v| d_end += bytes(v); + n.scratch = try device.buffer(@max(e_end, d_end), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + const views = [_]*Ref{ &n.x, &n.xs[0], &n.xs[1], &n.xs[2], &n.xs[3], &n.logits, &n.pooled }; + at = 0; + for (views, persistent) |v, len| v.* = n.take(&at, len); + const shared = at; + for ([_]*Ref{ &n.xh, &n.ta, &n.tb, &n.slots, &n.big, &n.att }, enc) |v, len| v.* = n.take(&at, len); + at = shared; + n.q48 = .{ .buf = n.scratch, .off = at }; + std.debug.assert(bytes(P.lin[0]) + bytes(P.lin[1]) + bytes(P.lin[2]) + bytes(P.lin[3]) >= P.q48 * 4); + for (&n.lin, P.lin) |*v, len| v.* = n.take(&at, len); + for (&n.out, P.out) |*v, len| v.* = n.take(&at, len); + for ([_]*Ref{ &n.a64, &n.offs, &n.cols, &n.cat, &n.gdt, &n.patches }, decs[8..]) |v, len| v.* = n.take(&at, len); + return n; + } + + fn take(n: *Net, at: *usize, len: usize) Ref { + const r: Ref = .{ .buf = n.scratch, .off = at.* }; + at.* += bytes(len); + return r; + } + + pub fn deinit(n: *Net) void { + n.scratch.deinit(); + n.maps.deinit(); + n.w.deinit(); + n.k.deinit(); + n.queue.deinit(); + n.gpa.destroy(n); + } + + /// Each layer's window slots to tokens, both passes, plain and shifted by half a window (torch.roll by -6). + fn buildMaps(n: *Net) !void { + var total: usize = 0; + for (grids) |g0| for (0..4) |i| { + const hp = padded(g0 >> @intCast(i)); + total += 2 * std.mem.alignForward(usize, hp * hp, 64); + }; + n.maps = try n.device.buffer(total * 4, mtl.ResourceOptions.shared); + const all = @as([*]i32, @ptrCast(@alignCast(n.maps.contents())))[0..total]; + var at: usize = 0; + for (grids, 0..) |g0, pass| for (0..4) |i| for (0..2) |s| { + const h = g0 >> @intCast(i); + const hp = padded(h); + n.map[pass][i][s] = .{ .buf = n.maps, .off = at * 4 }; + fillMap(all[at..][0 .. hp * hp], h, if (s == 1) wts.window / 2 else 0); + at += std.mem.alignForward(usize, hp * hp, 64); + }; + } + + pub fn floats(r: Ref, len: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..len]; + } + + /// The GPU work so far, waited for; its seconds on the GPU. + pub fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !f64 { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("idol birefnet: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + return cb.gpuSeconds(); + } + + /// Mask logits [1024, 1024] of the normalized picture `input` [1024 * 1024, 3] (remove_background's tensor). + pub fn forward(n: *Net, input: []const f32, logits: []f32) !Timing { + if (input.len != P.x or logits.len != P.logits) return error.PictureSize; + @memcpy(floats(n.x, P.x), input); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = n.queue.commandBuffer(); + const e = cb.compute(.serial); + n.encode(e); + dec.squeeze(n, e); + dec.decode(n, e); + const gpu = try finish(cb, e); + @memcpy(logits, floats(n.logits, P.logits)); + return .{ .wall = mtl.clock.seconds() - started, .gpu = gpu }; + } + + /// remove_background's alpha [h, w] of the resized photo `im` (its RGB): the input, the logits, then the alpha. + pub fn removeBackground(n: *Net, a: std.mem.Allocator, im: pre.Image, logits: []f32, alpha: []u8) !Timing { + const x = try a.alloc(f32, P.x); + defer a.free(x); + try pre.input(a, im, x); + const t = try n.forward(x, logits); + try pre.alpha(a, logits, im.width, im.height, alpha); + return t; + } + + /// _forward_enc: both backbone passes, concatenated into x1, x2, x3 and the squeeze's input [32^2, 5,760]. + pub fn encode(n: *Net, e: mtl.ComputeEncoder) void { + ks.resize(e, n.k.get("br_resize"), n.x, n.xh, .{ .h = size, .w = size }, .{ .h = size / 2, .w = size / 2 }, 3, 3, 0, 3, 0); + for (0..2) |pass| { + const side = @as(usize, size) >> @intCast(pass); + const g0 = grids[pass]; + ks.tiles(e, n.k.get("br_tiles"), if (pass == 0) n.x else n.xh, n.slots, .{ .h = g0, .w = g0 }, side, 3, 4, 4, false, 64); + n.conv(e, n.w.patch, n.slots, n.tb, .{ .rows = g0 * g0, .ldy = wts.embed }); + n.norm(e, n.tb, null, n.w.patch_norm, n.ta, g0 * g0, wts.embed, .rows, 0, wts.embed, 0); + for (0..4) |i| { + const h = g0 >> @intCast(i); + const C = wts.widths[i]; + n.layer(e, pass, i, n.ta, n.tb); + // the layer's features: norm_i, at the full grid's channels or resized to it beside them + const full = grids[0] >> @intCast(i); + const dst, const ld, const off = n.featureSlot(i); + if (pass == 0) { + n.norm(e, n.ta, null, n.w.norms[i], dst, h * h, C, .rows, 0, ld, off); + } else { + n.norm(e, n.ta, null, n.w.norms[i], n.att, h * h, C, .rows, 0, C, 0); + ks.resize(e, n.k.get("br_resize"), n.att, dst, .{ .h = h, .w = h }, .{ .h = full, .w = full }, C, C, 0, ld, off + C); + } + if (i < 3) std.mem.swap(Ref, &n.ta, &n.tb); + } + } + // the squeeze's input: x1, x2 and x3 resized to the coarsest grid, then x4 + var col: usize = 0; + for (0..3) |i| { + const g = grids[0] >> @intCast(i); + const C = 2 * wts.widths[i]; + ks.resize(e, n.k.get("br_resize"), n.xs[i], n.xs[3], .{ .h = g, .w = g }, .{ .h = 32, .w = 32 }, C, C, 0, wts.squeeze_in, col); + col += C; + } + } + + /// Where layer i's features go (buffer, row length, column): x1 to x3 [grid^2, 2 C], x4 in the squeeze's input. + fn featureSlot(n: *const Net, i: usize) struct { Ref, usize, usize } { + if (i < 3) return .{ n.xs[i], 2 * wts.widths[i], 0 }; + return .{ n.xs[3], wts.squeeze_in, wts.squeeze_in - wts.channels[0] }; + } + + /// Swin layer i of `pass` on t [h * h, C] in place, then (i < 3) PatchMerging into down [h^2 / 4, 2 C]. + pub fn layer(n: *Net, e: mtl.ComputeEncoder, pass: usize, i: usize, t: Ref, down: Ref) void { + const h = grids[pass] >> @intCast(i); + const hp = padded(h); + const C = wts.widths[i]; + const L = &n.w.layers[i]; + for (L.blocks[0..wts.depths[i]], 0..) |*b, j| { + const shift: usize = if (j % 2 == 1) wts.window / 2 else 0; + const map = n.map[pass][i][j % 2]; + const slots = hp * hp; + n.norm(e, t, map, b.norm1, n.slots, slots, C, .map, 0, C, 0); + n.conv(e, b.qkv, n.slots, n.big, .{ .rows = slots, .ldy = 3 * C }); + ks.window(e, n.k.get("br_window"), n.big, b.table, n.att, C, wts.heads[i], hp, hp, shift); + n.convWith(e, b.proj, n.att, t, t, map, .{ .rows = slots, .ldy = C, .flags = ks.flag_add | ks.flag_map }); + n.norm(e, t, null, b.norm2, n.slots, h * h, C, .rows, 0, C, 0); + n.conv(e, b.fc1, n.slots, n.big, .{ .rows = h * h, .ldy = 4 * C, .flags = ks.flag_gelu }); + n.convWith(e, b.fc2, n.big, t, t, null, .{ .rows = h * h, .ldy = C, .flags = ks.flag_add }); + } + if (L.merge) |m| { + n.norm(e, t, null, m, n.slots, h * h / 4, 4 * C, .merge, h, 4 * C, 0); + n.conv(e, L.reduction.?, n.slots, down, .{ .rows = h * h / 4, .ldy = 2 * C }); + } + } + + /// The shape fields `conv` fills from c. + pub const Out = struct { rows: usize, cols: usize = 1, ldy: usize, yoff: usize = 0, flags: u32 = 0, row0: usize = 0, ldx: ?usize = null }; + + /// y = c(x) with its batch norm when it has one (flags add bn), rows x cols at c's side. + pub fn conv(n: *const Net, e: mtl.ComputeEncoder, c: wts.Conv, x: Ref, y: Ref, o: Out) void { + n.convWith(e, c, x, null, y, null, o); + } + + pub fn convWith(n: *const Net, e: mtl.ComputeEncoder, c: wts.Conv, x: Ref, r: ?Ref, y: Ref, map: ?Ref, o: Out) void { + const flags = o.flags | (if (c.post != null) ks.flag_bn else 0); + ks.conv(e, n.k.get("br_conv"), x, c.w, c.b, c.post, r, map, y, .{ .rows = o.rows, .cols = o.cols, .cin = c.cin, .ldx = o.ldx orelse c.cin, .cout = c.cout, .ldw = c.ldw, .k = c.k, .ldy = o.ldy, .yoff = o.yoff, .flags = flags, .row0 = o.row0 }); + } + + pub fn norm(n: *const Net, e: mtl.ComputeEncoder, x: Ref, map: ?Ref, nm: wts.Norm, y: Ref, rows: usize, width: usize, mode: ks.NormMode, src_w: usize, ldy: usize, yoff: usize) void { + ks.norm(e, n.k.get("br_norm"), x, map, nm.g, nm.b, y, rows, width, mode, src_w, ldy, yoff, wts.ln_eps); + } +}; + +/// Grid h's window map shifted by `shift`: each slot (144 a window) to the token torch.roll puts there, -1 padding. +pub fn fillMap(out: []i32, h: usize, shift: usize) void { + const hp = padded(h); + const per = hp / wts.window; + for (out, 0..) |*m, slot| { + const win = slot / 144; + const t = slot % 144; + const sy = (win / per) * wts.window + t / wts.window; + const sx = (win % per) * wts.window + t % wts.window; + const oy = (sy + shift) % hp; + const ox = (sx + shift) % hp; + m.* = if (oy < h and ox < h) @intCast(oy * h + ox) else -1; + } +} + +test "a window map sends every token to one slot, the padding nowhere" { + const a = std.testing.allocator; + for ([_]usize{ 0, 6 }) |shift| { + const h = 16; + const hp = padded(h); + const m = try a.alloc(i32, hp * hp); + defer a.free(m); + fillMap(m, h, shift); + var seen: [h * h]u8 = @splat(0); + var pad: usize = 0; + for (m) |v| { + if (v < 0) pad += 1 else seen[@intCast(v)] += 1; + } + try std.testing.expectEqual(hp * hp - h * h, pad); + for (seen) |s| try std.testing.expectEqual(@as(u8, 1), s); + } + var m: [24 * 24]i32 = undefined; + fillMap(&m, 16, 6); + // the first window's first slot holds token (6, 6) once rolled by -6 + try std.testing.expectEqual(@as(i32, 6 * 16 + 6), m[0]); +} diff --git a/zig/src/families/idol/birefnet_decoder.zig b/zig/src/families/idol/birefnet_decoder.zig new file mode 100644 index 0000000000..6f3c0cbd90 --- /dev/null +++ b/zig/src/families/idol/birefnet_decoder.zig @@ -0,0 +1,102 @@ +//! BiRefNet's squeeze and decoder on Metal: the blocks, image patches, gradient gates, laterals and the head. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("birefnet_kernels.zig"); +const wts = @import("birefnet_weights.zig"); +const bn = @import("birefnet.zig"); +const Ref = ks.Ref; +const Net = bn.Net; +const C = wts.channels; + +/// Each level's input channels (the previous level's output plus its image patches' convolution) and grid side. +pub const level_in = [4]usize{ C[0] + C[0] / 8, C[1] + C[0] / 8, C[2] + C[1] / 8, C[3] + C[2] / 8 }; +pub const level_side = [4]usize{ 32, 64, 128, 256 }; +/// Each level's output channels (decoder_block4 .. 1), and the image patches' tile side at each level. +pub const level_out = [4]usize{ C[1], C[2], C[3], C[3] / 2 }; +const ipt_tile = [4]usize{ 32, 16, 8, 4 }; + +/// squeeze_module: the encoder's [32^2, 5,760] to the first level's input [32^2, 3,456], its first 3,072 channels. +pub fn squeeze(n: *Net, e: mtl.ComputeEncoder) void { + block(n, e, &n.w.squeeze, n.xs[3], wts.squeeze_in, 32, n.lin[0], level_in[0], 0); +} + +/// _BasicDecBlk on x [side^2, cin]: conv_in, BN, ReLU; ASPP; conv_out, BN into y [side^2, ldy] from column yoff. +pub fn block(n: *Net, e: mtl.ComputeEncoder, d: *const wts.DecBlock, x: Ref, cin: usize, side: usize, y: Ref, ldy: usize, yoff: usize) void { + n.conv(e, d.conv_in, x, n.a64, .{ .rows = side, .cols = side, .ldy = wts.inter, .flags = ks.flag_relu, .ldx = cin }); + aspp(n, e, &d.aspp, side); + n.conv(e, d.conv_out, n.gdt, y, .{ .rows = side, .cols = side, .ldy = ldy, .yoff = yoff, .ldx = wts.inter }); +} + +/// _ASPPDeformable, a64 into gdt [side^2, 64]: four deformable branches and the pooled one, concatenated, conv1. +fn aspp(n: *Net, e: mtl.ComputeEncoder, a: *const wts.Aspp, side: usize) void { + const pixels = side * side; + const ld = 5 * wts.aspp_planes; + const g: ks.Grid = .{ .h = side, .w = side }; + for (a.branches, wts.branch_k, 0..) |b, k, i| { + n.conv(e, b.offs, n.a64, n.offs, .{ .rows = side, .cols = side, .ldy = b.offs.ldw }); + var first: usize = 0; + while (first < pixels) : (first += bn.chunk) { + const count = @min(bn.chunk, pixels - first); + ks.deform(e, n.k.get("br_deform"), n.a64, n.offs, n.cols, g, k, first, count, b.offs.ldw); + n.conv(e, b.regular, n.cols, n.cat, .{ .rows = count, .ldy = ld, .yoff = i * wts.aspp_planes, .flags = ks.flag_relu, .row0 = first }); + } + } + ks.pool(e, n.k.get("br_pool"), n.a64, n.pooled, pixels, wts.inter); + const pooled_out = n.pooled.at(wts.inter * 4); + n.conv(e, a.pool, n.pooled, pooled_out, .{ .rows = 1, .ldy = wts.aspp_planes, .flags = ks.flag_relu }); + ks.resize(e, n.k.get("br_resize"), pooled_out, n.cat, .{ .h = 1, .w = 1 }, g, wts.aspp_planes, wts.aspp_planes, 0, ld, 4 * wts.aspp_planes); + n.conv(e, a.out, n.cat, n.gdt, .{ .rows = pixels, .ldy = wts.inter, .flags = ks.flag_relu }); +} + +/// ipt_blk at level l: x's tiles of the level's side, conv1 and conv_out into the level's input past its channels. +fn patches(n: *Net, e: mtl.ComputeEncoder, l: usize) void { + const side = level_side[l]; + const t = ipt_tile[l]; + const s = &n.w.ipt[l]; + const ch = 3 * t * t; + ks.tiles(e, n.k.get("br_tiles"), n.x, n.patches, .{ .h = side, .w = side }, bn.size, 3, t, t, true, s.conv1.cin); + n.conv(e, s.conv1, n.patches, n.a64, .{ .rows = side, .cols = side, .ldy = wts.inter }); + std.debug.assert(ch <= s.conv1.cin); + n.conv(e, s.conv_out, n.a64, n.lin[l], .{ .rows = side, .cols = side, .ldy = level_in[l], .yoff = level_in[l] - s.conv_out.cout }); +} + +/// The decoder from lin[0]'s first 3,072 channels (the squeezed x4), xs[0..3] and x: the mask logits [1024^2]. +pub fn decode(n: *Net, e: mtl.ComputeEncoder) void { + for (0..4) |l| { + if (l > 0) { + // _p = resize(p, this grid) + lateral(x_l), beside this level's patches + const prev = level_side[l - 1]; + const side = level_side[l]; + const cl = level_out[l - 1]; + ks.resize(e, n.k.get("br_resize"), n.out[l - 1], n.lin[l], .{ .h = prev, .w = prev }, .{ .h = side, .w = side }, cl, cl, 0, level_in[l], 0); + n.convWith(e, n.w.lateral[l - 1], n.xs[3 - l], n.lin[l], n.lin[l], null, .{ .rows = side * side, .ldy = level_in[l], .flags = ks.flag_add }); + } + patches(n, e, l); + level(n, e, l); + } + head(n, e); +} + +/// decoder_block at level l on lin[l] into out[l], then (levels 0 to 2) gated by its gradient attention. +pub fn level(n: *Net, e: mtl.ComputeEncoder, l: usize) void { + const side = level_side[l]; + block(n, e, &n.w.dec[l], n.lin[l], level_in[l], side, n.out[l], level_out[l], 0); + if (l == 3) return; + n.conv(e, n.w.gdt[l], n.out[l], n.gdt, .{ .rows = side, .cols = side, .ldy = 16, .flags = ks.flag_relu, .ldx = level_out[l] }); + ks.gate(e, n.k.get("br_gate"), n.gdt, n.w.gdt_attn[l].w, n.w.gdt_attn[l].b, n.out[l], side * side, level_out[l]); +} + +/// ipt_blk1 on the picture (3 to 64 to 48 channels at 1024^2) and conv_out1 over [resize(out[3]) | that]: the logits. +pub fn head(n: *Net, e: mtl.ComputeEncoder) void { + const s = &n.w.ipt[4]; + const full: ks.Grid = .{ .h = bn.size, .w = bn.size }; + ks.conv3(e, n.k.get("br_conv3"), n.x, s.conv1.w, s.conv1.b, n.cat, full, wts.inter); + n.conv(e, s.conv_out, n.cat, n.q48, .{ .rows = bn.size, .cols = bn.size, .ldy = s.conv_out.cout, .ldx = wts.inter }); + const side = level_side[3]; + ks.head(e, n.k.get("br_head"), n.out[3], n.q48, n.w.head.w, n.w.head.b, n.logits, .{ .h = side, .w = side }, full, level_out[3], s.conv_out.cout); +} + +test "the levels' widths are the decoder blocks' inputs" { + try std.testing.expectEqual([4]usize{ 3456, 1920, 960, 480 }, level_in); + try std.testing.expectEqual(@as(usize, 240), level_out[3] + C[3] / 8); +} diff --git a/zig/src/families/idol/birefnet_kernels.zig b/zig/src/families/idol/birefnet_kernels.zig new file mode 100644 index 0000000000..a97365f5bb --- /dev/null +++ b/zig/src/families/idol/birefnet_kernels.zig @@ -0,0 +1,161 @@ +//! BiRefNet's pipelines (idol_birefnet.metal, fp32) and their launches. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const Ref = @import("birefnet_weights.zig").Ref; + +/// Each pipeline by its function. +pub const names = [_][]const u8{ "br_conv", "br_norm", "br_window", "br_resize", "br_tiles", "br_conv3", "br_deform", "br_pool", "br_gate", "br_head" }; + +pub const Kernels = struct { + lib: mtl.Library, + p: [names.len]mtl.Pipeline, + + pub fn load(device: mtl.Device) !Kernels { + return loadSource(device, sources.idol_birefnet); + } + + /// The pipelines of `text` (the served source, or a contract control's copy of it). + pub fn loadSource(device: mtl.Device, text: []const u8) !Kernels { + var k: Kernels = undefined; + k.lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + errdefer k.lib.deinit(); + var made: usize = 0; + errdefer for (k.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + k.p[i] = try mtl.Pipeline.init(device, k.lib, name, false); + made += 1; + } + return k; + } + + pub fn deinit(k: *Kernels) void { + for (k.p) |p| p.deinit(); + k.lib.deinit(); + } + + pub fn get(k: *const Kernels, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return k.p[i]; + @compileError("no pipeline " ++ name); + } +}; + +/// One kernel's buffers in order, then its parameters at the next index. +pub fn bind(e: mtl.ComputeEncoder, p: mtl.Pipeline, refs: []const Ref, params: anytype) void { + e.setPipeline(p); + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(params, refs.len); +} + +pub const flag_bn: u32 = 1; +pub const flag_relu: u32 = 2; +pub const flag_gelu: u32 = 4; +pub const flag_add: u32 = 8; +pub const flag_map: u32 = 16; + +/// br_conv's launch: picture rows x cols, channels a tap (% 32), ldx, outputs, ldw (% 64), side, output, flags. +pub const ConvShape = struct { + rows: usize, + cols: usize = 1, + cin: usize, + ldx: usize, + cout: usize, + ldw: usize, + k: usize = 1, + ldy: usize, + yoff: usize = 0, + flags: u32 = 0, + row0: usize = 0, +}; + +pub const ConvParams = extern struct { rows: u32, cols: u32, cin: u32, ldx: u32, cout: u32, ldw: u32, k: u32, ldy: u32, yoff: u32, flags: u32, row0: u32 }; + +pub fn convParams(s: ConvShape) ConvParams { + return .{ .rows = @intCast(s.rows), .cols = @intCast(s.cols), .cin = @intCast(s.cin), .ldx = @intCast(s.ldx), .cout = @intCast(s.cout), .ldw = @intCast(s.ldw), .k = @intCast(s.k), .ldy = @intCast(s.ldy), .yoff = @intCast(s.yoff), .flags = s.flags, .row0 = @intCast(s.row0) }; +} + +/// br_conv: y = conv(x) + b and its epilogue; post, r and map are bound where the flags need them (else x). +pub fn conv(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, post: ?Ref, r: ?Ref, map: ?Ref, y: Ref, s: ConvShape) void { + std.debug.assert(s.cin % 32 == 0 and s.ldw % 64 == 0 and s.ldx % 4 == 0 and s.cout <= s.ldw); + bind(e, p, &.{ x, w, b, post orelse x, r orelse y, map orelse x, y }, convParams(s)); + e.dispatchGroups(mtl.Size.of(s.ldw / 64, (s.rows * s.cols + 63) / 64, 1), mtl.Size.of(128, 1, 1)); +} + +/// br_norm's modes: a row as it is, the row a map names (negative: zeros), PatchMerging's 2 x 2 patch. +pub const NormMode = enum(u32) { rows = 0, map = 1, merge = 2 }; +pub const NormParams = extern struct { rows: u32, width: u32, mode: u32, src_w: u32, ldy: u32, yoff: u32, eps: f32 }; + +/// br_norm: y [rows, ldy] from yoff = LayerNorm of each row (`width` wide) with g and beta. +pub fn norm(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, map: ?Ref, g: Ref, beta: Ref, y: Ref, rows: usize, width: usize, mode: NormMode, src_w: usize, ldy: usize, yoff: usize, eps: f32) void { + bind(e, p, &.{ x, map orelse x, g, beta, y }, NormParams{ .rows = @intCast(rows), .width = @intCast(width), .mode = @backingInt(mode), .src_w = @intCast(src_w), .ldy = @intCast(ldy), .yoff = @intCast(yoff), .eps = eps }); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); +} + +pub const WindowParams = extern struct { width: u32, heads: u32, rows: u32, cols: u32, shift: u32, scale: f32 }; + +/// br_window: o [windows * 144, width] = each window's attention of qkv [windows * 144, 3 * width], grid rows x cols. +pub fn window(e: mtl.ComputeEncoder, p: mtl.Pipeline, qkv: Ref, table: Ref, o: Ref, width: usize, heads: usize, rows: usize, cols: usize, shift: usize) void { + const scale: f32 = @floatCast(std.math.pow(f64, @floatFromInt(width / heads), -0.5)); + bind(e, p, &.{ qkv, table, o }, WindowParams{ .width = @intCast(width), .heads = @intCast(heads), .rows = @intCast(rows), .cols = @intCast(cols), .shift = @intCast(shift), .scale = scale }); + e.dispatchGroups(mtl.Size.of(3, heads, (rows / 12) * (cols / 12)), mtl.Size.of(192, 1, 1)); +} + +pub const Grid = struct { h: usize, w: usize }; +pub const ResizeParams = extern struct { in_h: u32, in_w: u32, out_h: u32, out_w: u32, channels: u32, ldx: u32, xoff: u32, ldy: u32, yoff: u32 }; + +/// br_resize: bilinear with align_corners of `channels` columns of x (row length ldx, from xoff) into y's (ldy, yoff). +pub fn resize(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, y: Ref, in: Grid, out: Grid, channels: usize, ldx: usize, xoff: usize, ldy: usize, yoff: usize) void { + bind(e, p, &.{ x, y }, ResizeParams{ .in_h = @intCast(in.h), .in_w = @intCast(in.w), .out_h = @intCast(out.h), .out_w = @intCast(out.w), .channels = @intCast(channels), .ldx = @intCast(ldx), .xoff = @intCast(xoff), .ldy = @intCast(ldy), .yoff = @intCast(yoff) }); + e.dispatchThreads(mtl.Size.of(channels, out.h * out.w, 1), mtl.Size.of(@min(channels, 64), @as(usize, 256) / @min(channels, 64), 1)); +} + +pub const TilesParams = extern struct { out_h: u32, out_w: u32, src_w: u32, channels: u32, a: u32, b: u32, mode: u32, cpad: u32 }; + +/// br_tiles: dst [out.h * out.w, cpad] from src [., src_w, channels], a x b patches (mode 0) or grid tiles (mode 1). +pub fn tiles(e: mtl.ComputeEncoder, p: mtl.Pipeline, src: Ref, dst: Ref, out: Grid, src_w: usize, channels: usize, a: usize, b: usize, grid: bool, cpad: usize) void { + bind(e, p, &.{ src, dst }, TilesParams{ .out_h = @intCast(out.h), .out_w = @intCast(out.w), .src_w = @intCast(src_w), .channels = @intCast(channels), .a = @intCast(a), .b = @intCast(b), .mode = @intFromBool(grid), .cpad = @intCast(cpad) }); + e.dispatchThreads(mtl.Size.of(cpad, out.h * out.w, 1), mtl.Size.of(@min(cpad, 64), @as(usize, 256) / @min(cpad, 64), 1)); +} + +/// br_conv3: y [pixels, cout] = conv3x3(x) + b of a 3-channel picture, w [27, cout]. +pub fn conv3(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, y: Ref, g: Grid, cout: usize) void { + bind(e, p, &.{ x, w, b, y }, [3]u32{ @intCast(g.h), @intCast(g.w), @intCast(cout) }); + e.dispatchThreads(mtl.Size.of(cout, g.h * g.w, 1), mtl.Size.of(64, 4, 1)); +} + +pub const DeformParams = extern struct { rows: u32, cols: u32, k: u32, first: u32, count: u32, ldo: u32 }; + +/// br_deform: the columns [count, k * k * 64] of pixels first .. first + count, offsets and modulators in offs. +pub fn deform(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, offs: Ref, col: Ref, g: Grid, k: usize, first: usize, count: usize, ldo: usize) void { + bind(e, p, &.{ x, offs, col }, DeformParams{ .rows = @intCast(g.h), .cols = @intCast(g.w), .k = @intCast(k), .first = @intCast(first), .count = @intCast(count), .ldo = @intCast(ldo) }); + const taps = k * k; + const tx = @min(taps, 64); + e.dispatchThreads(mtl.Size.of(taps, count, 1), mtl.Size.of(tx, @as(usize, 256) / tx, 1)); +} + +/// br_pool: mean [channels] over x's pixels. +pub fn pool(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, mean: Ref, pixels: usize, channels: usize) void { + bind(e, p, &.{ x, mean }, [2]u32{ @intCast(pixels), @intCast(channels) }); + e.dispatchGroups(mtl.Size.of(channels, 1, 1), mtl.Size.of(256, 1, 1)); +} + +/// br_gate: x [pixels, channels] *= sigmoid(g [pixels, 16] . w + b). +pub fn gate(e: mtl.ComputeEncoder, p: mtl.Pipeline, g: Ref, w: Ref, b: Ref, x: Ref, pixels: usize, channels: usize) void { + bind(e, p, &.{ g, w, b, x }, [2]u32{ @intCast(pixels), @intCast(channels) }); + e.dispatchThreads(mtl.Size.of(channels, pixels, 1), mtl.Size.of(64, 4, 1)); +} + +/// br_head: logits [out] = the head's 1x1 over [a resized from `in` | q]. +pub fn head(e: mtl.ComputeEncoder, p: mtl.Pipeline, a: Ref, q: Ref, w: Ref, b: Ref, y: Ref, in: Grid, out: Grid, cp: usize, cq: usize) void { + bind(e, p, &.{ a, q, w, b, y }, [6]u32{ @intCast(in.h), @intCast(in.w), @intCast(out.h), @intCast(out.w), @intCast(cp), @intCast(cq) }); + e.dispatchThreads(mtl.Size.of(out.h * out.w, 1, 1), mtl.Size.of(256, 1, 1)); +} + +test "every background remover pipeline compiles and loads on this GPU" { + const pool_ = mtl.objc.Pool.push(); + defer pool_.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + var k = try Kernels.load(device); + k.deinit(); +} diff --git a/zig/src/families/idol/birefnet_reference.zig b/zig/src/families/idol/birefnet_reference.zig new file mode 100644 index 0000000000..6f5aafb6a5 --- /dev/null +++ b/zig/src/families/idol/birefnet_reference.zig @@ -0,0 +1,102 @@ +//! tf-idol-birefnet's gates against the float64 ruler, each within twice torch's own distance from it. +const std = @import("std"); +const ref = @import("reference.zig"); + +pub const Match = ref.Match; + +/// Ours and torch's against the ruler in rows of `width`: worst row and 1 - cosine within twice torch's. +pub const Verdict = struct { ours: Match, torch: Match, pass: bool }; + +pub fn gate(ours: []const f32, torch: []const f32, ruler: []const f32, width: usize) Verdict { + const o = match(ours, ruler, width); + const t = match(torch, ruler, width); + return .{ .ours = o, .torch = t, .pass = ref.within(o, t) }; +} + +/// reference.match, 1 - cos as |a/|a| - b/|b||^2 / 2: a dot product cannot resolve 1e-12 and turns negative. +pub fn match(ours: []const f32, theirs: []const f32, width: usize) Match { + var m = ref.match(ours, theirs, width); + var aa: f64 = 0; + var bb: f64 = 0; + for (ours, theirs) |x, y| { + aa += @as(f64, x) * x; + bb += @as(f64, y) * y; + } + const na = @sqrt(@max(aa, 1e-300)); + const nb = @sqrt(@max(bb, 1e-300)); + var d: f64 = 0; + for (ours, theirs) |x, y| { + const e = @as(f64, x) / na - @as(f64, y) / nb; + d += e * e; + } + m.cosine = 1 - d / 2; + return m; +} + +/// Pixels (of `channels` bytes) where two 8-bit pictures differ, and the largest difference. +pub const Diff = struct { pixels: usize, max: u8 }; + +pub fn diff(a: []const u8, b: []const u8, channels: usize) Diff { + std.debug.assert(a.len == b.len and a.len % channels == 0); + var d: Diff = .{ .pixels = 0, .max = 0 }; + var i: usize = 0; + while (i < a.len) : (i += channels) { + var most: u8 = 0; + for (a[i..][0..channels], b[i..][0..channels]) |x, y| most = @max(most, if (x > y) x - y else y - x); + if (most > 0) d.pixels += 1; + d.max = @max(d.max, most); + } + return d; +} + +/// 8-bit against the ruler: differing pixels within twice torch's, largest difference within torch's (or one). +pub fn withinDiff(ours: Diff, torch: Diff) bool { + return ours.pixels <= 2 * torch.pixels and ours.max <= @max(torch.max, 1); +} + +/// The picture: pixels and largest difference within twice torch's (one alpha level moves a pixel up to two). +pub fn withinPicture(ours: Diff, torch: Diff) bool { + return ours.pixels <= 2 * torch.pixels and ours.max <= 2 * @max(torch.max, 1); +} + +test "a planted 2% error in one logits row fails the gate" { + const a = std.testing.allocator; + const n = 64 * 1024; + const ruler = try a.alloc(f32, n); + defer a.free(ruler); + const torch = try a.alloc(f32, n); + defer a.free(torch); + const ours = try a.alloc(f32, n); + defer a.free(ours); + var rng = std.Random.DefaultPrng.init(3); + for (ruler, torch, ours) |*r, *t, *o| { + r.* = 8 * rng.random().floatNorm(f32); + t.* = r.* * (1 + 1e-5 * rng.random().floatNorm(f32)); + o.* = r.* * (1 + 1e-6 * rng.random().floatNorm(f32)); + } + try std.testing.expect(gate(ours, torch, ruler, 1024).pass); + for (ours[17 * 1024 ..][0..1024]) |*o| o.* *= 1.02; + try std.testing.expect(!gate(ours, torch, ruler, 1024).pass); +} + +test "a planted mismatch in the alpha or the picture fails its gate" { + const base = [6]u8{ 0, 10, 200, 254, 128, 3 }; + var ruler: [24]u8 = undefined; + for (&ruler, 0..) |*r, i| r.* = base[i % 6]; + var torch = ruler; + torch[1] += 1; + torch[7] += 1; + var ours = ruler; + ours[2] += 1; + try std.testing.expect(withinDiff(diff(&ours, &ruler, 1), diff(&torch, &ruler, 1))); + ours[5] += 2; + try std.testing.expect(!withinDiff(diff(&ours, &ruler, 1), diff(&torch, &ruler, 1))); + ours[5] -= 2; + for ([_]usize{ 8, 9, 10, 11 }) |i| ours[i] +%= 1; + try std.testing.expect(!withinDiff(diff(&ours, &ruler, 1), diff(&torch, &ruler, 1))); + try std.testing.expect(withinPicture(.{ .pixels = 27, .max = 2 }, .{ .pixels = 26, .max = 1 })); + try std.testing.expect(!withinPicture(.{ .pixels = 27, .max = 3 }, .{ .pixels = 26, .max = 1 })); + try std.testing.expect(!withinPicture(.{ .pixels = 53, .max = 1 }, .{ .pixels = 26, .max = 1 })); + // three channels a pixel: one pixel, however many of its channels differ + try std.testing.expectEqual(@as(usize, 1), diff(&[_]u8{ 1, 2, 3 }, &[_]u8{ 2, 3, 4 }, 3).pixels); +} diff --git a/zig/src/families/idol/birefnet_weights.zig b/zig/src/families/idol/birefnet_weights.zig new file mode 100644 index 0000000000..38a335ea8f --- /dev/null +++ b/zig/src/families/idol/birefnet_weights.zig @@ -0,0 +1,314 @@ +//! BiRefNet's weights (speridlabs/idol's birefnet.safetensors, fp16) as fp32 GEMM operands in one Metal buffer. +const std = @import("std"); +const mtl = @import("metal"); +const ck = @import("../../core/checkpoint_metal.zig"); +pub const Ref = @import("../../core/qwen_vl/weights.zig").Ref; + +/// The released geometry: Swin-L (embed 192, depths 2 2 18 2, heads 6 12 24 48, windows of 12), decoder widths. +pub const embed = 192; +pub const depths = [4]usize{ 2, 2, 18, 2 }; +pub const heads = [4]usize{ 6, 12, 24, 48 }; +pub const window = 12; +pub const widths = [4]usize{ 192, 384, 768, 1536 }; +/// The decoder's channels after the backbone's two passes are concatenated (mul_scl_ipt='cat'). +pub const channels = [4]usize{ 3072, 1536, 768, 384 }; +pub const inter = 64; +/// The squeeze's input: the four concatenated features, the three finer ones resized to the coarsest grid. +pub const squeeze_in = channels[0] + channels[1] + channels[2] + channels[3]; +pub const aspp_planes = 256; +/// The tensors the file holds: 742, ZhengPeng7/BiRefNet's less its 12 training-only ones. +pub const tensor_count = 742; +pub const ln_eps: f32 = 1e-5; +pub const bn_eps: f64 = 1e-5; + +/// A conv or linear as a GEMM: w [k * k * cin, ldw] by tap, bias [ldw], post [mean | scale | shift] after a BN. +pub const Conv = struct { w: Ref, b: Ref, post: ?Ref = null, cin: usize, cout: usize, ldw: usize, k: usize }; +pub const Norm = struct { g: Ref, b: Ref }; +pub const Block = struct { norm1: Norm, qkv: Conv, table: Ref, proj: Conv, norm2: Norm, fc1: Conv, fc2: Conv }; +pub const Layer = struct { blocks: [18]Block, merge: ?Norm, reduction: ?Conv }; +/// A modulated deformable conv: offsets (2 k^2) and modulators (k^2) as one conv, then the conv over its columns. +pub const Deform = struct { offs: Conv, regular: Conv }; +pub const Aspp = struct { branches: [4]Deform, pool: Conv, out: Conv }; +pub const DecBlock = struct { conv_in: Conv, aspp: Aspp, conv_out: Conv }; +pub const Simple = struct { conv1: Conv, conv_out: Conv }; +pub const Dot = struct { w: Ref, b: Ref }; + +/// The deformable branches' kernel sides: aspp1, then aspp_deforms 0, 1, 2. +pub const branch_k = [4]usize{ 1, 1, 3, 7 }; + +pub const Weights = struct { + buf: mtl.Buffer, + patch: Conv, + patch_norm: Norm, + layers: [4]Layer, + norms: [4]Norm, + squeeze: DecBlock, + /// ipt_blk5 .. ipt_blk1; ipt_blk1's conv1 is [27, 64] for br_conv3. + ipt: [5]Simple, + /// decoder_block4 .. decoder_block1, then lateral_block4 .. 2, gdt_convs_4 .. 2 and their attention 1x1s. + dec: [4]DecBlock, + lateral: [3]Conv, + gdt: [3]Conv, + gdt_attn: [3]Dot, + head: Dot, + + pub fn deinit(w: *Weights) void { + w.buf.deinit(); + } +}; + +/// `path` (birefnet.safetensors) checked against the geometry and packed; Swin's position indices must match. +pub fn load(gpa: std.mem.Allocator, device: mtl.Device, path: [:0]const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + var c = ck.Checkpoint.init(gpa); + defer c.deinit(); + try c.addFile(device, path, ""); + if (c.tensors.count() != tensor_count) { + std.log.err("idol birefnet: {s} has {d} tensors; BiRefNet's Swin-L release has {d}", .{ path, c.tensors.count(), tensor_count }); + return error.UnexpectedTensors; + } + var w: Weights = undefined; + var p: Packer = .{ .c = &c, .arena = arena_state.allocator() }; + try layout(&w, &p); + w.buf = try device.buffer(p.at * 4, mtl.ResourceOptions.shared); + errdefer w.buf.deinit(); + p.out = @as([*]f32, @ptrCast(@alignCast(w.buf.contents())))[0..p.at]; + p.at = 0; + p.read = 0; + p.buf = w.buf; + try layout(&w, &p); + if (p.read != tensor_count) { + std.log.err("idol birefnet: read {d} of the {d} tensors", .{ p.read, tensor_count }); + return error.UnexpectedTensors; + } + return w; +} + +fn layout(w: *Weights, p: *Packer) !void { + var buf: [128]u8 = undefined; + w.patch = try p.conv("bb.patch_embed.proj", 3, embed, 4, .{ .cin_pad = 64, .tap_major = false }); + w.patch_norm = try p.norm("bb.patch_embed.norm", embed); + for (&w.layers, 0..) |*l, i| { + const C = widths[i]; + for (l.blocks[0..depths[i]], 0..) |*b, j| { + const pre = try std.fmt.bufPrint(&buf, "bb.layers.{d}.blocks.{d}", .{ i, j }); + var nb: [160]u8 = undefined; + b.norm1 = try p.norm(try cat(&nb, pre, ".norm1"), C); + b.qkv = try p.conv(try cat(&nb, pre, ".attn.qkv"), C, 3 * C, 1, .{}); + b.table = try p.copy(try cat(&nb, pre, ".attn.relative_position_bias_table"), &.{ 529, heads[i] }); + try p.index(try cat(&nb, pre, ".attn.relative_position_index")); + b.proj = try p.conv(try cat(&nb, pre, ".attn.proj"), C, C, 1, .{}); + b.norm2 = try p.norm(try cat(&nb, pre, ".norm2"), C); + b.fc1 = try p.conv(try cat(&nb, pre, ".mlp.fc1"), C, 4 * C, 1, .{}); + b.fc2 = try p.conv(try cat(&nb, pre, ".mlp.fc2"), 4 * C, C, 1, .{}); + } + if (i < 3) { + l.merge = try p.norm(try std.fmt.bufPrint(&buf, "bb.layers.{d}.downsample.norm", .{i}), 4 * C); + l.reduction = try p.conv(try std.fmt.bufPrint(&buf, "bb.layers.{d}.downsample.reduction", .{i}), 4 * C, 2 * C, 1, .{ .bias = false }); + } else { + l.merge = null; + l.reduction = null; + } + w.norms[i] = try p.norm(try std.fmt.bufPrint(&buf, "bb.norm{d}", .{i}), C); + } + w.squeeze = try p.decBlock("squeeze_module.0", squeeze_in, channels[0]); + const ipt_in = [5]usize{ 3 << 10, 3 << 8, 3 << 6, 3 << 4, 3 }; + const ipt_out = [5]usize{ channels[0] / 8, channels[0] / 8, channels[1] / 8, channels[2] / 8, channels[3] / 8 }; + for (&w.ipt, 0..) |*s, i| { + const n = 5 - i; + s.conv1 = if (n == 1) + try p.conv("decoder.ipt_blk1.conv1", 3, inter, 3, .{ .direct = true }) + else + try p.conv(try std.fmt.bufPrint(&buf, "decoder.ipt_blk{d}.conv1", .{n}), ipt_in[i], inter, 3, .{}); + s.conv_out = try p.conv(try std.fmt.bufPrint(&buf, "decoder.ipt_blk{d}.conv_out", .{n}), inter, ipt_out[i], 3, .{}); + } + const dec_in = [4]usize{ channels[0] + channels[0] / 8, channels[1] + channels[0] / 8, channels[2] + channels[1] / 8, channels[3] + channels[2] / 8 }; + const dec_out = [4]usize{ channels[1], channels[2], channels[3], channels[3] / 2 }; + for (&w.dec, 0..) |*d, i| d.* = try p.decBlock(try std.fmt.bufPrint(&buf, "decoder.decoder_block{d}", .{4 - i}), dec_in[i], dec_out[i]); + for (0..3) |i| { + const n = 4 - i; + const C = channels[i + 1]; + w.lateral[i] = try p.conv(try std.fmt.bufPrint(&buf, "decoder.lateral_block{d}.conv", .{n}), C, C, 1, .{}); + w.gdt[i] = try p.conv(try std.fmt.bufPrint(&buf, "decoder.gdt_convs_{d}.0", .{n}), C, 16, 3, .{}); + w.gdt[i].post = try p.batchNorm(try std.fmt.bufPrint(&buf, "decoder.gdt_convs_{d}.1", .{n}), 16, w.gdt[i].ldw); + w.gdt_attn[i] = try p.dot(try std.fmt.bufPrint(&buf, "decoder.gdt_convs_attn_{d}.0", .{n}), 16); + } + w.head = try p.dot("decoder.conv_out1.0", channels[3] / 2 + channels[3] / 8); +} + +fn cat(buf: []u8, a: []const u8, b: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{s}", .{ a, b }); +} + +const Options = struct { bias: bool = true, cin_pad: ?usize = null, tap_major: bool = true, direct: bool = false }; + +/// Packs the tensors in one fp32 buffer: sizes on the first pass (out null), copies on the second. +const Packer = struct { + c: *const ck.Checkpoint, + arena: std.mem.Allocator, + buf: mtl.Buffer = undefined, + out: ?[]f32 = null, + at: usize = 0, + read: usize = 0, + + fn alloc(p: *Packer, n: usize) Ref { + const r: Ref = .{ .buf = p.buf, .off = p.at * 4 }; + p.at = std.mem.alignForward(usize, p.at + n, 64); + return r; + } + + fn dst(p: *Packer, r: Ref, n: usize) ?[]f32 { + const out = p.out orelse return null; + return out[r.off / 4 ..][0..n]; + } + + /// Tensor `name` of `shape`, fp16 as the file holds it, widened to fp32 exactly (as the SDK's load does). + fn tensor(p: *Packer, name: []const u8, shape: []const usize) ![]const f32 { + const t = p.c.get(name) catch { + std.log.err("idol birefnet: {s} is missing", .{name}); + return error.MissingTensor; + }; + if (t.dtype != .f16 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("idol birefnet: {s} is {t} {any}; BiRefNet's is f16 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } + p.read += 1; + if (p.out == null) return &.{}; + const v = try p.arena.alloc(f32, t.count()); + for (v, t.host(f16)) |*o, h| o.* = h; + return v; + } + + fn copy(p: *Packer, name: []const u8, shape: []const usize) !Ref { + var n: usize = 1; + for (shape) |s| n *= s; + const v = try p.tensor(name, shape); + const r = p.alloc(n); + if (p.dst(r, n)) |d| @memcpy(d, v); + return r; + } + + fn norm(p: *Packer, name: []const u8, width: usize) !Norm { + var buf: [160]u8 = undefined; + return .{ .g = try p.copy(try cat(&buf, name, ".weight"), &.{width}), .b = try p.copy(try cat(&buf, name, ".bias"), &.{width}) }; + } + + /// relative_position_index must be (dy + 11) * 23 + dx + 11, as br_window computes it. + fn index(p: *Packer, name: []const u8) !void { + const t = try p.c.get(name); + if (t.dtype != .i64 or t.rank != 2 or t.shape[0] != 144 or t.shape[1] != 144) return error.UnexpectedTensor; + p.read += 1; + if (p.out == null) return; + for (t.host(i64), 0..) |v, at| { + const i = at / 144; + const j = at % 144; + const want: i64 = @intCast((i / 12 + 11 - j / 12) * 23 + (i % 12 + 11 - j % 12)); + if (v != want) return error.UnexpectedRelativeIndex; + } + } + + /// Batch norm `name` as post [mean | weight / sqrt(var + eps) | bias]; its counter is read and ignored. + fn batchNorm(p: *Packer, name: []const u8, n: usize, ldw: usize) !Ref { + var buf: [160]u8 = undefined; + const mean = try p.tensor(try cat(&buf, name, ".running_mean"), &.{n}); + const var_ = try p.tensor(try cat(&buf, name, ".running_var"), &.{n}); + const g = try p.tensor(try cat(&buf, name, ".weight"), &.{n}); + const b = try p.tensor(try cat(&buf, name, ".bias"), &.{n}); + const counter = try p.c.get(try cat(&buf, name, ".num_batches_tracked")); + if (counter.dtype != .i64 or counter.rank != 0) return error.UnexpectedTensor; + p.read += 1; + const r = p.alloc(3 * ldw); + if (p.dst(r, 3 * ldw)) |d| { + @memset(d, 0); + for (0..n) |i| { + d[i] = mean[i]; + d[ldw + i] = @floatCast(@as(f64, g[i]) / @sqrt(@as(f64, var_[i]) + bn_eps)); + d[2 * ldw + i] = b[i]; + } + } + return r; + } + + /// Conv `name` as w [taps * cin_pad, ldw] (row tap * cin_pad + c); tap_major false: torch's order; direct: [27, n]. + fn conv(p: *Packer, name: []const u8, cin: usize, cout: usize, k: usize, o: Options) !Conv { + var buf: [160]u8 = undefined; + const taps = k * k; + const cin_pad = o.cin_pad orelse if (o.direct) cin else std.mem.alignForward(usize, cin, 32); + const ldw = if (o.direct) cout else std.mem.alignForward(usize, cout, 64); + const linear = std.mem.indexOf(u8, name, "bb.layers") != null and k == 1; + const shape: []const usize = if (linear) &.{ cout, cin } else &.{ cout, cin, k, k }; + const g = try p.tensor(try cat(&buf, name, ".weight"), shape); + const rows = if (o.tap_major) taps * cin_pad else cin_pad; + const r = p.alloc(rows * ldw); + if (p.dst(r, rows * ldw)) |d| { + @memset(d, 0); + for (0..cout) |co| for (0..cin) |ci| for (0..taps) |tap| { + const row = if (o.tap_major) tap * cin_pad + ci else ci * taps + tap; + d[row * ldw + co] = g[(co * cin + ci) * taps + tap]; + }; + } + const b = p.alloc(ldw); + if (o.bias) { + const v = try p.tensor(try cat(&buf, name, ".bias"), &.{cout}); + if (p.dst(b, ldw)) |d| { + @memset(d, 0); + @memcpy(d[0..cout], v); + } + } else if (p.dst(b, ldw)) |d| @memset(d, 0); + return .{ .w = r, .b = b, .cin = cin_pad, .cout = cout, .ldw = ldw, .k = if (o.tap_major) k else 1 }; + } + + /// A one-output 1x1 [1, n, 1, 1] and its bias, as plain vectors. + fn dot(p: *Packer, name: []const u8, n: usize) !Dot { + var buf: [160]u8 = undefined; + return .{ .w = try p.copy(try cat(&buf, name, ".weight"), &.{ 1, n, 1, 1 }), .b = try p.copy(try cat(&buf, name, ".bias"), &.{1}) }; + } + + /// One deformable branch `name` (atrous_conv's three convolutions and bn) with kernel side k over 64 channels. + fn deform(p: *Packer, name: []const u8, k: usize) !Deform { + var buf: [160]u8 = undefined; + var nb: [160]u8 = undefined; + const taps = k * k; + const ldo = std.mem.alignForward(usize, 3 * taps, 64); + const rows = taps * inter; + const offs: Conv = .{ .w = p.alloc(rows * ldo), .b = p.alloc(ldo), .cin = inter, .cout = 3 * taps, .ldw = ldo, .k = k }; + if (p.dst(offs.w, rows * ldo)) |d| @memset(d, 0); + if (p.dst(offs.b, ldo)) |d| @memset(d, 0); + for ([_][]const u8{ ".atrous_conv.offset_conv", ".atrous_conv.modulator_conv" }, [_]usize{ 0, 2 * taps }, [_]usize{ 2 * taps, taps }) |part, first, n| { + const pre = try cat(&nb, name, part); + const g = try p.tensor(try cat(&buf, pre, ".weight"), &.{ n, inter, k, k }); + const b = try p.tensor(try cat(&buf, pre, ".bias"), &.{n}); + if (p.dst(offs.w, rows * ldo)) |d| for (0..n) |co| for (0..inter) |ci| for (0..taps) |tap| { + d[(tap * inter + ci) * ldo + first + co] = g[(co * inter + ci) * taps + tap]; + }; + if (p.dst(offs.b, ldo)) |d| @memcpy(d[first..][0..n], b); + } + var regular = try p.conv(try cat(&buf, name, ".atrous_conv.regular_conv"), inter, aspp_planes, k, .{ .bias = false }); + regular.cin = taps * inter; + regular.k = 1; + regular.post = try p.batchNorm(try cat(&buf, name, ".bn"), aspp_planes, regular.ldw); + return .{ .offs = offs, .regular = regular }; + } + + fn decBlock(p: *Packer, name: []const u8, cin: usize, cout: usize) !DecBlock { + var buf: [160]u8 = undefined; + var d: DecBlock = undefined; + d.conv_in = try p.conv(try cat(&buf, name, ".conv_in"), cin, inter, 3, .{}); + d.conv_in.post = try p.batchNorm(try cat(&buf, name, ".bn_in"), inter, d.conv_in.ldw); + const names = [4][]const u8{ ".dec_att.aspp1", ".dec_att.aspp_deforms.0", ".dec_att.aspp_deforms.1", ".dec_att.aspp_deforms.2" }; + for (&d.aspp.branches, names, branch_k) |*b, n, k| b.* = try p.deform(try cat(&buf, name, n), k); + d.aspp.pool = try p.conv(try cat(&buf, name, ".dec_att.global_avg_pool.1"), inter, aspp_planes, 1, .{ .bias = false }); + d.aspp.pool.post = try p.batchNorm(try cat(&buf, name, ".dec_att.global_avg_pool.2"), aspp_planes, d.aspp.pool.ldw); + d.aspp.out = try p.conv(try cat(&buf, name, ".dec_att.conv1"), 5 * aspp_planes, inter, 1, .{ .bias = false }); + d.aspp.out.post = try p.batchNorm(try cat(&buf, name, ".dec_att.bn1"), inter, d.aspp.out.ldw); + d.conv_out = try p.conv(try cat(&buf, name, ".conv_out"), inter, cout, 3, .{}); + d.conv_out.post = try p.batchNorm(try cat(&buf, name, ".bn_out"), cout, d.conv_out.ldw); + return d; + } +}; + +test "the decoder's widths follow the backbone's: 5,760 channels into the squeeze, 240 into the head" { + try std.testing.expectEqual(@as(usize, 5760), squeeze_in); + try std.testing.expectEqual(@as(usize, 240), channels[3] / 2 + channels[3] / 8); +} diff --git a/zig/src/families/idol/idol.zig b/zig/src/families/idol/idol.zig index e3292c44f6..112444eb73 100644 --- a/zig/src/families/idol/idol.zig +++ b/zig/src/families/idol/idol.zig @@ -19,6 +19,14 @@ pub const vae_kernels = @import("vae_kernels.zig"); pub const vae_weights = @import("vae_weights.zig"); /// Its DINOv3 ViT-H+/16 image encoder: an image to feature1. pub const dino = @import("dino.zig"); +/// Its background remover (BiRefNet) and the SDK's preprocess: a raw photo to the 1024 x 1024 picture. +pub const birefnet = @import("birefnet.zig"); +pub const birefnet_kernels = @import("birefnet_kernels.zig"); +pub const birefnet_weights = @import("birefnet_weights.zig"); +pub const birefnet_decoder = @import("birefnet_decoder.zig"); +pub const birefnet_reference = @import("birefnet_reference.zig"); +pub const preprocess = @import("preprocess.zig"); +pub const photo = @import("photo.zig"); pub const Model = model.Model; pub const Condition = model.Condition; pub const Timing = model.Timing; @@ -42,4 +50,11 @@ test { _ = vae_kernels; _ = vae_weights; _ = dino; + _ = birefnet; + _ = birefnet_kernels; + _ = birefnet_weights; + _ = birefnet_decoder; + _ = birefnet_reference; + _ = preprocess; + _ = photo; } diff --git a/zig/src/families/idol/photo.zig b/zig/src/families/idol/photo.zig new file mode 100644 index 0000000000..3a343365b0 --- /dev/null +++ b/zig/src/families/idol/photo.zig @@ -0,0 +1,120 @@ +//! A raw photo to IdolPipeline.preprocess's 1024 x 1024 picture (BiRefNet's alpha unless real), and a PNG writer. +const std = @import("std"); +const mtl = @import("metal"); +const pre = @import("preprocess.zig"); +const bn = @import("birefnet.zig"); + +/// Seconds of each step (preprocess: all of it but the load, set by the caller), the Metal bytes BiRefNet held. +pub const Timing = struct { decode: f64 = 0, load: f64 = 0, birefnet: f64 = 0, birefnet_gpu: f64 = 0, finish: f64 = 0, preprocess: f64 = 0, real_alpha: bool = false, metal_peak: u64 = 0 }; + +/// Whether `path` names a photo (.png, .jpg, .jpeg or .webp) rather than a safetensors picture. +pub fn isPhoto(path: []const u8) bool { + const ext = std.fs.path.extension(path); + for ([_][]const u8{ ".png", ".jpg", ".jpeg", ".webp" }) |e| if (std.ascii.eqlIgnoreCase(ext, e)) return true; + return false; +} + +/// The picture [1024 * 1024 * 3] of the photo at `path`; BiRefNet is loaded from `weights` only if needed. +pub fn preprocess(a: std.mem.Allocator, io: std.Io, device: mtl.Device, weights: [:0]const u8, path: []const u8, crop_margin: f64, out: []u8, t: *Timing) !void { + var started = mtl.clock.seconds(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 28)); + defer a.free(bytes); + var photo = try pre.decode(io, a, bytes); + defer photo.deinit(a); + var resized = try pre.resizeFirst(a, photo); + defer resized.deinit(a); + t.decode = mtl.clock.seconds() - started; + t.real_alpha = pre.hasRealAlpha(resized); + var alpha: ?[]u8 = null; + defer if (alpha) |al| a.free(al); + if (!t.real_alpha) { + started = mtl.clock.seconds(); + const net = try bn.Net.load(std.heap.page_allocator, device, weights); + defer net.deinit(); + t.load = mtl.clock.seconds() - started; + const logits = try a.alloc(f32, bn.size * bn.size); + defer a.free(logits); + alpha = try a.alloc(u8, resized.width * resized.height); + started = mtl.clock.seconds(); + const tm = try net.removeBackground(a, resized, logits, alpha.?); + t.birefnet = mtl.clock.seconds() - started; + t.birefnet_gpu = tm.gpu; + t.metal_peak = device.allocated(); + } + started = mtl.clock.seconds(); + try pre.finish(a, resized, alpha, crop_margin, out); + t.finish = mtl.clock.seconds() - started; +} + +/// An 8-bit RGB PNG of `rgb` [h, w, 3], its image data in stored (uncompressed) deflate blocks. +pub fn png(a: std.mem.Allocator, rgb: []const u8, w: usize, h: usize) ![]u8 { + var raw: std.ArrayList(u8) = .empty; + defer raw.deinit(a); + for (0..h) |y| { + try raw.append(a, 0); + try raw.appendSlice(a, rgb[y * w * 3 ..][0 .. w * 3]); + } + var z: std.ArrayList(u8) = .empty; + defer z.deinit(a); + try z.appendSlice(a, &.{ 0x78, 0x01 }); + var at: usize = 0; + while (true) { + const n = @min(raw.items.len - at, 65535); + const last = at + n == raw.items.len; + try z.append(a, @intFromBool(last)); + var len: [4]u8 = undefined; + std.mem.writeInt(u16, len[0..2], @intCast(n), .little); + std.mem.writeInt(u16, len[2..4], ~@as(u16, @intCast(n)), .little); + try z.appendSlice(a, &len); + try z.appendSlice(a, raw.items[at..][0..n]); + at += n; + if (last) break; + } + var adler: [4]u8 = undefined; + std.mem.writeInt(u32, &adler, std.hash.Adler32.hash(raw.items), .big); + try z.appendSlice(a, &adler); + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + try out.appendSlice(a, "\x89PNG\r\n\x1a\n"); + var hdr: [13]u8 = undefined; + std.mem.writeInt(u32, hdr[0..4], @intCast(w), .big); + std.mem.writeInt(u32, hdr[4..8], @intCast(h), .big); + hdr[8..13].* = .{ 8, 2, 0, 0, 0 }; + try chunk(a, &out, "IHDR", &hdr); + try chunk(a, &out, "IDAT", z.items); + try chunk(a, &out, "IEND", ""); + return out.toOwnedSlice(a); +} + +fn chunk(a: std.mem.Allocator, out: *std.ArrayList(u8), kind: *const [4]u8, body: []const u8) !void { + var be: [4]u8 = undefined; + std.mem.writeInt(u32, &be, @intCast(body.len), .big); + try out.appendSlice(a, &be); + try out.appendSlice(a, kind); + try out.appendSlice(a, body); + var crc = std.hash.Crc32.init(); + crc.update(kind); + crc.update(body); + std.mem.writeInt(u32, &be, crc.final(), .big); + try out.appendSlice(a, &be); +} + +test "the PNG writer's picture decodes back to its bytes" { + const a = std.testing.allocator; + const w = 300; + const h = 90; + const rgb = try a.alloc(u8, w * h * 3); + defer a.free(rgb); + for (rgb, 0..) |*v, i| v.* = @truncate(i * 7); + const bytes = try png(a, rgb, w, h); + defer a.free(bytes); + const s = try @import("../../core/picture/png.zig").decode(a, bytes); + defer a.free(s.data); + try std.testing.expectEqual(@as(usize, w), s.width); + for (0..w * h) |i| try std.testing.expectEqualSlices(u8, rgb[i * 3 ..][0..3], s.data[i * s.channels ..][0..3]); +} + +test "a photo is named by its extension, any case" { + try std.testing.expect(isPhoto("a/b.PNG") and isPhoto("x.jpeg") and isPhoto("y.webp") and isPhoto("z.JpG")); + try std.testing.expect(!isPhoto("case/vae.safetensors") and !isPhoto("-")); +} diff --git a/zig/src/families/idol/preprocess.zig b/zig/src/families/idol/preprocess.zig new file mode 100644 index 0000000000..64ff587053 --- /dev/null +++ b/zig/src/families/idol/preprocess.zig @@ -0,0 +1,373 @@ +//! idol-sdk's preprocess_image on the host, byte for byte as Pillow 13's LANCZOS and torch's CPU compute it. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const picture = @import("../../core/picture/picture.zig"); + +pub const size = 1024; +/// CROP_MARGIN by mode: renders and generated images ('synthetic'), photos ('real'). +pub const Mode = enum { synthetic, real }; +pub fn margin(m: Mode) f64 { + return switch (m) { + .synthetic => 1.2, + .real => 1.7, + }; +} + +/// 8-bit pixels: 3 (Pillow's RGB, or L read as RGB) or 4 (RGBA) channels a pixel. +pub const Image = struct { + width: usize, + height: usize, + channels: usize, + pixels: []u8, + + pub fn deinit(im: *Image, a: Allocator) void { + a.free(im.pixels); + im.* = undefined; + } +}; + +/// The words for a photo this path refuses, where Pillow's mode would take another path through the SDK. +pub const unsupported = "only 8-bit RGB, RGBA or grayscale PNGs, RGB or grayscale JPEGs and WebPs are supported"; + +/// PNG, JPEG or WebP bytes as Image.open gives them: no EXIF transpose, RGBA only where Pillow's mode is RGBA. +pub fn decode(io: std.Io, a: Allocator, bytes: []const u8) !Image { + const format = picture.sniff(bytes); + var rgba = false; + switch (format) { + .png => { + // IHDR: depth at 24, colour type at 25; Pillow's P, LA, I;16 and 1-bit modes take other paths + if (bytes.len < 33 or bytes[24] != 8) return error.UnsupportedPhoto; + switch (bytes[25]) { + 0, 2 => {}, + 6 => rgba = true, + else => return error.UnsupportedPhoto, + } + }, + .webp => rgba = webpAlpha(bytes), + .jpeg => {}, + .unknown => return error.UnsupportedPhoto, + } + const s = switch (format) { + .png => try picture.png.decode(a, bytes), + .jpeg => try picture.jpeg.decode(io, a, bytes), + .webp => try picture.webp.decode(io, a, bytes), + .unknown => unreachable, + }; + const ch: usize = if (rgba) 4 else 3; + if (s.channels == ch) return .{ .width = s.width, .height = s.height, .channels = ch, .pixels = s.data }; + defer a.free(s.data); + const out = try a.alloc(u8, s.width * s.height * ch); + for (0..s.width * s.height) |i| @memcpy(out[i * ch ..][0..ch], s.data[i * s.channels ..][0..ch]); + return .{ .width = s.width, .height = s.height, .channels = ch, .pixels = out }; +} + +/// Whether Pillow opens a WebP as RGBA: VP8X's alpha flag, or a lossless image's alpha bit. +fn webpAlpha(b: []const u8) bool { + if (b.len < 30) return false; + if (std.mem.eql(u8, b[12..16], "VP8X")) return b[20] & 0x10 != 0; + if (std.mem.eql(u8, b[12..16], "VP8L")) return (std.mem.readInt(u32, b[21..25], .little) >> 28) & 1 != 0; + return false; +} + +// --- Pillow's resample (Resample.c), LANCZOS, 8 bits: double coefficients, fixed point with 22 fraction bits ------- + +const precision = 22; + +extern "c" fn sin(x: f64) f64; + +fn sinc(x0: f64) f64 { + if (x0 == 0.0) return 1.0; + const x = x0 * std.math.pi; + return sin(x) / x; +} + +fn lanczos(x: f64) f64 { + return if (-3.0 <= x and x < 3.0) sinc(x) * sinc(x / 3) else 0.0; +} + +/// One axis's taps: each output's first input and fixed-point coefficients (precompute_coeffs, normalize_coeffs_8bpc). +const Taps = struct { + first: []usize, + count: []usize, + k: []i32, + ksize: usize, + + fn init(a: Allocator, in: usize, out: usize) !Taps { + const scale = @as(f64, @floatFromInt(in)) / @as(f64, @floatFromInt(out)); + const fs = @max(scale, 1.0); + const support = 3.0 * fs; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const t: Taps = .{ .first = try a.alloc(usize, out), .count = try a.alloc(usize, out), .k = try a.alloc(i32, out * ksize), .ksize = ksize }; + const w = try a.alloc(f64, ksize); + defer a.free(w); + for (0..out) |xx| { + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + const ss = 1.0 / fs; + 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))); + const n: usize = @intCast(hi - lo); + var ww: f64 = 0; + for (0..n) |x| { + w[x] = lanczos((@as(f64, @floatFromInt(@as(i64, @intCast(x)) + lo)) - center + 0.5) * ss); + ww += w[x]; + } + const k = t.k[xx * ksize ..][0..ksize]; + @memset(k, 0); + for (0..n) |x| { + const v = if (ww != 0.0) w[x] / ww else w[x]; + const f = v * @as(f64, 1 << precision); + k[x] = @intFromFloat(if (v < 0) -0.5 + f else 0.5 + f); + } + t.first[xx] = @intCast(lo); + t.count[xx] = n; + } + return t; + } + + fn deinit(t: *Taps, a: Allocator) void { + a.free(t.first); + a.free(t.count); + a.free(t.k); + } +}; + +fn clip8(ss: i32) u8 { + return @intCast(std.math.clamp(ss >> precision, 0, 255)); +} + +/// Image.resize(LANCZOS) of 8-bit pixels: width, then height; a pass keeping its size is skipped, as Pillow skips it. +pub fn resample(a: Allocator, src: []const u8, w: usize, h: usize, ch: usize, ow: usize, oh: usize) ![]u8 { + var cur = try a.dupe(u8, src); + var cw = w; + if (ow != w) { + var t = try Taps.init(a, w, ow); + defer t.deinit(a); + const out = try a.alloc(u8, ow * h * ch); + for (0..h) |y| for (0..ow) |x| for (0..ch) |c| { + var ss: i32 = 1 << (precision - 1); + const k = t.k[x * t.ksize ..]; + for (0..t.count[x]) |i| ss += @as(i32, cur[(y * w + t.first[x] + i) * ch + c]) * k[i]; + out[(y * ow + x) * ch + c] = clip8(ss); + }; + a.free(cur); + cur = out; + cw = ow; + } + if (oh != h) { + var t = try Taps.init(a, h, oh); + defer t.deinit(a); + const out = try a.alloc(u8, cw * oh * ch); + for (0..oh) |y| { + const k = t.k[y * t.ksize ..]; + for (0..cw * ch) |xc| { + var ss: i32 = 1 << (precision - 1); + for (0..t.count[y]) |i| ss += @as(i32, cur[(t.first[y] + i) * cw * ch + xc]) * k[i]; + out[y * cw * ch + xc] = clip8(ss); + } + } + a.free(cur); + cur = out; + } + return cur; +} + +fn muldiv255(v: u32, b: u32) u8 { + const t = v * b + 128; + return @intCast((t + (t >> 8)) >> 8); +} + +/// Image.resize of `im` as Pillow does it: RGBA through RGBa (premultiplied by MULDIV255, back by 255 v / a truncated). +pub fn resize(a: Allocator, im: Image, ow: usize, oh: usize) !Image { + if (im.width == ow and im.height == oh) return .{ .width = ow, .height = oh, .channels = im.channels, .pixels = try a.dupe(u8, im.pixels) }; + if (im.channels == 3) return .{ .width = ow, .height = oh, .channels = 3, .pixels = try resample(a, im.pixels, im.width, im.height, 3, ow, oh) }; + const pre = try a.dupe(u8, im.pixels); + defer a.free(pre); + for (0..im.width * im.height) |i| { + const px = pre[i * 4 ..][0..4]; + for (px[0..3]) |*c| c.* = muldiv255(c.*, px[3]); + } + const out = try resample(a, pre, im.width, im.height, 4, ow, oh); + for (0..ow * oh) |i| { + const px = out[i * 4 ..][0..4]; + const al: u32 = px[3]; + if (al == 0 or al == 255) continue; + for (px[0..3]) |*c| c.* = @intCast(@min(255, (255 * @as(u32, c.*)) / al)); + } + return .{ .width = ow, .height = oh, .channels = 4, .pixels = out }; +} + +/// The first step: the short side to 1,024, each side int(round(side * 1024 / short)) (Python's round: half to even). +pub fn resizedSize(w: usize, h: usize) [2]usize { + const scale = @as(f64, size) / @as(f64, @floatFromInt(@min(w, h))); + return .{ @max(1, roundEven(@as(f64, @floatFromInt(w)) * scale)), @max(1, roundEven(@as(f64, @floatFromInt(h)) * scale)) }; +} + +fn roundEven(v: f64) usize { + const f = @floor(v); + const d = v - f; + const r = if (d > 0.5) f + 1 else if (d < 0.5) f else if (@mod(f, 2) == 0) f else f + 1; + return @intFromFloat(r); +} + +/// has_real_alpha: Pillow's mode RGBA with some alpha under 255. +pub fn hasRealAlpha(im: Image) bool { + if (im.channels != 4) return false; + for (0..im.width * im.height) |i| if (im.pixels[i * 4 + 3] < 255) return true; + return false; +} + +/// BiRefNet's input [1024 * 1024, 3]: x / 255, torch's CPU bilinear (align_corners) to 1,024, (x - mean) / std. +pub fn input(a: Allocator, im: Image, out: []f32) !void { + std.debug.assert(out.len == size * size * 3); + const w = im.width; + const h = im.height; + const v = try a.alloc(f32, w * h * 3); + defer a.free(v); + for (0..w * h) |i| for (0..3) |c| { + v[i * 3 + c] = @as(f32, @floatFromInt(im.pixels[i * im.channels + c])) / 255.0; + }; + const mean = [3]f32{ 0.485, 0.456, 0.406 }; + const sd = [3]f32{ 0.229, 0.224, 0.225 }; + const ys = axis(h, size); + const xs = axis(w, size); + for (0..size) |oy| { + const y = ys.at(oy); + for (0..size) |ox| { + const x = xs.at(ox); + for (0..3) |c| { + const t = bilinear(v, w, c, y, x); + out[(oy * size + ox) * 3 + c] = (t - mean[c]) / sd[c]; + } + } + } +} + +/// One axis of an align_corners resize: scale (in - 1) / (out - 1) in fp32 (0 for one output). +pub const Axis = struct { + scale: f32, + in: usize, + pub const Tap = struct { lo: usize, hi: usize, l0: f32, l1: f32 }; + pub fn at(ax: Axis, o: usize) Tap { + const real = ax.scale * @as(f32, @floatFromInt(o)); + const lo = @min(@as(usize, @intFromFloat(real)), ax.in - 1); + const l1 = @min(@max(real - @as(f32, @floatFromInt(lo)), 0), 1); + return .{ .lo = lo, .hi = @min(lo + 1, ax.in - 1), .l0 = 1 - l1, .l1 = l1 }; + } +}; + +pub fn axis(in: usize, out: usize) Axis { + const scale: f32 = if (out > 1) @as(f32, @floatFromInt(in - 1)) / @as(f32, @floatFromInt(out - 1)) else 0; + return .{ .scale = scale, .in = in }; +} + +/// torch's CPU bilinear of channel c, fma(t0, w0, t1 w1) along each axis: its bytes on a photo resized both ways. +fn bilinear(v: []const f32, w: usize, c: usize, y: Axis.Tap, x: Axis.Tap) f32 { + const row0 = @mulAdd(f32, v[(y.lo * w + x.lo) * 3 + c], x.l0, v[(y.lo * w + x.hi) * 3 + c] * x.l1); + const row1 = @mulAdd(f32, v[(y.hi * w + x.lo) * 3 + c], x.l0, v[(y.hi * w + x.hi) * 3 + c] * x.l1); + return @mulAdd(f32, row0, y.l0, row1 * y.l1); +} + +/// remove_background's alpha [h, w]: sigmoid, bilinear (align_corners), clamp * 255 truncated, in fp32 as on MPS. +pub fn alpha(a: Allocator, logits: []const f32, w: usize, h: usize, out: []u8) !void { + std.debug.assert(logits.len == size * size and out.len == w * h); + const s = try a.alloc(f32, logits.len); + defer a.free(s); + for (s, logits) |*o, l| o.* = 1.0 / (1.0 + @exp(-l)); + const ys = axis(size, h); + const xs = axis(size, w); + for (0..h) |oy| { + const y = ys.at(oy); + for (0..w) |ox| { + const x = xs.at(ox); + const top = (1 - x.l1) * s[y.lo * size + x.lo] + x.l1 * s[y.lo * size + x.hi]; + const bottom = (1 - x.l1) * s[y.hi * size + x.lo] + x.l1 * s[y.hi * size + x.hi]; + const v = (1 - y.l1) * top + y.l1 * bottom; + out[oy * w + ox] = @intFromFloat(std.math.clamp(v, 0, 1) * 255); + } + } +} + +/// The steps after the alpha: erode, crop, resize and paste on black, into the 1024 x 1024 RGB the encoders see. +pub fn finish(a: Allocator, im: Image, given: ?[]const u8, crop_margin: f64, out: []u8) !void { + std.debug.assert(out.len == size * size * 3); + const w = im.width; + const h = im.height; + const al = try a.alloc(u8, w * h); + defer a.free(al); + for (al, 0..) |*o, i| o.* = if (given) |g| g[i] else im.pixels[i * 4 + 3]; + // ImageFilter.MinFilter(3): the minimum over the 3 x 3 neighbours inside the picture + const er = try a.alloc(u8, w * h); + defer a.free(er); + var left: usize = w; + var top: usize = h; + var right: usize = 0; + var bottom: usize = 0; + for (0..h) |y| for (0..w) |x| { + var m: u8 = 255; + for (@max(y, 1) - 1..@min(y + 2, h)) |yy| for (@max(x, 1) - 1..@min(x + 2, w)) |xx| { + m = @min(m, al[yy * w + xx]); + }; + er[y * w + x] = m; + if (m != 0) { + left = @min(left, x); + right = @max(right, x); + top = @min(top, y); + bottom = @max(bottom, y); + } + }; + if (left > right) return error.NoObject; + const cx = @as(f64, @floatFromInt(left + right)) / 2; + const cy = @as(f64, @floatFromInt(top + bottom)) / 2; + const half = @as(f64, @floatFromInt(@max(right - left, bottom - top))) / 2 * crop_margin; + // Python's int() truncates toward zero; Pillow's crop fills outside the picture with zeros + const box = [4]i64{ @intFromFloat(cx - half), @intFromFloat(cy - half), @intFromFloat(cx + half), @intFromFloat(cy + half) }; + if (box[2] <= box[0] or box[3] <= box[1]) return error.NoObject; + const cw: usize = @intCast(box[2] - box[0]); + const ch: usize = @intCast(box[3] - box[1]); + const crop = try a.alloc(u8, cw * ch * 4); + defer a.free(crop); + @memset(crop, 0); + for (0..ch) |y| { + const sy = box[1] + @as(i64, @intCast(y)); + if (sy < 0 or sy >= h) continue; + for (0..cw) |x| { + const sx = box[0] + @as(i64, @intCast(x)); + if (sx < 0 or sx >= w) continue; + const s: usize = @intCast(sy * @as(i64, @intCast(w)) + sx); + const d = crop[(y * cw + x) * 4 ..][0..4]; + @memcpy(d[0..3], im.pixels[s * im.channels ..][0..3]); + d[3] = er[s]; + } + } + var big = try resize(a, .{ .width = cw, .height = ch, .channels = 4, .pixels = crop }, size, size); + defer big.deinit(a); + // paste on black with the alpha as the mask: v * a / 255 by Pillow's BLEND + for (0..size * size) |i| { + const px = big.pixels[i * 4 ..][0..4]; + for (0..3) |c| out[i * 3 + c] = muldiv255(px[c], px[3]); + } +} + +/// preprocess's first step on a decoded photo: its short side to 1,024. +pub fn resizeFirst(a: Allocator, im: Image) !Image { + const s = resizedSize(im.width, im.height); + return resize(a, im, s[0], s[1]); +} + +test "Python's round halves to even, as preprocess's int(round(...)) sizes do" { + try std.testing.expectEqual(@as(usize, 2), roundEven(2.5)); + try std.testing.expectEqual(@as(usize, 4), roundEven(3.5)); + try std.testing.expectEqual(@as(usize, 1365), resizedSize(1280, 960)[0]); + try std.testing.expectEqual(@as(usize, 1024), resizedSize(1280, 960)[1]); +} + +test "Pillow's LANCZOS fixed point keeps a flat picture flat and its premultiply round-trips opaque pixels" { + const a = std.testing.allocator; + const src = try a.alloc(u8, 37 * 53 * 3); + defer a.free(src); + @memset(src, 200); + const out = try resample(a, src, 37, 53, 3, 100, 64); + defer a.free(out); + for (out) |v| try std.testing.expectEqual(@as(u8, 200), v); + for (0..256) |v| try std.testing.expectEqual(@as(u8, @intCast(v)), muldiv255(@intCast(v), 255)); +} diff --git a/zig/tests/contract/metal_idol_birefnet.zig b/zig/tests/contract/metal_idol_birefnet.zig new file mode 100644 index 0000000000..9d24cfa632 --- /dev/null +++ b/zig/tests/contract/metal_idol_birefnet.zig @@ -0,0 +1,202 @@ +//! Idol's background remover's kernels (idol_birefnet.metal, fp32) under the contract; each control breaks one gate. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const cv = @import("metal_idol_birefnet_conv.zig"); +const rw = @import("metal_idol_birefnet_rows.zig"); +const nm = @import("metal_idol_birefnet_norm.zig"); +const wn = @import("metal_idol_birefnet_window.zig"); + +const g = id.g; +const Ends = id.Ends; + +/// A control: the kernel it breaks, the gate that must catch it, and the one text change that breaks it. +const Control = struct { kernel: []const u8, gate: g.Gate, what: []const u8, from: []const u8, to: []const u8 }; + +const flush = "acc[i][j].thread_elements()[0] += part[i][j].thread_elements()[0];"; +const flush_half = "acc[i][j].thread_elements()[0] = float(half(acc[i][j].thread_elements()[0]) + " ++ + "half(part[i][j].thread_elements()[0]));"; +const gelu = "if (p.flags & 4) v = (v * 0.5f) * (1.0f + br_erf(v * 0.70710677f));"; +const gelu_tanh = "if (p.flags & 4) v = (v * 0.5f) * (1.0f + precise::tanh(0.7978846f * (v + 0.044715f * v * v * v)));"; +const bias_at = "const uint at = ((qi / 12 + 11 - j / 12) * 23 + (qi % 12 + 11 - j % 12)) * p.heads + h;"; +const bias_swapped = "const uint at = ((qi % 12 + 11 - j % 12) * 23 + (qi / 12 + 11 - j / 12)) * p.heads + h;"; +const norm_sum = "for (uint c = t; c < p.width; c += 256) s += br_norm_at(x, map, p, row, c);"; +const norm_sum_half = "for (uint c = t; c < p.width; c += 256) s = float(half(s) + half(br_norm_at(x, map, p, row, c)));"; +const pool_sum = "for (uint i = t; i < p.pixels; i += 256) s += x[ulong(i) * p.channels + ch];"; +const pool_sum_half = "for (uint i = t; i < p.pixels; i += 256) s = float(half(s) + half(x[ulong(i) * p.channels + ch]));"; +const outside = "if (!(h > -1.0f && h < float(H) && w > -1.0f && w < float(W))) {"; +const outside_floor = "if (!(h >= 0.0f && h < float(H) && w >= 0.0f && w < float(W))) {"; + +const controls = [_]Control{ + .{ .kernel = "br_conv 1x1", .gate = .shapes, .what = "stages its inputs in half", .from = "for (uint i = 0; i < 4; i++) dst[i] = src[i];", .to = "for (uint i = 0; i < 4; i++) dst[i] = float4(half4(src[i]));" }, + .{ .kernel = "br_conv 3x3", .gate = .shapes, .what = "reads its taps a row low", .from = "const int iy = oy + int(tap / p.k) - pad;", .to = "const int iy = oy + int(tap / p.k) - pad + 1;" }, + .{ .kernel = "br_conv 1x1", .gate = .sums, .what = "adds its blocks in half", .from = flush, .to = flush_half }, + .{ .kernel = "br_conv 1x1 gelu", .gate = .shapes, .what = "GELU's tanh form", .from = gelu, .to = gelu_tanh }, + .{ .kernel = "br_conv 1x1 bn relu", .gate = .nonfinite, .what = "ReLU by max (drops a NaN)", .from = "if (p.flags & 2) v = v < 0.0f ? 0.0f : v;", .to = "if (p.flags & 2) v = max(v, 0.0f);" }, + .{ .kernel = "br_conv 1x1 + residual by map", .gate = .shapes, .what = "subtracts its residual", .from = "if (p.flags & 8) v = r[at] + v;", .to = "if (p.flags & 8) v = r[at] - v;" }, + .{ .kernel = "br_conv3", .gate = .shapes, .what = "drops its bias", .from = "y[ulong(at.y) * p.cout + at.x] = s + b[at.x];", .to = "y[ulong(at.y) * p.cout + at.x] = s;" }, + .{ .kernel = "br_conv3", .gate = .sums, .what = "sums in half", .from = "s = fma(v[c], w[((ky * 3 + kx) * 3 + c) * p.cout + at.x], s);", .to = "s = float(half(s) + half(v[c] * w[((ky * 3 + kx) * 3 + c) * p.cout + at.x]));" }, + .{ .kernel = "br_norm", .gate = .shapes, .what = "the unbiased variance", .from = "const float rstd = 1.0f / precise::sqrt(all / float(p.width) + p.eps);", .to = "const float rstd = 1.0f / precise::sqrt(all / float(p.width - 1) + p.eps);" }, + .{ .kernel = "br_norm", .gate = .sums, .what = "sums in half", .from = norm_sum, .to = norm_sum_half }, + .{ .kernel = "br_norm map", .gate = .shapes, .what = "leaves the padding rows unwritten", .from = "for (uint c = t; c < p.width; c += 256) out[c] = 0.0f;", .to = "" }, + .{ .kernel = "br_norm merge", .gate = .shapes, .what = "swaps x1 and x2", .from = "const uint sy = (r / half_w) * 2 + (part & 1);", .to = "const uint sy = (r / half_w) * 2 + (part >> 1);" }, + .{ .kernel = "br_window shift 6", .gate = .shapes, .what = "drops the shifted windows' mask", .from = "if (p.shift) v = v + (br_region(j, wy, wx, p) == mine ? 0.0f : -100.0f);", .to = "" }, + .{ .kernel = "br_window", .gate = .shapes, .what = "swaps the bias table's axes", .from = bias_at, .to = bias_swapped }, + .{ .kernel = "br_resize", .gate = .shapes, .what = "weighs columns by the row's lambda", .from = "const float top = (1.0f - lx) * v00 + lx * v01;", .to = "const float top = (1.0f - ly) * v00 + ly * v01;" }, + .{ .kernel = "br_tiles grid", .gate = .shapes, .what = "takes the tile's row from its column", .from = "const uint sy = p.mode == 0 ? oy * p.a + i : i * p.out_h + oy;", .to = "const uint sy = p.mode == 0 ? oy * p.a + i : j * p.out_h + oy;" }, + .{ .kernel = "br_deform 3x3", .gate = .shapes, .what = "a modulator of sigmoid, not twice it", .from = "const float mask = 2.0f * br_sigmoid(off[2 * taps + at.x]);", .to = "const float mask = br_sigmoid(off[2 * taps + at.x]);" }, + .{ .kernel = "br_deform 7x7", .gate = .shapes, .what = "zeroes samples within a pixel of the border", .from = outside, .to = outside_floor }, + .{ .kernel = "br_pool", .gate = .sums, .what = "sums in half", .from = pool_sum, .to = pool_sum_half }, + .{ .kernel = "br_gate", .gate = .shapes, .what = "drops the gate's bias", .from = "x[o] = x[o] * br_sigmoid(s + b[0]);", .to = "x[o] = x[o] * br_sigmoid(s);" }, + .{ .kernel = "br_head", .gate = .shapes, .what = "drops ipt_blk1's channels", .from = "for (uint c = 0; c < p.cq; c++) s = fma(qr[c], w[p.cp + c], s);", .to = "" }, +}; + +/// The remover's source as the engine compiles it, with `ctl`'s change made where it occurs, exactly once. +fn compile(a: std.mem.Allocator, device: mtl.Device, ctl: ?Control) !mtl.Library { + const text = if (ctl) |x| blk: { + if (std.mem.count(u8, sources.idol_birefnet, x.from) != 1) { + std.debug.print("control {s}: text not found once: {s}\n", .{ x.what, x.from }); + return error.ControlTextNotFoundOnce; + } + break :blk try std.mem.replaceOwned(u8, a, sources.idol_birefnet, x.from, x.to); + } else try a.dupe(u8, sources.idol_birefnet); + defer a.free(text); + return mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); +} + +/// A test the suite made, and how it goes. +const Made = struct { ptr: *anyopaque, free: *const fn (std.mem.Allocator, *anyopaque) void }; + +pub const Suite = struct { + a: std.mem.Allocator, + pipes: std.ArrayList(mtl.Pipeline) = .empty, + names: std.ArrayList([]u8) = .empty, + libs: std.ArrayList(mtl.Library) = .empty, + made: std.ArrayList(Made) = .empty, + + fn pipeline(s: *Suite, device: mtl.Device, lib: mtl.Library, name: []const u8) !mtl.Pipeline { + const p = try mtl.Pipeline.init(device, lib, name, false); + errdefer p.deinit(); + try s.pipes.append(s.a, p); + return p; + } + + fn label(s: *Suite, comptime fmt: []const u8, args: anytype) ![]const u8 { + const name = try std.fmt.allocPrint(s.a, fmt, args); + errdefer s.a.free(name); + try s.names.append(s.a, name); + return name; + } + + fn make(s: *Suite, comptime T: type, value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + const f = struct { + fn free(a: std.mem.Allocator, p: *anyopaque) void { + const x: *T = @ptrCast(@alignCast(p)); + x.deinit(); + a.destroy(x); + } + }; + try s.made.append(s.a, .{ .ptr = t, .free = f.free }); + return t; + } + + /// Each kernel (or, with `ctl`, the one it breaks) on `lib`'s pipelines. + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: mtl.Library, ctl: ?Control, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + const suffix = if (ctl) |x| try s.label(" control: {s}", .{x.what}) else ""; + const gate: ?g.Gate = if (ctl) |x| x.gate else null; + const rig: rig_.Rig = .{ .a = s.a, .device = device, .queue = queue }; + const ends: Ends = .{ .a = s.a }; + const W = struct { + s: *Suite, + ctl: ?Control, + suffix: []const u8, + only: ?[]const u8, + /// The test's name when this pass registers `kernel`, else null. + fn name(w: @This(), kernel: []const u8) !?[]const u8 { + if (w.ctl) |x| if (!std.mem.eql(u8, x.kernel, kernel)) return null; + const full = try w.s.label("metal idol birefnet {s}{s}", .{ kernel, w.suffix }); + if (w.only) |needle| if (std.mem.indexOf(u8, full, needle) == null) return null; + return full; + } + }; + const w = W{ .s = s, .ctl = ctl, .suffix = suffix, .only = only }; + for (cv.conv_forms) |f| { + const name = (try w.name(f.label)) orelse continue; + const t = try s.make(cv.Conv, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_conv"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("br_conv3")) |name| { + const t = try s.make(cv.Conv3, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_conv3"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]nm.NormForm{ .rows, .map, .merge }, [_][]const u8{ "br_norm", "br_norm map", "br_norm merge" }) |f, label_| { + const name = (try w.name(label_)) orelse continue; + const t = try s.make(nm.Norm, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_norm"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]usize{ 0, 6 }, [_][]const u8{ "br_window", "br_window shift 6" }) |shift, label_| { + const name = (try w.name(label_)) orelse continue; + const t = try s.make(wn.Window, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_window"), .shift = shift, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("br_resize")) |name| { + const t = try s.make(rw.Resize, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_resize"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ false, true }, [_][]const u8{ "br_tiles", "br_tiles grid" }) |grid, label_| { + const name = (try w.name(label_)) orelse continue; + const t = try s.make(rw.Tiles, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_tiles"), .grid = grid, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]usize{ 1, 3, 7 }, [_][]const u8{ "br_deform 1x1", "br_deform 3x3", "br_deform 7x7" }) |side, label_| { + const name = (try w.name(label_)) orelse continue; + const t = try s.make(nm.Deform, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_deform"), .side = side, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("br_pool")) |name| { + const t = try s.make(rw.Pool, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_pool"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("br_gate")) |name| { + const t = try s.make(rw.Gate, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_gate"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + if (try w.name("br_head")) |name| { + const t = try s.make(rw.Head, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "br_head"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + for ([_]?Control{null} ++ controls) |x| { + const lib = try compile(a, device, x); + s.libs.append(a, lib) catch |err| { + lib.deinit(); + return err; + }; + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + for (s.made.items) |m| m.free(a, m.ptr); + s.made.deinit(a); + for (s.pipes.items) |p| p.deinit(); + s.pipes.deinit(a); + for (s.libs.items) |l| l.deinit(); + s.libs.deinit(a); + for (s.names.items) |n| a.free(n); + s.names.deinit(a); + a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_idol_birefnet_conv.zig b/zig/tests/contract/metal_idol_birefnet_conv.zig new file mode 100644 index 0000000000..5c84d87e30 --- /dev/null +++ b/zig/tests/contract/metal_idol_birefnet_conv.zig @@ -0,0 +1,411 @@ +//! br_conv's forms and br_conv3 under the contract: BiRefNet's fp32 sums against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); + +const c = id.c; +const g = id.g; +pub const ks = tf.idol.birefnet_kernels; +const uniform = st.uniform; +const plant = st.plant; + +pub fn ref(r: rig_.Ref) ks.Ref { + return .{ .buf = r.buf, .off = r.off }; +} + +/// GELU's exact form in float64, as torch writes it: x / 2 (1 + erf(x / sqrt 2)). +fn gelu(x: f64) f64 { + return 0.5 * x * (1 + erf(x * std.math.sqrt1_2)); +} + +extern "c" fn erf(x: f64) f64; + +// --- br_conv: an implicit GEMM, k x k at stride 1 (padding k / 2), fp32 sums, then its epilogue ------------------- + +/// A form the network launches; rows: rows (1x1) or picture rows of 9 columns (k x k); n outputs; k channels a tap. +pub const Form = struct { label: []const u8, side: usize = 1, flags: u32 = 0, served: []const c.Shape }; + +const bn = ks.flag_bn; +const relu = ks.flag_relu; +pub const conv_forms = [_]Form{ + .{ .label = "br_conv 1x1", .served = &.{ .{ .rows = 1100, .n = 576, .k = 192 }, .{ .rows = 300, .n = 1152, .k = 384 }, .{ .rows = 130, .n = 2304, .k = 768 }, .{ .rows = 70, .n = 4608, .k = 1536 }, .{ .rows = 1100, .n = 192, .k = 64 }, .{ .rows = 300, .n = 147, .k = 3136 } } }, + .{ .label = "br_conv 1x1 gelu", .flags = ks.flag_gelu, .served = &.{ .{ .rows = 1100, .n = 768, .k = 192 }, .{ .rows = 70, .n = 6144, .k = 1536 } } }, + .{ .label = "br_conv 1x1 + residual", .flags = ks.flag_add, .served = &.{ .{ .rows = 1100, .n = 192, .k = 768 }, .{ .rows = 130, .n = 1536, .k = 1536 } } }, + .{ .label = "br_conv 1x1 + residual by map", .flags = ks.flag_add | ks.flag_map, .served = &.{ .{ .rows = 1100, .n = 192, .k = 192 }, .{ .rows = 300, .n = 768, .k = 768 } } }, + .{ .label = "br_conv 1x1 bn relu", .flags = bn | relu, .served = &.{ .{ .rows = 300, .n = 256, .k = 3136 }, .{ .rows = 1100, .n = 64, .k = 1280 }, .{ .rows = 1, .n = 256, .k = 64 } } }, + .{ .label = "br_conv 3x3", .side = 3, .served = &.{ .{ .rows = 8, .n = 64, .k = 3072 }, .{ .rows = 8, .n = 384, .k = 64 }, .{ .rows = 8, .n = 48, .k = 64 } } }, + .{ .label = "br_conv 3x3 bn relu", .side = 3, .flags = bn | relu, .served = &.{ .{ .rows = 8, .n = 64, .k = 5760 }, .{ .rows = 8, .n = 16, .k = 1536 } } }, + .{ .label = "br_conv 3x3 bn", .side = 3, .flags = bn, .served = &.{ .{ .rows = 8, .n = 1536, .k = 64 }, .{ .rows = 8, .n = 3072, .k = 64 } } }, + .{ .label = "br_conv 7x7", .side = 7, .served = &.{.{ .rows = 8, .n = 147, .k = 64 }} }, +}; + +const cols9 = 9; + +pub const Conv = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + form: Form, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + b: []f32 = &.{}, + post: []f32 = &.{}, + r: []f32 = &.{}, + map: []i32 = &.{}, + ins: [6]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + fn depthOf(s: c.Shape, taps: usize) usize { + return taps * s.k + 2; + } + fn depth1(s: c.Shape) usize { + return depthOf(s, 1); + } + fn depth3(s: c.Shape) usize { + return depthOf(s, 9); + } + fn depth7(s: c.Shape) usize { + return depthOf(s, 49); + } + + pub fn spec(name: []const u8, f: Form) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = switch (f.side) { + 1 => depth1, + 3 => depth3, + else => depth7, + }, + .ops = if (f.flags & ks.flag_gelu != 0) 8 else if (f.flags & bn != 0) 3 else 1, + .out = .f32, + .tiles = .{ if (f.side == 1) .{ .size = 64 } else .{ .size = 8 }, .{ .size = 64 }, .{ .size = 32, .step = 32, .min = 32 } }, + .served = f.served, + }; + } + + pub fn deinit(t: *Conv) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Conv) void { + inline for (.{ &t.x, &t.w, &t.b, &t.post, &t.r }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + t.a.free(t.map); + t.map = &.{}; + } + pub fn kernel(t: *Conv, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Conv); + + fn ldw(t: *const Conv) usize { + return std.mem.alignForward(usize, t.shape.n, 64); + } + fn cols(t: *const Conv) usize { + return if (t.form.side == 1) 1 else cols9; + } + /// The GEMM's rows (a map's slots), and the output rows (a map's targets: one in nine slots is padding). + fn slots(t: *const Conv) usize { + return t.shape.rows * t.cols(); + } + fn outRows(t: *const Conv) usize { + return if (t.form.flags & ks.flag_map != 0) t.slots() - (t.slots() + 4) / 9 else t.slots(); + } + fn has(t: *const Conv, f: u32) bool { + return t.form.flags & f != 0; + } + + pub fn prepare(t: *Conv, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 32 != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + const L = t.ldw(); + const taps = t.form.side * t.form.side; + const K = taps * shape.k; + const px = t.slots(); + t.x = try t.a.alloc(f32, px * shape.k); + t.w = try t.a.alloc(f32, K * L); + t.b = try t.a.alloc(f32, L); + t.post = try t.a.alloc(f32, 3 * L); + t.r = try t.a.alloc(f32, if (t.has(ks.flag_add)) t.outRows() * shape.n else 0); + t.map = try t.a.alloc(i32, if (t.has(ks.flag_map)) px else 0); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + @memset(t.w, 0); + @memset(t.b, 0); + @memset(t.post, 0); + if (fill == .stress) { + // second half mirrors the first, weights 1 then -1: a half's sum is exact in fp32, past fp16, and cancels + const half = shape.k / 2; + const big: f32 = @floatFromInt(@min(32639, (1 << 23) / half - 2041)); + for (0..px) |p| for (0..shape.k) |ci| { + const at = if (ci < half) ci else ci - half; + const v: f32 = if (at % 2 == 0) big else 2041; + t.x[p * shape.k + ci] = if (ci < half) v else v + @as(f32, if (at < 4 and p % 3 == 0) 1 else 0); + }; + for (0..K) |k| for (0..shape.n) |o| { + t.w[k * L + o] = if (k % shape.k < half) 1 else -1; + }; + // GELU on small positive sums; batch norms exact in fp32 + for (t.b[0..shape.n]) |*v| v.* = if (t.has(ks.flag_gelu)) 2 else @floatFromInt(rng.intRangeAtMost(i32, -2, 2)); + for (0..shape.n) |o| { + t.post[o] = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + t.post[L + o] = if (o % 2 == 0) 0.5 else 2; + t.post[2 * L + o] = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + } + for (t.r) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, -3, 3)); + } else { + uniform(rng, t.x, -2, 2); + for (0..K) |k| uniform(rng, t.w[k * L ..][0..shape.n], -0.1, 0.1); + uniform(rng, t.b[0..shape.n], -1, 1); + uniform(rng, t.post[0..shape.n], -0.5, 0.5); + uniform(rng, t.post[L..][0..shape.n], 0.2, 3); + uniform(rng, t.post[2 * L ..][0..shape.n], -1, 1); + uniform(rng, t.r, -3, 3); + if (fill == .nonfinite) plant(t.x, shape.k); + } + if (t.map.len > 0) { + // every ninth slot padding, the rest a shuffled permutation of the output rows + const targets = try t.a.alloc(i32, t.outRows()); + defer t.a.free(targets); + for (targets, 0..) |*v, i| v.* = @intCast(i); + rng.shuffle(i32, targets); + var next: usize = 0; + for (t.map, 0..) |*m, s| { + if (s % 9 == 4) { + m.* = -1; + } else { + m.* = targets[next]; + next += 1; + } + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.post)); + if (t.r.len > 0) t.ins[4] = try t.rig.input(std.mem.sliceAsBytes(t.r)); + if (t.map.len > 0) t.ins[5] = try t.rig.input(std.mem.sliceAsBytes(t.map)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Conv) !void { + const p = t.rig.begin(); + const s: ks.ConvShape = .{ .rows = t.shape.rows, .cols = t.cols(), .cin = t.shape.k, .ldx = t.shape.k, .cout = t.shape.n, .ldw = t.ldw(), .k = t.form.side, .ldy = t.shape.n, .flags = t.form.flags }; + ks.conv(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.ins[3]), if (t.r.len > 0) ref(t.ins[4]) else null, if (t.map.len > 0) ref(t.ins[5]) else null, ref(t.out), s); + try t.rig.finish(p); + } + + pub fn count(t: *Conv) usize { + return t.outRows() * t.shape.n; + } + + /// Slot p's products for output o in the kernel's order (tap, then channel); taps past the picture are 0. + fn terms(t: *const Conv, p: usize, o: usize, out: []f64) void { + const K = t.shape.k; + const side = t.form.side; + const pad: isize = @intCast(side / 2); + const oy: isize = @intCast(p / t.cols()); + const ox: isize = @intCast(p % t.cols()); + for (0..side * side) |tap| { + const iy = oy + @as(isize, @intCast(tap / side)) - pad; + const ix = ox + @as(isize, @intCast(tap % side)) - pad; + const inside = iy >= 0 and ix >= 0 and iy < t.shape.rows and ix < t.cols(); + for (0..K) |ci| { + out[tap * K + ci] = if (inside) blk: { + const at = (@as(usize, @intCast(iy)) * t.cols() + @as(usize, @intCast(ix))) * K + ci; + break :blk @as(f64, t.x[at]) * t.w[(tap * K + ci) * t.ldw() + o]; + } else 0; + } + } + } + + /// Output row of slot p (a map's target), or null where the map pads. + fn target(t: *const Conv, p: usize) ?usize { + if (t.map.len == 0) return p; + return if (t.map[p] < 0) null else @intCast(t.map[p]); + } + + /// The epilogue on s = sum + bias: its value, and the magnitude its error scales with. + fn epilogue(t: *const Conv, s: f64, mag: f64, o: usize, i: usize) [2]f64 { + const L = t.ldw(); + var v = s; + var m = mag; + if (t.has(bn)) { + const scale: f64 = t.post[L + o]; + v = (v - t.post[o]) * scale + t.post[2 * L + o]; + m = @abs(scale) * (m + @abs(@as(f64, t.post[o]))) + @abs(@as(f64, t.post[2 * L + o])); + } + if (t.has(relu) and v < 0) v = 0; + if (t.has(ks.flag_gelu)) { + m = 1.13 * m + @abs(v); + v = gelu(v); + } + if (t.has(ks.flag_add)) { + v = t.r[i] + v; + m += @abs(@as(f64, t.r[i])); + } + return .{ v, m }; + } + + pub fn reference(t: *Conv, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, t.form.side * t.form.side * t.shape.k); + defer t.a.free(buf); + for (0..t.slots()) |p| { + const row = t.target(p) orelse continue; + for (0..t.shape.n) |o| { + t.terms(p, o, buf); + var sum: f64 = 0; + var m: f64 = 0; + for (buf) |v| { + sum += v; + m += @abs(v); + } + const i = row * t.shape.n + o; + const e = t.epilogue(sum + t.b[o], m + @abs(@as(f64, t.b[o])), o, i); + want[i] = e[0]; + mag[i] = if (std.math.isFinite(e[1])) e[1] else 0; + } + } + } + + pub fn simulate(t: *Conv, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, t.form.side * t.form.side * t.shape.k); + defer t.a.free(buf); + for (0..t.slots()) |p| { + const row = t.target(p) orelse continue; + for (0..t.shape.n) |o| { + t.terms(p, o, buf); + const i = row * t.shape.n + o; + out[i] = t.epilogue(contract.rows.sumIn(acc, order, buf) + t.b[o], 0, o, i)[0]; + } + } + } +}; + +// --- br_conv3: a 3x3 of the 3-channel picture, fp32 FMAs ------------------------------------------------------------- + +pub const Conv3 = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + w: []f32 = &.{}, + b: []f32 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + fn depth(_: c.Shape) usize { + return 27 + 1; + } + + /// Rows: the picture's rows; n: output channels; k: its columns. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = depth, .out = .f32, .tiles = .{ .{ .size = 4 }, .{ .size = 64 }, .{ .size = 4 } }, .served = &.{ .{ .rows = 12, .n = 64, .k = 1024 }, .{ .rows = 1024, .n = 64, .k = 12 }, .{ .rows = 37, .n = 64, .k = 41 } } }; + } + + pub fn deinit(t: *Conv3) void { + inline for (.{ &t.x, &t.w, &t.b }) |p| t.a.free(p.*); + t.rig.deinit(); + } + pub fn kernel(t: *Conv3, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Conv3); + + pub fn prepare(t: *Conv3, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0 or shape.n == 0) return false; + t.rig.reset(); + inline for (.{ &t.x, &t.w, &t.b }) |p| t.a.free(p.*); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.k * 3); + t.w = try t.a.alloc(f32, 27 * shape.n); + t.b = try t.a.alloc(f32, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // pixels 2041 and 127 by parity, weights 509, -509, 1 by channel: partial sums ~2^20, cancelling + for (0..shape.rows * shape.k) |p| { + const v: f32 = if (p % 2 == 0) 2041 else 127; + t.x[p * 3] = v; + t.x[p * 3 + 1] = v + 1; + t.x[p * 3 + 2] = 0; + } + for (0..27) |r| @memset(t.w[r * shape.n ..][0..shape.n], switch (r % 3) { + 0 => 509, + 1 => -509, + else => 1, + }); + @memset(t.b, 0); + } else { + uniform(rng, t.x, -2, 2); + uniform(rng, t.w, -0.5, 0.5); + uniform(rng, t.b, -1, 1); + if (fill == .nonfinite) plant(t.x, shape.k * 3); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.b)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Conv3) !void { + const p = t.rig.begin(); + ks.conv3(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.out), .{ .h = t.shape.rows, .w = t.shape.k }, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Conv3) usize { + return t.shape.rows * t.shape.k * t.shape.n; + } + + fn terms(t: *const Conv3, px: usize, o: usize, out: *[27]f64) void { + const py: isize = @intCast(px / t.shape.k); + const pxx: isize = @intCast(px % t.shape.k); + for (0..3) |ky| for (0..3) |kx| for (0..3) |ch| { + const iy = py + @as(isize, @intCast(ky)) - 1; + const ix = pxx + @as(isize, @intCast(kx)) - 1; + const at = (ky * 3 + kx) * 3 + ch; + out[at] = if (iy < 0 or ix < 0 or iy >= t.shape.rows or ix >= t.shape.k) 0 else @as(f64, t.x[(@as(usize, @intCast(iy)) * t.shape.k + @as(usize, @intCast(ix))) * 3 + ch]) * t.w[at * t.shape.n + o]; + }; + } + + pub fn reference(t: *Conv3, want: []f64, mag: []f64) !void { + var buf: [27]f64 = undefined; + for (0..t.shape.rows * t.shape.k) |px| for (0..t.shape.n) |o| { + t.terms(px, o, &buf); + var s: f64 = 0; + var m: f64 = 0; + for (buf) |v| { + s += v; + m += @abs(v); + } + const i = px * t.shape.n + o; + want[i] = s + t.b[o]; + m += @abs(@as(f64, t.b[o])); + mag[i] = if (std.math.isFinite(m)) m else 0; + }; + } + + pub fn simulate(t: *Conv3, acc: c.Accum, order: g.Order, out: []f64) !void { + var buf: [27]f64 = undefined; + for (0..t.shape.rows * t.shape.k) |px| for (0..t.shape.n) |o| { + t.terms(px, o, &buf); + out[px * t.shape.n + o] = contract.rows.sumIn(acc, order, &buf) + t.b[o]; + }; + } +}; diff --git a/zig/tests/contract/metal_idol_birefnet_norm.zig b/zig/tests/contract/metal_idol_birefnet_norm.zig new file mode 100644 index 0000000000..31882b4d2d --- /dev/null +++ b/zig/tests/contract/metal_idol_birefnet_norm.zig @@ -0,0 +1,391 @@ +//! br_norm (LayerNorm by rows, by a window map, and PatchMerging's) and br_deform under the contract against float64. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); +const cv = @import("metal_idol_birefnet_conv.zig"); + +const c = id.c; +const g = id.g; +const Range = id.Range; +const Ends = id.Ends; +const ks = cv.ks; +const ref = cv.ref; +const uniform = st.uniform; +const plant = st.plant; +const U = 0x1p-24; + +// --- br_norm: LayerNorm in two fp32 passes; a map's padding rows zeros; PatchMerging's 2 x 2 gather ----------------- + +pub const NormForm = enum { rows, map, merge }; + +/// PatchMerging's source grid in the test: 8 tokens a row, so 4 output rows a source row pair. +const merge_w = 8; + +pub const Norm = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + form: NormForm, + shape: c.Shape = undefined, + x: []f32 = &.{}, + map: []i32 = &.{}, + gamma: []f32 = &.{}, + beta: []f32 = &.{}, + allow: []f64 = &.{}, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// A thread's strided chain, the simdgroup's five levels, the eight simdgroups' chain. + fn depth(s: c.Shape) usize { + return (s.n + 255) / 256 + 5 + 8; + } + + /// Rows: output rows; n: the width (PatchMerging's 4 C); k: the same, the sums' length. + pub fn spec(name: []const u8, f: NormForm) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = if (f == .merge) .{ .{ .size = 8, .step = 4, .min = 4 }, .{}, .{} } else .{ .{ .size = 8 }, .{}, .{} }, + .served = switch (f) { + .rows => &.{ .{ .rows = 300, .n = 192, .k = 192 }, .{ .rows = 130, .n = 384, .k = 384 }, .{ .rows = 70, .n = 768, .k = 768 }, .{ .rows = 40, .n = 1536, .k = 1536 } }, + .map => &.{ .{ .rows = 288, .n = 192, .k = 192 }, .{ .rows = 144, .n = 1536, .k = 1536 } }, + .merge => &.{ .{ .rows = 64, .n = 768, .k = 768 }, .{ .rows = 32, .n = 1536, .k = 1536 }, .{ .rows = 16, .n = 3072, .k = 3072 } }, + }, + .sums = .{ .exact = false }, + }; + } + + pub fn deinit(t: *Norm) void { + t.free(); + t.rig.deinit(); + } + fn free(t: *Norm) void { + inline for (.{ &t.x, &t.gamma, &t.beta }) |p| { + t.a.free(p.*); + p.* = &.{}; + } + t.a.free(t.map); + t.map = &.{}; + t.a.free(t.allow); + t.allow = &.{}; + } + pub fn kernel(t: *Norm, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.form), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Norm); + + /// The map's padding: every seventh row (as torch pads after the norm, those rows are zeros). + fn pad(r: usize) bool { + return r % 7 == 3; + } + fn sourceRows(t: *const Norm) usize { + return switch (t.form) { + .rows => t.shape.rows, + .map => t.shape.rows - (t.shape.rows + 3) / 7, + .merge => 4 * t.shape.rows, + }; + } + fn width(t: *const Norm) usize { + return if (t.form == .merge) t.shape.n / 4 else t.shape.n; + } + + pub fn prepare(t: *Norm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 4 != 0 or shape.k != shape.n) return false; + if (t.form == .merge and shape.rows % (merge_w / 2) != 0) return false; + t.rig.reset(); + t.free(); + t.shape = shape; + t.x = try t.a.alloc(f32, t.sourceRows() * t.width()); + t.gamma = try t.a.alloc(f32, shape.n); + t.beta = try t.a.alloc(f32, shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // integers to 2041 a row: its sum exact in fp32, past fp16's range in a few dozen terms + for (t.x) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, 0, 2041)); + } else { + uniform(rng, t.x, -4, 4); + // a row far off zero: the mean's rounding then dominates the deviations + for (t.x[0..t.width()]) |*v| v.* += 300; + if (fill == .nonfinite) plant(t.x, t.width()); + } + uniform(rng, t.gamma, -2, 2); + uniform(rng, t.beta, -1, 1); + if (t.form == .map) { + const order = try t.a.alloc(i32, t.sourceRows()); + defer t.a.free(order); + for (order, 0..) |*v, i| v.* = @intCast(i); + rng.shuffle(i32, order); + t.map = try t.a.alloc(i32, shape.rows); + var next: usize = 0; + for (t.map, 0..) |*m, r| { + if (pad(r)) { + m.* = -1; + } else { + m.* = order[next]; + next += 1; + } + } + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(if (t.map.len > 0) t.map else &[_]i32{0})); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.gamma)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(t.beta)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Norm) !void { + const p = t.rig.begin(); + const mode: ks.NormMode = switch (t.form) { + .rows => .rows, + .map => .map, + .merge => .merge, + }; + ks.norm(p.e, t.pipe, ref(t.ins[0]), if (t.form == .map) ref(t.ins[1]) else null, ref(t.ins[2]), ref(t.ins[3]), ref(t.out), t.shape.rows, t.shape.n, mode, merge_w, t.shape.n, 0, 1e-5); + try t.rig.finish(p); + } + + pub fn count(t: *Norm) usize { + return t.shape.rows * t.shape.n; + } + + /// Row r's values in the kernel's order; false for a map's padding. + fn row(t: *const Norm, r: usize, out: []f64) bool { + const n = t.shape.n; + switch (t.form) { + .rows => for (out, 0..) |*v, ch| { + v.* = t.x[r * n + ch]; + }, + .map => { + if (t.map[r] < 0) return false; + const src: usize = @intCast(t.map[r]); + for (out, 0..) |*v, ch| v.* = t.x[src * n + ch]; + }, + .merge => { + // [x0 | x1 | x2 | x3]: (2i, 2j), (2i + 1, 2j), (2i, 2j + 1), (2i + 1, 2j + 1) + const q = n / 4; + const half = merge_w / 2; + for (out, 0..) |*v, ch| { + const part = ch / q; + const sy = (r / half) * 2 + (part & 1); + const sx = (r % half) * 2 + (part >> 1); + v.* = t.x[(sy * merge_w + sx) * q + ch % q]; + } + }, + } + return true; + } + + /// The outputs in float64, each one's allowance from fp32's analysis of both passes and the scale and shift. + fn compute(t: *Norm, acc: ?struct { c.Accum, g.Order }, want: []f64, allow: ?[]f64) !void { + const n = t.shape.n; + const buf = try t.a.alloc(f64, 2 * n); + defer t.a.free(buf); + const v = buf[0..n]; + const sq = buf[n..]; + const nf: f64 = @floatFromInt(n); + const D: f64 = @floatFromInt(depth(t.shape)); + for (0..t.shape.rows) |r| { + if (!t.row(r, v)) { + @memset(want[r * n ..][0..n], 0); + if (allow) |al| @memset(al[r * n ..][0..n], 0); + continue; + } + var s: f64 = 0; + var s1: f64 = 0; + for (v) |x| { + s += x; + s1 += @abs(x); + } + if (acc) |ao| s = contract.rows.sumIn(ao[0], ao[1], v); + const mu = s / nf; + var s2: f64 = 0; + var sd: f64 = 0; + var most: f64 = 0; + for (v, sq) |x, *q| { + q.* = (x - mu) * (x - mu); + s2 += q.*; + sd += @abs(x - mu); + most = @max(most, @abs(x - mu)); + } + if (acc) |ao| s2 = contract.rows.sumIn(ao[0], ao[1], sq); + const variance = s2 / nf; + const rs = 1 / @sqrt(variance + 1e-5); + for (v, 0..) |x, ch| { + const i = r * n + ch; + const gm: f64 = t.gamma[ch]; + const y = (x - mu) * rs * gm + t.beta[ch]; + want[i] = y; + const al = allow orelse continue; + const e_mu = c.gamma(D, U) * s1 / nf + 2 * id.div_ulps * U * @abs(mu); + const e_d = e_mu + U * most; + const e_var = (2 * e_d * sd + nf * e_d * e_d) / nf + c.gamma(D + 1, U) * s2 / nf + (2 * id.div_ulps + 1) * U * (variance + 1e-5); + const rel_r = e_var / (2 * (variance + 1e-5)) * 1.01 + 2 * (id.sqrt_ulps + id.div_ulps) * U; + const d = @abs(x - mu); + const e = @abs(gm) * rs * ((e_mu + U * d) * (1 + rel_r) + d * (rel_r + 2 * U)) + 2 * U * @abs(y) + U * @abs(@as(f64, t.beta[ch])); + al[i] = if (std.math.isFinite(y)) 1.01 * e + 0x1p-149 else 0; + } + } + } + + pub fn reference(t: *Norm, want: []f64, mag: []f64) !void { + t.a.free(t.allow); + t.allow = try t.a.alloc(f64, t.count()); + try t.compute(null, want, t.allow); + for (mag, want) |*m, w| m.* = if (std.math.isFinite(w)) @abs(w) else 0; + } + pub fn bounds(t: *Norm, _: []const f64, _: []const f64, allow: []f64) !void { + @memcpy(allow, t.allow); + } + pub fn simulate(t: *Norm, acc: c.Accum, order: g.Order, out: []f64) !void { + try t.compute(.{ acc, order }, out, null); + } +}; + +// --- br_deform: torchvision's modulated deform_conv2d columns (stride 1, padding k / 2) ------------------------------ + +const grid_h = 16; +const grid_w = 12; +const first = 5; + +pub const Deform = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + side: usize, + shape: c.Shape = undefined, + x: []f32 = &.{}, + offs: []f32 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: the chunk's pixels (from pixel 5 of a 16 x 12 grid); n: its columns, k x k taps of 64; k: 64. + pub fn spec(name: []const u8, side: usize) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .out = .f32, .tiles = .{ .{ .size = 32 }, .{}, .{} }, .served = switch (side) { + 1 => &comptime served(1), + 3 => &comptime served(3), + else => &comptime served(7), + } }; + } + fn served(comptime side: usize) [2]c.Shape { + return .{ .{ .rows = grid_h * grid_w - first, .n = side * side * 64, .k = 64 }, .{ .rows = 64, .n = side * side * 64, .k = 64 } }; + } + + pub fn deinit(t: *Deform) void { + t.a.free(t.x); + t.a.free(t.offs); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Deform, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.side), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Deform); + + fn taps(t: *const Deform) usize { + return t.side * t.side; + } + fn ldo(t: *const Deform) usize { + return std.mem.alignForward(usize, 3 * t.taps(), 64); + } + + pub fn prepare(t: *Deform, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > grid_h * grid_w - first or shape.k != 64 or shape.n != t.taps() * 64) return false; + t.rig.reset(); + t.a.free(t.x); + t.a.free(t.offs); + t.shape = shape; + t.x = try t.a.alloc(f32, grid_h * grid_w * 64); + t.offs = try t.a.alloc(f32, grid_h * grid_w * t.ldo()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + uniform(rng, t.x, -3, 3); + // offsets to 2.5 pixels: taps fall outside, on the border and between pixels; whole-pixel ones too + uniform(rng, t.offs, -2.5, 2.5); + for (t.offs, 0..) |*v, i| { + if (i % t.ldo() >= 3 * t.taps()) v.* = 0 else if (i % 11 == 0) v.* = @round(v.*); + } + if (fill == .nonfinite) plant(t.x, 64); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.offs)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Deform) !void { + const p = t.rig.begin(); + ks.deform(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.out), .{ .h = grid_h, .w = grid_w }, t.side, first, t.shape.rows, t.ldo()); + try t.rig.finish(p); + } + + pub fn count(t: *Deform) usize { + return t.shape.rows * t.shape.n; + } + + fn px(t: *const Deform, y: i64, x: i64, ch: usize) f64 { + return t.x[(@as(usize, @intCast(y)) * grid_w + @as(usize, @intCast(x))) * 64 + ch]; + } + + /// The sample's position and weights as the kernel forms them in fp32; products, sums and the mask's ulps. + pub fn reference(t: *Deform, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const nt = t.taps(); + const pad: i64 = @intCast(t.side / 2); + for (0..t.shape.rows) |r| { + const pix = first + r; + const oy: i64 = @intCast(pix / grid_w); + const ox: i64 = @intCast(pix % grid_w); + const off = t.offs[pix * t.ldo() ..]; + for (0..nt) |tap| { + const m: f64 = off[2 * nt + tap]; + const mask = 2 / (1 + @exp(-m)); + const er = Range.of(@exp(-m)).fp32(id.exp_ulps).add(Range.of(1)).fp32(0.5); + const mr = er.recip().fp32(id.div_ulps).mul(Range.of(2)); + const h: f32 = @as(f32, @floatFromInt(oy - pad + @as(i64, @intCast(tap / t.side)))) + off[2 * tap]; + const w: f32 = @as(f32, @floatFromInt(ox - pad + @as(i64, @intCast(tap % t.side)))) + off[2 * tap + 1]; + const inside = h > -1 and h < grid_h and w > -1 and w < grid_w; + const hl: i64 = @intFromFloat(@floor(h)); + const wl: i64 = @intFromFloat(@floor(w)); + const lh: f32 = h - @as(f32, @floatFromInt(hl)); + const lw: f32 = w - @as(f32, @floatFromInt(wl)); + const hh: f32 = 1 - lh; + const hw: f32 = 1 - lw; + const wt = [4]f32{ hh * hw, hh * lw, lh * hw, lh * lw }; + const ok = [4]bool{ hl >= 0 and wl >= 0, hl >= 0 and wl + 1 <= grid_w - 1, hl + 1 <= grid_h - 1 and wl >= 0, hl + 1 <= grid_h - 1 and wl + 1 <= grid_w - 1 }; + const ys = [4]i64{ hl, hl, hl + 1, hl + 1 }; + const xs = [4]i64{ wl, wl + 1, wl, wl + 1 }; + for (0..64) |ch| { + const i = r * t.shape.n + tap * 64 + ch; + if (!inside) { + t.ends.put(want, mag, i, mask * 0, Range.of(0)); + continue; + } + var s: f64 = 0; + var sr = Range.of(0); + for (0..4) |q| { + const v: f64 = if (ok[q]) t.px(ys[q], xs[q], ch) else 0; + const term = @as(f64, wt[q]) * v; + s += term; + sr = if (q == 0) Range.of(term).fp32(0.5) else sr.add(Range.of(term).fp32(0.5)).fp32(0.5); + } + t.ends.put(want, mag, i, mask * s, sr.mul(mr).fp32(0.5)); + } + } + } + } + pub fn bounds(t: *Deform, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Deform, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; diff --git a/zig/tests/contract/metal_idol_birefnet_rows.zig b/zig/tests/contract/metal_idol_birefnet_rows.zig new file mode 100644 index 0000000000..bd6ab44b80 --- /dev/null +++ b/zig/tests/contract/metal_idol_birefnet_rows.zig @@ -0,0 +1,541 @@ +//! br_gate, br_head, br_pool, br_resize and br_tiles under the contract: BiRefNet's small steps against float64. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); +const cv = @import("metal_idol_birefnet_conv.zig"); + +const c = id.c; +const g = id.g; +const Range = id.Range; +const Ends = id.Ends; +const ks = cv.ks; +const ref = cv.ref; +const uniform = st.uniform; +const plant = st.plant; + +fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +// --- br_gate: x *= sigmoid(g . w + b), the gradient attention's gate ------------------------------------------------ + +pub const Gate = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + gv: []f32 = &.{}, + w: [16]f32 = undefined, + b: [1]f32 = undefined, + x: []f32 = &.{}, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + fn depth(_: c.Shape) usize { + return 16 + 1; + } + + /// Rows: pixels; n: channels gated; k: the 16 gate channels. The output is x gated, x given as an input copy. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = depth, .ops = 1 + id.exp_ulps + id.div_ulps + 2, .out = .f32, .tiles = .{ .{ .size = 4 }, .{ .size = 64 }, .{} }, .served = &.{ .{ .rows = 1024, .n = 1536, .k = 16 }, .{ .rows = 4096, .n = 768, .k = 16 }, .{ .rows = 999, .n = 384, .k = 16 } }, .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = "16 terms into a sigmoid: no input makes fp16 sums miss by 100x and keep the gate's fp32 value" } }; + } + + pub fn deinit(t: *Gate) void { + t.a.free(t.gv); + t.a.free(t.x); + t.rig.deinit(); + } + pub fn kernel(t: *Gate, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Gate); + + pub fn prepare(t: *Gate, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k != 16) return false; + t.rig.reset(); + t.a.free(t.gv); + t.a.free(t.x); + t.shape = shape; + t.gv = try t.a.alloc(f32, shape.rows * 16); + t.x = try t.a.alloc(f32, shape.rows * shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + uniform(rng, t.gv, 0, 2); + uniform(rng, &t.w, -1, 1); + uniform(rng, &t.b, -1, 1); + uniform(rng, t.x, -3, 3); + if (fill == .nonfinite) { + plant(t.x, shape.n); + plant(t.gv, 16); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.gv)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(&t.w)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(&t.b)); + // the kernel gates in place: the output starts as x + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Gate) !void { + @memcpy(t.rig.bytes(t.out, t.count() * 4), std.mem.sliceAsBytes(t.x)); + const p = t.rig.begin(); + ks.gate(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.out), t.shape.rows, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Gate) usize { + return t.shape.rows * t.shape.n; + } + + /// Gate 4's stress fill is the random one: the gate's 16 terms summed in `acc`, the rest in float64. + pub fn simulate(t: *Gate, acc: c.Accum, order: g.Order, out: []f64) !void { + var terms: [16]f64 = undefined; + for (0..t.shape.rows) |p| { + for (&terms, 0..) |*v, i| v.* = @as(f64, t.gv[p * 16 + i]) * t.w[i]; + const sg = sigmoid(contract.rows.sumIn(acc, order, &terms) + t.b[0]); + for (0..t.shape.n) |ch| out[p * t.shape.n + ch] = @as(f64, t.x[p * t.shape.n + ch]) * sg; + } + } + + pub fn reference(t: *Gate, want: []f64, mag: []f64) !void { + for (0..t.shape.rows) |p| { + var s: f64 = 0; + var m: f64 = 0; + for (0..16) |i| { + const v = @as(f64, t.gv[p * 16 + i]) * t.w[i]; + s += v; + m += @abs(v); + } + const sg = sigmoid(s + t.b[0]); + for (0..t.shape.n) |ch| { + const x: f64 = t.x[p * t.shape.n + ch]; + want[p * t.shape.n + ch] = x * sg; + const mm = @abs(x) * 0.25 * (m + @abs(@as(f64, t.b[0]))); + mag[p * t.shape.n + ch] = if (std.math.isFinite(mm)) mm else 0; + } + } + } +}; + +// --- br_head: conv_out1 over [the last level resized (align_corners) | ipt_blk1's channels] -------------------------- + +pub const Head = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + src: []f32 = &.{}, + q: []f32 = &.{}, + w: []f32 = &.{}, + b: [1]f32 = undefined, + ins: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + const cq = 48; + const in_side = 7; + + fn depth(s: c.Shape) usize { + return s.k + 4; + } + + /// Rows: the output side (rows x rows from 7 x 7); n: 1; k: the resized channels plus 48. + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = depth, .out = .f32, .tiles = .{ .{ .size = 16 }, .{}, .{} }, .served = &.{ .{ .rows = 29, .n = 1, .k = 240 }, .{ .rows = 64, .n = 1, .k = 240 } }, .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = "each term is a bilinear sample rounded in fp32, so no input keeps fp32's sums exact" } }; + } + + pub fn deinit(t: *Head) void { + inline for (.{ &t.src, &t.q, &t.w }) |p| t.a.free(p.*); + t.rig.deinit(); + } + pub fn kernel(t: *Head, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Head); + + fn cp(t: *const Head) usize { + return t.shape.k - cq; + } + + pub fn prepare(t: *Head, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows < 2 or shape.n != 1 or shape.k <= cq) return false; + t.rig.reset(); + inline for (.{ &t.src, &t.q, &t.w }) |p| t.a.free(p.*); + t.shape = shape; + t.src = try t.a.alloc(f32, in_side * in_side * t.cp()); + t.q = try t.a.alloc(f32, shape.rows * shape.rows * cq); + t.w = try t.a.alloc(f32, shape.k); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + uniform(rng, t.src, -2, 2); + uniform(rng, t.q, -2, 2); + uniform(rng, t.w, -0.3, 0.3); + uniform(rng, &t.b, -1, 1); + if (fill == .nonfinite) plant(t.q, cq); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.src)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.q)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(t.w)); + t.ins[3] = try t.rig.input(std.mem.sliceAsBytes(&t.b)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Head) !void { + const p = t.rig.begin(); + const s = t.shape.rows; + ks.head(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.ins[2]), ref(t.ins[3]), ref(t.out), .{ .h = in_side, .w = in_side }, .{ .h = s, .w = s }, t.cp(), cq); + try t.rig.finish(p); + } + + pub fn count(t: *Head) usize { + return t.shape.rows * t.shape.rows; + } + + /// Output (oy, ox)'s terms and magnitude; a sample counts its four corners whole, covering its position's rounding. + fn terms(t: *const Head, oy: usize, ox: usize, out: []f64) f64 { + const s = t.shape.rows; + const ax = tf.idol.preprocess.axis(in_side, s); + const y = ax.at(oy); + const x = ax.at(ox); + var m: f64 = 0; + for (0..t.cp()) |ch| { + const v = [4]f64{ t.src[(y.lo * in_side + x.lo) * t.cp() + ch], t.src[(y.lo * in_side + x.hi) * t.cp() + ch], t.src[(y.hi * in_side + x.lo) * t.cp() + ch], t.src[(y.hi * in_side + x.hi) * t.cp() + ch] }; + const top = @as(f64, y.l0) * (@as(f64, x.l0) * v[0] + @as(f64, x.l1) * v[1]); + const bottom = @as(f64, y.l1) * (@as(f64, x.l0) * v[2] + @as(f64, x.l1) * v[3]); + out[ch] = (top + bottom) * t.w[ch]; + m += (@abs(v[0]) + @abs(v[1]) + @abs(v[2]) + @abs(v[3])) * @abs(@as(f64, t.w[ch])); + } + const at = oy * s + ox; + for (0..cq) |ch| { + out[t.cp() + ch] = @as(f64, t.q[at * cq + ch]) * t.w[t.cp() + ch]; + m += @abs(out[t.cp() + ch]); + } + return m; + } + + pub fn reference(t: *Head, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, t.shape.k); + defer t.a.free(buf); + const s = t.shape.rows; + for (0..s) |oy| for (0..s) |ox| { + const m = t.terms(oy, ox, buf); + var sum: f64 = t.b[0]; + for (buf) |v| sum += v; + want[oy * s + ox] = sum; + const mm = m + @abs(@as(f64, t.b[0])); + mag[oy * s + ox] = if (std.math.isFinite(mm)) mm else 0; + }; + } + + pub fn simulate(t: *Head, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, t.shape.k); + defer t.a.free(buf); + const s = t.shape.rows; + for (0..s) |oy| for (0..s) |ox| { + _ = t.terms(oy, ox, buf); + out[oy * s + ox] = contract.rows.sumIn(acc, order, buf) + t.b[0]; + }; + } +}; + +// --- br_pool: the ASPP's global average pool, a channel's mean over the pixels -------------------------------------- + +pub const Pool = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// A thread's strided chain, the simdgroup's five levels and the eight simdgroups' chain. + fn depth(s: c.Shape) usize { + return (s.k + 255) / 256 + 5 + 8; + } + + /// Rows: channels; n: 1; k: pixels (32^2 .. 256^2 at 64 channels). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .f32, .depth = depth, .ops = id.div_ulps, .out = .f32, .tiles = .{ .{ .size = 8 }, .{}, .{ .size = 256 } }, .served = &.{ .{ .rows = 64, .n = 1, .k = 1024 }, .{ .rows = 64, .n = 1, .k = 4096 }, .{ .rows = 64, .n = 1, .k = 65536 } } }; + } + + pub fn deinit(t: *Pool) void { + t.a.free(t.x); + t.rig.deinit(); + } + pub fn kernel(t: *Pool, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Pool); + + pub fn prepare(t: *Pool, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0) return false; + t.rig.reset(); + t.a.free(t.x); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.k); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + if (fill == .stress) { + // integers to 255: every partial sum under 2^24 is exact in fp32, past fp16's range in a few hundred + for (t.x) |*v| v.* = @floatFromInt(rng.intRangeAtMost(i32, 0, 255)); + } else { + uniform(rng, t.x, -3, 3); + if (fill == .nonfinite) plant(t.x, shape.rows); + } + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Pool) !void { + const p = t.rig.begin(); + ks.pool(p.e, t.pipe, ref(t.ins[0]), ref(t.out), t.shape.k, t.shape.rows); + try t.rig.finish(p); + } + + pub fn count(t: *Pool) usize { + return t.shape.rows; + } + + fn column(t: *const Pool, ch: usize, out: []f64) void { + for (out, 0..) |*v, i| v.* = t.x[i * t.shape.rows + ch]; + } + + pub fn reference(t: *Pool, want: []f64, mag: []f64) !void { + const n: f64 = @floatFromInt(t.shape.k); + for (0..t.shape.rows) |ch| { + var s: f64 = 0; + var m: f64 = 0; + for (0..t.shape.k) |i| { + const v: f64 = t.x[i * t.shape.rows + ch]; + s += v; + m += @abs(v); + } + want[ch] = s / n; + mag[ch] = if (std.math.isFinite(m)) m / n else 0; + } + } + + pub fn simulate(t: *Pool, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try t.a.alloc(f64, t.shape.k); + defer t.a.free(buf); + for (0..t.shape.rows) |ch| { + t.column(ch, buf); + out[ch] = contract.rows.sumIn(acc, order, buf) / @as(f64, @floatFromInt(t.shape.k)); + } + } +}; + +// --- br_resize: F.interpolate(bilinear, align_corners=True) between columns of row-major grids --------------------- + +/// One axis's source position as the kernel's fp32 computes it (fast-math division): its value and possible ends. +const Pos = struct { at: f64, lo: f64, hi: f64 }; + +fn position(in: usize, out: usize, o: usize) Pos { + if (out <= 1) return .{ .at = 0, .lo = 0, .hi = 0 }; + const scale: f32 = @as(f32, @floatFromInt(in - 1)) / @as(f32, @floatFromInt(out - 1)); + const at: f64 = scale * @as(f32, @floatFromInt(o)); + const e = 8 * 0x1p-24 * at; + return .{ .at = at, .lo = at - e, .hi = at + e }; +} + +/// x's bilinear value of channel ch at (fy, fx), the kernel's clamps, exactly. +fn sampleAt(x: []const f32, ld: usize, off: usize, h: usize, w: usize, ch: usize, fy: f64, fx: f64) f64 { + const y0 = @min(@as(usize, @intFromFloat(@max(fy, 0))), h - 1); + const x0 = @min(@as(usize, @intFromFloat(@max(fx, 0))), w - 1); + const y1 = @min(y0 + 1, h - 1); + const x1 = @min(x0 + 1, w - 1); + const ly = std.math.clamp(fy - @as(f64, @floatFromInt(y0)), 0, 1); + const lx = std.math.clamp(fx - @as(f64, @floatFromInt(x0)), 0, 1); + const v = [4]f64{ x[(y0 * w + x0) * ld + off + ch], x[(y0 * w + x1) * ld + off + ch], x[(y1 * w + x0) * ld + off + ch], x[(y1 * w + x1) * ld + off + ch] }; + const top = (1 - lx) * v[0] + lx * v[1]; + const bottom = (1 - lx) * v[2] + lx * v[3]; + return (1 - ly) * top + ly * bottom; +} + +/// The positions to try across [lo, hi]: both ends and the grid line between (a piece's extremes lie at its corners). +fn tries(p: Pos) [3]f64 { + const mid = @floor(p.hi); + return .{ p.lo, if (mid > p.lo) mid else p.lo, p.hi }; +} + +pub const Resize = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + const xoff = 2; + + /// Rows: the output's side; n: channels; k: the input's side (the input wider than its channels, from column 2). + pub fn spec(name: []const u8) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .out = .f32, .tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{} }, .served = &.{ .{ .rows = 512, .n = 3, .k = 1024 }, .{ .rows = 64, .n = 384, .k = 32 }, .{ .rows = 32, .n = 384, .k = 256 }, .{ .rows = 64, .n = 256, .k = 1 }, .{ .rows = 33, .n = 96, .k = 7 } } }; + } + + pub fn deinit(t: *Resize) void { + t.a.free(t.x); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Resize, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Resize); + + fn ldx(t: *const Resize) usize { + return t.shape.n + 4; + } + + pub fn prepare(t: *Resize, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + t.rig.reset(); + t.a.free(t.x); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.k * shape.k * t.ldx()); + var prng = std.Random.DefaultPrng.init(seed); + uniform(prng.random(), t.x, -3, 3); + if (fill == .nonfinite) plant(t.x[xoff..], t.ldx()); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Resize) !void { + const p = t.rig.begin(); + const k = t.shape.k; + const s = t.shape.rows; + ks.resize(p.e, t.pipe, ref(t.ins[0]), ref(t.out), .{ .h = k, .w = k }, .{ .h = s, .w = s }, t.shape.n, t.ldx(), xoff, t.shape.n, 0); + try t.rig.finish(p); + } + + pub fn count(t: *Resize) usize { + return t.shape.rows * t.shape.rows * t.shape.n; + } + + /// The value at the nominal positions; its range over fp32's possible positions, widened by the eight roundings. + pub fn reference(t: *Resize, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const k = t.shape.k; + const s = t.shape.rows; + const n = t.shape.n; + for (0..s) |oy| for (0..s) |ox| { + const py = position(k, s, oy); + const px = position(k, s, ox); + for (0..n) |ch| { + const w = sampleAt(t.x, t.ldx(), xoff, k, k, ch, py.at, px.at); + var lo = w; + var hi = w; + var size: f64 = 0; + for (tries(py)) |fy| for (tries(px)) |fx| { + const v = sampleAt(t.x, t.ldx(), xoff, k, k, ch, fy, fx); + lo = @min(lo, v); + hi = @max(hi, v); + }; + const y0 = @min(@as(usize, @intFromFloat(py.at)), k - 1); + const x0 = @min(@as(usize, @intFromFloat(px.at)), k - 1); + for ([_]usize{ y0, @min(y0 + 1, k - 1) }) |yy| for ([_]usize{ x0, @min(x0 + 1, k - 1) }) |xx| { + size = @max(size, @abs(@as(f64, t.x[(yy * k + xx) * t.ldx() + xoff + ch]))); + }; + const i = (oy * s + ox) * n + ch; + t.ends.put(want, mag, i, w, (Range{ .lo = lo, .hi = hi }).widen(8 * 0x1p-24 * size + 0x1p-149)); + } + }; + } + pub fn bounds(t: *Resize, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Resize, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; + +// --- br_tiles: the patch embedding's 4 x 4 patches and image2patches' grid of tiles, copies -------------------------- + +pub const Tiles = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + grid: bool, + shape: c.Shape = undefined, + x: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: the output's side; n: its channels (cpad, zero past 3 a^2); k: the tile's side a; 3 source channels. + pub fn spec(name: []const u8, grid: bool) c.Spec { + return .{ .name = name, .accum = .none, .depth = id.noSums, .out = .f32, .tiles = .{ .{ .size = 8 }, .{}, .{} }, .served = if (grid) &.{ .{ .rows = 32, .n = 3072, .k = 32 }, .{ .rows = 64, .n = 768, .k = 16 }, .{ .rows = 128, .n = 192, .k = 8 }, .{ .rows = 256, .n = 64, .k = 4 } } else &.{ .{ .rows = 256, .n = 64, .k = 4 }, .{ .rows = 128, .n = 64, .k = 4 } } }; + } + + pub fn deinit(t: *Tiles) void { + t.a.free(t.x); + t.rig.deinit(); + } + pub fn kernel(t: *Tiles, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name, t.grid), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Tiles); + + fn side(t: *const Tiles) usize { + return t.shape.rows * t.shape.k; + } + + pub fn prepare(t: *Tiles, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0 or shape.n < 3 * shape.k * shape.k) return false; + t.rig.reset(); + t.a.free(t.x); + t.shape = shape; + t.x = try t.a.alloc(f32, t.side() * t.side() * 3); + var prng = std.Random.DefaultPrng.init(seed); + uniform(prng.random(), t.x, -3, 3); + if (fill == .nonfinite) plant(t.x, 3); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Tiles) !void { + const p = t.rig.begin(); + const s = t.shape.rows; + ks.tiles(p.e, t.pipe, ref(t.ins[0]), ref(t.out), .{ .h = s, .w = s }, t.side(), 3, t.shape.k, t.shape.k, t.grid, t.shape.n); + try t.rig.finish(p); + } + + pub fn count(t: *Tiles) usize { + return t.shape.rows * t.shape.rows * t.shape.n; + } + + /// Channel ch * a^2 + i * a + j of (oy, ox): a patch's pixel, or image2patches' tile (i, j) as einops places it. + pub fn reference(t: *Tiles, want: []f64, mag: []f64) !void { + const s = t.shape.rows; + const a = t.shape.k; + for (0..s * s) |px| for (0..t.shape.n) |col| { + const i = px * t.shape.n + col; + mag[i] = 0; + if (col >= 3 * a * a) { + want[i] = 0; + continue; + } + const ch = col / (a * a); + const ti = (col % (a * a)) / a; + const tj = col % a; + const oy = px / s; + const ox = px % s; + const sy = if (t.grid) ti * s + oy else oy * a + ti; + const sx = if (t.grid) tj * s + ox else ox * a + tj; + want[i] = t.x[(sy * t.side() + sx) * 3 + ch]; + }; + } +}; diff --git a/zig/tests/contract/metal_idol_birefnet_window.zig b/zig/tests/contract/metal_idol_birefnet_window.zig new file mode 100644 index 0000000000..78a7043f55 --- /dev/null +++ b/zig/tests/contract/metal_idol_birefnet_window.zig @@ -0,0 +1,186 @@ +//! br_window under the contract: Swin's window attention (bias, shifted windows' mask) in fp32 against float64. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const id = @import("metal_idol_common.zig"); +const st = @import("metal_idol_vae_steps.zig"); +const cv = @import("metal_idol_birefnet_conv.zig"); + +const c = id.c; +const g = id.g; +const ks = cv.ks; +const ref = cv.ref; +const uniform = st.uniform; +const plant = st.plant; +const U = 0x1p-24; +const N = 144; +const D = 32; + +pub const Window = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shift: usize, + shape: c.Shape = undefined, + qkv: []f32 = &.{}, + table: []f32 = &.{}, + allow: []f64 = &.{}, + ins: [2]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// The longest chain: P V's 144 products (Q K^T's are 32, the softmax's total 38). + fn depth(_: c.Shape) usize { + return N; + } + + /// Rows: windows a side (the grid 12 rows on a side); n: the width C (C / 32 heads); k: the 144 keys. + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 2 }, .{}, .{} }, + .served = &.{ .{ .rows = 3, .n = 192, .k = N }, .{ .rows = 2, .n = 384, .k = N }, .{ .rows = 1, .n = 768, .k = N }, .{ .rows = 1, .n = 1536, .k = N } }, + .sums = .{ .f16 = 0, .bf16 = 0, .exact = false, .why = "a softmax's terms: no input keeps fp32's sums exact through exp and the division" }, + }; + } + + pub fn deinit(t: *Window) void { + t.a.free(t.qkv); + t.a.free(t.table); + t.a.free(t.allow); + t.rig.deinit(); + } + pub fn kernel(t: *Window, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Window); + + fn heads(t: *const Window) usize { + return t.shape.n / D; + } + fn windows(t: *const Window) usize { + return t.shape.rows * t.shape.rows; + } + + pub fn prepare(t: *Window, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % D != 0 or shape.k != N) return false; + t.rig.reset(); + t.a.free(t.qkv); + t.a.free(t.table); + t.shape = shape; + t.qkv = try t.a.alloc(f32, t.windows() * N * 3 * shape.n); + t.table = try t.a.alloc(f32, 529 * t.heads()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + uniform(rng, t.qkv, -2, 2); + uniform(rng, t.table, -3, 3); + if (fill == .nonfinite) plant(t.qkv, 3 * shape.n); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(t.table)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Window) !void { + const p = t.rig.begin(); + const side = 12 * t.shape.rows; + ks.window(p.e, t.pipe, ref(t.ins[0]), ref(t.ins[1]), ref(t.out), t.shape.n, t.heads(), side, side, t.shift); + try t.rig.finish(p); + } + + pub fn count(t: *Window) usize { + return t.windows() * N * t.shape.n; + } + + /// Token i's region of window (wy, wx) in the shifted grid: rows and columns split at side - 12 and side - shift. + fn region(t: *const Window, i: usize, wy: usize, wx: usize) usize { + const side = 12 * t.shape.rows; + const y = wy * 12 + i / 12; + const x = wx * 12 + i % 12; + const ry: usize = if (y < side - 12) 0 else if (y < side - t.shift) 1 else 2; + const rx: usize = if (x < side - 12) 0 else if (x < side - t.shift) 1 else 2; + return ry * 3 + rx; + } + + /// Logits (fp32-scaled q) . k + bias (+ mask), softmax, P V; the allowance tracks each logit's error through both. + fn compute(t: *Window, acc: ?struct { c.Accum, g.Order }, want: []f64, allow: ?[]f64) !void { + const C = t.shape.n; + const ld = 3 * C; + const scale: f32 = @floatCast(std.math.pow(f64, D, -0.5)); + var s: [N]f64 = undefined; + var err: [N]f64 = undefined; + var p: [N]f64 = undefined; + var terms: [N]f64 = undefined; + for (0..t.windows()) |win| for (0..t.heads()) |h| for (0..N) |qi| { + const wy = win / t.shape.rows; + const wx = win % t.shape.rows; + const base = win * N; + const q = t.qkv[(base + qi) * ld + h * D ..][0..D]; + const mine = t.region(qi, wy, wx); + var most = -std.math.inf(f64); + for (0..N) |j| { + const k = t.qkv[(base + j) * ld + C + h * D ..][0..D]; + var dot: f64 = 0; + var size: f64 = 0; + for (0..D) |d| { + const qs: f32 = q[d] * scale; + terms[d] = @as(f64, qs) * k[d]; + dot += terms[d]; + size += @abs(terms[d]); + } + if (acc) |ao| dot = contract.rows.sumIn(ao[0], ao[1], terms[0..D]); + const at = ((qi / 12 + 11 - j / 12) * 23 + (qi % 12 + 11 - j % 12)) * t.heads() + h; + const bias: f64 = t.table[at]; + const masked: f64 = if (t.shift > 0 and t.region(j, wy, wx) != mine) -100 else 0; + s[j] = dot + bias + masked; + err[j] = c.gamma(D, U) * size + U * (@abs(dot + bias) + @abs(s[j])); + most = @max(most, s[j]); + } + var total: f64 = 0; + var worst: f64 = 0; + for (0..N) |j| { + p[j] = @exp(s[j] - most); + total += p[j]; + err[j] += U * @abs(s[j] - most) + 2 * id.exp_ulps * U; + worst = @max(worst, err[j]); + } + if (acc) |ao| total = contract.rows.sumIn(ao[0], ao[1], &p); + for (&p) |*v| v.* /= total; + for (0..D) |d| { + var o: f64 = 0; + var spread: f64 = 0; + var size: f64 = 0; + for (0..N) |j| { + const v: f64 = t.qkv[(base + j) * ld + 2 * C + h * D + d]; + terms[j] = p[j] * v; + o += terms[j]; + size += @abs(terms[j]); + spread += @abs(terms[j]) * (err[j] + worst); + } + if (acc) |ao| o = contract.rows.sumIn(ao[0], ao[1], &terms); + const i = (base + qi) * C + h * D + d; + want[i] = o; + const al = allow orelse continue; + const e = 1.01 * (spread + (c.gamma(38, U) + 2 * id.div_ulps * U + 2 * U) * size) + c.gamma(N, U) * size + U * @abs(o); + al[i] = if (std.math.isFinite(o)) e + 0x1p-149 else 0; + } + }; + } + + pub fn reference(t: *Window, want: []f64, mag: []f64) !void { + t.a.free(t.allow); + t.allow = try t.a.alloc(f64, t.count()); + try t.compute(null, want, t.allow); + for (mag, want) |*m, w| m.* = if (std.math.isFinite(w)) @abs(w) else 0; + } + pub fn bounds(t: *Window, _: []const f64, _: []const f64, allow: []f64) !void { + @memcpy(allow, t.allow); + } + pub fn simulate(t: *Window, acc: c.Accum, order: g.Order, out: []f64) !void { + try t.compute(.{ acc, order }, out, null); + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 15752aadf3..58c4711319 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -105,6 +105,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "idol decoder ")) { text = sources.idol_decoder; frags = true; + } else if (std.mem.startsWith(u8, d.name, "idol birefnet ")) { + text = sources.idol_birefnet; } else if (std.mem.startsWith(u8, d.name, "idol vae ")) { text = sources.idol_vae; } else if (std.mem.startsWith(u8, d.name, "idol dino ")) { diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 4fccee9bcb..16d384888e 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -435,6 +435,8 @@ pub fn main(init: std.process.Init) !u8 { defer idol_vae.destroy(); const idol_dino = try @import("metal_idol_dino.zig").Suite.create(a, device, queue, &kernels, o.only); defer idol_dino.destroy(); + const idol_birefnet = try @import("metal_idol_birefnet.zig").Suite.create(a, device, queue, &kernels, o.only); + defer idol_birefnet.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { diff --git a/zig/tests/idol_birefnet.zig b/zig/tests/idol_birefnet.zig new file mode 100644 index 0000000000..7580dfd4c6 --- /dev/null +++ b/zig/tests/idol_birefnet.zig @@ -0,0 +1,268 @@ +//! tf-idol-birefnet: background removal and preprocessing on an idol_birefnet_case.py case, gated and timed. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const bn = tf.idol.birefnet; +const dec = tf.idol.birefnet_decoder; +const wts = tf.idol.birefnet_weights; +const pre = tf.idol.preprocess; +const gt = tf.idol.birefnet_reference; + +const usage = + \\tf-idol-birefnet SNAPSHOT_DIR CASE_DIR PHOTO + \\SNAPSHOT_DIR holds birefnet.safetensors (speridlabs/idol's snapshot). CASE_DIR is idol_birefnet_case.py's for PHOTO + \\(birefnet.safetensors and .json, and birefnet_stages.safetensors with --stages). Gates, each printed: + \\- the resized photo and the BiRefNet input: byte for byte the SDK's; + \\- the mask logits from our own input against the float64 ruler: worst image row and 1 - cosine within twice + \\ torch's own; with the stages, each stage on torch's own input within twice torch's own as well; + \\- the alpha step on the SDK's own logits: byte for byte the SDK's alpha; + \\- the uint8 alpha: differing pixels against the ruler's within twice torch's, largest difference within torch's; + \\- the picture from the SDK's alpha and from the ruler's: byte for byte the SDK's (Pillow's steps exactly); from our + \\ own alpha against the ruler's picture: differing pixels and largest difference within twice the SDK's; + \\- three more forwards byte for byte the first. CASE_DIR gets tf-idol-birefnet.json. +; + +const timed = 3; + +fn take(comptime T: type, a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8) ![]T { + const t = f.get(name) orelse { + std.debug.print("the case has no {s}\n", .{name}); + return error.MissingTensor; + }; + const want: tf.safetensors.DType = if (T == f32) .f32 else .u8; + if (t.dtype != want) return error.BadCaseTensor; + const out = try a.alloc(T, t.numel()); + @memcpy(std.mem.sliceAsBytes(out), t.bytes); + return out; +} + +fn print(list: *std.ArrayList(u8), a: std.mem.Allocator, comptime fmt: []const u8, args: anytype) !void { + try list.appendSlice(a, try std.fmt.allocPrint(a, fmt, args)); +} + +fn verdict(list: *std.ArrayList(u8), a: std.mem.Allocator, name: []const u8, v: gt.Verdict) !void { + std.debug.print("{s}: ours worst row {e:.3}, 1 - cos {e:.3}; torch {e:.3}, {e:.3}: {s}\n", .{ name, v.ours.worst_row, 1 - v.ours.cosine, v.torch.worst_row, 1 - v.torch.cosine, if (v.pass) "pass" else "FAIL" }); + try print(list, a, ", \"{s}\": {{\"ours_worst_row\": {e:.4}, \"ours_one_minus_cos\": {e:.4}, \"torch_worst_row\": {e:.4}, \"torch_one_minus_cos\": {e:.4}, \"pass\": {}}}", .{ name, v.ours.worst_row, 1 - v.ours.cosine, v.torch.worst_row, 1 - v.torch.cosine, v.pass }); +} + +fn diffLine(list: *std.ArrayList(u8), a: std.mem.Allocator, name: []const u8, d: gt.Diff) !void { + std.debug.print("{s}: {d} pixels differ, by at most {d}\n", .{ name, d.pixels, d.max }); + try print(list, a, ", \"{s}\": {{\"pixels\": {d}, \"max\": {d}}}", .{ name, d.pixels, d.max }); +} + +/// Rows [n, cols] of `src` (row length ld, from column off). +fn columns(a: std.mem.Allocator, src: []const f32, n: usize, ld: usize, off: usize, cols: usize) ![]f32 { + const out = try a.alloc(f32, n * cols); + for (0..n) |r| @memcpy(out[r * cols ..][0..cols], src[r * ld + off ..][0..cols]); + return out; +} + +fn put(dst: []f32, src: []const f32, ld: usize, off: usize, cols: usize) void { + for (0..src.len / cols) |r| @memcpy(dst[r * ld + off ..][0..cols], src[r * cols ..][0..cols]); +} + +const Stages = struct { + a: std.mem.Allocator, + f: *const tf.safetensors.File, + n: *bn.Net, + out: *std.ArrayList(u8), + pass: bool = true, + + fn check(s: *Stages, name: []const u8, ours: []const f32, width: usize) !void { + var buf: [64]u8 = undefined; + const torch = try take(f32, s.a, s.f, name); + const ruler = try take(f32, s.a, s.f, try std.fmt.bufPrint(&buf, "ruler.{s}", .{name})); + const v = gt.gate(ours, torch, ruler, width); + s.pass = s.pass and v.pass; + try verdict(s.out, s.a, name, v); + } + + fn run(s: *Stages, comptime f: anytype, args: anytype) !void { + const cb = s.n.queue.commandBuffer(); + const e = cb.compute(.serial); + @call(.auto, f, .{ s.n, e } ++ args); + _ = try bn.Net.finish(cb, e); + } + + fn layer(n: *bn.Net, e: mtl.ComputeEncoder, pass: usize, i: usize) void { + n.layer(e, pass, i, n.ta, n.tb); + } + + fn all(s: *Stages) !void { + const n = s.n; + const F = bn.Net.floats; + var buf: [64]u8 = undefined; + for (0..2) |pass| for (0..4) |i| { + const h = bn.grids[pass] >> @intCast(i); + const C = wts.widths[i]; + const in = try take(f32, s.a, s.f, try std.fmt.bufPrint(&buf, "layer.{d}.{d}.in", .{ i, pass })); + @memcpy(F(n.ta, in.len), in); + try s.run(layer, .{ pass, i }); + try s.check(try std.fmt.bufPrint(&buf, "layer.{d}.{d}.out", .{ i, pass }), F(n.ta, h * h * C), C); + if (i < 3) try s.check(try std.fmt.bufPrint(&buf, "layer.{d}.{d}.down", .{ i, pass }), F(n.tb, h * h * C / 2), 2 * C); + }; + @memcpy(F(n.x, bn.size * bn.size * 3), try take(f32, s.a, s.f, "enc.in")); + try s.run(bn.Net.encode, .{}); + const sides = [4]usize{ 256, 128, 64, 32 }; + const widths = [4]usize{ 384, 768, 1536, wts.squeeze_in }; + for (0..4) |i| try s.check(try std.fmt.bufPrint(&buf, "enc.out.x{d}", .{i + 1}), F(n.xs[i], sides[i] * sides[i] * widths[i]), widths[i]); + @memcpy(F(n.xs[3], 1024 * wts.squeeze_in), try take(f32, s.a, s.f, "squeeze.in")); + try s.run(dec.squeeze, .{}); + try s.check("squeeze.out", try columns(s.a, F(n.lin[0], 1024 * dec.level_in[0]), 1024, dec.level_in[0], 0, wts.channels[0]), wts.channels[0]); + for (0..4) |l| { + const side = dec.level_side[l]; + const name = try std.fmt.bufPrint(&buf, "dec{d}.in", .{4 - l}); + @memcpy(F(n.lin[l], side * side * dec.level_in[l]), try take(f32, s.a, s.f, name)); + try s.run(dec.block, .{ &n.w.dec[l], n.lin[l], dec.level_in[l], side, n.out[l], dec.level_out[l], @as(usize, 0) }); + try s.check(try std.fmt.bufPrint(&buf, "dec{d}.out", .{4 - l}), F(n.out[l], side * side * dec.level_out[l]), dec.level_out[l]); + } + @memcpy(F(n.x, bn.size * bn.size * 3), try take(f32, s.a, s.f, "decoder.in.x")); + for (0..3) |i| { + const t = try take(f32, s.a, s.f, try std.fmt.bufPrint(&buf, "decoder.in.x{d}", .{i + 1})); + @memcpy(F(n.xs[i], t.len), t); + } + put(F(n.lin[0], 1024 * dec.level_in[0]), try take(f32, s.a, s.f, "decoder.in.x4"), dec.level_in[0], 0, wts.channels[0]); + try s.run(dec.decode, .{}); + try s.check("decoder.out", F(n.logits, bn.size * bn.size), bn.size); + } +}; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 4) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + const cwd = std.Io.Dir.cwd(); + var case = try tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/birefnet.safetensors", .{args[2]})); + defer case.close(io); + const meta_text = try cwd.readFileAlloc(io, try std.fmt.allocPrint(a, "{s}/birefnet.json", .{args[2]}), a, .limited(1 << 22)); + const meta = try std.json.parseFromSliceLeaky(struct { crop_margin: f64, has_real_alpha: bool }, a, meta_text, .{ .ignore_unknown_fields = true }); + var out: std.ArrayList(u8) = .empty; + var pass = true; + + // the host's steps: decode and the first resize, against the SDK's resized photo + const started = mtl.clock.seconds(); + const bytes = try cwd.readFileAlloc(io, args[3], a, .limited(1 << 28)); + var photo = try pre.decode(io, a, bytes); + var resized = try pre.resizeFirst(a, photo); + const resize_s = mtl.clock.seconds() - started; + std.debug.print("photo {d} x {d}, {d} channels; resized to {d} x {d} in {d:.3} s\n", .{ photo.width, photo.height, photo.channels, resized.width, resized.height, resize_s }); + try print(&out, a, "{{\"photo\": [{d}, {d}, {d}], \"resized\": [{d}, {d}], \"resize_s\": {d:.4}, \"has_real_alpha\": {}", .{ photo.width, photo.height, photo.channels, resized.width, resized.height, resize_s, pre.hasRealAlpha(resized) }); + if (pre.hasRealAlpha(resized) != meta.has_real_alpha) { + std.debug.print("the photo's real alpha differs from the SDK's\n", .{}); + return 1; + } + const picture = try a.alloc(u8, pre.size * pre.size * 3); + const sdk_picture = try take(u8, a, &case, "picture"); + if (meta.has_real_alpha) { + const finishing = mtl.clock.seconds(); + try pre.finish(a, resized, null, meta.crop_margin, picture); + const d = gt.diff(picture, sdk_picture, 3); + try diffLine(&out, a, "picture_vs_sdk", d); + std.debug.print("preprocess (no BiRefNet: real alpha) {d:.3} s\n", .{resize_s + mtl.clock.seconds() - finishing}); + try print(&out, a, "}}\n", .{}); + try cwd.writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-birefnet.json", .{args[2]}), .data = out.items }); + return if (d.pixels == 0) 0 else 1; + } + const sdk_resized = try take(u8, a, &case, "resized"); + const rd = gt.diff(resized.pixels, sdk_resized, 3); + try diffLine(&out, a, "resized_vs_sdk", rd); + pass = pass and rd.pixels == 0; + const input = try a.alloc(f32, bn.size * bn.size * 3); + try pre.input(a, resized, input); + const sdk_input = try take(f32, a, &case, "input"); + var in_most: f32 = 0; + var in_count: usize = 0; + for (input, sdk_input) |x, y| { + in_most = @max(in_most, @abs(x - y)); + if (x != y) in_count += 1; + } + std.debug.print("input against the SDK's: {d} values differ, by at most {e:.3}\n", .{ in_count, in_most }); + try print(&out, a, ", \"input_vs_sdk\": {{\"values\": {d}, \"max_abs\": {e:.4}}}", .{ in_count, in_most }); + pass = pass and in_count == 0; + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const net = try bn.Net.load(std.heap.page_allocator, device, try std.fmt.allocPrintSentinel(a, "{s}/birefnet.safetensors", .{args[1]}, 0)); + defer net.deinit(); + const load_s = mtl.clock.seconds() - loading; + std.debug.print("{s}: BiRefNet loaded in {d:.2} s\n", .{ device.name(), load_s }); + try print(&out, a, ", \"device\": \"{s}\", \"load_s\": {d:.3}", .{ device.name(), load_s }); + + var stage_file: ?tf.safetensors.File = tf.safetensors.File.open(a, io, try std.fmt.allocPrint(a, "{s}/birefnet_stages.safetensors", .{args[2]})) catch null; + defer if (stage_file) |*f| f.close(io); + if (stage_file) |*f| { + var s: Stages = .{ .a = a, .f = f, .n = net, .out = &out }; + try s.all(); + pass = pass and s.pass; + } + + const n = bn.size * bn.size; + const logits = try a.alloc(f32, n); + const alpha = try a.alloc(u8, resized.width * resized.height); + const first = try net.removeBackground(a, resized, logits, alpha); + const lv = gt.gate(logits, try take(f32, a, &case, "logits"), try take(f32, a, &case, "logits_ruler"), bn.size); + try verdict(&out, a, "logits", lv); + pass = pass and lv.pass; + const sdk_alpha = try take(u8, a, &case, "alpha"); + const ruler_alpha = try take(u8, a, &case, "alpha_ruler"); + // the alpha step alone: the SDK's own logits through ours must give the SDK's alpha + const host_alpha = try a.alloc(u8, alpha.len); + try pre.alpha(a, try take(f32, a, &case, "logits"), resized.width, resized.height, host_alpha); + const hv = gt.diff(host_alpha, sdk_alpha, 1); + try diffLine(&out, a, "alpha_from_sdk_logits_vs_sdk", hv); + pass = pass and hv.pixels == 0; + const av: [3]gt.Diff = .{ gt.diff(alpha, sdk_alpha, 1), gt.diff(alpha, ruler_alpha, 1), gt.diff(sdk_alpha, ruler_alpha, 1) }; + try diffLine(&out, a, "alpha_vs_sdk", av[0]); + try diffLine(&out, a, "alpha_vs_ruler", av[1]); + try diffLine(&out, a, "sdk_alpha_vs_ruler", av[2]); + const alpha_pass = gt.withinDiff(av[1], av[2]); + std.debug.print("alpha gate: {s}\n", .{if (alpha_pass) "pass" else "FAIL"}); + pass = pass and alpha_pass; + + // the steps after the alpha, on the SDK's alpha and the ruler's (Pillow's exactly), then on ours + const ruler_picture = try take(u8, a, &case, "picture_ruler"); + try pre.finish(a, resized, sdk_alpha, meta.crop_margin, picture); + const p_sdk = gt.diff(picture, sdk_picture, 3); + try pre.finish(a, resized, ruler_alpha, meta.crop_margin, picture); + const p_ruler = gt.diff(picture, ruler_picture, 3); + try diffLine(&out, a, "picture_from_sdk_alpha_vs_sdk", p_sdk); + try diffLine(&out, a, "picture_from_ruler_alpha_vs_ruler", p_ruler); + pass = pass and p_sdk.pixels == 0 and p_ruler.pixels == 0; + const finishing = mtl.clock.seconds(); + try pre.finish(a, resized, alpha, meta.crop_margin, picture); + const finish_s = mtl.clock.seconds() - finishing; + const pv: [3]gt.Diff = .{ gt.diff(picture, sdk_picture, 3), gt.diff(picture, ruler_picture, 3), gt.diff(sdk_picture, ruler_picture, 3) }; + try diffLine(&out, a, "picture_vs_sdk", pv[0]); + try diffLine(&out, a, "picture_vs_ruler", pv[1]); + try diffLine(&out, a, "sdk_picture_vs_ruler", pv[2]); + const pic_pass = gt.withinPicture(pv[1], pv[2]); + std.debug.print("picture gate: {s}\n", .{if (pic_pass) "pass" else "FAIL"}); + pass = pass and pic_pass; + + const again = try a.alloc(f32, n); + var repeat = true; + var times: [timed]f64 = undefined; + var gpus: [timed]f64 = undefined; + for (×, &gpus) |*t, *g| { + const tm = try net.forward(input, again); + t.* = tm.wall; + g.* = tm.gpu; + repeat = repeat and std.mem.eql(u8, std.mem.sliceAsBytes(logits), std.mem.sliceAsBytes(again)); + } + std.mem.sort(f64, ×, {}, std.sort.asc(f64)); + std.mem.sort(f64, &gpus, {}, std.sort.asc(f64)); + pass = pass and repeat; + std.debug.print("{d} more forwards: {s}; first {d:.3} s, median {d:.4} s ({d:.4} s on the GPU); host finish {d:.3} s\n", .{ timed, if (repeat) "identical bytes" else "DIFFERENT bytes", first.wall, times[timed / 2], gpus[timed / 2], finish_s }); + try print(&out, a, ", \"alpha_pass\": {}, \"picture_pass\": {}, \"repeat_identical\": {}, \"first_s\": {d:.4}, \"median_s\": {d:.4}, \"median_gpu_s\": {d:.4}, \"finish_s\": {d:.4}, \"pass\": {}}}\n", .{ alpha_pass, pic_pass, repeat, first.wall, times[timed / 2], gpus[timed / 2], finish_s, pass }); + try cwd.writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-birefnet.json", .{args[2]}), .data = out.items }); + photo.deinit(a); + resized.deinit(a); + std.debug.print("{s}\n", .{if (pass) "PASS" else "FAIL"}); + return if (pass) 0 else 1; +} From b0642cb1d659e2f9b2943aedaf2f9b216eee4733 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 06:10:17 +0800 Subject: [PATCH 20/21] idol: tf-idol-generate takes a raw photo PICTURE may be a .png, .jpg, .jpeg or .webp photo, preprocessed natively with BiRefNet; --mode synthetic|real and --crop-margin X; a photo writes OUT_DIR/preprocessed.png and adds preprocess_s and birefnet_s to the JSON. Recipe section and CHANGELOG line. --- CHANGELOG.md | 1 + docs/recipes/idol.md | 228 +++++++++++++++++++++++++++++++++++- zig/tests/idol_generate.zig | 51 ++++++-- 3 files changed, 268 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b15466c84..2551fd0a81 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (#1774, after #1766, for #1757). Thanks to @jayleaton. - Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (#1775, after #1766, for #1757). Thanks to @jayleaton. - Idol runs end to end natively on M5 Metal: `tf-idol-generate` takes one preprocessed picture through DINOv3, the FLUX.2 VAE, the denoiser's 50 guided steps and both decoders to a 3DGS PLY and camera. On an M5 Ultra it is 29.1 s against 95.4 s for the SDK's IdolPipeline on MPS (excluding its preprocess), the denoiser 27.5 s against 93.9 s, at a 2.71 GB peak footprint against 12.88 GB. Given the SDK's starting and VAE noise the PLY's points sit as close to the SDK's as sampling resolves, and the camera's angles within 0.006 rad (#1780, after #1773, #1774 and #1775, for #1757). Thanks to @jayleaton. +- Idol's background removal and preprocessing run natively too: `tf-idol-generate` now takes a raw photo (`.png`, `.jpg`, `.jpeg` or `.webp`), with `--mode synthetic|real` and `--crop-margin X`, writes the picture the encoders saw as `preprocessed.png`, and adds `preprocess_s` and `birefnet_s` to its JSON. BiRefNet (MIT, from the Idol snapshot) runs in fp32 on Metal in 0.172 s against the SDK's 0.48 to 0.52 s preprocess, its logits within torch's own distance from a float64 ruler; Pillow's resizes, the erosion, crop and paste are byte for byte the SDK's, and the preprocessed picture differs from the SDK's in 2 pixels by 1. From a photo a generation takes 30.1 s against the SDK's 101.3 s at a 2.77 GB peak footprint against 12.88 GB (after #1780, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too. diff --git a/docs/recipes/idol.md b/docs/recipes/idol.md index 5aa0542761..961c3ce2ff 100644 --- a/docs/recipes/idol.md +++ b/docs/recipes/idol.md @@ -17,7 +17,9 @@ The sampler's defaults: 50 Euler steps, `t = 3u / (1 + 2u)`, guidance `v = 4 v_c with the condition, and with zeros). Both image encoders run natively too: the FLUX.2 VAE's ([below](#vae-encoder)) and DINOv3's -([below](#dinov3-encoder)). `tools/zig/idol_case.py` runs the publisher's SDK in torch, a development reference. +([below](#dinov3-encoder)), and so does the SDK's preprocessing with BiRefNet's background removal +([below](#background-removal-and-preprocessing)). `tools/zig/idol_case.py` runs the publisher's SDK in torch, a +development reference. It is a denoiser, not a token family: it does not register with the lane engine, the server or `families.detect`. It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them. @@ -477,6 +479,222 @@ The avocado example, the same machine and torch as above (`ATTN_BACKEND=sdpa`): - **Memory**: 2.28 GB peak process footprint (`/usr/bin/time -l`, `tf-idol-dino` with the case's feature rows), the bf16 weights 1.68 GB of it. The SDK's MPS encode holds 2.24 GB of MPS driver memory. +## Background removal and preprocessing + +`tf-idol-generate` takes a raw photo too. `preprocess.zig`, `photo.zig` and `birefnet*.zig` (kernels in +`idol_birefnet.metal`) do what idol-sdk's `IdolPipeline.preprocess` (`preprocess.py`) does: + +1. Resize the photo so its short side is 1,024 (Pillow LANCZOS). +2. Unless the photo has real alpha (RGBA with some alpha below 255), predict an alpha with BiRefNet + (`remove_background`). The input is `x / 255`, resized to 1024 x 1024 with torch's CPU bilinear (align_corners) and + normalized by ImageNet's mean and std. The mask's logits pass through a sigmoid, are resized back (bilinear, + align_corners), clamped and scaled by 255, then truncated to uint8. +3. Erode the alpha (`MinFilter(3)`). +4. Crop a square of `crop_margin` times the object's box: 1.2 for `synthetic`, 1.7 for `real`. +5. Resize it to 1,024 (LANCZOS, through Pillow's premultiplied RGBa) and paste it on black. + +The SDK applies no EXIF orientation, and neither does the port. + +The network is BiRefNet's Swin-L (embedding 192; depths 2, 2, 18, 2; windows of 12, shifted by 6), run twice: on the +1,024 input and on a 512 copy. Its features are concatenated, squeezed, and decoded by four ASPP-deformable blocks +with the image-patch convolutions, the gradient gates, the laterals and the head. + +### Weights, licence and identity + +`birefnet.safetensors` comes in the `speridlabs/idol` snapshot above. It is byte for byte `ZhengPeng7/BiRefNet` at +`e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4` (MIT), stored in fp16: all 742 tensors are equal, and 12 training-only +tensors are dropped. The loader requires exactly those 742 tensors at Swin-L's geometry. Every +`relative_position_index` must equal `(dy + 11) * 23 + dx + 11`, the formula `br_window` uses. It widens fp16 to fp32 +exactly and folds each batch norm as `(x - mean) * (weight / sqrt(var + 1e-5)) + bias`, the scale computed in float64. +The batch norms' `num_batches_tracked` counters are read and ignored. + +### Precision + +On a Mac, the SDK runs BiRefNet in fp32 on MPS; it calls `.half()` only on CUDA. A probe of MPS's fp32 matmul and conv +found true fp32: relative error about 7e-7, the CPU's. torchvision's `deform_conv2d` runs on MPS. + +The port therefore computes everything in fp32: + +- every convolution and linear layer as an fp32 simdgroup-matrix implicit GEMM, never the tensor units, whose fp32 + GEMM truncates its operands to 10 mantissa bits (#1772); +- sums in blocks of 256 terms added into the total, because one running sum over the squeeze's 51,840 terms drifts + to twice torch's distance from float64; +- LayerNorm in two passes; +- the exact `erf` GELU (musl's `erff`, within 0.76 ulp); +- `precise::exp`; +- the deformable columns as torchvision forms them. + +The host's steps are exact: + +- PNG and WebP decoding as Pillow 13.0.0 decodes them; +- LANCZOS with Pillow's double-precision coefficients and 22-bit fixed point; +- the RGBa premultiply and its inverse; +- `MinFilter`; +- the paste's `MULDIV255`; +- Python's `round` and `int`; +- torch's CPU bilinear as it computes it, `fma(t0, w0, t1 * w1)` along each axis; +- the alpha's fp32 sigmoid, bilinear and truncation as torch computes them on MPS. + +JPEG photos need the pinned libjpeg-turbo 3.1.4.1 and are refused without it. + +### Gates + +```sh +python3 tools/zig/idol_birefnet_case.py WEIGHTS SDK photo.png case/ [--mode real] [--stages] +zig build tf-idol-birefnet -Dcpu=apple_m1 +zig-out/bin/tf-idol-birefnet WEIGHTS case/ photo.png +``` + +`idol_birefnet_case.py` imports the SDK's `birefnet.py`, `preprocess.py` and `pipeline.py` only at their pinned +sha256. It also pins the weights' sha256 and Pillow 13.0.0. It runs `IdolPipeline.preprocess` on MPS and records: + +- the resized photo, BiRefNet's input, logits and alpha; +- the picture; +- a ruler: the same network in float64 on the CPU from the same resized photo. + +With `--stages`, it also records each stage's torch input and output and a float64 run of that stage on torch's own +input. The stages are each Swin layer of both passes, the encoder, the squeeze, each decoder block and the decoder. + +`tf-idol-birefnet` gates: + +- **Host steps:** the resized photo and BiRefNet's input are byte for byte the SDK's. The alpha step on the SDK's + own logits is byte for byte the SDK's alpha. +- **Logits** from our input: worst image row and `1 - cos` within twice torch's own distance from the ruler. The + cosine comes from the unit vectors' distance, since a dot product cannot resolve 1e-12. +- **Stages**, each on torch's own input: the same 2x gate. +- **uint8 alpha** against the ruler's: differing pixels within twice torch's, largest difference within torch's (or + one level). +- **Picture:** from the SDK's alpha and from the ruler's, byte for byte the SDK's picture, so Pillow's steps match + exactly. From our own alpha, differing pixels and largest difference within twice the SDK's. One alpha level moves + a pasted, resized pixel by up to two; the SDK's own picture sits two levels off the ruler's on the opaque photo. +- **Repeat:** three more forwards byte for byte. + +`birefnet_reference.zig`'s tests plant a 2% error in one logits row and mismatches in the alpha and the picture, and +require each gate to fail them. `preprocess.zig`'s and `photo.zig`'s tests cover the resampler, the PNG writer and +photo detection. + +### Measured + +The same M5 Ultra and torch 2.13. The photos: + +- the avocado reference input (1024 x 1024 PNG); +- the demo's `food-avocado.webp` (768 x 768); +- an opaque 1280 x 960 PNG of the avocado on a gradient, run in both modes; +- an RGBA PNG with real alpha. + +All pass. On every photo, the resized photo and BiRefNet's input are byte for byte the SDK's, and so is the alpha from +the SDK's logits. The picture from the SDK's alpha is byte for byte the SDK's. + +| photo | logits worst row, ours / torch | `1 - cos`, ours / torch | alpha vs ruler, ours / torch | picture vs SDK | +| --- | --- | --- | --- | --- | +| avocado input | 3.84e-6 / 4.15e-6 | 3.9e-13 / 4.0e-13 | 13 px max 1 / 11 px max 1 | 3 px, max 2 | +| food-avocado.webp | 8.15e-6 / 1.14e-5 | 5.8e-13 / 6.6e-13 | 5 px max 1 / 6 px max 1 | 2 px, max 1 | +| opaque, synthetic | 1.90e-5 / 1.65e-5 | 2.7e-12 / 2.2e-12 | 6 px max 1 / 5 px max 1 | 2 px, max 1 | +| opaque, real | 1.90e-5 / 1.65e-5 | 2.7e-12 / 2.2e-12 | 6 px max 1 / 5 px max 1 | 0 px | +| RGBA, real alpha | no BiRefNet | | | 0 px | + +- **Stages** (avocado input): all 24 pass. Our worst row is torch's times 0.2 to 1.13: layer 0's PatchMerging at + 1.13, the squeeze at 0.22. +- **Repeat:** identical bytes. +- **Speed:** BiRefNet's forward takes 0.172 s (median, 0.172 s on the GPU), the first 0.18 to 0.22 s. The host's + steps take 0.02 to 0.03 s for the resize and 0.02 s for the finish. The SDK's whole `preprocess` takes 0.26 to 0.28 s + warm and 0.48 to 0.52 s in a generation. Load: 0.7 to 0.9 s. +- **Memory:** `tf-idol-birefnet` peaks at a 2.33 GB footprint with its case; the fp32 weights are 0.89 GB of it. + +### In `tf-idol-generate` + +New to its command line, all else unchanged: + +- PICTURE may be a `.png`, `.jpg`, `.jpeg` or `.webp` photo (any case), preprocessed as above; +- `--mode synthetic|real` (default `synthetic`): the SDK's mode, crop margin 1.2 or 1.7; +- `--crop-margin X` overrides the mode's margin; +- a photo also writes `OUT_DIR/preprocessed.png`, the 1024 x 1024 RGB the encoders saw; +- a photo's `tf-idol-generate.json` adds `preprocess_s` (all of preprocessing but BiRefNet's load, which `load_s` + counts) and `birefnet_s` (BiRefNet's forward and alpha). + +BiRefNet loads from `SNAPSHOT_DIR/birefnet.safetensors` and is freed before DINOv3 loads. + +```sh +zig-out/bin/tf-idol-generate WEIGHTS KLEIN/vae DINO photo.png out/ [--mode real] [--crop-margin 1.5] +``` + +End to end against the SDK's `IdolPipeline.__call__` on MPS: + +- the same photo, seed 42, 50 steps, guidance 3, 262,144 Gaussians; +- each run in a fresh process under `/usr/bin/time -l`; +- the SDK's stages timed as in [End to end](#end-to-end). + +| stage | ours, opaque | SDK, opaque | ours, webp | SDK, webp | +| --- | ---: | ---: | ---: | ---: | +| load (all models) | 1.36 s (BiRefNet 0.72) | 2.10 s | 1.35 s (BiRefNet 0.72) | 2.16 s | +| preprocess (BiRefNet) | 0.295 s (0.224) | 0.519 s | 0.264 s (0.221) | 0.481 s | +| DINOv3 | 0.155 s | 0.276 s | 0.156 s | 0.271 s | +| VAE | 0.244 s | 0.117 s | 0.244 s | 0.107 s | +| denoise | 27.49 s | 93.94 s | 27.51 s | 93.92 s | +| octree | 0.139 s | 0.750 s | 0.142 s | 0.600 s | +| Gaussian decoder | 0.222 s | 0.491 s | 0.223 s | 0.492 s | +| total | 30.07 s | 101.27 s (96.11 s generating) | 30.04 s | 100.69 s (95.89 s generating) | +| peak footprint | 2.77 GB | 12.88 GB | 2.76 GB | 12.92 GB | +| max RSS | 2.27 GB | 6.19 GB | 2.25 GB | 6.24 GB | + +Our `preprocessed.png` differs from the SDK's preprocessed image in 2 pixels by at most 1 on each photo. + +The PLYs, as in [Quality](#quality). The NN columns are against the SDK's PLY of the same photo; the sampled baseline +is 0.00922 for the opaque photo and 0.00896 for the webp. + +| run | xyz std | opacity mean (q50) | NN, sampled | NN, all points | sample cosine | +| --- | --- | --- | ---: | ---: | ---: | +| SDK, opaque | 0.188, 0.066, 0.305 | 0.547 (0.674) | | 0 | 1 | +| ours, opaque, native noise | 0.175, 0.068, 0.318 | 0.448 (0.251) | 0.04518 | 0.0367 | 0.510 | +| ours, opaque, SDK's noises | 0.188, 0.066, 0.305 | 0.553 (0.700) | 0.00951 | 0.00175 | 0.99642 | +| SDK, webp | 0.189, 0.078, 0.304 | 0.885 (0.9986) | | 0 | 1 | +| ours, webp, native noise | 0.176, 0.073, 0.320 | 0.414 (0.166) | 0.04647 | 0.0365 | 0.478 | +| ours, webp, SDK's noises | 0.189, 0.078, 0.304 | 0.880 (0.9985) | 0.00937 | 0.00165 | 0.99906 | + +- **Every PLY:** 262,144 Gaussians, all finite. +- **Native noise:** the PLYs differ from the SDK's as much as with a preprocessed picture: its starting noise is + MPS's. +- **Parity** (`--noise` with the SDK's VAE and starting noise): the PLYs are as close as sampling resolves. + - Opaque: camera yaw 1.6946, pitch 0.5193, fov 41.45, against the SDK's 1.6887, 0.5208 and 41.60. + - Webp: 1.6990, 0.5358 and 0.1090, against 1.7006, 0.5413 and 0.1096. + +### Kernel contract + +`zig/tests/contract/metal_idol_birefnet*.zig` registers every kernel of `idol_birefnet.metal`: + +- `br_conv`'s nine forms (1x1 plain, GELU, residual, residual by map, BN ReLU; 3x3 plain, BN ReLU, BN; 7x7); +- `br_conv3`; +- `br_norm`'s three modes; +- `br_window` with and without the shift; +- `br_resize`; +- `br_tiles`' two modes; +- `br_deform` at 1x1, 3x3 and 7x7; +- `br_pool`, `br_gate` and `br_head`. + +Each is checked against float64 at its fp32 bound. Twenty-one controls each break one kernel, and each must go red at +its gate. They include: + +- sums or staging in fp16; +- a tap read a row low; +- the GELU's tanh form; +- ReLU by `max`, which drops a NaN; +- unbiased variance; +- the shifted windows' mask dropped; +- the bias table's axes swapped; +- the deformable modulator as `sigmoid` instead of `2 sigmoid`. + +`br_window`, `br_gate` and `br_head` declare inexact sums: fp16 sums would miss by only 61x, 84x and 3x. + +``` +SUMMARY 24 kernels: 24 pass gates 1-4, 0 fail, 0 known findings; 21 controls: 21 red as planted, 0 missed; 0 limit problems +``` + +With the rest of Idol (`--only "metal idol"`): + +``` +SUMMARY 81 kernels: 81 pass gates 1-4, 0 fail, 0 known findings; 78 controls: 78 red as planted, 0 missed; 0 limit problems +``` + ## End to end `tf-idol-generate` (`zig/tests/idol_generate.zig`) chains the pieces above, from one preprocessed picture to a PLY: @@ -487,9 +705,9 @@ The avocado example, the same machine and torch as above (`ATTN_BACKEND=sdpa`): 4. the octree and Gaussian decoders; 5. the PLY and the camera. -Each model is loaded when its stage starts and freed when it ends. Background removal and cropping (BiRefNet, the SDK's -`preprocess`) are not ported, so the input is the picture `preprocess` returns, as a `pixels` or `image` tensor that -`idol_case.py dino` or `vae` writes. +Each model is loaded when its stage starts and freed when it ends. The input is the picture `preprocess` returns, as +a `pixels` or `image` tensor that `idol_case.py dino` or `vae` writes, or a raw photo preprocessed natively +([Background removal and preprocessing](#background-removal-and-preprocessing)). ```sh zig build tf-idol-generate -Dcpu=apple_m1 @@ -523,7 +741,7 @@ Setup: | stage | ours, cold | ours, warm | SDK, MPS | | --- | ---: | ---: | ---: | | load (all models) | 0.67 s | 0.67 s | 2.11 s (and 2.75 s of imports) | -| preprocess (BiRefNet) | not ported | not ported | 0.51 s | +| preprocess (BiRefNet) | picture given | picture given | 0.51 s | | DINOv3 | 0.158 s | 0.155 s | 0.275 s | | VAE | 0.245 s | 0.241 s | 0.119 s | | condition | 0.072 s | 0.071 s | in the sampling | diff --git a/zig/tests/idol_generate.zig b/zig/tests/idol_generate.zig index f6db0c56ed..cdae5c0ec0 100644 --- a/zig/tests/idol_generate.zig +++ b/zig/tests/idol_generate.zig @@ -11,11 +11,16 @@ const gs = id.gaussian; const usage = \\tf-idol-generate SNAPSHOT_DIR VAE_DIR DINO_DIR PICTURE OUT_DIR [--seed N] [--steps N] [--cfg X] [--gaussians N] - \\ [--noise FILE] [--features FILE]... + \\ [--noise FILE] [--features FILE]... [--mode synthetic|real] [--crop-margin X] \\SNAPSHOT_DIR is a speridlabs/idol snapshot, VAE_DIR FLUX.2-klein-4B's vae/ folder, DINO_DIR \\timm/vit_huge_plus_patch16_dinov3.lvd1689m. PICTURE is a safetensors file holding the 1,024 x 1,024 picture \\IdolPipeline.preprocess returns, in [0, 1]: `image` [1024, 1024, 3] (idol_case.py vae) or `pixels` [3, 1024, 1024] \\(idol_case.py dino). Defaults: seed 42, 50 steps, guidance 3, 262,144 Gaussians, as IdolPipeline's. + \\PICTURE may instead be a raw photo (.png, .jpg, .jpeg or .webp), preprocessed natively as IdolPipeline.preprocess + \\does it: the short side resized to 1,024, the background removed by BiRefNet (SNAPSHOT_DIR/birefnet.safetensors) + \\unless the photo has real alpha, the alpha eroded, a square crop of the object, resized to 1,024, pasted on black. + \\--mode synthetic (default; crop 1.2 x the object's box) or real (1.7) as the SDK's mode; --crop-margin X overrides + \\it. A photo also writes OUT_DIR/preprocessed.png, the 1,024 x 1,024 RGB the encoders saw. \\The denoiser starts from torch.Generator("cpu").manual_seed(seed)'s randn [8192, 16] then [5], as idol_case.py \\case draws it (the SDK draws on its device's generator). The VAE posterior's noise and the octree's draws come from \\native seeded generators. --noise FILE gives the posterior's noise instead (`noise` [128, 128, 32], as idol_case.py @@ -23,6 +28,8 @@ const usage = \\(once or twice) gives feature1 [1, 4101, 1280] and feature2 [1, 4101, 128] and skips both encoders (PICTURE: -). \\Writes OUT_DIR/generate.ply, camera.json, start.f32 (the starting latent and camera), sample.f32 and \\tf-idol-generate.json (each stage's seconds). Each model is freed once its stage is done. + \\A photo's tf-idol-generate.json adds preprocess_s (all of preprocessing but BiRefNet's load, which load_s counts) + \\and birefnet_s (BiRefNet's forward and alpha). ; const N = cfg.latent_rows * cfg.latent_channels; @@ -122,6 +129,8 @@ const Options = struct { gaussians: usize = 262144, noise: ?[]const u8 = null, features: [2]?[]const u8 = .{ null, null }, + mode: id.preprocess.Mode = .synthetic, + crop_margin: ?f64 = null, }; fn parse(args: []const [:0]const u8) !Options { @@ -143,6 +152,11 @@ fn parse(args: []const [:0]const u8) !Options { o.noise = v; } else if (std.mem.eql(u8, k, "--features")) { if (o.features[0] == null) o.features[0] = v else if (o.features[1] == null) o.features[1] = v else return error.Usage; + } else if (std.mem.eql(u8, k, "--mode")) { + o.mode = std.meta.stringToEnum(id.preprocess.Mode, v) orelse return error.Usage; + } else if (std.mem.eql(u8, k, "--crop-margin")) { + o.crop_margin = try std.fmt.parseFloat(f64, v); + if (!(o.crop_margin.? > 0)) return error.Usage; } else return error.Usage; } if (o.steps == 0) return error.Usage; @@ -161,8 +175,28 @@ fn maybe(a: std.mem.Allocator, f: *const tf.safetensors.File, name: []const u8, return out; } -/// The picture as [3, 1024, 1024] (DINOv3's input) and [1024, 1024, 3] (the VAE's). -fn picture(a: std.mem.Allocator, io: std.Io, path: []const u8) ![2][]f32 { +/// The picture as [3, 1024, 1024] (DINOv3's) and [1024, 1024, 3] (the VAE's), from a photo or a safetensors file. +fn picture(a: std.mem.Allocator, io: std.Io, device: mtl.Device, args: []const [:0]const u8, o: Options, pt: *id.photo.Timing) ![2][]f32 { + const path = args[4]; + if (id.photo.isPhoto(path)) { + const started = mtl.clock.seconds(); + const rgb = try a.alloc(u8, side * side * 3); + const margin = o.crop_margin orelse id.preprocess.margin(o.mode); + const weights = try std.fmt.allocPrintSentinel(a, "{s}/birefnet.safetensors", .{args[1]}, 0); + try id.photo.preprocess(a, io, device, weights, path, margin, rgb, pt); + pt.preprocess = mtl.clock.seconds() - started - pt.load; + const png = try id.photo.png(a, rgb, side, side); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/preprocessed.png", .{args[5]}), .data = png }); + const chw = try a.alloc(f32, 3 * side * side); + const hwc = try a.alloc(f32, 3 * side * side); + for (0..side * side) |i| for (0..3) |c| { + const v = @as(f32, @floatFromInt(rgb[i * 3 + c])) / 255.0; + hwc[i * 3 + c] = v; + chw[c * side * side + i] = v; + }; + std.debug.print("preprocess: {s} in {d:.3} s ({s}; BiRefNet loaded in {d:.2} s, {d:.3} s, {d:.3} s on the GPU), crop margin {d}\n", .{ path, pt.preprocess, if (pt.real_alpha) "real alpha, no BiRefNet" else "background removed", pt.load, pt.birefnet, pt.birefnet_gpu, margin }); + return .{ chw, hwc }; + } var f = try tf.safetensors.File.open(a, io, path); defer f.close(io); const n = 3 * side * side; @@ -211,7 +245,7 @@ const Times = struct { }; /// feature1 and feature2 from the encoders, or from --features files. -fn encode(a: std.mem.Allocator, io: std.Io, device: mtl.Device, args: []const [:0]const u8, o: Options, t: *Times, metal_peak: *u64) ![2][]f32 { +fn encode(a: std.mem.Allocator, io: std.Io, device: mtl.Device, args: []const [:0]const u8, o: Options, t: *Times, pt: *id.photo.Timing, metal_peak: *u64) ![2][]f32 { if (o.features[0]) |first| { var f1: ?[]f32 = null; var f2: ?[]f32 = null; @@ -228,7 +262,8 @@ fn encode(a: std.mem.Allocator, io: std.Io, device: mtl.Device, args: []const [: std.debug.print("features: given, both encoders skipped\n", .{}); return .{ f1.?, f2.? }; } - const pics = try picture(a, io, args[4]); + const pics = try picture(a, io, device, args, o, pt); + metal_peak.* = @max(metal_peak.*, pt.metal_peak); const f1 = try a.alloc(f32, f1_n); const f2 = try a.alloc(f32, f2_n); { @@ -346,7 +381,8 @@ pub fn main(init: std.process.Init) !u8 { defer pool.pop(); const device = try mtl.Device.init(); std.debug.print("{s}: seed {d}, {d} steps, guidance {d}, {d} Gaussians; starting noise {s}\n", .{ device.name(), o.seed, o.steps, o.guidance, o.gaussians, if (start_given) "given" else "torch CPU-seeded" }); - const features = try encode(a, io, device, args, o, &t, &metal_peak); + var pt: id.photo.Timing = .{}; + const features = try encode(a, io, device, args, o, &t, &pt, &metal_peak); try denoise(a, io, device, args[1], features, x, o, &t, &metal_peak); try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/sample.f32", .{out_dir}), .data = std.mem.sliceAsBytes(x) }); @@ -372,13 +408,14 @@ pub fn main(init: std.process.Init) !u8 { t.ply = mtl.clock.seconds() - started; t.total = mtl.clock.seconds() - begun; - const load = t.load_dino + t.load_vae + t.load_denoiser + t.load_decoders; + const load = pt.load + t.load_dino + t.load_vae + t.load_denoiser + t.load_decoders; std.debug.print("PLY: {d} Gaussians, {d} bytes in {d:.3} s; camera {s}", .{ r.gaussians.count(), bytes.len, t.ply, cam_json }); std.debug.print("total {d:.2} s: load {d:.2}, dino {d:.3}, vae {d:.3}, condition {d:.3}, denoise {d:.2}, octree {d:.3}, Gaussian decoder {d:.3} (+ activation {d:.3}), PLY {d:.3}; Metal peak {d:.2} GB\n", .{ t.total, load, t.dino, t.vae, t.condition, t.denoise, t.octree, t.gaussian_decoder, t.activation, t.ply, gb(metal_peak) }); var out: std.ArrayList(u8) = .empty; try print(&out, a, "{{\"device\": \"{s}\", \"seed\": {d}, \"steps\": {d}, \"cfg\": {d}, \"gaussians\": {d}, \"start_noise\": \"{s}\", \"vae_noise\": \"{s}\", \"features_given\": {}", .{ device.name(), o.seed, o.steps, o.guidance, r.gaussians.count(), if (start_given) "given" else "torch_cpu", if (o.features[0] != null) "unused" else if (o.noise != null) "given" else "native", o.features[0] != null }); try print(&out, a, ", \"load_s\": {d:.4}", .{load}); inline for (@typeInfo(Times).@"struct".field_names) |name| try print(&out, a, ", \"{s}_s\": {d:.4}", .{ name, @field(t, name) }); + if (id.photo.isPhoto(args[4]) and o.features[0] == null) try print(&out, a, ", \"preprocess_s\": {d:.4}, \"birefnet_s\": {d:.4}", .{ pt.preprocess, pt.birefnet }); try print(&out, a, ", \"metal_peak_allocated_gb\": {d:.3}, \"camera\": {s}}}\n", .{ gb(metal_peak), std.mem.trimEnd(u8, cam_json, "\n") }); try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/tf-idol-generate.json", .{out_dir}), .data = out.items }); return 0; From 931d76ce72520125b80e22f4ea54d826904477ec Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 1 Nov 2026 06:11:58 +0800 Subject: [PATCH 21/21] changelog: #1787 for Idol's background removal --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2551fd0a81..db46ddc059 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Idol's FLUX.2 VAE image encoder runs natively on Metal in fp32: one 1,024 x 1,024 picture to the denoiser's 4,101 condition rows, with the posterior's noise as an input. It is 0.185 s an encode on an M5 Ultra against 0.095 s for the SDK's bf16 encoder on torch MPS, and lies far closer to a float64 run: worst row 5.1e-5 against the SDK's 2.7e-2. `tf-idol-vae` gates it from a case `tools/zig/idol_case.py vae` records, with the ungated FLUX.2-klein-4B weights (#1774, after #1766, for #1757). Thanks to @jayleaton. - Idol's DINOv3 ViT-H+/16 image encoder runs natively on M5 Metal: one 1024 x 1024 image to the denoiser's feature1 in bf16, as idol-sdk's ImageEncoder makes it, from the ungated timm weights (DINOv3 License), which `tools/zig/idol_case.py dino` shows to be Meta's gated checkpoint byte for byte and feature for feature. On an M5 Ultra an encode takes 0.120 s against the SDK's 0.261 s on torch MPS, with a 2.28 GB peak footprint. feature1 and each of the 32 blocks on torch's own input sit as close to a float32 ruler as torch's bf16 run, within the 2x gate; `tf-idol-dino` runs those gates, and its eight kernels pass the kernel contract (#1775, after #1766, for #1757). Thanks to @jayleaton. - Idol runs end to end natively on M5 Metal: `tf-idol-generate` takes one preprocessed picture through DINOv3, the FLUX.2 VAE, the denoiser's 50 guided steps and both decoders to a 3DGS PLY and camera. On an M5 Ultra it is 29.1 s against 95.4 s for the SDK's IdolPipeline on MPS (excluding its preprocess), the denoiser 27.5 s against 93.9 s, at a 2.71 GB peak footprint against 12.88 GB. Given the SDK's starting and VAE noise the PLY's points sit as close to the SDK's as sampling resolves, and the camera's angles within 0.006 rad (#1780, after #1773, #1774 and #1775, for #1757). Thanks to @jayleaton. -- Idol's background removal and preprocessing run natively too: `tf-idol-generate` now takes a raw photo (`.png`, `.jpg`, `.jpeg` or `.webp`), with `--mode synthetic|real` and `--crop-margin X`, writes the picture the encoders saw as `preprocessed.png`, and adds `preprocess_s` and `birefnet_s` to its JSON. BiRefNet (MIT, from the Idol snapshot) runs in fp32 on Metal in 0.172 s against the SDK's 0.48 to 0.52 s preprocess, its logits within torch's own distance from a float64 ruler; Pillow's resizes, the erosion, crop and paste are byte for byte the SDK's, and the preprocessed picture differs from the SDK's in 2 pixels by 1. From a photo a generation takes 30.1 s against the SDK's 101.3 s at a 2.77 GB peak footprint against 12.88 GB (after #1780, for #1757). Thanks to @jayleaton. +- Idol's background removal and preprocessing run natively too: `tf-idol-generate` now takes a raw photo (`.png`, `.jpg`, `.jpeg` or `.webp`), with `--mode synthetic|real` and `--crop-margin X`, writes the picture the encoders saw as `preprocessed.png`, and adds `preprocess_s` and `birefnet_s` to its JSON. BiRefNet (MIT, from the Idol snapshot) runs in fp32 on Metal in 0.172 s against the SDK's 0.48 to 0.52 s preprocess, its logits within torch's own distance from a float64 ruler; Pillow's resizes, the erosion, crop and paste are byte for byte the SDK's, and the preprocessed picture differs from the SDK's in 2 pixels by 1. From a photo a generation takes 30.1 s against the SDK's 101.3 s at a 2.77 GB peak footprint against 12.88 GB (#1787, after #1780, for #1757). Thanks to @jayleaton. - DeepSeek-V4.x: `tool_choice: "required"` and named tool calls are now constrained to the model's own DSML call markup, so those requests are held to a valid call on DeepSeek too.