diff --git a/CHANGELOG.md b/CHANGELOG.md index a04ea41b9f..bf33fb8927 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 +- MOSS-Image's image decoder (the checkpoint's FLUX.2 small VAE, in float32 as the pipeline runs it) runs on Metal inside the `moss_image` family, and `tf-moss-image` writes `picture.png` itself, so torch no longer decodes (#1687, #1504). Thanks to @BobClawblaw. + - On CUDA, a driver older than CUDA 13.0 is refused at startup with a plain message naming the driver to install (R580 or newer), before any kernel loads; an R550 driver had ended the process with SIGILL inside libcuda on the first module load (#1791, #1724). Thanks to @rootsbymenda. - Spark 2.5 on CUDA: `TF_KV_SHARE=n` keeps prompts as the KV pool's pages in the shared prompt store, with at most n kept states, so a conversation resumes at the store's last 256-token mark from pages it already holds instead of a full copy. Unset or 0, nothing changes. The native CUDA adapter it uses is family-neutral; Kolibri-1 follows. diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index c59c8ca519..89b63ce12c 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -11,21 +11,23 @@ transformer on the Metal runtime from the released bf16 diffusers checkpoint: The card's guidance runs two passes a step: the prompt, and the empty negative prompt with block 22 skipped. -The prompt encoder, the checkpoint's Qwen3.5 text model, runs on Metal too (below): `tf-moss-image --prompt` makes both -prompts' rows itself. The image decoder (FLUX.2's small decoder) is not ported. `tools/zig/moss_image_case.py` runs it, -and the reference pipeline, through the publisher's own `pipeline_moss_image.py` in torch. +The prompt encoder, the checkpoint's Qwen3.5 text model, and the image decoder, FLUX.2's small VAE decoder, run on +Metal too (below): `tf-moss-image --prompt` makes both prompts' rows itself, and `tf-moss-image` turns the final latents +into the picture. `tools/zig/moss_image_case.py` runs the reference pipeline through the publisher's own +`pipeline_moss_image.py` 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. +It runs on the M5 alone, the only chip it was checked on: `Model.load` and `Decoder.load` refuse a GPU without tensor +units. The same transformer runs on CUDA too: [moss-image-cuda.md](moss-image-cuda.md). ## Checkpoint `OpenMOSS-Team/MOSS-Image-Base` at `85975793d8b3d432f1b6e27f352abaa4c2a4c0d3`, Apache 2.0. `transformer/config.json` -must be `MOSSImageTransformer` at the released geometry; any other config is refused, and so is a checkpoint with -tensors this family does not read. +must be `MOSSImageTransformer` at the released geometry, and `vae/config.json` `AutoencoderKLFlux2` at its released +decoder geometry. Any other config is refused, and so is a checkpoint with tensors this family does not read. ## Precision @@ -37,6 +39,14 @@ tensors this family does not read. - The rotary tables in bf16, as the pipeline casts them; its frequencies are the checkpoint's `rope.bands`. - No int8. +The decoder is float32 throughout, as the pipeline decodes: its weights as published (float32), fp32 sums in every +convolution and in its attention, GroupNorm's mean and variance in two fp32 passes, then the pipeline's own 8-bit +conversion (clamp, `* 0.5 + 0.5`, `* 255`, round to even). Its 3x3 convolutions run as Winograd F(2x2, 3x3) +(`md_wino`): the weights' transform G g G^T is computed in float64 at load and rounded once, the input's and the +output's transforms in fp32 in the kernel, 2.25x fewer products than a direct convolution. Its error stays the size +of a direct convolution's (see Measured). Like any Winograd convolution, an infinite input can come out NaN where a +direct one gives an infinity; the contract checks the classes against the Winograd formula's. + The bf16 projections are the Qwen-VL tower's `tfp::gemm_nax`. Attention is the prefill `tfp::attention_nax` template's full (non-causal) form. The time MLP and both modulation projections run on the host, once a timestep. @@ -118,10 +128,28 @@ Gate 4: Gate 5 holds the 17 launches of the family's source and the two split-K projections. +`zig/tests/contract/metal_moss_image_decoder*.zig` registers every kernel of `moss_image_decoder.metal`: + +- `md_unpack` (the latents' batch norm undone, the patches unpacked); +- `md_wino` (64 output channels a threadgroup) and `md_wino_32` (32, for 96 and 3 channels) in the forms the decoder + launches: 3x3, 3x3 read nearest from half size, and 3x3 with a residual, against the float64 convolution through the + exact transform, within a bound on the transformed terms' magnitude; +- `md_conv`, the 1x1 convolutions, with and without a residual; +- `md_norm`, with and without SiLU, and GroupNorm's statistics: `md_group_part` (plain and centered) and + `md_group_finish`; +- `md_attention` and `md_pixels`. + +Each fp32 step's output must lie in the range float32 allows about float64. Gate 4 uses integer channels, pixels or +values whose sums pass fp16's range and cancel. `md_pixels` is checked byte for byte against the same fp32 steps, with +every tie between two bytes among its inputs. Nine controls each break one kernel: a swapped patch axis, a 1x1's +inputs staged in half, a Winograd patch staged in half, a term dropped from its output transform, a SiLU in half, +both statistics kernels summing in half, the attention's rescale dropped, and the pixels' halves rounded up. Each is +caught at the gate it names. Gate 5 holds its nine launches. + `tf-kernel-contract --only "moss image"` on the M5 Ultra: ``` -SUMMARY 22 kernels: 22 pass gates 1-4, 0 fail, 0 known findings; 16 controls: 16 red as planted, 0 missed; 0 limit problems +SUMMARY 37 kernels: 37 pass gates 1-4, 0 fail, 0 known findings; 25 controls: 25 red as planted, 0 missed; 0 limit problems ``` ## Run a case @@ -133,9 +161,8 @@ only when the file's sha256 is the pinned revision's (`1dac1aca…8b31`) and ref hf download OpenMOSS-Team/MOSS-Image-Base --revision 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 --local-dir MODEL_DIR python3 tools/zig/moss_image_case.py case MODEL_DIR case/ --prompt-file prompt.txt # needs torch, diffusers, transformers zig build tf-moss-image -Dcpu=apple_m1 -zig-out/bin/tf-moss-image MODEL_DIR/transformer case/case.safetensors case/ -python3 tools/zig/moss_image_case.py decode MODEL_DIR case/latents.f32 case/zig.png -python3 tools/zig/moss_image_case.py compare case/reference.png case/zig.png case/side.png +zig-out/bin/tf-moss-image MODEL_DIR/transformer case/case.safetensors case/ # case/picture.png +python3 tools/zig/moss_image_case.py compare case/reference.png case/picture.png case/side.png python3 tools/zig/moss_image_case.py exact MODEL_DIR case/ # the same schedule with a float32 transformer python3 tools/zig/moss_image_case.py exact MODEL_DIR case/ --fp32-encoder prompt.txt # and a float32 encoder @@ -155,10 +182,23 @@ The case holds: - with `--fp32-steps N`, a float32 run of the same transformer for the first N steps. `tf-moss-image` first runs both branches of each step on the pipeline's own input and compares the velocities. It then -runs the whole guided schedule from the starting latents, times it, and writes the final latents. A case made by -`start` holds only the prompts, the schedule and the noise; `tf-moss-image` then just makes the picture. With -`--prompt` (and `--negative`, else the empty negative prompt), it encodes both prompts natively in place of the case's -rows; `start --no-encoder` leaves them out of the case, and torch never loads the encoder. +runs the whole guided schedule from the starting latents, times it, frees the transformer, and decodes the final +latents with the checkpoint's `vae/` decoder into `picture.png`. A case made by `start` holds only the prompts, the +schedule and the noise; `tf-moss-image` then just makes the picture. With `--prompt` (and `--negative`, else the empty +negative prompt), it encodes both prompts natively in place of the case's rows; `start --no-encoder` leaves them out of +the case, and torch never loads the encoder. + +To check the decoder alone against a float64 decode of the same latents: + +```sh +# float64 on the CPU, and float32 on MPS as the pipeline decodes +python3 tools/zig/moss_image_case.py vae MODEL_DIR case/latents.f32 case/decode.f64 case/torch.f32 +zig build tf-moss-image-decoder -Dcpu=apple_m1 +zig-out/bin/tf-moss-image-decoder MODEL_DIR case/latents.f32 1024 1024 case/decode.f64 case/torch.f32 case/decoded.png +``` + +Ours must be within twice torch float32's distance from float64, in cosine error and in its worst row of pixels, and a +second decode must repeat the first byte for byte; it fails otherwise. On MPS, which has no float64, the case tool computes the time sinusoid on the host in float64 and moves its float32 result to the GPU; every other step of the publisher's code runs unchanged. @@ -226,3 +266,12 @@ Results: pipeline's own picture: 47.07 dB. - Speed, one session with no other GPU work: 0.054 to 0.067 s a prompt against torch MPS's 0.158 to 0.165 s (warm). Loading takes 0.6 s against 2.60 s. +- **Decoder**, on the transformer's final latents of both cases, against the publisher's decoder in float64 on the CPU: + - Ours: max |diff| 8.9e-6 (fox) and 5.2e-6 (sticky note) before the clamp; 162 of 3,145,728 and 45 of 3,168,000 + 8-bit values differ from float64's, each by 1. + - torch's own float32 decode on MPS, in the same session: 155 and 34 values differ from float64's, each by 1 (the + count moves between torch runs). Ours and torch's pictures differ by at most 1 (95.4 and 97.7 dB PSNR). + - Speed, back to back in one session with no other GPU work, the median of three: ours 0.181 s and 0.188 s on the + GPU; torch's warm float32 decode on MPS 0.186 s and 0.191 s. Ours loads in 0.3 to 0.4 s. Its scratch is five + float32 picture tensors at the widest stage: 4.7 GB at peak at 1024x1024, after the transformer is freed, so + `tf-moss-image` peaks at the transformer's 11.3 GB. diff --git a/tools/zig/moss_image_case.py b/tools/zig/moss_image_case.py index f1ef12b6f1..bc3ad002c7 100644 --- a/tools/zig/moss_image_case.py +++ b/tools/zig/moss_image_case.py @@ -20,6 +20,9 @@ with --fp32-encoder F, F's rows from a float32 encoder too (CASE_DIR/exact_fp32_encoder.png) decode MODEL_DIR LATENTS.f32 PNG raw float32 latents (rows x columns x 128) from tf-moss-image to a picture + vae MODEL_DIR LATENTS.f32 OUT.f64 TORCH.f32 the same decode in float64 on the CPU and in float32 as the + pipeline runs it, each output before the clamp [height, width, 3] raw, + for tf-moss-image-decoder encode MODEL_DIR OUT_DIR [--prompt-file F] the prompt encoder alone: token ids, embedding rows, and every layer's output and the final rows in bf16 and in float32, and each step alone in bf16 on the float32 run's input @@ -79,19 +82,24 @@ def rows(x): return x[0].float().permute(1, 2, 0).contiguous().cpu() -def to_png(pipe, latents, path: str) -> None: - """Latents [1, 128, H, W] through the pipeline's own decode path (its batch-norm undo, unpack, VAE) to a PNG.""" +def decoded(vae, latents): + """Latents [1, 128, H, W] through the pipeline's own decode path (its batch-norm undo, unpack, VAE), unclamped.""" import torch with torch.no_grad(): - z = latents.to(device=pipe.vae.device, dtype=pipe.vae.dtype) - bn = pipe.vae.bn - eps = pipe.vae.config.batch_norm_eps + z = latents.to(device=vae.device, dtype=vae.dtype) + bn = vae.bn + eps = vae.config.batch_norm_eps z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1) b, c, h, w = z.shape z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w) - image = pipe.vae.decode(z, return_dict=False)[0].clamp(-1, 1) - pipe.image_processor.postprocess(image, output_type="pil")[0].save(path) + return vae.decode(z, return_dict=False)[0] + + +def to_png(pipe, latents, path: str) -> None: + """Latents [1, 128, H, W] decoded as the pipeline does, clamped, to a PNG.""" + image = decoded(pipe.vae, latents).clamp(-1, 1) + pipe.image_processor.postprocess(image, output_type="pil")[0].save(path) def case(args) -> int: @@ -228,6 +236,28 @@ def decode(args) -> int: return 0 +def vae(args) -> int: + import numpy as np + import torch + from diffusers import AutoencoderKLFlux2 + + raw = np.fromfile(args.latents, dtype=np.float32) + h, w = args.height // 16, args.width // 16 + if raw.size != h * w * 128: + sys.exit(f"{args.latents} holds {raw.size} values; {args.width}x{args.height} needs {h * w * 128}") + latents = torch.from_numpy(raw).reshape(1, h, w, 128).permute(0, 3, 1, 2) + model = AutoencoderKLFlux2.from_pretrained(Path(args.model_dir) / "vae", torch_dtype=torch.float32).eval() + torch_rows = decoded(model.to(args.device), latents)[0].permute(1, 2, 0).contiguous().cpu().numpy() + torch_rows.astype(np.float32).tofile(args.torch_out) + model = model.cpu().double() + mark = time.perf_counter() + out = decoded(model, latents)[0].permute(1, 2, 0).contiguous().numpy() + out.astype(np.float64).tofile(args.out) + print(json.dumps({"out": args.out, "torch_out": args.torch_out, "shape": list(out.shape), + "decode_s": round(time.perf_counter() - mark, 1)})) + return 0 + + def compare(args) -> int: import numpy as np from PIL import Image @@ -399,6 +429,14 @@ def main() -> int: n.add_argument("out_dir") n.add_argument("--prompt-file", help="the prompt (default: the empty negative prompt)") n.add_argument("--device", default="mps") + v = sub.add_parser("vae") + v.add_argument("model_dir") + v.add_argument("latents") + v.add_argument("out") + v.add_argument("torch_out") + v.add_argument("--device", default="mps", help="where the float32 decode runs, as the pipeline would") + v.add_argument("--width", type=int, default=1024) + v.add_argument("--height", type=int, default=1024) p = sub.add_parser("compare") p.add_argument("a") p.add_argument("b") @@ -406,7 +444,7 @@ def main() -> int: n = sub.add_parser("noise") n.add_argument("out") args = parser.parse_args() - return {"case": case, "start": start, "exact": exact, "decode": decode, "encode": encode, "compare": compare, "noise": noise}[args.mode](args) + return {"case": case, "start": start, "exact": exact, "decode": decode, "encode": encode, "vae": vae, "compare": compare, "noise": noise}[args.mode](args) if __name__ == "__main__": diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 160447387b..cbffa0074e 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -316,8 +316,9 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-glm-video", .path = "zig/tests/glm_video.zig", .about = "GLM-5.3-Flash's video preparation on a clip (libtfvideo), written for tools/zig/glm_video_fixtures.py" }, .{ .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-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image on a reference case: each step's two branches compared, then the guided schedule timed and decoded" }, .{ .name = "tf-moss-image-encoder", .path = "zig/tests/moss_image_encoder.zig", .about = "MOSS-Image's prompt encoder on a reference case: the ids checked, each step gated against float32, then timed" }, + .{ .name = "tf-moss-image-decoder", .path = "zig/tests/moss_image_decoder.zig", .about = "MOSS-Image's image decoder on latents, gated against float64 beside torch float32, 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/moss_image_decoder.metal b/zig/kernels/metal/moss_image_decoder.metal new file mode 100644 index 0000000000..8e20dce47c --- /dev/null +++ b/zig/kernels/metal/moss_image_decoder.metal @@ -0,0 +1,421 @@ +// MOSS-Image's image decoder (FLUX.2's VAE decoder) in fp32, each picture tensor [rows, columns, channels]. + +#include +using namespace metal; + +#define TF_UNROLL_MD _Pragma("clang loop unroll(full)") + +// The latents [h * w, 128] with the batch norm undone (x * s + m), each 2 x 2 patch unpacked: z [2h, 2w, 32]. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void md_unpack( + const device float* x [[buffer(0)]], const device float* s [[buffer(1)]], const device float* m [[buffer(2)]], + device float* z [[buffer(3)]], constant uint2& hw [[buffer(4)]], uint2 at [[thread_position_in_grid]]) { + const uint w2 = hw.y * 2; + if (at.y >= hw.x * 2 * w2) return; + const uint c = at.x; + const uint py = at.y / w2; + const uint px = at.y % w2; + const uint k = c * 4 + (py % 2) * 2 + px % 2; + const float v = x[ulong((py / 2) * hw.y + px / 2) * 128 + k] * s[k]; + z[ulong(at.y) * 32 + c] = v + m[k]; +} + +struct MdConv { + uint rows; + uint cols; + uint cin; + uint cout; + uint ldw; + uint up; + uint add; +}; + +// y = x w + b (+ r): a 1x1 convolution, x [pixels, cin], w [cin, ldw]; a 64 x 64 tile a threadgroup, fp32 sums. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void md_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 MdConv& 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; + 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 sr = (sg / 2) * 32; + const uint sc = (sg % 2) * 32; + simdgroup_float8x8 acc[4][4]; + TF_UNROLL_MD + for (uint i = 0; i < 4; i++) { + TF_UNROLL_MD + for (uint j = 0; j < 4; j++) acc[i][j] = simdgroup_float8x8(0.0f); + } + for (uint k0 = 0; k0 < p.cin; k0 += BK) { + threadgroup float4* dst = (threadgroup float4*)(As + arow * lda + acol); + if (pix < pixels) { + const device float4* src = (const device float4*)(x + ulong(pix) * p.cin + k0 + acol); + TF_UNROLL_MD + for (uint i = 0; i < 4; i++) dst[i] = src[i]; + } else { + TF_UNROLL_MD + 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_MD + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 am[4]; + simdgroup_float8x8 bm[4]; + TF_UNROLL_MD + for (uint i = 0; i < 4; i++) simdgroup_load(am[i], As + (sr + i * 8) * lda + kk, lda); + TF_UNROLL_MD + for (uint j = 0; j < 4; j++) simdgroup_load(bm[j], wk + ulong(kk) * p.ldw + j * 8, p.ldw); + TF_UNROLL_MD + for (uint i = 0; i < 4; i++) { + TF_UNROLL_MD + 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_MD + for (uint i = 0; i < 4; i++) { + TF_UNROLL_MD + 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 md_wa(uint i) { return i == 2 ? 2 : (i == 0 ? 0 : 1); } +inline uint md_wb(uint i) { return i == 0 ? 2 : (i == 3 ? 3 : (i == 1 ? 2 : 1)); } +inline float md_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 8 NF channels a group. +template +inline void md_wino_body(const device float* x, const device float* u, const device float* b, const device float* r, + device float* y, constant MdConv& p, uint3 tg, uint t, uint lane, uint sg, + threadgroup float* buf) { + constexpr uint BK = 16; + constexpr uint TILES = 16; + constexpr uint PR = 6; + constexpr uint PC = 18; + const uint n0 = tg.x * NF * 8; + const int y0 = int(tg.y) * 4; + const int x0 = int(tg.z) * 16; + const uint in_cols = p.up ? p.cols / 2 : p.cols; + // 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_MD + for (uint i = 0; i < 2; i++) { + TF_UNROLL_MD + 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)) { + const uint sy = p.up ? uint(iy) / 2 : uint(iy); + const uint sx = p.up ? uint(ix) / 2 : uint(ix); + v = *(const device float4*)(x + ulong(sy * in_cols + sx) * p.cin + c0 + q * 4); + } + *(threadgroup float4*)(buf + px * BK + q * 4) = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + TF_UNROLL_MD + for (uint kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 vm[2]; + TF_UNROLL_MD + 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_MD + for (uint c = 0; c < 2; c++) { + const threadgroup float* q = d + c; + const float wa = q[(md_wa(pi) * PC + md_wa(pj)) * BK] + md_ws(pi) * q[(md_wb(pi) * PC + md_wa(pj)) * BK]; + const float wb = q[(md_wa(pi) * PC + md_wb(pj)) * BK] + md_ws(pi) * q[(md_wb(pi) * PC + md_wb(pj)) * BK]; + vm[i].thread_elements()[c] = wa + md_ws(pj) * wb; + } + } + const device float* up_ = u + (ulong(sg) * p.cin + c0 + kk) * p.ldw + n0; + TF_UNROLL_MD + for (uint j = 0; j < NF; j++) { + simdgroup_float8x8 um; + simdgroup_load(um, up_ + j * 8, p.ldw); + TF_UNROLL_MD + 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_MD + for (uint i = 0; i < 2; i++) { + TF_UNROLL_MD + 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_MD + 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_MD + 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_MD + 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); + } +} + +// the input patch [6, 18, 16]; the epilogue's [16 points, 16 tiles, 32 channels] after +[[max_total_threads_per_threadgroup(512)]] [[kernel]] void md_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 MdConv& 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]]) { + threadgroup float buf[16 * 16 * 32]; + md_wino_body<8>(x, u, b, r, y, p, tg, t, lane, sg, buf); +} + +// md_wino for 32 output channels a threadgroup: an output a multiple of 32 but not of 64 wastes none +[[max_total_threads_per_threadgroup(512)]] [[kernel]] void md_wino_32( + 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 MdConv& 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]]) { + threadgroup float buf[16 * 16 * 32]; + md_wino_body<4>(x, u, b, r, y, p, tg, t, lane, sg, buf); +} + +struct MdNorm { + 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 md_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 MdNorm& 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 MdGroups { + 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(384)]] [[kernel]] void md_group_part( + const device float* x [[buffer(0)]], const device float* stats [[buffer(1)]], device float* part [[buffer(2)]], + constant MdGroups& p [[buffer(3)]], uint tg [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]]) { + threadgroup float sums[384]; + 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 md_group_finish( + const device float* part [[buffer(0)]], device float* stats [[buffer(1)]], constant MdGroups& 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 MdAttention { + uint rows; + float scale; +}; + +// o [rows, 384] = softmax(q k^T * scale) v of qkv rows [rows, 1152]: 8 queries a threadgroup, keys 32 at a time. +[[max_total_threads_per_threadgroup(128)]] [[kernel]] void md_attention( + const device float* qkv [[buffer(0)]], device float* o [[buffer(1)]], constant MdAttention& 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 = 384; + 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_MD + 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_MD + for (uint j = 0; j < BK / 8; j++) { + simdgroup_float8x8 s = simdgroup_float8x8(0.0f); + TF_UNROLL_MD + 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_MD + 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_MD + 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_MD + 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_MD + 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_MD + 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_MD + for (uint j = 0; j < BK / 8; j++) { + simdgroup_float8x8 pm; + simdgroup_load(pm, Ps + j * 8, BK); + TF_UNROLL_MD + 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_MD + 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); + } +} + +// rgb = the 8-bit values of y: clamped to [-1, 1], * 0.5 + 0.5 clamped to [0, 1], * 255, rounded to even. +[[max_total_threads_per_threadgroup(256)]] [[kernel]] void md_pixels( + const device float* y [[buffer(0)]], device uchar* rgb [[buffer(1)]], constant uint& n [[buffer(2)]], + uint i [[thread_position_in_grid]]) { + if (i >= n) return; + float v = clamp(y[i], -1.0f, 1.0f); + v = clamp(v * 0.5f + 0.5f, 0.0f, 1.0f); + rgb[i] = uchar(metal::rint(v * 255.0f)); +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index f1a8258d04..91cef8c22c 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: the prefill attention template, then its transformer's and encoder's steps (gemms: core_qwen_vision). pub const moss_image = @embedFile("prefill/attention_nax.metal") ++ @embedFile("moss_image.metal") ++ @embedFile("moss_image_encoder.metal"); +/// MOSS-Image's image decoder: its own float32 steps, without a template. +pub const moss_image_decoder = @embedFile("moss_image_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/core/contract/moss_image_records.zig b/zig/src/core/contract/moss_image_records.zig index 5995c26fee..05c6acd50a 100644 --- a/zig/src/core/contract/moss_image_records.zig +++ b/zig/src/core/contract/moss_image_records.zig @@ -1,4 +1,4 @@ -//! MOSS-Image's launches: its transformer's and prompt encoder's own kernels, and the split-K GEMM it reads sums from. +//! MOSS-Image's launches: its transformer's, encoder's and decoder's kernels, and the split-K GEMM it reads sums from. const std = @import("std"); const M = @import("limits.zig").MetalLaunch; pub const records = [_]M{ @@ -21,6 +21,15 @@ pub const records = [_]M{ .{ .name = "moss image me_gdn_prep", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "moss image me_gdn", .threads = 128, .tg_memory = 1024, .reserved = 128 }, .{ .name = "moss image me_gated_norm", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image decoder md_unpack", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image decoder md_conv", .threads = 128, .tg_memory = 26624, .reserved = 128 }, + .{ .name = "moss image decoder md_wino", .threads = 512, .tg_memory = 32768, .reserved = 512 }, + .{ .name = "moss image decoder md_wino_32", .threads = 512, .tg_memory = 32768, .reserved = 512 }, + .{ .name = "moss image decoder md_norm", .threads = 256, .tg_memory = 0, .reserved = 256 }, + .{ .name = "moss image decoder md_group_part", .threads = 384, .tg_memory = 1536, .reserved = 384 }, + .{ .name = "moss image decoder md_group_finish", .threads = 1024, .tg_memory = 0, .reserved = 1024 }, + .{ .name = "moss image decoder md_attention", .threads = 128, .tg_memory = 5168, .reserved = 128 }, + .{ .name = "moss image decoder md_pixels", .threads = 256, .tg_memory = 0, .reserved = 256 }, }; test "the MOSS-Image launches fit every Apple family" { diff --git a/zig/src/families/moss_image/decoder.zig b/zig/src/families/moss_image/decoder.zig new file mode 100644 index 0000000000..bd3dd3e03b --- /dev/null +++ b/zig/src/families/moss_image/decoder.zig @@ -0,0 +1,553 @@ +//! MOSS-Image's image decoder on Metal: FLUX.2's VAE decoder (the checkpoint's vae/) in fp32, latents to 8-bit RGB. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const ck = @import("../../core/checkpoint_metal.zig"); +const cfg = @import("config.zig"); +const Ref = @import("../../core/qwen_vl/weights.zig").Ref; +const bind = @import("kernels.zig").Kernels.bind; + +/// The decoder's widths from its last stage to its first (decoder_block_out_channels reversed), and the rest. +pub const widths = [4]usize{ 384, 384, 192, 96 }; +pub const latent_channels = 32; +pub const groups = 32; +pub const resnets = 3; +pub const eps: f32 = 1e-6; +pub const batch_norm_eps: f32 = 1e-4; +/// Pixels a GroupNorm statistics part sums. +pub const chunk = 256; + +/// Accepts vae/config.json at FLUX.2's released decoder 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 = .{ .{ "latent_channels", latent_channels }, .{ "layers_per_block", resnets - 1 }, .{ "norm_num_groups", groups }, .{ "out_channels", 3 } }; + 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("moss image decoder: {s} is not {d}; the decoder runs the released geometry only", .{ kv[0], kv[1] }); + return error.UnsupportedGeometry; + } + } + const lists = .{ .{ "decoder_block_out_channels", [_]i64{ 96, 192, 384, 384 } }, .{ "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; + } + const flags = .{ "mid_block_add_attention", "use_post_quant_conv" }; + inline for (flags) |f| { + const v = o.get(f) orelse return error.BadConfig; + if (v != .bool or !v.bool) 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 (a 1x1's [cin, ldw], a 3x3's G g G^T [16, cin, ldw], ldw cout padded to 64) and bias. +pub const Conv = struct { w: Ref, b: Ref, cin: usize, cout: usize, ldw: usize, taps: usize }; +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, upsample: ?Conv }; + +/// The decoder's pipelines, by function name. +pub const names = [_][]const u8{ "md_unpack", "md_conv", "md_wino", "md_wino_32", "md_norm", "md_group_part", "md_group_finish", "md_attention", "md_pixels" }; + +pub const Timing = struct { wall: f64, gpu: f64 }; + +pub const Decoder = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + lib: mtl.Library, + p: [names.len]mtl.Pipeline, + /// Every weight, reordered: 1x1 convolutions [cin, cout padded], 3x3s G g G^T, the attention's q | k | v as one. + weights: mtl.Buffer, + bn_scale: Ref, + bn_mean: Ref, + post_quant: Conv, + conv_in: Conv, + mid: [2]Resnet, + attn_norm: Norm, + qkv: Conv, + attn_out: Conv, + stages: [4]Stage, + norm_out: Norm, + conv_out: Conv, + /// The picture this decoder's scratch is for. + width: usize, + height: usize, + scratch: mtl.Buffer, + lat: Ref, + t: [5]Ref, + qkv_rows: Ref, + att: Ref, + stats: Ref, + part: Ref, + out: Ref, + rgb: Ref, + + fn get(d: *const Decoder, comptime name: []const u8) mtl.Pipeline { + inline for (names, 0..) |n, i| if (comptime std.mem.eql(u8, n, name)) return d.p[i]; + @compileError("no pipeline " ++ name); + } + + /// The vae/ folder `dir`, with scratch for one width x height picture. + pub fn load(gpa: std.mem.Allocator, io: std.Io, device: mtl.Device, dir: []const u8, width: usize, height: usize) !*Decoder { + const g = try cfg.grid(width, height); + 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); + // qualified on the M5 alone, as the transformer is + if (!device.tensorUnits()) { + std.log.err("moss image decoder: {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 d = try gpa.create(Decoder); + errdefer gpa.destroy(d); + d.* = undefined; + d.gpa = gpa; + d.device = device; + d.width = width; + d.height = height; + d.queue = try device.queue(); + errdefer d.queue.deinit(); + const src = try frags.source(device, gpa, sources.moss_image_decoder); + defer gpa.free(src); + d.lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + errdefer d.lib.deinit(); + var made: usize = 0; + errdefer for (d.p[0..made]) |p| p.deinit(); + for (names, 0..) |name, i| { + d.p[i] = try mtl.Pipeline.init(device, d.lib, name, false); + made += 1; + } + 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), ""); + try d.pack(&c); + errdefer d.weights.deinit(); + try d.plan(g.rows, g.columns); + return d; + } + + pub fn deinit(d: *Decoder) void { + d.scratch.deinit(); + d.weights.deinit(); + for (d.p) |p| p.deinit(); + d.lib.deinit(); + d.queue.deinit(); + d.gpa.destroy(d); + } + + /// The decoder's tensors of `c`, each checked for its shape, copied into one buffer in the kernels' layouts. + fn pack(d: *Decoder, c: *const ck.Checkpoint) !void { + var p: Packer = .{ .c = c }; + try d.layout(&p); + if (p.read != p.expected(c)) { + std.log.err("moss image decoder: the checkpoint has {d} decoder tensors; the decoder reads {d}", .{ p.expected(c), p.read }); + return error.UnexpectedTensors; + } + d.weights = try d.device.buffer(p.at * 4, mtl.ResourceOptions.shared); + errdefer d.weights.deinit(); + p.out = @as([*]f32, @ptrCast(@alignCast(d.weights.contents())))[0..p.at]; + p.at = 0; + p.read = 0; + p.buf = d.weights; + try d.layout(&p); + } + + /// Walks every tensor in the order it is packed: sizes it on the first pass, copies it on the second. + fn layout(d: *Decoder, p: *Packer) !void { + const var_ = try p.tensor("bn.running_var", &.{128}); + const mean = try p.tensor("bn.running_mean", &.{128}); + d.bn_scale = p.alloc(128); + d.bn_mean = p.alloc(128); + if (p.out) |out| { + const s = out[d.bn_scale.off / 4 ..][0..128]; + const m = out[d.bn_mean.off / 4 ..][0..128]; + for (s, m, var_, mean) |*so, *mo, v, mv| { + so.* = @sqrt(v + batch_norm_eps); + mo.* = mv; + } + } + d.post_quant = try p.conv("post_quant_conv", latent_channels, latent_channels, 1); + d.conv_in = try p.conv("decoder.conv_in", latent_channels, widths[0], 9); + var name: [80]u8 = undefined; + for (&d.mid, 0..) |*r, i| r.* = try p.resnet(try std.fmt.bufPrint(&name, "decoder.mid_block.resnets.{d}", .{i}), widths[0], widths[0]); + d.attn_norm = try p.norm("decoder.mid_block.attentions.0.group_norm", widths[0]); + d.qkv = try p.linears(&.{ "decoder.mid_block.attentions.0.to_q", "decoder.mid_block.attentions.0.to_k", "decoder.mid_block.attentions.0.to_v" }, widths[0]); + d.attn_out = try p.linears(&.{"decoder.mid_block.attentions.0.to_out.0"}, widths[0]); + var prev = widths[0]; + for (&d.stages, widths, 0..) |*s, out, i| { + for (&s.resnets, 0..) |*r, j| { + r.* = try p.resnet(try std.fmt.bufPrint(&name, "decoder.up_blocks.{d}.resnets.{d}", .{ i, j }), if (j == 0) prev else out, out); + } + s.upsample = if (i + 1 < widths.len) try p.conv(try std.fmt.bufPrint(&name, "decoder.up_blocks.{d}.upsamplers.0.conv", .{i}), out, out, 9) else null; + prev = out; + } + d.norm_out = try p.norm("decoder.conv_norm_out", widths[3]); + d.conv_out = try p.conv("decoder.conv_out", widths[3], 3, 9); + } + + /// Scratch for a rows x columns latent grid: five picture tensors at the widest stage, the attention's rows. + fn plan(d: *Decoder, rows: usize, columns: usize) !void { + const R = rows * 2; + const C = columns * 2; + var most: usize = R * C * widths[0]; + var scale: usize = 1; + for (widths, 0..) |w, i| { + most = @max(most, R * C * scale * scale * w); + if (i + 1 < widths.len) { + scale *= 2; + most = @max(most, R * C * scale * scale * w); + } + } + const picture = R * C * 64; + const attn_rows = std.mem.alignForward(usize, R * C, 32); + const parts = [_]usize{ rows * columns * cfg.channels * 4, most * 4, most * 4, most * 4, most * 4, most * 4, attn_rows * 3 * widths[0] * 4, attn_rows * widths[0] * 4, 64 * 4, groups * chunks(picture) * 4, picture * 3 * 4, picture * 3 }; + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n, 256); + d.scratch = try d.device.buffer(total, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + @memset(d.scratch.contents()[0..total], 0); + var at: usize = 0; + const views = [_]*Ref{ &d.lat, &d.t[0], &d.t[1], &d.t[2], &d.t[3], &d.t[4], &d.qkv_rows, &d.att, &d.stats, &d.part, &d.out, &d.rgb }; + for (views, parts) |v, n| { + v.* = .{ .buf = d.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + } + + /// The 8-bit RGB picture [height, width, 3] of `latents` (rows x columns, 128), the transformer's output. + pub fn decode(d: *Decoder, latents: []const f32, rgb: []u8) !Timing { + const g = try cfg.grid(d.width, d.height); + if (latents.len != g.rows * g.columns * cfg.channels) return error.LatentRows; + if (rgb.len != d.width * d.height * 3) return error.PictureSize; + @memcpy(floats(d.lat, latents.len), latents); + const started = mtl.clock.seconds(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = d.queue.commandBuffer(); + const e = cb.compute(.serial); + var R = g.rows * 2; + var C = g.columns * 2; + unpack(e, d.get("md_unpack"), d.lat, d.bn_scale, d.bn_mean, d.t[0], g.rows, g.columns); + d.conv(e, d.post_quant, d.t[0], d.t[1], R, C, null); + d.conv(e, d.conv_in, d.t[1], d.t[0], R, C, null); + var x: usize = 0; + x = d.resnet(e, d.mid[0], x, R, C); + x = d.attention(e, x, R * C); + x = d.resnet(e, d.mid[1], x, R, C); + for (d.stages) |s| { + for (s.resnets) |r| x = d.resnet(e, r, x, R, C); + if (s.upsample) |u| { + R *= 2; + C *= 2; + const y = (x + 1) % d.t.len; + d.convUp(e, u, d.t[x], d.t[y], R, C); + x = y; + } + } + const n = (x + 1) % d.t.len; + d.norm(e, d.t[x], d.t[n], R * C, widths[3], d.norm_out, true); + d.conv(e, d.conv_out, d.t[n], d.out, R, C, null); + pixels(e, d.get("md_pixels"), d.out, d.rgb, R * C * 3); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("moss image decoder: a pass failed on the GPU: {s}", .{msg}); + return error.GpuFailure; + } + @memcpy(rgb, d.rgb.addr()[0..rgb.len]); + return .{ .wall = mtl.clock.seconds() - started, .gpu = cb.gpuSeconds() }; + } + + /// The last decode's output before the clamp [height, width, 3]. + pub fn output(d: *const Decoder) []const f32 { + return floats(d.out, d.width * d.height * 3); + } + + /// One residual block on t[x], its output in another of t; returns that one's index. + fn resnet(d: *const Decoder, e: mtl.ComputeEncoder, r: Resnet, x: usize, R: usize, C: usize) usize { + const n = (x + 1) % d.t.len; + const h = (x + 2) % d.t.len; + const o = (x + 3) % d.t.len; + const s = (x + 4) % d.t.len; + d.norm(e, d.t[x], d.t[n], R * C, r.conv1.cin, r.norm1, true); + d.conv(e, r.conv1, d.t[n], d.t[h], R, C, null); + d.norm(e, d.t[h], d.t[n], R * C, r.conv1.cout, r.norm2, true); + const skip = if (r.shortcut) |sc| blk: { + d.conv(e, sc, d.t[x], d.t[s], R, C, null); + break :blk d.t[s]; + } else d.t[x]; + d.conv(e, r.conv2, d.t[n], d.t[o], R, C, skip); + return o; + } + + /// The middle block's self-attention on t[x] (`n` positions): GroupNorm, q | k | v, one head, out + t[x]. + fn attention(d: *const Decoder, e: mtl.ComputeEncoder, x: usize, n: usize) usize { + const xn = (x + 1) % d.t.len; + const o = (x + 2) % d.t.len; + const W = widths[0]; + d.norm(e, d.t[x], d.t[xn], n, W, d.attn_norm, false); + d.conv(e, d.qkv, d.t[xn], d.qkv_rows, n, 1, null); + attend(e, d.get("md_attention"), d.qkv_rows, d.att, n); + d.conv(e, d.attn_out, d.att, d.t[o], n, 1, d.t[x]); + return o; + } + + /// y = GroupNorm(x) with `nm` (then SiLU when `silu`), x and y [pixels, channels]. + fn norm(d: *const Decoder, e: mtl.ComputeEncoder, x: Ref, y: Ref, n: usize, channels: usize, nm: Norm, silu: bool) void { + for ([_]bool{ false, true }) |centered| { + groupPart(e, d.get("md_group_part"), x, d.stats, d.part, n, channels, centered); + groupFinish(e, d.get("md_group_finish"), d.part, d.stats, n, channels, centered); + } + normed(e, d.get("md_norm"), x, d.stats, nm.g, nm.b, y, n, channels, silu); + } + + /// y = conv(x) + b (+ r) at R x C. + fn conv(d: *const Decoder, e: mtl.ComputeEncoder, c: Conv, x: Ref, y: Ref, R: usize, C: usize, r: ?Ref) void { + const s: ConvShape = .{ .rows = R, .cols = C, .cin = c.cin, .cout = c.cout, .ldw = c.ldw }; + if (c.taps == 9) d.wino(e, c, x, r, y, s) else convolve(e, d.get("md_conv"), x, c.w, c.b, r, y, s); + } + + /// md_wino_32 where md_wino's 64-channel groups would compute padding, else md_wino. + fn wino(d: *const Decoder, e: mtl.ComputeEncoder, c: Conv, x: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + if (wide(c.cout)) winograd(e, d.get("md_wino"), 64, x, c.w, c.b, r, y, s) else winograd(e, d.get("md_wino_32"), 32, x, c.w, c.b, r, y, s); + } + + /// y = conv(x read nearest from R/2 x C/2) + b at R x C. + fn convUp(d: *const Decoder, e: mtl.ComputeEncoder, c: Conv, x: Ref, y: Ref, R: usize, C: usize) void { + d.wino(e, c, x, null, y, .{ .rows = R, .cols = C, .cin = c.cin, .cout = c.cout, .ldw = c.ldw, .up = true }); + } +}; + +/// md_unpack: latents x [rows * columns, 128] to z [2 rows, 2 columns, 32], the batch norm undone (x * s + m). +pub fn unpack(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, s: Ref, m: Ref, z: Ref, rows: usize, columns: usize) void { + bind(e, p, &.{ x, s, m, z }, [2]u32{ @intCast(rows), @intCast(columns) }); + e.dispatchThreads(mtl.Size.of(latent_channels, rows * columns * 4, 1), mtl.Size.of(32, 8, 1)); +} + +/// A convolution's size: the output picture, its channels, the weights' row length, and a 3x3's read from half size. +pub const ConvShape = struct { rows: usize, cols: usize, cin: usize, cout: usize, ldw: usize, up: bool = false }; + +fn convParams(s: ConvShape, add: bool) extern struct { rows: u32, cols: u32, cin: u32, cout: u32, ldw: u32, up: u32, add: u32 } { + return .{ .rows = @intCast(s.rows), .cols = @intCast(s.cols), .cin = @intCast(s.cin), .cout = @intCast(s.cout), .ldw = @intCast(s.ldw), .up = @intFromBool(s.up), .add = @intFromBool(add) }; +} + +/// md_conv: y = x w + b (+ r when given), a 1x1 convolution of x [rows * cols, cin]. +pub fn convolve(e: mtl.ComputeEncoder, p: mtl.Pipeline, x: Ref, w: Ref, b: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + const params = convParams(s, r != null); + bind(e, p, &.{ x, w, b, r orelse y, y }, params); + e.dispatchGroups(mtl.Size.of(s.ldw / 64, (s.rows * s.cols + 63) / 64, 1), mtl.Size.of(128, 1, 1)); +} + +/// True when md_wino's 64 channels a threadgroup compute no padding past `cout`'s last 32. +pub fn wide(cout: usize) bool { + return cout % 64 == 0 or cout % 64 > 32; +} + +/// md_wino (`group` 64) or md_wino_32 (32): y = conv3x3(x) + b (+ r when given) by Winograd F(2x2, 3x3). +pub fn winograd(e: mtl.ComputeEncoder, p: mtl.Pipeline, group: usize, x: Ref, u: Ref, b: Ref, r: ?Ref, y: Ref, s: ConvShape) void { + const params = convParams(s, r != null); + bind(e, p, &.{ x, u, b, r orelse y, y }, params); + e.dispatchGroups(mtl.Size.of((s.cout + group - 1) / group, (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; +} + +/// md_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 of x [n, channels]: 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 }; +} + +/// md_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 { + 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)); +} + +/// md_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)); +} + +/// md_attention: o [rows, 384] = softmax(q k^T / sqrt(384)) v of qkv rows [rows, 1152], 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, widths[0])) }; + bind(e, p, &.{ qkv, o }, params); + e.dispatchGroups(mtl.Size.of((rows + 7) / 8, 1, 1), mtl.Size.of(128, 1, 1)); +} + +/// md_pixels: rgb = the n 8-bit values of y. +pub fn pixels(e: mtl.ComputeEncoder, p: mtl.Pipeline, y: Ref, rgb: Ref, n: usize) void { + bind(e, p, &.{ y, rgb }, @as(u32, @intCast(n))); + e.dispatchThreads(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1)); +} + +fn floats(r: Ref, n: usize) []f32 { + return @as([*]f32, @ptrCast(@alignCast(r.addr())))[0..n]; +} + +/// Lays the decoder'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, + 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 tensor(p: *Packer, name: []const u8, shape: []const usize) ![]const f32 { + const t = try p.c.get(name); + if (t.dtype != .f32 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) { + std.log.err("moss image decoder: {s} is {t} {any}; the decoder reads f32 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + } + p.read += 1; + return t.host(f32); + } + + /// A copied vector `name` of n. + fn copy(p: *Packer, name: []const u8, n: usize) !Ref { + const v = try p.tensor(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): a 1x1 as [cin, ldw], a 3x3 as G g G^T [16, cin, ldw]. + fn conv(p: *Packer, name: []const u8, cin: usize, cout: usize, taps: usize) !Conv { + var buf: [96]u8 = undefined; + const k: usize = if (taps == 9) 3 else 1; + const ldw = std.mem.alignForward(usize, cout, 64); + const points: usize = if (taps == 9) 16 else 1; + const w = try p.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{name}), &.{ cout, cin, k, k }); + const r = p.alloc(points * cin * ldw); + if (p.out) |out| { + const dst = out[r.off / 4 ..][0 .. points * cin * ldw]; + @memset(dst, 0); + for (0..cout) |o| for (0..cin) |i| { + const g = w[(o * cin + i) * taps ..][0..taps]; + if (taps == 9) { + for (transform(g[0..9]), 0..) |v, pt| dst[(pt * cin + i) * ldw + o] = v; + } else dst[i * ldw + o] = g[0]; + }; + } + 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, .taps = taps }; + } + + /// 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 ldw = std.mem.alignForward(usize, n, 64); + const r = p.alloc(width * ldw); + 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 * ldw]; + for (0..width) |o| for (0..width) |i| { + dst[i * ldw + part * width + o] = w[o * width + i]; + }; + } + } + const b = p.alloc(n); + for (names_, 0..) |name, i| { + const v = try p.tensor(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 = ldw, .taps = 1 }; + } + + 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, 9), + .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, 9), + .shortcut = if (cin != cout) try p.conv(try std.fmt.bufPrint(&buf, "{s}.conv_shortcut", .{name}), cin, cout, 1) else null, + }; + } + + /// The checkpoint's decoder-side tensors: the decoder's, post_quant_conv's and the batch norm's two statistics. + fn expected(_: *const Packer, c: *const ck.Checkpoint) usize { + var n: usize = 0; + var it = c.tensors.keyIterator(); + while (it.next()) |k| { + if (std.mem.startsWith(u8, k.*, "decoder.") or std.mem.startsWith(u8, k.*, "post_quant_conv.") or + std.mem.eql(u8, k.*, "bn.running_mean") or std.mem.eql(u8, k.*, "bn.running_var")) n += 1; + } + return n; + } +}; + +test "only FLUX.2's released decoder geometry passes" { + const a = std.testing.allocator; + const sample = + \\{"_class_name": "AutoencoderKLFlux2", "act_fn": "silu", "batch_norm_eps": 0.0001, "decoder_block_out_channels": [96, 192, 384, 384], + \\ "latent_channels": 32, "layers_per_block": 2, "mid_block_add_attention": true, "norm_num_groups": 32, "out_channels": 3, + \\ "patch_size": [2, 2], "use_post_quant_conv": true} + ; + try check(a, sample); + const narrow = try std.mem.replaceOwned(u8, a, sample, "[96, 192, 384, 384]", "[128, 256, 512, 512]"); + defer a.free(narrow); + try std.testing.expectError(error.UnsupportedGeometry, check(a, narrow)); + const no_attn = try std.mem.replaceOwned(u8, a, sample, "\"mid_block_add_attention\": true", "\"mid_block_add_attention\": false"); + defer a.free(no_attn); + try std.testing.expectError(error.UnsupportedGeometry, check(a, no_attn)); + const other = try std.mem.replaceOwned(u8, a, sample, "AutoencoderKLFlux2", "AutoencoderKL"); + defer a.free(other); + try std.testing.expectError(error.NotFlux2Vae, check(a, other)); +} diff --git a/zig/src/families/moss_image/moss_image.zig b/zig/src/families/moss_image/moss_image.zig index a4a96a1e0d..72108db490 100644 --- a/zig/src/families/moss_image/moss_image.zig +++ b/zig/src/families/moss_image/moss_image.zig @@ -10,9 +10,12 @@ pub const encoder_weights = @import("encoder_weights.zig"); pub const encoder = @import("encoder.zig"); pub const Encoder = encoder.Encoder; pub const reference = @import("reference.zig"); +pub const decoder = @import("decoder.zig"); +pub const png = @import("png.zig"); pub const Model = model.Model; pub const Prompt = model.Prompt; pub const Timing = model.Timing; +pub const Decoder = decoder.Decoder; test { _ = config; @@ -25,4 +28,6 @@ test { _ = encoder_weights; _ = encoder; _ = reference; + _ = decoder; + _ = png; } diff --git a/zig/src/families/moss_image/png.zig b/zig/src/families/moss_image/png.zig new file mode 100644 index 0000000000..2c68ed24a6 --- /dev/null +++ b/zig/src/families/moss_image/png.zig @@ -0,0 +1,55 @@ +//! An 8-bit RGB picture as a PNG: one IDAT of zlib-compressed rows, each with filter 0 (none). +const std = @import("std"); +const flate = std.compress.flate; + +/// The PNG of `rgb` [height, width, 3], allocated with `a`. +pub fn encode(a: std.mem.Allocator, rgb: []const u8, width: usize, height: usize) ![]u8 { + if (width == 0 or height == 0 or rgb.len != width * height * 3) return error.PictureSize; + if (width > std.math.maxInt(u31) or height > std.math.maxInt(u31)) return error.PictureSize; + var z: std.Io.Writer.Allocating = try .initCapacity(a, rgb.len / 2 + 64); + defer z.deinit(); + var window: [flate.max_window_len]u8 = undefined; + var c = try flate.Compress.init(&z.writer, &window, .zlib, .default); + const row = width * 3; + for (0..height) |y| { + try c.writer.writeByte(0); + try c.writer.writeAll(rgb[y * row ..][0..row]); + } + try c.finish(); + var out: std.Io.Writer.Allocating = .init(a); + errdefer out.deinit(); + try out.writer.writeAll("\x89PNG\r\n\x1a\n"); + var header: [13]u8 = undefined; + std.mem.writeInt(u32, header[0..4], @intCast(width), .big); + std.mem.writeInt(u32, header[4..8], @intCast(height), .big); + header[8..13].* = .{ 8, 2, 0, 0, 0 }; + try chunk(&out.writer, "IHDR", &header); + try chunk(&out.writer, "IDAT", z.written()); + try chunk(&out.writer, "IEND", ""); + return out.toOwnedSlice(); +} + +fn chunk(w: *std.Io.Writer, kind: *const [4]u8, data: []const u8) !void { + try w.writeInt(u32, @intCast(data.len), .big); + try w.writeAll(kind); + try w.writeAll(data); + var crc = std.hash.Crc32.init(); + crc.update(kind); + crc.update(data); + try w.writeInt(u32, crc.final(), .big); +} + +test "a PNG of a small picture decodes back to its pixels" { + const a = std.testing.allocator; + var rgb: [5 * 3 * 3]u8 = undefined; + for (&rgb, 0..) |*v, i| v.* = @truncate(i * 37); + const png = try encode(a, &rgb, 5, 3); + defer a.free(png); + const picture = @import("../../core/picture/png.zig"); + const s = try picture.decode(a, png); + defer a.free(s.data); + try std.testing.expectEqual(@as(usize, 5), s.width); + try std.testing.expectEqual(@as(usize, 3), s.height); + for (0..15) |i| try std.testing.expectEqualSlices(u8, rgb[i * 3 ..][0..3], s.data[i * s.channels ..][0..3]); + try std.testing.expectError(error.PictureSize, encode(a, rgb[0..10], 5, 3)); +} diff --git a/zig/src/families/moss_image/reference.zig b/zig/src/families/moss_image/reference.zig index c2e8e81ed8..0e80be3236 100644 --- a/zig/src/families/moss_image/reference.zig +++ b/zig/src/families/moss_image/reference.zig @@ -1,11 +1,11 @@ -//! A run's rows against float32 ones, gated by the bf16 reference's own distance from them. +//! A run's rows against exact ones, gated by the reference implementation's own distance from them. const std = @import("std"); const host = @import("host.zig"); pub const Match = struct { max_abs: f64, cosine: f64, worst_row: f64, worst_at: usize }; /// Largest absolute difference, cosine, and the largest relative L2 error of one row of `width` (and its row). -pub fn match(comptime T: type, ours: []const T, theirs: []const f32, width: usize) Match { +pub fn match(comptime T: type, ours: []const T, theirs: anytype, width: usize) Match { return measure(T, ours, theirs, width, false); } @@ -14,7 +14,7 @@ pub fn rounding(theirs: []const f32, width: usize) Match { return measure(f32, theirs, theirs, width, true); } -fn measure(comptime T: type, ours: []const T, theirs: []const f32, width: usize, comptime rounded: bool) Match { +fn measure(comptime T: type, ours: []const T, theirs: anytype, width: usize, comptime rounded: bool) Match { var dot: f64 = 0; var aa: f64 = 0; var bb: f64 = 0; @@ -43,15 +43,20 @@ fn measure(comptime T: type, ours: []const T, theirs: []const f32, width: usize, return m; } -/// How far a result may be from float32: this many times the bf16 reference's distance, or bf16 rounding's. +/// How far a result may be from exact: this many times the reference's distance (or bf16 rounding's, for bf16). pub const factor = 2.0; pub const Verdict = struct { pass: bool, cosine_err: f64, cosine_limit: f64, row: f64, row_at: usize, row_limit: f64 }; /// Ours against float32: its cosine error and worst row each within `factor` of the bf16 reference's or rounding's. pub fn gate(ours: Match, reference: Match, rounded: Match) Verdict { - const cosine_limit = factor * @max(1 - reference.cosine, 1 - rounded.cosine); - const row_limit = factor * @max(reference.worst_row, rounded.worst_row); + return within(ours, .{ .max_abs = 0, .cosine = @min(reference.cosine, rounded.cosine), .worst_row = @max(reference.worst_row, rounded.worst_row), .worst_at = 0 }); +} + +/// Ours against exact rows: its cosine error and worst row each within `factor` of the reference's. +pub fn within(ours: Match, reference: Match) Verdict { + const cosine_limit = factor * (1 - reference.cosine); + const row_limit = factor * reference.worst_row; const cosine_err = 1 - ours.cosine; return .{ .pass = cosine_err <= cosine_limit and ours.worst_row <= row_limit, @@ -116,3 +121,22 @@ test "a different repeat fails even when every gate passes" { const v: Verdict = .{ .pass = true, .cosine_err = 0, .cosine_limit = 1, .row = 0, .row_at = 0, .row_limit = 1 }; try std.testing.expect(!accept(&.{ v, v }, false)); } + +test "against float64, a row further off than the reference allows fails" { + var exact: [cols * n]f64 = undefined; + var f: [cols * n]f32 = undefined; + synthetic(&f, 3); + for (&exact, f) |*e, x| e.* = x; + var theirs: [cols * n]f32 = undefined; + var ours: [cols * n]f32 = undefined; + for (&theirs, &ours, exact, 0..) |*t, *o, x, i| { + const s: f64 = if (i % 2 == 0) 1 else -1; + t.* = @floatCast(x * (1 + 1e-5 * s)); + o.* = @floatCast(x * (1 - 1e-5 * s)); + } + try std.testing.expect(within(match(f32, &ours, &exact, cols), match(f32, &theirs, &exact, cols)).pass); + for (ours[3 * cols ..][0..cols], exact[3 * cols ..][0..cols]) |*o, x| o.* = @floatCast(x * (1 + 1e-4)); + const v = within(match(f32, &ours, &exact, cols), match(f32, &theirs, &exact, cols)); + try std.testing.expect(!v.pass); + try std.testing.expectEqual(@as(usize, 3), v.row_at); +} diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 796ed74d6b..9a85c19b50 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -99,6 +99,8 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La } else if (std.mem.startsWith(u8, d.name, "moss image gemm ")) { text = sources.core_qwen_vision; frags = true; + } else if (std.mem.startsWith(u8, d.name, "moss image decoder ")) { + text = sources.moss_image_decoder; } else if (std.mem.startsWith(u8, d.name, "moss image ")) { text = sources.moss_image; frags = true; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index b32e2baa9e..999a332645 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 moss_image_decoder = try @import("metal_moss_image_decoder.zig").Suite.create(a, device, queue, &kernels, o.only); + defer moss_image_decoder.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { diff --git a/zig/tests/contract/metal_moss_image_decoder.zig b/zig/tests/contract/metal_moss_image_decoder.zig new file mode 100644 index 0000000000..ad50974619 --- /dev/null +++ b/zig/tests/contract/metal_moss_image_decoder.zig @@ -0,0 +1,303 @@ +//! MOSS-Image's image decoder kernels under the contract, and md_pixels; controls each break one gate. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); +const mi = @import("metal_moss_image_common.zig"); +const st = @import("metal_moss_image_decoder_steps.zig"); + +const c = mi.c; +const g = mi.g; +const Range = mi.Range; +const Ends = mi.Ends; +const dec = tf.moss_image.decoder; +const ref = st.ref; +const uniform = st.uniform; +const plant = st.plant; +const nan = st.nan; +const inf = st.inf; +const rw = @import("metal_moss_image_decoder_rows.zig"); +const wn = @import("metal_moss_image_decoder_wino.zig"); +const Wino = wn.Wino; +const Unpack = st.Unpack; +const Conv = st.Conv; +const Norm = st.Norm; +const conv_forms = st.conv_forms; +const GroupPart = rw.GroupPart; +const GroupFinish = rw.GroupFinish; +const Attention = rw.Attention; + +// --- md_pixels: the 8-bit picture, as the pipeline's clamp, * 0.5 + 0.5, * 255 and round to even make it ----------- + +pub const Pixels = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + y: []f32 = &.{}, + ins: [1]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + /// Rows: pixels; n: the 3 channels (k unused). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = mi.noSums, + .out = .raw, + .nan = .{ .stated = "+Inf gives 255 and -Inf 0; a NaN's byte is unspecified, as numpy's cast of NaN to uint8 is" }, + .tiles = .{ .{ .size = 256 }, .{}, .{} }, + .served = &.{ .{ .rows = 1024 * 1024, .n = 3, .k = 1 }, .{ .rows = 1200 * 880, .n = 3, .k = 1 }, .{ .rows = 1344 * 768, .n = 3, .k = 1 } }, + }; + } + + pub fn deinit(t: *Pixels) void { + t.a.free(t.y); + t.rig.deinit(); + } + pub fn kernel(t: *Pixels, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Pixels); + + /// The byte of y by the same fp32 steps. + fn byte(y: f32) u8 { + const v = std.math.clamp(std.math.clamp(y, -1, 1) * 0.5 + 0.5, 0, 1) * 255; + const r = @round(v); + const tie = r - v == 0.5 and @mod(r, 2) == 1; + return @intFromFloat(if (tie) r - 1 else r); + } + + /// The decoder outputs whose * 255 lands exactly halfway between two bytes, one for each byte up to 254. + fn ties(out: *std.ArrayList(f32), a: std.mem.Allocator) !void { + for (0..255) |k| { + const want: f32 = @as(f32, @floatFromInt(k)) + 0.5; + var v: f32 = want / 255; + for (0..64) |_| { + if (v * 255 == want) break; + v = if (v * 255 < want) std.math.nextAfter(f32, v, 2) else std.math.nextAfter(f32, v, -2); + } + if (v * 255 != want) continue; + const y = v * 2 - 1; + if (std.math.clamp(y, -1, 1) * 0.5 + 0.5 == v) try out.append(a, y); + } + } + + pub fn prepare(t: *Pixels, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n != 3) return false; + t.rig.reset(); + t.a.free(t.y); + t.shape = shape; + t.y = try t.a.alloc(f32, shape.rows * 3); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + uniform(rng, t.y, -1.5, 1.5); + var tie: std.ArrayList(f32) = .empty; + defer tie.deinit(t.a); + try ties(&tie, t.a); + for (tie.items, 0..) |v, i| t.y[(i * 37) % t.y.len] = v; + if (fill == .nonfinite) plant(t.y, 3); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.y)); + t.out = try t.rig.output(t.y.len); + return true; + } + + pub fn run(t: *Pixels) !void { + const p = t.rig.begin(); + dec.pixels(p.e, t.pipe, ref(t.ins[0]), ref(t.out), t.y.len); + try t.rig.finish(p); + } + + pub fn count(t: *Pixels) usize { + return t.y.len; + } + + pub fn reference(t: *Pixels, want: []f64, mag: []f64) !void { + for (want, mag, t.y) |*w, *m, y| { + w.* = if (std.math.isNan(y)) 0 else @floatFromInt(byte(y)); + m.* = 0; + } + } + + /// Every byte equals the fp32 steps' (NaN's aside), ties to even included. + pub fn check(t: *Pixels, _: g.Gate) !?[]const u8 { + const got = t.rig.bytes(t.out, t.y.len); + for (got, t.y) |b, y| { + if (std.math.isNan(y)) continue; + if (b != byte(y)) return "a byte differs from the fp32 steps' (clamp, * 0.5 + 0.5, * 255, round to even)"; + } + return null; + } +}; + +// --- the suite: the served source's pipelines, and each control's own copy of the source ------------------------ + +/// 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 controls = [_]Control{ + .{ .kernel = "md_unpack", .gate = .shapes, .what = "swaps the patch's axes", .from = "const uint k = c * 4 + (py % 2) * 2 + px % 2;", .to = "const uint k = c * 4 + (px % 2) * 2 + py % 2;" }, + .{ .kernel = "md_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 = "md_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 = "md_wino_32 3x3", .gate = .shapes, .what = "drops a term of the output transform", .from = "s[j][1] = (m1 - m2) - m3;", .to = "s[j][1] = (m1 - m2);" }, + .{ .kernel = "md_norm", .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 = "md_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 = "md_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 = "md_attention", .gate = .shapes, .what = "keeps stale maxima (no rescale)", .from = " if (rescale) {", .to = " if (false) {" }, + .{ .kernel = "md_pixels", .gate = .shapes, .what = "rounds halves up", .from = "rgb[i] = uchar(metal::rint(v * 255.0f));", .to = "rgb[i] = uchar(floor(v * 255.0f + 0.5f));" }, +}; + +/// The decoder'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 served = try tf.frags.source(device, a, sources.moss_image_decoder); + defer a.free(served); + const text = if (ctl) |x| blk: { + if (std.mem.count(u8, served, x.from) != 1) return error.ControlTextNotFoundOnce; + break :blk try std.mem.replaceOwned(u8, a, served, x.from, x.to); + } else try a.dupe(u8, served); + defer a.free(text); + return mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); +} + +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, + unpacks: std.ArrayList(*Unpack) = .empty, + convs: std.ArrayList(*Conv) = .empty, + winos: std.ArrayList(*Wino) = .empty, + norms: std.ArrayList(*Norm) = .empty, + parts: std.ArrayList(*GroupPart) = .empty, + finishes: std.ArrayList(*GroupFinish) = .empty, + attns: std.ArrayList(*Attention) = .empty, + pixels: std.ArrayList(*Pixels) = .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, list: *std.ArrayList(*T), value: T) !*T { + const t = try s.a.create(T); + errdefer s.a.destroy(t); + t.* = value; + try list.append(s.a, t); + 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 want = struct { + fn is(x: ?Control, kernel: []const u8) bool { + return if (x) |y| std.mem.eql(u8, y.kernel, kernel) else true; + } + fn keep(o: ?[]const u8, name: []const u8) bool { + return if (o) |needle| std.mem.indexOf(u8, name, needle) != null else true; + } + }; + { + const name = try s.label("metal moss image decoder md_unpack{s}", .{suffix}); + if (want.is(ctl, "md_unpack") and want.keep(only, name)) { + const t = try s.make(Unpack, &s.unpacks, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_unpack"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + for (conv_forms) |f| { + const kind = try s.label("md_conv {s}", .{f.label}); + const name = try s.label("metal moss image decoder {s}{s}", .{ kind, suffix }); + if (!want.is(ctl, kind) or !want.keep(only, name)) continue; + const t = try s.make(Conv, &s.convs, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_conv"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for (wn.forms) |f| { + const name = try s.label("metal moss image decoder {s}{s}", .{ f.label, suffix }); + if (!want.is(ctl, f.label) or !want.keep(only, name)) continue; + const t = try s.make(Wino, &s.winos, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, if (f.group == 64) "md_wino" else "md_wino_32"), .form = f, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ true, false }) |silu| { + const name = try s.label("metal moss image decoder md_norm{s}{s}", .{ if (silu) " silu" else "", suffix }); + if (!want.is(ctl, "md_norm") or !want.keep(only, name) or (ctl != null and !silu)) continue; + const t = try s.make(Norm, &s.norms, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_norm"), .silu = silu, .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + for ([_]bool{ false, true }) |centered| { + const name = try s.label("metal moss image decoder md_group_part{s}{s}", .{ if (centered) " centered" else "", suffix }); + if (!want.is(ctl, "md_group_part") or !want.keep(only, name) or (ctl != null and centered)) continue; + const t = try s.make(GroupPart, &s.parts, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_group_part"), .centered = centered, .name = name }); + try out.append(s.a, t.kernel(gate)); + } + { + const name = try s.label("metal moss image decoder md_group_finish{s}", .{suffix}); + if (want.is(ctl, "md_group_finish") and want.keep(only, name)) { + const t = try s.make(GroupFinish, &s.finishes, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_group_finish"), .name = name, .ends = ends }); + try out.append(s.a, t.kernel(gate)); + } + } + { + const name = try s.label("metal moss image decoder md_attention{s}", .{suffix}); + if (want.is(ctl, "md_attention") and want.keep(only, name)) { + const t = try s.make(Attention, &s.attns, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_attention"), .name = name }); + try out.append(s.a, t.kernel(gate)); + } + } + { + const name = try s.label("metal moss image decoder md_pixels{s}", .{suffix}); + if (want.is(ctl, "md_pixels") and want.keep(only, name)) { + const t = try s.make(Pixels, &s.pixels, .{ .a = s.a, .rig = rig, .pipe = try s.pipeline(device, lib, "md_pixels"), .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 moss image decoder: the family runs on the M5 alone; its kernels unchecked here\n", .{}); + return s; + } + for ([_]?Control{null} ++ controls) |x| { + const lib = try compile(a, device, x); + try s.libs.append(a, lib); + try s.add(device, queue, lib, x, out, only); + } + return s; + } + + pub fn destroy(s: *Suite) void { + const a = s.a; + inline for (.{ &s.unpacks, &s.convs, &s.winos, &s.norms, &s.parts, &s.finishes, &s.attns, &s.pixels }) |list| { + for (list.items) |t| { + t.deinit(); + a.destroy(t); + } + list.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_moss_image_decoder_rows.zig b/zig/tests/contract/metal_moss_image_decoder_rows.zig new file mode 100644 index 0000000000..04900d83a1 --- /dev/null +++ b/zig/tests/contract/metal_moss_image_decoder_rows.zig @@ -0,0 +1,437 @@ +//! md_group_part, md_group_finish and md_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_moss_image_common.zig"); +const st = @import("metal_moss_image_decoder_steps.zig"); + +const c = mi.c; +const g = mi.g; +const Range = mi.Range; +const Ends = mi.Ends; +const dec = tf.moss_image.decoder; +const ref = st.ref; +const uniform = st.uniform; +const plant = st.plant; +const nan = st.nan; +const inf = st.inf; + +// --- md_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 / 256 pixels of the group's channels in turn, then the simdgroup tree and the 8 simdgroups. + fn depth(s: c.Shape) usize { + return (@min(s.k, dec.chunk) + 255) / 256 * (s.n / 32) + 5 + 8 + 2; + } + + /// 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 = 96, .step = 32, .min = 32 }, .{ .size = dec.chunk } }, + .served = &.{ .{ .rows = 1, .n = 384, .k = 16384 }, .{ .rows = 1, .n = 384, .k = 16500 }, .{ .rows = 1, .n = 192, .k = 262144 }, .{ .rows = 1, .n = 96, .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]); + }; + } +}; + +// --- md_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, 384 / 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 = 4 }, .{ .rows = 1, .n = 32, .k = 5 }, .{ .rows = 1, .n = 32, .k = 64 }, .{ .rows = 1, .n = 32, .k = 256 } }, + }; + } + + 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(), 384, 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() * 12); + 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() * 12); + 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); + } + } +}; + +// --- md_attention: softmax(q k^T / sqrt(384)) v, one head of 384, 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 = 384; + const ld = 3 * D; + + /// The scores' 384-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_moss_image_decoder_steps.zig b/zig/tests/contract/metal_moss_image_decoder_steps.zig new file mode 100644 index 0000000000..2c9285ed5d --- /dev/null +++ b/zig/tests/contract/metal_moss_image_decoder_steps.zig @@ -0,0 +1,378 @@ +//! md_unpack, md_conv (1x1) and md_norm under the contract: 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 mi = @import("metal_moss_image_common.zig"); + +const c = mi.c; +const g = mi.g; +const Range = mi.Range; +const Ends = mi.Ends; +const dec = tf.moss_image.decoder; + +pub fn ref(r: rig_.Ref) mi.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) { + x[3 * apart] = inf; + x[3 * apart + 1] = -inf; + } +} + +// --- md_unpack: z = x s + m, each 2 x 2 patch of a latent row unpacked --------------------------------------------- + +pub const Unpack = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + name: []const u8, + shape: c.Shape = undefined, + x: []f32 = &.{}, + s: [128]f32 = undefined, + m: [128]f32 = undefined, + ins: [3]rig_.Ref = undefined, + out: rig_.Ref = undefined, + ends: Ends, + + /// Rows: the latent grid's rows; n: its columns (k unused). + pub fn spec(name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = mi.noSums, + .ops = 2, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{ .size = 8 }, .{} }, + .served = &.{ .{ .rows = 64, .n = 64, .k = 1 }, .{ .rows = 55, .n = 75, .k = 1 }, .{ .rows = 48, .n = 84, .k = 1 } }, + }; + } + + pub fn deinit(t: *Unpack) void { + t.a.free(t.x); + t.ends.deinit(); + t.rig.deinit(); + } + pub fn kernel(t: *Unpack, control: ?g.Gate) g.Kernel { + return .{ .spec = spec(t.name), .ctx = t, .vt = &vt, .control = control }; + } + const vt = g.hooks(Unpack); + + pub fn prepare(t: *Unpack, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + t.a.free(t.x); + t.shape = shape; + t.x = try t.a.alloc(f32, shape.rows * shape.n * 128); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + uniform(rng, t.x, -3, 3); + uniform(rng, &t.s, 0.5, 2); + uniform(rng, &t.m, -1, 1); + if (fill == .nonfinite) plant(t.x, 128); + t.ins[0] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.ins[1] = try t.rig.input(std.mem.sliceAsBytes(&t.s)); + t.ins[2] = try t.rig.input(std.mem.sliceAsBytes(&t.m)); + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Unpack) !void { + const p = t.rig.begin(); + dec.unpack(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: *Unpack) usize { + return t.shape.rows * t.shape.n * 128; + } + + pub fn reference(t: *Unpack, want: []f64, mag: []f64) !void { + try t.ends.size(t.count()); + const w2 = t.shape.n * 2; + for (0..t.shape.rows * 2 * w2) |px| for (0..32) |ch| { + const py = px / w2; + const pxx = px % w2; + const k = ch * 4 + (py % 2) * 2 + pxx % 2; + const x: f64 = t.x[((py / 2) * t.shape.n + pxx / 2) * 128 + k]; + const r = Range.of(x * t.s[k]).fp32(0.5).add(Range.of(t.m[k])).fp32(0.5); + t.ends.put(want, mag, px * 32 + ch, x * t.s[k] + t.m[k], r); + }; + } + pub fn bounds(t: *Unpack, want: []const f64, _: []const f64, allow: []f64) !void { + t.ends.bounds(want, allow); + } + pub fn intervals(t: *Unpack, low: []f64, high: []f64) !void { + t.ends.intervals(low, high); + } +}; + +// --- md_conv: a 1x1 convolution, fp32 sums, + bias (+ a residual) ------------------------------------------------ + +/// A 1x1 form the decoder launches: whether it adds a residual, and its served shapes. +pub const Form = struct { label: []const u8, add: bool = false, served: []const c.Shape }; + +pub const conv_forms = [_]Form{ + .{ .label = "1x1", .served = &.{ .{ .rows = 4100, .n = 1152, .k = 384 }, .{ .rows = 517, .n = 32, .k = 32 }, .{ .rows = 517, .n = 192, .k = 384 }, .{ .rows = 517, .n = 96, .k = 192 } } }, + .{ .label = "1x1 + residual", .add = true, .served = &.{.{ .rows = 4100, .n = 384, .k = 384 }} }, +}; + +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 cin products, then the bias and the residual. + fn depth(s: c.Shape) usize { + return s.k + 2; + } + + /// Rows: pixels; 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 = 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); + } + + 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(); + t.x = try t.a.alloc(f32, shape.rows * shape.k); + t.w = try t.a.alloc(f32, shape.k * L); + t.b = try t.a.alloc(f32, shape.n); + t.r = try t.a.alloc(f32, if (t.form.add) shape.rows * 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..shape.rows) |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..shape.k) |k| for (0..shape.n) |o| { + t.w[k * L + o] = if (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..shape.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: dec.ConvShape = .{ .rows = t.shape.rows, .cols = 1, .cin = t.shape.k, .cout = t.shape.n, .ldw = t.ldw() }; + dec.convolve(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.shape.rows * t.shape.n; + } + + /// Each output's products, in the kernel's order over the input channels. + fn terms(t: *const Conv, p: usize, o: usize, out: []f64) void { + for (out, 0..) |*v, ci| v.* = @as(f64, t.x[p * t.shape.k + ci]) * t.w[ci * t.ldw() + o]; + } + + pub fn reference(t: *Conv, want: []f64, mag: []f64) !void { + const buf = try t.a.alloc(f64, t.shape.k); + defer t.a.free(buf); + for (0..t.shape.rows) |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.shape.k); + defer t.a.free(buf); + for (0..t.shape.rows) |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); + }; + } +}; + +// --- md_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) c.Spec { + return .{ + .name = name, + .accum = .none, + .depth = mi.noSums, + .ops = 4 + mi.exp_ulps + mi.div_ulps + 1, + .out = .f32, + .tiles = .{ .{ .size = 8 }, .{ .size = 96, .step = 32, .min = 32 }, .{} }, + .served = &.{ .{ .rows = 4100, .n = 384, .k = 1 }, .{ .rows = 4100, .n = 192, .k = 1 }, .{ .rows = 4100, .n = 96, .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), .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(); + dec.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(mi.exp_ulps); + vr = vr.mul(e.add(Range.of(1)).fp32(0.5).recip()).fp32(mi.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); + } +}; diff --git a/zig/tests/contract/metal_moss_image_decoder_wino.zig b/zig/tests/contract/metal_moss_image_decoder_wino.zig new file mode 100644 index 0000000000..6ccb0f3068 --- /dev/null +++ b/zig/tests/contract/metal_moss_image_decoder_wino.zig @@ -0,0 +1,268 @@ +//! md_wino and md_wino_32 under the contract: Winograd F(2x2, 3x3) against the float64 convolution. +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_moss_image_common.zig"); +const st = @import("metal_moss_image_decoder_steps.zig"); + +const c = mi.c; +const g = mi.g; +const dec = tf.moss_image.decoder; +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 decoder launches through one instance: its group width, read from half size or not, a residual. +pub const Form = struct { label: []const u8, group: usize, up: bool = false, add: bool = false, served: []const c.Shape }; + +/// The pictures here are 8 x 18 (the decoder's up to 1024 x 1024): the sums' depth is the input channels alone. +pub const forms = [_]Form{ + .{ .label = "md_wino 3x3", .group = 64, .served = &.{ .{ .rows = 8, .n = 384, .k = 384 }, .{ .rows = 8, .n = 384, .k = 32 }, .{ .rows = 8, .n = 192, .k = 384 }, .{ .rows = 8, .n = 192, .k = 192 } } }, + .{ .label = "md_wino 3x3 upsampled", .group = 64, .up = true, .served = &.{ .{ .rows = 8, .n = 384, .k = 384 }, .{ .rows = 8, .n = 192, .k = 192 } } }, + .{ .label = "md_wino 3x3 + residual", .group = 64, .add = true, .served = &.{ .{ .rows = 8, .n = 384, .k = 384 }, .{ .rows = 8, .n = 192, .k = 192 } } }, + .{ .label = "md_wino_32 3x3", .group = 32, .served = &.{ .{ .rows = 8, .n = 96, .k = 192 }, .{ .rows = 8, .n = 96, .k = 96 }, .{ .rows = 8, .n = 3, .k = 96 } } }, + .{ .label = "md_wino_32 3x3 + residual", .group = 32, .add = true, .served = &.{.{ .rows = 8, .n = 96, .k = 96 }} }, +}; + +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 = f.group }, .{ .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; + } + fn inCols(t: *const Wino) usize { + return if (t.form.up) cols / 2 else 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 = if (t.form.up) t.pixels() / 4 else 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 (dec.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: dec.ConvShape = .{ .rows = t.shape.rows, .cols = cols, .cin = t.shape.k, .cout = t.shape.n, .ldw = t.ldw(), .up = t.form.up }; + dec.winograd(p.e, t.pipe, t.form.group, 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); + const p = if (t.form.up) (uy / 2) * t.inCols() + ux / 2 else uy * cols + ux; + return t.x[p * 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/moss_image.zig b/zig/tests/moss_image.zig index 1bd6d7be08..66cf5ad535 100644 --- a/zig/tests/moss_image.zig +++ b/zig/tests/moss_image.zig @@ -8,6 +8,7 @@ const Metal = struct { pub const mi = tf.moss_image; pub const st = tf.safetensors; pub const seconds = mtl.clock.seconds; + pub const Picture = struct { png: []const u8, load_s: f64, wall: f64, gpu: f64 }; pub const Device = struct { pool: mtl.objc.Pool, @@ -45,6 +46,17 @@ const Metal = struct { } return mtl.clock.seconds() - encoding; } + + /// `latents` as a PNG by the checkpoint's vae/ decoder, freed before this returns; its load and decode times. + pub fn decode(g: *Device, a: std.mem.Allocator, io: std.Io, checkpoint: []const u8, latents: []const f32, width: usize, height: usize) !Picture { + const loading = mtl.clock.seconds(); + const decoder = try mi.Decoder.load(std.heap.page_allocator, io, g.device, try std.fs.path.join(a, &.{ checkpoint, "vae" }), width, height); + defer decoder.deinit(); + const load_s = mtl.clock.seconds() - loading; + const rgb = try a.alloc(u8, width * height * 3); + const t = try decoder.decode(latents, rgb); + return .{ .png = try mi.png.encode(a, rgb, width, height), .load_s = load_s, .wall = t.wall, .gpu = t.gpu }; + } }; }; diff --git a/zig/tests/moss_image_decoder.zig b/zig/tests/moss_image_decoder.zig new file mode 100644 index 0000000000..dc45d15665 --- /dev/null +++ b/zig/tests/moss_image_decoder.zig @@ -0,0 +1,93 @@ +//! tf-moss-image-decoder: MOSS-Image's image decoder on a latents file, compared with a float64 decode of it. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const mi = tf.moss_image; + +const usage = + \\tf-moss-image-decoder MODEL_DIR LATENTS WIDTH HEIGHT REFERENCE TORCH OUT + \\MODEL_DIR is the OpenMOSS-Team/MOSS-Image-Base folder; LATENTS the transformer's output (height/16 x width/16 rows of + \\128 float32, as tf-moss-image writes it). REFERENCE and TORCH hold the publisher's decoder's output before the + \\clamp, [height, width, 3], in float64 and in float32 as the pipeline runs it (moss_image_case.py vae). The decode + \\must be as close to float64 as twice torch float32's cosine error and worst pixel row, and a second decode must + \\repeat the first byte for byte; else it fails. It is written to OUT as a PNG and timed. +; + +const check = mi.reference; + +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 != 8) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + const width = try std.fmt.parseInt(usize, args[3], 10); + const height = try std.fmt.parseInt(usize, args[4], 10); + const cwd = std.Io.Dir.cwd(); + const lat_bytes = try cwd.readFileAlloc(io, args[2], a, .limited(1 << 30)); + const latents = try a.alloc(f32, lat_bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(latents), lat_bytes[0 .. latents.len * 4]); + const ref_bytes = try cwd.readFileAlloc(io, args[5], a, .limited(1 << 32)); + const ref = try a.alloc(f64, ref_bytes.len / 8); + @memcpy(std.mem.sliceAsBytes(ref), ref_bytes[0 .. ref.len * 8]); + const torch_bytes = try cwd.readFileAlloc(io, args[6], a, .limited(1 << 31)); + const torch = try a.alloc(f32, torch_bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(torch), torch_bytes[0 .. torch.len * 4]); + for ([_]struct { []const u8, usize }{ .{ args[5], ref.len }, .{ args[6], torch.len } }) |f| if (f[1] != width * height * 3) { + std.debug.print("{s} holds {d} values; {d}x{d} needs {d}\n", .{ f[0], f[1], width, height, width * height * 3 }); + return 2; + }; + + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const loading = mtl.clock.seconds(); + const vae = try std.fs.path.join(a, &.{ args[1], "vae" }); + const d = try mi.Decoder.load(std.heap.page_allocator, io, device, vae, width, height); + defer d.deinit(); + const load_s = mtl.clock.seconds() - loading; + const rgb = try a.alloc(u8, width * height * 3); + const first = try d.decode(latents, rgb); + const out = d.output(); + std.debug.print("{s}: {d}x{d} decoded; loaded in {d:.2} s, first decode {d:.3} s\n", .{ device.name(), width, height, load_s, first.wall }); + const ours = check.match(f32, out, ref, width * 3); + const theirs = check.match(f32, torch, ref, width * 3); + const v = check.within(ours, theirs); + const torch_rgb = try a.alloc(u8, rgb.len); + for (torch_rgb, torch) |*p, x| p.* = pixel(x); + for ([_]struct { []const u8, []const u8, check.Match }{ .{ "ours", rgb, ours }, .{ "torch float32", torch_rgb, theirs } }) |c| { + const px = pixels(c[1], ref); + std.debug.print("{s} against float64: max |diff| {e:.3}, 1-cosine {e:.2}, worst row {e:.3} (row {d}); 8-bit values differing {d} of {d}, by at most {d}\n", .{ c[0], c[2].max_abs, 1 - c[2].cosine, c[2].worst_row, c[2].worst_at, px.differ, rgb.len, px.off }); + } + std.debug.print("gate: 1-cosine {e:.2} of {e:.2}, worst row {e:.3} of {e:.3}: {s}\n", .{ v.cosine_err, v.cosine_limit, v.row, v.row_limit, if (v.pass) "pass" else "FAIL" }); + try cwd.writeFile(io, .{ .sub_path = args[7], .data = try mi.png.encode(a, rgb, width, height) }); + const again = try a.alloc(u8, rgb.len); + const t = try d.decode(latents, again); + const same = std.mem.eql(u8, rgb, again); + std.debug.print("decoded again in {d:.3} s ({d:.3} s on the GPU), the same bytes: {}\n", .{ t.wall, t.gpu, same }); + const pass = check.accept(&.{v}, same); + std.debug.print("{s}\n", .{if (pass) "PASS within the gate, repeated byte for byte" else "FAIL"}); + return if (pass) 0 else 1; +} + +/// The 8-bit values in `got` that differ from float64's, and by how much at most. +fn pixels(got_all: []const u8, ref: []const f64) struct { differ: usize, off: u8 } { + var differ: usize = 0; + var off: u8 = 0; + for (got_all, ref) |got, r| { + const want = pixel(r); + if (want != got) { + differ += 1; + off = @max(off, if (want > got) want - got else got - want); + } + } + return .{ .differ = differ, .off = off }; +} + +/// The 8-bit value of a decoder output in float64: clamped, (x * 0.5 + 0.5) * 255 rounded half to even. +fn pixel(x: f64) u8 { + const v = std.math.clamp(std.math.clamp(x, -1, 1) * 0.5 + 0.5, 0, 1) * 255; + return @intFromFloat(@round(v) - @as(f64, if (@abs(v - @trunc(v)) == 0.5 and @mod(@round(v), 2) == 1) 1 else 0)); +} diff --git a/zig/tests/moss_image_run.zig b/zig/tests/moss_image_run.zig index 04ec331627..da0e873f7d 100644 --- a/zig/tests/moss_image_run.zig +++ b/zig/tests/moss_image_run.zig @@ -7,9 +7,11 @@ const usage = \\and noise levels and the starting latents. A `case` also holds a reference run (each step's input and both \\branches' velocities): each step then runs both branches on the reference's input and compares them (max |diff|, \\cosine, worst row). Then the whole schedule runs from the starting latents with the card's guidance (scale 3, the - \\negative branch skipping block 22), timed, and OUT_DIR gets latents.f32 (rows x columns x 128) and run.json. - \\With --prompt, the encoder in TRANSFORMER_DIR's checkpoint folder makes both prompts' rows (the negative from - \\--negative's file, else empty) in place of the case's; only the Metal build has the encoder. + \\negative branch skipping block 22), timed, and OUT_DIR gets latents.f32 (rows x columns x 128) and run.json. On + \\Metal the transformer is then freed and the checkpoint's image decoder (the vae/ folder beside TRANSFORMER_DIR) + \\turns the latents into OUT_DIR's picture.png. With --prompt, the encoder in TRANSFORMER_DIR's checkpoint folder + \\makes both prompts' rows (the negative from --negative's file, else empty) in place of the case's; only the Metal + \\build has the encoder. ; /// The card's sampling defaults: the guidance scale, and the block the negative branch skips. @@ -166,42 +168,55 @@ pub fn run(comptime B: type, init: std.process.Init) !u8 { const negative = both[1]; const text = prompt.len / mi.config.context; const negative_rows = negative.len / mi.config.context; - const loading = B.seconds(); - const model = try gpu.load(std.heap.page_allocator, io, args[1], @max(text, negative_rows), grid[1] * 16, grid[0] * 16); - defer model.deinit(); - const load_s = B.seconds() - loading; - std.debug.print("{s}: {d} image rows, prompts of {d} and {d} rows, {d} steps; loaded in {d:.1} s\n", .{ device, rows, text, negative_rows, steps, load_s }); - var out: std.ArrayList(u8) = .empty; - const refining = B.seconds(); - var branches: [2]mi.Prompt = undefined; - branches[0] = try model.prompt(prompt); - defer branches[0].deinit(); - branches[1] = try model.prompt(negative); - defer branches[1].deinit(); - const refine_s = B.seconds() - refining; - std.debug.print("both prompts refined: {d:.3} s\n", .{refine_s}); - try print(&out, a, "{{\"device\": \"{s}\", \"image_rows\": {d}, \"prompt_rows\": {d}, \"negative_rows\": {d}, \"load_s\": {d:.2}, \"refine_s\": {d:.4}", .{ device, rows, text, negative_rows, load_s, refine_s }); - if (encode_s) |s| try print(&out, a, ", \"native_encoder\": true, \"encode_s\": {d:.4}", .{s}); - const repeat = if (case.get("inputs") == null) true else try teacher(B, a, &case, model, &branches, timesteps, count, &out); - try print(&out, a, ", \"run\": [", .{}); + // The transformer is freed when its run ends, before the decoder loads. + const x, const repeat = denoise: { + const loading = B.seconds(); + const model = try gpu.load(std.heap.page_allocator, io, args[1], @max(text, negative_rows), grid[1] * 16, grid[0] * 16); + defer model.deinit(); + const load_s = B.seconds() - loading; + std.debug.print("{s}: {d} image rows, prompts of {d} and {d} rows, {d} steps; loaded in {d:.1} s\n", .{ device, rows, text, negative_rows, steps, load_s }); - const x = try a.dupe(f32, start); - const cond = try a.alloc(f32, count); - const uncond = try a.alloc(f32, count); - var total: f64 = 0; - for (0..steps) |s| { - const c = try model.forward(&branches[0], x, timesteps[s], null, cond); - const u = try model.forward(&branches[1], x, timesteps[s], negative_skip, uncond); - const by = sigmas[s + 1] - sigmas[s]; - for (x, cond, uncond) |*o, pc, pu| o.* += (pu + guidance * (pc - pu)) * by; - const wall = c.wall + u.wall; - total += wall; - std.debug.print("step {d}/{d}: {d:.3} s ({d:.3} s on the GPU, {d:.3} s modulation on the host)\n", .{ s + 1, steps, wall, c.gpu + u.gpu, c.host + u.host }); - try print(&out, a, "{s}{{\"wall_s\": {d:.4}, \"gpu_s\": {d:.4}, \"host_s\": {d:.4}}}", .{ if (s == 0) "" else ", ", wall, c.gpu + u.gpu, c.host + u.host }); + const refining = B.seconds(); + var branches: [2]mi.Prompt = undefined; + branches[0] = try model.prompt(prompt); + defer branches[0].deinit(); + branches[1] = try model.prompt(negative); + defer branches[1].deinit(); + const refine_s = B.seconds() - refining; + std.debug.print("both prompts refined: {d:.3} s\n", .{refine_s}); + try print(&out, a, "{{\"device\": \"{s}\", \"image_rows\": {d}, \"prompt_rows\": {d}, \"negative_rows\": {d}, \"load_s\": {d:.2}, \"refine_s\": {d:.4}", .{ device, rows, text, negative_rows, load_s, refine_s }); + if (encode_s) |s| try print(&out, a, ", \"native_encoder\": true, \"encode_s\": {d:.4}", .{s}); + const repeat = if (case.get("inputs") == null) true else try teacher(B, a, &case, model, &branches, timesteps, count, &out); + try print(&out, a, ", \"run\": [", .{}); + + const x = try a.dupe(f32, start); + const cond = try a.alloc(f32, count); + const uncond = try a.alloc(f32, count); + var total: f64 = 0; + for (0..steps) |s| { + const c = try model.forward(&branches[0], x, timesteps[s], null, cond); + const u = try model.forward(&branches[1], x, timesteps[s], negative_skip, uncond); + const by = sigmas[s + 1] - sigmas[s]; + for (x, cond, uncond) |*o, pc, pu| o.* += (pu + guidance * (pc - pu)) * by; + const wall = c.wall + u.wall; + total += wall; + std.debug.print("step {d}/{d}: {d:.3} s ({d:.3} s on the GPU, {d:.3} s modulation on the host)\n", .{ s + 1, steps, wall, c.gpu + u.gpu, c.host + u.host }); + try print(&out, a, "{s}{{\"wall_s\": {d:.4}, \"gpu_s\": {d:.4}, \"host_s\": {d:.4}}}", .{ if (s == 0) "" else ", ", wall, c.gpu + u.gpu, c.host + u.host }); + } + try print(&out, a, "], \"denoise_s\": {d:.3}", .{total}); + std.debug.print("denoise: {d:.2} s for {d} steps ({d} forwards)\n", .{ total, steps, 2 * steps }); + break :denoise .{ x, repeat }; + }; + if (comptime @hasDecl(B.Device, "decode")) { + const width = grid[1] * 16; + const height = grid[0] * 16; + const picture = try gpu.decode(a, io, std.fs.path.dirname(args[1]) orelse ".", x, width, height); + std.debug.print("decoded {d}x{d}: {d:.3} s ({d:.3} s on the GPU), the decoder loaded in {d:.2} s\n", .{ width, height, picture.wall, picture.gpu, picture.load_s }); + try print(&out, a, ", \"decoder_load_s\": {d:.3}, \"decode_s\": {d:.4}", .{ picture.load_s, picture.wall }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/picture.png", .{args[3]}), .data = picture.png }); } - try print(&out, a, "], \"denoise_s\": {d:.3}}}\n", .{total}); - std.debug.print("denoise: {d:.2} s for {d} steps ({d} forwards)\n", .{ total, steps, 2 * steps }); + try print(&out, a, "}}\n", .{}); try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/latents.f32", .{args[3]}), .data = std.mem.sliceAsBytes(x) }); try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = try std.fmt.allocPrint(a, "{s}/run.json", .{args[3]}), .data = out.items }); return if (repeat) 0 else 1;