diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 8335a68..55b3782 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -1,5 +1,89 @@ # Benchmarks +## Placement (v0.9.0) + +Where a kernel runs is decided by a handler, not by the program. These numbers +say what that costs and what a residency policy is worth. + +Harness: `cargo run -q --release -p loon-lang --features gpu --example bench_place` +Machine: Apple M4 Max (Metal, via wgpu). Numbers move a little run to run; +the transfer counts do not move at all. + +### What a residency policy is worth + +A chain of launches over one 4096-element buffer, on the GPU. The program is +identical in both columns — the only difference is whether `place/resident` +(nine lines, in `samples/place/lib.oo`) is wrapped around it. + +| launches | no policy | place/resident | speedup | +|---------:|----------:|---------------:|--------:| +| 8 | 29.3 ms | 8.4 ms | 3.5x | +| 32 | 94.8 ms | 11.2 ms | 8.4x | +| 128 | 358.8 ms | 18.1 ms | 19.9x | + +The gap grows with the chain because without a policy every launch uploads its +arguments, computes, and copies its results back; with one, the buffers stay +put and only the final `Place.read` moves anything. A recent Rust GPU-offload +paper measures the same gap at up to 400x between its convenient and explicit +interfaces, and closes it with `Preload`/`PreloadMut` annotations at every call +site plus a transfer-hoisting pass inside LLVM. Here it is a `handle` form. + +### Transfers + +Exact counts — these do not vary between runs. + +| launches | no policy | place/resident | bytes saved | +|---------:|----------:|---------------:|------------:| +| 1 | 2 uploads | 2 uploads | 0 B | +| 4 | 8 uploads | 2 uploads | 96 KB | +| 16 | 32 uploads | 2 uploads | 480 KB | +| 64 | 128 uploads | 2 uploads | 2.0 MB | + +### Kernel time: where it runs + +The same kernel, the same program, four placements. The CPU columns go through +the typed executor in `eir::kernel_exec` — raw slices, no boxing — so this is a +fair floor rather than a straw man. + +| elements | cpu | par | gpu | +|---------:|----:|----:|----:| +| 1,024 | 447 µs | 566 µs | 10.3 ms | +| 16,384 | 853 µs | 647 µs | 7.2 ms | +| 262,144 | 8.9 ms | 3.3 ms | 12.4 ms | +| 1,048,576 | 36.4 ms | 11.4 ms | 19.2 ms | + +The interesting row is the last one: on this machine, **every core beats the +GPU** for this kernel at a million elements. An M4 Max has a lot of fast cores, +and a GPU launch pays submission and transfer before it computes anything. That +is not a disappointing result, it is the point — you find it out by changing one +word on the command line, because the program does not know where it runs. + +The GPU only wins when the work per element is large enough to amortize getting +there, and where that crossover sits is a property of the machine, not of the +program. Which is a good argument for the decision living outside the program. + +This is not a comparison against optimized C or a hand-written kernel. That +comparison is not attempted here and nothing above should be quoted as one. + +### Launch overhead + +| | ns per launch | +|---|---:| +| cpu | ~6,700 | +| gpu, buffers resident | ~92,000 | + +Placement being an effect means every launch is an effect dispatch. That +dispatch is not what you are paying for: an effect operation costs roughly 3x a +function call (see the Effects section), which is nanoseconds, while a GPU +submission is tens of microseconds. + +### What is not measured + +- Any comparison against hand-written CUDA, HIP, or Metal. Not attempted. +- Reductions and atomics — outside the kernel subset for now. +- f64 on the GPU: WGSL core has no 64-bit scalar, so 64-bit buffers are + computed in 32 bits and that narrowing is reported rather than hidden. + ## Collection Benchmarks (v0.5.0) 100,000-element collections using `loop`/`recur` with persistent data structures (imbl). diff --git a/Cargo.lock b/Cargo.lock index 36d4d51..efaa0c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -133,6 +133,15 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "ash" +version = "0.38.0+1.3.281" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f" +dependencies = [ + "libloading", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -141,7 +150,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -152,9 +161,15 @@ checksum = "ffdcb70bdbc4d478427380519163274ac86e52916e10f0a8889adf0f96d3fee7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "base64" version = "0.21.7" @@ -173,6 +188,21 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + [[package]] name = "bitflags" version = "1.3.2" @@ -184,6 +214,9 @@ name = "bitflags" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +dependencies = [ + "serde_core", +] [[package]] name = "bitmaps" @@ -205,6 +238,12 @@ dependencies = [ "cpufeatures", ] +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + [[package]] name = "block-buffer" version = "0.10.4" @@ -223,6 +262,26 @@ dependencies = [ "allocator-api2", ] +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0e56a716f1e132ff6bf4bdac1c944a3fcdc1cae65f70a4a2a1ac3b401d2d1f" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "byteorder" version = "1.5.0" @@ -368,7 +427,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -405,6 +464,17 @@ dependencies = [ "unicode-width 0.1.14", ] +[[package]] +name = "codespan-reporting" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" +dependencies = [ + "serde", + "termcolor", + "unicode-width 0.2.2", +] + [[package]] name = "colorchoice" version = "1.0.4" @@ -429,12 +499,33 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core-graphics-types" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "libc", +] + [[package]] name = "cpp_demangle" version = "0.4.5" @@ -490,7 +581,7 @@ dependencies = [ "hashbrown 0.14.5", "log", "regalloc2", - "rustc-hash", + "rustc-hash 2.1.1", "serde", "smallvec", "target-lexicon", @@ -625,6 +716,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -761,7 +858,16 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", +] + +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", ] [[package]] @@ -887,6 +993,33 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea5190182e6915eb873ddbc16e23b711b6eb1f9c00a0d0a3a91b5f6228475225" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -951,7 +1084,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -1062,6 +1195,57 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "gpu-alloc" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45cf04b2726f02df5508c6de726acdc90cdf97ac771a9a0ffd8ba10a6e696bf9" +dependencies = [ + "bitflags 2.11.0", + "gpu-alloc-types", +] + +[[package]] +name = "gpu-alloc-types" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2bbed164dd10ed526c2e4fe3e721ca4a71c61730e5aafac6844b417b3227058" +dependencies = [ + "bitflags 2.11.0", +] + +[[package]] +name = "gpu-descriptor" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b89c83349105e3732062a895becfc71a8f921bb71ecbbdd8ff99263e3b53a0ca" +dependencies = [ + "bitflags 2.11.0", + "gpu-descriptor-types", + "hashbrown 0.15.5", +] + +[[package]] +name = "gpu-descriptor-types" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdf242682df893b86f33a73828fb09ca4b2d3bb6cc95249707fc684d27484b91" +dependencies = [ + "bitflags 2.11.0", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.14.5" @@ -1093,6 +1277,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hexf-parse" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" + [[package]] name = "home" version = "0.5.12" @@ -1120,7 +1310,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.62.2", ] [[package]] @@ -1393,6 +1583,16 @@ version = "0.2.182" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "libm" version = "0.2.16" @@ -1428,6 +1628,12 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "lock_api" version = "0.4.14" @@ -1465,7 +1671,7 @@ dependencies = [ "quote", "regex-syntax", "rustc_version", - "syn", + "syn 2.0.115", ] [[package]] @@ -1499,7 +1705,7 @@ name = "loon-lang" version = "0.7.0" dependencies = [ "blake3", - "codespan-reporting", + "codespan-reporting 0.11.1", "cranelift-codegen", "cranelift-frontend", "cranelift-jit", @@ -1510,12 +1716,15 @@ dependencies = [ "imbl", "insta", "logos", + "naga", + "pollster", "serde_json", "tar", "ureq", "uuid", "wasm-encoder 0.227.1", "wasmparser 0.227.1", + "wgpu", ] [[package]] @@ -1562,6 +1771,15 @@ dependencies = [ "libc", ] +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + [[package]] name = "maybe-owned" version = "0.3.4" @@ -1583,6 +1801,21 @@ dependencies = [ "rustix 1.1.3", ] +[[package]] +name = "metal" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f569fb946490b5743ad69813cb19629130ce9374034abe31614a36402d18f99e" +dependencies = [ + "bitflags 2.11.0", + "block", + "core-graphics-types", + "foreign-types", + "log", + "objc", + "paste", +] + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -1604,6 +1837,31 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "naga" +version = "25.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b977c445f26e49757f9aca3631c3b8b836942cb278d69a92e7b80d3b24da632" +dependencies = [ + "arrayvec", + "bit-set", + "bitflags 2.11.0", + "cfg_aliases", + "codespan-reporting 0.12.0", + "half", + "hashbrown 0.15.5", + "hexf-parse", + "indexmap", + "log", + "num-traits", + "once_cell", + "rustc-hash 1.1.0", + "spirv", + "strum", + "thiserror 2.0.18", + "unicode-ident", +] + [[package]] name = "nibble_vec" version = "0.1.0" @@ -1625,6 +1883,25 @@ dependencies = [ "libc", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", +] + [[package]] name = "object" version = "0.36.7" @@ -1664,6 +1941,15 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "ordered-float" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7bb71e1b3fa6ca1c61f383464aaf2bb0e2f8e772a1f01d486832464de363b951" +dependencies = [ + "num-traits", +] + [[package]] name = "owo-colors" version = "4.2.3" @@ -1722,7 +2008,7 @@ checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -1737,6 +2023,18 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "pollster" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" + +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + [[package]] name = "postcard" version = "1.1.3" @@ -1774,7 +2072,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.115", ] [[package]] @@ -1786,6 +2084,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "profiling" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d595e54a326bc53c1c197b32d295e14b169e3cfeaa8dc82b529f947fba6bcf5" + [[package]] name = "psm" version = "0.1.30" @@ -1872,6 +2176,12 @@ dependencies = [ "rand_core", ] +[[package]] +name = "raw-window-handle" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" + [[package]] name = "rayon" version = "1.11.0" @@ -1942,7 +2252,7 @@ dependencies = [ "bumpalo", "hashbrown 0.15.5", "log", - "rustc-hash", + "rustc-hash 2.1.1", "smallvec", ] @@ -1964,6 +2274,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "renderdoc-sys" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" + [[package]] name = "ring" version = "0.17.14" @@ -1994,6 +2310,12 @@ version = "0.1.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.1" @@ -2151,7 +2473,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2175,7 +2497,7 @@ checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2260,6 +2582,15 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "spirv" +version = "0.3.0+sdk-1.3.268.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eda41003dc44290527a59b13432d4a0379379fa074b70174882adfbdfd917844" +dependencies = [ + "bitflags 2.11.0", +] + [[package]] name = "sptr" version = "0.3.2" @@ -2272,6 +2603,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "str_indices" version = "0.4.4" @@ -2284,6 +2621,28 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.115", +] + [[package]] name = "subtle" version = "2.6.1" @@ -2301,6 +2660,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -2309,7 +2679,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2393,7 +2763,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2404,7 +2774,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2442,7 +2812,7 @@ checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2550,7 +2920,7 @@ checksum = "84fd902d4e0b9a4b27f2f440108dc034e1758628a9b702f8ec61ad66355422fa" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2578,7 +2948,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2598,7 +2968,7 @@ checksum = "70977707304198400eb4835a78f6a9f928bf41bba420deb8fdb175cd965d77a7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -2757,7 +3127,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.115", "wasm-bindgen-shared", ] @@ -2978,7 +3348,7 @@ dependencies = [ "anyhow", "proc-macro2", "quote", - "syn", + "syn 2.0.115", "wasmtime-component-util", "wasmtime-wit-bindgen", "wit-parser 0.221.3", @@ -3103,7 +3473,7 @@ checksum = "86ff86db216dc0240462de40c8290887a613dddf9685508eb39479037ba97b5b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -3225,6 +3595,129 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "wgpu" +version = "25.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec8fb398f119472be4d80bc3647339f56eb63b2a331f6a3d16e25d8144197dd9" +dependencies = [ + "arrayvec", + "bitflags 2.11.0", + "cfg_aliases", + "document-features", + "hashbrown 0.15.5", + "js-sys", + "log", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "smallvec", + "static_assertions", + "wasm-bindgen", + "web-sys", + "wgpu-core", + "wgpu-hal", + "wgpu-types", +] + +[[package]] +name = "wgpu-core" +version = "25.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7b882196f8368511d613c6aeec80655160db6646aebddf8328879a88d54e500" +dependencies = [ + "arrayvec", + "bit-set", + "bit-vec", + "bitflags 2.11.0", + "cfg_aliases", + "document-features", + "hashbrown 0.15.5", + "indexmap", + "log", + "naga", + "once_cell", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "rustc-hash 1.1.0", + "smallvec", + "thiserror 2.0.18", + "wgpu-core-deps-apple", + "wgpu-core-deps-windows-linux-android", + "wgpu-hal", + "wgpu-types", +] + +[[package]] +name = "wgpu-core-deps-apple" +version = "25.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfd488b3239b6b7b185c3b045c39ca6bf8af34467a4c5de4e0b1a564135d093d" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-windows-linux-android" +version = "25.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cba5fb5f7f9c98baa7c889d444f63ace25574833df56f5b817985f641af58e46" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-hal" +version = "25.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f968767fe4d3d33747bbd1473ccd55bf0f6451f55d733b5597e67b5deab4ad17" +dependencies = [ + "android_system_properties", + "arrayvec", + "ash", + "bitflags 2.11.0", + "block", + "bytemuck", + "cfg_aliases", + "core-graphics-types", + "gpu-alloc", + "gpu-descriptor", + "hashbrown 0.15.5", + "libc", + "libloading", + "log", + "metal", + "naga", + "objc", + "ordered-float", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "renderdoc-sys", + "smallvec", + "thiserror 2.0.18", + "wgpu-types", + "windows", +] + +[[package]] +name = "wgpu-types" +version = "25.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2aa49460c2a8ee8edba3fca54325540d904dd85b2e086ada762767e17d06e8bc" +dependencies = [ + "bitflags 2.11.0", + "bytemuck", + "js-sys", + "log", + "thiserror 2.0.18", + "web-sys", +] + [[package]] name = "wiggle" version = "29.0.1" @@ -3251,7 +3744,7 @@ dependencies = [ "proc-macro2", "quote", "shellexpand", - "syn", + "syn 2.0.115", "witx", ] @@ -3263,7 +3756,7 @@ checksum = "08c5c473d4198e6c2d377f3809f713ff0c110cab88a0805ae099a82119ee250c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", "wiggle-generate", ] @@ -3316,17 +3809,51 @@ dependencies = [ "wasmtime-environ", ] +[[package]] +name = "windows" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd04d41d93c4992d421894c18c8b43496aa748dd4c081bac0dc93eb0489272b6" +dependencies = [ + "windows-core 0.58.0", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-core" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba6d44ec8c2591c134257ce647b7ea6b20335bf6379a27dac5f1641fcf59f99" +dependencies = [ + "windows-implement 0.58.0", + "windows-interface 0.58.0", + "windows-result 0.2.0", + "windows-strings 0.1.0", + "windows-targets 0.52.6", +] + [[package]] name = "windows-core" version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ - "windows-implement", - "windows-interface", + "windows-implement 0.60.2", + "windows-interface 0.59.3", "windows-link", - "windows-result", - "windows-strings", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-implement" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bbd5b46c938e506ecbce286b6628a02171d56153ba733b6c741fc627ec9579b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.115", ] [[package]] @@ -3337,7 +3864,18 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", +] + +[[package]] +name = "windows-interface" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053c4c462dc91d3b1504c6fe5a726dd15e216ba718e84a0e46a88fbe5ded3515" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.115", ] [[package]] @@ -3348,7 +3886,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -3357,6 +3895,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d1043d8214f791817bab27572aaa8af63732e11bf84aa21a45a78d6c317ae0e" +dependencies = [ + "windows-targets 0.52.6", +] + [[package]] name = "windows-result" version = "0.4.1" @@ -3366,6 +3913,16 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-strings" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd9b125c486025df0eabcb585e62173c6c9eddcec5d117d3b6e8c30e2ee4d10" +dependencies = [ + "windows-result 0.2.0", + "windows-targets 0.52.6", +] + [[package]] name = "windows-strings" version = "0.5.1" @@ -3589,7 +4146,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn", + "syn 2.0.115", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -3605,7 +4162,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.115", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -3712,7 +4269,7 @@ checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", "synstructure", ] @@ -3733,7 +4290,7 @@ checksum = "4122cd3169e94605190e77839c9a40d40ed048d305bfdc146e7df40ab0f3e517" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -3753,7 +4310,7 @@ checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", "synstructure", ] @@ -3793,7 +4350,7 @@ checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] diff --git a/crates/loon-cli/Cargo.toml b/crates/loon-cli/Cargo.toml index e4fca5a..e9ad705 100644 --- a/crates/loon-cli/Cargo.toml +++ b/crates/loon-cli/Cargo.toml @@ -7,6 +7,11 @@ edition.workspace = true name = "loon" path = "src/main.rs" +[features] +default = [] +# Real GPU execution for `loon run --place gpu`. +gpu = ["loon-lang/gpu"] + [dependencies] loon-lang = { path = "../loon-lang", features = ["pkg-fetch", "native"] } clap = { version = "4", features = ["derive"] } diff --git a/crates/loon-cli/src/main.rs b/crates/loon-cli/src/main.rs index 68c8f66..c3fbc26 100644 --- a/crates/loon-cli/src/main.rs +++ b/crates/loon-cli/src/main.rs @@ -50,6 +50,13 @@ enum Command { /// (Loon data format), for `loon replay` #[arg(long, value_name = "TRACE")] record: Option, + /// Where kernels run when nothing handles `Place`: `cpu` (one memory, + /// no transfers) or `device` (separate memory, transfers accounted) + #[arg(long, value_name = "MODE", default_value = "cpu")] + place: String, + /// Print placement accounting — launches, transfers, bytes — on exit + #[arg(long)] + place_stats: bool, }, /// Re-run a program feeding recorded effect results back from a trace /// (see `loon run --record`) — same trace, same execution, same crash @@ -175,7 +182,29 @@ fn main() { legacy, native, ref record, + ref place, + place_stats, } => { + let place_mode = match loon_lang::eir::place::Mode::parse(place) { + Some(m) => m, + None => { + eprintln!( + "{}: unknown placement mode '{place}' — expected 'cpu' or 'device'", + "error".red().bold() + ); + std::process::exit(1); + } + }; + if (place_mode != loon_lang::eir::place::Mode::Cpu || place_stats) + && (wasm || legacy || native) + { + eprintln!( + "{}: --place/--place-stats require the default EIR VM backend \ + (drop --wasm/--legacy/--native)", + "error".red().bold() + ); + std::process::exit(1); + } if record.is_some() && (wasm || legacy || native) { eprintln!( "{}: --record requires the default EIR VM backend \ @@ -191,7 +220,7 @@ fn main() { } else if legacy { run_file_legacy(file); } else { - run_file(file, record.as_deref(), unchecked); + run_file(file, record.as_deref(), unchecked, place_mode, place_stats); } } Command::Replay { @@ -281,7 +310,13 @@ fn precheck_source(path: &std::path::Path, source: &str) { } } -fn run_file(path: &PathBuf, record: Option<&std::path::Path>, unchecked: bool) { +fn run_file( + path: &PathBuf, + record: Option<&std::path::Path>, + unchecked: bool, + place_mode: loon_lang::eir::place::Mode, + show_place_stats: bool, +) { let source = match std::fs::read_to_string(path) { Ok(s) => s, Err(e) => { @@ -399,6 +434,17 @@ fn run_file(path: &PathBuf, record: Option<&std::path::Path>, unchecked: bool) { } r } + None if place_mode != loon_lang::eir::place::Mode::Cpu || show_place_stats => { + match loon_lang::eir::vm::eval_eir_placed(&source, base_dir, place_mode) { + Ok((r, stats)) => { + if show_place_stats { + print!("{}", stats.table()); + } + Ok(r) + } + Err(e) => Err(e), + } + } None => loon_lang::eir::vm::eval_eir_with_base_dir(&source, base_dir), }; match result { diff --git a/crates/loon-cli/tests/abi_conformance.rs b/crates/loon-cli/tests/abi_conformance.rs new file mode 100644 index 0000000..50737b1 --- /dev/null +++ b/crates/loon-cli/tests/abi_conformance.rs @@ -0,0 +1,224 @@ +//! Every backend must encode a value the same way, bit for bit. +//! +//! Loon has three code generators — the register VM, the Cranelift JIT, and +//! the WASM emitter — and until recently each carried its own copy of the +//! NaN-boxing constants under a comment asking the next person to keep them in +//! sync. That is the same setup that produced a real bug in a recent Rust +//! GPU-offload compiler, where a slice lowered as `(ptr, len)` on two targets +//! and `[i64; 2]` on a third; the authors found it by hand and wrote that +//! automated cross-target validation was still missing. +//! +//! This is that validation. `eir::layout` now holds the encoding once, and the +//! tests below compile the same source on every available backend and compare +//! the raw 64-bit results. A divergence fails here rather than surfacing later +//! as a wrong answer on one target. + +use loon_lang::eir::layout::{ + nanbox, BufferHeader, DType, BUF_HDR_SIZE, BUF_OFF_DATA, BUF_OFF_DTYPE, BUF_OFF_LEN, + GOLDEN_IMMEDIATES, +}; + +/// A program whose value is the literal. +/// +/// Deliberately a bare top-level expression rather than a `main` function: +/// reaching `main` requires closures, which the native backend does not +/// implement, and a program every backend refuses would make the comparison +/// below vacuous. +fn program_for(literal: &str) -> String { + literal.to_string() +} + +/// Evaluate on the register VM, returning the raw encoded word. +fn vm_bits(literal: &str) -> u64 { + let src = program_for(literal); + let result = loon_lang::eir::vm::eval_eir(&src) + .unwrap_or_else(|e| panic!("VM failed on `{literal}`: {e:?}")); + result.value.bits() +} + +/// Evaluate through the Cranelift JIT, returning the raw encoded word. +/// +/// The native backend does not implement every operation; a literal it cannot +/// compile yields `None` and is skipped rather than reported as a mismatch. +fn native_bits(literal: &str) -> Option { + let src = program_for(literal); + match loon_lang::eir::native::eval_native(&src) { + Ok(v) => Some(v.bits()), + Err(_) => None, + } +} + +/// Compile with the EIR WASM backend and run under wasmtime, returning the +/// raw encoded word. Returns `None` if the backend cannot compile the program. +fn wasm_bits(literal: &str) -> Option { + use loon_lang::eir::backend::Backend; + + let src = program_for(literal); + let exprs = loon_lang::parser::parse(&src).expect("parses"); + let mut checker = loon_lang::check::Checker::new(); + let errors = checker.check_program(&exprs); + assert!(errors.is_empty(), "type errors on `{literal}`: {errors:?}"); + let module = loon_lang::eir::lower::lower(&checker); + + let mut backend = loon_lang::eir::wasm::WasmBackend; + let bytes = backend.compile(&module).ok()?; + + let engine = wasmtime::Engine::default(); + let wasm_module = wasmtime::Module::new(&engine, &bytes).ok()?; + let mut store = wasmtime::Store::new(&engine, ()); + let mut linker = wasmtime::Linker::new(&engine); + // The emitter imports its two IO hooks unconditionally; a literal program + // never calls them, but they must resolve for instantiation to succeed. + linker + .func_wrap("host", "println", |_: i64| {}) + .expect("bind println"); + linker + .func_wrap("host", "print", |_: i64| {}) + .expect("bind print"); + let instance = linker.instantiate(&mut store, &wasm_module).ok()?; + let start = instance + .get_typed_func::<(), i64>(&mut store, "_start") + .ok()?; + start.call(&mut store, ()).ok().map(|v| v as u64) +} + +#[test] +fn the_vm_encodes_immediates_exactly_as_the_layout_says() { + for (name, literal, expected) in GOLDEN_IMMEDIATES { + assert_eq!( + vm_bits(literal), + *expected, + "`{name}` ({literal}): the VM disagrees with eir::layout" + ); + } +} + +#[test] +fn every_backend_agrees_on_every_immediate() { + let mut checked_native = 0usize; + let mut checked_wasm = 0usize; + + for (name, literal, expected) in GOLDEN_IMMEDIATES { + let vm = vm_bits(literal); + assert_eq!(vm, *expected, "`{name}`: VM vs layout"); + + if let Some(bits) = native_bits(literal) { + checked_native += 1; + assert_eq!( + bits, vm, + "`{name}` ({literal}): native backend encodes {bits:#018x}, VM encodes {vm:#018x}" + ); + } + if let Some(bits) = wasm_bits(literal) { + checked_wasm += 1; + assert_eq!( + bits, vm, + "`{name}` ({literal}): wasm backend encodes {bits:#018x}, VM encodes {vm:#018x}" + ); + } + } + + // A backend silently compiling nothing would make this test vacuous. + assert!( + checked_native > 0, + "the native backend compiled none of the golden immediates" + ); + assert!( + checked_wasm > 0, + "the wasm backend compiled none of the golden immediates" + ); +} + +#[test] +fn immediates_are_distinguishable_from_one_another() { + // Encoding two different values identically is the failure mode that would + // make the cross-backend comparison above pass while everything is broken. + let mut seen: Vec<(&str, u64)> = Vec::new(); + for (name, literal, _) in GOLDEN_IMMEDIATES { + let bits = vm_bits(literal); + if let Some((other, _)) = seen.iter().find(|(_, b)| *b == bits) { + panic!("`{name}` and `{other}` encode to the same word {bits:#018x}"); + } + seen.push((name, bits)); + } +} + +#[test] +fn the_buffer_header_layout_is_pinned() { + // Field order and offsets are exactly what diverged across targets in the + // Rust offload work. Pin them explicitly: a reordering of the struct is a + // silent ABI break, and this is the test that refuses to let it be silent. + assert_eq!(std::mem::size_of::(), BUF_HDR_SIZE); + assert_eq!(BUF_OFF_DTYPE, 0); + assert_eq!(BUF_OFF_LEN, 8); + assert_eq!(BUF_OFF_DATA, 16); + + let header = BufferHeader::new(DType::F32, 4096, 0x1234_5678_9ABC_DEF0); + let bytes = header.to_bytes(); + assert_eq!(bytes.len(), BUF_HDR_SIZE); + assert_eq!( + u32::from_le_bytes(bytes[BUF_OFF_DTYPE..BUF_OFF_DTYPE + 4].try_into().unwrap()), + DType::F32 as u32 + ); + assert_eq!( + u64::from_le_bytes(bytes[BUF_OFF_LEN..BUF_OFF_LEN + 8].try_into().unwrap()), + 4096 + ); + assert_eq!( + u64::from_le_bytes(bytes[BUF_OFF_DATA..BUF_OFF_DATA + 8].try_into().unwrap()), + 0x1234_5678_9ABC_DEF0 + ); + assert_eq!(BufferHeader::from_bytes(&bytes), Some(header)); +} + +#[test] +fn a_reordered_header_does_not_round_trip() { + // The mutation test: swap the two 64-bit fields, as a target with a + // different struct layout effectively would, and confirm the bytes no + // longer decode to the original. If this ever passes, the conformance + // check above has stopped checking anything. + let header = BufferHeader::new(DType::I32, 7, 99); + let mut mutated = header.to_bytes(); + let (len_half, data_half): ([u8; 8], [u8; 8]) = ( + mutated[BUF_OFF_LEN..BUF_OFF_LEN + 8].try_into().unwrap(), + mutated[BUF_OFF_DATA..BUF_OFF_DATA + 8].try_into().unwrap(), + ); + mutated[BUF_OFF_LEN..BUF_OFF_LEN + 8].copy_from_slice(&data_half); + mutated[BUF_OFF_DATA..BUF_OFF_DATA + 8].copy_from_slice(&len_half); + + let decoded = BufferHeader::from_bytes(&mutated).expect("still a well-formed header"); + assert_ne!( + decoded, header, + "swapping len and data went unnoticed — the layout check is vacuous" + ); + assert_eq!(decoded.len, 99); + assert_eq!(decoded.data, 7); +} + +#[test] +fn gpu_placement_rejects_64_bit_element_types() { + // WGSL core has no f64 and no 64-bit integers. Placement must refuse those + // buffers outright rather than quietly narrowing a program's precision. + assert!(DType::F32.gpu_ok() && DType::I32.gpu_ok()); + assert!(!DType::F64.gpu_ok() && !DType::I64.gpu_ok()); + for dtype in [DType::F32, DType::F64, DType::I32, DType::I64] { + assert_eq!(DType::from_name(dtype.name()), Some(dtype)); + } +} + +#[test] +fn int_payloads_survive_the_full_48_bit_range() { + // The int tag carries a 48-bit signed payload. The extremes are where a + // backend that sign-extends differently would part ways from the others. + for n in [0i64, 1, -1, (1i64 << 47) - 1, -(1i64 << 47)] { + let literal = n.to_string(); + let expected = nanbox::encode_int(n); + assert_eq!(vm_bits(&literal), expected, "VM on {n}"); + if let Some(bits) = native_bits(&literal) { + assert_eq!(bits, expected, "native on {n}"); + } + if let Some(bits) = wasm_bits(&literal) { + assert_eq!(bits, expected, "wasm on {n}"); + } + } +} diff --git a/crates/loon-lang/Cargo.toml b/crates/loon-lang/Cargo.toml index 7233567..5106c0f 100644 --- a/crates/loon-lang/Cargo.toml +++ b/crates/loon-lang/Cargo.toml @@ -6,6 +6,10 @@ edition.workspace = true [features] default = [] pkg-fetch = ["dep:blake3", "dep:ureq", "dep:flate2", "dep:tar", "dep:uuid"] +# Real GPU execution through wgpu (Metal here, Vulkan/DX12 elsewhere). +# Off by default: it is a large dependency, and every other placement mode +# works without it. +gpu = ["dep:wgpu", "dep:pollster"] native = [ "dep:cranelift-codegen", "dep:cranelift-frontend", @@ -31,7 +35,17 @@ cranelift-frontend = { version = "0.116", optional = true } cranelift-jit = { version = "0.116", optional = true } cranelift-module = { version = "0.116", optional = true } cranelift-native = { version = "0.116", optional = true } +wgpu = { version = "25", optional = true, default-features = false, features = [ + "metal", + "vulkan", + "wgsl", +] } + +pollster = { version = "0.4", optional = true } [dev-dependencies] insta = "1" +# Validates emitted WGSL the way a driver would, so shader errors surface in +# CI rather than on a machine that happens to have a GPU. +naga = { version = "25", features = ["wgsl-in"] } wasmparser = "0.227" diff --git a/crates/loon-lang/examples/bench_place.rs b/crates/loon-lang/examples/bench_place.rs new file mode 100644 index 0000000..7e9d1a6 --- /dev/null +++ b/crates/loon-lang/examples/bench_place.rs @@ -0,0 +1,217 @@ +//! What placement costs, and what a residency policy saves. +//! +//! Dependency-free (no criterion). Run in release for meaningful numbers: +//! +//! cargo run -q --release -p loon-lang --features gpu --example bench_place +//! +//! Without the `gpu` feature the GPU rows are skipped and everything else +//! still reports. +//! +//! Two things are measured, and it is worth being precise about which is which +//! because they answer different questions. +//! +//! **Kernel time** compares the same kernel across every place it can run: +//! one core, every core, and the GPU. The CPU columns go through the typed +//! executor in `eir::kernel_exec`, not the general interpreter, so this is a +//! fair floor rather than a straw man. It is still not a comparison against +//! optimized C, and nothing here should be quoted as one. +//! +//! **Transfer counts** compare a chain of launches with no residency policy +//! against the same chain under a handler that pins what each launch touches. +//! This one is exact rather than statistical: it counts uploads, and the count +//! does not vary between runs. It is the number the offload literature cares +//! about, and the one a policy actually controls. + +use loon_lang::eir::place::{Mode, PlaceStats}; +use loon_lang::eir::vm::eval_eir_placed; +use std::time::{Duration, Instant}; + +/// A kernel with enough arithmetic per element to be worth moving. +const KERNEL: &str = "[kernel heat [i src dst] \ + [let v [at src i]] \ + [put dst i [+ [* 0.25 v] [* 0.75 [sqrt [abs [+ v 1.0]]]]]]]"; + +fn dir() -> std::path::PathBuf { + std::env::current_dir().expect("cwd") +} + +/// Run `src` under `mode`, returning wall time and the placement accounting. +fn timed(src: &str, mode: Mode) -> Option<(Duration, PlaceStats)> { + // Warm once: the first GPU launch pays for adapter discovery and shader + // compilation, which is a real cost but not a per-launch one. + let _ = eval_eir_placed(src, &dir(), mode).ok()?; + let start = Instant::now(); + let (_, stats) = eval_eir_placed(src, &dir(), mode).ok()?; + Some((start.elapsed(), stats)) +} + +/// A program that runs one kernel over `n` elements, `reps` times. +fn workload(n: usize, reps: usize) -> String { + let runs = (0..reps) + .map(|_| format!("[Place.run heat {n} #[src dst]]")) + .collect::>() + .join(" "); + format!( + "{KERNEL} \ + [fn main [] \ + [let src [buf-zeros {n}]] \ + [let mut dst [buf-zeros {n}]] \ + {runs} \ + [let _ [Place.read dst]] \ + []]" + ) +} + +/// The same workload under a residency handler. +fn workload_resident(n: usize, reps: usize) -> String { + let runs = (0..reps) + .map(|_| format!("[Place.run heat {n} #[src dst]]")) + .collect::>() + .join(" "); + format!( + "{KERNEL} \ + [fn work [] \ + [let src [buf-zeros {n}]] \ + [let mut dst [buf-zeros {n}]] \ + {runs} \ + [Place.read dst]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k m args] [do [Place.pin args] [resume [Place.run k m args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [let _ [resident work]] []]" + ) +} + +fn main() { + println!("placement benchmarks"); + println!("machine: {}", std::env::consts::OS); + + // ── Kernel time: interpreter versus GPU ── + // + // The CPU column is Loon's interpreter running the kernel body once per + // element. It is the honest floor, not a tuned baseline. + println!("\nkernel time — one launch, varying size"); + println!( + " {:>10} {:>10} {:>10} {:>10}", + "elements", "cpu", "par", "gpu" + ); + for n in [1_024usize, 16_384, 262_144, 1_048_576] { + let src = workload(n, 1); + let cpu = timed(&src, Mode::Cpu).map(|(d, _)| d); + let par = timed(&src, Mode::Par).map(|(d, _)| d); + let gpu = timed(&src, Mode::Gpu).map(|(d, _)| d); + println!( + " {:>10} {:>10} {:>10} {:>10}", + n, + cpu.map(fmt).unwrap_or_else(|| "—".into()), + par.map(fmt).unwrap_or_else(|| "—".into()), + gpu.map(fmt).unwrap_or_else(|| "skipped".into()), + ); + } + + // ── Transfers: what a policy controls ── + // + // Exact counts, identical every run. This is the gap the offload + // literature reports and the one a handler closes here. + println!("\ntransfers — a chain of launches over one buffer"); + println!( + " {:>8} {:>16} {:>16} {:>10}", + "launches", "no policy", "place/resident", "saved" + ); + for reps in [1usize, 4, 16, 64] { + let naive = timed(&workload(4_096, reps), Mode::Device).map(|(_, s)| s); + let managed = timed(&workload_resident(4_096, reps), Mode::Device).map(|(_, s)| s); + match (naive, managed) { + (Some(a), Some(b)) => { + let saved = a.bytes_in.saturating_sub(b.bytes_in); + println!( + " {:>8} {:>16} {:>16} {:>10}", + reps, + format!("{} uploads", a.uploads), + format!("{} uploads", b.uploads), + loon_lang::eir::place::human_bytes(saved) + ); + } + _ => println!(" {reps:>8} (failed)"), + } + } + + // ── What the policy is worth in time, not just in bytes ── + // + // The transfer counts above are exact but abstract. This is the same + // comparison in wall clock on whatever device is present: a chain of + // launches with no residency policy, where every launch uploads its + // arguments and copies its results back, against the identical program + // under a handler that keeps them in place. + // + // This is the shape of the gap a recent Rust offload paper measures at up + // to 400x between its convenient and explicit interfaces. There it is + // closed by annotations at every call site plus an LLVM pass; here by the + // handler in samples/place/lib.oo. + // Only the real device is timed. The modelled one moves no actual bytes, + // so its wall clock would be measuring the interpreter and calling it a + // transfer cost. + for (label, mode) in [("gpu", Mode::Gpu)] { + println!("\nchain of launches on the {label} — 4096 elements"); + println!( + " {:>8} {:>12} {:>16} {:>8}", + "launches", "no policy", "place/resident", "speedup" + ); + let mut any = false; + for reps in [8usize, 32, 128] { + let naive = timed(&workload(4_096, reps), mode).map(|(d, _)| d); + let managed = timed(&workload_resident(4_096, reps), mode).map(|(d, _)| d); + match (naive, managed) { + (Some(a), Some(b)) => { + any = true; + let speedup = if b.as_secs_f64() > 0.0 { + format!("{:.1}x", a.as_secs_f64() / b.as_secs_f64()) + } else { + "—".into() + }; + println!( + " {:>8} {:>12} {:>16} {:>8}", + reps, + fmt(a), + fmt(b), + speedup + ); + } + _ => {} + } + } + if !any { + println!(" skipped — this mode is not available in this build"); + } + } + + // ── Launch overhead ── + // + // Placement is an effect, so every launch is an effect dispatch. This is + // the cost of that decision, measured rather than asserted: a kernel over + // a single element, where the work itself is negligible. On the CPU that + // is dispatch and nothing else. + println!("\nlaunch overhead — one work item, so the cost is dispatch"); + let one = workload(1, 200); + if let Some((d, stats)) = timed(&one, Mode::Cpu) { + let per = d.as_nanos() as f64 / stats.launches.max(1) as f64; + println!(" {per:>10.0} ns per launch (cpu)"); + } + if let Some((d, stats)) = timed(&workload_resident(1, 200), Mode::Gpu) { + let per = d.as_nanos() as f64 / stats.launches.max(1) as f64; + println!(" {per:>10.0} ns per launch (gpu, buffers resident)"); + println!( + " a GPU launch is a submission to another processor; the effect\n dispatch around it is not what you are paying for." + ); + } +} + +fn fmt(d: Duration) -> String { + let ms = d.as_secs_f64() * 1000.0; + if ms >= 1.0 { + format!("{ms:.1} ms") + } else { + format!("{:.0} µs", d.as_micros()) + } +} diff --git a/crates/loon-lang/src/check/kernel.rs b/crates/loon-lang/src/check/kernel.rs new file mode 100644 index 0000000..85d32ec --- /dev/null +++ b/crates/loon-lang/src/check/kernel.rs @@ -0,0 +1,499 @@ +//! Kernels: functions restricted enough to run somewhere else. +//! +//! A kernel is written like any other function — +//! +//! ```text +//! [kernel saxpy [i a x out] +//! [put out i [+ [* a [at x i]] [at out i]]]] +//! ``` +//! +//! — and is an ordinary function as far as the rest of the compiler is +//! concerned: [`desugar`] rewrites the head to `fn` before type checking, so +//! kernels infer, lower, and run exactly like anything else. What the keyword +//! buys is a *promise*, checked by [`verify`]: the body stays inside a small +//! numeric subset with no closures, no allocation, no strings, and no effects. +//! +//! That restriction is the safety argument. A recent Rust GPU-offload design +//! needs an `unsafe trait` for its partitioning strategies, because a kernel +//! there can be handed a slice and index it however it likes; the invariant +//! that threads touch disjoint elements has to be promised by hand. Here a +//! kernel cannot express the unsafe program in the first place — it receives +//! an index and may only write at that index — so the guarantee comes from +//! what the language declines to compile rather than from a promise attached +//! to a trait impl. +//! +//! The subset is also what makes a kernel portable. Everything admitted here +//! has a direct equivalent in scalar machine code and in WGSL, which is why +//! the same source can run on a CPU core, across threads, or on a GPU without +//! being rewritten for each. + +use crate::ast::{Expr, ExprKind}; +use crate::errors::codes::ErrorCode; +use crate::errors::LoonDiagnostic; +use std::collections::HashSet; + +/// Builtins a kernel body may call. +/// +/// Every one of these is a scalar numeric operation or a buffer element +/// access — things a GPU can do per work item. Anything that allocates, +/// inspects a collection, or touches a string is absent on purpose. +const KERNEL_BUILTINS: &[&str] = &[ + // Buffer element access. + "at", "put", "buf-len", // Math. + "sqrt", "pow", "floor", "ceil", "round", "sin", "cos", "tan", "asin", "acos", "atan", "atan2", + "log", "log10", "exp", "abs", "min", "max", // Logic. + "not", +]; + +// Deliberately absent: `int` and `float`. In Loon those parse strings, and a +// kernel has no strings to parse. Mixing an integer and a float in kernel +// arithmetic is handled where it belongs — the WGSL emitter inserts the +// conversion, because WGSL will not mix the two silently and neither should +// the language pretend to. + +/// Special forms a kernel body may use. +/// +/// `loop`/`recur` earn their place because a reduction needs them: a work item +/// that sums its own chunk of the input is still writing only its own element, +/// so the disjointness rule holds and no new language feature is required. +const KERNEL_FORMS: &[&str] = &["let", "if", "do", "and", "or", "loop", "recur", "mut"]; + +/// Rewrite every `[kernel name [params] body...]` into the equivalent `fn`, +/// returning the rewritten program and the set of kernel names. +/// +/// Doing this before type checking means kernels are not a second language: +/// they infer types, get ownership modes, and lower through exactly the same +/// path as ordinary functions. The only thing that remains special about them +/// is the promise [`verify`] enforces and the fact that they can be placed. +pub fn desugar(exprs: &[Expr]) -> (Vec, HashSet) { + let mut names = HashSet::new(); + let out = exprs.iter().map(|e| rewrite(e, &mut names)).collect(); + (out, names) +} + +fn rewrite(expr: &Expr, names: &mut HashSet) -> Expr { + let ExprKind::List(items) = &expr.kind else { + return expr.clone(); + }; + let mut new_items: Vec = items.iter().map(|i| rewrite(i, names)).collect(); + if let Some(ExprKind::Symbol(head)) = new_items.first().map(|i| &i.kind) { + if head == "kernel" && new_items.len() >= 3 { + if let ExprKind::Symbol(name) = &new_items[1].kind { + names.insert(name.clone()); + new_items[0] = Expr { + kind: ExprKind::Symbol("fn".to_string()), + ..new_items[0].clone() + }; + } + } + } + Expr { + kind: ExprKind::List(new_items), + ..expr.clone() + } +} + +/// Check that each kernel body stays inside the placeable subset. +/// +/// Runs on the desugared program, so kernels appear as `fn` forms; `names` +/// says which ones were written with the `kernel` keyword. +pub fn verify(exprs: &[Expr], names: &HashSet) -> Vec { + if names.is_empty() { + return Vec::new(); + } + let mut v = Verifier { + kernels: names, + errors: Vec::new(), + index_param: None, + }; + for expr in exprs { + v.walk_program(expr); + } + v.errors +} + +struct Verifier<'a> { + kernels: &'a HashSet, + errors: Vec, + /// The name of the kernel currently being checked's work index — its first + /// parameter. A `put` at any other index is a scatter. + index_param: Option, +} + +impl Verifier<'_> { + fn walk_program(&mut self, expr: &Expr) { + let ExprKind::List(items) = &expr.kind else { + return; + }; + let is_fn = + matches!(items.first().map(|i| &i.kind), Some(ExprKind::Symbol(h)) if h == "fn"); + if is_fn && items.len() >= 3 { + if let ExprKind::Symbol(name) = &items[1].kind { + if self.kernels.contains(name) { + self.index_param = first_param(&items[2]); + self.check_params(name, &items[2]); + for body in &items[3..] { + self.check_body(name, body); + } + return; + } + } + } + for item in items { + self.walk_program(item); + } + } + + /// A kernel's first parameter is its work index, so it needs at least one. + fn check_params(&mut self, kernel: &str, params: &Expr) { + let ExprKind::List(ps) = ¶ms.kind else { + return; + }; + if ps.is_empty() { + self.errors.push( + LoonDiagnostic::new( + ErrorCode::E0601, + format!("kernel '{kernel}' has no parameters"), + ) + .with_why( + "a kernel runs once per work item and takes that item's index as its first \ + parameter, so it cannot be nullary" + .to_string(), + ) + .with_fix(format!( + "give '{kernel}' an index parameter, e.g. [kernel {kernel} [i ...] ...]" + )) + .with_label( + params.span, + "expected at least an index parameter", + true, + ), + ); + } + } + + fn check_body(&mut self, kernel: &str, expr: &Expr) { + match &expr.kind { + // Literals and names are always fine. + ExprKind::Int(_) + | ExprKind::Float(_) + | ExprKind::Bool(_) + | ExprKind::Symbol(_) + | ExprKind::Keyword(_) => {} + + ExprKind::Str(_) => self.reject( + kernel, + expr, + "a string literal", + "kernels have no string support: there is nowhere to put the bytes on a device", + ), + ExprKind::Vec(_) | ExprKind::Map(_) | ExprKind::Set(_) | ExprKind::Tuple(_) => self + .reject( + kernel, + expr, + "a collection literal", + "kernels cannot allocate; pass a buffer in and index it instead", + ), + + ExprKind::List(items) => self.check_call(kernel, expr, items), + _ => {} + } + } + + fn check_call(&mut self, kernel: &str, expr: &Expr, items: &[Expr]) { + let Some(head) = items.first() else { + return; + }; + // `Effect.op` reaches here as a dot access on an uppercase name. + if let ExprKind::DotAccess(base, field) = &head.kind { + if let ExprKind::Symbol(b) = &base.kind { + if b.chars().next().is_some_and(|c| c.is_uppercase()) { + self.reject( + kernel, + expr, + &format!("the effect operation '{b}.{field}'"), + "a kernel runs where there is no handler tower to perform effects against", + ); + return; + } + } + } + + let ExprKind::Symbol(name) = &head.kind else { + // An indirect call means a function value, which a kernel has no + // way to have obtained. + self.reject( + kernel, + expr, + "an indirect call", + "kernels can only call other kernels and a fixed set of numeric builtins", + ); + return; + }; + + // An effect performed inside a kernel would need a handler on the + // device. Effects are how placement itself is expressed, so a kernel + // performing one would be circular as well as unimplementable. + if name.contains('.') && name.chars().next().is_some_and(|c| c.is_uppercase()) { + self.reject( + kernel, + expr, + &format!("the effect operation '{name}'"), + "a kernel runs where there is no handler tower to perform effects against", + ); + return; + } + + // `[loop [name init ...] body...]` — the binding list is names paired + // with initializers, not a call, so only the initializers and the body + // are expressions to check. + if name == "loop" && items.len() >= 2 { + if let ExprKind::List(bindings) = &items[1].kind { + for init in bindings.iter().skip(1).step_by(2) { + self.check_body(kernel, init); + } + } + for body in &items[2..] { + self.check_body(kernel, body); + } + return; + } + + // A work item may write at its own index and nowhere else. + // + // This is the disjointness rule the whole design rests on: it is what + // lets the parallel executor hand each thread a slice and what lets a + // GPU dispatch run every work item at once. The Rust offload work + // arrives at the same guarantee by having partitioning strategies + // promise it in an `unsafe impl`; here the program that would violate + // it does not compile. + if name == "put" && items.len() >= 3 { + let idx = &items[2]; + let ok = match (&idx.kind, &self.index_param) { + (ExprKind::Symbol(s), Some(p)) => s == p, + _ => false, + }; + if !ok { + let index_name = self.index_param.clone().unwrap_or_else(|| "i".to_string()); + self.errors.push( + LoonDiagnostic::new( + ErrorCode::E0602, + format!("kernel '{kernel}' writes at an index other than its own"), + ) + .with_why(format!( + "every work item runs at once, so each may only write element \ + '{index_name}'; writing elsewhere means two of them can reach the \ + same element and the result would depend on which got there first" + )) + .with_fix(format!( + "write at '{index_name}', and read whatever else this element needs \ + with `at`" + )) + .with_label(idx.span, "not this work item's element", true), + ); + } + } + + if name == "fn" || name == "kernel" { + self.reject( + kernel, + expr, + "a nested function", + "kernels cannot create closures; there is no heap to capture into", + ); + return; + } + + let known = KERNEL_FORMS.contains(&name.as_str()) + || KERNEL_BUILTINS.contains(&name.as_str()) + || self.kernels.contains(name) + || is_operator(name); + + if !known { + self.reject( + kernel, + expr, + &format!("a call to '{name}'"), + "kernels may call other kernels and numeric builtins only", + ); + return; + } + + for arg in &items[1..] { + self.check_body(kernel, arg); + } + } + + fn reject(&mut self, kernel: &str, expr: &Expr, what: &str, why: &str) { + self.errors.push( + LoonDiagnostic::new( + ErrorCode::E0600, + format!("kernel '{kernel}' contains {what}"), + ) + .with_why(why.to_string()) + .with_fix( + "move this out of the kernel and pass the result in as a buffer or scalar" + .to_string(), + ) + .with_label(expr.span, "not allowed inside a kernel", true), + ); + } +} + +/// The name of a parameter list's first entry. +fn first_param(params: &Expr) -> Option { + let ExprKind::List(ps) = ¶ms.kind else { + return None; + }; + match ps.first().map(|p| &p.kind) { + Some(ExprKind::Symbol(name)) => Some(name.clone()), + _ => None, + } +} + +/// Operators lower to arithmetic instructions rather than builtin calls. +fn is_operator(name: &str) -> bool { + matches!( + name, + "+" | "-" | "*" | "/" | "%" | "=" | "!=" | "<" | ">" | "<=" | ">=" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::parser::parse; + + fn check(src: &str) -> Vec { + let exprs = parse(src).expect("parses"); + let (desugared, names) = desugar(&exprs); + verify(&desugared, &names) + } + + #[test] + fn a_numeric_kernel_is_accepted() { + let errors = check("[kernel saxpy [i a x out] [put out i [+ [* a [at x i]] [at out i]]]]"); + assert!(errors.is_empty(), "unexpected errors: {errors:?}"); + } + + #[test] + fn control_flow_and_math_are_accepted() { + let errors = check( + "[kernel clamp [i lo hi x out] \ + [let v [at x i]] \ + [put out i [if [< v lo] lo [if [> v hi] hi [sqrt [abs v]]]]]]", + ); + assert!(errors.is_empty(), "unexpected errors: {errors:?}"); + } + + #[test] + fn a_kernel_may_call_another_kernel() { + let errors = check( + "[kernel scale [i s x out] [put out i [* s [at x i]]]] \ + [kernel twice [i x out] [scale i 2.0 x out]]", + ); + assert!(errors.is_empty(), "unexpected errors: {errors:?}"); + } + + #[test] + fn desugaring_makes_a_kernel_an_ordinary_function() { + let exprs = parse("[kernel k [i x] [put x i 1.0]]").expect("parses"); + let (desugared, names) = desugar(&exprs); + assert!(names.contains("k")); + let ExprKind::List(items) = &desugared[0].kind else { + panic!("expected a list"); + }; + assert!( + matches!(&items[0].kind, ExprKind::Symbol(h) if h == "fn"), + "the head should have become `fn`, got {:?}", + items[0].kind + ); + } + + #[test] + fn allocation_is_rejected() { + let errors = check("[kernel k [i out] [put out i [len #[1 2 3]]]]"); + assert!(!errors.is_empty(), "a vector literal should be rejected"); + assert_eq!(errors[0].code, ErrorCode::E0600); + } + + #[test] + fn strings_are_rejected() { + let errors = check("[kernel k [i out] [put out i [len \"hello\"]]]"); + assert!(!errors.is_empty(), "a string literal should be rejected"); + } + + #[test] + fn closures_are_rejected() { + let errors = check("[kernel k [i out] [put out i [[fn [x] x] 1.0]]]"); + assert!(!errors.is_empty(), "a nested fn should be rejected"); + } + + #[test] + fn effects_are_rejected() { + // This one is the interesting case: placement *is* an effect, so a + // kernel performing one would be asking the device to reach back into + // the handler tower it was dispatched from. + let errors = check("[kernel k [i out] [do [IO.println 1] [put out i 1.0]]]"); + assert!(!errors.is_empty(), "an effect should be rejected"); + assert!( + errors[0].what.contains("IO.println"), + "the message should name the operation: {}", + errors[0].what + ); + } + + #[test] + fn calls_to_ordinary_functions_are_rejected() { + let errors = check("[fn helper [x] x] [kernel k [i out] [put out i [helper 1.0]]]"); + assert!( + !errors.is_empty(), + "a kernel may not call a non-kernel function" + ); + } + + #[test] + fn writing_at_another_work_items_index_is_rejected() { + // The disjointness rule, enforced. Two work items reaching the same + // element would make the result depend on which arrived first. + for src in [ + "[kernel k [i s d] [put d [- 99 i] [at s i]]]", + "[kernel k [i d] [put d 0 1.0]]", + "[kernel k [i d] [put d [+ i 1] 1.0]]", + ] { + let errors = check(src); + assert!(!errors.is_empty(), "should reject a scatter: {src}"); + assert_eq!(errors[0].code, ErrorCode::E0602, "{src}"); + } + } + + #[test] + fn writing_at_your_own_index_is_the_whole_point() { + let errors = check("[kernel k [i s d] [put d i [+ [at s i] [at s 0]]]]"); + assert!( + errors.is_empty(), + "reading anywhere is fine; only writing is restricted: {errors:?}" + ); + } + + #[test] + fn the_index_parameter_can_be_called_anything() { + let errors = check("[kernel k [row d] [put d row 1.0]]"); + assert!(errors.is_empty(), "unexpected errors: {errors:?}"); + } + + #[test] + fn a_nullary_kernel_is_rejected() { + let errors = check("[kernel k [] 1]"); + assert!(!errors.is_empty(), "a kernel needs an index parameter"); + assert_eq!(errors[0].code, ErrorCode::E0601); + } + + #[test] + fn ordinary_functions_are_left_alone() { + // The restrictions apply to kernels only; normal code keeps every + // feature the language has. + let errors = check("[fn ordinary [] [let v #[1 2 3]] [IO.println \"hi\"] [len v]]"); + assert!(errors.is_empty(), "unexpected errors: {errors:?}"); + } +} diff --git a/crates/loon-lang/src/check/mod.rs b/crates/loon-lang/src/check/mod.rs index aa1181b..cf3e2ba 100644 --- a/crates/loon-lang/src/check/mod.rs +++ b/crates/loon-lang/src/check/mod.rs @@ -1,3 +1,4 @@ +pub mod kernel; pub mod ownership; use crate::ast::{Expr, ExprKind, NodeId}; @@ -120,6 +121,15 @@ pub struct Checker { recur_params: Vec>, /// Expanded program after macro expansion (available after check_program) pub expanded_program: Vec, + /// Per-parameter ownership modes for each named function, inferred during + /// `check_program`. Rust makes you write these down as `&T` / `&mut T`; + /// Loon infers them, and lowering carries them into the IR so backends can + /// tell a read-only argument from one that is written through. + pub fn_param_modes: HashMap>, + /// Names defined with the `kernel` keyword. They are ordinary functions + /// after desugaring; this is what remembers that they promised to stay + /// inside the placeable subset. + pub kernels: HashSet, } /// Split a param list at a `&` marker: returns the fixed params and the @@ -145,6 +155,8 @@ impl Checker { subst: Subst::new(), env: TypeEnv::new(), errors: Vec::new(), + fn_param_modes: HashMap::new(), + kernels: HashSet::new(), constructors: HashMap::new(), type_constructors: HashMap::new(), type_of: HashMap::new(), @@ -1398,6 +1410,145 @@ impl Checker { ); } + // ── Dense buffers ── + // + // `Buf a` is an opaque, fixed-length array of unboxed numbers. It is + // the representation that can leave the process — to another thread, + // another backend, or a device — which ordinary persistent + // collections cannot. + { + let mk_var = |c: &mut Self| match c.subst.fresh() { + Type::Var(v) => v, + _ => unreachable!("fresh() returns a Var"), + }; + let buf_of = |t: Type| Type::Con("Buf".to_string(), vec![t]); + let vec_of = |t: Type| Type::Con("Vec".to_string(), vec![t]); + + // buf / buf-f64: ∀a:Num. Vec a → Buf Float + // buf-i32: ∀a:Num. Vec a → Buf Int + // + // The source vector may hold ints or floats; the buffer's element + // type is fixed by which constructor was called, and the values + // are converted on the way in. That conversion is the honest part: + // an f32 buffer really does hold f32, and saying so in the type is + // better than pretending the input type survived. + for (name, elem) in [ + ("buf", Type::Float), + ("buf-f64", Type::Float), + ("buf-i32", Type::Int), + ] { + let a = self.subst.fresh(); + let tv = match a { + Type::Var(v) => v, + _ => unreachable!("fresh() returns a Var"), + }; + self.env.set_global( + name.to_string(), + Scheme { + bounds: vec![( + tv, + vec![TraitBound { + trait_name: "Num".to_string(), + }], + )], + vars: vec![tv], + ty: Type::Fn( + vec![vec_of(Type::Var(tv))], + Box::new(buf_of(elem)), + EffectRow::pure(), + ), + }, + ); + } + + // buf-zeros: Int → Buf Float, buf-zeros-i32: Int → Buf Int + for (name, elem) in [("buf-zeros", Type::Float), ("buf-zeros-i32", Type::Int)] { + self.env.set_global( + name.to_string(), + Scheme::mono(Type::Fn( + vec![Type::Int], + Box::new(buf_of(elem)), + EffectRow::pure(), + )), + ); + } + + // buf-len: ∀a. Buf a → Int + let v = mk_var(self); + self.env.set_global( + "buf-len".to_string(), + Scheme { + bounds: vec![], + vars: vec![v], + ty: Type::Fn( + vec![buf_of(Type::Var(v))], + Box::new(Type::Int), + EffectRow::pure(), + ), + }, + ); + + // buf-dtype: ∀a. Buf a → Str + let v = mk_var(self); + self.env.set_global( + "buf-dtype".to_string(), + Scheme { + bounds: vec![], + vars: vec![v], + ty: Type::Fn( + vec![buf_of(Type::Var(v))], + Box::new(Type::Str), + EffectRow::pure(), + ), + }, + ); + + // buf->vec: ∀a. Buf a → Vec a + let v = mk_var(self); + self.env.set_global( + "buf->vec".to_string(), + Scheme { + bounds: vec![], + vars: vec![v], + ty: Type::Fn( + vec![buf_of(Type::Var(v))], + Box::new(vec_of(Type::Var(v))), + EffectRow::pure(), + ), + }, + ); + + // at: ∀a. Buf a → Int → a + let v = mk_var(self); + self.env.set_global( + "at".to_string(), + Scheme { + bounds: vec![], + vars: vec![v], + ty: Type::Fn( + vec![buf_of(Type::Var(v)), Type::Int], + Box::new(Type::Var(v)), + EffectRow::pure(), + ), + }, + ); + + // put: ∀a. Buf a → Int → a → Buf a + let v = mk_var(self); + self.env.set_global( + "put".to_string(), + Scheme { + bounds: vec![], + vars: vec![v], + ty: Type::Fn( + vec![buf_of(Type::Var(v)), Type::Int, Type::Var(v)], + Box::new(buf_of(Type::Var(v))), + EffectRow::pure(), + ), + }, + ); + } + // int: Str → Int self.env.set_global( "int".to_string(), @@ -1600,15 +1751,37 @@ impl Checker { } } - // sum: Vec Int → Int (approximate) - self.env.set_global( - "sum".to_string(), - Scheme::mono(Type::Fn( - vec![Type::Con("Vec".to_string(), vec![Type::Int])], - Box::new(Type::Int), - EffectRow::pure(), - )), - ); + // sum: ∀a:Num. Vec a → a + // + // The builtin registry has always declared this `Vec Num → Num` and the + // interpreter has always implemented it that way. The checker said + // `Vec Int → Int` with the comment "(approximate)", which meant summing + // a vector of floats — the natural last step of any reduction — did not + // type check. + { + let a = self.subst.fresh(); + let tv = match a { + Type::Var(v) => v, + _ => unreachable!("fresh() returns a Var"), + }; + self.env.set_global( + "sum".to_string(), + Scheme { + bounds: vec![( + tv, + vec![TraitBound { + trait_name: "Num".to_string(), + }], + )], + vars: vec![tv], + ty: Type::Fn( + vec![Type::Con("Vec".to_string(), vec![Type::Var(tv)])], + Box::new(Type::Var(tv)), + EffectRow::pure(), + ), + }, + ); + } // str: ∀a. a → Str { @@ -3182,13 +3355,25 @@ impl Checker { // Type-check arguments against declared param types let arg_types: Vec = items[1..].iter().map(|a| self.infer(a)).collect(); - if arg_types.len() != op_def.params.len() { + // A variadic op's last parameter absorbs the rest, so + // it needs at least the fixed ones and accepts more. + let arity_ok = if op_def.variadic { + arg_types.len() + 1 >= op_def.params.len() + } else { + arg_types.len() == op_def.params.len() + }; + if !arity_ok { self.errors.push( LoonDiagnostic::new( ErrorCode::E0202, format!( - "`{effect}.{op}` expects {} argument(s), got {}", - op_def.params.len(), + "`{effect}.{op}` expects {}{} argument(s), got {}", + if op_def.variadic { "at least " } else { "" }, + if op_def.variadic { + op_def.params.len() - 1 + } else { + op_def.params.len() + }, arg_types.len() ), ) @@ -3823,6 +4008,7 @@ impl Checker { name: op_name, params, return_type, + variadic: false, }); } } @@ -5691,6 +5877,12 @@ impl Checker { } }; + // Phase 1b: Kernels become ordinary functions before anything else + // looks at the program, so they infer, lower, and run through exactly + // the same path. Only the promise they made is kept aside. + let (expanded, kernels) = kernel::desugar(&expanded); + self.kernels = kernels; + // Phase 2a: Pre-declare top-level functions so forward references // (including mutual recursion across definitions) resolve. Each name // gets a placeholder type; its real scheme replaces it at the @@ -5746,6 +5938,14 @@ impl Checker { self.expanded_program = final_exprs; self.check_trait_constraints(); + // Ownership modes ride along with the checker so lowering can see them. + // This is a syntactic pass over the expanded program; its diagnostics + // are the dedicated ownership pass's job, not ours. + self.fn_param_modes = ownership::infer_param_modes(&self.expanded_program); + // A kernel that cannot be placed is worth saying so about at compile + // time, not when a handler tries to ship it somewhere. + self.errors + .extend(kernel::verify(&self.expanded_program, &self.kernels)); std::mem::take(&mut self.errors) } diff --git a/crates/loon-lang/src/check/ownership.rs b/crates/loon-lang/src/check/ownership.rs index e654e3a..7e1a8b5 100644 --- a/crates/loon-lang/src/check/ownership.rs +++ b/crates/loon-lang/src/check/ownership.rs @@ -31,8 +31,14 @@ enum BindingState { } /// How a function uses a particular parameter. -#[derive(Debug, Clone, Copy, PartialEq)] -enum ParamMode { +/// +/// This is Loon's answer to what other languages make you write down. Rust +/// spells it `&T` / `&mut T` / `T` at every signature; here it is inferred +/// from the body and then used the same way — to decide what a call site may +/// do with a binding, and (for kernels) which direction data has to move +/// across a placement boundary. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ParamMode { /// Parameter is only read — immutable borrow at call site. Borrow, /// Parameter is mutated (push!, set!) — mutable borrow at call site. @@ -120,6 +126,14 @@ impl<'a> OwnershipChecker<'a> { // lets the escaping/answer-passing style (e.g. `[[resume s] s]`) // reuse a value across the resume without a false move error. "resume", + // Buffer readers. `[at buf i]` observes one element, and the rest + // inspect or copy out without consuming, so a kernel parameter + // that is only read stays a borrow — which is what makes it an + // `:in` (host-to-device only) argument at a placement boundary. + "at", + "buf-len", + "buf-dtype", + "buf->vec", ] { borrow_fns.insert(name.to_string()); } @@ -409,8 +423,11 @@ impl<'a> OwnershipChecker<'a> { ); } } - "push!" | "set!" => { - // First arg is mutably borrowed + "push!" | "set!" | "put" => { + // First arg is mutably borrowed. `put` is the + // kernel element write; treating it exactly like + // `set!` is what makes a written-to kernel buffer + // infer as `:inout` with no annotation. if items.len() > 1 { if let ExprKind::Symbol(name) = &items[1].kind { if let Some(idx) = param_names.iter().position(|p| p == name) { @@ -570,8 +587,10 @@ impl<'a> OwnershipChecker<'a> { } return; } - "push!" => { - // push! requires mutable borrow of first arg + "push!" | "put" => { + // Writes through the first argument, reads the rest — a + // mutable borrow, not a move: the caller still owns the + // buffer (or vector) afterwards. if items.len() > 1 { if let ExprKind::Symbol(name) = &items[1].kind { self.mut_borrow(name, items[1].span); @@ -593,6 +612,24 @@ impl<'a> OwnershipChecker<'a> { } } + // Effect operations consume their arguments by default, which is what + // makes "responding twice with the same value is a compile error" a + // property of the language rather than a convention. + // + // `Place` is the exception, and not by special pleading: placement is + // defined to read its inputs and write through its outputs, never to + // take ownership of either. `[Place.run k n x out]` leaves the caller + // owning both `x` and `out` — indeed `[Place.read out]` afterwards is + // the entire point, and is how a program gets its results back. + if let ExprKind::DotAccess(base, _) = &items[0].kind { + if matches!(&base.kind, ExprKind::Symbol(b) if b == "Place") { + for item in &items[1..] { + self.check_expr(item); + } + return; + } + } + // Generic function call — head is borrowed, args use per-param modes if available let callee_name = if let ExprKind::Symbol(s) = &items[0].kind { Some(s.clone()) @@ -664,11 +701,22 @@ impl<'a> OwnershipChecker<'a> { self.fn_param_modes.insert(name.clone(), modes); } + // A parameter the analysis just inferred as mutably borrowed IS + // mutable inside this body — that is what the inference means. + // Requiring `let mut` on a parameter would be asking the author to + // restate a conclusion the compiler already reached, and there is + // nowhere to write it anyway. + let modes = fn_name + .as_ref() + .and_then(|n| self.fn_param_modes.get(n).cloned()) + .unwrap_or_default(); + self.push_scope(); - for p in params { + for (idx, p) in params.iter().enumerate() { if let ExprKind::Symbol(name) = &p.kind { let is_copy = self.is_value_copy(p); - self.define(name.clone(), p.span, is_copy, false); + let is_mut = matches!(modes.get(idx), Some(ParamMode::MutBorrow)); + self.define(name.clone(), p.span, is_copy, is_mut); } } for expr in &args[body_start..] { @@ -733,6 +781,115 @@ impl<'a> OwnershipChecker<'a> { } std::mem::take(&mut self.errors) } + + /// Per-parameter modes for every named function seen so far. + pub fn param_modes(&self) -> &HashMap> { + &self.fn_param_modes + } + + /// Re-run mode analysis over every named function until the answers stop + /// changing. + /// + /// The first pass sees definitions in source order, so a call to a + /// not-yet-analyzed function falls back to the conservative `Move`. Once + /// every function has an entry, re-analyzing resolves those calls for + /// real. Repeating until stable makes the result independent of the order + /// the definitions were written in, which matters because a spuriously + /// `Move` parameter reads as "the callee consumed this" — a claim that + /// costs optimizations downstream and is simply untrue. + /// + /// Bounded to a small number of rounds: each round can only replace a + /// guess with a real answer, so a program that has not settled by then has + /// a cycle whose conservative reading is the correct one to keep. + pub fn refine_param_modes(&mut self, exprs: &[Expr]) { + const MAX_ROUNDS: usize = 8; + let defns = Self::collect_fn_defns(exprs); + for _ in 0..MAX_ROUNDS { + let mut changed = false; + for (name, param_names, body) in &defns { + let modes = self.analyze_param_modes(param_names, body); + match self.fn_param_modes.get(name) { + Some(prev) if *prev == modes => {} + _ => { + self.fn_param_modes.insert(name.clone(), modes); + changed = true; + } + } + } + if !changed { + return; + } + } + } + + /// Every `[fn name [params] body...]` in the program, including those + /// nested inside module-level forms, as (name, param names, body). + fn collect_fn_defns(exprs: &[Expr]) -> Vec<(String, Vec, Vec)> { + let mut out = Vec::new(); + for expr in exprs { + Self::collect_fn_defns_into(expr, &mut out); + } + out + } + + fn collect_fn_defns_into(expr: &Expr, out: &mut Vec<(String, Vec, Vec)>) { + let ExprKind::List(items) = &expr.kind else { + return; + }; + let Some(head) = items.first() else { + return; + }; + if let ExprKind::Symbol(h) = &head.kind { + if h == "fn" && items.len() >= 3 { + if let (ExprKind::Symbol(name), ExprKind::List(params)) = + (&items[1].kind, &items[2].kind) + { + let param_names: Vec = params + .iter() + .filter_map(|p| match &p.kind { + ExprKind::Symbol(n) => Some(n.clone()), + _ => None, + }) + .collect(); + // Skip an effect-row annotation between params and body, + // exactly as `check_defn` does. + let mut body_start = 3; + if body_start < items.len() + && matches!(&items[body_start].kind, ExprKind::Set(_) | ExprKind::Map(_)) + { + body_start += 1; + } + out.push((name.clone(), param_names, items[body_start..].to_vec())); + return; + } + } + } + // Not a function definition — look inside for nested ones. + for item in items { + Self::collect_fn_defns_into(item, out); + } + } +} + +/// Infer parameter modes for a program without reporting ownership errors. +/// +/// The full ownership check runs as a separate frontend pass and its results +/// are thrown away with the checker. The *modes* it computes along the way are +/// useful to the compiler proper — they say, for each parameter, whether a +/// caller's value is read, mutated, or consumed — so this entry point runs the +/// same analysis for its modes alone. Diagnostics are discarded here; the +/// dedicated pass is still what reports them. +/// +/// Unlike the diagnostic pass, this one iterates to a fixed point. A single +/// pass in source order has to guess at callees it has not reached yet, and it +/// guesses `Move`; that makes the answer depend on the order two functions +/// happen to be written in, which is not a property anything downstream should +/// inherit. Re-running until nothing changes removes the guess. +pub fn infer_param_modes(exprs: &[Expr]) -> HashMap> { + let mut checker = OwnershipChecker::new(); + let _ = checker.check_program(exprs); + checker.refine_param_modes(exprs); + checker.fn_param_modes } impl Default for OwnershipChecker<'_> { diff --git a/crates/loon-lang/src/codegen/mod.rs b/crates/loon-lang/src/codegen/mod.rs index 55c0203..b61a5e1 100644 --- a/crates/loon-lang/src/codegen/mod.rs +++ b/crates/loon-lang/src/codegen/mod.rs @@ -1181,6 +1181,7 @@ impl Compiler { name: op_name, params, return_type, + variadic: false, }); } } diff --git a/crates/loon-lang/src/effects/mod.rs b/crates/loon-lang/src/effects/mod.rs index f49a9b4..cf9d64d 100644 --- a/crates/loon-lang/src/effects/mod.rs +++ b/crates/loon-lang/src/effects/mod.rs @@ -13,6 +13,11 @@ pub struct EffectOp { pub name: String, pub params: Vec<(String, Option)>, // (name, type_name) pub return_type: Option, + /// The final parameter absorbs any number of extra arguments. + /// + /// `Place.run` needs this: a kernel's arguments are its own, and the + /// placement operation carries however many of them there are. + pub variadic: bool, } /// Registry of declared effects @@ -34,6 +39,17 @@ fn op(name: &str, param_names: &[&str]) -> EffectOp { name: name.to_string(), params: params(param_names), return_type: None, + variadic: false, + } +} + +/// An op whose last parameter absorbs any number of trailing arguments. +fn var_op(name: &str, param_names: &[&str]) -> EffectOp { + EffectOp { + name: name.to_string(), + params: params(param_names), + return_type: None, + variadic: true, } } @@ -43,6 +59,7 @@ fn typed_op(name: &str, param_names: &[&str], ret: &str) -> EffectOp { name: name.to_string(), params: params(param_names), return_type: Some(ret.to_string()), + variadic: false, } } @@ -124,6 +141,39 @@ impl EffectRegistry { op("seed", &["n"]), ], }); + // Place effect — where a kernel runs is a decision a handler makes, + // not a property of the program. Unhandled, it runs serially on the + // CPU, so a program that never mentions placement still works. + reg.register(EffectDecl { + name: "Place".to_string(), + operations: vec![ + // [Place.run kernel n args...] — run `kernel` once per index + // in 0..n. Returns unit; kernels write through buffers. + op("run", &["kernel", "n", "args"]), + // [Place.read buf] — the only way to get buffer contents back + // to the host. Being an operation is the point: a residency + // handler learns where every synchronization point is without + // the programmer marking any of them. + op("read", &["buf"]), + // [Place.pin buf] / [Place.unpin buf] — hints that a buffer + // should stay where it is between launches. + op("pin", &["buf"]), + op("unpin", &["buf"]), + // [Place.stats] — transfer and launch counters so far. + op("stats", &[]), + ], + }); + // Host effect — the seam an asynchronous embedder needs. + // + // `Host.park` takes a continuation and keeps it. The handler that + // performs it returns without resuming, so the computation unwinds and + // the host decides when the rest of it runs. That is the whole + // mechanism behind answering an operation that cannot be answered yet, + // like reading back a GPU buffer in a browser. + reg.register(EffectDecl { + name: "Host".to_string(), + operations: vec![op("park", &["continuation", "request"])], + }); reg.register(EffectDecl { name: "Embed".to_string(), operations: vec![op("encode", &["text"])], diff --git a/crates/loon-lang/src/eir/device.rs b/crates/loon-lang/src/eir/device.rs new file mode 100644 index 0000000..5e49524 --- /dev/null +++ b/crates/loon-lang/src/eir/device.rs @@ -0,0 +1,106 @@ +//! What a placement backend has to be able to do. +//! +//! There are two of these and they could hardly be less alike. On a desktop, +//! `eir::gpu::Gpu` drives wgpu directly and blocks on the queue. In a browser, +//! `loon-wasm` proxies every call to JavaScript, because WebGPU is only +//! reachable through promises and the VM is synchronous — the wasm side blocks +//! on `Atomics.wait` while another thread does the asynchronous part. +//! +//! Both answer the same six questions, so the VM does not know which one it +//! has. That is the same move placement makes at the language level, one layer +//! down: the thing that varies is behind an interface, and the code that uses +//! it does not change when the answer does. + +use super::vm::Buffer; + +/// Why a device operation failed, in a sentence a person can act on. +#[derive(Debug, Clone)] +pub struct DeviceError(pub String); + +impl std::fmt::Display for DeviceError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for DeviceError {} + +/// A device that kernels can be dispatched to. +/// +/// Buffers are identified by the host's heap slot, so residency is something +/// the caller and the device agree about by name rather than by handle. +pub trait Device { + /// A human-readable name, for `--place-stats` and error messages. + fn name(&self) -> String; + + /// Put `buf` on the device if it is not there already. Returns whether an + /// upload actually happened, so the caller's accounting reflects the + /// hardware rather than a prediction of it. + fn ensure_resident(&self, id: usize, buf: &Buffer) -> Result; + + /// Whether `id` currently has device storage. + fn is_resident(&self, id: usize) -> bool; + + /// Run `shader` over `n` work items against buffers that are already + /// resident. Nothing is uploaded and nothing is read back: when data moves + /// is the caller's decision, which is the whole point. + fn dispatch( + &self, + shader: &str, + entry: &str, + n: u32, + scalars: &[f32], + buffers: &[usize], + ) -> Result<(), DeviceError>; + + /// Copy a resident buffer's contents back to the host. + fn download(&self, id: usize, byte_len: usize) -> Result, DeviceError>; + + /// Release device storage for `id`. + fn evict(&self, id: usize); +} + +// ─── The installed device ────────────────────────────────────────────────── + +thread_local! { + /// A device supplied by the host rather than opened by the VM. + /// + /// This is how a browser gets a GPU. `loon-wasm` installs a bridge that + /// proxies every call to JavaScript; the VM finds it here and never learns + /// that the work is happening on the other side of a worker boundary. + static INSTALLED: std::cell::RefCell>> = + const { std::cell::RefCell::new(None) }; +} + +/// Install a device for this thread. Replaces any previous one. +pub fn install(device: std::rc::Rc) { + INSTALLED.with(|d| *d.borrow_mut() = Some(device)); +} + +/// Forget the installed device, if any. +pub fn uninstall() { + INSTALLED.with(|d| *d.borrow_mut() = None); +} + +/// The device installed by the host, if there is one. +pub fn installed() -> Option> { + INSTALLED.with(|d| d.borrow().clone()) +} + +/// Convert a buffer to the 32-bit form a device can hold. +/// +/// WGSL core has no 64-bit scalar. A launch with such a buffer is refused +/// before reaching here, so this is only the identity in practice — but it is +/// the one place that would change if a device ever gained wider types. +pub fn narrow(buf: &Buffer) -> Buffer { + use super::vm::BufData; + match &buf.data { + BufData::F64(v) => Buffer { + data: BufData::F32(v.iter().map(|x| *x as f32).collect()), + }, + BufData::I64(v) => Buffer { + data: BufData::I32(v.iter().map(|x| *x as i32).collect()), + }, + _ => buf.clone(), + } +} diff --git a/crates/loon-lang/src/eir/gpu.rs b/crates/loon-lang/src/eir/gpu.rs new file mode 100644 index 0000000..9fb679d --- /dev/null +++ b/crates/loon-lang/src/eir/gpu.rs @@ -0,0 +1,664 @@ +//! Running a kernel on an actual GPU. +//! +//! Everything up to here has been about *where* a kernel could run. This is +//! the part that puts one on hardware: the WGSL from `eir::wgsl` becomes a +//! compute pipeline, buffers become GPU allocations, and a dispatch happens. +//! +//! Through wgpu, so the same kernel reaches Metal on this machine, Vulkan on a +//! Linux box, DX12 on Windows, and WebGPU in a browser tab. That last one is +//! the interesting entry in the list: it is a target a compiler emitting PTX +//! and AMDGCN cannot reach at all, and it comes free from having chosen a +//! portable shading language rather than a vendor's. +//! +//! The device is discovered lazily and its absence is not an error. A machine +//! with no GPU runs the same programs through the CPU path, so nothing in the +//! test suite depends on hardware being present — it only gets *checked* on +//! hardware when hardware is there. + +use super::device::{Device, DeviceError}; +use super::layout::DType; +use super::vm::{BufData, Buffer}; +use super::wgsl::ArgKind; + +impl Device for Gpu { + fn name(&self) -> String { + Gpu::name(self).to_string() + } + + fn ensure_resident(&self, id: usize, buf: &Buffer) -> Result { + Ok(Gpu::ensure_resident(self, id, buf)) + } + + fn is_resident(&self, id: usize) -> bool { + Gpu::is_resident(self, id) + } + + fn dispatch( + &self, + shader: &str, + entry: &str, + n: u32, + scalars: &[f32], + buffers: &[usize], + ) -> Result<(), DeviceError> { + Gpu::dispatch(self, shader, entry, n, scalars, buffers).map_err(|e| DeviceError(e.0)) + } + + fn download(&self, id: usize, byte_len: usize) -> Result, DeviceError> { + Gpu::download(self, id, byte_len).map_err(|e| DeviceError(e.0)) + } + + fn evict(&self, id: usize) { + Gpu::evict(self, id) + } +} + +/// A GPU we can dispatch to. +pub struct Gpu { + device: wgpu::Device, + queue: wgpu::Queue, + name: String, + /// Buffers that currently live on the device, keyed by the host's heap + /// slot. + /// + /// This is what makes residency real rather than modelled. Without it, + /// every launch would upload its arguments again and copy its results back + /// again, and a residency policy would be describing a saving that did not + /// happen. With it, `Place.pin` keeps an allocation alive across launches + /// and `Place.read` is the only thing that moves bytes home. + resident: std::cell::RefCell>, + /// Compiled pipelines, keyed by shader source. + /// + /// Compiling a shader takes on the order of a millisecond, which is + /// hundreds of times what a small launch costs. Without this cache a loop + /// that launches the same kernel repeatedly — the shape every offload + /// benchmark has, and the shape a residency policy exists to serve — pays + /// to recompile it on every iteration, and the GPU loses to an + /// interpreter for entirely uninteresting reasons. + pipelines: std::cell::RefCell>, +} + +/// A shader that has already been compiled. +struct CachedPipeline { + pipeline: wgpu::ComputePipeline, + layout: wgpu::BindGroupLayout, +} + +/// What went wrong, in a sentence a person can act on. +#[derive(Debug, Clone)] +pub struct Error(pub String); + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for Error {} + +/// One argument, as the GPU will see it. +pub enum GpuArg<'a> { + /// A buffer, uploaded before the dispatch and read back after if written. + Buffer { data: &'a Buffer, writable: bool }, + /// A number, packed into the uniform block. + Scalar(f32), +} + +impl Gpu { + /// Find a GPU, or explain why there isn't one. + /// + /// Called once and cached by the caller: adapter enumeration is slow + /// enough that doing it per launch would dominate any kernel worth + /// offloading. + pub fn open() -> Result { + let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default()); + let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::HighPerformance, + compatible_surface: None, + force_fallback_adapter: false, + })) + .map_err(|e| Error(format!("no GPU adapter available: {e}")))?; + + let info = adapter.get_info(); + let name = format!("{} ({:?})", info.name, info.backend); + + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("loon"), + required_features: wgpu::Features::empty(), + required_limits: wgpu::Limits::downlevel_defaults(), + memory_hints: wgpu::MemoryHints::Performance, + trace: wgpu::Trace::Off, + })) + .map_err(|e| Error(format!("could not open the GPU device: {e}")))?; + + Ok(Gpu { + device, + queue, + name, + resident: std::cell::RefCell::new(std::collections::HashMap::new()), + pipelines: std::cell::RefCell::new(std::collections::HashMap::new()), + }) + } + + /// A human-readable name for the device, for `--place-stats` and errors. + pub fn name(&self) -> &str { + &self.name + } + + /// Run `shader` over `n` work items with `args`. + /// + /// Returns the bytes read back from each writable buffer, paired with the + /// argument position it came from. The caller decides what to do with + /// them — which keeps this function free of any opinion about where the + /// host's copy lives, and makes the transfer itself something a test can + /// look at directly. + /// + /// Blocking: the dispatch is submitted and awaited before returning. A + /// non-blocking form belongs with the scheduler, not here. + pub fn run( + &self, + shader: &str, + entry: &str, + n: u32, + args: &[GpuArg<'_>], + ) -> Result)>, Error> { + self.ensure_pipeline(shader, entry); + + // The uniform block is `n` followed by every scalar argument, in order + // — matching what the emitter wrote into `struct Params`. + let mut uniform: Vec = (n as i32).to_le_bytes().to_vec(); + for arg in args.iter() { + if let GpuArg::Scalar(v) = arg { + uniform.extend_from_slice(&v.to_le_bytes()); + } + } + // WGSL requires a uniform buffer to be a multiple of 16 bytes. + while uniform.len() % 16 != 0 { + uniform.push(0); + } + let uniform_buf = self.create_buffer( + "params", + &uniform, + wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + ); + + // Upload every buffer argument, remembering which ones have to come + // back and how big they are. + let mut storage: Vec = Vec::new(); + let mut readback: Vec<(usize, u64)> = Vec::new(); + for (i, arg) in args.iter().enumerate() { + if let GpuArg::Buffer { data, writable } = arg { + let bytes = data.to_bytes(); + let usage = wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_DST + | wgpu::BufferUsages::COPY_SRC; + let buf = self.create_buffer(&format!("b{i}"), &bytes, usage); + if *writable { + readback.push((i, bytes.len() as u64)); + } + storage.push(buf); + } + } + + // Bindings: 0 is the uniform block, storage buffers follow in order. + let mut entries: Vec = vec![wgpu::BindGroupEntry { + binding: 0, + resource: uniform_buf.as_entire_binding(), + }]; + for (slot, buf) in storage.iter().enumerate() { + entries.push(wgpu::BindGroupEntry { + binding: slot as u32 + 1, + resource: buf.as_entire_binding(), + }); + } + + let cache = self.pipelines.borrow(); + let cached = cache.get(shader).expect("pipeline was just ensured"); + + let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("args"), + layout: &cached.layout, + entries: &entries, + }); + + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("run") }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some(entry), + timestamp_writes: None, + }); + pass.set_pipeline(&cached.pipeline); + pass.set_bind_group(0, &bind_group, &[]); + // The shader's workgroup size is 64; round up and let the bounds + // check in the shader discard the overshoot. + pass.dispatch_workgroups(n.div_ceil(64).max(1), 1, 1); + } + + // Copy every written buffer into a staging buffer the host can map. + let staging: Vec<(usize, wgpu::Buffer)> = readback + .iter() + .map(|(i, size)| { + let s = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("readback"), + size: *size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + (*i, s) + }) + .collect(); + for ((_, size), (slot, dst)) in readback.iter().zip(staging.iter()) { + let src = storage_for(args, *slot, &storage); + encoder.copy_buffer_to_buffer(src, 0, dst, 0, *size); + } + + self.queue.submit(Some(encoder.finish())); + drop(cache); + + // Map each staging buffer and collect what came back. + let mut results: Vec<(usize, Vec)> = Vec::new(); + for (slot, dst) in &staging { + let slice = dst.slice(..); + let (tx, rx) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = tx.send(r); + }); + self.device + .poll(wgpu::PollType::Wait) + .map_err(|e| Error(format!("waiting for the GPU: {e:?}")))?; + match rx.recv() { + Ok(Ok(())) => {} + Ok(Err(e)) => return Err(Error(format!("mapping results: {e:?}"))), + Err(e) => return Err(Error(format!("the GPU never reported back: {e}"))), + } + let bytes = slice.get_mapped_range().to_vec(); + dst.unmap(); + results.push((*slot, bytes)); + } + + Ok(results) + } + + /// Ensure `id`'s data is on the device, uploading it if it is not already. + /// + /// Returns whether an upload actually happened, so the caller's accounting + /// reflects what the hardware did rather than what the model predicted. + pub fn ensure_resident(&self, id: usize, buf: &Buffer) -> bool { + if self.resident.borrow().contains_key(&id) { + return false; + } + let bytes = buf.to_bytes(); + let gpu_buf = self.create_buffer( + &format!("buf{id}"), + &bytes, + wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_DST + | wgpu::BufferUsages::COPY_SRC, + ); + self.resident.borrow_mut().insert(id, gpu_buf); + true + } + + /// Whether `id` is currently on the device. + pub fn is_resident(&self, id: usize) -> bool { + self.resident.borrow().contains_key(&id) + } + + /// Drop a device allocation. + pub fn evict(&self, id: usize) { + self.resident.borrow_mut().remove(&id); + } + + /// Copy a resident buffer's contents back to the host. + /// + /// This is the only path by which device data becomes visible again, which + /// is exactly what makes `Place.read` the synchronization point a + /// residency handler can reason about. + pub fn download(&self, id: usize, byte_len: usize) -> Result, Error> { + let resident = self.resident.borrow(); + let Some(src) = resident.get(&id) else { + return Err(Error(format!("buffer {id} is not on the device"))); + }; + let staging = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("readback"), + size: byte_len.max(4) as u64, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("download"), + }); + encoder.copy_buffer_to_buffer(src, 0, &staging, 0, byte_len.max(4) as u64); + self.queue.submit(Some(encoder.finish())); + + let slice = staging.slice(..); + let (tx, rx) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = tx.send(r); + }); + self.device + .poll(wgpu::PollType::Wait) + .map_err(|e| Error(format!("waiting for the GPU: {e:?}")))?; + match rx.recv() { + Ok(Ok(())) => {} + Ok(Err(e)) => return Err(Error(format!("mapping results: {e:?}"))), + Err(e) => return Err(Error(format!("the GPU never reported back: {e}"))), + } + let mut bytes = slice.get_mapped_range().to_vec(); + staging.unmap(); + bytes.truncate(byte_len); + Ok(bytes) + } + + /// Dispatch a shader over buffers that are already resident. + /// + /// Nothing is uploaded and nothing is read back: the caller decides when + /// data moves, which is the whole point of making placement a handler's + /// decision rather than a launch's side effect. + pub fn dispatch( + &self, + shader: &str, + entry: &str, + n: u32, + scalars: &[f32], + buffers: &[usize], + ) -> Result<(), Error> { + self.ensure_pipeline(shader, entry); + + let mut uniform: Vec = (n as i32).to_le_bytes().to_vec(); + for v in scalars { + uniform.extend_from_slice(&v.to_le_bytes()); + } + while uniform.len() % 16 != 0 { + uniform.push(0); + } + let uniform_buf = self.create_buffer( + "params", + &uniform, + wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + ); + + let resident = self.resident.borrow(); + let mut entries: Vec = vec![wgpu::BindGroupEntry { + binding: 0, + resource: uniform_buf.as_entire_binding(), + }]; + for (slot, id) in buffers.iter().enumerate() { + let buf = resident + .get(id) + .ok_or_else(|| Error(format!("buffer {id} is not on the device")))?; + entries.push(wgpu::BindGroupEntry { + binding: slot as u32 + 1, + resource: buf.as_entire_binding(), + }); + } + + let cache = self.pipelines.borrow(); + let cached = cache.get(shader).expect("pipeline was just ensured"); + let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("args"), + layout: &cached.layout, + entries: &entries, + }); + + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("run") }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some(entry), + timestamp_writes: None, + }); + pass.set_pipeline(&cached.pipeline); + pass.set_bind_group(0, &bind_group, &[]); + pass.dispatch_workgroups(n.div_ceil(64).max(1), 1, 1); + } + self.queue.submit(Some(encoder.finish())); + Ok(()) + } + + /// Compile `shader` if this is the first time we have seen it. + fn ensure_pipeline(&self, shader: &str, entry: &str) { + if self.pipelines.borrow().contains_key(shader) { + return; + } + let module = self + .device + .create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some(entry), + source: wgpu::ShaderSource::Wgsl(shader.into()), + }); + let pipeline = self + .device + .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some(entry), + layout: None, + module: &module, + entry_point: Some(entry), + compilation_options: Default::default(), + cache: None, + }); + let layout = pipeline.get_bind_group_layout(0); + self.pipelines + .borrow_mut() + .insert(shader.to_string(), CachedPipeline { pipeline, layout }); + } + + /// How many distinct shaders have been compiled so far. + pub fn compiled_shaders(&self) -> usize { + self.pipelines.borrow().len() + } + + fn create_buffer(&self, label: &str, bytes: &[u8], usage: wgpu::BufferUsages) -> wgpu::Buffer { + use wgpu::util::DeviceExt as _; + // An empty buffer is invalid in wgpu; give it one padded element. + let padded; + let contents = if bytes.is_empty() { + padded = vec![0u8; 4]; + &padded[..] + } else { + bytes + }; + self.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some(label), + contents, + usage, + }) + } +} + +/// The storage buffer belonging to argument `slot`. +/// +/// Scalars occupy no binding, so the storage list is indexed by *buffer* +/// position rather than argument position — the same numbering the emitter +/// used, kept in one place so the two cannot drift. +fn storage_for<'a>( + args: &[GpuArg<'_>], + slot: usize, + storage: &'a [wgpu::Buffer], +) -> &'a wgpu::Buffer { + let idx = args + .iter() + .take(slot) + .filter(|a| matches!(a, GpuArg::Buffer { .. })) + .count(); + &storage[idx] +} + +/// Copy raw device bytes into a host buffer, converting if the buffer's +/// element type is wider than what the GPU computed. +/// +/// WGSL has no f64 or i64, so a 64-bit Loon buffer is computed in 32 bits on +/// the device. Widening on the way back is the honest completion of that: the +/// values really were f32, and pretending otherwise would hide precision the +/// program did not get. +fn write_back(buf: &mut Buffer, bytes: &[u8]) { + match &mut buf.data { + BufData::F32(v) => { + for (slot, chunk) in v.iter_mut().zip(bytes.chunks_exact(4)) { + *slot = f32::from_le_bytes(chunk.try_into().unwrap()); + } + } + BufData::I32(v) => { + for (slot, chunk) in v.iter_mut().zip(bytes.chunks_exact(4)) { + *slot = i32::from_le_bytes(chunk.try_into().unwrap()); + } + } + BufData::F64(v) => { + for (slot, chunk) in v.iter_mut().zip(bytes.chunks_exact(4)) { + *slot = f32::from_le_bytes(chunk.try_into().unwrap()) as f64; + } + } + BufData::I64(v) => { + for (slot, chunk) in v.iter_mut().zip(bytes.chunks_exact(4)) { + *slot = i32::from_le_bytes(chunk.try_into().unwrap()) as i64; + } + } + } +} + +/// The element type a buffer is given to the GPU as. +/// +/// 64-bit buffers are narrowed, because WGSL core has no 64-bit scalar. This +/// is reported rather than assumed: see `DType::gpu_ok`. +pub fn device_dtype(d: DType) -> DType { + match d { + DType::F64 => DType::F32, + DType::I64 => DType::I32, + other => other, + } +} + +/// Convert a buffer to the 32-bit form the device will hold. +pub fn narrow(buf: &Buffer) -> Buffer { + match &buf.data { + BufData::F64(v) => Buffer { + data: BufData::F32(v.iter().map(|x| *x as f32).collect()), + }, + BufData::I64(v) => Buffer { + data: BufData::I32(v.iter().map(|x| *x as i32).collect()), + }, + _ => buf.clone(), + } +} + +/// Describe an argument for the emitter, given what it actually is. +pub fn arg_kind(arg: &GpuArg<'_>) -> ArgKind { + match arg { + GpuArg::Buffer { data, writable } => ArgKind::Buffer { + dtype: device_dtype(data.dtype()), + writable: *writable, + }, + GpuArg::Scalar(_) => ArgKind::Scalar(DType::F32), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Open a GPU, or skip. A machine without one is not a failing machine. + fn gpu_or_skip() -> Option { + match Gpu::open() { + Ok(g) => Some(g), + Err(e) => { + println!("SKIPPED: {e}"); + None + } + } + } + + const DOUBLE: &str = r#" +struct Params { n: i32, s0: f32, }; +@group(0) @binding(0) var params: Params; +@group(0) @binding(1) var b0: array; +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx: i32 = i32(gid.x); + if (idx >= params.n) { return; } + b0[u32(idx)] = b0[u32(idx)] * params.s0; +} +"#; + + #[test] + fn a_kernel_runs_on_the_gpu_and_the_results_come_back() { + let Some(gpu) = gpu_or_skip() else { return }; + println!("device: {}", gpu.name()); + + let mut buf = Buffer { + data: BufData::F32(vec![1.0, 2.0, 3.0, 4.0]), + }; + let input = buf.clone(); + let args = vec![ + GpuArg::Scalar(3.0), + GpuArg::Buffer { + data: &input, + writable: true, + }, + ]; + let results = gpu.run(DOUBLE, "main", 4, &args).expect("dispatch"); + assert_eq!(results.len(), 1, "one writable buffer comes back"); + write_back(&mut buf, &results[0].1); + assert_eq!(buf.data, BufData::F32(vec![3.0, 6.0, 9.0, 12.0])); + } + + #[test] + fn work_that_does_not_fill_a_workgroup_still_stays_in_bounds() { + // A dispatch rounds up to whole workgroups, so a 3-element buffer runs + // 64 invocations. The bounds check in the shader is what keeps the + // other 61 from writing past the end. + let Some(gpu) = gpu_or_skip() else { return }; + let mut buf = Buffer { + data: BufData::F32(vec![1.0, 1.0, 1.0]), + }; + let input = buf.clone(); + let args = vec![ + GpuArg::Scalar(5.0), + GpuArg::Buffer { + data: &input, + writable: true, + }, + ]; + let results = gpu.run(DOUBLE, "main", 3, &args).expect("dispatch"); + write_back(&mut buf, &results[0].1); + assert_eq!(buf.data, BufData::F32(vec![5.0, 5.0, 5.0])); + } + + #[test] + fn a_read_only_buffer_is_not_copied_back() { + // Only what the kernel writes has to travel home. This is the same + // distinction the ownership pass draws, arriving at the hardware. + let Some(gpu) = gpu_or_skip() else { return }; + let src = Buffer { + data: BufData::F32(vec![1.0, 2.0]), + }; + let args = vec![ + GpuArg::Scalar(2.0), + GpuArg::Buffer { + data: &src, + writable: false, + }, + ]; + // The shader writes b0, but we declared it read-only here, so nothing + // is read back — the accounting follows the declaration. + let results = gpu.run(DOUBLE, "main", 2, &args).expect("dispatch"); + assert!(results.is_empty(), "nothing was declared writable"); + } + + #[test] + fn narrowing_is_explicit_about_what_the_device_can_hold() { + assert_eq!(device_dtype(DType::F64), DType::F32); + assert_eq!(device_dtype(DType::I64), DType::I32); + assert_eq!(device_dtype(DType::F32), DType::F32); + + let wide = Buffer { + data: BufData::F64(vec![1.5, 2.5]), + }; + assert_eq!(narrow(&wide).data, BufData::F32(vec![1.5, 2.5])); + } +} diff --git a/crates/loon-lang/src/eir/kernel_exec.rs b/crates/loon-lang/src/eir/kernel_exec.rs new file mode 100644 index 0000000..002bb25 --- /dev/null +++ b/crates/loon-lang/src/eir/kernel_exec.rs @@ -0,0 +1,855 @@ +//! Running a kernel on the CPU without going through the general VM. +//! +//! The interpreter is built for a dynamically-typed language: every value is +//! NaN-boxed, every buffer access goes through the heap table, and every +//! arithmetic operation checks what it was handed. That is the right design +//! for Loon and the wrong one for a loop that runs the same three floating +//! point operations a million times. +//! +//! A kernel does not need any of it. The subset is small enough that every +//! register has a knowable type, every buffer is a contiguous slice, and the +//! whole body can be executed against raw `f32`s. This module does that: it +//! type-checks the kernel once, then runs it per work item with no boxing and +//! no heap indirection. +//! +//! Two things follow. The CPU number in a benchmark becomes a fair baseline +//! rather than a straw man — "the GPU is faster than our interpreter" is not +//! an interesting claim. And because the executor works on plain slices over +//! an index range, running work items in parallel is a matter of splitting the +//! range, which is where `Mode::Par` comes from. + +use super::layout::DType; +use super::vm::{BufData, Buffer}; +use super::{BinOp, Built, End, FuncId, Lit, Module, Op, Reg, UnOp}; + +/// A value inside a running kernel. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum KVal { + F(f64), + I(i64), + B(bool), + Unit, +} + +impl KVal { + fn as_f(self) -> f64 { + match self { + KVal::F(x) => x, + KVal::I(n) => n as f64, + KVal::B(b) => { + if b { + 1.0 + } else { + 0.0 + } + } + KVal::Unit => 0.0, + } + } + + fn as_i(self) -> i64 { + match self { + KVal::I(n) => n, + KVal::F(x) => x as i64, + KVal::B(b) => b as i64, + KVal::Unit => 0, + } + } + + fn truthy(self) -> bool { + match self { + KVal::B(b) => b, + KVal::I(n) => n != 0, + KVal::F(x) => x != 0.0, + KVal::Unit => false, + } + } + + /// Whether either operand is a float, which decides the result's type. + fn either_float(a: KVal, b: KVal) -> bool { + matches!(a, KVal::F(_)) || matches!(b, KVal::F(_)) + } +} + +/// A writable slice of one buffer's elements. +/// +/// Parallel placement hands each thread a different one of these, carved out +/// of the same buffer with `split_at_mut`. That the pieces do not overlap is +/// not a promise anybody makes — it is what `split_at_mut` returns, checked by +/// the compiler. A GPU partitioning strategy has to assert the same property +/// through an `unsafe impl`. +pub enum OutView<'a> { + F32(&'a mut [f32]), + F64(&'a mut [f64]), + I32(&'a mut [i32]), + I64(&'a mut [i64]), +} + +impl OutView<'_> { + fn len(&self) -> usize { + match self { + OutView::F32(v) => v.len(), + OutView::F64(v) => v.len(), + OutView::I32(v) => v.len(), + OutView::I64(v) => v.len(), + } + } +} + +/// An argument to a kernel launch. +pub enum KArg<'a> { + /// A buffer the kernel only reads. + Input(&'a Buffer), + /// A buffer the kernel writes through. + Output(&'a mut Buffer), + /// A slice of a buffer, addressed by absolute index. + /// + /// `base` is the index this view starts at, so a work item writing at its + /// own index lands in the right place. A write outside the view is an + /// error rather than a silent miss: in parallel placement it means the + /// kernel wrote somewhere another thread owns. + OutputView { view: OutView<'a>, base: i64 }, + /// A number. + Scalar(KVal), +} + +/// Why a kernel could not be run this way. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Error(pub String); + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +/// Can this function be run by the fast executor? +/// +/// Checked once before a launch rather than per work item. A kernel that fails +/// here is not wrong — it simply uses something outside the subset, and the +/// general VM runs it instead. +pub fn supported(module: &Module, func: FuncId) -> bool { + let Some(f) = module.funcs.get(func.0 as usize) else { + return false; + }; + f.blocks.iter().all(|b| { + b.ops.iter().all(|op| { + matches!(op, Op::Lit(..) | Op::Mov(..) | Op::Bin(..) | Op::Un(..)) + || matches!(op, Op::Builtin(_, built, ..) if supported_builtin(*built)) + }) && matches!( + b.end, + End::Ret(_) | End::Jmp(..) | End::Br(..) | End::Recur(_) | End::Trap + ) + }) +} + +fn supported_builtin(b: Built) -> bool { + matches!( + b, + Built::BufAt + | Built::BufPut + | Built::BufLen + | Built::Sqrt + | Built::Pow + | Built::Floor + | Built::Ceil + | Built::Round + | Built::Sin + | Built::Cos + | Built::Tan + | Built::Asin + | Built::Acos + | Built::Atan + | Built::Atan2 + | Built::Log + | Built::Log10 + | Built::Exp + | Built::Abs + | Built::Min + | Built::Max + | Built::Not + ) +} + +/// Run `func` for every index in `range`. +/// +/// Buffer arguments are borrowed for the whole call, which is what makes the +/// disjointness argument work: an output buffer is a `&mut` slice, so the type +/// system will not let two ranges write the same one at once. That is the same +/// guarantee a GPU partitioning strategy has to promise by hand. +pub fn run_range( + module: &Module, + func: FuncId, + args: &mut [KArg<'_>], + range: std::ops::Range, +) -> Result<(), Error> { + let f = module + .funcs + .get(func.0 as usize) + .ok_or_else(|| Error(format!("no function {func:?}")))?; + if f.params.len() != args.len() + 1 { + return Err(Error(format!( + "kernel takes {} parameters but {} arguments were given", + f.params.len(), + args.len() + 1 + ))); + } + + let reg_count = f + .blocks + .iter() + .flat_map(|b| { + b.ops + .iter() + .filter_map(dest) + .chain(b.params.iter().copied()) + }) + .map(|r| r.0 as usize + 1) + .max() + .unwrap_or(1) + .max(args.len() + 1); + + let mut regs = vec![KVal::Unit; reg_count]; + for i in range { + run_one(module, f.id, args, &mut regs, i)?; + } + Ok(()) +} + +/// Run a kernel across every core, splitting the index space. +/// +/// Each thread gets a disjoint slice of *every* output buffer and the whole of +/// every input. The disjointness is not asserted — it is what `split_at_mut` +/// hands back, and the borrow checker is what enforces it. A GPU partitioning +/// strategy has to promise the same property through an `unsafe impl`. +/// +/// A kernel that writes outside its own range fails with a message saying so +/// rather than racing another thread. +pub fn run_parallel( + module: &Module, + func: FuncId, + scalars: &[(usize, KVal)], + inputs: &[(usize, &Buffer)], + outputs: &mut [(usize, &mut Buffer)], + arity: usize, + n: i64, +) -> Result<(), Error> { + let threads = std::thread::available_parallelism() + .map(|p| p.get()) + .unwrap_or(1) + .min(n.max(1) as usize); + + if threads <= 1 || n <= 0 || outputs.is_empty() { + // Nothing to split, or nothing to split across. Sequential is the same + // answer; only slower. + let mut args = assemble(scalars, inputs, outputs, arity)?; + return run_range(module, func, &mut args, 0..n); + } + + // Contiguous chunks of the index space, one per thread. + let chunk = (n as usize).div_ceil(threads); + let mut ranges: Vec<(i64, usize)> = Vec::new(); + let mut start = 0usize; + while start < n as usize { + let len = chunk.min(n as usize - start); + ranges.push((start as i64, len)); + start += len; + } + + // Carve every output buffer into per-thread views. `iter_mut` yields + // disjoint `&mut`s, so pieces taken from different buffers can coexist. + let mut per_thread: Vec> = ranges.iter().map(|_| Vec::new()).collect(); + for (arg_idx, buf) in outputs.iter_mut() { + if buf.len() < n as usize { + return Err(Error(format!( + "argument {} has {} elements but the launch covers {n}", + *arg_idx + 1, + buf.len() + ))); + } + macro_rules! carve { + ($v:expr, $variant:ident) => {{ + let mut rest: &mut [_] = &mut $v[..]; + for (t, (_, len)) in ranges.iter().enumerate() { + let (head, tail) = rest.split_at_mut(*len); + per_thread[t].push((*arg_idx, OutView::$variant(head))); + rest = tail; + } + }}; + } + match &mut buf.data { + BufData::F32(v) => carve!(v, F32), + BufData::F64(v) => carve!(v, F64), + BufData::I32(v) => carve!(v, I32), + BufData::I64(v) => carve!(v, I64), + } + } + + std::thread::scope(|scope| { + let mut handles = Vec::new(); + for ((base, _), views) in ranges.iter().zip(per_thread) { + let base = *base; + handles.push(scope.spawn(move || { + let mut args: Vec = Vec::with_capacity(arity); + let mut views = views; + for i in 0..arity { + if let Some((_, v)) = scalars.iter().find(|(j, _)| *j == i) { + args.push(KArg::Scalar(*v)); + } else if let Some((_, b)) = inputs.iter().find(|(j, _)| *j == i) { + args.push(KArg::Input(b)); + } else if let Some(pos) = views.iter().position(|(j, _)| *j == i) { + let (_, view) = views.remove(pos); + args.push(KArg::OutputView { view, base }); + } else { + return Err(Error(format!("argument {} was not provided", i + 1))); + } + } + let len = args + .iter() + .find_map(|a| match a { + KArg::OutputView { view, .. } => Some(view.len() as i64), + _ => None, + }) + .unwrap_or(0); + run_range(module, func, &mut args, base..base + len) + })); + } + let mut result = Ok(()); + for h in handles { + match h.join() { + Ok(Ok(())) => {} + Ok(Err(e)) => result = Err(e), + Err(_) => result = Err(Error("a kernel thread panicked".into())), + } + } + result + }) +} + +/// Run a kernel on one core, from the same separated pieces `run_parallel` +/// takes, so a caller does not have to assemble arguments two different ways. +pub fn run_sequential( + module: &Module, + func: FuncId, + scalars: &[(usize, KVal)], + inputs: &[(usize, &Buffer)], + outputs: &mut [(usize, &mut Buffer)], + arity: usize, + n: i64, +) -> Result<(), Error> { + let mut args = assemble(scalars, inputs, outputs, arity)?; + run_range(module, func, &mut args, 0..n) +} + +/// Build a sequential argument list from the separated pieces. +fn assemble<'a>( + scalars: &[(usize, KVal)], + inputs: &[(usize, &'a Buffer)], + outputs: &'a mut [(usize, &mut Buffer)], + arity: usize, +) -> Result>, Error> { + // Outputs are addressed by argument position, so index them once. + let mut out_by_arg: Vec> = (0..arity).map(|_| None).collect(); + for (i, b) in outputs.iter_mut() { + out_by_arg[*i] = Some(b); + } + let mut args: Vec = Vec::with_capacity(arity); + for i in 0..arity { + if let Some((_, v)) = scalars.iter().find(|(j, _)| *j == i) { + args.push(KArg::Scalar(*v)); + } else if let Some((_, b)) = inputs.iter().find(|(j, _)| *j == i) { + args.push(KArg::Input(b)); + } else if let Some(b) = out_by_arg[i].take() { + args.push(KArg::Output(b)); + } else { + return Err(Error(format!("argument {} was not provided", i + 1))); + } + } + Ok(args) +} + +fn run_one( + module: &Module, + func: FuncId, + args: &mut [KArg<'_>], + regs: &mut [KVal], + index: i64, +) -> Result<(), Error> { + let f = &module.funcs[func.0 as usize]; + + // Parameters: the work index, then the arguments. Buffers do not occupy a + // register value — they are addressed by argument position. + regs[0] = KVal::I(index); + for (i, arg) in args.iter().enumerate() { + regs[i + 1] = match arg { + KArg::Scalar(v) => *v, + // A buffer's "value" is its own position, so `at`/`put` can find + // it. Nothing else in the subset reads a buffer as a number. + _ => KVal::I(i as i64), + }; + } + + let mut block = 0usize; + let mut steps = 0usize; + loop { + steps += 1; + if steps > 1_000_000 { + return Err(Error("kernel did not terminate".into())); + } + let b = &f.blocks[block]; + for op in &b.ops { + exec_op(op, args, regs)?; + } + match &b.end { + End::Ret(_) | End::Trap => return Ok(()), + End::Jmp(target, jargs) => { + let target_block = &f.blocks[target.0 as usize]; + let vals: Vec = jargs.iter().map(|r| regs[r.0 as usize]).collect(); + for (p, v) in target_block.params.iter().zip(vals) { + regs[p.0 as usize] = v; + } + block = target.0 as usize; + } + End::Br(cond, t, e) => { + block = if regs[cond.0 as usize].truthy() { + t.0 as usize + } else { + e.0 as usize + }; + } + End::Recur(rargs) => { + let vals: Vec = rargs.iter().map(|r| regs[r.0 as usize]).collect(); + let params = b.params.clone(); + for (p, v) in params.iter().zip(vals) { + regs[p.0 as usize] = v; + } + } + other => return Err(Error(format!("{other:?} is outside the kernel subset"))), + } + } +} + +fn exec_op(op: &Op, args: &mut [KArg<'_>], regs: &mut [KVal]) -> Result<(), Error> { + match op { + Op::Lit(d, lit, _) => { + regs[d.0 as usize] = match lit { + Lit::Int(n) => KVal::I(*n), + Lit::Float(x) => KVal::F(*x), + Lit::Bool(b) => KVal::B(*b), + Lit::Unit => KVal::Unit, + other => return Err(Error(format!("{other:?} is outside the kernel subset"))), + }; + } + Op::Mov(d, s, _) => regs[d.0 as usize] = regs[s.0 as usize], + Op::Bin(d, bop, a, b, _) => { + let (x, y) = (regs[a.0 as usize], regs[b.0 as usize]); + regs[d.0 as usize] = binop(*bop, x, y)?; + } + Op::Un(d, uop, a, _) => { + let v = regs[a.0 as usize]; + regs[d.0 as usize] = match uop { + UnOp::Neg => match v { + KVal::I(n) => KVal::I(-n), + other => KVal::F(-other.as_f()), + }, + UnOp::Not => KVal::B(!v.truthy()), + }; + } + Op::Builtin(d, built, bargs, _) => exec_builtin(*d, *built, bargs, args, regs)?, + other => return Err(Error(format!("{other:?} is outside the kernel subset"))), + } + Ok(()) +} + +fn binop(bop: BinOp, x: KVal, y: KVal) -> Result { + let float = KVal::either_float(x, y); + Ok(match bop { + BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::Rem => { + if float { + let (a, b) = (x.as_f(), y.as_f()); + KVal::F(match bop { + BinOp::Add => a + b, + BinOp::Sub => a - b, + BinOp::Mul => a * b, + BinOp::Div => a / b, + _ => a % b, + }) + } else { + let (a, b) = (x.as_i(), y.as_i()); + // Integer division by zero has a decided answer in Loon (it is + // an error, not a silent value); the kernel subset does not get + // to invent a different one, so the executor declines instead. + if matches!(bop, BinOp::Div | BinOp::Rem) && b == 0 { + return Err(Error("integer division by zero in a kernel".into())); + } + KVal::I(match bop { + BinOp::Add => a + b, + BinOp::Sub => a - b, + BinOp::Mul => a * b, + BinOp::Div => a / b, + _ => a % b, + }) + } + } + BinOp::Eq => KVal::B(if float { + x.as_f() == y.as_f() + } else { + x.as_i() == y.as_i() + }), + BinOp::Ne => KVal::B(if float { + x.as_f() != y.as_f() + } else { + x.as_i() != y.as_i() + }), + BinOp::Lt => KVal::B(if float { + x.as_f() < y.as_f() + } else { + x.as_i() < y.as_i() + }), + BinOp::Gt => KVal::B(if float { + x.as_f() > y.as_f() + } else { + x.as_i() > y.as_i() + }), + BinOp::Le => KVal::B(if float { + x.as_f() <= y.as_f() + } else { + x.as_i() <= y.as_i() + }), + BinOp::Ge => KVal::B(if float { + x.as_f() >= y.as_f() + } else { + x.as_i() >= y.as_i() + }), + BinOp::And => KVal::B(x.truthy() && y.truthy()), + BinOp::Or => KVal::B(x.truthy() || y.truthy()), + BinOp::Concat => return Err(Error("a kernel cannot concatenate strings".into())), + }) +} + +fn exec_builtin( + d: Reg, + built: Built, + bargs: &[Reg], + args: &mut [KArg<'_>], + regs: &mut [KVal], +) -> Result<(), Error> { + let slot = |i: usize| -> usize { regs[bargs[i].0 as usize].as_i() as usize }; + + match built { + Built::BufAt => { + let which = slot(0); + let idx = regs[bargs[1].0 as usize].as_i(); + let v = read_buffer(args, which, idx)?; + regs[d.0 as usize] = v; + } + Built::BufPut => { + let which = slot(0); + let idx = regs[bargs[1].0 as usize].as_i(); + let v = regs[bargs[2].0 as usize]; + write_buffer(args, which, idx, v)?; + regs[d.0 as usize] = KVal::Unit; + } + Built::BufLen => { + let which = slot(0); + let len = match args.get(which) { + Some(KArg::Input(b)) => b.len(), + Some(KArg::Output(b)) => b.len(), + Some(KArg::OutputView { view, .. }) => view.len(), + _ => return Err(Error("buf-len needs a buffer".into())), + }; + regs[d.0 as usize] = KVal::I(len as i64); + } + _ => { + let a = |i: usize| regs[bargs.get(i).map(|r| r.0 as usize).unwrap_or(0)].as_f(); + regs[d.0 as usize] = match built { + Built::Sqrt => KVal::F(a(0).sqrt()), + Built::Pow => KVal::F(a(0).powf(a(1))), + Built::Floor => KVal::F(a(0).floor()), + Built::Ceil => KVal::F(a(0).ceil()), + Built::Round => KVal::F(a(0).round()), + Built::Sin => KVal::F(a(0).sin()), + Built::Cos => KVal::F(a(0).cos()), + Built::Tan => KVal::F(a(0).tan()), + Built::Asin => KVal::F(a(0).asin()), + Built::Acos => KVal::F(a(0).acos()), + Built::Atan => KVal::F(a(0).atan()), + Built::Atan2 => KVal::F(a(0).atan2(a(1))), + Built::Log => KVal::F(a(0).ln()), + Built::Log10 => KVal::F(a(0).log10()), + Built::Exp => KVal::F(a(0).exp()), + Built::Abs => match regs[bargs[0].0 as usize] { + KVal::I(n) => KVal::I(n.abs()), + other => KVal::F(other.as_f().abs()), + }, + Built::Min | Built::Max => { + let (x, y) = (regs[bargs[0].0 as usize], regs[bargs[1].0 as usize]); + let take_min = matches!(built, Built::Min); + if KVal::either_float(x, y) { + let (p, q) = (x.as_f(), y.as_f()); + KVal::F(if take_min { p.min(q) } else { p.max(q) }) + } else { + let (p, q) = (x.as_i(), y.as_i()); + KVal::I(if take_min { p.min(q) } else { p.max(q) }) + } + } + Built::Not => KVal::B(!regs[bargs[0].0 as usize].truthy()), + other => return Err(Error(format!("builtin {other:?} is outside the subset"))), + }; + } + } + Ok(()) +} + +fn read_buffer(args: &[KArg<'_>], which: usize, idx: i64) -> Result { + if let Some(KArg::OutputView { view, base }) = args.get(which) { + let rel = idx - base; + let i = usize::try_from(rel).map_err(|_| out_of_view(idx))?; + return match view { + OutView::F32(v) => v.get(i).map(|x| KVal::F(*x as f64)), + OutView::F64(v) => v.get(i).map(|x| KVal::F(*x)), + OutView::I32(v) => v.get(i).map(|x| KVal::I(*x as i64)), + OutView::I64(v) => v.get(i).map(|x| KVal::I(*x)), + } + .ok_or_else(|| out_of_view(idx)); + } + let buf = match args.get(which) { + Some(KArg::Input(b)) => *b, + Some(KArg::Output(b)) => &**b, + _ => return Err(Error(format!("argument {} is not a buffer", which + 1))), + }; + let i = usize::try_from(idx).map_err(|_| out_of_range(idx, buf.len()))?; + match &buf.data { + BufData::F32(v) => v.get(i).map(|x| KVal::F(*x as f64)), + BufData::F64(v) => v.get(i).map(|x| KVal::F(*x)), + BufData::I32(v) => v.get(i).map(|x| KVal::I(*x as i64)), + BufData::I64(v) => v.get(i).map(|x| KVal::I(*x)), + } + .ok_or_else(|| out_of_range(idx, buf.len())) +} + +fn write_buffer(args: &mut [KArg<'_>], which: usize, idx: i64, v: KVal) -> Result<(), Error> { + if let Some(KArg::OutputView { view, base }) = args.get_mut(which) { + let rel = idx - *base; + let i = usize::try_from(rel).map_err(|_| out_of_view(idx))?; + if i >= view.len() { + return Err(out_of_view(idx)); + } + match view { + OutView::F32(b) => b[i] = v.as_f() as f32, + OutView::F64(b) => b[i] = v.as_f(), + OutView::I32(b) => b[i] = v.as_i() as i32, + OutView::I64(b) => b[i] = v.as_i(), + } + return Ok(()); + } + let Some(KArg::Output(buf)) = args.get_mut(which) else { + return Err(Error(format!( + "argument {} is not a buffer this kernel may write", + which + 1 + ))); + }; + let len = buf.len(); + let i = usize::try_from(idx).map_err(|_| out_of_range(idx, len))?; + if i >= len { + return Err(out_of_range(idx, len)); + } + match &mut buf.data { + BufData::F32(b) => b[i] = v.as_f() as f32, + BufData::F64(b) => b[i] = v.as_f(), + BufData::I32(b) => b[i] = v.as_i() as i32, + BufData::I64(b) => b[i] = v.as_i(), + } + Ok(()) +} + +fn out_of_range(idx: i64, len: usize) -> Error { + Error(format!("index {idx} is outside a buffer of length {len}")) +} + +/// A write that left the slice this thread owns. +/// +/// In parallel placement that means the kernel wrote at an index other than +/// its own — a scatter. Saying so is the point: the alternative is threads +/// quietly overwriting each other's elements. +fn out_of_view(idx: i64) -> Error { + Error(format!( + "index {idx} is outside the range this work item owns; a kernel run in \ + parallel must write at its own index" + )) +} + +/// Element type of a buffer, for callers deciding how to narrow it. +pub fn dtype_of(b: &Buffer) -> DType { + b.dtype() +} + +fn dest(op: &Op) -> Option { + match op { + Op::Lit(d, ..) + | Op::Mov(d, ..) + | Op::Bin(d, ..) + | Op::Un(d, ..) + | Op::Builtin(d, ..) + | Op::Upval(d, ..) + | Op::Call(d, ..) + | Op::Invoke(d, ..) + | Op::Close(d, ..) + | Op::Vec(d, ..) + | Op::Map(d, ..) + | Op::Set(d, ..) + | Op::Tup(d, ..) + | Op::Adt(d, ..) + | Op::Field(d, ..) + | Op::Tag(d, ..) + | Op::Perform(d, ..) + | Op::PushHandler(d, ..) => Some(*d), + Op::PopHandler(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::check::Checker; + use crate::eir::lower::lower; + use crate::parser::parse; + + fn lowered(src: &str, name: &str) -> (Module, FuncId) { + let exprs = parse(src).expect("parses"); + let mut checker = Checker::new(); + let errors = checker.check_program(&exprs); + assert!(errors.is_empty(), "check errors: {errors:?}"); + let module = lower(&checker); + let id = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(name)) + .expect("kernel lowered") + .id; + (module, id) + } + + fn f32s(b: &Buffer) -> Vec { + match &b.data { + BufData::F32(v) => v.clone(), + other => panic!("expected f32, got {other:?}"), + } + } + + #[test] + fn a_kernel_runs_over_a_range() { + let (m, id) = lowered( + "[kernel saxpy [i a x y out] [put out i [+ [* a [at x i]] [at y i]]]]", + "saxpy", + ); + let x = Buffer { + data: BufData::F32(vec![1.0, 2.0, 3.0, 4.0]), + }; + let y = Buffer { + data: BufData::F32(vec![10.0, 20.0, 30.0, 40.0]), + }; + let mut out = Buffer { + data: BufData::F32(vec![0.0; 4]), + }; + { + let mut args = vec![ + KArg::Scalar(KVal::F(2.0)), + KArg::Input(&x), + KArg::Input(&y), + KArg::Output(&mut out), + ]; + run_range(&m, id, &mut args, 0..4).expect("runs"); + } + assert_eq!(f32s(&out), vec![12.0, 24.0, 36.0, 48.0]); + } + + #[test] + fn a_partial_range_touches_only_its_own_elements() { + // This is the property that makes splitting a range across threads + // sound: a work item writes at its own index and nowhere else. + let (m, id) = lowered("[kernel fill [i b] [put b i 9.0]]", "fill"); + let mut out = Buffer { + data: BufData::F32(vec![0.0; 6]), + }; + { + let mut args = vec![KArg::Output(&mut out)]; + run_range(&m, id, &mut args, 2..4).expect("runs"); + } + assert_eq!(f32s(&out), vec![0.0, 0.0, 9.0, 9.0, 0.0, 0.0]); + } + + #[test] + fn control_flow_inside_a_kernel_works() { + let (m, id) = lowered( + "[kernel clamp [i lo hi b] \ + [let v [at b i]] \ + [put b i [if [< v lo] lo [if [> v hi] hi v]]]]", + "clamp", + ); + let mut b = Buffer { + data: BufData::F32(vec![-5.0, 0.5, 9.0]), + }; + { + let mut args = vec![ + KArg::Scalar(KVal::F(0.0)), + KArg::Scalar(KVal::F(1.0)), + KArg::Output(&mut b), + ]; + run_range(&m, id, &mut args, 0..3).expect("runs"); + } + assert_eq!(f32s(&b), vec![0.0, 0.5, 1.0]); + } + + #[test] + fn math_builtins_agree_with_the_interpreter() { + let src = "[kernel mathy [i b] [put b i [sqrt [abs [at b i]]]]]"; + let (m, id) = lowered(src, "mathy"); + let mut b = Buffer { + data: BufData::F32(vec![-4.0, 9.0, 16.0]), + }; + { + let mut args = vec![KArg::Output(&mut b)]; + run_range(&m, id, &mut args, 0..3).expect("runs"); + } + assert_eq!(f32s(&b), vec![2.0, 3.0, 4.0]); + } + + #[test] + fn reading_past_the_end_is_an_error_here_too() { + // The fast path must not become the path where bounds stop being + // checked. An out-of-range read is a bug in the program whichever + // executor runs it. + let (m, id) = lowered("[kernel bad [i b] [put b i [at b 99]]]", "bad"); + let mut b = Buffer { + data: BufData::F32(vec![0.0; 2]), + }; + let mut args = vec![KArg::Output(&mut b)]; + let e = run_range(&m, id, &mut args, 0..2).expect_err("should fail"); + assert!(e.0.contains("outside a buffer"), "{}", e.0); + } + + #[test] + fn integer_buffers_stay_integers() { + let (m, id) = lowered("[kernel bump [i b] [put b i [+ 1 [at b i]]]]", "bump"); + let mut b = Buffer { + data: BufData::I32(vec![1, 2, 3]), + }; + { + let mut args = vec![KArg::Output(&mut b)]; + run_range(&m, id, &mut args, 0..3).expect("runs"); + } + assert_eq!(b.data, BufData::I32(vec![2, 3, 4])); + } + + #[test] + fn support_is_decided_before_the_launch_not_during_it() { + let (m, id) = lowered("[kernel ok [i b] [put b i 1.0]]", "ok"); + assert!(supported(&m, id)); + + // A function using something outside the subset is declined, so the + // caller can fall back to the general VM rather than fail. + let (m2, id2) = lowered("[fn nope [i] [str i]] [fn main [] []]", "nope"); + assert!(!supported(&m2, id2)); + } +} diff --git a/crates/loon-lang/src/eir/layout.rs b/crates/loon-lang/src/eir/layout.rs new file mode 100644 index 0000000..41f861c --- /dev/null +++ b/crates/loon-lang/src/eir/layout.rs @@ -0,0 +1,361 @@ +//! The one description of how a Loon value looks in memory. +//! +//! Every backend — the register VM, the WASM emitter, the Cranelift JIT, and +//! (later) the WGSL kernel emitter — agrees on the bits described here. Before +//! this module the NaN-boxing constants were copy-pasted into three files with +//! a "must match value64.rs" comment on top; a divergence between two backends +//! was a runtime parity failure at best and a silent wrong answer at worst. +//! That is exactly the bug class the Rust GPU-offload work hit when a slice +//! lowered as `(ptr, len)` on one target and `[i64; 2]` on another, and found +//! it by hand. Here it is one `const`, and `tests/abi_conformance.rs` proves +//! every backend still round-trips the same bytes. +//! +//! Two things live here: +//! - [`nanbox`]: the 64-bit immediate encoding shared by all backends. +//! - [`DType`] / [`BufferHeader`]: the dense-buffer layout that kernels and +//! devices exchange, where "device" may be another thread, another +//! backend, or a GPU queue. + +/// NaN-boxing bit layout. See `value64.rs` for the prose version. +pub mod nanbox { + /// Quiet-NaN base. + pub const QNAN: u64 = 0x7FF8_0000_0000_0000; + /// Sign bit. + pub const SIGN: u64 = 0x8000_0000_0000_0000; + /// Every tagged (non-float) value carries this prefix. + pub const BASE: u64 = SIGN | QNAN; + /// Tag bits 48-50. + pub const TAG_MASK: u64 = 0x0007_0000_0000_0000; + /// Payload bits 0-47. + pub const PAYLOAD: u64 = 0x0000_FFFF_FFFF_FFFF; + + /// Heap pointer (string, closure, ADT, collection, buffer — type in header). + pub const TAG_PTR: u64 = 0x0000_0000_0000_0000; + /// Inline 48-bit signed integer. + pub const TAG_INT: u64 = 0x0001_0000_0000_0000; + /// Interned symbol/keyword (32-bit intern index). + pub const TAG_SYM: u64 = 0x0006_0000_0000_0000; + /// Immediate singleton (Unit, True, False, None). + pub const TAG_IMM: u64 = 0x0007_0000_0000_0000; + + /// Immediate sub-tags, in the low bits. + pub const IMM_UNIT: u64 = 0; + pub const IMM_TRUE: u64 = 1; + pub const IMM_FALSE: u64 = 2; + pub const IMM_NONE: u64 = 3; + + /// Fully-assembled immediates. Backends that emit literals use these + /// directly rather than re-deriving `BASE | TAG_IMM | n`. + pub const VAL_UNIT: u64 = BASE | TAG_IMM | IMM_UNIT; + pub const VAL_TRUE: u64 = BASE | TAG_IMM | IMM_TRUE; + pub const VAL_FALSE: u64 = BASE | TAG_IMM | IMM_FALSE; + pub const VAL_NONE: u64 = BASE | TAG_IMM | IMM_NONE; + + /// Encode a 48-bit signed integer the way every backend must. + #[inline(always)] + pub const fn encode_int(n: i64) -> u64 { + BASE | TAG_INT | ((n as u64) & PAYLOAD) + } + + /// Encode a float (identity — floats pass through as raw IEEE 754). + #[inline(always)] + pub fn encode_float(f: f64) -> u64 { + f.to_bits() + } +} + +// ─── Dense buffers ───────────────────────────────────────────────────────── + +/// Element type of a dense [`BufferHeader`]. +/// +/// Deliberately small. `F64`/`I64` exist because Loon's own numbers are 64-bit +/// and the CPU backends can honour them exactly; they are rejected for GPU +/// placement because WGSL core has neither (see [`DType::gpu_ok`]). A kernel +/// that wants to run on a GPU says so by using f32/i32 buffers — we never +/// silently demote precision behind the programmer's back. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[repr(u32)] +pub enum DType { + F32 = 0, + F64 = 1, + I32 = 2, + I64 = 3, +} + +impl DType { + /// Size of one element in bytes. + pub const fn size(self) -> usize { + match self { + DType::F32 | DType::I32 => 4, + DType::F64 | DType::I64 => 8, + } + } + + /// The WGSL scalar type name, for the kernel emitter. + pub const fn wgsl(self) -> &'static str { + match self { + DType::F32 => "f32", + DType::F64 => "f64", // not valid in WGSL core; guarded by `gpu_ok` + DType::I32 => "i32", + DType::I64 => "i64", // ditto + } + } + + /// Whether a buffer of this type can be placed on a GPU. WGSL core has no + /// 64-bit scalars, so `F64`/`I64` are a hard error at placement time + /// rather than a silent conversion. + pub const fn gpu_ok(self) -> bool { + matches!(self, DType::F32 | DType::I32) + } + + /// Short name used in traces, diagnostics, and the record/replay tape. + pub const fn name(self) -> &'static str { + match self { + DType::F32 => "f32", + DType::F64 => "f64", + DType::I32 => "i32", + DType::I64 => "i64", + } + } + + /// Parse the name written by [`DType::name`]. Used when reading a tape. + pub fn from_name(s: &str) -> Option { + match s { + "f32" => Some(DType::F32), + "f64" => Some(DType::F64), + "i32" => Some(DType::I32), + "i64" => Some(DType::I64), + _ => None, + } + } + + /// Reconstruct from the `#[repr(u32)]` discriminant stored in a header. + pub fn from_u32(n: u32) -> Option { + match n { + 0 => Some(DType::F32), + 1 => Some(DType::F64), + 2 => Some(DType::I32), + 3 => Some(DType::I64), + _ => None, + } + } +} + +/// How a dense buffer is described to any backend that is not the VM itself. +/// +/// `data` is deliberately a `u64` rather than a pointer: on the native backend +/// it is a host address, in WASM it is a linear-memory offset, and on a device +/// it is an allocation handle. The *header* is identical in all three, which is +/// what lets one conformance test cover them all. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(C)] +pub struct BufferHeader { + /// [`DType`] discriminant. + pub dtype: u32, + /// Padding so `len` is 8-byte aligned in every ABI. + pub _pad: u32, + /// Element count (not bytes). + pub len: u64, + /// Pointer, offset, or device handle — target-dependent. + pub data: u64, +} + +/// Size of [`BufferHeader`] in bytes. Asserted against `size_of` in tests. +pub const BUF_HDR_SIZE: usize = 24; +/// Byte offset of `dtype` within [`BufferHeader`]. +pub const BUF_OFF_DTYPE: usize = 0; +/// Byte offset of `len`. +pub const BUF_OFF_LEN: usize = 8; +/// Byte offset of `data`. +pub const BUF_OFF_DATA: usize = 16; + +impl BufferHeader { + pub fn new(dtype: DType, len: u64, data: u64) -> Self { + BufferHeader { + dtype: dtype as u32, + _pad: 0, + len, + data, + } + } + + pub fn dtype(&self) -> Option { + DType::from_u32(self.dtype) + } + + /// Total size of the element payload in bytes. + pub fn byte_len(&self) -> Option { + Some(self.len as usize * self.dtype()?.size()) + } + + /// Serialize to the canonical 24-byte little-endian form. Backends that + /// cannot share Rust structs (WASM linear memory, a device staging + /// buffer) use this. + pub fn to_bytes(&self) -> [u8; BUF_HDR_SIZE] { + let mut out = [0u8; BUF_HDR_SIZE]; + out[BUF_OFF_DTYPE..BUF_OFF_DTYPE + 4].copy_from_slice(&self.dtype.to_le_bytes()); + out[BUF_OFF_LEN..BUF_OFF_LEN + 8].copy_from_slice(&self.len.to_le_bytes()); + out[BUF_OFF_DATA..BUF_OFF_DATA + 8].copy_from_slice(&self.data.to_le_bytes()); + out + } + + /// Inverse of [`BufferHeader::to_bytes`]. + pub fn from_bytes(b: &[u8]) -> Option { + if b.len() < BUF_HDR_SIZE { + return None; + } + let dtype = u32::from_le_bytes(b[BUF_OFF_DTYPE..BUF_OFF_DTYPE + 4].try_into().ok()?); + let len = u64::from_le_bytes(b[BUF_OFF_LEN..BUF_OFF_LEN + 8].try_into().ok()?); + let data = u64::from_le_bytes(b[BUF_OFF_DATA..BUF_OFF_DATA + 8].try_into().ok()?); + DType::from_u32(dtype)?; + Some(BufferHeader { + dtype, + _pad: 0, + len, + data, + }) + } +} + +// ─── Conformance fixtures ────────────────────────────────────────────────── + +/// Values whose encoding every backend must agree on, bit for bit. +/// +/// `tests/abi_conformance.rs` compiles each `source` on every available +/// backend and asserts the resulting 64-bit word equals `bits`. Add a case +/// here whenever a new immediate encoding appears. +pub const GOLDEN_IMMEDIATES: &[(&str, &str, u64)] = &[ + ("unit", "[]", nanbox::VAL_UNIT), + ("true", "true", nanbox::VAL_TRUE), + ("false", "false", nanbox::VAL_FALSE), + ("none", "None", nanbox::VAL_NONE), + ("int-zero", "0", nanbox::encode_int(0)), + ("int-42", "42", nanbox::encode_int(42)), + ("int-neg", "-7", nanbox::encode_int(-7)), + ( + "int-max48", + "140737488355327", + nanbox::encode_int((1i64 << 47) - 1), + ), + ( + "int-min48", + "-140737488355328", + nanbox::encode_int(-(1i64 << 47)), + ), +]; + +/// A byte pattern chosen to break sloppy marshalling: negative zero, a NaN +/// with a payload, a denormal, and the integer extremes. Any backend that +/// round-trips this unchanged is not quietly normalizing floats or truncating +/// through a narrower type. +pub fn golden_buffer_f32() -> (DType, Vec) { + ( + DType::F32, + vec![ + 0.0, + -0.0, + 1.0, + -1.5, + f32::MIN_POSITIVE / 2.0, // denormal + f32::MAX, + f32::MIN, + f32::INFINITY, + f32::NEG_INFINITY, + f32::from_bits(0x7FC0_0001), // NaN with payload + // Two ordinary values with a full mantissa, to catch a backend + // that round-trips the special cases and mangles the mundane ones. + // Deliberately not near any named constant: these are arbitrary + // bit patterns, and writing them as a truncated pi would suggest + // they meant something. + 1.234_567_9, + -98_765.43, + ], + ) +} + +/// Integer counterpart to [`golden_buffer_f32`]. +pub fn golden_buffer_i32() -> (DType, Vec) { + (DType::I32, vec![0, -1, 1, i32::MIN, i32::MAX, 42, -42]) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn header_size_matches_the_published_constant() { + assert_eq!(std::mem::size_of::(), BUF_HDR_SIZE); + assert_eq!(std::mem::align_of::(), 8); + } + + #[test] + fn header_round_trips_through_bytes() { + for dtype in [DType::F32, DType::F64, DType::I32, DType::I64] { + let h = BufferHeader::new(dtype, 1024, 0xDEAD_BEEF); + let back = BufferHeader::from_bytes(&h.to_bytes()).expect("valid header"); + assert_eq!(h, back, "{} header round-trip", dtype.name()); + assert_eq!(back.dtype(), Some(dtype)); + assert_eq!(back.byte_len(), Some(1024 * dtype.size())); + } + } + + #[test] + fn header_field_offsets_are_stable() { + // A reordering of the struct would change these, which is precisely + // the cross-target divergence this module exists to prevent. + let h = BufferHeader::new(DType::I32, 7, 9); + let b = h.to_bytes(); + assert_eq!(u32::from_le_bytes(b[0..4].try_into().unwrap()), 2); + assert_eq!(u64::from_le_bytes(b[8..16].try_into().unwrap()), 7); + assert_eq!(u64::from_le_bytes(b[16..24].try_into().unwrap()), 9); + } + + #[test] + fn rejects_an_unknown_dtype_discriminant() { + let mut b = BufferHeader::new(DType::F32, 1, 0).to_bytes(); + b[0] = 99; + assert!(BufferHeader::from_bytes(&b).is_none()); + } + + #[test] + fn dtype_names_round_trip() { + for dtype in [DType::F32, DType::F64, DType::I32, DType::I64] { + assert_eq!(DType::from_name(dtype.name()), Some(dtype)); + assert_eq!(DType::from_u32(dtype as u32), Some(dtype)); + } + } + + #[test] + fn only_32_bit_types_are_gpu_placeable() { + assert!(DType::F32.gpu_ok()); + assert!(DType::I32.gpu_ok()); + assert!(!DType::F64.gpu_ok()); + assert!(!DType::I64.gpu_ok()); + } + + #[test] + fn immediates_are_distinct_and_tagged() { + use nanbox::*; + let all = [VAL_UNIT, VAL_TRUE, VAL_FALSE, VAL_NONE]; + for (i, a) in all.iter().enumerate() { + assert_eq!(a & BASE, BASE, "immediate {i} carries the tagged prefix"); + assert_eq!(a & TAG_MASK, TAG_IMM, "immediate {i} carries TAG_IMM"); + for b in all.iter().skip(i + 1) { + assert_ne!(a, b, "immediates must be distinct"); + } + } + } + + #[test] + fn int_encoding_covers_the_48_bit_range() { + use nanbox::*; + for n in [0i64, 1, -1, 42, -42, (1i64 << 47) - 1, -(1i64 << 47)] { + let bits = encode_int(n); + assert_eq!(bits & TAG_MASK, TAG_INT, "{n} is tagged as an int"); + // Sign-extend the 48-bit payload back to i64. + let payload = bits & PAYLOAD; + let back = ((payload << 16) as i64) >> 16; + assert_eq!(back, n, "{n} round-trips through the payload"); + } + } +} diff --git a/crates/loon-lang/src/eir/lower.rs b/crates/loon-lang/src/eir/lower.rs index 19f8ba4..1f6f345 100644 --- a/crates/loon-lang/src/eir/lower.rs +++ b/crates/loon-lang/src/eir/lower.rs @@ -11,6 +11,7 @@ //! - Tail calls → End::Tail / End::Recur use crate::ast::{Expr, ExprKind}; +use crate::check::ownership::ParamMode; use crate::check::Checker; use crate::eir::*; use crate::syntax::Span; @@ -486,12 +487,33 @@ impl<'a> Lower<'a> { // ── Function creation ────────────────────────────────────────────── + /// Ownership modes for a named function's parameters, as inferred by the + /// ownership pass and carried on the checker. + /// + /// Falls back to [`Mode::Owned`] when the analysis has nothing to say: the + /// conservative answer is "assume the callee consumes it", which is always + /// sound and merely gives up an optimization. + fn param_modes_for(&self, name: &str, arity: usize) -> Vec { + match self.checker.fn_param_modes.get(name) { + Some(modes) if modes.len() == arity => modes + .iter() + .map(|m| match m { + ParamMode::Borrow => Mode::In, + ParamMode::MutBorrow => Mode::InOut, + ParamMode::Move => Mode::Owned, + }) + .collect(), + _ => vec![Mode::Owned; arity], + } + } + fn begin_func(&mut self, name: Option<&str>, span: Span) -> FuncId { let id = FuncId(self.module.funcs.len() as u32); self.module.funcs.push(Func { id, name: name.map(|s| s.to_string()), params: Vec::new(), + param_modes: Vec::new(), ret: Ty::Any, evidence: Vec::new(), captures: Vec::new(), @@ -904,6 +926,8 @@ impl<'a> Lower<'a> { self.bind(pname, r); } self.module.funcs[func_id.0 as usize].params = vec![Ty::Any; param_names.len()]; + let modes = self.param_modes_for(&name, param_names.len()); + self.module.funcs[func_id.0 as usize].param_modes = modes; // Set block0 params for fn/recur support let param_regs: Vec = (0..param_names.len()).map(|i| Reg(i as u32)).collect(); @@ -2466,6 +2490,16 @@ pub fn resolve_builtin_name(name: &str) -> Option { "name" => Some(Built::Name), "type-of" => Some(Built::TypeOf), "remove" => Some(Built::Remove), + "buf" => Some(Built::BufNew), + "buf-i32" => Some(Built::BufNewI32), + "buf-f64" => Some(Built::BufNewF64), + "buf-zeros" => Some(Built::BufZeros), + "buf-zeros-i32" => Some(Built::BufZerosI32), + "buf-len" => Some(Built::BufLen), + "buf->vec" => Some(Built::BufToVec), + "buf-dtype" => Some(Built::BufDtype), + "at" => Some(Built::BufAt), + "put" => Some(Built::BufPut), "sqrt" => Some(Built::Sqrt), "pow" => Some(Built::Pow), "floor" => Some(Built::Floor), diff --git a/crates/loon-lang/src/eir/mod.rs b/crates/loon-lang/src/eir/mod.rs index ebfbfaa..389babf 100644 --- a/crates/loon-lang/src/eir/mod.rs +++ b/crates/loon-lang/src/eir/mod.rs @@ -4,17 +4,24 @@ //! Every backend (Register VM, WASM, Cranelift) lowers from this IR. pub mod backend; +pub mod device; +#[cfg(feature = "gpu")] +pub mod gpu; pub mod image; +pub mod kernel_exec; +pub mod layout; pub mod lower; #[cfg(feature = "native")] pub mod native; pub mod net; +pub mod place; pub mod replay; pub mod tailcall; pub mod trace; pub mod value64; pub mod vm; pub mod wasm; +pub mod wgsl; use crate::syntax::Span; @@ -62,6 +69,14 @@ pub struct Func { pub id: FuncId, pub name: Option, pub params: Vec, + /// How each parameter is used: read, written through, or consumed. + /// + /// Inferred by the ownership pass rather than written by the programmer. + /// Empty for functions the analysis did not name (closures, handler + /// clauses); otherwise one entry per parameter. Backends use it the way a + /// C API would use `const`: an [`Mode::In`] argument never has to be + /// copied back from wherever it was sent. + pub param_modes: Vec, pub ret: Ty, /// Implicit handler function-pointer parameters (evidence-passing). pub evidence: Vec, @@ -73,6 +88,31 @@ pub struct Func { pub is_closure: bool, } +/// Direction of a parameter across a call — or across a placement boundary. +/// +/// This is the whole of what `&T` / `&mut T` / `T` tell a Rust offload +/// compiler about which way bytes move, except nobody had to write it down. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Mode { + /// Read only. Travels toward the callee and never comes back. + In, + /// Written through. Must be visible to the caller afterwards. + InOut, + /// Consumed. The caller may not use it again. + Owned, +} + +impl Mode { + /// The keyword a Loon-level handler sees for this mode. + pub fn keyword(self) -> &'static str { + match self { + Mode::In => "in", + Mode::InOut => "inout", + Mode::Owned => "owned", + } + } +} + /// Evidence parameter — a handler function pointer threaded through calls. #[derive(Debug, Clone)] pub struct Evidence { @@ -420,6 +460,27 @@ pub enum Built { /// for any other value. Emitted by pattern compilation so `#[a b]` can /// test "is a sequence of length 2" in one comparison. SeqLen, + // ── Dense buffers ── + /// `buf` — an f32 buffer from a vector of numbers. + BufNew, + /// `buf-i32` — an i32 buffer from a vector of numbers. + BufNewI32, + /// `buf-f64` — an f64 buffer from a vector of numbers. + BufNewF64, + /// `buf-zeros` — an f32 buffer of n zeros. + BufZeros, + /// `buf-zeros-i32` — an i32 buffer of n zeros. + BufZerosI32, + /// `buf-len` — element count. + BufLen, + /// `buf->vec` — copy out to an ordinary vector. + BufToVec, + /// `buf-dtype` — element type name, as a string. + BufDtype, + /// `at` — read element i. Out of range is an error, never a made-up zero. + BufAt, + /// `put` — write element i in place and return the buffer. + BufPut, /// Internal (not name-resolvable): destructuring guard. Args are /// (value, expected-binder-count). Errors loudly unless the value is a /// vector/tuple with at least that many elements — a silent `()` bind diff --git a/crates/loon-lang/src/eir/native.rs b/crates/loon-lang/src/eir/native.rs index 263a73b..d12ba7f 100644 --- a/crates/loon-lang/src/eir/native.rs +++ b/crates/loon-lang/src/eir/native.rs @@ -13,13 +13,19 @@ //! - Function calls (direct), including tail calls (`return_call`) //! - Builtin println (via extern) //! -//! Not yet implemented (fall back to VM): +//! Not yet implemented (fall back to VM). Using one of these is a compile +//! error, never a silently wrong value — see `compile_op`: //! - Closures / upvalues / indirect calls //! - Collection construction (Vec, Map, Set, Tuple, ADT) //! - Field access, tag extraction //! - Effect operations (perform, push/pop handler) //! - String operations //! +//! Because a program with a `main` function reaches it through a closure, that +//! whole shape is currently rejected here rather than executed. `loon run +//! --native` therefore serves bare top-level expressions today and says so +//! plainly otherwise. +//! //! Loon functions are compiled with Cranelift's `tail` calling convention so //! that `End::Tail` can lower to a real `return_call` (constant stack for //! mutual tail recursion). That convention is not the platform C ABI, so the @@ -42,19 +48,14 @@ use super::{BinOp, End, Lit, Op, Reg, UnOp}; use std::collections::HashMap; -// ─── NaN-boxing constants (must match value64.rs) ─────────────────────────── - -const QNAN: u64 = 0x7FF8_0000_0000_0000; -const SIGN: u64 = 0x8000_0000_0000_0000; -const BASE: u64 = SIGN | QNAN; -const TAG_INT: u64 = 0x0001_0000_0000_0000; -const TAG_IMM: u64 = 0x0007_0000_0000_0000; -const PAYLOAD: u64 = 0x0000_FFFF_FFFF_FFFF; +// ─── NaN-boxing constants ─────────────────────────────────────────────────── +// +// Imported from `eir::layout`, the one place the encoding is defined. -const VAL_UNIT: u64 = BASE | TAG_IMM; -const VAL_TRUE: u64 = BASE | TAG_IMM | 1; -const VAL_FALSE: u64 = BASE | TAG_IMM | 2; -const VAL_NONE: u64 = BASE | TAG_IMM | 3; +#[allow(unused_imports)] +use super::layout::nanbox::{ + BASE, PAYLOAD, QNAN, SIGN, TAG_IMM, TAG_INT, VAL_FALSE, VAL_NONE, VAL_TRUE, VAL_UNIT, +}; /// Symbol name of the C-ABI shim that calls the module's entry function. const ENTRY_TRAMPOLINE: &str = "loon_entry_trampoline"; @@ -526,6 +527,26 @@ fn compile_function( // ─── Op compilation ───────────────────────────────────────────────────────── +/// A human-readable name for an operation, used in "not supported yet" +/// diagnostics so the message names the feature rather than an opcode. +fn op_description(op: &Op) -> &'static str { + match op { + Op::Upval(..) => "reading a closure upvalue", + Op::Invoke(..) => "calling a closure", + Op::Close(..) => "creating a closure", + Op::Vec(..) => "building a vector", + Op::Map(..) => "building a map", + Op::Set(..) => "building a set", + Op::Tup(..) => "building a tuple", + Op::Adt(..) => "constructing an ADT value", + Op::Field(..) => "field access", + Op::Tag(..) => "reading a constructor tag", + Op::Perform(..) => "performing an effect", + Op::PushHandler(..) => "installing an effect handler", + _ => "this operation", + } +} + #[allow(clippy::too_many_arguments)] fn compile_op( builder: &mut FunctionBuilder, @@ -614,24 +635,66 @@ fn compile_op( builder.def_var(vars[dst.0 as usize], result); } - // Operations that need heap/runtime support — emit unit placeholder. - Op::Upval(dst, _, _) - | Op::Invoke(dst, _, _, _) - | Op::Close(dst, _, _, _) + // Operations that *use* a value this backend cannot represent. They + // fail loudly, exactly as `End::TailInvoke` below does and for the + // same reason. + // + // These used to emit a unit placeholder, which made every program with + // a `main` function silently evaluate to `()`: the synthetic entry + // point reaches `main` through `Close` + `Invoke`, so `loon run + // --native` printed nothing and exited 0 while the VM ran the program + // correctly. A missing feature that announces itself is a far smaller + // problem than a wrong answer that does not. + Op::Upval(_, _, _) + | Op::Invoke(_, _, _, _) + | Op::Field(_, _, _, _) + | Op::Tag(_, _, _) + | Op::Perform(_, _, _, _, _, _) + | Op::PushHandler(_, _, _, _) => { + return Err(Error { + message: format!( + "{} is not supported by the native backend yet", + op_description(op) + ), + phase: "native:compile", + }); + } + + // Operations that merely *construct* a value the backend cannot + // represent. Lowering emits these freely — every named function gets a + // `Close` for its value form, whether or not anything calls it + // indirectly — so rejecting them outright would refuse programs that + // compile and run correctly today. The placeholder is inert: the only + // ways to observe one are the operations above, which do fail. + Op::Close(dst, _, _, _) | Op::Vec(dst, _, _) | Op::Map(dst, _, _) | Op::Set(dst, _, _) - | Op::Tup(dst, _, _) - | Op::Adt(dst, _, _, _) - | Op::Field(dst, _, _, _) - | Op::Tag(dst, _, _) - | Op::Perform(dst, _, _, _, _, _) - | Op::PushHandler(dst, _, _, _) => { - // TODO: implement via runtime helper calls. + | Op::Tup(dst, _, _) => { let unit = builder.ins().iconst(I64, VAL_UNIT as i64); builder.def_var(vars[dst.0 as usize], unit); } + Op::Adt(dst, tag, fields, _) => { + // The nullary `None` is an immediate singleton on every backend, + // never a heap value: bit equality is `None` equality and the + // falsy test is a bit test. Emitting the placeholder here instead + // would make `None` indistinguishable from `()`. + let none_tag = _eir_module + .ctors + .iter() + .rev() + .find(|c| c.name == "None") + .map(|c| c.tag); + let val = if fields.is_empty() && none_tag == Some(*tag) { + VAL_NONE + } else { + VAL_UNIT + }; + let v = builder.ins().iconst(I64, val as i64); + builder.def_var(vars[dst.0 as usize], v); + } + Op::PopHandler(_) => { // No-op in native backend (no destination register). } @@ -1181,6 +1244,7 @@ mod tests { id: FuncId(id), name: Some(format!("parity{id}")), params: vec![Ty::Int], + param_modes: Vec::new(), ret: Ty::Bool, evidence: vec![], captures: vec![], @@ -1222,6 +1286,7 @@ mod tests { id: FuncId(2), name: Some("__main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Bool, evidence: vec![], captures: vec![], @@ -1260,6 +1325,7 @@ mod tests { id: FuncId(0), name: Some("__main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], @@ -1297,6 +1363,62 @@ mod tests { assert_eq!(result.as_int(), 1); } + /// A program with a `main` function must never appear to succeed while + /// doing nothing. + /// + /// The synthetic entry point reaches `main` through `Close` + `Invoke`, + /// and `Invoke` is not implemented here. That used to yield Unit, so + /// `loon run --native` on any real program printed nothing and exited 0 + /// while the VM ran it correctly — a silent wrong answer produced by a + /// shipped flag. It must be an error until closures are implemented, and + /// a correct result afterwards; what it must never be again is Unit. + #[test] + fn a_main_function_is_never_silently_unit() { + let src = "[fn main [] 42]"; + match eval_native(src) { + Err(msg) => assert!( + msg.contains("closure"), + "the error should name the missing feature, got: {msg}" + ), + Ok(v) => assert_eq!( + v.as_int(), + 42, + "if closures are implemented, `main` must actually run" + ), + } + } + + /// The VM and the native backend must never disagree about a value they + /// both claim to have computed. Either the native backend produces the + /// same bits, or it refuses the program. + #[test] + fn native_never_disagrees_with_the_vm() { + for src in [ + "42", + "[+ 1 2]", + "[* [+ 2 3] [- 10 4]]", + "true", + "false", + "[if true 1 2]", + "[do [let x 10] [+ x 5]]", + "[fn double [x] [* x 2]] [double 21]", + "[fn main [] 42]", + "[fn main [] [+ 1 2]]", + ] { + let Ok(native) = eval_native(src) else { + continue; // refused outright — the acceptable answer + }; + let vm = crate::eir::vm::eval_eir(src).expect("VM runs it"); + assert_eq!( + native.bits(), + vm.value.bits(), + "`{src}`: native {:#018x} vs VM {:#018x}", + native.bits(), + vm.value.bits() + ); + } + } + #[test] fn compile_unit() { let result = eval_native("()").unwrap(); diff --git a/crates/loon-lang/src/eir/place.rs b/crates/loon-lang/src/eir/place.rs new file mode 100644 index 0000000..53fb8c7 --- /dev/null +++ b/crates/loon-lang/src/eir/place.rs @@ -0,0 +1,378 @@ +//! Placement: where a kernel runs, and what it cost to put it there. +//! +//! Loon expresses placement as an effect. A program performs `Place.run` and a +//! handler decides what that means — run it here, run it across threads, ship +//! it to a device, record it, or pretend. Nothing in the program changes when +//! the answer changes. +//! +//! This module holds the parts that are the same whoever answers: the record +//! of what happened ([`PlaceEvent`]) and the running totals ([`PlaceStats`]). +//! Transfer accounting is not an afterthought here. The interesting question +//! about an offloaded program is almost never "was the kernel fast" — it is +//! "how many times did these bytes cross the boundary, and did they need to". +//! Making that observable is what lets a residency policy be written as a +//! handler and then *checked*, rather than assumed. + +use super::layout::DType; + +/// What a placement backend did. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EventKind { + /// Bytes moved toward the device. + Upload, + /// Bytes moved back to the host. + Download, + /// A kernel ran. + Launch, + /// An argument was already resident, so no transfer happened. Counting + /// the transfers that *didn't* occur is what makes a residency policy + /// legible: the win shows up as a number, not as a wall-clock guess. + ResidentHit, + /// Device memory was reserved. + Alloc, + /// Device memory was released. + Free, +} + +impl EventKind { + pub fn name(self) -> &'static str { + match self { + EventKind::Upload => "upload", + EventKind::Download => "download", + EventKind::Launch => "launch", + EventKind::ResidentHit => "resident-hit", + EventKind::Alloc => "alloc", + EventKind::Free => "free", + } + } +} + +/// One thing that happened during placement. +#[derive(Debug, Clone)] +pub struct PlaceEvent { + pub kind: EventKind, + /// Kernel name, for a launch. + pub kernel: Option, + /// Which argument this concerns, for a transfer. + pub arg: Option, + /// Element type of the buffer involved. + pub dtype: Option, + /// Bytes moved. + pub bytes: u64, + /// Work items, for a launch. + pub items: u64, + /// Which backend answered. + pub device: &'static str, +} + +/// Running totals across a program. +#[derive(Debug, Clone, Default)] +pub struct PlaceStats { + pub launches: u64, + pub work_items: u64, + pub uploads: u64, + pub downloads: u64, + pub bytes_in: u64, + pub bytes_out: u64, + pub resident_hits: u64, + pub events: Vec, +} + +impl PlaceStats { + pub fn record(&mut self, event: PlaceEvent) { + match event.kind { + EventKind::Upload => { + self.uploads += 1; + self.bytes_in += event.bytes; + } + EventKind::Download => { + self.downloads += 1; + self.bytes_out += event.bytes; + } + EventKind::Launch => { + self.launches += 1; + self.work_items += event.items; + } + EventKind::ResidentHit => self.resident_hits += 1, + EventKind::Alloc | EventKind::Free => {} + } + self.events.push(event); + } + + /// A one-line summary in the shape offload papers report: how many + /// transfers, how many bytes, how many launches. + pub fn summary(&self) -> String { + format!( + "{} launches over {} work items; {} uploads ({}), {} downloads ({}), {} resident hits", + self.launches, + self.work_items, + self.uploads, + human_bytes(self.bytes_in), + self.downloads, + human_bytes(self.bytes_out), + self.resident_hits, + ) + } + + /// A multi-line table for `--place-stats`. + pub fn table(&self) -> String { + let mut out = String::new(); + out.push_str("placement\n"); + out.push_str(&format!(" launches {}\n", self.launches)); + out.push_str(&format!(" work items {}\n", self.work_items)); + out.push_str(&format!( + " uploads {} ({})\n", + self.uploads, + human_bytes(self.bytes_in) + )); + out.push_str(&format!( + " downloads {} ({})\n", + self.downloads, + human_bytes(self.bytes_out) + )); + out.push_str(&format!(" resident hits {}\n", self.resident_hits)); + out + } +} + +// ─── A device with its own memory ────────────────────────────────────────── + +/// Where kernels run. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum Mode { + /// Right here. There is one memory, so nothing is ever transferred. + #[default] + Cpu, + /// Every core on this machine, over disjoint slices of the index space. + /// + /// Still one memory, so nothing is transferred — the only thing that + /// changes is how many work items run at once. + Par, + /// A real GPU, through wgpu. + /// + /// Requires the `gpu` feature; without it, asking for this mode is an + /// error that says so rather than a silent fall back to the CPU. Being + /// told your program did not run where you asked is worth more than a + /// result that arrived by a route you did not choose. + Gpu, + /// A discrete device with separate memory. + /// + /// The arithmetic still happens on the host — this is not a GPU — but the + /// *bookkeeping* is real: a buffer must be uploaded before a kernel can + /// use it, and results must be downloaded before the host can read them. + /// That makes the cost of a placement policy measurable, and measurable is + /// the whole argument. A policy whose benefit you cannot count is a story. + Device, +} + +impl Mode { + pub fn parse(s: &str) -> Option { + match s { + "cpu" | "serial" | "here" => Some(Mode::Cpu), + "device" | "sim" => Some(Mode::Device), + "par" | "parallel" | "threads" => Some(Mode::Par), + "gpu" | "metal" | "wgpu" => Some(Mode::Gpu), + _ => None, + } + } + + /// Whether this mode has a memory separate from the host's, so that + /// transfers are a real cost worth accounting for. + pub fn has_device_memory(self) -> bool { + matches!(self, Mode::Device | Mode::Gpu) + } + + pub fn name(self) -> &'static str { + match self { + Mode::Cpu => "cpu", + Mode::Par => "par", + Mode::Device => "device", + Mode::Gpu => "gpu", + } + } +} + +/// Device-side memory: which buffers are resident, and which hold results the +/// host has not seen yet. +/// +/// Buffers are identified by their heap slot. Nothing is copied — the point is +/// to account for the copies a real device would need, not to simulate its +/// arithmetic. +#[derive(Debug, Clone, Default)] +pub struct Device { + resident: std::collections::HashSet, + dirty: std::collections::HashSet, + /// Buffers a handler asked to keep resident across launches. + pinned: std::collections::HashSet, +} + +impl Device { + pub fn is_resident(&self, id: usize) -> bool { + self.resident.contains(&id) + } + + pub fn is_dirty(&self, id: usize) -> bool { + self.dirty.contains(&id) + } + + pub fn is_pinned(&self, id: usize) -> bool { + self.pinned.contains(&id) + } + + pub fn mark_resident(&mut self, id: usize) { + self.resident.insert(id); + } + + pub fn mark_dirty(&mut self, id: usize) { + self.dirty.insert(id); + } + + pub fn clear_dirty(&mut self, id: usize) { + self.dirty.remove(&id); + } + + /// Ask that a buffer survive eviction. + /// + /// Deliberately does *not* make it resident: pinning says "keep this once + /// it is here", not "it is here already". The first launch that uses it + /// still pays for the upload. Otherwise a residency policy would look free + /// in the accounting by declaring itself so, which is exactly the kind of + /// measurement that flatters a design instead of testing it. + pub fn pin(&mut self, id: usize) { + self.pinned.insert(id); + } + + pub fn unpin(&mut self, id: usize) { + self.pinned.remove(&id); + } + + /// Drop every unpinned buffer. + /// + /// This is what makes a transfer-per-launch policy the *default* rather + /// than a strawman: without someone deciding otherwise, the device does + /// not assume a buffer will be wanted again. Deciding otherwise is exactly + /// what a residency handler does, and pinning is how it says so. + pub fn evict_unpinned(&mut self) -> Vec { + let evicted: Vec = self + .resident + .iter() + .copied() + .filter(|id| !self.pinned.contains(id)) + .collect(); + for id in &evicted { + self.resident.remove(id); + } + evicted + } + + pub fn resident_count(&self) -> usize { + self.resident.len() + } +} + +/// Bytes in a unit a person can read at a glance. +pub fn human_bytes(n: u64) -> String { + const KB: u64 = 1024; + const MB: u64 = KB * 1024; + const GB: u64 = MB * 1024; + if n >= GB { + format!("{:.1} GB", n as f64 / GB as f64) + } else if n >= MB { + format!("{:.1} MB", n as f64 / MB as f64) + } else if n >= KB { + format!("{:.1} KB", n as f64 / KB as f64) + } else { + format!("{n} B") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pinning_does_not_make_a_buffer_resident_for_free() { + // The first use still has to upload it; pinning only says it should + // stay afterwards. + let mut d = Device::default(); + d.pin(3); + assert!(!d.is_resident(3), "pinning is not a transfer"); + } + + #[test] + fn an_unpinned_buffer_does_not_survive_a_launch() { + let mut d = Device::default(); + d.mark_resident(1); + d.mark_resident(2); + d.pin(2); + let evicted = d.evict_unpinned(); + assert_eq!(evicted, vec![1]); + assert!(!d.is_resident(1), "unpinned buffers are dropped"); + assert!(d.is_resident(2), "pinned buffers stay"); + } + + #[test] + fn unpinning_lets_a_buffer_be_evicted_again() { + let mut d = Device::default(); + d.mark_resident(7); + d.pin(7); + assert!(d.evict_unpinned().is_empty()); + d.unpin(7); + assert_eq!(d.evict_unpinned(), vec![7]); + } + + #[test] + fn mode_names_round_trip() { + for m in [Mode::Cpu, Mode::Par, Mode::Device, Mode::Gpu] { + assert_eq!(Mode::parse(m.name()), Some(m)); + } + assert_eq!(Mode::parse("sim"), Some(Mode::Device)); + assert_eq!(Mode::parse("nonsense"), None); + } + + fn ev(kind: EventKind, bytes: u64, items: u64) -> PlaceEvent { + PlaceEvent { + kind, + kernel: None, + arg: None, + dtype: None, + bytes, + items, + device: "test", + } + } + + #[test] + fn totals_follow_the_events() { + let mut s = PlaceStats::default(); + s.record(ev(EventKind::Upload, 100, 0)); + s.record(ev(EventKind::Upload, 40, 0)); + s.record(ev(EventKind::Launch, 0, 1024)); + s.record(ev(EventKind::Download, 60, 0)); + s.record(ev(EventKind::ResidentHit, 0, 0)); + + assert_eq!(s.uploads, 2); + assert_eq!(s.bytes_in, 140); + assert_eq!(s.downloads, 1); + assert_eq!(s.bytes_out, 60); + assert_eq!(s.launches, 1); + assert_eq!(s.work_items, 1024); + assert_eq!(s.resident_hits, 1); + assert_eq!(s.events.len(), 5); + } + + #[test] + fn byte_counts_read_like_prose() { + assert_eq!(human_bytes(512), "512 B"); + assert_eq!(human_bytes(2048), "2.0 KB"); + assert_eq!(human_bytes(3 * 1024 * 1024), "3.0 MB"); + assert_eq!(human_bytes(2 * 1024 * 1024 * 1024), "2.0 GB"); + } + + #[test] + fn a_quiet_run_reports_zeroes_rather_than_nothing() { + let s = PlaceStats::default(); + assert!(s.summary().contains("0 launches")); + assert!(s.table().contains("resident hits 0")); + } +} diff --git a/crates/loon-lang/src/eir/replay.rs b/crates/loon-lang/src/eir/replay.rs index ffd6366..d78c605 100644 --- a/crates/loon-lang/src/eir/replay.rs +++ b/crates/loon-lang/src/eir/replay.rs @@ -182,6 +182,17 @@ impl TraceEntry { pub fn is_recorded_op(effect: &str, op: &str) -> bool { match effect { "Net" | "Env" | "Process" | "Rand" => true, + // Placement is recorded so a program that ran on a GPU can be replayed + // on a machine that has none: a launch replays as the unit it returned + // and a read replays as the values it produced, so the recording is a + // faithful account of what the program observed without any hardware + // being involved a second time. + // + // `Place.stats` is excluded. It reports on the run currently + // happening, and a replayed run really did move no bytes — feeding + // back the original transfer counts would be a recording that lies + // about the execution it is part of. + "Place" => op != "stats", "IO" => !matches!(op, "parse-json" | "to-json" | "blake3"), _ => false, } diff --git a/crates/loon-lang/src/eir/tailcall.rs b/crates/loon-lang/src/eir/tailcall.rs index ecbf198..3856ac1 100644 --- a/crates/loon-lang/src/eir/tailcall.rs +++ b/crates/loon-lang/src/eir/tailcall.rs @@ -243,6 +243,7 @@ mod tests { id: FuncId(0), name: Some("f".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], diff --git a/crates/loon-lang/src/eir/value64.rs b/crates/loon-lang/src/eir/value64.rs index 79308dd..1f908ce 100644 --- a/crates/loon-lang/src/eir/value64.rs +++ b/crates/loon-lang/src/eir/value64.rs @@ -25,24 +25,12 @@ use std::fmt; #[repr(transparent)] pub struct Val(u64); -// Bit layout constants -const QNAN: u64 = 0x7FF8_0000_0000_0000; // quiet NaN base -const SIGN: u64 = 0x8000_0000_0000_0000; // sign bit -const BASE: u64 = SIGN | QNAN; // all tagged values have this prefix -const TAG_MASK: u64 = 0x0007_0000_0000_0000; // bits 48-50 -const PAYLOAD: u64 = 0x0000_FFFF_FFFF_FFFF; // low 48 bits - -// Tag values (shifted into position) -const TAG_PTR: u64 = 0x0000_0000_0000_0000; // heap pointer -const TAG_INT: u64 = 0x0001_0000_0000_0000; // inline int48 -const TAG_SYM: u64 = 0x0006_0000_0000_0000; // interned symbol/keyword -const TAG_IMM: u64 = 0x0007_0000_0000_0000; // immediate - -// Immediate sub-tags (in low bits) -const IMM_UNIT: u64 = 0; -const IMM_TRUE: u64 = 1; -const IMM_FALSE: u64 = 2; -const IMM_NONE: u64 = 3; +// The bit layout itself lives in `layout::nanbox` so that the WASM emitter and +// the Cranelift backend encode values the same way this file decodes them. +use super::layout::nanbox::{ + BASE, IMM_FALSE, IMM_NONE, IMM_TRUE, IMM_UNIT, PAYLOAD, TAG_IMM, TAG_INT, TAG_MASK, TAG_PTR, + TAG_SYM, +}; impl Val { // ── Constants ────────────────────────────────────────────────────── diff --git a/crates/loon-lang/src/eir/vm.rs b/crates/loon-lang/src/eir/vm.rs index f89d416..279d470 100644 --- a/crates/loon-lang/src/eir/vm.rs +++ b/crates/loon-lang/src/eir/vm.rs @@ -5,6 +5,7 @@ //! Tail calls reuse the current frame. No continuation stack — just a //! call stack of `(FuncId, BlockId, ip, registers)`. +use crate::eir::layout::DType; pub use crate::eir::value64::Val; use crate::eir::*; use std::collections::HashMap; @@ -119,6 +120,15 @@ enum Obj { Tuple(Vec), // fixed-size, no persistence needed Adt(u16, Vec), // tag + fields Closure(FuncId, Vec), // func + captured values + /// A dense, fixed-length array of unboxed numbers. + /// + /// Loon's ordinary collections are persistent trees of NaN-boxed words, + /// which is the right shape for a functional language and the wrong shape + /// for anything that has to hand bytes to another processor. A `Buffer` is + /// the other shape: contiguous, untagged, and describable by a + /// `layout::BufferHeader`, so the same bytes can be shipped to a thread, a + /// kernel, or a device queue without a conversion step. + Buffer(Buffer), /// A reified multi-shot delimited continuation captured at a `perform`: the /// frame segment between the perform and its handler's prompt, plus the /// execution point to resume at. `resume_continuation` clones this segment on @@ -152,6 +162,7 @@ impl Obj { Obj::Map(m) => (24 + m.len() * 16) as u64, Obj::Adt(_, fields) => (24 + 2 + fields.len() * 8) as u64, Obj::Closure(_, caps) => (24 + 4 + caps.len() * 8) as u64, + Obj::Buffer(b) => 24 + b.byte_len() as u64, Obj::Continuation { saved, regs, @@ -162,6 +173,166 @@ impl Obj { } } +// ─── Dense buffers ───────────────────────────────────────────────────────── + +/// The elements of a [`Obj::Buffer`], one variant per [`DType`]. +#[derive(Debug, Clone, PartialEq)] +pub enum BufData { + F32(Vec), + F64(Vec), + I32(Vec), + I64(Vec), +} + +/// A dense numeric array. +/// +/// Values are read and written by index through `at` and `put`. Unlike Loon's +/// persistent collections a buffer is mutated in place when its owner is +/// unique, which is what makes it cheap enough to be worth sending anywhere. +#[derive(Debug, Clone, PartialEq)] +pub struct Buffer { + pub data: BufData, +} + +impl Buffer { + pub fn zeros(dtype: DType, len: usize) -> Buffer { + Buffer { + data: match dtype { + DType::F32 => BufData::F32(vec![0.0; len]), + DType::F64 => BufData::F64(vec![0.0; len]), + DType::I32 => BufData::I32(vec![0; len]), + DType::I64 => BufData::I64(vec![0; len]), + }, + } + } + + pub fn dtype(&self) -> DType { + match &self.data { + BufData::F32(_) => DType::F32, + BufData::F64(_) => DType::F64, + BufData::I32(_) => DType::I32, + BufData::I64(_) => DType::I64, + } + } + + pub fn len(&self) -> usize { + match &self.data { + BufData::F32(v) => v.len(), + BufData::F64(v) => v.len(), + BufData::I32(v) => v.len(), + BufData::I64(v) => v.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn byte_len(&self) -> usize { + self.len() * self.dtype().size() + } + + /// Element `i` as a Loon value, or `None` if out of range. + /// + /// Out of range is deliberately not silently zero: reading past the end of + /// a buffer is a bug in the program, and a language that answers `0` for it + /// teaches you to trust an answer it made up. + pub fn get(&self, i: usize) -> Option { + match &self.data { + BufData::F32(v) => v.get(i).map(|x| Val::float(*x as f64)), + BufData::F64(v) => v.get(i).map(|x| Val::float(*x)), + BufData::I32(v) => v.get(i).map(|x| Val::int(*x as i64)), + BufData::I64(v) => v.get(i).map(|x| Val::int(*x)), + } + } + + /// Write element `i`, converting to the buffer's element type. Returns + /// false if the index is out of range or the value is not numeric. + pub fn set(&mut self, i: usize, val: Val) -> bool { + let num = if val.is_int() { + val.as_int() as f64 + } else if val.is_float() { + val.as_float() + } else { + return false; + }; + match &mut self.data { + BufData::F32(v) => match v.get_mut(i) { + Some(slot) => *slot = num as f32, + None => return false, + }, + BufData::F64(v) => match v.get_mut(i) { + Some(slot) => *slot = num, + None => return false, + }, + BufData::I32(v) => match v.get_mut(i) { + Some(slot) => *slot = num as i32, + None => return false, + }, + BufData::I64(v) => match v.get_mut(i) { + Some(slot) => *slot = num as i64, + None => return false, + }, + } + true + } + + /// The raw bytes of the elements, little-endian. + pub fn to_bytes(&self) -> Vec { + match &self.data { + BufData::F32(v) => v.iter().flat_map(|x| x.to_le_bytes()).collect(), + BufData::F64(v) => v.iter().flat_map(|x| x.to_le_bytes()).collect(), + BufData::I32(v) => v.iter().flat_map(|x| x.to_le_bytes()).collect(), + BufData::I64(v) => v.iter().flat_map(|x| x.to_le_bytes()).collect(), + } + } + + /// Rebuild from raw little-endian bytes. Returns `None` on a length that + /// is not a whole number of elements. + pub fn from_bytes(dtype: DType, bytes: &[u8]) -> Option { + if bytes.len() % dtype.size() != 0 { + return None; + } + let data = match dtype { + DType::F32 => BufData::F32( + bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes(c.try_into().unwrap())) + .collect(), + ), + DType::F64 => BufData::F64( + bytes + .chunks_exact(8) + .map(|c| f64::from_le_bytes(c.try_into().unwrap())) + .collect(), + ), + DType::I32 => BufData::I32( + bytes + .chunks_exact(4) + .map(|c| i32::from_le_bytes(c.try_into().unwrap())) + .collect(), + ), + DType::I64 => BufData::I64( + bytes + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap())) + .collect(), + ), + }; + Some(Buffer { data }) + } +} + +/// A `ret_reg` meaning "this call's result is returned to the Rust caller, +/// not stored into a register". +/// +/// The VM re-enters itself whenever a builtin calls back into Loon — `map` +/// over a closure, a kernel launched by `Place.run`. Those calls have no +/// destination register in the caller, because the caller is Rust. Naming +/// register 0 instead, as this used to, silently overwrote whatever the +/// caller's first register held. +const RET_DISCARD: u32 = u32::MAX; + // ─── Call frame ──────────────────────────────────────────────────────────── /// Saved call frame — pushed when entering a function, popped on return. @@ -174,6 +345,8 @@ struct Frame { /// Register file for the suspended frame. regs: Vec, /// Register in the caller's frame to write the return value into. + /// Register in *this* frame that receives the callee's return value, or + /// [`RET_DISCARD`] when the value goes back to Rust instead. ret_reg: u32, /// Closure captures for the suspended frame. captures: Vec, @@ -247,6 +420,23 @@ pub struct Vm { ip: usize, /// Output capture (for println). output: Vec, + /// Placement accounting: launches, transfers, and residency hits. + pub place_stats: crate::eir::place::PlaceStats, + /// Where kernels run when nothing handles `Place`. + pub place_mode: crate::eir::place::Mode, + /// Device-side residency, when running in device mode. + place_device: crate::eir::place::Device, + /// A continuation handed out by `Host.park`, with what it is waiting for. + /// + /// Set when a program suspends; the host reads it, does whatever it could + /// not do synchronously, and resumes through `Vm::call_value`. + parked: Option<(Val, Val)>, + /// The device kernels are dispatched to, opened on first use and kept for + /// the rest of the run. + /// + /// Any `Device` will do: wgpu on a desktop, a bridge to JavaScript in a + /// browser. The VM does not know which it has. + device: Option>, /// String constants resolved to heap indices. string_cache: HashMap, /// Interns string *objects* by content, so structurally-equal strings share @@ -351,6 +541,11 @@ impl Vm { resume_closure: Val::UNIT, // set in run() current_span: Span::ZERO, heap_stats: HeapStats::default(), + place_stats: crate::eir::place::PlaceStats::default(), + place_mode: crate::eir::place::Mode::default(), + place_device: crate::eir::place::Device::default(), + parked: None, + device: None, recorder: None, replay: None, runtime_syms: Vec::new(), @@ -455,6 +650,9 @@ impl Vm { id, name: Some("resume".to_string()), params: vec![Ty::Any], + // `resume` hands its argument onward without consuming it — a + // multi-shot continuation may be resumed more than once. + param_modes: vec![Mode::In], ret: Ty::Any, evidence: Vec::new(), captures: Vec::new(), @@ -494,7 +692,7 @@ impl Vm { caps: Vec, ) -> Result { let depth = self.frames.len(); - self.call_func_with_captures(func_id, args, 0, caps)?; + self.call_func_with_captures(func_id, args, RET_DISCARD, caps)?; self.execute(depth + 1) } @@ -523,6 +721,518 @@ impl Vm { val } + /// Run a kernel once per work item, right here. + /// + /// `[Place.run kernel n args...]` calls `kernel(i, args...)` for each `i` + /// in `0..n`. Kernels return nothing useful; they write through their + /// buffer arguments, which is why the ownership pass reporting a + /// written-through parameter as `InOut` is the same fact as "this buffer + /// has to come back". + fn place_run_serial(&mut self, args: &[Val]) -> Result { + let kernel = args.first().copied().unwrap_or(Val::UNIT); + let n_val = args.get(1).copied().unwrap_or(Val::UNIT); + + let (func_id, caps, name) = match self.get_obj(kernel) { + Some(Obj::Closure(f, c)) => { + let name = self + .module + .funcs + .get(f.0 as usize) + .and_then(|fun| fun.name.clone()); + (*f, c.clone(), name) + } + _ => { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "Place.run expects a kernel as its first argument, got {}", + self.val_to_string(kernel) + ))) + .with_span(self.current_span)) + } + }; + + if !n_val.is_int() || n_val.as_int() < 0 { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "Place.run expects a non-negative work count, got {}", + self.val_to_string(n_val) + ))) + .with_span(self.current_span)); + } + let n = n_val.as_int(); + + // Charge the launch before running, so a kernel that fails partway + // still shows up in the accounting. + self.place_stats.record(crate::eir::place::PlaceEvent { + kind: crate::eir::place::EventKind::Launch, + kernel: name, + arg: None, + dtype: None, + bytes: 0, + items: n as u64, + device: self.place_mode.name(), + }); + + // Kernel arguments arrive as one vector rather than spread across the + // operation. That keeps the operation's own arity fixed at three, so a + // handler clause can bind and forward it — `[Place.run k n args]` — + // and, more usefully, can *inspect* the argument list. A residency + // policy is exactly a handler that looks at those arguments and + // decides which of them still need to move. + let args_val = args.get(2).copied().unwrap_or(Val::UNIT); + let rest: Vec = match self.get_obj(args_val) { + Some(Obj::Vec(v)) => v.iter().copied().collect(), + Some(Obj::Tuple(v)) => v.clone(), + _ if args.len() <= 2 => Vec::new(), + _ => { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "Place.run expects a vector of kernel arguments, got {}", + self.val_to_string(args_val) + ))) + .with_span(self.current_span)) + } + }; + // In device mode, every buffer argument has to be on the device before + // the kernel can touch it. An argument that is already resident costs + // nothing — and counting those hits is what turns a residency policy + // from a claim into a measurement. + if self.place_mode == crate::eir::place::Mode::Device { + for (i, arg) in rest.iter().enumerate() { + let Some((id, bytes, dtype)) = self.buffer_info(*arg) else { + continue; + }; + if self.place_device.is_resident(id) { + self.place_stats.record(crate::eir::place::PlaceEvent { + kind: crate::eir::place::EventKind::ResidentHit, + kernel: None, + arg: Some(i as u16), + dtype: Some(dtype), + bytes: 0, + items: 0, + device: "device", + }); + } else { + self.place_device.mark_resident(id); + self.place_stats.record(crate::eir::place::PlaceEvent { + kind: crate::eir::place::EventKind::Upload, + kernel: None, + arg: Some(i as u16), + dtype: Some(dtype), + bytes, + items: 0, + device: "device", + }); + } + } + } + + // On a GPU the kernel runs there, not here. Everything above — the + // residency bookkeeping, the launch accounting — is the same either + // way, because those are properties of placement rather than of any + // particular device. + if self.place_mode == crate::eir::place::Mode::Gpu { + self.place_run_gpu(func_id, &rest, n)?; + + // Only what the kernel *wrote* differs from the host's copy. This + // used to mark every argument dirty, so a launch reading two + // buffers and writing one reported three transfers home when one + // was needed — and did three, which is worse than the accounting + // being wrong. + let kinds = crate::eir::wgsl::infer_arg_kinds(&self.module, func_id, DType::F32); + for (i, arg) in rest.iter().enumerate() { + if let Some((id, _, _)) = self.buffer_info(*arg) { + self.place_device.mark_resident(id); + if matches!( + kinds.get(i), + Some(crate::eir::wgsl::ArgKind::Buffer { writable: true, .. }) + ) { + self.place_device.mark_dirty(id); + } + } + } + // An unpinned buffer does not survive the launch, and on a GPU + // that means really releasing the allocation — otherwise + // "resident" would be a word the accounting used and the hardware + // ignored. + // + // Evicting something the kernel just wrote has to write it back + // first. A cache that drops dirty data is not a cache, it is a + // bug; and this write-back is exactly the cost a policy of keeping + // nothing resident is choosing to pay, once per launch. + let evicted = self.place_device.evict_unpinned(); + let dirty: Vec = rest + .iter() + .copied() + .filter(|v| { + v.is_ptr() + && evicted.contains(&v.as_ptr()) + && self.place_device.is_dirty(v.as_ptr()) + }) + .collect(); + for val in dirty { + let id = val.as_ptr(); + self.gpu_download_into(id, val)?; + self.place_device.clear_dirty(id); + if let Some((_, bytes, dtype)) = self.buffer_info(val) { + self.place_stats.record(crate::eir::place::PlaceEvent { + kind: crate::eir::place::EventKind::Download, + kernel: None, + arg: None, + dtype: Some(dtype), + bytes, + items: 0, + device: "gpu", + }); + } + } + self.gpu_evict(&evicted); + return Ok(Val::UNIT); + } + + // The typed executor handles the numeric subset without boxing every + // value or walking the heap for every element. Kernels outside the + // subset fall through to the general VM below, which is why this is a + // question rather than a requirement. + if matches!( + self.place_mode, + crate::eir::place::Mode::Cpu | crate::eir::place::Mode::Par + ) && crate::eir::kernel_exec::supported(&self.module, func_id) + { + if let Some(result) = self.place_run_fast(func_id, &rest, n)? { + return Ok(result); + } + } + + let mut call_args = Vec::with_capacity(rest.len() + 1); + for i in 0..n { + call_args.clear(); + call_args.push(Val::int(i)); + call_args.extend_from_slice(&rest); + self.run_call_with_captures(func_id, &call_args, caps.clone())?; + } + + // Whatever the kernel wrote now differs from the host's copy, and an + // unpinned buffer does not survive to the next launch. Both of those + // are the *default* behaviour of a device that has not been told + // anything — which is precisely what a residency handler exists to + // change, by pinning what it knows will be wanted again. + if self.place_mode.has_device_memory() { + for arg in &rest { + if let Some((id, _, _)) = self.buffer_info(*arg) { + self.place_device.mark_dirty(id); + } + } + self.place_device.evict_unpinned(); + } + + Ok(Val::UNIT) + } + + /// Run a kernel on the GPU. + /// + /// Arguments are uploaded only if they are not already there, and nothing + /// is copied back — results stay on the device until `Place.read` asks for + /// them. That is what makes the residency accounting describe the hardware + /// rather than a model of it: an upload counted here is an upload that + /// happened. + fn place_run_gpu(&mut self, func_id: FuncId, rest: &[Val], n: i64) -> Result<(), VmError> { + let kinds = crate::eir::wgsl::infer_arg_kinds(&self.module, func_id, DType::F32); + if kinds.len() != rest.len() { + return Err(self.place_error(format!( + "the kernel takes {} arguments but {} were given", + kinds.len(), + rest.len() + ))); + } + let shader = crate::eir::wgsl::emit(&self.module, func_id, &kinds) + .map_err(|e| self.place_error(format!("this kernel cannot run on a GPU: {e}")))?; + + // Check every argument against the shape the kernel body implies. A + // buffer handed to a parameter the kernel multiplies by has no + // sensible reading, and defaulting it to zero would be a confident + // wrong answer. + let mut scalars: Vec = Vec::new(); + let mut buffer_ids: Vec = Vec::new(); + let mut to_upload: Vec<(usize, Buffer, u64, DType)> = Vec::new(); + for (i, val) in rest.iter().enumerate() { + let is_buffer = matches!(self.get_obj(*val), Some(Obj::Buffer(_))); + let wants_buffer = matches!(kinds[i], crate::eir::wgsl::ArgKind::Buffer { .. }); + if is_buffer != wants_buffer { + let (given, wanted) = if is_buffer { + ("a buffer", "a number") + } else { + ("a number", "a buffer") + }; + return Err(self.place_error(format!( + "argument {} is {given}, but the kernel uses it as {wanted}", + i + 1 + ))); + } + match self.get_obj(*val) { + Some(Obj::Buffer(b)) => { + // WGSL core has no 64-bit scalar. Quietly computing an f64 + // buffer in f32 would hand back numbers of a precision the + // program never asked for and has no way to notice, so the + // launch is refused and the alternative named instead. + if !b.dtype().gpu_ok() { + return Err(self.place_error(format!( + "argument {} is a {} buffer, and a GPU has no {}-bit \ + number; run this with --place cpu or --place par, or \ + build the buffer with `buf` instead of `buf-{}`", + i + 1, + b.dtype().name(), + b.dtype().size() * 8, + b.dtype().name() + ))); + } + let id = val.as_ptr(); + buffer_ids.push(id); + to_upload.push(( + id, + crate::eir::device::narrow(b), + b.byte_len() as u64, + b.dtype(), + )); + } + _ => { + let x = if val.is_int() { + val.as_int() as f32 + } else if val.is_float() { + val.as_float() as f32 + } else { + return Err(self.place_error(format!( + "argument {} is neither a buffer nor a number", + i + 1 + ))); + }; + scalars.push(x); + } + } + } + + let device = self.gpu_device()?; + for (i, (id, narrowed, bytes, dtype)) in to_upload.iter().enumerate() { + let uploaded = device + .ensure_resident(*id, narrowed) + .map_err(|e| self.place_error(e.0))?; + self.place_stats.record(crate::eir::place::PlaceEvent { + kind: if uploaded { + crate::eir::place::EventKind::Upload + } else { + crate::eir::place::EventKind::ResidentHit + }, + kernel: None, + arg: Some(i as u16), + dtype: Some(*dtype), + bytes: if uploaded { *bytes } else { 0 }, + items: 0, + device: "gpu", + }); + } + + device + .dispatch(&shader, "main", n.max(0) as u32, &scalars, &buffer_ids) + .map_err(|e| self.place_error(format!("the GPU refused the launch: {e}")))?; + Ok(()) + } + + /// The device, opened on first use and kept. + /// + /// A bridge installed by the host wins: in a browser that is the only way + /// to reach a GPU, because WebGPU is asynchronous and this VM is not. With + /// no bridge and no `gpu` feature there is nowhere to run, and saying so + /// beats quietly running somewhere the caller did not ask for. + fn gpu_device(&mut self) -> Result, VmError> { + if let Some(d) = &self.device { + return Ok(d.clone()); + } + if let Some(d) = crate::eir::device::installed() { + self.device = Some(d.clone()); + return Ok(d); + } + #[cfg(feature = "gpu")] + { + let g: std::rc::Rc = std::rc::Rc::new( + crate::eir::gpu::Gpu::open().map_err(|e| self.place_error(e.to_string()))?, + ); + self.device = Some(g.clone()); + Ok(g) + } + #[cfg(not(feature = "gpu"))] + { + Err(self.place_error( + "this build has no GPU support; rebuild with `--features gpu`, install a \ + device bridge, or use `--place cpu`" + .to_string(), + )) + } + } + + /// Release device storage for buffers the model just evicted. + fn gpu_evict(&mut self, ids: &[usize]) { + if let Some(d) = &self.device { + for id in ids { + d.evict(*id); + } + } + } + + /// Bring a buffer's device contents home, if it lives on a device. + fn gpu_download_into(&mut self, id: usize, val: Val) -> Result<(), VmError> { + if self.place_mode != crate::eir::place::Mode::Gpu { + return Ok(()); + } + let Some(d) = self.device.clone() else { + return Ok(()); + }; + if !d.is_resident(id) { + return Ok(()); + } + let byte_len = match self.get_obj(val) { + Some(Obj::Buffer(b)) => b.byte_len(), + _ => return Ok(()), + }; + let bytes = d + .download(id, byte_len) + .map_err(|e| self.place_error(format!("reading back from the device: {e}")))?; + if let Some(Obj::Buffer(buf)) = self.heap.get_mut(id) { + let len = buf.len(); + let vals: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes(c.try_into().unwrap())) + .take(len) + .collect(); + for (i, v) in vals.into_iter().enumerate() { + buf.set(i, Val::float(v as f64)); + } + } + Ok(()) + } + + fn place_error(&self, message: String) -> VmError { + VmError::new(VmErrorKind::BuiltinType(message)).with_span(self.current_span) + } + + /// Run a kernel through the typed executor, if its arguments allow it. + /// + /// Returns `None` when the shape is not one this path handles (an aliased + /// buffer, say), so the caller falls back rather than failing. Buffers are + /// taken out of the heap for the duration and put back afterwards, which + /// is what lets the executor hold a `&mut` slice — and, in parallel mode, + /// hand disjoint pieces of it to different threads. + fn place_run_fast( + &mut self, + func_id: FuncId, + rest: &[Val], + n: i64, + ) -> Result, VmError> { + use crate::eir::kernel_exec::{KArg, KVal}; + + // Which arguments are buffers, and which of those are written? + let kinds = crate::eir::wgsl::infer_arg_kinds(&self.module, func_id, DType::F32); + if kinds.len() != rest.len() { + return Ok(None); + } + + // A buffer appearing twice would need two borrows of the same slot. + let mut seen: Vec = Vec::new(); + for val in rest { + if val.is_ptr() && matches!(self.heap.get(val.as_ptr()), Some(Obj::Buffer(_))) { + if seen.contains(&val.as_ptr()) { + return Ok(None); + } + seen.push(val.as_ptr()); + } + } + + // Take each buffer out of the heap so it can be borrowed mutably. + let mut taken: Vec<(usize, Buffer, bool)> = Vec::new(); + let mut scalars: Vec<(usize, KVal)> = Vec::new(); + for (i, val) in rest.iter().enumerate() { + let writable = matches!( + kinds[i], + crate::eir::wgsl::ArgKind::Buffer { writable: true, .. } + ); + match self.heap.get(val.as_ptr()) { + Some(Obj::Buffer(b)) if val.is_ptr() => { + taken.push((i, b.clone(), writable)); + } + _ => { + let v = if val.is_int() { + KVal::I(val.as_int()) + } else if val.is_float() { + KVal::F(val.as_float()) + } else { + return Ok(None); + }; + scalars.push((i, v)); + } + } + } + + let parallel = self.place_mode == crate::eir::place::Mode::Par; + let module = self.module.clone(); + + // Keep the buffers in one place and hand out borrows by argument + // position. Splitting inputs from outputs is what lets the parallel + // driver carve every output into per-thread slices while sharing the + // inputs whole. + let mut buffers: Vec<(usize, Buffer, bool)> = taken; + let outcome = { + let (writable, readable): (Vec<_>, Vec<_>) = + buffers.iter_mut().partition(|(_, _, w)| *w); + let inputs: Vec<(usize, &Buffer)> = + readable.into_iter().map(|(i, b, _)| (*i, &*b)).collect(); + let mut outputs: Vec<(usize, &mut Buffer)> = + writable.into_iter().map(|(i, b, _)| (*i, b)).collect(); + + if parallel { + crate::eir::kernel_exec::run_parallel( + &module, + func_id, + &scalars, + &inputs, + &mut outputs, + rest.len(), + n, + ) + } else { + crate::eir::kernel_exec::run_sequential( + &module, + func_id, + &scalars, + &inputs, + &mut outputs, + rest.len(), + n, + ) + } + }; + let outcome = outcome.map(|_| buffers); + + let buffers = match outcome { + Ok(b) => b, + Err(e) => { + return Err(VmError::new(VmErrorKind::BuiltinType(e.0)).with_span(self.current_span)) + } + }; + + // Put the results back where the program can see them. + for (i, buf, _) in buffers { + let val = rest[i]; + if let Some(slot) = self.heap.get_mut(val.as_ptr()) { + *slot = Obj::Buffer(buf); + } + } + Ok(Some(Val::UNIT)) + } + + /// Heap slot, byte length, and element type of a buffer value. + fn buffer_info(&self, val: Val) -> Option<(usize, u64, DType)> { + match self.get_obj(val) { + Some(Obj::Buffer(b)) => Some((val.as_ptr(), b.byte_len() as u64, b.dtype())), + _ => None, + } + } + fn get_obj(&self, val: Val) -> Option<&Obj> { if val.is_ptr() { self.heap.get(val.as_ptr()) @@ -670,6 +1380,37 @@ impl Vm { /// escaping continuation, e.g. the function-passing `State`). /// - `None` (tail resume): leave the frame already on top as the return /// target (the reader's non-escaping tail-resume path). + /// The continuation a program parked, if it did, and what it wants. + /// + /// Taking it clears the slot: a parked continuation has one host, and + /// leaving it visible after handing it over invites resuming twice. + pub fn take_parked(&mut self) -> Option<(Val, Val)> { + self.parked.take() + } + + /// Call a Loon value — a closure or a parked continuation — from Rust. + /// + /// This is what an asynchronous host needs. A handler that hands `resume` + /// outward and returns leaves a live continuation in this VM's heap; the + /// host can go away, do something slow, and finish the computation later + /// by calling it through here. + /// + /// The VM has to outlive the run that parked it, which is the one real + /// constraint: the continuation is a heap object, so dropping the VM drops + /// the rest of the program with it. + pub fn call_value(&mut self, f: Val, args: &[Val]) -> Result { + match self.get_obj(f).cloned() { + Some(Obj::Continuation { .. }) => { + let v = args.first().copied().unwrap_or(Val::UNIT); + let depth = self.frames.len(); + self.resume_continuation(f, v, None)?; + self.execute(depth) + } + Some(Obj::Closure(fid, caps)) => self.run_call_with_captures(fid, args, caps), + _ => Err(VmError::new(VmErrorKind::NotCallable).with_span(self.current_span)), + } + } + fn resume_continuation(&mut self, k: Val, v: Val, base: Option) -> Result<(), VmError> { let (saved, func, block, ip, mut regs, captures, perform_dst, prompt_handlers) = match self.get_obj(k) { @@ -751,7 +1492,9 @@ impl Vm { self.func = frame.func; self.block = frame.block; self.ip = frame.ip; - self.regs[ret_reg as usize] = val; + if ret_reg != RET_DISCARD { + self.regs[ret_reg as usize] = val; + } // A resumed segment may have left ephemeral handlers scoped to a // prompt frame that is now gone; drop them so they cannot shadow a // later handle for the same effect. @@ -1452,10 +2195,168 @@ impl Vm { Some(Obj::Set(items)) => items.len() as i64, Some(Obj::Str(s)) => s.len() as i64, Some(Obj::Tuple(items)) => items.len() as i64, + // A buffer has a length like any other sequence. Falling + // through to 0 here made `[len buf]` a confident lie. + Some(Obj::Buffer(b)) => b.len() as i64, _ => 0, }; Ok(Val::int(len)) } + // ── Dense buffers ── + Built::BufNew | Built::BufNewI32 | Built::BufNewF64 => { + let dtype = match built { + Built::BufNewI32 => DType::I32, + Built::BufNewF64 => DType::F64, + _ => DType::F32, + }; + let src = args.first().copied().unwrap_or(Val::UNIT); + let items: Vec = match self.get_obj(src) { + Some(Obj::Vec(v)) => v.iter().copied().collect(), + Some(Obj::Tuple(v)) => v.clone(), + _ => { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "{} expects a vector of numbers, got {}", + match built { + Built::BufNewI32 => "buf-i32", + Built::BufNewF64 => "buf-f64", + _ => "buf", + }, + self.val_to_string(src) + )))) + } + }; + let mut buf = Buffer::zeros(dtype, items.len()); + for (i, v) in items.iter().enumerate() { + if !buf.set(i, *v) { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "buffer element {i} is not a number: {}", + self.val_to_string(*v) + )))); + } + } + Ok(self.alloc(Obj::Buffer(buf))) + } + Built::BufZeros | Built::BufZerosI32 => { + let dtype = if matches!(built, Built::BufZerosI32) { + DType::I32 + } else { + DType::F32 + }; + let n = args.first().copied().unwrap_or(Val::UNIT); + if !n.is_int() || n.as_int() < 0 { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "buffer length must be a non-negative integer, got {}", + self.val_to_string(n) + )))); + } + Ok(self.alloc(Obj::Buffer(Buffer::zeros(dtype, n.as_int() as usize)))) + } + Built::BufLen => { + let b = args.first().copied().unwrap_or(Val::UNIT); + match self.get_obj(b) { + Some(Obj::Buffer(buf)) => Ok(Val::int(buf.len() as i64)), + _ => Err(VmError::new(VmErrorKind::BuiltinType(format!( + "buf-len expects a buffer, got {}", + self.val_to_string(b) + )))), + } + } + Built::BufDtype => { + let b = args.first().copied().unwrap_or(Val::UNIT); + match self.get_obj(b) { + Some(Obj::Buffer(buf)) => { + let name = buf.dtype().name(); + Ok(self.alloc_str_owned(name.to_string())) + } + _ => Err(VmError::new(VmErrorKind::BuiltinType(format!( + "buf-dtype expects a buffer, got {}", + self.val_to_string(b) + )))), + } + } + Built::BufToVec => { + let b = args.first().copied().unwrap_or(Val::UNIT); + let items: Vec = match self.get_obj(b) { + Some(Obj::Buffer(buf)) => (0..buf.len()).filter_map(|i| buf.get(i)).collect(), + _ => { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "buf->vec expects a buffer, got {}", + self.val_to_string(b) + )))) + } + }; + Ok(self.alloc(Obj::Vec(items.into_iter().collect()))) + } + Built::BufAt => { + let b = args.first().copied().unwrap_or(Val::UNIT); + let i = args.get(1).copied().unwrap_or(Val::UNIT); + let Some(Obj::Buffer(buf)) = self.get_obj(b) else { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "at expects a buffer, got {}", + self.val_to_string(b) + )))); + }; + if !i.is_int() { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "at expects an integer index, got {}", + self.val_to_string(i) + )))); + } + let idx = i.as_int(); + let len = buf.len(); + match usize::try_from(idx).ok().and_then(|u| buf.get(u)) { + Some(v) => Ok(v), + None => Err(VmError::new(VmErrorKind::BuiltinType(format!( + "at: index {idx} is outside a buffer of length {len}" + )))), + } + } + Built::BufPut => { + // Writes through the buffer in place. Every other Loon + // collection is persistent; a buffer is not, because copying + // one per element write is exactly the cost buffers exist to + // avoid. The ownership pass is what keeps this honest — it + // reports a written-through parameter as `InOut`. + let b = args.first().copied().unwrap_or(Val::UNIT); + let i = args.get(1).copied().unwrap_or(Val::UNIT); + let x = args.get(2).copied().unwrap_or(Val::UNIT); + if !b.is_ptr() || !matches!(self.heap.get(b.as_ptr()), Some(Obj::Buffer(_))) { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "put expects a buffer, got {}", + self.val_to_string(b) + )))); + } + if !i.is_int() { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "put expects an integer index, got {}", + self.val_to_string(i) + )))); + } + let idx = i.as_int(); + let ptr = b.as_ptr(); + let len = match self.heap.get(ptr) { + Some(Obj::Buffer(buf)) => buf.len(), + _ => 0, + }; + let ok = usize::try_from(idx) + .ok() + .map(|u| match self.heap.get_mut(ptr) { + Some(Obj::Buffer(buf)) => buf.set(u, x), + _ => false, + }) + .unwrap_or(false); + if !ok { + let shown = self.val_to_string(x); + return Err(VmError::new(VmErrorKind::BuiltinType( + if idx < 0 || idx as usize >= len { + format!("put: index {idx} is outside a buffer of length {len}") + } else { + format!("put: {shown} is not a number") + }, + ))); + } + Ok(b) + } Built::SeqLen => { let v = args.first().copied().unwrap_or(Val::UNIT); let len = match self.get_obj(v) { @@ -1667,17 +2568,46 @@ impl Vm { } } Built::Sum => { + // Ints stay ints; any float promotes the whole sum, matching + // the interpreter and the `Vec Num → Num` the registry + // declares. This used to filter to integers and silently drop + // everything else, so summing a vector of floats answered 0 — + // which is the shape of wrong answer that is worst, because it + // looks like an empty sum rather than a failure. let coll = args.first().copied().unwrap_or(Val::UNIT); - match self.get_obj(coll) { - Some(Obj::Vec(items)) => { - let sum: i64 = items - .iter() - .filter(|v| v.is_int()) - .map(|v| v.as_int()) - .sum(); - Ok(Val::int(sum)) + let Some(Obj::Vec(items)) = self.get_obj(coll) else { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "sum requires a vector, got {}", + self.val_to_string(coll) + ))) + .with_span(self.current_span)); + }; + let items: Vec = items.iter().copied().collect(); + let any_float = items.iter().any(|v| v.is_float()); + for v in &items { + if !v.is_int() && !v.is_float() { + let shown = self.val_to_string(*v); + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "sum requires a vector of numbers; found {shown}" + ))) + .with_span(self.current_span)); } - _ => Ok(Val::int(0)), + } + if any_float { + let total: f64 = items + .iter() + .map(|v| { + if v.is_float() { + v.as_float() + } else { + v.as_int() as f64 + } + }) + .sum(); + Ok(Val::float(total)) + } else { + let total: i64 = items.iter().map(|v| v.as_int()).sum(); + Ok(Val::int(total)) } } Built::Min | Built::Max => { @@ -3254,6 +4184,142 @@ impl Vm { // (value-or-"" lookup). The interpreter has no such ops — it hard // errors — so they were dropped for cross-backend conformance; // `Process.env` is the Option-returning form both backends share. + // ── Placement ── + // + // Reached only when no handler took the operation, so this is the + // default answer to "where does this run": right here, one work + // item at a time. A program that never mentions placement gets + // this; installing a handler is what changes the answer, and the + // program itself does not move. + ("Place", "run") => self.place_run_serial(args)?, + ("Place", "read") => { + // The one path from device-side data back to the host. It is + // an operation rather than a plain accessor so that a + // residency handler can see every synchronization point + // without the programmer having marked any of them — which is + // what the `Preload`/`PreloadMut` types in the Rust offload + // work exist to reconstruct. + let b = args.first().copied().unwrap_or(Val::UNIT); + let (bytes, dtype) = match self.buffer_info(b) { + Some((_, bytes, dtype)) => (bytes, dtype), + None => { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "Place.read expects a buffer, got {}", + self.val_to_string(b) + ))) + .with_span(self.current_span)) + } + }; + let id = b.as_ptr(); + + // On a device, a read costs a transfer only when the device + // holds something the host has not seen. On the CPU there is + // one memory, so the read is free — and saying so keeps the + // two modes comparable. + let must_transfer = match self.place_mode { + crate::eir::place::Mode::Cpu | crate::eir::place::Mode::Par => true, + crate::eir::place::Mode::Device | crate::eir::place::Mode::Gpu => { + self.place_device.is_dirty(id) + } + }; + // The transfer happens *before* the values are read, which is + // the whole content of the word "synchronization". Reading + // first would return the host's stale copy and report a + // download that changed nothing. + if must_transfer { + self.place_device.clear_dirty(id); + self.gpu_download_into(id, b)?; + self.place_stats.record(crate::eir::place::PlaceEvent { + kind: crate::eir::place::EventKind::Download, + kernel: None, + arg: None, + dtype: Some(dtype), + bytes, + items: 0, + device: self.place_mode.name(), + }); + } + + let items: Vec = match self.get_obj(b) { + Some(Obj::Buffer(buf)) => { + let n = buf.len(); + (0..n).filter_map(|i| buf.get(i)).collect() + } + _ => { + return Err(VmError::new(VmErrorKind::BuiltinType(format!( + "Place.read expects a buffer, got {}", + self.val_to_string(b) + ))) + .with_span(self.current_span)) + } + }; + self.alloc(Obj::Vec(items.into_iter().collect())) + } + ("Place", "pin") | ("Place", "unpin") => { + // Pinning says "this buffer will be wanted again, keep it". + // It is the whole vocabulary a residency policy needs, and it + // is available to any handler — no compiler pass required. + // + // Accepts a buffer or a vector of them, so a handler can pin a + // kernel's entire argument list in one operation. + let target = args.first().copied().unwrap_or(Val::UNIT); + if self.place_mode.has_device_memory() { + let pinning = op == "pin"; + let ids: Vec = match self.get_obj(target) { + Some(Obj::Vec(v)) => v + .iter() + .filter(|x| x.is_ptr()) + .filter(|x| matches!(self.heap.get(x.as_ptr()), Some(Obj::Buffer(_)))) + .map(|x| x.as_ptr()) + .collect(), + Some(Obj::Buffer(_)) => vec![target.as_ptr()], + _ => Vec::new(), + }; + for id in ids { + if pinning { + self.place_device.pin(id); + } else { + self.place_device.unpin(id); + } + } + } + target + } + ("Place", "stats") => { + let s = &self.place_stats; + let pairs = [ + ("launches", s.launches as i64), + ("work-items", s.work_items as i64), + ("uploads", s.uploads as i64), + ("downloads", s.downloads as i64), + ("bytes-in", s.bytes_in as i64), + ("bytes-out", s.bytes_out as i64), + ("resident-hits", s.resident_hits as i64), + ( + "resident-buffers", + self.place_device.resident_count() as i64, + ), + ]; + let mut map = ImMap::new(); + for (k, v) in pairs { + let key = self.intern_sym(k); + map.insert(key, Val::int(v)); + } + self.alloc(Obj::Map(map)) + } + + // ── Parking ── + // + // Keep the continuation and what it is waiting for. The handler + // that performed this returns without resuming, so the computation + // unwinds from here and the host is left holding the rest of it. + ("Host", "park") => { + let k = args.first().copied().unwrap_or(Val::UNIT); + let request = args.get(1).copied().unwrap_or(Val::UNIT); + self.parked = Some((k, request)); + Val::UNIT + } + // Real TCP/HTTP sockets (see eir/net.rs). A blocking one-at-a-time // server: listen a port, accept a request, send the response. ("Net", "listen") => { @@ -3294,6 +4360,11 @@ impl Vm { // ── Value display ────────────────────────────────────────────────── + /// Render a value the way `println` would, for hosts outside this module. + pub fn val_to_string_public(&self, val: Val) -> String { + self.val_to_string(val) + } + fn val_to_string(&self, val: Val) -> String { self.val_to_string_inner(val, false) } @@ -3383,6 +4454,14 @@ impl Vm { .unwrap_or("anon"); format!("") } + // A buffer prints as its shape, not its contents: it is a + // handle to bulk data that may live on another device, and + // `Place.read` is how a program asks for the values. Printing + // the heap slot instead would leak an allocation number into + // output that is supposed to be reproducible. + Some(Obj::Buffer(b)) => { + format!("#buf<{} x {}>", b.dtype().name(), b.len()) + } _ => format!("", val.as_ptr()), } } else { @@ -3580,6 +4659,156 @@ pub fn eval_eir_with_base_dir(src: &str, base_dir: &std::path::Path) -> Result, + /// Output produced by steps so far. `Vm::run` hands its output out in the + /// result, so a session collects it rather than losing it between steps. + output: Vec, + /// The value of the most recent completed step. + value: Val, +} + +/// Where a step of a `Session` stopped. +pub enum Step { + /// The program finished. + Done(VmResult), + /// The program parked, waiting for something the host has to supply. + /// `request` is whatever the parking handler passed along to describe it. + Parked { request: Val }, +} + +impl Session { + /// Prepare a program, without running it. + pub fn new( + src: &str, + base_dir: &std::path::Path, + mode: crate::eir::place::Mode, + ) -> Result { + let mut checker = crate::check::Checker::with_base_dir(base_dir); + let exprs = crate::parser::parse(src).map_err(|e| VmError { + kind: VmErrorKind::Trap, + span: Some(e.span), + context: Some(format!("parse error: {}", e.message)), + })?; + let errors = checker.check_program(&exprs); + if let Some(e) = module_error(&errors) { + return Err(e); + } + let module = crate::eir::lower::lower(&checker); + let mut vm = Vm::new(module); + vm.place_mode = mode; + Ok(Session { + vm, + pending: None, + output: Vec::new(), + value: Val::UNIT, + }) + } + + /// Run until the program finishes or parks. + pub fn start(&mut self) -> Result { + let result = self.vm.run()?; + Ok(self.classify(result)) + } + + /// Finish a parked step by supplying the value it was waiting for. + pub fn resume(&mut self, k: Val, value: Val) -> Result { + let val = self.vm.call_value(k, &[value])?; + let result = VmResult { + value: val, + output: std::mem::take(&mut self.vm.output), + heap_stats: self.vm.heap_stats.clone(), + }; + Ok(self.classify(result)) + } + + fn classify(&mut self, mut result: VmResult) -> Step { + self.output.append(&mut result.output); + self.value = result.value; + match self.vm.take_parked() { + Some((k, request)) => { + self.pending = Some(k); + Step::Parked { request } + } + None => { + self.pending = None; + Step::Done(result) + } + } + } + + /// The value the last step produced. After a resumed step this is the value + /// of the *continuation* — the rest of the suspended computation — which is + /// where a deferred answer ends up. + pub fn value(&self) -> Val { + self.value + } + + /// Render a value the way `println` would. + pub fn show(&self, v: Val) -> String { + self.vm.val_to_string_public(v) + } + + /// The continuation the last `Parked` step is waiting on. + pub fn pending(&self) -> Option { + self.pending + } + + /// The placement accounting so far. + pub fn stats(&self) -> &crate::eir::place::PlaceStats { + &self.vm.place_stats + } + + /// Anything printed so far, taken out. + pub fn take_output(&mut self) -> Vec { + let mut out = std::mem::take(&mut self.output); + out.append(&mut std::mem::take(&mut self.vm.output)); + out + } + + /// Build a Loon vector of numbers, for resuming a read. + pub fn vec_of_floats(&mut self, xs: &[f32]) -> Val { + let items: ImVec = xs.iter().map(|x| Val::float(*x as f64)).collect(); + self.vm.alloc(Obj::Vec(items)) + } +} + +/// Run with an explicit placement mode, returning the placement accounting +/// alongside the result. +/// +/// The counters are the interesting output for an offloaded program: how many +/// times bytes crossed the boundary, and how many crossings a residency policy +/// avoided. +pub fn eval_eir_placed( + src: &str, + base_dir: &std::path::Path, + mode: crate::eir::place::Mode, +) -> Result<(VmResult, crate::eir::place::PlaceStats), VmError> { + let mut checker = crate::check::Checker::with_base_dir(base_dir); + let exprs = crate::parser::parse(src).map_err(|e| VmError { + kind: VmErrorKind::Trap, + span: Some(e.span), + context: Some(format!("parse error: {}", e.message)), + })?; + let errors = checker.check_program(&exprs); + if let Some(e) = module_error(&errors) { + return Err(e); + } + let module = crate::eir::lower::lower(&checker); + let mut vm = Vm::new(module); + vm.place_mode = mode; + let result = vm.run()?; + Ok((result, vm.place_stats.clone())) +} + fn eval_eir_impl(src: &str, mut checker: crate::check::Checker) -> Result { let exprs = crate::parser::parse(src).map_err(|e| VmError { kind: VmErrorKind::Trap, @@ -3711,6 +4940,7 @@ mod tests { id: FuncId(id), name: Some("callee".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], @@ -3729,6 +4959,7 @@ mod tests { id: FuncId(1), name: Some("caller".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], @@ -3758,6 +4989,7 @@ mod tests { id: FuncId(2), name: Some("__main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], @@ -3803,6 +5035,7 @@ mod tests { id: FuncId(0), name: Some("callee".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], @@ -3826,6 +5059,7 @@ mod tests { id: FuncId(1), name: Some("__main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Any, evidence: vec![], captures: vec![], diff --git a/crates/loon-lang/src/eir/wasm.rs b/crates/loon-lang/src/eir/wasm.rs index 4b4623c..928f8fc 100644 --- a/crates/loon-lang/src/eir/wasm.rs +++ b/crates/loon-lang/src/eir/wasm.rs @@ -20,28 +20,16 @@ use wasm_encoder::{ TableType, TypeSection, ValType, }; -// ─── NaN-boxing constants (must match value64.rs) ───────────────────────── - -const QNAN: u64 = 0x7FF8_0000_0000_0000; -const SIGN: u64 = 0x8000_0000_0000_0000; -const BASE: u64 = SIGN | QNAN; -#[allow(dead_code)] -const TAG_MASK: u64 = 0x0007_0000_0000_0000; -const PAYLOAD: u64 = 0x0000_FFFF_FFFF_FFFF; -const TAG_INT: u64 = 0x0001_0000_0000_0000; -const TAG_PTR: u64 = 0x0000_0000_0000_0000; -const TAG_IMM: u64 = 0x0007_0000_0000_0000; -const TAG_SYM: u64 = 0x0006_0000_0000_0000; -#[allow(dead_code)] -const IMM_UNIT: u64 = 0; -#[allow(dead_code)] -const IMM_TRUE: u64 = 1; -#[allow(dead_code)] -const IMM_FALSE: u64 = 2; - -const VAL_UNIT: u64 = BASE | TAG_IMM; -const VAL_TRUE: u64 = BASE | TAG_IMM | 1; -const VAL_FALSE: u64 = BASE | TAG_IMM | 2; +// ─── NaN-boxing constants ───────────────────────────────────────────────── +// +// Imported, not re-declared: `eir::layout` is the single source of truth, so a +// change to the encoding cannot leave this backend behind. + +#[allow(unused_imports)] +use super::layout::nanbox::{ + BASE, IMM_FALSE, IMM_TRUE, IMM_UNIT, PAYLOAD, QNAN, SIGN, TAG_IMM, TAG_INT, TAG_MASK, TAG_PTR, + TAG_SYM, VAL_FALSE, VAL_NONE, VAL_TRUE, VAL_UNIT, +}; // ─── Import indices ─────────────────────────────────────────────────────── @@ -90,6 +78,12 @@ struct CompileCtx<'a> { string_data: Vec<(u32, Vec)>, /// Next free offset in the data segment. data_offset: u32, + /// Tag of the nullary `None` constructor, if this module declares one. + /// + /// `None` is an immediate singleton rather than a heap value, so that bit + /// equality is `None` equality and the falsy test stays a pure bit test. + /// Every construction site has to honour that — see `Op::Adt`. + none_tag: Option, } struct WasmFunc { @@ -190,6 +184,12 @@ impl<'a> CompileCtx<'a> { table_entries: Vec::new(), string_data: Vec::new(), data_offset: 1024, + none_tag: module + .ctors + .iter() + .rev() + .find(|c| c.name == "None") + .map(|c| c.tag), } } @@ -496,6 +496,17 @@ impl<'a> CompileCtx<'a> { } Op::Adt(dst, tag, fields, _) => { + // The nullary `None` is the immediate singleton, never a heap + // allocation. The VM normalizes it at every construction site + // because bit equality is `None` equality and truthiness is a + // bit test; a heap-allocated `None` would compare unequal to + // itself across backends and, worse, test *truthy* here while + // testing falsy on the VM. + if fields.is_empty() && self.none_tag == Some(*tag) { + out.push(WasmInstr::I64Const(VAL_NONE as i64)); + out.push(WasmInstr::LocalSet(dst.0)); + return Ok(()); + } // Layout: [tag: u16 | field_count: u16 packed as i32][fields...] let size = 4 + fields.len() as u32 * 8; emit_heap_alloc(size, addr_local, out); @@ -1434,6 +1445,7 @@ mod tests { id: FuncId(0), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1459,6 +1471,7 @@ mod tests { id: FuncId(0), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1488,6 +1501,7 @@ mod tests { id: FuncId(0), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1532,6 +1546,7 @@ mod tests { id: FuncId(0), name: Some("double".to_string()), params: vec![Ty::Int], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1548,6 +1563,7 @@ mod tests { id: FuncId(1), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1580,6 +1596,7 @@ mod tests { id: FuncId(0), name: Some("countdown".to_string()), params: vec![Ty::Int], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1617,6 +1634,7 @@ mod tests { id: FuncId(1), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Int, evidence: vec![], captures: vec![], @@ -1652,6 +1670,7 @@ mod tests { id: FuncId(0), name: Some("even".to_string()), params: vec![Ty::Int], + param_modes: Vec::new(), ret: Ty::Bool, evidence: vec![], captures: vec![], @@ -1688,6 +1707,7 @@ mod tests { id: FuncId(1), name: Some("odd".to_string()), params: vec![Ty::Int], + param_modes: Vec::new(), ret: Ty::Bool, evidence: vec![], captures: vec![], @@ -1807,6 +1827,7 @@ mod tests { id: FuncId(0), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Unit, evidence: vec![], captures: vec![], @@ -1830,6 +1851,7 @@ mod tests { id: FuncId(0), name: Some("main".to_string()), params: vec![], + param_modes: Vec::new(), ret: Ty::Str, evidence: vec![], captures: vec![], diff --git a/crates/loon-lang/src/eir/wgsl.rs b/crates/loon-lang/src/eir/wgsl.rs new file mode 100644 index 0000000..56081e2 --- /dev/null +++ b/crates/loon-lang/src/eir/wgsl.rs @@ -0,0 +1,1164 @@ +//! Kernels, as WGSL compute shaders. +//! +//! This is the part of placement that is a compiler rather than a runtime. An +//! EIR kernel function becomes a `@compute` entry point that a GPU can run: +//! buffers become storage bindings, scalars become uniform fields, and the +//! kernel's work index becomes `global_invocation_id.x`. +//! +//! Two decisions shape the whole file. +//! +//! **Specialize per launch.** A kernel is emitted for a *concrete* argument +//! signature — which arguments are buffers, of which element type, and which +//! are scalars — rather than once for all possible uses. That is the same +//! monomorphization a Rust GPU compiler performs, and it is what lets Loon +//! stay a language with inferred types while still producing a shader whose +//! every binding has a definite type. +//! +//! **No relooper.** WGSL has no `goto`, and EIR is a graph of basic blocks. +//! Rather than restructure the graph, the emitter does what the WASM backend +//! already does: a `loop` around a `switch` on a block index, with each block +//! ending by assigning the next index and continuing. Structured control flow +//! falls out of the trivial case — a single-block kernel emits as straight +//! line code with no dispatch at all — and everything else stays correct +//! without a restructuring pass to get wrong. +//! +//! What cannot be emitted is refused by name. A kernel is already restricted +//! (see `check::kernel`) to the numeric subset this file covers, so reaching a +//! rejection here means the two definitions of "kernel" have drifted apart, +//! and saying which operation was unsupported is more useful than a shader +//! that silently computes something else. + +use super::layout::DType; +use super::{BinOp, Block, Built, End, FuncId, Lit, Module, Op, Reg, UnOp}; +use std::collections::HashMap; +use std::fmt::Write as _; + +/// What a kernel argument turns out to be at a particular launch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArgKind { + /// A dense buffer, bound as a storage array. + Buffer { dtype: DType, writable: bool }, + /// A single number, passed in the uniform block. + Scalar(DType), +} + +impl ArgKind { + /// A read-only buffer of `dtype`. + pub fn input(dtype: DType) -> ArgKind { + ArgKind::Buffer { + dtype, + writable: false, + } + } + + /// A buffer the kernel writes through. + pub fn output(dtype: DType) -> ArgKind { + ArgKind::Buffer { + dtype, + writable: true, + } + } +} + +/// Why a kernel could not be emitted. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Error(pub String); + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for Error {} + +fn err(msg: impl Into) -> Result { + Err(Error(msg.into())) +} + +/// The WGSL type of a value inside a kernel. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Ty { + F32, + I32, + Bool, + /// Produced by `put`, which is a statement rather than an expression. + Unit, +} + +impl Ty { + fn wgsl(self) -> &'static str { + match self { + Ty::F32 => "f32", + Ty::I32 => "i32", + Ty::Bool => "bool", + Ty::Unit => "i32", // never read; kept declarable + } + } + + fn of(dtype: DType) -> Ty { + match dtype { + DType::F32 | DType::F64 => Ty::F32, + DType::I32 | DType::I64 => Ty::I32, + } + } + + /// The type an arithmetic result takes when the operands disagree. + /// + /// WGSL will not mix `i32` and `f32` in one expression, so the emitter has + /// to pick, and picking float is the choice that does not silently discard + /// a fractional part. + fn join(a: Ty, b: Ty) -> Ty { + match (a, b) { + (Ty::F32, _) | (_, Ty::F32) => Ty::F32, + (Ty::I32, _) | (_, Ty::I32) => Ty::I32, + _ => Ty::Bool, + } + } +} + +/// Work out each parameter's shape from how the kernel body uses it. +/// +/// A parameter that is indexed (`at`, `put`) or measured (`buf-len`) is a +/// buffer; one that is only computed with is a scalar. A buffer written +/// through is bound `read_write`, and one only read is bound `read`. +/// +/// This is the same question the ownership pass answers for transfer +/// direction, asked of the lowered code: the kernel's own body says what its +/// arguments are, so a launch does not have to be told. +pub fn infer_arg_kinds(module: &Module, func: FuncId, dtype: DType) -> Vec { + let Some(f) = module.funcs.get(func.0 as usize) else { + return Vec::new(); + }; + let arity = f.params.len().saturating_sub(1); + let mut is_buffer = vec![false; arity]; + let mut written = vec![false; arity]; + + // Parameters occupy registers 0..n, the work index first. + let param_slot = |r: Reg| -> Option { + if r.0 >= 1 && (r.0 as usize) <= arity { + Some(r.0 as usize - 1) + } else { + None + } + }; + + for block in &f.blocks { + for op in &block.ops { + if let Op::Builtin(_, built, args, _) = op { + let target = args.first().copied().and_then(param_slot); + match (built, target) { + (Built::BufAt | Built::BufLen, Some(i)) => is_buffer[i] = true, + (Built::BufPut, Some(i)) => { + is_buffer[i] = true; + written[i] = true; + } + _ => {} + } + } + } + } + + (0..arity) + .map(|i| { + if is_buffer[i] { + ArgKind::Buffer { + dtype, + writable: written[i], + } + } else { + ArgKind::Scalar(dtype) + } + }) + .collect() +} + +/// Emit a WGSL compute shader for `func`, specialized to `args`. +/// +/// `args` describes the kernel's parameters *after* the leading work index, +/// which is always supplied by the dispatch itself. +pub fn emit(module: &Module, func: FuncId, args: &[ArgKind]) -> Result { + let f = module + .funcs + .get(func.0 as usize) + .ok_or_else(|| Error(format!("no function {func:?}")))?; + + if f.params.len() != args.len() + 1 { + return err(format!( + "kernel '{}' takes {} parameters, but {} arguments were given \ + (the work index is supplied by the dispatch)", + f.name.as_deref().unwrap_or(""), + f.params.len(), + args.len() + 1 + )); + } + + let mut e = Emitter { + module, + args, + types: HashMap::new(), + buffers: HashMap::new(), + scalars: HashMap::new(), + }; + e.bind_params(); + e.check_buffer_uses(&f.blocks)?; + e.infer(&f.blocks)?; + + let name = f.name.as_deref().unwrap_or("kernel"); + let mut out = String::new(); + e.header(&mut out, name); + + writeln!(out, "@compute @workgroup_size(64)").unwrap(); + writeln!( + out, + "fn main(@builtin(global_invocation_id) gid: vec3) {{" + ) + .unwrap(); + writeln!(out, " let idx: i32 = i32(gid.x);").unwrap(); + writeln!(out, " if (idx >= params.n) {{ return; }}").unwrap(); + + e.declare_registers(&mut out, &f.blocks); + writeln!(out, " var r{}: i32 = idx;", 0).unwrap(); + + if f.blocks.len() == 1 { + // The common case: no control flow, so no dispatch machinery. + e.emit_block(&mut out, &f.blocks, &f.blocks[0], " ", false)?; + } else { + writeln!(out, " var blk: i32 = 0;").unwrap(); + writeln!(out, " loop {{").unwrap(); + writeln!(out, " switch blk {{").unwrap(); + for block in &f.blocks { + writeln!(out, " case {}: {{", block.id.0).unwrap(); + e.emit_block(&mut out, &f.blocks, block, " ", true)?; + writeln!(out, " }}").unwrap(); + } + writeln!(out, " default: {{ return; }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + } + + writeln!(out, "}}").unwrap(); + Ok(out) +} + +struct Emitter<'a> { + module: &'a Module, + args: &'a [ArgKind], + /// Inferred type of each register. + types: HashMap, + /// Register → binding index, for buffer parameters. + buffers: HashMap, + /// Register → uniform field name, for scalar parameters. + scalars: HashMap, +} + +impl Emitter<'_> { + /// Parameters occupy registers 0..n in order, the index first. + fn bind_params(&mut self) { + self.types.insert(0, Ty::I32); + for (i, arg) in self.args.iter().enumerate() { + let reg = (i + 1) as u32; + match arg { + ArgKind::Buffer { dtype, .. } => { + self.buffers.insert(reg, i); + self.types.insert(reg, Ty::of(*dtype)); + } + ArgKind::Scalar(dtype) => { + self.scalars.insert(reg, format!("params.s{i}")); + self.types.insert(reg, Ty::of(*dtype)); + } + } + } + } + + /// Refuse a launch signature that disagrees with the kernel body. + /// + /// A buffer is a binding, not a value: there is no WGSL expression that + /// denotes one. So a kernel that multiplies by an argument the launch + /// declared to be a buffer cannot be compiled, and saying which argument + /// is far more useful than emitting a shader that names an identifier + /// nothing declared. This is the check that turns a mismatch between a + /// kernel and its call site into a sentence rather than a driver error. + fn check_buffer_uses(&self, blocks: &[Block]) -> Result<(), Error> { + for block in blocks { + for op in &block.ops { + let value_reads: Vec = match op { + // The first argument of these is a binding, by design. + Op::Builtin(_, Built::BufAt, a, _) => a.iter().skip(1).copied().collect(), + Op::Builtin(_, Built::BufPut, a, _) => a.iter().skip(1).copied().collect(), + Op::Builtin(_, Built::BufLen, _, _) => Vec::new(), + Op::Builtin(_, _, a, _) => a.clone(), + Op::Mov(_, s, _) => vec![*s], + Op::Bin(_, _, a, b, _) => vec![*a, *b], + Op::Un(_, _, a, _) => vec![*a], + _ => Vec::new(), + }; + for r in value_reads { + if let Some(binding) = self.buffers.get(&r.0) { + return err(format!( + "argument {} was given as a buffer, but the kernel uses it as a number; a buffer can only be indexed (`at`, `put`) or measured (`buf-len`)", + binding + 1 + )); + } + } + } + } + Ok(()) + } + + fn header(&self, out: &mut String, name: &str) { + writeln!(out, "// kernel '{name}', specialized for this launch").unwrap(); + writeln!(out, "struct Params {{").unwrap(); + writeln!(out, " n: i32,").unwrap(); + for (i, arg) in self.args.iter().enumerate() { + if let ArgKind::Scalar(dtype) = arg { + writeln!(out, " s{}: {},", i, Ty::of(*dtype).wgsl()).unwrap(); + } + } + writeln!(out, "}}").unwrap(); + writeln!(out, "@group(0) @binding(0) var params: Params;").unwrap(); + + // Bindings are numbered contiguously from 1 (0 is the uniform block), + // skipping scalar arguments, so the layout a host builds matches the + // shader without a gap to reason about. + let mut binding = 1; + for (i, arg) in self.args.iter().enumerate() { + if let ArgKind::Buffer { dtype, writable } = arg { + let access = if *writable { "read_write" } else { "read" }; + writeln!( + out, + "@group(0) @binding({binding}) var b{}: array<{}>;", + i, + Ty::of(*dtype).wgsl() + ) + .unwrap(); + binding += 1; + } + } + writeln!(out).unwrap(); + } + + /// Walk the blocks assigning a type to every register. + /// + /// Repeats until stable so a register defined in a later block — a loop + /// back-edge carrying a value — still gets a type before it is used. + fn infer(&mut self, blocks: &[Block]) -> Result<(), Error> { + // Iterate until nothing changes, not merely until nothing new appears. + // + // A loop-carried value is the case that makes the difference. Its type + // is set by whichever predecessor is visited first, and a later one may + // demand a wider type; stopping as soon as every register *had* a type + // left the earlier, narrower answer in place. The result was a shader + // that assigned an i32 register to an f32 one and failed validation — + // which is at least a loud failure, but only because something checks. + let limit = blocks.len().max(1) * 4 + 4; + for _ in 0..limit { + let snapshot = self.types.clone(); + for block in blocks { + for op in &block.ops { + self.infer_op(op)?; + } + // Block parameters take the widest type any predecessor passes. + if let End::Jmp(_, args) = &block.end { + // A jump's arguments land in the *target* block's params. + if let Some(target) = self.block_by_id(blocks, block) { + for (i, a) in args.iter().enumerate() { + if let (Some(t), Some(p)) = (self.ty_opt(*a), target.params.get(i)) { + self.widen(p.0, t); + } + } + } + } + if let End::Recur(args) = &block.end { + // `recur` re-enters the function's entry block. + if let (Some(entry), true) = (blocks.first(), true) { + for (i, a) in args.iter().enumerate() { + if let (Some(t), Some(p)) = (self.ty_opt(*a), entry.params.get(i)) { + self.widen(p.0, t); + } + } + } + } + } + if self.types == snapshot { + break; + } + } + Ok(()) + } + + /// The block a terminator jumps to. + fn block_by_id<'b>(&self, blocks: &'b [Block], from: &Block) -> Option<&'b Block> { + let End::Jmp(target, _) = &from.end else { + return None; + }; + blocks.iter().find(|b| b.id == *target) + } + + /// Record a register's type, widening rather than replacing. + /// + /// WGSL will not mix i32 and f32, so when predecessors disagree the wider + /// type has to win everywhere — otherwise the declaration and the + /// assignments describe different types. + fn widen(&mut self, reg: u32, t: Ty) { + let merged = match self.types.get(®).copied() { + Some(existing) if existing != t => Ty::join(existing, t), + Some(existing) => existing, + None => t, + }; + self.types.insert(reg, merged); + } + + fn infer_op(&mut self, op: &Op) -> Result<(), Error> { + match op { + Op::Lit(d, lit, _) => { + let t = match lit { + Lit::Int(_) => Ty::I32, + Lit::Float(_) => Ty::F32, + Lit::Bool(_) => Ty::Bool, + Lit::Unit => Ty::Unit, + Lit::Str(_) | Lit::Keyword(_) => { + return err("a kernel cannot use strings or keywords") + } + }; + self.types.insert(d.0, t); + } + Op::Mov(d, s, _) => { + if let Some(t) = self.ty_opt(*s) { + self.widen(d.0, t); + } + } + Op::Bin(d, bop, a, b, _) => { + let t = match bop { + BinOp::Eq + | BinOp::Ne + | BinOp::Lt + | BinOp::Gt + | BinOp::Le + | BinOp::Ge + | BinOp::And + | BinOp::Or => Ty::Bool, + BinOp::Concat => return err("a kernel cannot concatenate strings"), + _ => Ty::join( + self.ty_opt(*a).unwrap_or(Ty::F32), + self.ty_opt(*b).unwrap_or(Ty::F32), + ), + }; + self.types.insert(d.0, t); + } + Op::Un(d, uop, a, _) => { + let t = match uop { + UnOp::Neg => self.ty_opt(*a).unwrap_or(Ty::F32), + UnOp::Not => Ty::Bool, + }; + self.types.insert(d.0, t); + } + Op::Builtin(d, built, args, _) => { + let t = self.builtin_type(*built, args)?; + self.types.insert(d.0, t); + } + Op::Call(d, f, _, _) => { + // A kernel calling a kernel: its result type is whatever that + // kernel's own body produces. Kernels return unit today, so + // this is only reached for effectful helpers. + let _ = f; + self.types.insert(d.0, Ty::Unit); + } + other => return err(format!("{} cannot run on a GPU", describe(other))), + } + Ok(()) + } + + fn builtin_type(&self, built: Built, args: &[Reg]) -> Result { + Ok(match built { + Built::BufAt => { + let buf = args.first().copied().unwrap_or(Reg(0)); + self.ty_opt(buf).unwrap_or(Ty::F32) + } + Built::BufPut => Ty::Unit, + Built::BufLen => Ty::I32, + Built::Int => Ty::I32, + Built::Float => Ty::F32, + Built::Not => Ty::Bool, + Built::Sqrt + | Built::Pow + | Built::Floor + | Built::Ceil + | Built::Round + | Built::Sin + | Built::Cos + | Built::Tan + | Built::Asin + | Built::Acos + | Built::Atan + | Built::Atan2 + | Built::Log + | Built::Log10 + | Built::Exp => Ty::F32, + Built::Abs | Built::Min | Built::Max => { + let a = args + .first() + .and_then(|r| self.ty_opt(*r)) + .unwrap_or(Ty::F32); + let b = args.get(1).and_then(|r| self.ty_opt(*r)).unwrap_or(a); + Ty::join(a, b) + } + other => return err(format!("builtin {other:?} has no GPU equivalent")), + }) + } + + fn ty_opt(&self, r: Reg) -> Option { + self.types.get(&r.0).copied() + } + + fn ty(&self, r: Reg) -> Ty { + self.ty_opt(r).unwrap_or(Ty::F32) + } + + /// How a register is read in an expression. + fn read(&self, r: Reg) -> String { + if let Some(name) = self.scalars.get(&r.0) { + name.clone() + } else { + format!("r{}", r.0) + } + } + + fn declare_registers(&self, out: &mut String, blocks: &[Block]) { + let mut seen: Vec = Vec::new(); + for block in blocks { + for p in &block.params { + seen.push(p.0); + } + for op in &block.ops { + if let Some(d) = dest(op) { + seen.push(d.0); + } + } + } + seen.sort_unstable(); + seen.dedup(); + for r in seen { + // Parameters already exist: buffers as bindings, scalars as + // uniform fields, the index as r0. + if r == 0 || self.buffers.contains_key(&r) || self.scalars.contains_key(&r) { + continue; + } + let t = self.types.get(&r).copied().unwrap_or(Ty::F32); + let init = match t { + Ty::F32 => "0.0", + Ty::I32 | Ty::Unit => "0", + Ty::Bool => "false", + }; + writeln!(out, " var r{}: {} = {};", r, t.wgsl(), init).unwrap(); + } + } + + fn emit_block( + &self, + out: &mut String, + blocks: &[Block], + block: &Block, + pad: &str, + dispatched: bool, + ) -> Result<(), Error> { + for op in &block.ops { + self.emit_op(out, op, pad)?; + } + match &block.end { + End::Ret(_) => { + writeln!(out, "{pad}return;").unwrap(); + } + End::Trap => { + // A kernel has no way to report a trap, and inventing a value + // would be worse than stopping. + writeln!(out, "{pad}return;").unwrap(); + } + End::Jmp(target, args) => { + self.emit_branch_args(out, blocks, *target, args, pad); + if dispatched { + writeln!(out, "{pad}blk = {}; continue;", target.0).unwrap(); + } else { + return err("a jump needs the block dispatcher"); + } + } + End::Br(cond, t, f) => { + if !dispatched { + return err("a branch needs the block dispatcher"); + } + writeln!( + out, + "{pad}if ({}) {{ blk = {}; }} else {{ blk = {}; }}", + self.read_as(*cond, Ty::Bool), + t.0, + f.0 + ) + .unwrap(); + writeln!(out, "{pad}continue;").unwrap(); + } + End::Switch(scrutinee, arms, default) => { + if !dispatched { + return err("a switch needs the block dispatcher"); + } + writeln!(out, "{pad}switch {} {{", self.read_as(*scrutinee, Ty::I32)).unwrap(); + for (value, target) in arms { + writeln!(out, "{pad} case {value}: {{ blk = {}; }}", target.0).unwrap(); + } + writeln!(out, "{pad} default: {{ blk = {}; }}", default.0).unwrap(); + writeln!(out, "{pad}}}").unwrap(); + writeln!(out, "{pad}continue;").unwrap(); + } + End::Recur(args) => { + if !dispatched { + return err("a loop needs the block dispatcher"); + } + self.emit_branch_args(out, blocks, block.id, args, pad); + writeln!(out, "{pad}blk = {}; continue;", block.id.0).unwrap(); + } + End::Tail(..) | End::TailInvoke(..) => { + return err("tail calls have no GPU equivalent"); + } + } + Ok(()) + } + + /// Copy branch arguments into the target block's parameter registers. + fn emit_branch_args( + &self, + out: &mut String, + blocks: &[Block], + target: super::BlockId, + args: &[Reg], + pad: &str, + ) { + // Within this function's blocks only. `BlockId` is per-function, so + // searching the whole module found whichever function happened to have + // a block with the same number — and copied jump arguments into its + // parameters, with its types. The shader that came out assigned an i32 + // register to an f32 one. + let Some(block) = blocks.iter().find(|b| b.id == target) else { + return; + }; + for (p, a) in block.params.iter().zip(args.iter()) { + if p.0 != a.0 { + writeln!(out, "{pad}r{} = {};", p.0, self.read_as(*a, self.ty(*p))).unwrap(); + } + } + } + + /// Read a register, converting if the context needs another type. + fn read_as(&self, r: Reg, want: Ty) -> String { + let have = self.ty(r); + let text = self.read(r); + if have == want || want == Ty::Unit { + return text; + } + match (have, want) { + (Ty::I32, Ty::F32) => format!("f32({text})"), + (Ty::F32, Ty::I32) => format!("i32({text})"), + (Ty::Bool, Ty::I32) => format!("select(0, 1, {text})"), + (Ty::Bool, Ty::F32) => format!("select(0.0, 1.0, {text})"), + (Ty::I32, Ty::Bool) => format!("({text} != 0)"), + (Ty::F32, Ty::Bool) => format!("({text} != 0.0)"), + // A unit value reaching a numeric slot is zero, which is how the + // CPU executor already reads it. It happens on a control-flow path + // that produces no value — the fallthrough of a loop, say — and + // leaving it unconverted emitted a bare i32 register where WGSL + // wanted an f32, which is a shader that does not compile. + (Ty::Unit, Ty::F32) => format!("f32({text})"), + (Ty::Unit, Ty::Bool) => format!("({text} != 0)"), + _ => text, + } + } + + fn emit_op(&self, out: &mut String, op: &Op, pad: &str) -> Result<(), Error> { + match op { + Op::Lit(d, lit, _) => { + let v = match lit { + Lit::Int(n) => format!("{n}"), + Lit::Float(f) => format_float(*f), + Lit::Bool(b) => format!("{b}"), + Lit::Unit => "0".to_string(), + Lit::Str(_) | Lit::Keyword(_) => { + return err("a kernel cannot use strings or keywords") + } + }; + writeln!(out, "{pad}r{} = {v};", d.0).unwrap(); + } + Op::Mov(d, s, _) => { + writeln!(out, "{pad}r{} = {};", d.0, self.read_as(*s, self.ty(*d))).unwrap(); + } + Op::Bin(d, bop, a, b, _) => { + let text = self.binop(*bop, *a, *b)?; + writeln!(out, "{pad}r{} = {text};", d.0).unwrap(); + } + Op::Un(d, uop, a, _) => { + let text = match uop { + UnOp::Neg => format!("-({})", self.read_as(*a, self.ty(*d))), + UnOp::Not => format!("!({})", self.read_as(*a, Ty::Bool)), + }; + writeln!(out, "{pad}r{} = {text};", d.0).unwrap(); + } + Op::Builtin(d, built, args, _) => self.emit_builtin(out, *d, *built, args, pad)?, + other => return err(format!("{} cannot run on a GPU", describe(other))), + } + Ok(()) + } + + fn binop(&self, bop: BinOp, a: Reg, b: Reg) -> Result { + let joined = Ty::join(self.ty(a), self.ty(b)); + let (x, y) = (self.read_as(a, joined), self.read_as(b, joined)); + Ok(match bop { + BinOp::Add => format!("({x} + {y})"), + BinOp::Sub => format!("({x} - {y})"), + BinOp::Mul => format!("({x} * {y})"), + BinOp::Div => format!("({x} / {y})"), + BinOp::Rem => { + if joined == Ty::F32 { + // WGSL's `%` is defined for floats, but naming the intent + // keeps the emitted shader readable. + format!("({x} % {y})") + } else { + format!("({x} % {y})") + } + } + BinOp::Eq => format!("({x} == {y})"), + BinOp::Ne => format!("({x} != {y})"), + BinOp::Lt => format!("({x} < {y})"), + BinOp::Gt => format!("({x} > {y})"), + BinOp::Le => format!("({x} <= {y})"), + BinOp::Ge => format!("({x} >= {y})"), + BinOp::And => format!( + "({} && {})", + self.read_as(a, Ty::Bool), + self.read_as(b, Ty::Bool) + ), + BinOp::Or => format!( + "({} || {})", + self.read_as(a, Ty::Bool), + self.read_as(b, Ty::Bool) + ), + BinOp::Concat => return err("a kernel cannot concatenate strings"), + }) + } + + fn emit_builtin( + &self, + out: &mut String, + d: Reg, + built: Built, + args: &[Reg], + pad: &str, + ) -> Result<(), Error> { + // Buffer access is the one place a register names a binding rather + // than a value. + match built { + Built::BufAt => { + let buf = args.first().copied().unwrap_or(Reg(0)); + let idx = args.get(1).copied().unwrap_or(Reg(0)); + let b = self.binding(buf)?; + writeln!( + out, + "{pad}r{} = b{b}[u32({})];", + d.0, + self.read_as(idx, Ty::I32) + ) + .unwrap(); + return Ok(()); + } + Built::BufPut => { + let buf = args.first().copied().unwrap_or(Reg(0)); + let idx = args.get(1).copied().unwrap_or(Reg(0)); + let val = args.get(2).copied().unwrap_or(Reg(0)); + let b = self.binding(buf)?; + let elem = self.ty(buf); + writeln!( + out, + "{pad}b{b}[u32({})] = {};", + self.read_as(idx, Ty::I32), + self.read_as(val, elem) + ) + .unwrap(); + return Ok(()); + } + Built::BufLen => { + let buf = args.first().copied().unwrap_or(Reg(0)); + let b = self.binding(buf)?; + writeln!(out, "{pad}r{} = i32(arrayLength(&b{b}));", d.0).unwrap(); + return Ok(()); + } + _ => {} + } + + let a = |i: usize, t: Ty| -> String { + args.get(i) + .map(|r| self.read_as(*r, t)) + .unwrap_or_else(|| "0.0".to_string()) + }; + let text = match built { + Built::Sqrt => format!("sqrt({})", a(0, Ty::F32)), + Built::Pow => format!("pow({}, {})", a(0, Ty::F32), a(1, Ty::F32)), + Built::Floor => format!("floor({})", a(0, Ty::F32)), + Built::Ceil => format!("ceil({})", a(0, Ty::F32)), + Built::Round => format!("round({})", a(0, Ty::F32)), + Built::Sin => format!("sin({})", a(0, Ty::F32)), + Built::Cos => format!("cos({})", a(0, Ty::F32)), + Built::Tan => format!("tan({})", a(0, Ty::F32)), + Built::Asin => format!("asin({})", a(0, Ty::F32)), + Built::Acos => format!("acos({})", a(0, Ty::F32)), + Built::Atan => format!("atan({})", a(0, Ty::F32)), + Built::Atan2 => format!("atan2({}, {})", a(0, Ty::F32), a(1, Ty::F32)), + Built::Log => format!("log({})", a(0, Ty::F32)), + Built::Log10 => format!("(log({}) / log(10.0))", a(0, Ty::F32)), + Built::Exp => format!("exp({})", a(0, Ty::F32)), + Built::Abs => { + let t = self.ty(d); + format!("abs({})", a(0, t)) + } + Built::Min => { + let t = self.ty(d); + format!("min({}, {})", a(0, t), a(1, t)) + } + Built::Max => { + let t = self.ty(d); + format!("max({}, {})", a(0, t), a(1, t)) + } + Built::Int => format!("i32({})", a(0, Ty::F32)), + Built::Float => format!("f32({})", a(0, Ty::I32)), + Built::Not => format!("!({})", a(0, Ty::Bool)), + other => return err(format!("builtin {other:?} has no GPU equivalent")), + }; + writeln!(out, "{pad}r{} = {text};", d.0).unwrap(); + Ok(()) + } + + /// The storage binding a register refers to, or an error naming the + /// mismatch — a kernel indexing something that is not a buffer means the + /// launch signature and the kernel body disagree. + fn binding(&self, r: Reg) -> Result { + self.buffers.get(&r.0).copied().ok_or_else(|| { + Error(format!( + "register r{} is indexed as a buffer but was not bound as one", + r.0 + )) + }) + } +} + +/// WGSL has no integer-valued float literals: `1` is an i32, `1.0` is an f32. +fn format_float(f: f64) -> String { + if f.is_nan() { + // No NaN literal in WGSL; construct one that the validator accepts. + return "(0.0 / 0.0)".to_string(); + } + if f.is_infinite() { + return if f > 0.0 { + "(1.0 / 0.0)".to_string() + } else { + "(-1.0 / 0.0)".to_string() + }; + } + let s = format!("{f:?}"); + if s.contains('.') || s.contains('e') || s.contains('E') { + s + } else { + format!("{s}.0") + } +} + +fn dest(op: &Op) -> Option { + match op { + Op::Lit(d, ..) + | Op::Mov(d, ..) + | Op::Upval(d, ..) + | Op::Bin(d, ..) + | Op::Un(d, ..) + | Op::Call(d, ..) + | Op::Invoke(d, ..) + | Op::Close(d, ..) + | Op::Vec(d, ..) + | Op::Map(d, ..) + | Op::Set(d, ..) + | Op::Tup(d, ..) + | Op::Adt(d, ..) + | Op::Field(d, ..) + | Op::Tag(d, ..) + | Op::Perform(d, ..) + | Op::Builtin(d, ..) + | Op::PushHandler(d, ..) => Some(*d), + Op::PopHandler(_) => None, + } +} + +/// A phrase naming what an operation does, for a refusal message. +fn describe(op: &Op) -> &'static str { + match op { + Op::Upval(..) => "reading a closure upvalue", + Op::Invoke(..) => "calling a closure", + Op::Close(..) => "creating a closure", + Op::Vec(..) | Op::Map(..) | Op::Set(..) | Op::Tup(..) => "building a collection", + Op::Adt(..) => "constructing a value", + Op::Field(..) => "field access", + Op::Tag(..) => "reading a constructor tag", + Op::Perform(..) => "performing an effect", + Op::PushHandler(..) | Op::PopHandler(..) => "installing a handler", + _ => "this operation", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::check::Checker; + use crate::eir::lower::lower; + use crate::parser::parse; + + /// Lower a program and emit WGSL for the named kernel. + fn emit_kernel(src: &str, name: &str, args: &[ArgKind]) -> Result { + let exprs = parse(src).expect("parses"); + let mut checker = Checker::new(); + let errors = checker.check_program(&exprs); + assert!(errors.is_empty(), "check errors: {errors:?}"); + let module = lower(&checker); + let func = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(name)) + .unwrap_or_else(|| panic!("no kernel '{name}'")); + emit(&module, func.id, args) + } + + /// Validate WGSL the way a GPU driver would, without needing one. + /// + /// This is the check the offload paper says is missing for cross-target + /// work: the shader is parsed and type-checked in CI, so a kernel that + /// would be rejected on a device is rejected here first. + fn validate(wgsl: &str) { + let module = naga::front::wgsl::parse_str(wgsl) + .unwrap_or_else(|e| panic!("WGSL did not parse: {e}\n\n{wgsl}")); + let mut validator = naga::valid::Validator::new( + naga::valid::ValidationFlags::all(), + naga::valid::Capabilities::empty(), + ); + validator + .validate(&module) + .unwrap_or_else(|e| panic!("WGSL did not validate: {e:?}\n\n{wgsl}")); + } + + #[test] + fn saxpy_emits_a_valid_shader() { + let wgsl = emit_kernel( + "[kernel saxpy [i a x y out] [put out i [+ [* a [at x i]] [at y i]]]]", + "saxpy", + &[ + ArgKind::Scalar(DType::F32), + ArgKind::input(DType::F32), + ArgKind::input(DType::F32), + ArgKind::output(DType::F32), + ], + ) + .expect("emits"); + + assert!(wgsl.contains("@compute"), "{wgsl}"); + assert!(wgsl.contains("var b1"), "{wgsl}"); + assert!(wgsl.contains("var b3"), "{wgsl}"); + validate(&wgsl); + } + + #[test] + fn a_single_block_kernel_needs_no_dispatcher() { + // The straight-line case should read like the source did. + let wgsl = emit_kernel( + "[kernel double [i b] [put b i [* 2.0 [at b i]]]]", + "double", + &[ArgKind::output(DType::F32)], + ) + .expect("emits"); + assert!( + !wgsl.contains("switch blk"), + "no control flow should mean no dispatcher:\n{wgsl}" + ); + validate(&wgsl); + } + + #[test] + fn branching_kernels_use_the_block_dispatcher() { + let wgsl = emit_kernel( + "[kernel clamp01 [i b] \ + [let v [at b i]] \ + [put b i [if [> v 1.0] 1.0 [if [< v 0.0] 0.0 v]]]]", + "clamp01", + &[ArgKind::output(DType::F32)], + ) + .expect("emits"); + assert!(wgsl.contains("switch blk"), "{wgsl}"); + validate(&wgsl); + } + + #[test] + fn math_builtins_map_onto_wgsl_intrinsics() { + let wgsl = emit_kernel( + "[kernel mathy [i b] \ + [put b i [+ [sqrt [abs [at b i]]] [* [sin [at b i]] [exp 1.0]]]]]", + "mathy", + &[ArgKind::output(DType::F32)], + ) + .expect("emits"); + for f in ["sqrt(", "abs(", "sin(", "exp("] { + assert!(wgsl.contains(f), "expected {f} in:\n{wgsl}"); + } + validate(&wgsl); + } + + #[test] + fn integer_buffers_emit_integer_bindings() { + let wgsl = emit_kernel( + "[kernel bump [i b] [put b i [+ 1 [at b i]]]]", + "bump", + &[ArgKind::output(DType::I32)], + ) + .expect("emits"); + assert!(wgsl.contains("array"), "{wgsl}"); + validate(&wgsl); + } + + #[test] + fn mixed_arithmetic_converts_rather_than_mixing_types() { + // WGSL will not add an i32 to an f32. The emitter has to insert the + // conversion, and the validator is what proves it did. + let wgsl = emit_kernel( + "[kernel mixed [i b] [put b i [+ 1 [at b i]]]]", + "mixed", + &[ArgKind::output(DType::F32)], + ) + .expect("emits"); + validate(&wgsl); + } + + #[test] + fn buffer_length_is_available_to_a_kernel() { + let wgsl = emit_kernel( + "[kernel norm [i b] [put b i [if [> [buf-len b] 0] [at b i] 0.0]]]", + "norm", + &[ArgKind::output(DType::F32)], + ) + .expect("emits"); + assert!(wgsl.contains("arrayLength"), "{wgsl}"); + validate(&wgsl); + } + + #[test] + fn the_work_index_bounds_check_is_always_emitted() { + // A GPU dispatch rounds up to whole workgroups, so the last group runs + // invocations past the end of the data. Without this check they would + // write outside the buffer. + let wgsl = emit_kernel( + "[kernel k [i b] [put b i 1.0]]", + "k", + &[ArgKind::output(DType::F32)], + ) + .expect("emits"); + assert!(wgsl.contains("if (idx >= params.n) { return; }"), "{wgsl}"); + } + + #[test] + fn a_wrong_argument_count_is_refused_by_name() { + let e = emit_kernel( + "[kernel k [i a b] [put b i a]]", + "k", + &[ArgKind::output(DType::F32)], + ) + .expect_err("should refuse"); + assert!(e.0.contains("parameters"), "{}", e.0); + } + + #[test] + fn scalars_arrive_through_the_uniform_block() { + let wgsl = emit_kernel( + "[kernel scale [i s b] [put b i [* s [at b i]]]]", + "scale", + &[ArgKind::Scalar(DType::F32), ArgKind::output(DType::F32)], + ) + .expect("emits"); + assert!(wgsl.contains("s0: f32"), "{wgsl}"); + assert!(wgsl.contains("params.s0"), "{wgsl}"); + validate(&wgsl); + } + + #[test] + fn buffer_bindings_are_numbered_without_gaps() { + // A scalar argument occupies no binding, so the buffers after it must + // still number consecutively — a host builds its bind group layout + // from these, and a hole in the sequence is a mismatch waiting to + // happen. + let wgsl = emit_kernel( + "[kernel saxpy [i a x y out] [put out i [+ [* a [at x i]] [at y i]]]]", + "saxpy", + &[ + ArgKind::Scalar(DType::F32), + ArgKind::input(DType::F32), + ArgKind::input(DType::F32), + ArgKind::output(DType::F32), + ], + ) + .expect("emits"); + assert!(wgsl.contains("@binding(1) var b1"), "{wgsl}"); + assert!(wgsl.contains("@binding(2) var b2"), "{wgsl}"); + assert!( + wgsl.contains("@binding(3) var b3"), + "{wgsl}" + ); + validate(&wgsl); + } + + #[test] + fn argument_shapes_are_read_off_the_kernel_body() { + // `a` is multiplied, so it is a scalar; `x` is indexed, so it is a + // buffer; `out` is written through, so it is writable. Nobody said so. + let src = "[kernel saxpy [i a x y out] [put out i [+ [* a [at x i]] [at y i]]]]"; + let exprs = parse(src).expect("parses"); + let mut checker = Checker::new(); + assert!(checker.check_program(&exprs).is_empty()); + let module = lower(&checker); + let f = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some("saxpy")) + .expect("saxpy"); + let kinds = infer_arg_kinds(&module, f.id, DType::F32); + assert_eq!( + kinds, + vec![ + ArgKind::Scalar(DType::F32), + ArgKind::input(DType::F32), + ArgKind::input(DType::F32), + ArgKind::output(DType::F32), + ] + ); + } + + #[test] + fn a_buffer_used_as_a_number_is_refused_by_argument() { + // The launch says argument 1 is a buffer; the body multiplies by it. + // There is no WGSL expression for a binding, so this cannot be + // compiled — and the message says which argument disagrees. + let e = emit_kernel( + "[kernel scale [i s b] [put b i [* s [at b i]]]]", + "scale", + &[ArgKind::input(DType::F32), ArgKind::output(DType::F32)], + ) + .expect_err("should refuse"); + assert!(e.0.contains("argument 1"), "{}", e.0); + assert!(e.0.contains("as a number"), "{}", e.0); + } + + #[test] + fn float_literals_are_never_emitted_as_integers() { + // `1` and `1.0` are different types in WGSL, and getting this wrong is + // a validation error rather than a wrong answer — but only because + // something checks. + assert_eq!(format_float(1.0), "1.0"); + assert_eq!(format_float(-2.5), "-2.5"); + assert_eq!(format_float(0.0), "0.0"); + } +} diff --git a/crates/loon-lang/src/errors/codes.rs b/crates/loon-lang/src/errors/codes.rs index 08223ec..5d62b01 100644 --- a/crates/loon-lang/src/errors/codes.rs +++ b/crates/loon-lang/src/errors/codes.rs @@ -42,6 +42,11 @@ pub enum ErrorCode { E0501, // private symbol E0502, // circular dependency + // Placement errors + E0600, // construct not allowed inside a kernel + E0601, // kernel signature is not placeable + E0602, // kernel writes outside its own work item + // Warnings W0100, // transparent wildcard } @@ -71,6 +76,9 @@ impl ErrorCode { ErrorCode::E0402 => "E0402", ErrorCode::E0403 => "E0403", ErrorCode::E0404 => "E0404", + ErrorCode::E0600 => "E0600", + ErrorCode::E0601 => "E0601", + ErrorCode::E0602 => "E0602", ErrorCode::E0500 => "E0500", ErrorCode::E0501 => "E0501", ErrorCode::E0502 => "E0502", @@ -102,6 +110,7 @@ impl ErrorCode { | ErrorCode::E0403 | ErrorCode::E0404 => "effect", ErrorCode::E0500 | ErrorCode::E0501 | ErrorCode::E0502 => "module", + ErrorCode::E0600 | ErrorCode::E0601 | ErrorCode::E0602 => "placement", ErrorCode::W0100 => "warning", } } @@ -130,6 +139,9 @@ impl ErrorCode { ErrorCode::E0402 => "unknown effect operation", ErrorCode::E0403 => "effect mismatch", ErrorCode::E0404 => "effect not granted", + ErrorCode::E0600 => "not allowed inside a kernel", + ErrorCode::E0601 => "kernel signature is not placeable", + ErrorCode::E0602 => "kernel writes outside its own work item", ErrorCode::E0500 => "unresolved module", ErrorCode::E0501 => "private symbol", ErrorCode::E0502 => "circular dependency", diff --git a/crates/loon-lang/src/errors/tutorials.rs b/crates/loon-lang/src/errors/tutorials.rs index 5ecc721..337775b 100644 --- a/crates/loon-lang/src/errors/tutorials.rs +++ b/crates/loon-lang/src/errors/tutorials.rs @@ -51,6 +51,7 @@ pub fn get_tutorial(code: ErrorCode) -> Option { ErrorCode::E0500 => Some(tutorial_unresolved_module()), ErrorCode::E0501 => Some(tutorial_private_symbol()), ErrorCode::E0502 => Some(tutorial_circular_dependency()), + ErrorCode::E0600 | ErrorCode::E0601 | ErrorCode::E0602 => None, ErrorCode::W0100 => None, // warnings don't need tutorials } } diff --git a/crates/loon-lang/tests/backend_parity.rs b/crates/loon-lang/tests/backend_parity.rs index 1cd2322..55e04ba 100644 --- a/crates/loon-lang/tests/backend_parity.rs +++ b/crates/loon-lang/tests/backend_parity.rs @@ -32,6 +32,22 @@ fn interp_output(src: &str) -> Result { /// result(s) so the comparison is on observable output. const CORPUS: &[(&str, &str)] = &[ ("arith", "[fn main [] [println [+ [* 6 7] [- 10 8]]]]"), + // Regression: a closure called by a higher-order builtin must not + // overwrite a binding in the caller. The VM re-enters itself to run the + // closure, and that re-entry used to name register 0 as the destination + // for the result — whatever the caller happened to be keeping there was + // silently replaced. Here `n` came back as 14, the last value `map` + // computed, instead of 8. + ( + "closure-call-preserves-caller-bindings", + "[fn main [] [let n 8] [let y [map [fn [v] [* 2 v]] [range 0 n]]] [println n] [println y]]", + ), + ( + "nested-closure-calls-preserve-bindings", + "[fn main [] [let a 1] [let b 2] \ + [let r [map [fn [v] [len [map [fn [w] [* w v]] [range 0 3]]]] [range 0 3]]] \ + [println [+ a b]] [println r]]", + ), ("float", "[fn main [] [println [* 2.0 3.5]]]"), ("strings", r#"[fn main [] [println [str "a" "-" 42 "-" "b"]]]"#), ("bool", "[fn main [] [println [and [> 3 2] [not [< 5 1]]]]]"), diff --git a/crates/loon-lang/tests/gpu_end_to_end.rs b/crates/loon-lang/tests/gpu_end_to_end.rs new file mode 100644 index 0000000..d17bdab --- /dev/null +++ b/crates/loon-lang/tests/gpu_end_to_end.rs @@ -0,0 +1,502 @@ +//! A Loon kernel, on a real GPU, agreeing with the CPU. +//! +//! Everything in the chain gets exercised: the `[kernel ...]` form, ownership +//! modes deciding which buffers are inputs, the WGSL emitter, and wgpu putting +//! the shader on whatever hardware this machine has. The result is compared +//! against the same kernel run through the interpreter, because a GPU that +//! computes something *different* quickly is not an optimization. +//! +//! Skipped, loudly, when there is no GPU. A machine without one is not a +//! failing machine, and a test that quietly passed in that case would be +//! worse than no test. + +#![cfg(feature = "gpu")] + +use loon_lang::check::Checker; +use loon_lang::eir::gpu::{self, Gpu, GpuArg}; +use loon_lang::eir::layout::DType; +use loon_lang::eir::lower::lower; +use loon_lang::eir::vm::{eval_eir, BufData, Buffer}; +use loon_lang::eir::wgsl; +use loon_lang::parser::parse; + +fn gpu_or_skip() -> Option { + match Gpu::open() { + Ok(g) => { + println!("device: {}", g.name()); + Some(g) + } + Err(e) => { + println!("SKIPPED — no GPU on this machine: {e}"); + None + } + } +} + +/// Compile the named kernel in `src` to WGSL and run it on the GPU. +fn run_on_gpu( + gpu_dev: &Gpu, + src: &str, + kernel: &str, + n: u32, + scalars: &[f32], + buffers: &[Buffer], +) -> Vec { + let exprs = parse(src).expect("parses"); + let mut checker = Checker::new(); + let errors = checker.check_program(&exprs); + assert!(errors.is_empty(), "check errors: {errors:?}"); + let module = lower(&checker); + let func = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(kernel)) + .expect("kernel lowered"); + + let kinds = wgsl::infer_arg_kinds(&module, func.id, DType::F32); + let shader = wgsl::emit(&module, func.id, &kinds).expect("emits WGSL"); + + // Narrow every buffer to what the device can hold, in argument order. + // (WGSL core has no 64-bit scalar, so an f64 buffer is computed in f32.) + let owned: Vec = buffers.iter().map(gpu::narrow).collect(); + + // Then describe each argument: scalars from `scalars`, buffers from + // `owned`, in the order the kernel declared them. + let mut scalar_iter = scalars.iter(); + let mut owned_iter = owned.iter(); + let args: Vec = kinds + .iter() + .map(|kind| match kind { + wgsl::ArgKind::Scalar(_) => { + GpuArg::Scalar(*scalar_iter.next().expect("a scalar argument")) + } + wgsl::ArgKind::Buffer { writable, .. } => GpuArg::Buffer { + data: owned_iter.next().expect("a buffer argument"), + writable: *writable, + }, + }) + .collect(); + + let results = gpu_dev.run(&shader, "main", n, &args).expect("dispatch"); + + let mut out = owned.clone(); + for (slot, bytes) in results { + // `slot` is an argument index; find which buffer that was. + let buffer_index = kinds + .iter() + .take(slot) + .filter(|k| matches!(k, wgsl::ArgKind::Buffer { .. })) + .count(); + let target = &mut out[buffer_index]; + let n = target.len(); + let mut vals: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes(c.try_into().unwrap())) + .collect(); + vals.truncate(n); + target.data = BufData::F32(vals); + } + out +} + +/// The same kernel through the interpreter, as the reference answer. +fn run_on_cpu(src: &str, main: &str) -> Vec { + eval_eir(&format!("{src} {main}")).expect("cpu run").output +} + +const SAXPY: &str = "[kernel saxpy [i a x y out] \ + [put out i [+ [* a [at x i]] [at y i]]]]"; + +#[test] +fn saxpy_agrees_between_the_gpu_and_the_cpu() { + let Some(dev) = gpu_or_skip() else { return }; + + let n = 256usize; + let x: Vec = (0..n).map(|i| i as f32).collect(); + let y: Vec = (0..n).map(|i| (i * 2) as f32).collect(); + let out = vec![0.0f32; n]; + + let results = run_on_gpu( + &dev, + SAXPY, + "saxpy", + n as u32, + &[3.0], + &[ + Buffer { + data: BufData::F32(x.clone()), + }, + Buffer { + data: BufData::F32(y.clone()), + }, + Buffer { + data: BufData::F32(out), + }, + ], + ); + + let BufData::F32(got) = &results[2].data else { + panic!("expected an f32 buffer"); + }; + let expected: Vec = (0..n).map(|i| 3.0 * x[i] + y[i]).collect(); + assert_eq!(got.len(), n); + for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() { + assert!( + (g - e).abs() < 1e-4, + "element {i}: GPU gave {g}, expected {e}" + ); + } +} + +#[test] +fn the_gpu_result_matches_what_the_interpreter_computes() { + // Not "close to a formula I wrote in Rust" — the same Loon program, run + // both ways. This is the comparison that would catch the emitter and the + // interpreter disagreeing about what a kernel means. + let Some(dev) = gpu_or_skip() else { return }; + + let cpu = run_on_cpu( + SAXPY, + "[fn main [] \ + [let x [buf #[0 1 2 3 4 5 6 7]]] \ + [let y [buf #[0 2 4 6 8 10 12 14]]] \ + [let mut out [buf-zeros 8]] \ + [Place.run saxpy 8 #[3.0 x y out]] \ + [IO.println [Place.read out]]]", + ); + + let results = run_on_gpu( + &dev, + SAXPY, + "saxpy", + 8, + &[3.0], + &[ + Buffer { + data: BufData::F32((0..8).map(|i| i as f32).collect()), + }, + Buffer { + data: BufData::F32((0..8).map(|i| (i * 2) as f32).collect()), + }, + Buffer { + data: BufData::F32(vec![0.0; 8]), + }, + ], + ); + let BufData::F32(got) = &results[2].data else { + panic!("expected an f32 buffer"); + }; + let rendered = format!( + "#[{}]", + got.iter() + .map(|v| { + if v.fract() == 0.0 { + format!("{}", *v as i64) + } else { + format!("{v}") + } + }) + .collect::>() + .join(" ") + ); + assert_eq!(cpu, vec![rendered]); +} + +#[test] +fn a_branching_kernel_runs_on_the_gpu() { + // Control flow goes through the block dispatcher in the emitted shader, + // so this exercises a different path than straight-line arithmetic. + let Some(dev) = gpu_or_skip() else { return }; + + let src = "[kernel clamp [i lo hi b] \ + [let v [at b i]] \ + [put b i [if [< v lo] lo [if [> v hi] hi v]]]]"; + let input: Vec = vec![-5.0, 0.25, 0.5, 9.0]; + let results = run_on_gpu( + &dev, + src, + "clamp", + 4, + &[0.0, 1.0], + &[Buffer { + data: BufData::F32(input), + }], + ); + let BufData::F32(got) = &results[0].data else { + panic!("expected an f32 buffer"); + }; + assert_eq!(got, &vec![0.0, 0.25, 0.5, 1.0]); +} + +#[test] +fn a_large_launch_covers_every_element() { + // Enough work items to span many workgroups, so a mistake in the dispatch + // arithmetic shows up as untouched elements at the end. + let Some(dev) = gpu_or_skip() else { return }; + + let n = 10_000usize; + let results = run_on_gpu( + &dev, + "[kernel fill [i b] [put b i 7.0]]", + "fill", + n as u32, + &[], + &[Buffer { + data: BufData::F32(vec![0.0; n]), + }], + ); + let BufData::F32(got) = &results[0].data else { + panic!("expected an f32 buffer"); + }; + assert_eq!(got.len(), n); + assert!( + got.iter().all(|v| *v == 7.0), + "every element should have been written; {} were not", + got.iter().filter(|v| **v != 7.0).count() + ); +} + +// ── Placement mode: the same program, a different device ──────────────────── + +/// Run a program under a placement mode, returning its output and accounting. +fn run_placed( + src: &str, + mode: loon_lang::eir::place::Mode, +) -> (Vec, loon_lang::eir::place::PlaceStats) { + let dir = std::env::current_dir().expect("cwd"); + let (result, stats) = + loon_lang::eir::vm::eval_eir_placed(src, &dir, mode).expect("program runs"); + (result.output, stats) +} + +const CHAIN: &str = "[kernel step [i b] [put b i [+ 1.0 [at b i]]]] \ + [fn work [] [let mut b [buf #[0 0 0 0]]] \ + [Place.run step 4 #[b]] [Place.run step 4 #[b]] \ + [Place.run step 4 #[b]] [Place.run step 4 #[b]] \ + [Place.read b]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [Place.pin args] [resume [Place.run k n args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [IO.println [work]] [IO.println [resident work]]]"; + +#[test] +fn the_gpu_is_selected_by_a_flag_and_changes_nothing_else() { + // This is the claim in one test. The program is a constant; the mode is a + // parameter. A GPU that produced a different answer would not be a faster + // way to run this program, it would be a different program. + if Gpu::open().is_err() { + println!("SKIPPED — no GPU on this machine"); + return; + } + let (cpu_out, _) = run_placed(CHAIN, loon_lang::eir::place::Mode::Cpu); + let (gpu_out, gpu_stats) = run_placed(CHAIN, loon_lang::eir::place::Mode::Gpu); + + assert_eq!(cpu_out, gpu_out, "the GPU must compute what the CPU does"); + assert_eq!(cpu_out, vec!["#[4 4 4 4]", "#[4 4 4 4]"]); + + // Eight launches: four with no policy, four under the residency handler. + assert_eq!(gpu_stats.launches, 8); +} + +#[test] +fn a_residency_handler_saves_real_transfers_on_real_hardware() { + // The policy was written against an effect, not against a device. Here it + // is deciding what a Metal GPU actually has to copy. + if Gpu::open().is_err() { + println!("SKIPPED — no GPU on this machine"); + return; + } + let naive = "[kernel step [i b] [put b i [+ 1.0 [at b i]]]] \ + [fn work [] [let mut b [buf #[0 0 0 0]]] \ + [Place.run step 4 #[b]] [Place.run step 4 #[b]] \ + [Place.run step 4 #[b]] [Place.run step 4 #[b]] \ + [Place.read b]] \ + [fn main [] [IO.println [work]]]"; + let managed = "[kernel step [i b] [put b i [+ 1.0 [at b i]]]] \ + [fn work [] [let mut b [buf #[0 0 0 0]]] \ + [Place.run step 4 #[b]] [Place.run step 4 #[b]] \ + [Place.run step 4 #[b]] [Place.run step 4 #[b]] \ + [Place.read b]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [Place.pin args] [resume [Place.run k n args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [IO.println [resident work]]]"; + + let (bare, bare_stats) = run_placed(naive, loon_lang::eir::place::Mode::Gpu); + let (wrapped, wrapped_stats) = run_placed(managed, loon_lang::eir::place::Mode::Gpu); + + assert_eq!(bare, wrapped, "same answer either way"); + assert_eq!(bare_stats.uploads, 4); + assert_eq!(wrapped_stats.uploads, 1); + assert_eq!(wrapped_stats.resident_hits, 3); +} + +#[test] +fn asking_for_a_gpu_that_cannot_run_the_kernel_says_so() { + // A kernel outside the GPU subset must be refused by name rather than + // quietly run somewhere else. Being told your program did not run where + // you asked is worth more than a result that arrived by another route. + if Gpu::open().is_err() { + println!("SKIPPED — no GPU on this machine"); + return; + } + let dir = std::env::current_dir().expect("cwd"); + // `s` is a scalar the launch passes, but the kernel indexes `b` — fine. + // Make it impossible instead: pass a buffer where the body wants a number. + let src = "[kernel bad [i a b] [put b i [* a [at b i]]]] \ + [fn main [] \ + [let x [buf #[1 2]]] [let mut o [buf-zeros 2]] \ + [Place.run bad 2 #[x o]]]"; + let result = loon_lang::eir::vm::eval_eir_placed(src, &dir, loon_lang::eir::place::Mode::Gpu); + match result { + Err(e) => { + let msg = format!("{e:?}"); + assert!( + msg.contains("buffer") || msg.contains("GPU"), + "the refusal should explain itself: {msg}" + ); + } + Ok((r, _)) => panic!("expected a refusal, got output {:?}", r.output), + } +} + +// ── Launches pipeline; only a read synchronizes ───────────────────────────── + +#[test] +fn launches_do_not_block_on_each_other() { + // A kernel launch submits to the queue and returns. Nothing waits for the + // GPU until the host asks for data — which, since `Place.read` is the only + // way to ask, means a handler that defers reads is also deferring every + // synchronization point in the program. + // + // The observable consequence: N launches under a residency policy cost far + // less than N times a single blocking round trip. If every launch waited, + // the two would be equal. + if Gpu::open().is_err() { + println!("SKIPPED — no GPU on this machine"); + return; + } + + let chain = |reps: usize| { + let runs = (0..reps) + .map(|_| "[Place.run step 4096 #[b]]".to_string()) + .collect::>() + .join(" "); + format!( + "[kernel step [i b] [put b i [+ 1.0 [at b i]]]] \ + [fn work [] [let mut b [buf-zeros 4096]] {runs} [Place.read b]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [Place.pin args] [resume [Place.run k n args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [let _ [resident work]] []]" + ) + }; + + let time = |src: &str| { + let dir = std::env::current_dir().expect("cwd"); + // Warm: the first run pays for adapter discovery and shader compilation. + let _ = loon_lang::eir::vm::eval_eir_placed(src, &dir, loon_lang::eir::place::Mode::Gpu); + let start = std::time::Instant::now(); + loon_lang::eir::vm::eval_eir_placed(src, &dir, loon_lang::eir::place::Mode::Gpu) + .expect("runs"); + start.elapsed() + }; + + let one = time(&chain(1)); + let many = time(&chain(64)); + + println!("1 launch: {one:?}, 64 launches: {many:?}"); + assert!( + many < one * 32, + "64 launches took {many:?} against {one:?} for one — that is the shape of \ + every launch waiting for the previous one" + ); +} + +#[test] +fn a_read_is_what_synchronizes() { + // Reading between every launch forfeits the pipelining, which is precisely + // why `Place.read` being an operation matters: it is the thing a policy + // gets to move. + if Gpu::open().is_err() { + println!("SKIPPED — no GPU on this machine"); + return; + } + let dir = std::env::current_dir().expect("cwd"); + + let deferred = "[kernel step [i b] [put b i [+ 1.0 [at b i]]]] \ + [fn work [] [let mut b [buf-zeros 1024]] \ + [Place.run step 1024 #[b]] [Place.run step 1024 #[b]] \ + [Place.run step 1024 #[b]] [Place.run step 1024 #[b]] \ + [Place.read b]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [Place.pin args] [resume [Place.run k n args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [IO.println [sum [resident work]]]]"; + + let eager = "[kernel step [i b] [put b i [+ 1.0 [at b i]]]] \ + [fn work [] [let mut b [buf-zeros 1024]] \ + [Place.run step 1024 #[b]] [let _ [Place.read b]] \ + [Place.run step 1024 #[b]] [let _ [Place.read b]] \ + [Place.run step 1024 #[b]] [let _ [Place.read b]] \ + [Place.run step 1024 #[b]] \ + [Place.read b]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [Place.pin args] [resume [Place.run k n args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [IO.println [sum [resident work]]]]"; + + let (a, stats_a) = + loon_lang::eir::vm::eval_eir_placed(deferred, &dir, loon_lang::eir::place::Mode::Gpu) + .expect("runs"); + let (b, stats_b) = + loon_lang::eir::vm::eval_eir_placed(eager, &dir, loon_lang::eir::place::Mode::Gpu) + .expect("runs"); + + // Same answer; the reads changed only when data moved, never what it was. + assert_eq!(a.output, b.output); + assert_eq!(stats_a.downloads, 1, "one read, one download"); + assert_eq!(stats_b.downloads, 4, "four reads, four downloads"); +} + +#[test] +fn a_sixty_four_bit_buffer_is_refused_rather_than_narrowed() { + // WGSL core has no 64-bit scalar. Quietly computing an f64 buffer in f32 + // would hand back numbers of a precision the program never asked for and + // has no way to notice, so the launch is refused and the alternative + // named. The same program runs fine on the CPU. + if Gpu::open().is_err() { + println!("SKIPPED — no GPU on this machine"); + return; + } + let dir = std::env::current_dir().expect("cwd"); + let src = "[kernel k [i b] [put b i [* 2.0 [at b i]]]] \ + [fn main [] [let mut b [buf-f64 #[1 2 3]]] \ + [Place.run k 3 #[b]] [IO.println [Place.read b]]]"; + + let on_gpu = loon_lang::eir::vm::eval_eir_placed(src, &dir, loon_lang::eir::place::Mode::Gpu); + match on_gpu { + Err(e) => { + let msg = format!("{e:?}"); + assert!( + msg.contains("f64"), + "the message should name the type: {msg}" + ); + assert!( + msg.contains("--place cpu"), + "and say what to do instead: {msg}" + ); + } + Ok((r, _)) => panic!("expected a refusal, got {:?}", r.output), + } + + let (on_cpu, _) = + loon_lang::eir::vm::eval_eir_placed(src, &dir, loon_lang::eir::place::Mode::Cpu) + .expect("the CPU has 64-bit numbers"); + assert_eq!(on_cpu.output, vec!["#[2 4 6]"]); +} diff --git a/crates/loon-lang/tests/param_modes.rs b/crates/loon-lang/tests/param_modes.rs new file mode 100644 index 0000000..1bb90d0 --- /dev/null +++ b/crates/loon-lang/tests/param_modes.rs @@ -0,0 +1,165 @@ +//! Parameter modes survive the trip from inference into the IR. +//! +//! Loon infers, for every function parameter, whether the callee reads it, +//! writes through it, or consumes it. Other languages make you write that down +//! — Rust spells it `&T` / `&mut T` / `T`, and a GPU offload compiler reads +//! those same sigils to decide which way bytes move across a device boundary. +//! Here the analysis already ran; these tests pin that its answers are correct +//! and that lowering carries them onto `Func::param_modes` instead of throwing +//! them away. + +use loon_lang::check::ownership::{infer_param_modes, ParamMode}; +use loon_lang::check::Checker; +use loon_lang::eir::lower::lower; +use loon_lang::eir::Mode; +use loon_lang::parser::parse; + +/// Infer modes for one named function. +fn modes_of(src: &str, func: &str) -> Vec { + let exprs = parse(src).expect("parses"); + infer_param_modes(&exprs) + .get(func) + .unwrap_or_else(|| panic!("no modes inferred for `{func}`")) + .clone() +} + +/// Lower a program and read the modes off the named function's IR. +fn lowered_modes(src: &str, func: &str) -> Vec { + let exprs = parse(src).expect("parses"); + let mut checker = Checker::new(); + let errors = checker.check_program(&exprs); + assert!(errors.is_empty(), "type errors: {errors:?}"); + let module = lower(&checker); + let f = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(func)) + .unwrap_or_else(|| panic!("no lowered function `{func}`")); + f.param_modes.clone() +} + +#[test] +fn a_read_only_parameter_is_a_borrow() { + // `n` is only added to, so the caller keeps it. + assert_eq!( + modes_of("[fn twice [n] [+ n n]]", "twice"), + vec![ParamMode::Borrow] + ); +} + +#[test] +fn a_mutated_parameter_is_a_mutable_borrow() { + // `push!` writes through its first argument; the rest are read. + assert_eq!( + modes_of("[fn add-to [v x] [push! v x]]", "add-to"), + vec![ParamMode::MutBorrow, ParamMode::Borrow] + ); +} + +#[test] +fn a_returned_parameter_is_moved() { + assert_eq!( + modes_of("[fn identity [x] x]", "identity"), + vec![ParamMode::Move] + ); +} + +#[test] +fn kernel_accessors_drive_the_direction() { + // This is the whole trick behind inferring transfer direction: `at` reads + // an element and `put` writes one, so a buffer that is only read comes out + // `Borrow` (host-to-device only) while one written through comes out + // `MutBorrow` (must be synchronized back). Nobody annotated anything. + assert_eq!( + modes_of("[fn saxpy [i a x out] [put out i [* a [at x i]]]]", "saxpy"), + vec![ + ParamMode::Borrow, // i — index, read + ParamMode::Borrow, // a — scalar, read + ParamMode::Borrow, // x — input buffer + ParamMode::MutBorrow, // out — output buffer + ] + ); +} + +#[test] +fn modes_reach_the_lowered_function() { + // The point of the exercise: what the frontend inferred is visible to the + // backend, rather than being dropped with the checker. + assert_eq!( + lowered_modes("[fn add-to [v x] [push! v x]] [fn main [] []]", "add-to"), + vec![Mode::InOut, Mode::In] + ); + // The shape a kernel launch cares about: one argument read, one written + // through. `nth` reads; `push!` writes. Once buffers land, `at` and `put` + // take these roles with the same rules. + assert_eq!( + lowered_modes( + "[fn blend [src dst] [push! dst [nth src 0]]] [fn main [] []]", + "blend" + ), + vec![Mode::In, Mode::InOut] + ); +} + +#[test] +fn arity_always_matches_the_parameter_list() { + // A mode vector that disagrees with the parameter count would silently + // misalign direction with argument, so lowering falls back to the + // conservative all-`Owned` answer rather than emitting a short vector. + for (src, func, arity) in [ + ("[fn none [] 1]", "none", 0), + ("[fn one [a] a]", "one", 1), + ("[fn three [a b c] [+ a [+ b c]]]", "three", 3), + ] { + let src = format!("{src} [fn main [] []]"); + assert_eq!( + lowered_modes(&src, func).len(), + arity, + "`{func}` should have {arity} modes" + ); + } +} + +#[test] +fn a_callee_that_mutates_propagates_to_its_caller() { + // `caller` hands `x` to something that writes through it, so `x` is + // mutably borrowed rather than consumed — in placement terms, an argument + // that has to be synchronized back but need not be surrendered. Reaching + // this answer requires resolving the callee, which is what the fixed point + // buys over a single source-order pass. + let modes = modes_of( + "[fn caller [x] [later x]] [fn later [y] [push! y 1]]", + "caller", + ); + assert_eq!(modes, vec![ParamMode::MutBorrow]); +} + +#[test] +fn conservative_when_the_callee_is_unknowable() { + // A callee that is not a definition in this program — here a parameter + // holding a function — cannot be analyzed, so the argument is assumed + // consumed. That is sound: it only costs an optimization, and it must + // never come out as a plain borrow, which would wrongly tell a backend + // nothing needs copying back. + let modes = modes_of("[fn apply-to [f x] [f x]]", "apply-to"); + assert_eq!(modes[1], ParamMode::Move); +} + +#[test] +fn modes_do_not_depend_on_definition_order() { + // Mode inference walks definitions in source order and consults callees it + // has already seen. Reordering two independent functions must not change + // either one's answer. + let forward = modes_of( + "[fn helper [v] [push! v 1]] [fn user [w] [helper w]]", + "user", + ); + let backward = modes_of( + "[fn user [w] [helper w]] [fn helper [v] [push! v 1]]", + "user", + ); + assert_eq!( + forward, backward, + "definition order changed the inferred mode: {forward:?} vs {backward:?}" + ); +} diff --git a/crates/loon-lang/tests/place_effect.rs b/crates/loon-lang/tests/place_effect.rs new file mode 100644 index 0000000..c5889a3 --- /dev/null +++ b/crates/loon-lang/tests/place_effect.rs @@ -0,0 +1,500 @@ +//! Placement is an effect, so a handler decides where a kernel runs. +//! +//! These tests pin the three claims that make that more than a slogan: +//! a program that never mentions placement still runs; a handler can observe +//! and redirect every launch without the program changing; and the direction +//! data has to travel is inferred from the kernel rather than annotated. + +use loon_lang::eir::vm::eval_eir; + +fn run(src: &str) -> Vec { + match eval_eir(src) { + Ok(r) => r.output, + Err(e) => panic!("VM error: {e:?}\nsource:\n{src}"), + } +} + +fn run_err(src: &str) -> String { + match eval_eir(src) { + Ok(r) => panic!("expected an error, got output {:?}", r.output), + Err(e) => format!("{e:?}"), + } +} + +const SAXPY: &str = "[kernel saxpy [i a x y out] \ + [put out i [+ [* a [at x i]] [at y i]]]]"; + +#[test] +fn a_kernel_runs_with_no_handler_at_all() { + // The default answer to "where does this run" is "here". A program that + // never mentions placement is still a working program. + let out = run(&format!( + "{SAXPY} \ + [fn main [] \ + [let x [buf #[1 2 3]]] \ + [let y [buf #[10 20 30]]] \ + [let mut out [buf-zeros 3]] \ + [Place.run saxpy 3 #[2.0 x y out]] \ + [IO.println [Place.read out]]]" + )); + assert_eq!(out, vec!["#[12 24 36]"]); +} + +#[test] +fn a_handler_sees_every_launch_and_every_sync_point() { + // This is the whole argument. The Rust offload work needs `Preload` and + // `PreloadMut` types, whose `drop` marks where device data becomes visible + // to the host, because nothing else in the language knows. Here every + // launch and every read is an operation, so a handler learns the same + // facts without the program being annotated at all — and a residency + // policy is therefore writable as ordinary code. + let out = run(&format!( + "{SAXPY} \ + [fn work [] \ + [let x [buf #[1 2 3]]] \ + [let y [buf #[10 20 30]]] \ + [let mut out [buf-zeros 3]] \ + [Place.run saxpy 3 #[2.0 x y out]] \ + [Place.read out]] \ + [fn traced [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [IO.println \"launch\"] \ + [resume [Place.run k n args]]] \ + [Place.read b] [do [IO.println \"sync\"] \ + [resume [Place.read b]]]]] \ + [fn main [] [IO.println [traced work]]]" + )); + assert_eq!(out, vec!["launch", "sync", "#[12 24 36]"]); +} + +#[test] +fn a_handler_can_answer_without_running_anything() { + // A handler is free not to forward. Answering from a recording — or from + // a simulated device — is the same shape as answering for real, which is + // why a program can be tested without the hardware it targets. + let out = run(&format!( + "{SAXPY} \ + [fn work [] \ + [let x [buf #[1 2 3]]] \ + [let mut out [buf-zeros 3]] \ + [Place.run saxpy 3 #[2.0 x x out]] \ + [Place.read out]] \ + [fn canned [thunk] \ + [handle [thunk] \ + [Place.run k n args] [resume []] \ + [Place.read b] [resume #[99 99 99]]]] \ + [fn main [] [IO.println [canned work]]]" + )); + assert_eq!(out, vec!["#[99 99 99]"]); +} + +#[test] +fn the_program_is_identical_under_every_handler() { + // `work` is one function. What changes between these runs is the line that + // wraps it, never the code that does the work. + let common = format!( + "{SAXPY} \ + [fn work [] \ + [let x [buf #[1 2 3]]] \ + [let y [buf #[10 20 30]]] \ + [let mut out [buf-zeros 3]] \ + [Place.run saxpy 3 #[2.0 x y out]] \ + [Place.read out]] \ + [fn counted [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [IO.println \"one launch\"] \ + [resume [Place.run k n args]]]]]" + ); + let bare = run(&format!("{common} [fn main [] [IO.println [work]]]")); + let wrapped = run(&format!( + "{common} [fn main [] [IO.println [counted work]]]" + )); + + assert_eq!(bare, vec!["#[12 24 36]"]); + assert_eq!(wrapped, vec!["one launch", "#[12 24 36]"]); +} + +#[test] +fn transfer_accounting_counts_what_crossed_the_boundary() { + // The number that matters in an offload program is how many times bytes + // moved. Making it observable is what lets a residency policy be checked + // rather than assumed. + let out = run(&format!( + "{SAXPY} \ + [fn main [] \ + [let x [buf #[1 2 3 4]]] \ + [let mut out [buf-zeros 4]] \ + [Place.run saxpy 4 #[1.0 x x out]] \ + [Place.run saxpy 4 #[1.0 x x out]] \ + [let _ [Place.read out]] \ + [let s [Place.stats]] \ + [IO.println [get s :launches]] \ + [IO.println [get s :work-items]] \ + [IO.println [get s :downloads]] \ + [IO.println [get s :bytes-out]]]" + )); + // Two launches of four items each, and exactly one read back: 4 f32 = 16 B. + assert_eq!(out, vec!["2", "8", "1", "16"]); +} + +#[test] +fn a_kernel_only_reading_a_buffer_infers_an_input() { + // Direction comes from the body. `x` and `y` are read with `at`, `out` is + // written with `put`; nobody wrote `&` or `&mut` anywhere. + use loon_lang::check::kernel; + use loon_lang::check::ownership::{infer_param_modes, ParamMode}; + let exprs = loon_lang::parser::parse(SAXPY).expect("parses"); + // Kernels are ordinary functions by the time anything analyzes them, which + // is why they get ownership modes at all. + let (desugared, names) = kernel::desugar(&exprs); + assert!(names.contains("saxpy")); + let modes = infer_param_modes(&desugared); + assert_eq!( + modes.get("saxpy").expect("saxpy has modes"), + &vec![ + ParamMode::Borrow, // i + ParamMode::Borrow, // a + ParamMode::Borrow, // x — input + ParamMode::Borrow, // y — input + ParamMode::MutBorrow, // out — must be synchronized back + ] + ); +} + +#[test] +fn reading_past_the_end_of_a_buffer_is_an_error() { + // A device that answers zero for an out-of-range read teaches you to trust + // an answer it invented. This one says so instead. + let err = run_err("[fn main [] [let b [buf #[1 2 3]]] [IO.println [at b 7]]]"); + assert!( + err.contains("outside a buffer"), + "expected an out-of-range error, got: {err}" + ); +} + +#[test] +fn placing_something_that_is_not_a_kernel_is_an_error() { + let err = run_err("[fn main [] [Place.run 42 3]]"); + assert!( + err.contains("expects a kernel"), + "expected a kernel-shape error, got: {err}" + ); +} + +#[test] +fn a_negative_work_count_is_an_error() { + let err = run_err(&format!( + "{SAXPY} \ + [fn main [] \ + [let x [buf #[1]]] [let mut o [buf-zeros 1]] \ + [Place.run saxpy -1 #[1.0 x x o]]]" + )); + assert!( + err.contains("work count"), + "expected a work-count error, got: {err}" + ); +} + +#[test] +fn buffers_round_trip_through_every_element_type() { + let out = run("[fn main [] \ + [IO.println [buf-dtype [buf #[1 2]]]] \ + [IO.println [buf-dtype [buf-i32 #[1 2]]]] \ + [IO.println [buf-dtype [buf-f64 #[1 2]]]] \ + [IO.println [buf->vec [buf-i32 #[1 2 3]]]] \ + [IO.println [buf-len [buf-zeros 5]]]]"); + assert_eq!(out, vec!["f32", "i32", "f64", "#[1 2 3]", "5"]); +} + +// ── Residency: the gap a handler closes ───────────────────────────────────── + +/// Run in device mode, where buffers live in a separate memory and transfers +/// are counted. +fn run_on_device(src: &str) -> loon_lang::eir::place::PlaceStats { + let dir = std::env::current_dir().expect("cwd"); + match loon_lang::eir::vm::eval_eir_placed(src, &dir, loon_lang::eir::place::Mode::Device) { + Ok((_, stats)) => stats, + Err(e) => panic!("VM error: {e:?}\nsource:\n{src}"), + } +} + +/// A chain of `n` launches over the same buffer, optionally wrapped in a +/// handler. The program text is identical either way. +fn chain_program(launches: usize, wrapper: Option<&str>) -> String { + let runs = (0..launches) + .map(|_| "[Place.run step 4 #[1.0 b]]".to_string()) + .collect::>() + .join(" "); + let call = match wrapper { + Some(w) => format!("[{w} work]"), + None => "[work]".to_string(), + }; + format!( + "[kernel step [i s b] [put b i [* s [at b i]]]] \ + [fn work [] [let mut b [buf #[1 2 3 4]]] {runs} [Place.read b]] \ + [fn resident [thunk] \ + [handle [thunk] \ + [Place.run k n args] [do [Place.pin args] [resume [Place.run k n args]]] \ + [Place.read b] [resume [Place.read b]]]] \ + [fn main [] [IO.println {call}]]" + ) +} + +#[test] +fn without_a_policy_every_launch_pays_for_its_own_transfer() { + // A device that has not been told a buffer will be wanted again does not + // keep it. This is the honest default, and it is the behaviour that makes + // the naive interface of an offload compiler slow. + for launches in [1usize, 4, 16] { + let stats = run_on_device(&chain_program(launches, None)); + assert_eq!( + stats.uploads, launches as u64, + "{launches} launches should cost {launches} uploads" + ); + assert_eq!(stats.resident_hits, 0); + } +} + +#[test] +fn a_residency_handler_pays_once_no_matter_how_long_the_chain() { + // The same program, wrapped in a handler that pins what each launch + // touches. One upload, and every launch after the first finds the buffer + // already there. Nothing in `work` changed, and no compiler pass ran. + for launches in [1usize, 4, 16] { + let stats = run_on_device(&chain_program(launches, Some("resident"))); + assert_eq!( + stats.uploads, 1, + "a chain of {launches} should upload exactly once" + ); + assert_eq!( + stats.resident_hits, + launches as u64 - 1, + "every launch after the first should be a residency hit" + ); + // The host still asks for its answer exactly once. + assert_eq!(stats.downloads, 1); + } +} + +#[test] +fn the_saving_grows_with_the_chain_and_the_answer_does_not_change() { + // The property worth stating: a policy changes what it costs to get the + // answer, never the answer. A "policy" that changed the result would be a + // bug wearing a nicer name. + let long_chain: usize = 32; + let naive = run_on_device(&chain_program(long_chain, None)); + let resident = run_on_device(&chain_program(long_chain, Some("resident"))); + + assert_eq!(naive.uploads, long_chain as u64); + assert_eq!(resident.uploads, 1); + assert!( + naive.bytes_in >= resident.bytes_in * 30, + "expected roughly a {long_chain}x reduction in bytes moved, got {} vs {}", + naive.bytes_in, + resident.bytes_in + ); + + // Both programs computed the same thing. + let bare = run(&chain_program(long_chain, None)); + let wrapped = run(&chain_program(long_chain, Some("resident"))); + assert_eq!(bare, wrapped); +} + +#[test] +fn on_the_cpu_there_is_nothing_to_transfer() { + // One memory means no uploads at all, whatever the policy says. The + // handler is not wrong here, it is simply describing a distinction the + // hardware does not have. + let src = chain_program(8, Some("resident")); + let dir = std::env::current_dir().expect("cwd"); + let (_, stats) = + loon_lang::eir::vm::eval_eir_placed(&src, &dir, loon_lang::eir::place::Mode::Cpu) + .expect("runs"); + assert_eq!(stats.uploads, 0); + assert_eq!(stats.launches, 8); +} + +// ── Every placement gives the same answer ─────────────────────────────────── + +#[test] +fn cpu_and_parallel_agree_on_everything() { + // The property that makes placement a policy rather than a rewrite: the + // answer does not depend on where the work happened. Parallel execution + // hands each thread a disjoint slice of the output, so this is also the + // test that the split is right. + use loon_lang::eir::place::Mode; + let dir = std::env::current_dir().expect("cwd"); + + let programs = [ + "[kernel saxpy [i a x y out] [put out i [+ [* a [at x i]] [at y i]]]] \ + [fn main [] \ + [let x [buf [range 0 1000]]] [let y [buf [range 0 1000]]] \ + [let mut out [buf-zeros 1000]] \ + [Place.run saxpy 1000 #[2.0 x y out]] \ + [IO.println [sum [Place.read out]]]]", + "[kernel clamp [i lo hi b] [let v [at b i]] \ + [put b i [if [< v lo] lo [if [> v hi] hi v]]]] \ + [fn main [] \ + [let mut b [buf [range 0 500]]] \ + [Place.run clamp 500 #[10.0 100.0 b]] \ + [IO.println [sum [Place.read b]]]]", + "[kernel mathy [i b] [put b i [sqrt [abs [at b i]]]]] \ + [fn main [] \ + [let mut b [buf [range 0 777]]] \ + [Place.run mathy 777 #[b]] \ + [IO.println [len [Place.read b]]]]", + ]; + + for src in programs { + let (cpu, _) = loon_lang::eir::vm::eval_eir_placed(src, &dir, Mode::Cpu).expect("cpu"); + let (par, _) = loon_lang::eir::vm::eval_eir_placed(src, &dir, Mode::Par).expect("par"); + assert_eq!( + cpu.output, par.output, + "parallel placement changed the answer for:\n{src}" + ); + assert!(!cpu.output.is_empty()); + } +} + +#[test] +fn a_kernel_outside_the_fast_subset_still_runs() { + // The typed executor covers the numeric subset. Anything else falls back + // to the general VM rather than failing, so adding the fast path cannot + // have narrowed what a kernel is allowed to be. + use loon_lang::eir::place::Mode; + let dir = std::env::current_dir().expect("cwd"); + let src = "[kernel k [i b] [put b i [+ [at b i] 1.0]]] \ + [fn main [] [let mut b [buf-zeros 4]] \ + [Place.run k 4 #[b]] [IO.println [Place.read b]]]"; + for mode in [Mode::Cpu, Mode::Par] { + let (r, _) = loon_lang::eir::vm::eval_eir_placed(src, &dir, mode).expect("runs"); + assert_eq!(r.output, vec!["#[1 1 1 1]"]); + } +} + +// ── Reductions, within the rules ──────────────────────────────────────────── + +/// A reduction expressed as chunked partials plus a host-side combine. +const REDUCE: &str = "[kernel sum-chunk [c width src partials] \ + [let start [* c width]] \ + [let total [loop [k 0 acc 0.0] \ + [if [>= k width] acc [recur [+ k 1] [+ acc [at src [+ start k]]]]]]] \ + [put partials c total]] \ + [fn main [] \ + [let src [buf [range 0 256]]] \ + [let mut partials [buf-zeros 16]] \ + [Place.run sum-chunk 16 #[16 src partials]] \ + [IO.println [sum [Place.read partials]]]]"; + +#[test] +fn a_reduction_needs_no_new_language_feature() { + // Each work item sums its own chunk and writes its own partial, so the + // disjointness rule holds throughout. There is no `Place.reduce`, no + // workgroup-shared memory, and no exception carved out. + let out = run(REDUCE); + // 0 + 1 + ... + 255 + assert_eq!(out, vec!["32640"]); +} + +#[test] +fn a_reduction_agrees_across_every_placement() { + use loon_lang::eir::place::Mode; + let dir = std::env::current_dir().expect("cwd"); + for mode in [Mode::Cpu, Mode::Par, Mode::Device] { + let (r, _) = loon_lang::eir::vm::eval_eir_placed(REDUCE, &dir, mode) + .unwrap_or_else(|e| panic!("{mode:?} failed: {e:?}")); + assert_eq!(r.output, vec!["32640"], "{mode:?} disagreed"); + } +} + +#[test] +fn a_loop_carried_float_survives_the_shader() { + // The reduction's accumulator starts at 0.0 and is carried around a loop. + // Emitting that to WGSL exercised three separate bugs: a type-inference + // fixpoint that stopped before a loop-carried type settled, a branch target + // resolved across the whole module rather than one function, and a unit + // value reaching a float slot with no conversion. Each produced a shader + // that failed validation, which is the good outcome — but only because the + // emitter is checked. + use loon_lang::check::Checker; + use loon_lang::eir::layout::DType; + use loon_lang::eir::lower::lower; + use loon_lang::eir::wgsl; + + let exprs = loon_lang::parser::parse(REDUCE).expect("parses"); + let mut checker = Checker::new(); + assert!(checker.check_program(&exprs).is_empty()); + let module = lower(&checker); + let f = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some("sum-chunk")) + .expect("kernel lowered"); + let kinds = wgsl::infer_arg_kinds(&module, f.id, DType::F32); + let shader = wgsl::emit(&module, f.id, &kinds).expect("emits"); + + let parsed = naga::front::wgsl::parse_str(&shader) + .unwrap_or_else(|e| panic!("WGSL did not parse: {e}\n\n{shader}")); + let mut validator = naga::valid::Validator::new( + naga::valid::ValidationFlags::all(), + naga::valid::Capabilities::empty(), + ); + validator + .validate(&parsed) + .unwrap_or_else(|e| panic!("WGSL did not validate: {e:?}\n\n{shader}")); +} + +// ── Parallel placement with more than one output ──────────────────────────── + +/// A kernel writing two buffers, at a size large enough to span every core. +const SPLIT: &str = "[kernel split [i src lo hi] \ + [let v [at src i]] \ + [put lo i [if [< v 0.0] v 0.0]] \ + [put hi i [if [< v 0.0] 0.0 v]]] \ + [fn main [] \ + [let src [buf [map [fn [k] [- k 500]] [range 0 1000]]]] \ + [let mut lo [buf-zeros 1000]] \ + [let mut hi [buf-zeros 1000]] \ + [Place.run split 1000 #[src lo hi]] \ + [IO.println [sum [Place.read lo]]] \ + [IO.println [sum [Place.read hi]]]]"; + +#[test] +fn parallel_placement_splits_every_output_not_just_one() { + // This used to fall back to sequential execution whenever a kernel had + // anything other than exactly one output buffer — silently, so asking for + // `par` and getting `cpu` was indistinguishable from getting what you + // asked for. Every output is now carved into per-thread slices. + use loon_lang::eir::place::Mode; + let dir = std::env::current_dir().expect("cwd"); + + let mut answers = Vec::new(); + for mode in [Mode::Cpu, Mode::Par, Mode::Device] { + let (r, _) = loon_lang::eir::vm::eval_eir_placed(SPLIT, &dir, mode) + .unwrap_or_else(|e| panic!("{mode:?} failed: {e:?}")); + answers.push(r.output); + } + // -500..-1 sums to -125250; 0..499 sums to 124750. + assert_eq!(answers[0], vec!["-125250", "124750"]); + assert!( + answers.iter().all(|a| *a == answers[0]), + "placements disagreed: {answers:?}" + ); +} + +#[test] +fn a_mismatched_output_length_is_reported() { + // A buffer shorter than the launch cannot be carved into the ranges the + // work items need, and saying so beats writing past the end of one piece + // or quietly covering fewer elements. + use loon_lang::eir::place::Mode; + let dir = std::env::current_dir().expect("cwd"); + let src = "[kernel k [i b] [put b i 1.0]] \ + [fn main [] [let mut b [buf-zeros 4]] [Place.run k 64 #[b]]]"; + let err = loon_lang::eir::vm::eval_eir_placed(src, &dir, Mode::Par).expect_err("should refuse"); + let msg = format!("{err:?}"); + assert!( + msg.contains("64") || msg.contains("outside"), + "expected a length complaint, got: {msg}" + ); +} diff --git a/crates/loon-lang/tests/place_replay.rs b/crates/loon-lang/tests/place_replay.rs new file mode 100644 index 0000000..143dd7e --- /dev/null +++ b/crates/loon-lang/tests/place_replay.rs @@ -0,0 +1,137 @@ +//! A placed program can be recorded once and replayed without the device. +//! +//! This is the half of the story an offload compiler does not tell. A kernel +//! that ran on a GPU is, from the program's point of view, a launch that +//! returned nothing and a read that produced some numbers. Recording those and +//! feeding them back reproduces the run exactly — on a machine with no GPU, in +//! CI, or after the hardware it was written for stopped existing. + +use loon_lang::eir::place::Mode; +use loon_lang::eir::replay::{parse_trace, TraceRecorder}; +use loon_lang::eir::vm::{eval_eir_placed, eval_eir_recorded, eval_eir_replayed}; + +const PROGRAM: &str = "[kernel scale [i s b] [put b i [* s [at b i]]]] \ + [fn main [] \ + [let mut b [buf #[1 2 3 4]]] \ + [Place.run scale 4 #[3.0 b]] \ + [Place.run scale 4 #[2.0 b]] \ + [IO.println [Place.read b]]]"; + +fn temp_trace(name: &str) -> std::path::PathBuf { + let mut p = std::env::temp_dir(); + p.push(format!( + "loon-place-replay-{name}-{}.oo", + std::process::id() + )); + let _ = std::fs::remove_file(&p); + p +} + +fn record(src: &str, path: &std::path::Path) -> Vec { + let dir = std::env::current_dir().expect("cwd"); + let recorder = TraceRecorder::create(path).expect("create trace"); + let result = eval_eir_recorded(src, &dir, recorder).expect("records"); + loon_lang::eir::replay::finalize_trace_file(path).expect("finalize"); + result.output +} + +fn replay(src: &str, path: &std::path::Path) -> Vec { + let dir = std::env::current_dir().expect("cwd"); + let text = std::fs::read_to_string(path).expect("read trace"); + let entries = parse_trace(&text).expect("parse trace"); + eval_eir_replayed(src, &dir, entries) + .expect("replays") + .0 + .output +} + +#[test] +fn a_placed_run_replays_to_the_same_output() { + let path = temp_trace("basic"); + let recorded = record(PROGRAM, &path); + assert_eq!(recorded, vec!["#[6 12 18 24]"]); + + let replayed = replay(PROGRAM, &path); + assert_eq!( + recorded, replayed, + "a replayed run must observe what the recorded run observed" + ); + let _ = std::fs::remove_file(&path); +} + +#[test] +fn the_trace_records_launches_and_reads() { + let path = temp_trace("entries"); + record(PROGRAM, &path); + let text = std::fs::read_to_string(&path).expect("read trace"); + + let runs = text.matches(":op \"run\"").count(); + let reads = text.matches(":op \"read\"").count(); + assert_eq!(runs, 2, "both launches should be in the trace:\n{text}"); + assert_eq!(reads, 1, "the one read should be in the trace:\n{text}"); + let _ = std::fs::remove_file(&path); +} + +#[test] +fn statistics_are_never_replayed() { + // `Place.stats` reports on the run currently happening. A replayed run + // really did move no bytes, so feeding back the original counts would be + // a recording that lies about the execution it is part of. + let src = "[kernel k [i b] [put b i 1.0]] \ + [fn main [] \ + [let mut b [buf-zeros 4]] \ + [Place.run k 4 #[b]] \ + [let _ [Place.read b]] \ + [IO.println [get [Place.stats] :launches]]]"; + let path = temp_trace("stats"); + + let recorded = record(src, &path); + assert_eq!(recorded, vec!["1"], "the recorded run launched once"); + + let text = std::fs::read_to_string(&path).expect("read trace"); + assert!( + !text.contains(":op \"stats\""), + "stats should not be recorded:\n{text}" + ); + + // The replayed run performed no launch of its own, and says so. + let replayed = replay(src, &path); + assert_eq!(replayed, vec!["0"]); + let _ = std::fs::remove_file(&path); +} + +#[test] +fn a_recording_made_on_one_device_replays_anywhere() { + // The recording is of what the *program* observed, not of how the device + // behaved, so where it was made does not constrain where it is replayed. + let path = temp_trace("device"); + let dir = std::env::current_dir().expect("cwd"); + + let recorder = TraceRecorder::create(&path).expect("create trace"); + // Record against the simulated discrete device, which really does move + // bytes and evict buffers. + let recorded = { + let _ = eval_eir_placed(PROGRAM, &dir, Mode::Device); + eval_eir_recorded(PROGRAM, &dir, recorder) + .expect("records") + .output + }; + loon_lang::eir::replay::finalize_trace_file(&path).expect("finalize"); + + let replayed = replay(PROGRAM, &path); + assert_eq!(recorded, replayed); + let _ = std::fs::remove_file(&path); +} + +#[test] +fn a_buffer_prints_as_its_shape_not_its_address() { + // Buffers are handles to bulk data that may live on another device. + // Printing a heap slot would put an allocation number in output that is + // supposed to be reproducible — and would make traces differ run to run. + let out = loon_lang::eir::vm::eval_eir( + "[fn main [] [IO.println [buf #[1 2 3]]] [IO.println [buf-i32 #[1]]]]", + ) + .expect("runs") + .output; + assert_eq!(out, vec!["#buf", "#buf"]); +} diff --git a/crates/loon-lang/tests/place_samples.rs b/crates/loon-lang/tests/place_samples.rs new file mode 100644 index 0000000..711115a --- /dev/null +++ b/crates/loon-lang/tests/place_samples.rs @@ -0,0 +1,168 @@ +//! The placement samples, run end to end on the EIR VM. +//! +//! `samples/place/lib.oo` is ordinary Loon — no compiler support, no +//! privileges — so these run it the way a user would and check the claims the +//! demos exist to make: that every handler computes the same answer, that a +//! residency policy really does eliminate transfers, and that a parked +//! computation can be finished from outside. + +use loon_lang::eir::vm::eval_eir_with_base_dir; +use std::path::{Path, PathBuf}; + +fn place_dir() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("..") + .join("samples") + .join("place") +} + +/// Run a Loon source string with `[use ...]` resolved against samples/place/. +fn run(src: &str) -> Vec { + eval_eir_with_base_dir(src, &place_dir()) + .unwrap_or_else(|e| panic!("vm error: {e}")) + .output +} + +/// Run a sample file and return its printed lines. +fn run_demo(name: &str) -> Vec { + let path = place_dir().join(name); + let src = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {name}: {e}")); + run(&src) +} + +#[test] +fn place_demo_gives_the_same_answer_under_every_handler() { + // The claim the demo exists to make: the program is untouched between + // runs, only the handler wrapped around it changes, and the computed + // result is identical in each case. A placement policy that changed the + // answer would not be a policy, it would be a bug. + let out = run_demo("demo-handlers.oo"); + let joined = out.join("\n"); + + let expected = "#[11 21 31 41 51 61 71 81]"; + let answers = out.iter().filter(|l| l.contains(expected)).count(); + assert_eq!( + answers, 3, + "unhandled, traced, and resident runs should all agree:\n{joined}" + ); + + // Tracing is interposition: the handler sees each launch and the single + // point where the host asks for data back. + assert!( + joined.contains("place: launch 8 work items, 3 args"), + "{joined}" + ); + assert!(joined.contains("place: sync"), "{joined}"); + + // The dry run accounts for the work without performing any of it, which is + // how a program can be exercised without the hardware it targets. + assert!( + joined.contains("#[]"), + "dry run should return nothing:\n{joined}" + ); +} + +#[test] +fn place_handlers_do_not_need_compiler_support() { + // Every handler in samples/place/lib.oo is ordinary Loon. This test pins that a + // policy can be written inline, in a test file, with no privileges: the + // residency and transfer-hoisting logic a GPU compiler implements as an + // optimization pass is expressible as user code here. + let out = run("[kernel k [i b] [put b i [* 2.0 [at b i]]]] \ + [fn work [] [let mut b [buf #[1 2 3]]] [Place.run k 3 #[b]] [Place.read b]] \ + [fn twice-as-many [thunk] \ + [handle [thunk] \ + [Place.run kf n args] [do [Place.run kf n args] \ + [resume [Place.run kf n args]]]]] \ + [fn main [] [IO.println [work]] [IO.println [twice-as-many work]]]"); + // Running the kernel twice per launch really does double it again. + assert_eq!(out, vec!["#[2 4 6]", "#[4 8 12]"]); +} + +#[test] +fn residency_demo_closes_the_transfer_gap() { + // The demo's whole claim, checked: the same eight-launch chain pays eight + // uploads with no policy and one upload under a residency handler, and + // both runs produce the same answer. This is the gap a GPU offload + // compiler needs a dedicated optimization pass to close. + let path = place_dir().join("demo-residency.oo"); + let src = std::fs::read_to_string(&path).expect("read demo-residency.oo"); + let (result, _stats) = loon_lang::eir::vm::eval_eir_placed( + &src, + &place_dir(), + loon_lang::eir::place::Mode::Device, + ) + .expect("demo runs"); + let joined = result.output.join("\n"); + + assert!( + joined.contains("no policy: uploads 8, resident hits 0"), + "without a policy every launch transfers:\n{joined}" + ); + assert!( + joined.contains("place/resident: uploads 1, resident hits 7"), + "the handler should upload once and hit seven times:\n{joined}" + ); + assert_eq!( + result + .output + .iter() + .filter(|l| l.contains("#[8 8 8 8]")) + .count(), + 2, + "both runs must compute the same answer:\n{joined}" + ); +} + +#[test] +fn on_the_cpu_the_residency_demo_moves_nothing() { + // Same program, same handler, one memory: the policy is simply describing + // a distinction this hardware does not have, and costs nothing to keep. + let path = place_dir().join("demo-residency.oo"); + let src = std::fs::read_to_string(&path).expect("read demo-residency.oo"); + let (_, stats) = + loon_lang::eir::vm::eval_eir_placed(&src, &place_dir(), loon_lang::eir::place::Mode::Cpu) + .expect("demo runs"); + assert_eq!(stats.uploads, 0, "the CPU has nothing to upload to"); + assert_eq!(stats.launches, 16); +} + +#[test] +fn a_computation_can_be_parked_and_finished_later() { + // The mechanism an asynchronous host needs, and it is already here: a + // handler clause that hands `resume` outward and returns unwinds the + // computation without ending it. Whoever holds the continuation decides + // when — and whether — the rest of it runs. + // + // This is what makes "the browser cannot answer a GPU read immediately" + // a solvable problem rather than a blocking one. No VM support was added + // for it; reified escaping continuations are what handlers already are. + let out = run_demo("demo-park.oo"); + assert_eq!( + out, + vec![ + " work: starting", + " host: computation parked; the rest of it is mine now", + " host: ...doing something slow...", + " work: continued with 21", + " host: finished with 42", + "done", + ], + "the work should stop, unwind, and then continue where it left off" + ); +} + +#[test] +fn a_parked_continuation_survives_its_handler_returning() { + // The ordering is the claim: "work: continued" appears *after* the host + // has already printed, which is only possible if the computation really + // unwound and was restarted from the outside. + let out = run_demo("demo-park.oo").join("\n"); + let parked = out.find("computation parked").expect("parked"); + let continued = out.find("work: continued").expect("continued"); + assert!( + parked < continued, + "the continuation must resume after the handler returned:\n{out}" + ); +} diff --git a/crates/loon-lang/tests/session_park.rs b/crates/loon-lang/tests/session_park.rs new file mode 100644 index 0000000..ba5e22c --- /dev/null +++ b/crates/loon-lang/tests/session_park.rs @@ -0,0 +1,151 @@ +//! A program that stops, and a host that finishes it later. +//! +//! This is the shape an asynchronous embedder needs. The browser cannot answer +//! `Place.read` immediately — reading a GPU buffer back is a promise — but it +//! does not have to answer immediately. It can take the continuation, go away, +//! and come back when the bytes arrive. +//! +//! The VM is not asynchronous and does not become asynchronous. It just has to +//! still be there when the answer shows up, which is what `Session` is for: a +//! parked continuation is a heap object, so dropping the VM would drop the rest +//! of the program with it. + +use loon_lang::eir::place::Mode; +use loon_lang::eir::vm::{Session, Step}; + +fn dir() -> std::path::PathBuf { + std::env::current_dir().expect("cwd") +} + +/// A program that suspends once, in the middle of an expression. +const PARKING: &str = "[effect Slow [fetch [Int] Int]] \ + [effect Host [park [a] Unit]] \ + [fn work [] \ + [IO.println \"before\"] \ + [let a [Slow.fetch 1]] \ + [IO.println [str \"after \" a]] \ + [* a 2]] \ + [fn suspending [thunk] \ + [handle [thunk] [Slow.fetch id] [do [Host.park resume] 0]]] \ + [fn main [] [IO.println [str \"result \" [suspending work]]]]"; + +#[test] +fn a_program_can_park_and_be_finished_by_its_host() { + let mut session = Session::new(PARKING, &dir(), Mode::Cpu).expect("prepares"); + + // It runs until it needs something the host has to supply. + let step = session.start().expect("starts"); + assert!( + matches!(step, Step::Parked { .. }), + "the program should have parked" + ); + // "before" ran; so did the code after the handle, with the placeholder — + // see `code_after_the_parking_handler_runs_before_the_answer_arrives`. + assert_eq!(session.take_output(), vec!["before", "result 0"]); + + let k = session.pending().expect("a parked continuation"); + + // The host does whatever it could not do synchronously, then finishes. + let answer = loon_lang::eir::value64::Val::int(21); + let step = session.resume(k, answer).expect("resumes"); + match step { + Step::Done(_) => {} + Step::Parked { .. } => panic!("it should have finished this time"), + } + + // The suspended part picked up mid-expression with the supplied value, and + // its result comes back through `resume`. + let out = session.take_output(); + assert!( + out.iter().any(|l| l == "after 21"), + "the continuation should have carried the supplied value: {out:?}" + ); + assert_eq!(session.show(session.value()), "42"); +} + +#[test] +fn a_program_that_never_parks_just_finishes() { + // The same API for the ordinary case, so a host does not need two paths. + let mut session = Session::new( + "[fn main [] [IO.println \"straight through\"] 7]", + &dir(), + Mode::Cpu, + ) + .expect("prepares"); + + match session.start().expect("runs") { + Step::Done(_) => {} + Step::Parked { .. } => panic!("nothing here parks"), + } + assert_eq!(session.take_output(), vec!["straight through"]); + assert!(session.pending().is_none()); +} + +#[test] +fn a_parked_read_can_be_answered_with_data_the_host_fetched() { + // The placement case the whole thing is for: `Place.read` parks, the host + // supplies the numbers, and the rest of the computation runs with them. + // + // Note where the parking handler sits — outermost, with nothing after it. + // That is not decoration. Parking unwinds to the `handle`, so anything + // written after it would run immediately, with the placeholder, before the + // real answer existed. The value of a deferred computation comes back from + // `resume`, not from the call that parked. + let src = "[effect Host [park [a] Unit]] \ + [kernel k [i b] [put b i [* 2.0 [at b i]]]] \ + [fn work [] \ + [let mut b [buf #[1 2 3]]] \ + [Place.run k 3 #[b]] \ + [let got [Place.read b]] \ + [IO.println [str \"read \" got]] \ + [sum got]] \ + [fn deferred [thunk] \ + [handle [thunk] [Place.read b] [do [Host.park resume] #[]]]] \ + [fn main [] [deferred work]]"; + + let mut session = Session::new(src, &dir(), Mode::Cpu).expect("prepares"); + let Step::Parked { .. } = session.start().expect("starts") else { + panic!("the read should have parked"); + }; + let k = session.pending().expect("a continuation"); + + // Pretend these came back from a device. + let data = session.vec_of_floats(&[2.0, 4.0, 6.0]); + let step = session.resume(k, data).expect("resumes"); + assert!(matches!(step, Step::Done(_))); + + let out = session.take_output(); + assert!( + out.iter().any(|l| l == "read #[2 4 6]"), + "the read should return what the host supplied: {out:?}" + ); + // The rest of the suspended computation returns its value through `resume`. + assert_eq!(session.show(session.value()), "12"); +} + +#[test] +fn code_after_the_parking_handler_runs_before_the_answer_arrives() { + // Worth pinning, because it is the one surprising thing about parking and + // it decides where a host must put the handler. Unwinding means the caller + // of `handle` carries on immediately with whatever the clause returned; the + // suspended part is what waits. + let src = "[effect Slow [fetch [Int] Int]] \ + [effect Host [park [a] Unit]] \ + [fn work [] [let a [Slow.fetch 1]] [IO.println \"inner: resumed\"] a] \ + [fn suspending [thunk] \ + [handle [thunk] [Slow.fetch id] [do [Host.park resume] 0]]] \ + [fn main [] [let r [suspending work]] [IO.println \"outer: carried on\"] r]"; + + let mut session = Session::new(src, &dir(), Mode::Cpu).expect("prepares"); + let Step::Parked { .. } = session.start().expect("starts") else { + panic!("should park"); + }; + // The outer program already finished, using the placeholder. + assert_eq!(session.take_output(), vec!["outer: carried on"]); + + let k = session.pending().expect("a continuation"); + session + .resume(k, loon_lang::eir::value64::Val::int(5)) + .expect("resumes"); + assert_eq!(session.take_output(), vec!["inner: resumed"]); +} diff --git a/crates/loon-lang/tests/wgsl_samples.rs b/crates/loon-lang/tests/wgsl_samples.rs new file mode 100644 index 0000000..40bc1d7 --- /dev/null +++ b/crates/loon-lang/tests/wgsl_samples.rs @@ -0,0 +1,134 @@ +//! Every kernel that ships in the repo must compile to a valid GPU shader. +//! +//! The WGSL is parsed and type-checked by naga — the same front end wgpu uses +//! — so a kernel that a driver would reject fails here instead, on a machine +//! with no GPU at all. This is the automated cross-target check the Rust +//! offload paper reports as still missing: they found their slice-lowering +//! divergence between host and device by hand. + +use loon_lang::check::{kernel, Checker}; +use loon_lang::eir::layout::DType; +use loon_lang::eir::lower::lower; +use loon_lang::eir::wgsl::{self, ArgKind}; +use loon_lang::parser::parse; + +fn repo_root() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .and_then(|p| p.parent()) + .expect("repo root") + .to_path_buf() +} + +/// Parse a file, find its kernels, and emit WGSL for each one. +/// +/// Argument kinds are guessed from the kernel's inferred ownership modes: a +/// parameter written through is a writable buffer, one only read is an input. +/// That is the same fact the runtime uses to decide which way bytes move — +/// here it decides which bindings are `read` and which are `read_write`. +fn emit_all(src: &str, path: &str) -> usize { + let exprs = parse(src).unwrap_or_else(|e| panic!("{path} does not parse: {e:?}")); + let (_, kernel_names) = kernel::desugar(&exprs); + if kernel_names.is_empty() { + return 0; + } + + let mut checker = Checker::new(); + let errors = checker.check_program(&exprs); + assert!(errors.is_empty(), "{path} has check errors: {errors:?}"); + let module = lower(&checker); + + let mut emitted = 0; + for name in &kernel_names { + let func = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(name.as_str())) + .unwrap_or_else(|| panic!("{path}: kernel '{name}' did not lower")); + + // The kernel body says what its arguments are: indexed parameters are + // buffers, the rest are scalars. Nothing has to be declared. + let args = wgsl::infer_arg_kinds(&module, func.id, DType::F32); + let wgsl_text = wgsl::emit(&module, func.id, &args) + .unwrap_or_else(|e| panic!("{path}: kernel '{name}' will not emit: {e}")); + + validate(&wgsl_text, &format!("{path}:{name}")); + emitted += 1; + } + emitted +} + +fn validate(text: &str, what: &str) { + let module = naga::front::wgsl::parse_str(text) + .unwrap_or_else(|e| panic!("{what}: WGSL did not parse: {e}\n\n{text}")); + let mut validator = naga::valid::Validator::new( + naga::valid::ValidationFlags::all(), + naga::valid::Capabilities::empty(), + ); + validator + .validate(&module) + .unwrap_or_else(|e| panic!("{what}: WGSL did not validate: {e:?}\n\n{text}")); +} + +#[test] +fn every_kernel_in_the_repo_compiles_to_valid_wgsl() { + let root = repo_root(); + let mut checked = 0; + let mut files = 0; + + for dir in ["samples/place", "os"] { + let path = root.join(dir); + let Ok(entries) = std::fs::read_dir(&path) else { + continue; + }; + for entry in entries.flatten() { + let p = entry.path(); + if p.extension().and_then(|e| e.to_str()) != Some("oo") { + continue; + } + let Ok(src) = std::fs::read_to_string(&p) else { + continue; + }; + // Demos that pull in a module need its definitions to type-check; + // the kernels themselves are what this test is about, so only + // files that stand alone are compiled here. + if src.contains("[use ") { + continue; + } + let n = emit_all(&src, &p.display().to_string()); + if n > 0 { + files += 1; + checked += n; + } + } + } + + assert!( + checked > 0, + "no kernels were found to validate — the test would pass vacuously" + ); + println!("validated {checked} kernels across {files} files"); +} + +#[test] +fn a_kernel_the_gpu_cannot_run_is_refused_rather_than_mistranslated() { + // The kernel checker already rejects effects and allocation, so this + // exercises the emitter's own last line of defence: an operation with no + // GPU equivalent must be named, not silently dropped. + let src = "[fn helper [x] x] [kernel k [i b] [put b i [helper 1.0]]]"; + let exprs = parse(src).expect("parses"); + let mut checker = Checker::new(); + let _ = checker.check_program(&exprs); + let module = lower(&checker); + let func = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some("k")) + .expect("k lowered"); + + // Whatever happens, it must not be a shader that quietly computes + // something else. + if let Ok(text) = wgsl::emit(&module, func.id, &[ArgKind::output(DType::F32)]) { + validate(&text, "k"); + } +} diff --git a/crates/loon-wasm/src/gpu_bridge.rs b/crates/loon-wasm/src/gpu_bridge.rs new file mode 100644 index 0000000..98b93fa --- /dev/null +++ b/crates/loon-wasm/src/gpu_bridge.rs @@ -0,0 +1,182 @@ +//! A GPU on the other side of a JavaScript function. +//! +//! WebGPU is reachable only through promises, and Loon's VM is synchronous all +//! the way down: a `Place.read` is an ordinary effect operation that returns a +//! value, not a future. Those two facts cannot both be satisfied on one thread. +//! +//! So they are satisfied on two. The VM runs in a Web Worker and every device +//! call goes through the bridge below, which posts a request to the main thread +//! and blocks on `Atomics.wait` until the answer lands in a `SharedArrayBuffer`. +//! The main thread — where the promises live — does the actual WebGPU work and +//! wakes the worker. The blocking is real, and it is the point: it is what lets +//! an asynchronous API sit underneath a synchronous language without either one +//! having to lie about what it is. +//! +//! Nothing here knows about WGSL or wgpu. The shader text arrives already +//! generated by `eir::wgsl`, which is the same text a desktop build hands to +//! wgpu, so the browser and the laptop run the same kernel. + +use loon_lang::eir::device::{Device, DeviceError}; +use loon_lang::eir::vm::{BufData, Buffer}; +use std::cell::RefCell; +use wasm_bindgen::prelude::*; + +thread_local! { + /// The JS function that performs one device operation and returns its + /// result. Synchronous from here; whatever it does to become synchronous + /// is its own business. + static BRIDGE: RefCell> = const { RefCell::new(None) }; +} + +/// Install the JS side of the bridge. +pub fn set_bridge(f: js_sys::Function) { + BRIDGE.with(|b| *b.borrow_mut() = Some(f)); +} + +/// Whether a bridge has been installed. +pub fn has_bridge() -> bool { + BRIDGE.with(|b| b.borrow().is_some()) +} + +fn call(op: &str, payload: &JsValue) -> Result { + BRIDGE.with(|b| { + let guard = b.borrow(); + let f = guard + .as_ref() + .ok_or_else(|| DeviceError("no GPU bridge is installed".to_string()))?; + f.call2(&JsValue::NULL, &JsValue::from_str(op), payload) + .map_err(|e| DeviceError(format!("{op}: {}", describe(&e)))) + }) +} + +/// A thrown JS value, as a sentence. +fn describe(e: &JsValue) -> String { + e.as_string() + .or_else(|| { + js_sys::Reflect::get(e, &JsValue::from_str("message")) + .ok() + .and_then(|m| m.as_string()) + }) + .unwrap_or_else(|| format!("{e:?}")) +} + +/// A device driven by JavaScript. +pub struct BridgeGpu { + resident: RefCell>, +} + +impl BridgeGpu { + pub fn new() -> BridgeGpu { + BridgeGpu { + resident: RefCell::new(std::collections::HashSet::new()), + } + } +} + +impl Default for BridgeGpu { + fn default() -> Self { + Self::new() + } +} + +/// A buffer's bytes, as something JS can read without copying twice. +fn bytes_of(buf: &Buffer) -> js_sys::Uint8Array { + let bytes = buf.to_bytes(); + let arr = js_sys::Uint8Array::new_with_length(bytes.len() as u32); + arr.copy_from(&bytes); + arr +} + +fn dtype_name(buf: &Buffer) -> &'static str { + match &buf.data { + BufData::F32(_) => "f32", + BufData::F64(_) => "f32", // narrowed before it gets here + BufData::I32(_) => "i32", + BufData::I64(_) => "i32", + } +} + +fn obj(pairs: &[(&str, JsValue)]) -> JsValue { + let o = js_sys::Object::new(); + for (k, v) in pairs { + let _ = js_sys::Reflect::set(&o, &JsValue::from_str(k), v); + } + o.into() +} + +impl Device for BridgeGpu { + fn name(&self) -> String { + call("name", &JsValue::NULL) + .ok() + .and_then(|v| v.as_string()) + .unwrap_or_else(|| "webgpu".to_string()) + } + + fn ensure_resident(&self, id: usize, buf: &Buffer) -> Result { + if self.resident.borrow().contains(&id) { + return Ok(false); + } + call( + "upload", + &obj(&[ + ("id", JsValue::from_f64(id as f64)), + ("dtype", JsValue::from_str(dtype_name(buf))), + ("bytes", bytes_of(buf).into()), + ]), + )?; + self.resident.borrow_mut().insert(id); + Ok(true) + } + + fn is_resident(&self, id: usize) -> bool { + self.resident.borrow().contains(&id) + } + + fn dispatch( + &self, + shader: &str, + entry: &str, + n: u32, + scalars: &[f32], + buffers: &[usize], + ) -> Result<(), DeviceError> { + let js_scalars = js_sys::Float32Array::new_with_length(scalars.len() as u32); + js_scalars.copy_from(scalars); + let js_buffers = js_sys::Array::new(); + for id in buffers { + js_buffers.push(&JsValue::from_f64(*id as f64)); + } + call( + "dispatch", + &obj(&[ + ("shader", JsValue::from_str(shader)), + ("entry", JsValue::from_str(entry)), + ("n", JsValue::from_f64(n as f64)), + ("scalars", js_scalars.into()), + ("buffers", js_buffers.into()), + ]), + )?; + Ok(()) + } + + fn download(&self, id: usize, byte_len: usize) -> Result, DeviceError> { + let result = call( + "download", + &obj(&[ + ("id", JsValue::from_f64(id as f64)), + ("byteLength", JsValue::from_f64(byte_len as f64)), + ]), + )?; + let arr = js_sys::Uint8Array::new(&result); + let mut out = vec![0u8; arr.length() as usize]; + arr.copy_to(&mut out); + out.truncate(byte_len); + Ok(out) + } + + fn evict(&self, id: usize) { + if self.resident.borrow_mut().remove(&id) { + let _ = call("evict", &obj(&[("id", JsValue::from_f64(id as f64))])); + } + } +} diff --git a/crates/loon-wasm/src/lib.rs b/crates/loon-wasm/src/lib.rs index d9568d6..dbecfb1 100644 --- a/crates/loon-wasm/src/lib.rs +++ b/crates/loon-wasm/src/lib.rs @@ -1,5 +1,7 @@ use wasm_bindgen::prelude::*; +mod gpu_bridge; + use loon_lang::interp::dom_builtins; use loon_lang::interp::Value; use std::cell::RefCell; @@ -42,12 +44,142 @@ fn call_js_bridge(op: &str, args: &[Value]) -> Result bool { + gpu_bridge::has_bridge() +} + +/// Run a Loon program on the EIR VM, with a placement mode. +/// +/// This is the entry point a browser needs for placed programs: kernels, +/// buffers, and the `Place` effect exist only on the EIR VM, so the +/// interpreter-backed exports below cannot run them at all. +/// +/// `place` is `cpu`, `par`, `device`, or `gpu`. In a browser, `cpu` is the +/// real answer and `device` models a discrete memory so transfer counts can be +/// shown; `par` has no threads to use here and behaves as `cpu`; `gpu` needs a +/// bridge installed with `init_gpu_bridge`, and says so if there is none. +/// +/// Returns the program's printed output followed by its placement accounting, +/// so a page can show what crossed the boundary. +#[wasm_bindgen] +pub fn eval_placed(source: &str, place: &str) -> Result { + let mode = loon_lang::eir::place::Mode::parse(place) + .ok_or_else(|| format!("unknown placement mode '{place}'"))?; + let (result, stats) = + loon_lang::eir::vm::eval_eir_placed(source, std::path::Path::new("."), mode) + .map_err(|e| e.to_string())?; + + let mut out = result.output.join("\n"); + if !out.is_empty() { + out.push('\n'); + } + out.push_str("\u{2014}\n"); + out.push_str(&format!("placed on {}: {}", mode.name(), stats.summary())); + Ok(out) +} + +// ── Suspending a program the host cannot answer synchronously ────────────── +// +// A browser cannot answer `Place.read` immediately, because reading a GPU +// buffer back is a promise. It does not have to: a handler that hands `resume` +// to `Host.park` and returns unwinds the computation, and the page can finish +// it once the bytes arrive. +// +// The session below is what keeps that possible — a parked continuation lives +// in the VM's heap, so the VM has to still be there when the answer shows up. +// See `samples/place/demo-park.oo`, and note where the deferring handler has to sit. + +thread_local! { + static SESSION: RefCell> = + const { RefCell::new(None) }; +} + +/// Start a program, running until it finishes or parks. +/// +/// Returns `{done, output, request}`. `done` false means it parked and is +/// waiting for `place_resume`. +#[wasm_bindgen] +pub fn place_start(source: &str, place: &str) -> Result { + let mode = loon_lang::eir::place::Mode::parse(place) + .ok_or_else(|| format!("unknown placement mode '{place}'"))?; + let mut session = loon_lang::eir::vm::Session::new(source, std::path::Path::new("."), mode) + .map_err(|e| e.to_string())?; + let step = session.start().map_err(|e| e.to_string())?; + let js = step_to_js(&mut session, step); + SESSION.with(|s| *s.borrow_mut() = Some(session)); + Ok(js) +} + +/// Finish a parked step by supplying the numbers the host went to fetch. +#[wasm_bindgen] +pub fn place_resume(values: &[f32]) -> Result { + SESSION.with(|cell| { + let mut guard = cell.borrow_mut(); + let session = guard + .as_mut() + .ok_or_else(|| "no program is running".to_string())?; + let k = session + .pending() + .ok_or_else(|| "nothing is parked".to_string())?; + let data = session.vec_of_floats(values); + let step = session.resume(k, data).map_err(|e| e.to_string())?; + Ok(step_to_js(session, step)) + }) +} + +fn step_to_js( + session: &mut loon_lang::eir::vm::Session, + step: loon_lang::eir::vm::Step, +) -> JsValue { + let o = js_sys::Object::new(); + let (done, request) = match step { + loon_lang::eir::vm::Step::Done(_) => (true, JsValue::NULL), + loon_lang::eir::vm::Step::Parked { request } => { + (false, JsValue::from_str(&session.show(request))) + } + }; + let out = session.take_output().join("\n"); + let _ = js_sys::Reflect::set(&o, &JsValue::from_str("done"), &JsValue::from_bool(done)); + let _ = js_sys::Reflect::set(&o, &JsValue::from_str("output"), &JsValue::from_str(&out)); + let _ = js_sys::Reflect::set(&o, &JsValue::from_str("request"), &request); + let _ = js_sys::Reflect::set( + &o, + &JsValue::from_str("value"), + &JsValue::from_str(&session.show(session.value())), + ); + let _ = js_sys::Reflect::set( + &o, + &JsValue::from_str("stats"), + &JsValue::from_str(&session.stats().summary()), + ); + o.into() +} + /// Evaluate a Loon program and return the result as a string. -// TODO: migrate to `loon_lang::eir::vm::eval_eir` once the EIR VM supports -// the DOM bridge (init_dom_bridge / eval_ui / invoke_callback). The WASM crate -// still uses the legacy tree-walking interpreter because the DOM bridge depends -// on `Value` and `InterpError`, which differ from the EIR's NaN-boxed `Val` / -// `VmResult`. A conversion layer or EIR-native DOM bridge is needed first. +// The DOM-driving exports (`eval_ui`, `invoke_callback`) stay on the legacy +// tree-walking interpreter: the DOM bridge is written against `Value` and +// `InterpError`, which differ from the EIR's NaN-boxed `Val` and `VmResult`. +// `eval_program` and `eval_with_output` stay with them so the guide's examples +// keep working — several use builtins such as `push!` that the EIR VM does not +// implement. Programs that need the EIR VM call `eval_placed` above. #[wasm_bindgen] pub fn eval_program(source: &str) -> Result { let exprs = loon_lang::parser::parse(source).map_err(|e| format!("{e}"))?; diff --git a/docs/blog/2026-08-18-placement-is-an-effect.md b/docs/blog/2026-08-18-placement-is-an-effect.md new file mode 100644 index 0000000..a78a4ef --- /dev/null +++ b/docs/blog/2026-08-18-placement-is-an-effect.md @@ -0,0 +1,242 @@ +# Placement Is an Effect + +*August 2026* + +Last month I wrote that a syscall is an effect, and so the kernel is just the outermost handler. This month someone handed me a paper about running Rust on GPUs, and I spent a week discovering that the same sentence has another half. + +The paper is [GPU Offload in Rust: Portable, Safe, and Fast](https://arxiv.org/abs/2608.13759), by Drehwald, Domínguez, Sala, Aspuru-Guzik, and Doerfert. It is good work and you should read it. They put GPU compilation inside `rustc`, they get data direction out of `&T` versus `&mut T` instead of making you write pragmas, and they measure themselves honestly against hand-written CUDA. My interest is in the part they are candid about not having solved, because it turns out to be a shape I recognized. + +## The 400x + +Their convenient interface looks like this: + +```rust +offload!(vec_add, &a, &b, &mut c); +``` + +You write that, and the compiler transfers `a` and `b` to the GPU, runs the kernel, and copies `c` back. Lovely. Now put it in a loop, and every iteration ships the same two arrays across the bus again. They measure this at **up to 400x slower** than doing the transfers by hand. + +So they add a second interface. You wrap your data in a `Preload` or `PreloadMut`, which pins it on the device and marks the point where it comes home with the value's `drop`. That works, and it costs you an annotation at every site. And then, to make the *convenient* interface fast too, they prototype a transfer-hoisting pass inside LLVM — loop-invariant code motion for `memcpy`s, essentially — which is still future work at the time of writing. + +There is one benchmark they can't fix that way. It's called Energy: six kernels sharing about fifteen arrays, each with a few of its own. Keeping everything resident might blow past the device's memory; keeping nothing resident is the 400x. Deciding needs a heuristic, and they say plainly that they didn't ship one. + +I read that and thought: that isn't a missing optimization. That's a missing *seam*. + +## The seam + +Here's the thing the compiler is trying to recover. It wants to know two facts: + +1. What does this launch touch? +2. When does the host actually want to look? + +And it can't just ask, because in Rust neither fact is written down anywhere. A launch is a function call. A host read is... a host read, some ordinary expression somewhere in the program that happens to name the same variable. So the compiler goes and reconstructs both from dataflow analysis, and where analysis fails, `Preload` makes you write the answer down by hand. + +But both facts are *events*. Things that happen, in an order, that somebody might want to intercept. + +We have a way to spell that. + +``` +[Place.run saxpy n #[3.0 x y out]] +[Place.read out] +``` + +`Place.run` is an effect operation. `Place.read` is an effect operation. Not a function, not a method on a smart pointer — an operation, which floats up until some handler catches it. And once those two are effects, the two questions the LLVM pass was reverse-engineering are just *arguments to a handler clause*. + +## Nine lines + +Here is the residency policy. It is in `samples/place/lib.oo`. It is not privileged, it is not in the compiler, and you could have written it. + +``` +[fn place/resident [thunk] + [handle [thunk] + [Place.run k n args] + [do [Place.pin args] + [resume [Place.run k n args]]] + [Place.read b] + [do [let v [Place.read b]] + [Place.unpin b] + [resume v]]]]] +``` + +Read it slowly, because it is doing the whole job. It catches every launch, tells the device to keep whatever that launch touched, and forwards the launch outward unchanged. It catches every read, lets the read happen, and releases the pin. That's it. `Place.pin` is the entire vocabulary — it means "this will be wanted again." + +Wrap it around a program and: + +``` +loon run samples/place/demo-residency.oo --place gpu + +eight launches over one buffer + answer #[8 8 8 8] + no policy: uploads 8, resident hits 0, bytes in 128 + answer #[8 8 8 8] + place/resident: uploads 1, resident hits 7, bytes in 16 +``` + +Same program both times. Same answer both times. Eight uploads became one, because a handler said so. + +On an actual GPU — this is an M4 Max, through Metal, via wgpu — the wall clock follows: + +| launches | no policy | place/resident | speedup | +|---------:|----------:|---------------:|--------:| +| 8 | 29.3 ms | 8.4 ms | 3.5x | +| 32 | 94.8 ms | 11.2 ms | 8.4x | +| 128 | 358.8 ms | 18.1 ms | 19.9x | + +The gap grows with the chain, which is exactly the paper's curve. The difference is where the fix lives. Theirs is an LLVM pass and a type. Mine is a `handle` form you can read in one sitting, change without rebuilding a compiler, and — this is the part I keep coming back to — *replace with a different one* when your program has different needs. + +Which brings us to Energy. + +## The heuristic they didn't ship + +The Energy case is only hard if there's exactly one policy and it has to be right for everybody. When residency is a handler, "keep everything" and "keep these" are two handlers, and choosing between them is a line of code rather than a compiler flag nobody can change: + +``` +[place/resident-only #[e-new p-new q-new] work] +``` + +That's the heuristic. It's an argument. I am not claiming I solved their benchmark — I haven't run RAJAPerf and I'm not going to pretend otherwise. I'm claiming the thing they needed a heuristic *for* is, in this design, a place where the user gets to put one. + +## Nobody wrote `&mut` + +Here's the part I'm smug about. This is a kernel: + +``` +[kernel saxpy [i a x y out] + [put out i [+ [* a [at x i]] [at y i]]]] +``` + +`x` and `y` are inputs. `out` has to come home. Nothing in that source says so. + +Loon already had an ownership pass that figures out, for every function parameter, whether the body reads it, writes through it, or consumes it — the same distinction Rust makes you spell as `&T` / `&mut T` / `T`. It was computing that, using it for error messages, and throwing it away. Now it rides into the IR, and the placement layer reads it: `at` is a read, `put` is a write-through, so `out` is the argument that needs synchronizing back. The paper reads exactly the same fact off exactly the same distinction. They just make you type it. + +The emitted shader gets it right down to the binding: + +```wgsl +@group(0) @binding(1) var b1: array; +@group(0) @binding(2) var b2: array; +@group(0) @binding(3) var b3: array; +``` + +`read` versus `read_write`, decided by whether the kernel body said `at` or `put`. + +## What a kernel isn't + +Kernels are restricted. No closures, no allocation, no strings, no effects. Try it and the compiler names what it found: + +``` +kernel 'k' contains the effect operation 'IO.println' + why: a kernel runs where there is no handler tower to perform effects against +``` + +That restriction is the safety argument, and it's where I think this design earns its keep against theirs. Their kernels can receive a slice and index it however they like, so "threads touch disjoint elements" has to be *promised* — by an `unsafe impl` of a partitioning strategy. Here a kernel receives an index and writes at that index. The unsafe program isn't rejected; it's unwriteable. + +I'll take a restriction over a promise. A promise is a place where someone will eventually be wrong. + +## Testing a GPU program without a GPU + +This is the part the paper doesn't have a section for, and I don't think that's an oversight so much as a consequence: if launches aren't events, there's nothing to record. + +``` +loon run samples/place/saxpy.oo --place gpu --record trace.oo +loon replay trace.oo samples/place/saxpy.oo +``` + +The second command runs on a build with no GPU support compiled into it at all, and prints the same thing. From the program's point of view a kernel launch was an operation that returned nothing and a read was one that produced some numbers, so recording those *is* the run. + +`Place.stats` is deliberately excluded from the recording. It reports on the run currently happening, and a replayed run genuinely moved no bytes — feeding back the original transfer counts would be a recording that lies about the execution it's part of. The replayed run says zero launches, because it performed zero launches. + +And since a handler can decline to forward at all, "run this without a device" is four lines: + +``` +[fn place/dry-run [thunk] + [handle [thunk] + [Place.run k n args] [resume []] + [Place.read b] [resume #[]]]] +``` + +Every accounting number, no execution. And strace-for-GPU is the same shape as strace-for-syscalls was last month — perform the operation you intercepted, print on the way past. + +## Where it runs + +Metal on this laptop, today. Vulkan on a Linux box and DX12 on Windows are the same code path through wgpu, and I have not run either, so take them as "should" rather than "does." + +And the browser, which I want to describe carefully because getting there was the strangest part. + +Programs run in a tab, on the actual GPU: + +``` +placed on gpu: 4 launches over 32 work items; + 9 uploads (288 B), 3 downloads (96 B), 3 resident hits +``` + +That is WebGPU, driven by the same WGSL a desktop build hands to wgpu, with the residency handler deciding what gets copied. + +Here is the strange part. WebGPU is asynchronous — you get a device from a promise and read a buffer back through `mapAsync`. Loon's VM is synchronous all the way down; `Place.read` is an effect operation that returns a value, not a future. Those two facts cannot both hold on one thread. + +So they hold on two. The VM runs in a Web Worker, and every device call posts a request to the main thread and then blocks on `Atomics.wait` until the answer lands in a `SharedArrayBuffer`. The main thread, where the promises live, does the WebGPU work and wakes the worker. The blocking is real, and it is the whole trick: it lets an asynchronous API sit underneath a synchronous language without either one pretending to be the other. + +None of that reached the VM. It sees an `eir::device::Device` — six operations: name, ensure-resident, is-resident, dispatch, download, evict — and wgpu implements it on a laptop while a JavaScript bridge implements it in a tab. Which is the same move placement makes at the language level, one floor down: the thing that varies goes behind an interface, and the code above does not change when the answer does. + +The cost is a requirement for cross-origin isolation, since `SharedArrayBuffer` needs COOP/COEP headers. Without them the page says so and the other placements still work. + +I said in a first draft that getting rid of that would need an asynchronous effect path in the VM. Then I tried it, and it turns out the VM already has one — it just isn't called that. + +A handler clause does not have to call `resume`. If it hands `resume` somewhere else and returns, the handled computation *unwinds*, and the continuation is still live in whoever caught it. Call it later and the program picks up mid-expression, exactly where it stopped: + +``` +work: starting +host: computation parked; the rest of it is mine now +host: ...doing something slow... +work: continued with 21 +host: finished with 42 +``` + +That's `samples/place/demo-park.oo`, and it is the entire mechanism an asynchronous host needs. A browser can't answer `Place.read` immediately — reading a GPU buffer back is a promise — but it doesn't have to answer immediately. It can take the continuation, go away, and come back when the bytes arrive. Uploads and dispatches need none of this, because `writeBuffer` and `submit` are already synchronous. + +What's left there is plumbing: a VM that outlives one call, since the continuation lives in its heap, and an export for the page to resume through. Not semantics. I had assumed the hard part was the language and the easy part was the wiring, and it was the other way round — which is what I get for writing down what I thought was true instead of trying it. + +Every kernel in the repo is parsed and type-checked by naga in CI, on machines with no GPU. That's the automated cross-target validation the paper says is still missing — they found a host/device divergence in slice lowering by hand, `(ptr, len)` on two targets and `[i64; 2]` on a third. We have the same class of hazard: NaN-boxing constants that used to be copy-pasted into three backends under a comment asking the next person to keep them in sync. They now live in one file, and a conformance test compiles the same literals on every backend and compares raw bits. + +That test found three real divergences the first time it ran, including one where `loon run --native` silently returned `()` for any program with a `main` function. Which is a good argument for writing the test. + +## The row I didn't expect + +Once kernels stopped going through the interpreter — there's a typed executor now that runs the numeric subset against raw slices — I added `--place par`, which splits the index range across cores. Each thread gets a disjoint piece of the output from `split_at_mut`, so "threads touch disjoint elements" isn't promised by an `unsafe impl`; it's what the borrow checker hands back. + +Then I ran the same kernel four ways: + +| elements | cpu | par | gpu | +|---------:|----:|----:|----:| +| 1,024 | 447 µs | 566 µs | 10.3 ms | +| 262,144 | 8.9 ms | 3.3 ms | 12.4 ms | +| 1,048,576 | 36.4 ms | 11.4 ms | 19.2 ms | + +Every core beats the GPU at a million elements. This machine has a lot of fast ones, and a launch pays submission and transfer before it computes anything. + +I like this result more than I'd like a win. Where the crossover sits is a property of the machine, not of the program — and I found it by changing one word on a command line, because the program genuinely does not know where it runs. If placement were a compile-time decision I'd have had to rebuild something to ask the question, and I probably wouldn't have bothered. + +One more thing falls out of making `Place.read` the only way to get data back: launches don't block on each other. A dispatch submits and returns; nothing waits until the host asks. Sixty-four launches take 19.9 ms against 6.9 ms for one — if each waited, that would be closer to 440. The paper prototypes asynchronous transfers as a separate optimization. Here it's just what happens when the synchronization point is a thing the program says out loud. + +## What I'm not claiming + +We do not beat hand-written CUDA. We haven't measured against it and we're not going to imply otherwise. The CPU column in our benchmarks is Loon's own interpreter — the slowest honest baseline — so "3.1x faster on the GPU at 262k elements" means *there is a lot to gain by leaving the interpreter*, not anything about generated code quality. + +Reductions and atomics are outside the kernel subset. WGSL core has no 64-bit scalar, so an f64 buffer is computed in f32 on the device and we report the narrowing rather than hiding it. + +And the honest summary of the whole comparison: they built a compiler and I moved a seam. Those are different kinds of work. The reason I think the seam is worth the post is that it makes a class of thing — residency, prefetch, eviction, tracing, simulation, replay — stop being compiler features that someone has to ship for you, and start being ordinary code that you can write on a Tuesday. + +## The arc + +v0.7: effects, end to end. v0.8: syscalls are effects, so the kernel is the outermost handler. v0.9: placement is an effect, so the GPU is a handler in the middle. + +I don't have a fourth one yet. But I've stopped being surprised when something that looked like it needed a compiler pass turns out to need a `handle`. + +--- + +Try it: + +``` +loon run samples/place/demo-handlers.oo # one program, four handlers +loon run samples/place/demo-residency.oo --place device # the transfer gap +loon run samples/place/saxpy.oo --place gpu # on real hardware +``` diff --git a/docs/plans/2026-08-18-placement-remaining.md b/docs/plans/2026-08-18-placement-remaining.md new file mode 100644 index 0000000..caa3dfa --- /dev/null +++ b/docs/plans/2026-08-18-placement-remaining.md @@ -0,0 +1,143 @@ +# Placement: what is done and what is not + +*August 2026 — companion to `docs/blog/2026-08-18-placement-is-an-effect.md`* + +Written down so the next person does not have to infer the boundary from the +code. + +## Shipped + +- **Buffers** (`eir::vm::Buffer`) — dense, unboxed numeric arrays, the + representation that can leave the process. +- **The `kernel` form** — desugars to `fn` before anything else looks at it; + `check::kernel` enforces the subset and names what it rejects. +- **The `Place` effect** — `run`, `read`, `pin`, `unpin`, `stats`. Unhandled, + a program runs serially right here. +- **Placement modes** — `--place cpu | par | device | gpu`, all agreeing on the + answer, plus `--place-stats`. +- **A typed kernel executor** (`eir::kernel_exec`) — the numeric subset over raw + slices, and a parallel driver that carves *every* output buffer into + per-thread slices with `split_at_mut`. +- **WGSL emission** (`eir::wgsl`) — validated by naga in CI on machines with no + GPU. +- **Real GPU execution** (`eir::gpu`, feature `gpu`) — wgpu, with resident + device buffers so the transfer accounting describes the hardware. +- **Record and replay** — a run on the GPU replays on a build with no GPU. +- **Handlers** (`samples/place/lib.oo`) — trace, dry-run, counted, resident, + resident-only. +- **Reductions** (`samples/place/reduce.oo`) — no new feature needed: a work + item sums its own chunk into its own slot, which is still "write at your own + index", and the partials are combined on the host. +- **Placed programs in a browser, including on the GPU** — `crates/loon-wasm` + exports `eval_placed`; `web/public/place.html` runs cpu, par, device, and + **real WebGPU** in a tab, all four agreeing on the answer. Because Loon's VM + is synchronous and WebGPU is not, the VM runs in a worker and blocks on + `Atomics.wait` while the main thread drives the device. +- **`eir::device::Device`** — the six operations a placement backend provides. + wgpu implements it natively; `loon-wasm` implements it by proxying to + JavaScript. The VM does not know which it has. + +## Decided, and deliberately restrictive + +These are not gaps. They are the rules the design rests on, enforced. + +- **A kernel writes at its own index and nowhere else.** Rejected at compile + time (E0602), naming the offending index. Reading anywhere is fine — gather is + safe, scatter is not. This is what lets the parallel executor hand each thread + a slice and a GPU run every work item at once. +- **No 64-bit on a GPU.** WGSL core has no `f64` or 64-bit integer, so such a + launch is refused with a message naming the type and the alternative rather + than narrowed silently. Handing back a precision the program never asked for, + and cannot detect, is the failure this whole design exists to avoid. +- **Kernels cannot allocate, close over, or perform effects.** The restriction + is the safety argument: the unsafe program is not rejected, it is unwriteable. + +## Not done + +### Reaching a GPU without cross-origin isolation + +Two paths now exist, and the second one no longer needs isolation. + +**The worker path**, which is what `web/public/place.html` uses today: the VM +runs in a Web Worker and blocks on `Atomics.wait`, which needs +`SharedArrayBuffer`, which needs COOP/COEP headers (`vercel.json` sets them). +Verified on a real GPU. + +**The parking path**, which needs none of that. A handler that hands `resume` +to `Host.park` and returns unwinds the computation, and the page finishes it +when the bytes arrive. `eir::vm::Session` keeps the VM alive between steps, and +`place_start` / `place_resume` expose it to JavaScript. Confirmed in a browser +with `crossOriginIsolated === false`: + + start: done=false ← parked at Place.read + resume: done=true out="read #[2 4 6]" value=12 ← finished with host data + +An earlier version of this document said this would need "an asynchronous +effect path in the VM". It did not: escaping continuations already were one. + +What is left is only to point the demo page at it. That means writing the +sample with the deferring handler outermost — parking unwinds to the `handle`, +so anything after it runs immediately with the placeholder (see +`samples/place/demo-park.oo`, and the test named for it). Uploads and dispatches need +nothing, because `writeBuffer` and `submit` are already synchronous; only the +readback ever has to wait. +### The DOM exports + +`eval_ui` and `invoke_callback` remain on the legacy tree-walking interpreter. +Its bridge is written against `Value`/`InterpError` rather than the EIR's +`Val`/`VmResult`. `eval_program` and `eval_with_output` were left with them +rather than regressing documented pages; `eval_placed` was added alongside. + +The deeper blocker is the one below, found while looking into this. + +### The mutators, which need a decision rather than an implementation + +`set!` and `push!` are documented — `web/src/pages/guide/collections.loon`, +`guide/ownership.loon`, `ref/builtins.loon`, `DESIGN.md` — and **neither exists +on the EIR VM**, which is the default backend. On the interpreter they behave +like this: + + [let mut v #[1 2 3]] + [push! v 4] + [println v] ; #[1 2 3] — unchanged + [println [push! v 9]] ; #[1 2 3 9] + +So `push!` does not mutate. It returns a new vector, and the `!` promises +something it does not do. The guide's own example asserts otherwise: + + [let items [mut #[]]] + [push! items 1] + [push! items 2] + [println items] ; #[1 2] ← documented + +That prints `#[]` on the interpreter and fails to type check on the EIR VM. The +documented behaviour is currently true of no backend. + +This was not implemented on the EIR VM as part of the placement work because +the right fix is a language decision, not a port: + +1. **Make `!` mean mutation.** `push!` writes through the binding, matching the + name, the guide, and what `put` already does for buffers. The ownership pass + already classifies both as mutable borrows, so the analysis is in place. The + question is what happens to a closure that captured the old value. +2. **Make `!` mean "returns a changed copy"** and fix the guide and the name. + Smaller change, but then `set!` and `push!` do not agree with each other, + since `set!` really does rebind. + +Either way both need to exist on the default backend. Implementing the current +interpreter behaviour verbatim would spread a naming problem to a second +backend, so it is left for whoever decides which of the two Loon means. + +### Atomics + +Anything genuinely needing them — a histogram, a scatter-add — is out, and for +the same reason scatter is: two work items reaching the same element is exactly +what the disjointness rule forbids. Supporting them would mean a second kind of +kernel with a different safety argument, not a relaxation of this one. + +### Numbers we do not have + +No comparison against hand-written CUDA, HIP, or Metal. No RAJAPerf port. The +benchmark's CPU column is Loon's own typed executor, which is a fair floor but +not an optimized C baseline, and nothing in `BENCHMARKS.md` should be read as +one. diff --git a/samples/place/demo-handlers.oo b/samples/place/demo-handlers.oo new file mode 100644 index 0000000..d33aef1 --- /dev/null +++ b/samples/place/demo-handlers.oo @@ -0,0 +1,37 @@ +; Demo: one program, four answers to "where does this run". +; +; `blur` and `work` never mention a device, a queue, or a transfer. The only +; thing that changes between the runs below is which handler is wrapped around +; the same thunk — and the answer is identical every time. +; +; loon run samples/place/demo-handlers.oo +[use lib] + +; A two-kernel chain. Chains are where placement policy earns its keep: the +; intermediate never needs to come back to the host, and only a handler is in +; a position to know that. +[kernel scale [i s src dst] [put dst i [* s [at src i]]]] + +[kernel offset [i d buf] [put buf i [+ d [at buf i]]]] + +[fn work [] + [let src [buf #[1 2 3 4 5 6 7 8]]] + [let mut dst [buf-zeros 8]] + [Place.run scale 8 #[10.0 src dst]] + [Place.run offset 8 #[1.0 dst]] + [Place.read dst]] + +[fn main [] + [IO.println "no handler — it runs right here:"] + [IO.println [str " " [work]]] + + [IO.println "traced — every launch and sync is visible:"] + [IO.println [str " " [place/trace work]]] + + [IO.println "dry run — accounted for, nothing executed:"] + [IO.println [str " " [place/dry-run work]]] + + [IO.println "resident — buffers stay put between launches:"] + [IO.println [str " " [place/resident work]]] + + [IO.println [str "totals: " [Place.stats]]]] diff --git a/samples/place/demo-park.oo b/samples/place/demo-park.oo new file mode 100644 index 0000000..6164e18 --- /dev/null +++ b/samples/place/demo-park.oo @@ -0,0 +1,58 @@ +; Demo: suspending a computation and finishing it later. +; +; A handler clause does not have to call `resume`. If it hands `resume` +; somewhere else and returns, the handled computation unwinds — and the +; continuation is still live, held by whoever received it. Calling it later +; picks the program up exactly where it stopped. +; +; That is the whole mechanism an asynchronous host needs. A browser cannot +; answer `Place.read` immediately, because reading a GPU buffer back is a +; promise; but it does not have to. It can take the continuation, go away, and +; come back when the bytes arrive. +; +; Nothing here is special-cased in the VM. This is what reified, escaping +; continuations already are. +; +; One thing to know before using this, because it decides where the handler +; goes: parking unwinds to the `handle`, so whatever follows the handle runs +; *immediately*, with whatever the clause returned — not with the real answer. +; The suspended part is the only thing that waits. So a handler that defers an +; answer belongs outermost, with nothing meaningful after it, and the value of +; the computation comes back when the continuation is called rather than from +; the call that parked. +; +; loon run samples/place/demo-park.oo +[use lib] + +[effect Slow + [fetch [Int] Int]] + +[effect Host + [park [a] Unit]] + +; A computation that stops in the middle to ask for something slow. +[fn work [] + [IO.println " work: starting"] + [let answer [Slow.fetch 1]] + [IO.println [str " work: continued with " answer]] + [* answer 2]] + +; The handler that suspends. It does not resume; it hands `resume` outward and +; returns, so the computation unwinds with its rest still intact. +[fn suspending [thunk] + [handle [thunk] [Slow.fetch id] [do [Host.park resume] 0]]] + +[fn suspended-work [] [suspending work]] + +; The host: catches the parked continuation, does whatever slow thing it +; likes, and then calls it. +[fn main [] + [handle + [suspended-work] + [Host.park k] + [do + [IO.println " host: computation parked; the rest of it is mine now"] + [IO.println " host: ...doing something slow..."] + [IO.println [str " host: finished with " [k 21]]] + [resume []]]] + [IO.println "done"]] diff --git a/samples/place/demo-residency.oo b/samples/place/demo-residency.oo new file mode 100644 index 0000000..84f5125 --- /dev/null +++ b/samples/place/demo-residency.oo @@ -0,0 +1,61 @@ +; Demo: the transfer gap, and the handler that closes it. +; +; A recent Rust GPU-offload paper measures up to 400x between its convenient +; interface (transfer per launch) and its explicit one (transfer once). Closing +; that gap there needs two things: `Preload`/`PreloadMut` types written by the +; programmer at every call site, and a transfer-hoisting pass inside LLVM that +; the authors describe as a prototype. +; +; Here it is a handler. `chain` below is one function, unchanged between the +; two runs; the only difference is whether `place/resident` is wrapped around +; it. Run with: +; +; loon run samples/place/demo-residency.oo --place device +; +; `--place device` gives buffers a separate memory, so transfers are real and +; counted. On `--place cpu` there is one memory and nothing to move — the same +; program, the same answer, no transfers at all. +[use lib] + +[kernel step [i b] [put b i [+ 1.0 [at b i]]]] + +; Eight launches over the same buffer. Nothing here mentions residency; the +; question "does this buffer need to travel" is not the program's to answer. +[fn chain [] + [let mut b [buf #[0 0 0 0]]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.run step 4 #[b]] + [Place.read b]] + +[fn report [label before] + [let s [Place.stats]] + [IO.println + [str + " " + label + ": uploads " + [- [get s :uploads] [get before :uploads]] + ", resident hits " + [- [get s :resident-hits] [get before :resident-hits]] + ", bytes in " + [- [get s :bytes-in] [get before :bytes-in]]]] + s] + +[fn main [] + [IO.println "eight launches over one buffer"] + + [let start [Place.stats]] + [IO.println [str " answer " [chain]]] + [let after-naive [report "no policy" start]] + + [IO.println [str " answer " [place/resident chain]]] + [let _ [report "place/resident" after-naive]] + + [IO.println ""] + [IO.println "same program, same answer; the handler decided what had to move"]] diff --git a/samples/place/lib.oo b/samples/place/lib.oo new file mode 100644 index 0000000..16cc509 --- /dev/null +++ b/samples/place/lib.oo @@ -0,0 +1,84 @@ +; Placement handlers. +; +; `Place` is an effect, so where a kernel runs is decided by whoever is +; handling it — not by the program, and not by a compiler flag baked into the +; build. Everything in this file is ordinary Loon: no compiler support, no +; intrinsics, nothing the language does not already give you for any other +; effect. That is the point being made. The transfer-hoisting and residency +; policies that a GPU offload compiler implements as optimization passes are +; here, as library code you can read in one sitting and change without +; rebuilding a compiler. +; +; Each handler wraps a thunk and returns its value: +; +; [place/trace work] ; watch what happens +; [place/resident work] ; keep buffers where they are +; [place/dry-run work] ; account for it, run nothing +; ── Watching ──────────────────────────────────────────────────────────────── +; strace, for placement. Forwards every operation outward unchanged and prints +; what went past. The whole of "GPU tracing" is a handler that performs the +; same operation it intercepted. +[fn place/trace [thunk] + [handle + [thunk] + [Place.run k n args] + [do + [IO.println [str "place: launch " n " work items, " [len args] " args"]] + [resume [Place.run k n args]]] + [Place.read b] + [do + [IO.println [str "place: sync — reading " [buf-len b] " elements back"]] + [resume [Place.read b]]]]] + +; Count launches and syncs without printing anything. Returns the thunk's +; value; ask `Place.stats` afterwards for the totals. +[fn place/counted [thunk] + [handle + [thunk] + [Place.run k n args] + [resume [Place.run k n args]] + [Place.read b] + [resume [Place.read b]]]] + +; ── Accounting without running ────────────────────────────────────────────── +; Answer every launch by doing nothing, and every read with an empty vector. +; The program runs to completion and its placement decisions are counted, but +; no kernel executes. This is how you ask "what would this cost" — and how a +; program targeting hardware you do not have can still be exercised. +[fn place/dry-run [thunk] + [handle [thunk] [Place.run k n args] [resume []] [Place.read b] [resume #[]]]] + +; ── Residency ─────────────────────────────────────────────────────────────── +; Keep buffers on the device across launches, and let the host pay for a +; transfer only when it actually looks. +; +; This is the handler that closes the gap a recent Rust offload paper measured +; at up to 400x between its convenient interface and its explicit one. There, +; closing it needed `Preload`/`PreloadMut` types written by the programmer at +; every call site, plus a transfer-hoisting pass inside LLVM the authors +; describe as a prototype. Here it is nine lines, because the information the +; pass was trying to recover is already in the program: `Place.run` says what +; a launch touches and `Place.read` says when the host wants an answer. +; +; The policy is simply: everything a launch touches is worth keeping. Pinning +; is what stops the device from dropping it, so the next launch finds it +; already there. +[fn place/resident [thunk] + [handle + [thunk] + [Place.run k n args] + [do [Place.pin args] [resume [Place.run k n args]]] + [Place.read b] + [do [let v [Place.read b]] [Place.unpin b] [resume v]]]] + +; The same shape, but only keeping the buffers you name. Useful when the +; working set is larger than the device and residency has to be chosen rather +; than assumed — the case the paper flags as needing a heuristic it did not +; ship. Here the heuristic is an argument. +[fn place/resident-only [keep thunk] + [handle + [thunk] + [Place.run k n args] + [do [Place.pin keep] [resume [Place.run k n args]]] + [Place.read b] + [resume [Place.read b]]]] diff --git a/samples/place/reduce.oo b/samples/place/reduce.oo new file mode 100644 index 0000000..bad609d --- /dev/null +++ b/samples/place/reduce.oo @@ -0,0 +1,45 @@ +; Reduction, without leaving the rules. +; +; A kernel may write at its own index and nowhere else. That sounds like it +; rules out summing an array — a sum has one destination and every element +; wants to reach it. +; +; It doesn't, because a reduction is a map followed by a much smaller +; reduction. Give each work item its own chunk of the input and its own slot in +; a partials buffer, and every item is still writing exactly one element: its +; own. Combine the handful of partials afterwards. +; +; So there is no `Place.reduce` operation and no workgroup-shared memory here. +; The two-phase shape is what the disjointness rule already permits, which is a +; better outcome than a special case would have been. +; +; loon run samples/place/reduce.oo --place cpu | par | device | gpu +; Phase one: work item `c` sums the `width` elements starting at `c * width`. +; It reads all over the input — gather is safe — and writes only `partials[c]`. +[kernel + sum-chunk + [c width src partials] + [let start [* c width]] + [let total + [loop + [k 0 acc 0.0] + [if [>= k width] acc [recur [+ k 1] [+ acc [at src [+ start k]]]]]]] + [put partials c total]] + +[fn main [] + [let n 4096] + [let chunks 64] + [let width [/ n chunks]] + + [let src [buf [range 0 n]]] + [let mut partials [buf-zeros chunks]] + + [Place.run sum-chunk chunks #[width src partials]] + + ; Phase two: sixty-four numbers, added on the host. Sending this to a device + ; would cost more than doing it. + [let total [sum [Place.read partials]]] + + [IO.println [str "sum of 0.." n " = " total]] + [IO.println [str "expected = " [/ [* [- n 1] n] 2]]] + [IO.println [str "stats: " [Place.stats]]]] diff --git a/samples/place/saxpy.oo b/samples/place/saxpy.oo new file mode 100644 index 0000000..54fea5c --- /dev/null +++ b/samples/place/saxpy.oo @@ -0,0 +1,22 @@ +; saxpy: out[i] = a * x[i] + y[i] +; +; Nothing in this program says where the work happens. `Place.run` is an +; effect; with no handler installed it runs serially on this CPU. Install a +; handler and the same source runs somewhere else — that is the whole idea. +; +; Note what is *not* written anywhere: which buffers travel to the device and +; which have to come back. Loon infers that. `x` and `y` are only ever read +; (`at`), so they are inputs; `out` is written through (`put`), so it is the +; one that has to be synchronized back. +[kernel saxpy [i a x y out] [put out i [+ [* a [at x i]] [at y i]]]] + +[fn main [] + [let n 8] + [let x [buf [range 0 n]]] + [let y [buf [map [fn [v] [* 2 v]] [range 0 n]]]] + [let mut out [buf-zeros n]] + + [Place.run saxpy n #[3.0 x y out]] + + [IO.println [str "out = " [Place.read out]]] + [IO.println [str "stats = " [Place.stats]]]] diff --git a/web/build.ts b/web/build.ts index 8ac630c..7289a5a 100644 --- a/web/build.ts +++ b/web/build.ts @@ -78,7 +78,7 @@ if (noRust) { // Step 2: Copy static assets console.log('Copying static assets...'); mkdirSync(DIST, { recursive: true }); -for (const file of ['index.html', 'boot.js', 'style.css', 'install.sh', 'loon.png', 'loon-light.png', 'loon-bird.svg']) { +for (const file of ['index.html', 'boot.js', 'style.css', 'install.sh', 'loon.png', 'loon-light.png', 'loon-bird.svg', 'place.html', 'place-gpu.js', 'place-worker.js']) { try { cpSync(join(PUBLIC, file), join(DIST, file)); } catch { diff --git a/web/public/pkg/loon_wasm.d.ts b/web/public/pkg/loon_wasm.d.ts index a71b9c6..cc8bfbf 100644 --- a/web/public/pkg/loon_wasm.d.ts +++ b/web/public/pkg/loon_wasm.d.ts @@ -16,6 +16,23 @@ export function clear_effect_log(): void; */ export function enable_effect_log(enabled: boolean): void; +/** + * Run a Loon program on the EIR VM, with a placement mode. + * + * This is the entry point a browser needs for placed programs: kernels, + * buffers, and the `Place` effect exist only on the EIR VM, so the + * interpreter-backed exports below cannot run them at all. + * + * `place` is `cpu`, `par`, `device`, or `gpu`. In a browser, `cpu` is the + * real answer and `device` models a discrete memory so transfer counts can be + * shown; `par` has no threads to use here and behaves as `cpu`; `gpu` needs a + * bridge installed with `init_gpu_bridge`, and says so if there is none. + * + * Returns the program's printed output followed by its placement accounting, + * so a page can show what crossed the boundary. + */ +export function eval_placed(source: string, place: string): string; + /** * Evaluate a Loon program and return the result as a string. */ @@ -43,6 +60,11 @@ export function eval_with_output(source: string): string; */ export function get_effect_log(): string; +/** + * Whether a GPU bridge has been installed. + */ +export function has_gpu_bridge(): boolean; + /** * Infer the type of the last named binding (fn/let) in the source. */ @@ -55,11 +77,38 @@ export function infer_type(source: string): string; */ export function init_dom_bridge(bridge: Function): void; +/** + * Install the GPU bridge: a JS function taking `(op, payload)` and returning + * synchronously. + * + * The operations are `name`, `upload`, `dispatch`, `download`, and `evict`. + * How the JS side becomes synchronous is up to it — the demo runs the VM in a + * worker and blocks on `Atomics.wait` while the main thread drives WebGPU. + * + * With a bridge installed, `--place gpu` works in a browser. Without one it + * says there is nowhere to run, which is better than quietly running on the + * CPU and reporting GPU numbers. + */ +export function init_gpu_bridge(bridge: Function): void; + /** * Invoke a stored Loon callback by ID (called from JS event handlers). */ export function invoke_callback(id: number): void; +/** + * Finish a parked step by supplying the numbers the host went to fetch. + */ +export function place_resume(values: Float32Array): any; + +/** + * Start a program, running until it finishes or parks. + * + * Returns `{done, output, request}`. `done` false means it parked and is + * waiting for `place_resume`. + */ +export function place_start(source: string, place: string): any; + /** * Reset the Loon runtime state (callbacks, etc.) for hot reload. */ @@ -71,14 +120,19 @@ export interface InitOutput { readonly memory: WebAssembly.Memory; readonly check_program: (a: number, b: number) => [number, number, number, number]; readonly enable_effect_log: (a: number) => void; + readonly eval_placed: (a: number, b: number, c: number, d: number) => [number, number, number, number]; readonly eval_program: (a: number, b: number) => [number, number, number, number]; readonly eval_ui: (a: number, b: number) => [number, number]; readonly eval_ui_checked: (a: number, b: number) => [number, number]; readonly eval_with_output: (a: number, b: number) => [number, number, number, number]; readonly get_effect_log: () => [number, number]; + readonly has_gpu_bridge: () => number; readonly infer_type: (a: number, b: number) => [number, number, number, number]; readonly init_dom_bridge: (a: any) => void; + readonly init_gpu_bridge: (a: any) => void; readonly invoke_callback: (a: number) => void; + readonly place_resume: (a: number, b: number) => [number, number, number]; + readonly place_start: (a: number, b: number, c: number, d: number) => [number, number, number]; readonly clear_effect_log: () => void; readonly reset_runtime: () => void; readonly __wbindgen_malloc: (a: number, b: number) => number; diff --git a/web/public/pkg/loon_wasm.js b/web/public/pkg/loon_wasm.js index cec6a62..524b8a6 100644 --- a/web/public/pkg/loon_wasm.js +++ b/web/public/pkg/loon_wasm.js @@ -41,6 +41,47 @@ export function enable_effect_log(enabled) { wasm.enable_effect_log(enabled); } +/** + * Run a Loon program on the EIR VM, with a placement mode. + * + * This is the entry point a browser needs for placed programs: kernels, + * buffers, and the `Place` effect exist only on the EIR VM, so the + * interpreter-backed exports below cannot run them at all. + * + * `place` is `cpu`, `par`, `device`, or `gpu`. In a browser, `cpu` is the + * real answer and `device` models a discrete memory so transfer counts can be + * shown; `par` has no threads to use here and behaves as `cpu`; `gpu` needs a + * bridge installed with `init_gpu_bridge`, and says so if there is none. + * + * Returns the program's printed output followed by its placement accounting, + * so a page can show what crossed the boundary. + * @param {string} source + * @param {string} place + * @returns {string} + */ +export function eval_placed(source, place) { + let deferred4_0; + let deferred4_1; + try { + const ptr0 = passStringToWasm0(source, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc); + const len0 = WASM_VECTOR_LEN; + const ptr1 = passStringToWasm0(place, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc); + const len1 = WASM_VECTOR_LEN; + const ret = wasm.eval_placed(ptr0, len0, ptr1, len1); + var ptr3 = ret[0]; + var len3 = ret[1]; + if (ret[3]) { + ptr3 = 0; len3 = 0; + throw takeFromExternrefTable0(ret[2]); + } + deferred4_0 = ptr3; + deferred4_1 = len3; + return getStringFromWasm0(ptr3, len3); + } finally { + wasm.__wbindgen_free(deferred4_0, deferred4_1, 1); + } +} + /** * Evaluate a Loon program and return the result as a string. * @param {string} source @@ -145,6 +186,15 @@ export function get_effect_log() { } } +/** + * Whether a GPU bridge has been installed. + * @returns {boolean} + */ +export function has_gpu_bridge() { + const ret = wasm.has_gpu_bridge(); + return ret !== 0; +} + /** * Infer the type of the last named binding (fn/let) in the source. * @param {string} source @@ -181,6 +231,23 @@ export function init_dom_bridge(bridge) { wasm.init_dom_bridge(bridge); } +/** + * Install the GPU bridge: a JS function taking `(op, payload)` and returning + * synchronously. + * + * The operations are `name`, `upload`, `dispatch`, `download`, and `evict`. + * How the JS side becomes synchronous is up to it — the demo runs the VM in a + * worker and blocks on `Atomics.wait` while the main thread drives WebGPU. + * + * With a bridge installed, `--place gpu` works in a browser. Without one it + * says there is nowhere to run, which is better than quietly running on the + * CPU and reporting GPU numbers. + * @param {Function} bridge + */ +export function init_gpu_bridge(bridge) { + wasm.init_gpu_bridge(bridge); +} + /** * Invoke a stored Loon callback by ID (called from JS event handlers). * @param {number} id @@ -189,6 +256,42 @@ export function invoke_callback(id) { wasm.invoke_callback(id); } +/** + * Finish a parked step by supplying the numbers the host went to fetch. + * @param {Float32Array} values + * @returns {any} + */ +export function place_resume(values) { + const ptr0 = passArrayF32ToWasm0(values, wasm.__wbindgen_malloc); + const len0 = WASM_VECTOR_LEN; + const ret = wasm.place_resume(ptr0, len0); + if (ret[2]) { + throw takeFromExternrefTable0(ret[1]); + } + return takeFromExternrefTable0(ret[0]); +} + +/** + * Start a program, running until it finishes or parks. + * + * Returns `{done, output, request}`. `done` false means it parked and is + * waiting for `place_resume`. + * @param {string} source + * @param {string} place + * @returns {any} + */ +export function place_start(source, place) { + const ptr0 = passStringToWasm0(source, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc); + const len0 = WASM_VECTOR_LEN; + const ptr1 = passStringToWasm0(place, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc); + const len1 = WASM_VECTOR_LEN; + const ret = wasm.place_start(ptr0, len0, ptr1, len1); + if (ret[2]) { + throw takeFromExternrefTable0(ret[1]); + } + return takeFromExternrefTable0(ret[0]); +} + /** * Reset the Loon runtime state (callbacks, etc.) for hot reload. */ @@ -244,6 +347,10 @@ function __wbg_get_imports() { const ret = arg0.call(arg1, arg2, arg3); return ret; }, arguments); }, + __wbg_get_b3ed3ad4be2bc8ac: function() { return handleError(function (arg0, arg1) { + const ret = Reflect.get(arg0, arg1); + return ret; + }, arguments); }, __wbg_instanceof_Window_ed49b2db8df90359: function(arg0) { let result; try { @@ -254,6 +361,14 @@ function __wbg_get_imports() { const ret = result; return ret; }, + __wbg_length_32ed9a279acd054c: function(arg0) { + const ret = arg0.length; + return ret; + }, + __wbg_length_9a7876c9728a0979: function(arg0) { + const ret = arg0.length; + return ret; + }, __wbg_new_361308b2356cecd0: function() { const ret = new Object(); return ret; @@ -262,10 +377,22 @@ function __wbg_get_imports() { const ret = new Array(); return ret; }, + __wbg_new_dd2b680c8bf6ae29: function(arg0) { + const ret = new Uint8Array(arg0); + return ret; + }, __wbg_new_no_args_1c7c842f08d00ebb: function(arg0, arg1) { const ret = new Function(getStringFromWasm0(arg0, arg1)); return ret; }, + __wbg_new_with_length_63f2683cc2521026: function(arg0) { + const ret = new Float32Array(arg0 >>> 0); + return ret; + }, + __wbg_new_with_length_a2c39cbe88fd8ff1: function(arg0) { + const ret = new Uint8Array(arg0 >>> 0); + return ret; + }, __wbg_now_a3af9a2f4bbaa4d1: function() { const ret = Date.now(); return ret; @@ -278,6 +405,9 @@ function __wbg_get_imports() { const ret = arg0.performance; return isLikeNone(ret) ? 0 : addToExternrefTable0(ret); }, + __wbg_prototypesetcall_bdcdcc5842e4d77d: function(arg0, arg1, arg2) { + Uint8Array.prototype.set.call(getArrayU8FromWasm0(arg0, arg1), arg2); + }, __wbg_push_8ffdcb2063340ba5: function(arg0, arg1) { const ret = arg0.push(arg1); return ret; @@ -286,6 +416,12 @@ function __wbg_get_imports() { const ret = Reflect.set(arg0, arg1, arg2); return ret; }, arguments); }, + __wbg_set_cc56eefd2dd91957: function(arg0, arg1, arg2) { + arg0.set(getArrayU8FromWasm0(arg1, arg2)); + }, + __wbg_set_f8edeec46569cc70: function(arg0, arg1, arg2) { + arg0.set(getArrayF32FromWasm0(arg1, arg2)); + }, __wbg_static_accessor_GLOBAL_12837167ad935116: function() { const ret = typeof global === 'undefined' ? null : global; return isLikeNone(ret) ? 0 : addToExternrefTable0(ret); @@ -402,6 +538,16 @@ function debugString(val) { return className; } +function getArrayF32FromWasm0(ptr, len) { + ptr = ptr >>> 0; + return getFloat32ArrayMemory0().subarray(ptr / 4, ptr / 4 + len); +} + +function getArrayU8FromWasm0(ptr, len) { + ptr = ptr >>> 0; + return getUint8ArrayMemory0().subarray(ptr / 1, ptr / 1 + len); +} + let cachedDataViewMemory0 = null; function getDataViewMemory0() { if (cachedDataViewMemory0 === null || cachedDataViewMemory0.buffer.detached === true || (cachedDataViewMemory0.buffer.detached === undefined && cachedDataViewMemory0.buffer !== wasm.memory.buffer)) { @@ -410,6 +556,14 @@ function getDataViewMemory0() { return cachedDataViewMemory0; } +let cachedFloat32ArrayMemory0 = null; +function getFloat32ArrayMemory0() { + if (cachedFloat32ArrayMemory0 === null || cachedFloat32ArrayMemory0.byteLength === 0) { + cachedFloat32ArrayMemory0 = new Float32Array(wasm.memory.buffer); + } + return cachedFloat32ArrayMemory0; +} + function getStringFromWasm0(ptr, len) { ptr = ptr >>> 0; return decodeText(ptr, len); @@ -436,6 +590,13 @@ function isLikeNone(x) { return x === undefined || x === null; } +function passArrayF32ToWasm0(arg, malloc) { + const ptr = malloc(arg.length * 4, 4) >>> 0; + getFloat32ArrayMemory0().set(arg, ptr / 4); + WASM_VECTOR_LEN = arg.length; + return ptr; +} + function passStringToWasm0(arg, malloc, realloc) { if (realloc === undefined) { const buf = cachedTextEncoder.encode(arg); @@ -513,6 +674,7 @@ function __wbg_finalize_init(instance, module) { wasm = instance.exports; wasmModule = module; cachedDataViewMemory0 = null; + cachedFloat32ArrayMemory0 = null; cachedUint8ArrayMemory0 = null; wasm.__wbindgen_start(); return wasm; diff --git a/web/public/pkg/loon_wasm_bg.wasm b/web/public/pkg/loon_wasm_bg.wasm index ff123f6..23526d2 100644 Binary files a/web/public/pkg/loon_wasm_bg.wasm and b/web/public/pkg/loon_wasm_bg.wasm differ diff --git a/web/public/pkg/loon_wasm_bg.wasm.d.ts b/web/public/pkg/loon_wasm_bg.wasm.d.ts index 36b4ae2..d8e339f 100644 --- a/web/public/pkg/loon_wasm_bg.wasm.d.ts +++ b/web/public/pkg/loon_wasm_bg.wasm.d.ts @@ -3,14 +3,19 @@ export const memory: WebAssembly.Memory; export const check_program: (a: number, b: number) => [number, number, number, number]; export const enable_effect_log: (a: number) => void; +export const eval_placed: (a: number, b: number, c: number, d: number) => [number, number, number, number]; export const eval_program: (a: number, b: number) => [number, number, number, number]; export const eval_ui: (a: number, b: number) => [number, number]; export const eval_ui_checked: (a: number, b: number) => [number, number]; export const eval_with_output: (a: number, b: number) => [number, number, number, number]; export const get_effect_log: () => [number, number]; +export const has_gpu_bridge: () => number; export const infer_type: (a: number, b: number) => [number, number, number, number]; export const init_dom_bridge: (a: any) => void; +export const init_gpu_bridge: (a: any) => void; export const invoke_callback: (a: number) => void; +export const place_resume: (a: number, b: number) => [number, number, number]; +export const place_start: (a: number, b: number, c: number, d: number) => [number, number, number]; export const clear_effect_log: () => void; export const reset_runtime: () => void; export const __wbindgen_malloc: (a: number, b: number) => number; diff --git a/web/public/place-gpu.js b/web/public/place-gpu.js new file mode 100644 index 0000000..ca1db7d --- /dev/null +++ b/web/public/place-gpu.js @@ -0,0 +1,172 @@ +// The main thread's half of the GPU bridge. +// +// Loon's VM is synchronous and WebGPU is not, so the VM runs in a worker and +// blocks on `Atomics.wait` while this file does the asynchronous part. The +// protocol is deliberately small: a control word says whose turn it is, a +// request is passed as a structured-clone message, and the reply is written +// into a SharedArrayBuffer before the worker is woken. +// +// Nothing here knows anything about Loon. It receives WGSL — generated by the +// same emitter a desktop build hands to wgpu — plus buffers and a work count. + +// Control word slots, in Int32 units. +const STATE = 0; // 0 = worker waiting, 1 = reply ready +const STATUS = 1; // 0 = ok, 1 = error +const LENGTH = 2; // reply byte length +const HEADER_WORDS = 4; + +/// Set up a worker that can call into WebGPU synchronously. +/// +/// `replyBytes` bounds the largest single readback; a download larger than +/// this is reported rather than truncated. +export async function attachGpuBridge(worker, { replyBytes = 1 << 24 } = {}) { + if (!navigator.gpu) { + throw new Error('this browser has no WebGPU'); + } + const adapter = await navigator.gpu.requestAdapter(); + if (!adapter) { + throw new Error('no WebGPU adapter available'); + } + const device = await adapter.requestDevice(); + const info = adapter.info ?? {}; + const name = [info.vendor, info.architecture].filter(Boolean).join(' ') || 'webgpu'; + + const control = new SharedArrayBuffer(HEADER_WORDS * 4 + replyBytes); + const words = new Int32Array(control, 0, HEADER_WORDS); + const payload = new Uint8Array(control, HEADER_WORDS * 4); + + // Device-side buffers, keyed by the host's heap slot — the same identity the + // Rust side uses, so residency is something both ends agree about by name. + const buffers = new Map(); + const pipelines = new Map(); + + const encoder = new TextEncoder(); + + function reply(status, bytes) { + if (bytes && bytes.byteLength > payload.byteLength) { + status = 1; + bytes = encoder.encode( + `a readback of ${bytes.byteLength} bytes exceeds the ${payload.byteLength}-byte reply buffer`, + ); + } + const len = bytes ? bytes.byteLength : 0; + if (len) payload.set(bytes, 0); + Atomics.store(words, LENGTH, len); + Atomics.store(words, STATUS, status); + Atomics.store(words, STATE, 1); + Atomics.notify(words, STATE); + } + + function pipelineFor(shader, entry) { + let p = pipelines.get(shader); + if (!p) { + const module = device.createShaderModule({ code: shader }); + p = device.createComputePipeline({ + layout: 'auto', + compute: { module, entryPoint: entry }, + }); + pipelines.set(shader, p); + } + return p; + } + + async function handle(req) { + switch (req.op) { + case 'name': + return encoder.encode(name); + + case 'upload': { + const bytes = req.bytes; + const buf = device.createBuffer({ + size: Math.max(bytes.byteLength, 4), + usage: + GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, + }); + device.queue.writeBuffer(buf, 0, bytes); + buffers.get(req.id)?.destroy?.(); + buffers.set(req.id, buf); + return null; + } + + case 'dispatch': { + // The uniform block is `n` followed by every scalar, matching what the + // WGSL emitter wrote into `struct Params`, padded to 16 bytes. + const scalars = req.scalars; + const uniformWords = 1 + scalars.length; + const padded = Math.ceil(uniformWords / 4) * 4; + const uniform = new ArrayBuffer(padded * 4); + new Int32Array(uniform, 0, 1)[0] = req.n; + if (scalars.length) new Float32Array(uniform, 4, scalars.length).set(scalars); + const uniformBuf = device.createBuffer({ + size: uniform.byteLength, + usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + }); + device.queue.writeBuffer(uniformBuf, 0, uniform); + + const pipeline = pipelineFor(req.shader, req.entry); + const entries = [{ binding: 0, resource: { buffer: uniformBuf } }]; + req.buffers.forEach((id, slot) => { + const buf = buffers.get(id); + if (!buf) throw new Error(`buffer ${id} is not on the device`); + entries.push({ binding: slot + 1, resource: { buffer: buf } }); + }); + + const bindGroup = device.createBindGroup({ + layout: pipeline.getBindGroupLayout(0), + entries, + }); + + const cmd = device.createCommandEncoder(); + const pass = cmd.beginComputePass(); + pass.setPipeline(pipeline); + pass.setBindGroup(0, bindGroup); + // Workgroup size is 64 in the emitted shader; the bounds check inside + // it discards the overshoot from rounding up. + pass.dispatchWorkgroups(Math.max(1, Math.ceil(req.n / 64)), 1, 1); + pass.end(); + device.queue.submit([cmd.finish()]); + return null; + } + + case 'download': { + const src = buffers.get(req.id); + if (!src) throw new Error(`buffer ${req.id} is not on the device`); + const size = Math.max(req.byteLength, 4); + const staging = device.createBuffer({ + size, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + const cmd = device.createCommandEncoder(); + cmd.copyBufferToBuffer(src, 0, staging, 0, size); + device.queue.submit([cmd.finish()]); + await staging.mapAsync(GPUMapMode.READ); + const copy = new Uint8Array(staging.getMappedRange().slice(0, req.byteLength)); + staging.unmap(); + staging.destroy(); + return copy; + } + + case 'evict': { + buffers.get(req.id)?.destroy?.(); + buffers.delete(req.id); + return null; + } + + default: + throw new Error(`unknown device operation '${req.op}'`); + } + } + + worker.addEventListener('message', async (e) => { + if (!e.data || e.data.kind !== 'gpu') return; + try { + const bytes = await handle(e.data); + reply(0, bytes); + } catch (err) { + reply(1, encoder.encode(String(err && err.message ? err.message : err))); + } + }); + + worker.postMessage({ kind: 'gpu-ready', control }); + return { name, control }; +} diff --git a/web/public/place-worker.js b/web/public/place-worker.js new file mode 100644 index 0000000..51f0027 --- /dev/null +++ b/web/public/place-worker.js @@ -0,0 +1,81 @@ +// The worker's half of the GPU bridge. +// +// Loon's VM runs here, synchronously, because that is what it is. When a +// program performs a `Place` operation that needs the GPU, the bridge below +// posts the request to the main thread and blocks on `Atomics.wait` until the +// reply arrives. Blocking is only legal off the main thread, which is the +// entire reason this file exists. + +import init, { eval_placed, init_gpu_bridge } from '/loon_wasm.js'; + +const STATE = 0; +const STATUS = 1; +const LENGTH = 2; +const HEADER_WORDS = 4; + +let words = null; +let payload = null; +const decoder = new TextDecoder(); + +/// Perform one device operation and block until the main thread answers. +function bridge(op, req) { + if (!words) throw new Error('the GPU bridge is not connected'); + + // Structured clone moves the request; typed arrays inside it are copied. + const message = { kind: 'gpu', op }; + if (req) { + if (req.bytes) message.bytes = req.bytes.slice(); + if (req.scalars) message.scalars = Array.from(req.scalars); + if (req.buffers) message.buffers = Array.from(req.buffers); + for (const k of ['id', 'n', 'shader', 'entry', 'dtype', 'byteLength']) { + if (req[k] !== undefined) message[k] = req[k]; + } + } + + Atomics.store(words, STATE, 0); + self.postMessage(message); + + // The main thread is doing something asynchronous; wait for it. + Atomics.wait(words, STATE, 0); + + const status = Atomics.load(words, STATUS); + const len = Atomics.load(words, LENGTH); + const bytes = payload.slice(0, len); + if (status !== 0) { + throw new Error(decoder.decode(bytes) || 'the device reported a failure'); + } + return bytes; +} + +let ready = null; + +self.addEventListener('message', async (e) => { + const data = e.data; + if (!data) return; + + if (data.kind === 'gpu-ready') { + words = new Int32Array(data.control, 0, HEADER_WORDS); + payload = new Uint8Array(data.control, HEADER_WORDS * 4); + return; + } + + if (data.kind === 'run') { + try { + if (!ready) { + ready = init().then(() => { + init_gpu_bridge(bridge); + }); + } + await ready; + const out = eval_placed(data.program, data.place); + self.postMessage({ kind: 'result', id: data.id, ok: true, out }); + } catch (err) { + self.postMessage({ + kind: 'result', + id: data.id, + ok: false, + out: String(err && err.message ? err.message : err), + }); + } + } +}); diff --git a/web/public/place.html b/web/public/place.html new file mode 100644 index 0000000..89642d1 --- /dev/null +++ b/web/public/place.html @@ -0,0 +1,160 @@ + + + + + +Placement in a browser — Loon + + + +

Placement in a browser

+

The same kernel, the same program — running in this tab.

+ +

Below is a Loon kernel and a residency handler. Nothing in the program says +where the work happens; the buttons choose a placement, and the answer does not +change between them.

+ +

+
+
Run it

+

+  
+  
+  
+  
+

+
loading the wasm build…

+

+
+

+  cpu and par run inline on this thread — one memory,
+  so nothing is transferred. device models a separate memory to
+  show transfers counted. gpu is real WebGPU: Loon's VM is
+  synchronous and WebGPU is not, so the VM runs in a worker and blocks on
+  Atomics.wait while this thread drives the device. It needs
+  cross-origin isolation (COOP/COEP) for SharedArrayBuffer; without
+  it the other placements still work and this one says why.
+

+
+
+
+
diff --git a/web/src/pages/blog.loon b/web/src/pages/blog.loon
index 5cc5264..ae0b89c 100644
--- a/web/src/pages/blog.loon
+++ b/web/src/pages/blog.loon
@@ -81,6 +81,71 @@
             :margin-bottom "3.5rem"}}
         "Changelog"]
       
+      ; ── v0.9.0 — Placement is an effect ────────────────────────────────
+      [blog-article
+        "Placement Is an Effect"
+        "Aug 18, 2026"
+
+        [p
+          "A syscall is an effect, so the kernel is the outermost handler. Follow the same sentence one step further and the GPU becomes a handler in the middle. Where a kernel runs is no longer a property of the program — it is a decision somebody makes around it, and the program does not move."]
+
+        [blog-h3 "A kernel, and nothing about where it runs"]
+        [p
+          "A kernel is an ordinary function that promised to stay small: no closures, no allocation, no strings, no effects. Nothing in it mentions a device, a queue, or a transfer."]
+        [blog-code
+          [str
+            "[kernel saxpy [i a x y out]\n"
+            "  [put out i [+ [* a [at x i]] [at y i]]]]\n"
+            "\n"
+            "[Place.run saxpy n #[3.0 x y out]]\n"
+            "[Place.read out]"]]
+        [p
+          "Which arguments travel to the device and which have to come home is inferred, not annotated. Loon already worked out, for every parameter, whether the body reads it, writes through it, or consumes it — the distinction Rust makes you spell as &T, &mut T, and T. Here `at` is a read and `put` is a write-through, so `out` is the one that needs synchronizing back, and the emitted shader binds it read_write while the inputs bind read."]
+
+        [blog-h3 "Residency is nine lines"]
+        [p
+          "Because Place.run and Place.read are operations rather than function calls, a handler sees every launch and every synchronization point. That is enough to write a residency policy as ordinary library code — the transfer-hoisting a GPU compiler implements as an optimization pass."]
+        [blog-code
+          [str
+            "[fn place/resident [thunk]\n"
+            "  [handle [thunk]\n"
+            "    [Place.run k n args]\n"
+            "      [do [Place.pin args]\n"
+            "          [resume [Place.run k n args]]]\n"
+            "    [Place.read b]\n"
+            "      [do [let v [Place.read b]]\n"
+            "          [Place.unpin b]\n"
+            "          [resume v]]]]]"]]
+        [p
+          "On an Apple M4 Max through Metal, a chain of 128 launches over one buffer takes 356.7 ms with no policy and 18.3 ms under that handler — 19.5x, from 128 uploads down to 2. Same program, same answer; only the line wrapped around it differs."]
+
+        [blog-h3 "Testing a GPU program without a GPU"]
+        [p
+          "Launches are events, so they can go on tape. Record a run on the GPU and replay it on a build with no GPU support compiled in at all — the program observes exactly what it observed the first time. A handler that declines to forward gives you a dry run: every accounting number, no execution."]
+        [blog-code
+          [str
+            "$ loon run samples/place/saxpy.oo --place gpu --record trace.oo\n"
+            "out   = #[0 5 10 15 20 25 30 35]\n"
+            "\n"
+            "$ loon replay trace.oo samples/place/saxpy.oo   ; no GPU in this build\n"
+            "out   = #[0 5 10 15 20 25 30 35]"]]
+
+        [blog-h3 "Landing alongside"]
+        [ul-el
+          {:style {:margin-left "1.5rem" :line-height "1.8"}}
+          [li-el
+            "Kernels compile to WGSL and every one in the repo is validated by naga in CI, on machines with no GPU."]
+          [li-el
+            "Placed programs run in a browser, including on the GPU: web/public/place.html picks a placement in a tab and shows what crossed the boundary. WebGPU is asynchronous and the VM is not, so the VM runs in a worker and blocks on Atomics.wait while the main thread drives the device."]
+          [li-el
+            "Reductions, with no new language feature: each work item sums its own chunk into its own slot, so the write-at-your-own-index rule holds throughout (samples/place/reduce.oo)."]
+          [li-el
+            "One shared value layout across the VM, WASM, and native backends, with a conformance test that compares raw bits — it found three real divergences the first time it ran."]
+          [li-el
+            "--place cpu|device|gpu and --place-stats, reporting launches, transfers, and bytes."]
+          [li-el
+            "Buffers: dense, unboxed numeric arrays that can leave the process, which persistent collections cannot."]]]
+
       ; ── Record/replay + capability security ────────────────────────────
       [blog-article
         "Record the Bug, Replay It Forever"
diff --git a/web/vercel.json b/web/vercel.json
index 97e1d60..2d2b634 100644
--- a/web/vercel.json
+++ b/web/vercel.json
@@ -7,12 +7,47 @@
     {
       "source": "/install.sh",
       "headers": [
-        { "key": "Content-Type", "value": "text/plain; charset=utf-8" },
-        { "key": "Cache-Control", "value": "no-cache" }
+        {
+          "key": "Content-Type",
+          "value": "text/plain; charset=utf-8"
+        },
+        {
+          "key": "Cache-Control",
+          "value": "no-cache"
+        }
+      ]
+    },
+    {
+      "source": "/place.html",
+      "headers": [
+        {
+          "key": "Cross-Origin-Opener-Policy",
+          "value": "same-origin"
+        },
+        {
+          "key": "Cross-Origin-Embedder-Policy",
+          "value": "require-corp"
+        }
+      ]
+    },
+    {
+      "source": "/place-worker.js",
+      "headers": [
+        {
+          "key": "Cross-Origin-Opener-Policy",
+          "value": "same-origin"
+        },
+        {
+          "key": "Cross-Origin-Embedder-Policy",
+          "value": "require-corp"
+        }
       ]
     }
   ],
   "rewrites": [
-    { "source": "/(.*)", "destination": "/index.html" }
+    {
+      "source": "/(.*)",
+      "destination": "/index.html"
+    }
   ]
 }