Dispatch through mruby-gpu-narray's dispatch_pipeline()

README.md

@@ -26,6 +26,17 @@ Written as a block, the same expression becomes one shader, one submit, one wait
 Adding more operators to the block does not make it slower — the cost was the round
 trips, not the arithmetic.
 
+Since `mruby-gpu-narray` gained deferred submission (operators are recorded and submitted
+together when a result is read), the operator chain also waits only once — but it still
+makes four passes over memory. Fusion removes those too, which is where it keeps winning
+at scale:
+
+| Apple M5, 100-run mean | 1,024 elements | 1,048,576 elements |
+|---|---|---|
+| `a * 2 + 1 - 3 + 4`, one submit per operator (`GPU.sync_mode = :eager`) | 0.95 ms | 2.33 ms |
+| same chain, one submit for the batch (default) | 0.25 ms | 2.12 ms |
+| one traced kernel | 0.22 ms | 0.35 ms |
+
 ## How — tracing, not parsing
 
 mruby has no `RubyVM::AbstractSyntaxTree`, and driving the mruby parser from C to read a
@@ -69,20 +80,67 @@ void main() {
 ## API
 
 ```ruby
-a.map { |x| x * 2 + 1 }        # trace, compile, run -> a new GPU::SFloat
-
-k = GPU.kernel { |x| x * 2 }   # a reusable kernel; identical source compiles once
-k.call(a)
-k.glsl                         # the exact GLSL that was compiled
-
-GPU.kernel_source { |x| x }    # trace only, no compilation (tests, slides)
+a.map { |x| x * 2 + 1 }              # trace, compile, run -> a new GPU::SFloat
+a.map2(b) { |x, y| x * y + 1 }       # two inputs, one dispatch
+a.sum { |x| x * 2 + 1 }              # expression + reduction in one dispatch -> Float
+a.dot(b)                             # inner product, one dispatch -> Float
+
+k = GPU.kernel { |x| x * 2 }         # a reusable kernel; identical source compiles once
+k = GPU.kernel { |x, y| x - y }      # two inputs, from the block's arity (or arity: 2)
+k = GPU.kernel(reduce: :sum) { |x, y| x * y }   # :sum, :min or :max; #call returns a Float
+k.call(a)                            # or k.call(a, b)
+k.glsl                               # the exact GLSL that was compiled
+k.inputs                             # 1 or 2
+k.reduce                             # nil, :sum, :min or :max
+
+GPU.kernel_source { |x| x }          # trace only, no compilation (tests, slides)
 ```
 
 Inside a block you have GLSL's element-wise built-ins under their GLSL names: `sin`,
 `cos`, `tan`, `asin`, `acos`, `atan`, `sinh`, `cosh`, `tanh`, `exp`, `exp2`, `log`,
 `log2`, `sqrt`, `inversesqrt`, `abs`, `sign`, `floor`, `ceil`, `fract`, `radians`,
 `degrees`, `min`, `max`, `pow`, `mod`, `step`, `clamp`, `mix`, `smoothstep`.
 
+### Control flow: where and iterate
+
+Ruby's own `if`, `while`, `&&` and `||` cannot be traced. The block runs once with a
+proxy, and a proxy is always truthy, so `if x > 0` would take the same branch for every
+element and `while` would never end. The shader gets its control flow from two
+functions instead:
+
+```ruby
+a.map { |x| where(x > 0, x, 0) }                     # branch: GLSL's `cond ? a : b`
+a.map { |x| where((x > -2) & (x < 2), x, 0) }        # combine conditions with & | !
+a.map { |x| iterate(x, 20) { |g, i| (g + x / g) * 0.5 } }   # loop: a real GLSL `for`
+```
+
+- Comparisons (`> < >= <= == !=`) produce a *condition*, which only `where`, `&`, `|` and
+  `!` accept. Using one in arithmetic raises `TypeError` before glslang ever sees it.
+- `&` and `|` bind tighter than comparisons in Ruby, so parenthesise: `(x > 1) & (x < 3)`.
+- `iterate(init, count) { |acc, i| ... }` runs the block `count` times with `acc`
+  carrying the value across rounds; `i` is the round number as a float. `count` must be
+  an Integer when the block is traced -- a loop bound that depends on the data is not
+  expressible. Loops nest.
+- A Ruby loop with a count known at trace time -- `3.times { y = y * 2 }` -- needs no
+  special support: it is simply unrolled into the expression.
+
+```glsl
+    float x = a[idx];
+    float acc1 = x;
+    for (int i2 = 0; i2 < 20; i2++) {
+        acc1 = ((acc1 + (x / acc1)) * 0.5);
+    }
+    b[idx] = acc1;
+```
+
+### Reductions
+
+`a.sum { |x| ... }` and `GPU.kernel(reduce: :sum | :min | :max)` evaluate the block per
+element and reduce inside the same shader: a tree reduction in workgroup shared memory,
+one partial per 256-element workgroup, and the partials combined on the host in double
+precision -- the same arrangement as the base gem's `#sum`. Compared with
+`a.map { }.sum` that is one dispatch instead of two and no intermediate array.
+
 ### Two rules
 
 - The block runs via `instance_exec` — that is what lets you write `sin(x)` instead of
@@ -93,9 +151,9 @@ Inside a block you have GLSL's element-wise built-ins under their GLSL names: `s
 
 ### Not yet
 
-One input and one output, element-wise only. No control flow, no reductions, no
-multi-array kernels. A shader that fails to compile raises `GPU::CompileError` carrying
-glslang's own message.
+At most two inputs and one output. No loop whose count depends on the data, no
+indexing into neighbouring elements, no 2-D. A shader that fails to compile raises
+`GPU::CompileError` carrying glslang's own message.
 
 ## Requirements
 
@@ -128,6 +186,7 @@ end
 cd /path/to/mruby && MRUBY_CONFIG=/path/to/mruby-gpu-kernel/build_config.reference.rb rake
 ./build/host/bin/mruby /path/to/mruby-gpu-kernel/test/kernel_test.rb    # ALL TESTS PASSED
 ./build/host/bin/mruby /path/to/mruby-gpu-kernel/examples/kernel_dsl.rb
+./build/host/bin/mruby /path/to/mruby-gpu-kernel/examples/kernel_dsl_2.rb   # map2, sum {}, where, iterate
 ```
 
 To develop against a local checkout of the base gem instead of the published one, list
@@ -136,24 +195,26 @@ it before this gem in `build_config.rb`.
 ## How it fits together
 
 ```
-mrblib/gpu_kernel.rb   Expr / KernelContext / GPU.kernel / NArray#map
-   │                   block -> expression tree -> GLSL
+mrblib/gpu_kernel.rb   Expr / Codegen / KernelContext / GPU.kernel / NArray#map, #map2, #sum, #dot
+   │                   block -> expression tree -> GLSL (statements for loops, an expression for the rest)
    ▼
-src/gpu_kernel.c       GPU::Kernel: glslangValidator -> SPIR-V -> VkPipeline -> dispatch
+src/gpu_kernel.c       GPU::Kernel: glslangValidator -> SPIR-V -> VkPipeline -> dispatch_pipeline (2 or 3 buffers)
    │
    ├─ uses mruby-gpu-narray's g_ctx (device, queue, layouts, pools)
-   └─ uses its create_buffer / wrap_buffer / gpu_buffer_type
+   ├─ uses its create_buffer / wrap_buffer / gpu_buffer_type
+   └─ dispatches through its dispatch_pipeline() (same batch as the operators)
 ```
 
 This gem compiles against the base gem's `src/gpu_internal.h` and links against symbols
-it already exports, so **the base gem needs no changes**. The dispatch routine here is a
-copy of the base gem's `dispatch_compute()`: upstream's version is keyed to its internal
-`PipeId` enum and cannot bind a pipeline built elsewhere. When upstream grows a
-pipeline-taking variant, this copy should go away.
+it already exports. A kernel is dispatched through the base gem's `dispatch_pipeline()`,
+so it is recorded into the same deferred batch as the built-in operators: `a.map { … }`
+queues one dispatch, and nothing is submitted until a result is read. This needs a
+`mruby-gpu-narray` that exports `dispatch_pipeline()` (its `deferred-submit` branch or
+later).
 
 ## Verified on
 
-- **macOS** — Apple M5 GPU via MoltenVK, Vulkan 1.1: 20/20 tests pass.
+- **macOS** — Apple M5 GPU via MoltenVK, Vulkan 1.1: 58/58 tests pass.
 - **Raspberry Pi 5** — not yet run.
 
 ## License

build_config.dsl.rb

@@ -0,0 +1,10 @@
+# Build config for the dsl-control-flow branch (development only).
+#   cd ~/dev/personal/mruby-dsl && MRUBY_CONFIG=~/dev/personal/mruby-gpu-kernel-dsl/build_config.dsl.rb rake -j8
+MRuby::Build.new do |conf|
+  toolchain :clang
+  conf.gembox 'default'
+  conf.gem File.expand_path('~/dev/personal/mruby-gpu-narray-base')   # listed first so mgem-list is not consulted
+  conf.gem File.expand_path('~/dev/personal/mruby-gpu-kernel-dsl')
+  conf.cc.include_paths     << '/opt/homebrew/include'
+  conf.linker.library_paths << '/opt/homebrew/lib'
+end

examples/bench_reduce.rb

@@ -0,0 +1,40 @@
+# How much a fused reduction saves: a.map { }.sum is two dispatches (and two
+# intermediate buffers -- the mapped array and the partials); a.sum { } is one.
+#
+#   mruby examples/bench_reduce.rb
+
+def bench(iters)
+  yield
+  t0 = Time.now.to_f
+  iters.times { yield }
+  ((Time.now.to_f - t0) * 1000.0) / iters
+end
+
+def ms(x)
+  ((x * 1000).round.to_f / 1000).to_s
+end
+
+puts "device: #{GPU.info[:device]}"
+puts
+# Both sides trace their block on every call, so the comparison is fair. The
+# third column reuses a kernel traced once (GPU.kernel outside the loop): the
+# gap between columns two and three is the cost of tracing + cache lookup in
+# mruby, which is why a hot loop should hoist its kernel.
+puts "size        map{}.sum (2 dispatch)   sum{} (1 dispatch)   hoisted kernel   speedup (col 2 / col 3)"
+[1024, 1_048_576].each do |n|
+  a = GPU::SFloat.new(n).seq
+  iters = n > 100_000 ? 30 : 200
+  k = GPU.kernel(reduce: :sum) { |x| x * 2 + 1 }
+  two = bench(iters) { a.map { |x| x * 2 + 1 }.sum }
+  one = bench(iters) { a.sum { |x| x * 2 + 1 } }
+  hoisted = bench(iters) { k.call(a) }
+  puts "#{n.to_s.rjust(9)}   #{ms(two).rjust(10)} ms           #{ms(one).rjust(8)} ms      #{ms(hoisted).rjust(8)} ms     #{ms(two / one)}x / #{ms(two / hoisted)}x"
+end
+puts
+puts "dot product, 1M elements"
+a = GPU::SFloat.new(1_048_576).fill(1.0)
+b = GPU::SFloat.new(1_048_576).fill(2.0)
+two = bench(30) { (a * b).sum }
+one = bench(30) { a.dot(b) }
+puts "  (a * b).sum   #{ms(two)} ms   (2 dispatches)"
+puts "  a.dot(b)      #{ms(one)} ms   (1 dispatch)   #{ms(two / one)}x"

examples/kernel_dsl.rb

@@ -32,8 +32,11 @@
 puts
 
 # ---- 3. why it is worth doing --------------------------------------------
-# Each operator is its own dispatch: a * 2 + 1 - 3 + 4 submits four times and
-# waits for the GPU four times. Traced into one kernel, it submits once.
+# Each operator is its own dispatch. With GPU.sync_mode = :eager that is also
+# a submit and a fence wait per operator; by default mruby-gpu-narray records
+# the four dispatches and submits them together when a result is read. The
+# traced kernel is one dispatch either way -- and one pass over memory instead
+# of four, which is what still matters once the arrays are large.
 def bench(iters)
   yield
   t0 = Time.now.to_f
@@ -48,10 +51,14 @@ def ms(x)
 big = GPU::SFloat.new(1024).seq
 fused = GPU.kernel { |x| x * 2 + 1 - 3 + 4 }
 
-chained_ms = bench(100) { big * 2 + 1 - 3 + 4 }
-fused_ms   = bench(100) { fused.call(big) }
+GPU.sync_mode = :eager
+eager_ms    = bench(100) { big * 2 + 1 - 3 + 4 }
+GPU.sync_mode = :deferred
+deferred_ms = bench(100) { (big * 2 + 1 - 3 + 4); GPU.sync }
+fused_ms    = bench(100) { fused.call(big); GPU.sync }
 
 puts "-- four operators, 1024 elements --"
-puts "  a * 2 + 1 - 3 + 4      #{ms(chained_ms)} ms   (4 dispatches, 4 waits)"
+puts "  a * 2 + 1 - 3 + 4      #{ms(eager_ms)} ms   (4 dispatches, 4 waits -- GPU.sync_mode = :eager)"
+puts "  same, one batch        #{ms(deferred_ms)} ms   (4 dispatches, 1 wait)"
 puts "  one traced kernel      #{ms(fused_ms)} ms   (1 dispatch,  1 wait)"
-puts "  speedup                #{ms(chained_ms / fused_ms)}x"
+puts "  speedup                #{ms(eager_ms / fused_ms)}x over eager, #{ms(deferred_ms / fused_ms)}x over the batch"

examples/kernel_dsl_2.rb

@@ -0,0 +1,42 @@
+# mruby-gpu-kernel -- two inputs, reductions, and control flow.
+#
+# Continues examples/kernel_dsl.rb. Everything below is still one traced block
+# per dispatch: two input arrays, a reduction fused with its expression, a
+# branch (where) and a loop (iterate) inside the shader.
+
+puts "device : #{GPU.info[:device]}"
+puts
+
+# ---- 1. two inputs -------------------------------------------------------
+a = GPU::SFloat[1, 2, 3, 4]
+b = GPU::SFloat[10, 20, 30, 40]
+puts "-- map2: a.map2(b) { |x, y| x * y + 1 } --"
+puts "  #{a.map2(b) { |x, y| x * y + 1 }.inspect}"
+puts
+
+# ---- 2. the expression and the reduction in one dispatch ------------------
+sig = GPU::SFloat.new(1024).seq
+puts "-- sum { |x| x * 2 + 1 } -- one shader: evaluate, then reduce in shared memory --"
+puts "  #{sig.sum { |x| x * 2 + 1 }}"
+puts "-- a.dot(b) --"
+puts "  #{a.dot(b)}"
+puts "-- kernel(reduce: :max) { |x| abs(x) } --"
+puts "  #{GPU.kernel(reduce: :max) { |x| abs(x) }.call(GPU::SFloat[-5, 3, 2])}"
+puts
+
+# ---- 3. a branch: where(cond, a, b) --------------------------------------
+# Ruby's `if` cannot be traced (the proxy is always truthy), so the branch is
+# a function. Conditions combine with & | ! -- parenthesise them.
+v = GPU::SFloat.new(9).seq(-4, 1)          # -4 .. 4
+puts "-- where((x > -2) & (x < 2), x, 0): keep the middle, zero the rest --"
+puts "  #{v.map { |x| where((x > -2) & (x < 2), x, 0) }.inspect}"
+puts
+
+# ---- 4. a loop: iterate(init, count) { |acc, i| ... } ----------------------
+# A real `for` in the shader. Newton's method for sqrt, 20 rounds:
+puts "-- iterate(x, 20) { |g, i| (g + x / g) * 0.5 } -- Newton's sqrt --"
+puts "  #{GPU::SFloat[2, 9, 100].map { |x| iterate(x, 20) { |g, i| (g + x / g) * 0.5 } }.inspect}"
+puts "  sqrt(x) for comparison: #{GPU::SFloat[2, 9, 100].map { |x| sqrt(x) }.inspect}"
+puts
+puts "-- and the shader it became --"
+puts GPU.kernel_source { |x| iterate(x, 20) { |g, i| (g + x / g) * 0.5 } }