Skip to content

GPU: portable unified memory, launch caching, shared memory, and a Metal benchmark#303

Open
PhilippGrulich wants to merge 7 commits into
mainfrom
claude/gpu-portable-memory-and-shared-mem
Open

GPU: portable unified memory, launch caching, shared memory, and a Metal benchmark#303
PhilippGrulich wants to merge 7 commits into
mainfrom
claude/gpu-portable-memory-and-shared-mem

Conversation

@PhilippGrulich

Copy link
Copy Markdown
Member

Summary

Makes the Metal GPU backend usable for larger kernels over larger data batches, with all new APIs designed to work for the CUDA backend too. Three subsystems plus a benchmark.

§1 — Portable unified-memory model (no fixed buffer size)

  • gpu::allocUnified / gpu::copy / gpu::wrapgpu::Array<T>; kernels take gpu::Array<T> for pointer args.
  • A small arg_factory hook in Engine.hpp lets a parameter type build itself from its value ref (default unchanged for every existing val<T>).
  • A runtime size table records each allocation's page-rounded length; generated Metal host code wraps the caller's memory with newBufferWithBytesNoCopy (zero copy, exact size) and resolves the length via nautilus_gpu_buffer_bytes at dlopen (-undefined dynamic_lookup, deterministic codegen). CUDA passes the managed pointer straight through.
  • Deletes NAUTILUS_BUFFER_SIZE and the gpu.metal.bufferSize option — the fixed 4096-byte ceiling that silently truncated buffers is gone.

§2 — Metal launch-context caching

  • The generated host dispatch no longer rebuilds MTLDevice/library/pipeline/queue every launch; they live in function-local statics (one-time init). Only the command buffer, encoder, binding, and dispatch run per call.

§3 — Portable static shared memory

  • gpu::sharedArray<T, N>()threadgroup T[N] (Metal) / __shared__ T[N] (CUDA); pair with gpu::syncThreads() for block-cooperative kernels.
  • Metal requires per-pointer address spaces, so the device lowering runs a per-function analysis that seeds from sharedArray results and propagates threadgroup through pointer arithmetic, casts, and block-argument phis to a fixpoint, then retags declarations, phi-copy temporaries, and load/store casts (byte-identical output for kernels without shared memory).

Benchmark

  • New nautilus-gpu-benchmarks (under ENABLE_BENCHMARKS): grid-stride kernels run the same launch function as a parallel GPU dispatch on Metal and a full serial loop on the CPU fallback, giving a fair CPU-vs-Metal comparison; plus a shared-memory reduction and one-time compile cost.
  • On an M-series machine: bandwidth-bound saxpy/vecadd are roughly on par CPU vs GPU (unified memory, ~3 ms dispatch floor), while the compute-bound logistic-map kernel is 6.7×–13× faster on the GPU.

Bugs found and fixed along the way

  • A traceCopy tag collision that aliased distinct buffer args (b read as a) — fixed by reconstructing Array's val<T*> from value+state without traceCopy.
  • Codegen non-determinism from an embedded function-pointer literal — switched to a named symbol resolved at load.
  • A latent Metal codegen bug: a scalar kernel arg (constant T&, const) reassigned by an SSA phi when live inside a conditional — fixed by copying scalar args into mutable locals.

Testing

  • Real Metal GPU: execution 95 assertions (incl. a 4096-element large-batch that previously truncated, and a shared-memory block reduction).
  • tracing 102, codegen 24 (Metal + CUDA, deterministic), core 8209 (the arg_factory hook is a strict no-op).
  • CUDA paths are codegen-verified (no nvcc needed for codegen) but not executed here (no CUDA device).

Design doc: docs/superpowers/specs/2026-06-18-gpu-portable-memory-and-kernels-design.md

🤖 Generated with Claude Code

PhilippGrulich and others added 7 commits June 18, 2026 23:28
… memory

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… size)

Replaces the two hardcoded GPU memory models (Metal's fixed NAUTILUS_BUFFER_SIZE
copy + CUDA's raw managed-pointer pass-through) with one portable unified model:

- gpu::allocUnified/copy/wrap return a gpu::Array<T>, a single-slot newtype over
  val<T*> that kernels and host entries accept for pointer arguments.
- A small arg_factory hook in Engine.hpp lets a parameter type build itself from
  its value ref (default unchanged for every existing val<T>).
- A runtime size table records each allocation's page-rounded length; generated
  Metal host code wraps the caller's memory with newBufferWithBytesNoCopy (zero
  copy, exact size) resolving the length via nautilus_gpu_buffer_bytes at dlopen
  (-undefined dynamic_lookup, deterministic codegen). CUDA passes the managed
  pointer through unchanged.
- Deletes NAUTILUS_BUFFER_SIZE and the gpu.metal.bufferSize option.

Array reconstructs its val<T*> from value+state on copy to avoid traceCopy tag
collisions that would otherwise alias distinct buffer arguments.

Migrates the demo and GPU test kernels to gpu::Array; adds a 4096-element
large-batch test (previously truncated at the 1024-float ceiling). Regenerates
trace/SSA/IR and Metal/CUDA codegen references (cleaner traces, no fixed size).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The generated host dispatch rebuilt MTLDevice/library/pipeline/queue on every
launch — the pipeline-state compile and on-disk metallib reload dominate looped
and multi-kernel batch workloads. Hoist them into function-local statics
(thread-safe one-time init); only the command buffer, encoder, buffer binding
and dispatch run per call. Regenerates Metal host-code references.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds block-shared memory: gpu::sharedArray<T,N>() returns a pointer to a
statically-sized block-shared array, lowering to threadgroup T[N] (Metal) and
__shared__ T[N] (CUDA); combine with syncThreads() for block-cooperative
kernels (tiled reduction/scan). CPU fallback uses a thread-local arena.

Metal requires per-pointer address spaces, so the device lowering now runs an
address-space analysis per function: it seeds from sharedArray results and
propagates threadgroup through pointer arithmetic, casts, and block-argument
phis to a fixpoint, then retags declarations, phi-copy temporaries, and
load/store casts accordingly (replacing the blanket 'device' post-process;
output is byte-identical for kernels without shared memory).

Adds a block-reduction kernel exercised by the Metal execution test (sum of
1..256 = 32896) and Metal/CUDA codegen + tracing references.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…flow

Adds nautilus-gpu-benchmarks (ENABLE_BENCHMARKS): steady-state execution
throughput for saxpy/vecadd/block-reduction over 1M/4M/16M-float unified
buffers, plus the one-time compile cost. On this M-series machine: ~3 ms
dispatch floor scaling to ~7-30 ms at 16M; compilation ~640 ms (external
metal/metallib/cc toolchain), which launch caching amortizes.

The benchmark surfaced a latent Metal codegen bug: a scalar kernel arg passed
as 'constant T&' (const) cannot be reassigned, but when such a value is live
inside a conditional it becomes an SSA phi whose trampoline emits 'var = temp;'
-> MSL compile error. Existing kernels only used scalar args in branch
conditions, never inside a guarded block, so they never hit it. Fix: bind the
buffer to a '<var>_arg' parameter and copy it into a mutable local of the SSA
name. Regenerates the affected metal device references.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rewrites the saxpy/vecadd benchmark kernels as grid-stride loops so the same
launch function runs both ways from one definition: threads stride in parallel
on Metal, and on a non-Metal backend the fallback's gridDim=blockDim=1 makes it
a full serial loop over all N on the CPU. Benchmarks each at 1M/4M/16M floats
on both mlir (optimized native CPU) and metal (GPU), an apples-to-apples
comparison of identical work over the same unified buffers.

Finding on this M-series machine: for these bandwidth-bound kernels the CPU and
GPU are roughly on par (e.g. saxpy 16M ~8.5 ms CPU vs ~7.9 ms GPU) — unified
memory gives both the same bandwidth and the GPU's ~3 ms dispatch floor isn't
amortized by so little compute per element. The GPU advantage shows on
compute-heavier kernels.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds a compute-bound benchmark: per element, iters logistic-map steps
(v = 3.9*v*(1-v)) over 1 read + 1 write, so it is arithmetic-bound and the
chaotic recurrence cannot be optimized away. Grid-stride, so it compares CPU
(serial) vs Metal (parallel) like the other kernels.

This is the regime where the GPU wins: on this M-series machine over 256K
elements the GPU is ~6.7x faster at 256 iters and ~13x at 1024 iters (the
dispatch floor amortizes as arithmetic intensity grows), versus the
roughly-on-par result for the bandwidth-bound saxpy/vecadd.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tracing Benchmark

Details
Benchmark suite Current: a11fdef Previous: d63bd8f Ratio
exec_bc_addOne 28.2694 ns (± 3.03521) 36.9564 ns (± 5.11982) 0.76
exec_mlir_addOne 214.291 ns (± 5.9849) 282.96 ns (± 7.86376) 0.76
exec_cpp_addOne 3.37508 ns (± 0.691129) 3.98915 ns (± 0.469858) 0.85
exec_interpreted_addOne 30.8891 ns (± 1.23967) 39.6342 ns (± 2.2298) 0.78
ir_add 547.449 ns (± 36.9648) 795.982 ns (± 87.5329) 0.69
ir_ifThenElse 1.2072 us (± 88.7315) 1.63138 us (± 141.868) 0.74
ir_deeplyNestedIfElse 2.59962 us (± 148.518) 3.34702 us (± 226.142) 0.78
ir_loop 1.21298 us (± 74.3156) 1.64605 us (± 113.623) 0.74
ir_ifInsideLoop 2.2362 us (± 221.049) 2.88247 us (± 245.938) 0.78
ir_loopDirectCall 1.39689 us (± 100.841) 1.85863 us (± 162.199) 0.75
ir_pointerLoop 1.4939 us (± 89.6933) 2.0172 us (± 179.225) 0.74
ir_staticLoop 1.13289 us (± 84.252) 1.4698 us (± 172.512) 0.77
ir_fibonacci 1.32318 us (± 121.705) 1.74373 us (± 101.681) 0.76
ir_gcd 1.13213 us (± 82.6718) 1.5363 us (± 98.0465) 0.74
ir_nestedIf10 6.0616 us (± 490.537) 7.53336 us (± 572.678) 0.80
ir_nestedIf100 72.8324 us (± 3.41112) 86.6056 us (± 6.84222) 0.84
ir_chainedIf10 9.4055 us (± 683.313) 11.706 us (± 1.05245) 0.80
ir_chainedIf100 132.717 us (± 3.51767) 162.563 us (± 10.9684) 0.82
trace_add 1.87507 us (± 235.593) 2.34496 us (± 232.961) 0.80
completing_trace_add 1.82703 us (± 188.356) 2.44207 us (± 308.893) 0.75
trace_ifThenElse 6.50987 us (± 1.23348) 9.02627 us (± 1.38021) 0.72
completing_trace_ifThenElse 3.67544 us (± 583.637) 4.71302 us (± 618.171) 0.78
trace_deeplyNestedIfElse 19.4309 us (± 3.24688) 26.141 us (± 3.38505) 0.74
completing_trace_deeplyNestedIfElse 10.0746 us (± 1.70474) 13.1745 us (± 1.88355) 0.76
trace_loop 6.46981 us (± 1.35706) 8.63063 us (± 1.1372) 0.75
completing_trace_loop 3.66515 us (± 484.581) 4.67873 us (± 394.54) 0.78
trace_ifInsideLoop 12.5641 us (± 3.16367) 16.9711 us (± 2.80042) 0.74
completing_trace_ifInsideLoop 6.72638 us (± 1.29262) 8.5433 us (± 1.04142) 0.79
trace_loopDirectCall 6.54736 us (± 1.627) 8.80748 us (± 968.435) 0.74
completing_trace_loopDirectCall 3.70674 us (± 540.194) 6.22222 us (± 2.08068) 0.60
trace_pointerLoop 10.9279 us (± 2.39788) 14.2787 us (± 2.09311) 0.77
completing_trace_pointerLoop 9.1247 us (± 2.82042) 10.4844 us (± 1.6894) 0.87
trace_staticLoop 5.93182 us (± 734.067) 7.36281 us (± 790.484) 0.81
completing_trace_staticLoop 5.89582 us (± 625.456) 7.34517 us (± 621.764) 0.80
trace_fibonacci 7.63102 us (± 1.67065) 10.1657 us (± 1.81837) 0.75
completing_trace_fibonacci 4.72481 us (± 657.618) 6.22191 us (± 551.559) 0.76
trace_gcd 5.90647 us (± 968.195) 8.12246 us (± 1.16253) 0.73
completing_trace_gcd 3.04341 us (± 410.607) 4.14409 us (± 578.076) 0.73
trace_nestedIf10 31.5552 us (± 7.69674) 40.2189 us (± 5.18296) 0.78
completing_trace_nestedIf10 31.6524 us (± 7.86377) 38.9548 us (± 4.13451) 0.81
trace_nestedIf100 1.15108 ms (± 25.8208) 1.38063 ms (± 47.1346) 0.83
completing_trace_nestedIf100 1.17289 ms (± 31.991) 1.39143 ms (± 27.2891) 0.84
trace_chainedIf10 74.7706 us (± 10.6303) 99.3271 us (± 8.40838) 0.75
completing_trace_chainedIf10 40.5501 us (± 8.51752) 51.5743 us (± 8.42379) 0.79
trace_chainedIf100 3.44566 ms (± 63.3545) 4.43705 ms (± 51.9703) 0.78
completing_trace_chainedIf100 1.84358 ms (± 117.111) 2.31297 ms (± 41.3848) 0.80
ssa_add 131.144 ns (± 4.37471) 189.998 ns (± 20.8108) 0.69
ssa_ifThenElse 322.532 ns (± 18.5345) 498.381 ns (± 36.2383) 0.65
ssa_deeplyNestedIfElse 0.856016 us (± 0.0786729) 1.2078 us (± 82.3775) 0.71
ssa_loop 343.834 ns (± 22.0639) 536.005 ns (± 34.4901) 0.64
ssa_ifInsideLoop 669.857 ns (± 51.7115) 949.222 ns (± 72.7998) 0.71
ssa_loopDirectCall 353.498 ns (± 24.189) 545.93 ns (± 41.961) 0.65
ssa_pointerLoop 428.485 ns (± 23.0651) 635.604 ns (± 42.4872) 0.67
ssa_staticLoop 333.899 ns (± 18.1036) 511.315 ns (± 49.3922) 0.65
ssa_fibonacci 357.915 ns (± 22.0075) 565.925 ns (± 55.8251) 0.63
ssa_gcd 320.27 ns (± 19.7429) 498.53 ns (± 44.5426) 0.64
tiered_compile_addOne 36.3155 us (± 15.4992) 54.1272 us (± 7.73487) 0.67
single_compile_mlir_addOne 2.91712 ms (± 119.534) 3.22997 ms (± 116.608) 0.90
single_compile_cpp_addOne 28.1675 ms (± 21.8743) 24.7532 ms (± 459.823) 1.14
single_compile_bc_addOne 37.2243 us (± 15.798) 56.0543 us (± 12.3518) 0.66
tiered_compile_sumLoop 49.5577 us (± 15.8884) 77.2887 us (± 16.7111) 0.64
single_compile_mlir_sumLoop 4.5666 ms (± 300.573) 5.19219 ms (± 79.0478) 0.88
single_compile_cpp_sumLoop 29.6654 ms (± 26.83) 25.472 ms (± 447.056) 1.16
single_compile_bc_sumLoop 49.8412 us (± 16.1178) 76.3645 us (± 15.2721) 0.65
comp_mlir_add 5.02455 ms (± 200.625) 5.59561 ms (± 178.594) 0.90
comp_mlir_ifThenElse 5.46357 ms (± 174.019) 6.2488 ms (± 412.914) 0.87
comp_mlir_deeplyNestedIfElse 4.56082 ms (± 143.514) 5.04225 ms (± 200.911) 0.90
comp_mlir_loop 6.27033 ms (± 182.527) 7.12491 ms (± 237.471) 0.88
comp_mlir_ifInsideLoop 25.068 ms (± 959.518) 28.6722 ms (± 476.575) 0.87
comp_mlir_loopDirectCall 10.2086 ms (± 908.468) 11.808 ms (± 265.45) 0.86
comp_mlir_pointerLoop 23.7211 ms (± 657.758) 27.8439 ms (± 674.969) 0.85
comp_mlir_staticLoop 4.51181 ms (± 146.301) 4.99347 ms (± 156.303) 0.90
comp_mlir_fibonacci 9.06594 ms (± 318.346) 10.3822 ms (± 196.241) 0.87
comp_mlir_gcd 8.12359 ms (± 632.839) 9.44927 ms (± 264.935) 0.86
comp_mlir_nestedIf10 8.80645 ms (± 151.676) 10.4284 ms (± 206.842) 0.84
comp_mlir_nestedIf100 20.6257 ms (± 293.222) 24.8611 ms (± 502.505) 0.83
comp_mlir_chainedIf10 9.84736 ms (± 164.922) 9.5372 ms (± 230.498) 1.03
comp_mlir_chainedIf100 48.0965 ms (± 354.931) 20.5244 ms (± 490.216) 2.34
comp_cpp_add 28.8232 ms (± 28.6193)
comp_cpp_ifThenElse 34.6031 ms (± 35.4673)
comp_cpp_deeplyNestedIfElse 23.0601 ms (± 4.70782)
comp_cpp_loop 28.2878 ms (± 17.9001)
comp_cpp_ifInsideLoop 34.6707 ms (± 34.0731)
comp_cpp_loopDirectCall 22.2504 ms (± 3.77396)
comp_cpp_pointerLoop 35.1556 ms (± 34.9575)
comp_cpp_staticLoop 26.0297 ms (± 18.5796)
comp_cpp_fibonacci 36.6623 ms (± 38.5356)
comp_cpp_gcd 26.6215 ms (± 13.7338)
comp_cpp_nestedIf10 24.3381 ms (± 6.67863)
comp_cpp_nestedIf100 50.9686 ms (± 8.39991)
comp_cpp_chainedIf10 30.8227 ms (± 19.3445)
comp_cpp_chainedIf100 74.4901 ms (± 6.30909)
comp_bc_add 11.118 us (± 2.28922)
comp_bc_ifThenElse 18.1199 us (± 4.81205)
comp_bc_deeplyNestedIfElse 17.0555 us (± 3.31096)
comp_bc_loop 12.9305 us (± 3.11371)
comp_bc_ifInsideLoop 14.7405 us (± 3.0333)
comp_bc_loopDirectCall 13.2449 us (± 2.14815)
comp_bc_pointerLoop 13.8977 us (± 3.05529)
comp_bc_staticLoop 12.8765 us (± 3.01276)
comp_bc_fibonacci 13.1778 us (± 3.09928)
comp_bc_gcd 12.861 us (± 2.98057)
comp_bc_nestedIf10 24.1182 us (± 4.36238)
comp_bc_nestedIf100 152.157 us (± 22.2457)
comp_bc_chainedIf10 33.3486 us (± 6.6259)
comp_bc_chainedIf100 239.194 us (± 9.46087)
comp_asmjit_add 13.6439 us (± 4.25452)
comp_asmjit_ifThenElse 21.0181 us (± 5.46818)
comp_asmjit_deeplyNestedIfElse 37.7607 us (± 8.75816)
comp_asmjit_loop 23.2627 us (± 5.38253)
comp_asmjit_ifInsideLoop 38.8383 us (± 9.95966)
comp_asmjit_loopDirectCall 25.5455 us (± 5.11012)
comp_asmjit_pointerLoop 27.8923 us (± 7.2225)
comp_asmjit_staticLoop 19.019 us (± 4.5078)
comp_asmjit_fibonacci 24.8794 us (± 6.0319)
comp_asmjit_gcd 22.7344 us (± 4.93875)
comp_asmjit_nestedIf10 76.1474 us (± 12.6064)
comp_asmjit_nestedIf100 816.402 us (± 20.0323)
comp_asmjit_chainedIf10 113.233 us (± 11.8286)
comp_asmjit_chainedIf100 1.71431 ms (± 19.0418)
exec_mlir_add 8.54981 ns (± 1.50989) 10.5656 ns (± 1.08602) 0.81
exec_mlir_fibonacci 12.4245 us (± 877.178) 13.0323 us (± 1.09201) 0.95
exec_mlir_sum 500.677 us (± 55.2445) 544.744 us (± 29.5581) 0.92
exec_cpp_add 3.45914 ns (± 0.268001) 4.75823 ns (± 1.02643) 0.73
exec_cpp_fibonacci 83.8254 us (± 5.06277) 96.1834 us (± 7.67007) 0.87
exec_cpp_sum 18.2544 ms (± 118.48) 35.9165 ms (± 549.897) 0.51
exec_bc_add 33.8129 ns (± 2.45011) 43.6174 ns (± 3.93352) 0.78
exec_bc_fibonacci 533.115 us (± 7.38852) 842.262 us (± 14.6154) 0.63
exec_bc_sum 126.731 ms (± 15.8683) 179.187 ms (± 955.653) 0.71
exec_asmjit_add 2.73003 ns (± 0.131805) 3.60782 ns (± 0.497292) 0.76
exec_asmjit_fibonacci 17.3051 us (± 2.2021) 21.5143 us (± 4.60828) 0.80
exec_asmjit_sum 4.09924 ms (± 11.9848) 4.85738 ms (± 33.8584) 0.84
exec_bc_add_noRegAlloc 34.0619 ns (± 3.90621) 45.4957 ns (± 7.38134) 0.75
exec_bc_add_regAlloc 34.4084 ns (± 5.54983) 43.8958 ns (± 4.00479) 0.78
exec_bc_fibonacci_noRegAlloc 532.86 us (± 8.92711) 844.035 us (± 14.8666) 0.63
exec_bc_fibonacci_regAlloc 528.961 us (± 7.74434) 842.01 us (± 14.1375) 0.63
exec_bc_sum_noRegAlloc 125.419 ms (± 8.08883) 178.948 ms (± 362.678) 0.70
exec_bc_sum_regAlloc 123.625 ms (± 9.01236) 179.176 ms (± 555.823) 0.69
e2e_tiered_bc_to_mlir 35.8279 us (± 13.4359) 55.2826 us (± 11.295) 0.65
e2e_single_mlir 4.95462 ms (± 189.158) 5.51468 ms (± 173.212) 0.90

This comment was automatically generated by workflow using github-action-benchmark.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant