Two inputs, fused reductions, and control flow in the traced DSL

[yujiteshima]

What

Three things the README listed under Not yet:

• Two inputs. a.map2(b) { |x, y| x * y + 1 }, or GPU.kernel { |x, y| ... }.call(a, b). The input count comes from Proc#arity (GPU.kernel(arity: 2) to say it explicitly); two inputs use LAYOUT_3BUF. A size mismatch raises ArgumentError.
• Fused reductions. a.sum { |x| x * 2 + 1 } evaluates the expression and reduces it in shared memory in the same shader; the per-workgroup partials are combined on the host in double, like the base gem's sum. GPU.kernel(reduce: :sum | :min | :max) is the general form, a.dot(b) is a one-dispatch dot product. sum without a block still goes to the C implementation.
• Control flow. Comparisons (> < >= <= == !=) yield boolean nodes, & | ! combine them, where(cond, a, b) (alias select) becomes a GLSL ternary, and iterate(init, count) { |acc, i| ... } becomes a real GLSL for loop — count is a trace-time Integer, loops nest. Ruby's own if / while still cannot be traced (the condition is a proxy, so it is always truthy); that is documented, and it is the same trade-off cupy.where makes. Constant-count Ruby loops unroll at trace time.

iterate(x, 20) { |y, i| (y + x / y) * 0.5 } generates:

float x = a[idx];
float acc1 = x;
for (int i2 = 0; i2 < 20; i2++) {
    acc1 = ((acc1 + (x / acc1)) * 0.5);
}
b[idx] = acc1;

The one-input map shader is byte-identical to main's.

Tests and numbers

58/58 — the 20 existing tests plus 38 new ones (map2, size and arity errors, sum { }, dot, min/max across workgroups, where / comparisons / & | !, Newton sqrt == sqrt(x), nested iterate, unrolled 3.times, type errors for bool arithmetic and non-Integer counts).

examples/bench_reduce.rb, Apple M5, 1,024 elements: a.map { }.sum (2 dispatches) 0.52 ms → a.sum { } (1 dispatch) 0.22 ms.

Notes

• src/gpu_kernel.c's dispatch now takes 2 or 3 buffers and does the reduction combine. The companion deferred-submit PR replaces the dispatch body with mruby-gpu-narray's dispatch_pipeline(); merge this one first.
• examples/kernel_dsl.rb is unchanged (it is a talk script); the new material is in examples/kernel_dsl_2.rb.

🤖 Generated with Claude Code