From 1aa6f51414d169701d2a934ed203030314caf679 Mon Sep 17 00:00:00 2001 From: yujiteshima Date: Sat, 12 Sep 2026 22:18:12 +0900 Subject: [PATCH 1/2] Two inputs, fused reductions, and control flow in the traced DSL - map2 / GPU.kernel { |x, y| }: a second input array (LAYOUT_3BUF), arity taken from the block or given as arity: 2; size mismatch is ArgumentError - sum { |x| ... }, dot(b), GPU.kernel(reduce: :sum | :min | :max): the expression and a shared-memory tree reduction in one shader, partials combined on the host in double (same arrangement as narray's #sum) - where(cond, a, b), comparisons > < >= <= == != and & | ! on Expr; a condition used as a number (or vice versa) is a TypeError before glslang - iterate(init, count) { |acc, i| } emits a real GLSL for loop via a small statement-emitting Codegen; loops nest; count must be an Integer at trace time. Ruby loops with a known count are simply unrolled (documented) - one-input shaders are byte-identical to before; 58/58 tests Merge note: dispatch_kernel() in src/gpu_kernel.c now takes a layout and 2 or 3 buffers. The deferred-submit branch replaces that function with narray's dispatch_pipeline(); when merging, keep the layout/3-buffer call sites here and drop the local dispatch body in favour of the upstream entry point. Co-Authored-By: Claude Fable 5.1 --- README.md | 76 ++++++-- build_config.dsl.rb | 10 ++ examples/bench_reduce.rb | 40 +++++ examples/kernel_dsl_2.rb | 42 +++++ mrblib/gpu_kernel.rb | 370 +++++++++++++++++++++++++++++++++++---- src/gpu_kernel.c | 198 ++++++++++++++++----- test/kernel_test.rb | 100 +++++++++++ 7 files changed, 741 insertions(+), 95 deletions(-) create mode 100644 build_config.dsl.rb create mode 100644 examples/bench_reduce.rb create mode 100644 examples/kernel_dsl_2.rb diff --git a/README.md b/README.md index 0b4927c..507d78e 100644 --- a/README.md +++ b/README.md @@ -69,13 +69,20 @@ 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`, @@ -83,6 +90,46 @@ Inside a block you have GLSL's element-wise built-ins under their GLSL names: `s `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 +140,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 +175,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,10 +184,10 @@ 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 (2 or 3 buffers) │ ├─ uses mruby-gpu-narray's g_ctx (device, queue, layouts, pools) └─ uses its create_buffer / wrap_buffer / gpu_buffer_type @@ -153,7 +201,7 @@ pipeline-taking variant, this copy should go away. ## 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 diff --git a/build_config.dsl.rb b/build_config.dsl.rb new file mode 100644 index 0000000..1fee6c4 --- /dev/null +++ b/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 diff --git a/examples/bench_reduce.rb b/examples/bench_reduce.rb new file mode 100644 index 0000000..1919736 --- /dev/null +++ b/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" diff --git a/examples/kernel_dsl_2.rb b/examples/kernel_dsl_2.rb new file mode 100644 index 0000000..1034848 --- /dev/null +++ b/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 } } diff --git a/mrblib/gpu_kernel.rb b/mrblib/gpu_kernel.rb index 2960f4b..77a8f6e 100644 --- a/mrblib/gpu_kernel.rb +++ b/mrblib/gpu_kernel.rb @@ -1,6 +1,8 @@ # Turn a Ruby block into a GPU kernel, on top of mruby-gpu-narray. # # na.map { |x| x * 2 + sin(x) } +# na.map2(nb) { |x, y| x * y + 1 } +# na.sum { |x| x * 2 + 1 } # # mruby has no way to read a block's AST, so we trace instead. The block is # called once with a proxy (Expr) in place of the data. Every operator and math @@ -9,24 +11,37 @@ # printed as GLSL, compiled to SPIR-V by glslangValidator, and run as a compute # shader -- see src/gpu_kernel.c for that half. # -# Two consequences worth knowing: +# Consequences worth knowing: # * The block runs on a KernelContext (instance_exec), which is what lets you # write sin(x) rather than x.sin. Methods of the enclosing object are not # visible inside; local variables still are. # * mruby has no numeric coercion, so `2 * x` cannot work -- same rule as the # arrays themselves: keep the expression on the left, as in `x * 2`. +# * Ruby's own `if`, `while`, `&&` and `||` cannot be traced: the proxy is +# always truthy, so `if x > 0` takes the same branch for every element. +# Use where(cond, a, b) for a branch and iterate(init, n) { |acc, i| } for +# a loop; combine conditions with `&`, `|` and `!`. A Ruby loop with a +# count known at trace time (3.times { ... }) is simply unrolled. module GPU # A node in the traced expression tree. class Expr + # Nodes whose GLSL type is bool rather than float. + BOOL_OPS = [:>, :<, :>=, :<=, :==, :!=, :and, :or, :not, :bconst].freeze + def initialize(op, args = []) @op = op @args = args end attr_reader :op, :args + def bool? + BOOL_OPS.include?(@op) + end + def self.wrap(value) return value if value.is_a?(Expr) + return new(:bconst, [value]) if true.equal?(value) || false.equal?(value) unless value.is_a?(Numeric) raise TypeError, "a GPU kernel expression takes Numeric or Expr, got #{value.class}" @@ -34,6 +49,28 @@ def self.wrap(value) new(:const, [value.to_f]) end + # wrap, then insist on a float-typed node. `what` names the operation for + # the error message. + def self.numeric(value, what) + e = wrap(value) + if e.bool? + raise TypeError, + "#{what} needs a number but got a condition; " \ + "use where(cond, a, b) to turn a condition into a number" + end + e + end + + # wrap, then insist on a bool-typed node. + def self.boolean(value, what) + e = wrap(value) + unless e.bool? + raise TypeError, + "#{what} needs a condition (x > 0, x == 1, ...) but got a number" + end + e + end + # GLSL reads a bare 2 as an int, which will not implicitly convert in every # position, so every literal is emitted as a float. def self.literal(f) @@ -43,39 +80,142 @@ def self.literal(f) [:+, :-, :*, :/].each do |operator| define_method(operator) do |other| - Expr.new(operator, [self, Expr.wrap(other)]) + Expr.new(operator, [Expr.numeric(self, operator.to_s), + Expr.numeric(other, operator.to_s)]) end end + # Comparisons produce bool nodes: feed them to where / & / | / !. + [:>, :<, :>=, :<=, :==, :!=].each do |operator| + define_method(operator) do |other| + Expr.new(operator, [Expr.numeric(self, operator.to_s), + Expr.numeric(other, operator.to_s)]) + end + end + + # && and || cannot be overloaded, so the logical operators are the + # bitwise ones. They bind tighter than comparisons in Ruby, so + # parenthesise: (x > 1) & (x < 3). + def &(other) + Expr.new(:and, [Expr.boolean(self, "&"), Expr.boolean(other, "&")]) + end + + def |(other) + Expr.new(:or, [Expr.boolean(self, "|"), Expr.boolean(other, "|")]) + end + + def ! + Expr.new(:not, [Expr.boolean(self, "!")]) + end + def -@ - Expr.new(:neg, [self]) + Expr.new(:neg, [Expr.numeric(self, "-")]) end + # GLSL for a tree without loops. Kernel.source_for handles the general + # case, where a loop needs statements as well as an expression. def to_glsl - case @op - when :var then @args[0] - when :const then Expr.literal(@args[0]) - when :neg then "(-#{@args[0].to_glsl})" - when :+, :-, :*, :/ - "(#{@args[0].to_glsl} #{@op} #{@args[1].to_glsl})" - else - "#{@op}(#{@args.map { |a| a.to_glsl }.join(', ')})" - end + Codegen.new.expr(self) end # S-expression view of the tree, handy when showing what tracing produced. def inspect case @op - when :var then @args[0] - when :const then Expr.literal(@args[0]) + when :var then @args[0] + when :const then Expr.literal(@args[0]) + when :bconst then @args[0].to_s + when :acc then "acc" + when :idx then "i" + when :and then "(& #{@args[0].inspect} #{@args[1].inspect})" + when :or then "(| #{@args[0].inspect} #{@args[1].inspect})" + when :not then "(! #{@args[0].inspect})" + when :loop + init, count, _acc, _idx, body = @args + "(iterate #{init.inspect} #{count} #{body.inspect})" else "(#{@op} #{@args.map { |a| a.inspect }.join(' ')})" end end alias to_s inspect end + # Expression tree -> GLSL. An expression may need statements ahead of it (a + # loop is a `for` in GLSL, not an expression), so the generator returns the + # expression string and collects those statements in #lines, indented for + # the block they belong in. + class Codegen + attr_reader :lines + + def initialize(depth = 1) + @lines = [] + @depth = depth + @tmp = 0 + @scopes = [{}] # per open block: loop node -> its accumulator name + @names = {} # acc / idx node -> GLSL name, set when its loop is emitted + end + + def expr(e) + case e.op + when :var then e.args[0] + when :const then Expr.literal(e.args[0]) + when :bconst then e.args[0] ? "true" : "false" + when :acc, :idx + @names[e.object_id] || + raise(ArgumentError, "a loop variable was used outside its iterate block") + when :neg then "(-#{expr(e.args[0])})" + when :not then "(!#{expr(e.args[0])})" + when :and then "(#{expr(e.args[0])} && #{expr(e.args[1])})" + when :or then "(#{expr(e.args[0])} || #{expr(e.args[1])})" + when :+, :-, :*, :/, :>, :<, :>=, :<=, :==, :!= + "(#{expr(e.args[0])} #{e.op} #{expr(e.args[1])})" + when :where + "(#{expr(e.args[0])} ? #{expr(e.args[1])} : #{expr(e.args[2])})" + when :loop then loop(e) + else "#{e.op}(#{e.args.map { |a| expr(a) }.join(', ')})" + end + end + + private + + def emit(line) + @lines << (" " * @depth) + line + end + + def fresh(prefix) + @tmp += 1 + "#{prefix}#{@tmp}" + end + + # iterate(init, count) { |acc, i| body } becomes + # float accN = ; + # for (int iN = 0; iN < count; iN++) { accN = ; } + # and the expression is accN. A node reached twice inside one block is + # emitted once; reached again in a different block it is emitted afresh + # there, because the first accumulator is out of scope. + def loop(e) + @scopes.reverse_each { |s| return s[e.object_id] if s[e.object_id] } + init, count, acc, idx, body = e.args + acc_name = fresh("acc") + i_name = fresh("i") + init_s = expr(init) + emit "float #{acc_name} = #{init_s};" + @names[acc.object_id] = acc_name + @names[idx.object_id] = "float(#{i_name})" + emit "for (int #{i_name} = 0; #{i_name} < #{count}; #{i_name}++) {" + @depth += 1 + @scopes.push({}) + body_s = expr(body) + emit "#{acc_name} = #{body_s};" + @scopes.pop + @depth -= 1 + emit "}" + @scopes.last[e.object_id] = acc_name + acc_name + end + end + # The block is instance_exec'd here, so these become the vocabulary available - # inside a kernel. Every name below is a GLSL built-in of the same name. + # inside a kernel. Every name below is a GLSL built-in of the same name, + # plus where / iterate for control flow. class KernelContext UNARY = [:sin, :cos, :tan, :asin, :acos, :atan, :sinh, :cosh, :tanh, :exp, :exp2, :log, :log2, :sqrt, :inversesqrt, :abs, :sign, @@ -84,33 +224,90 @@ class KernelContext TERNARY = [:clamp, :mix, :smoothstep].freeze UNARY.each do |f| - define_method(f) { |a| Expr.new(f, [Expr.wrap(a)]) } + define_method(f) { |a| Expr.new(f, [Expr.numeric(a, f.to_s)]) } end BINARY.each do |f| - define_method(f) { |a, b| Expr.new(f, [Expr.wrap(a), Expr.wrap(b)]) } + define_method(f) do |a, b| + Expr.new(f, [Expr.numeric(a, f.to_s), Expr.numeric(b, f.to_s)]) + end end TERNARY.each do |f| define_method(f) do |a, b, c| - Expr.new(f, [Expr.wrap(a), Expr.wrap(b), Expr.wrap(c)]) + Expr.new(f, [Expr.numeric(a, f.to_s), Expr.numeric(b, f.to_s), + Expr.numeric(c, f.to_s)]) end end + + # where(cond, a, b) -> a where cond holds, else b. GLSL's `cond ? a : b`. + # This is the branch: Ruby's `if` cannot be traced. + def where(cond, a, b) + Expr.new(:where, [Expr.boolean(cond, "where"), + Expr.numeric(a, "where"), Expr.numeric(b, "where")]) + end + alias select where + + # iterate(init, count) { |acc, i| ... } -> the value of acc after `count` + # rounds of acc = block(acc, i). A real GLSL `for` loop; `count` must be + # an Integer at trace time, `i` is the round as a float (0, 1, ...). + # Nests freely. This is the loop: Ruby's `while` cannot be traced. + def iterate(init, count, &block) + raise ArgumentError, "iterate requires a block" unless block + unless count.is_a?(Integer) + raise TypeError, + "iterate needs an Integer count known when the block is traced, " \ + "got #{count.class}" + end + raise ArgumentError, "iterate count must be >= 0, got #{count}" if count < 0 + acc = Expr.new(:acc, []) + idx = Expr.new(:idx, []) + body = Expr.numeric(instance_exec(acc, idx, &block), "iterate body") + Expr.new(:loop, [Expr.numeric(init, "iterate init"), count, acc, idx, body]) + end end class Kernel + # Reductions: GLSL identity, GLSL combine step, and the code src/gpu_kernel.c + # uses to combine the per-workgroup partials on the host. + REDUCE = { + sum: ["0.0", "sdata[tid] + sdata[tid + s]", 1], + min: ["uintBitsToFloat(0x7F800000u)", "min(sdata[tid], sdata[tid + s])", 2], + max: ["uintBitsToFloat(0xFF800000u)", "max(sdata[tid], sdata[tid + s])", 3] + }.freeze + # Set by GPU::Kernel.compile in C: the exact GLSL that was compiled. # Showing it is half the fun. def glsl @glsl end - def self.source_for(body) + # 1 or 2 input arrays. + def inputs + @inputs + end + + # nil for an element-wise kernel, else :sum / :min / :max. + def reduce + @reduce + end + + def self.source_for(tree, inputs = 1, reduce = nil) + reduce ? reduce_source(tree, inputs, reduce) : map_source(tree, inputs) + end + + # Element-wise: one thread per element, out[idx] = expr. + def self.map_source(tree, inputs) + gen = Codegen.new(1) + result = gen.expr(tree) + prelude = gen.lines.map { |l| "#{l}\n" }.join + loads = " float x = a[idx];\n" + loads += " float y = b[idx];\n" if inputs == 2 + out = inputs == 2 ? "c" : "b" <<~GLSL #version 450 layout(local_size_x = 256) in; - layout(set = 0, binding = 0) buffer BufA { float a[]; }; - layout(set = 0, binding = 1) buffer BufB { float b[]; }; + #{bindings(inputs, out)} layout(push_constant) uniform PushConstants { uint n; } pc; @@ -118,44 +315,151 @@ def self.source_for(body) void main() { uint idx = gl_GlobalInvocationID.x; if (idx >= pc.n) { return; } - float x = a[idx]; - b[idx] = #{body}; + #{loads}#{prelude} #{out}[idx] = #{result}; + } + GLSL + end + + # Reduction: expr per element, then a tree reduction in shared memory, one + # partial per workgroup. The host combines the partials (gpu_kernel.c). + # The bounds check is an `if`, not a `return`: every thread must reach + # the barriers. + def self.reduce_source(tree, inputs, reduce) + identity, combine, _code = REDUCE[reduce] + gen = Codegen.new(2) + result = gen.expr(tree) + prelude = gen.lines.map { |l| "#{l}\n" }.join + loads = " float x = a[idx];\n" + loads += " float y = b[idx];\n" if inputs == 2 + <<~GLSL + #version 450 + + layout(local_size_x = 256) in; + + #{bindings(inputs, "partial")} + + layout(push_constant) uniform PushConstants { uint n; } pc; + + shared float sdata[256]; + + // generated by GPU.kernel (reduce: #{reduce}) + void main() { + uint tid = gl_LocalInvocationID.x; + uint idx = gl_GlobalInvocationID.x; + float v = #{identity}; + if (idx < pc.n) { + #{loads}#{prelude} v = #{result}; + } + sdata[tid] = v; + barrier(); + for (uint s = 128u; s > 0u; s >>= 1u) { + if (tid < s) { sdata[tid] = #{combine}; } + barrier(); + } + if (tid == 0u) { partial[gl_WorkGroupID.x] = sdata[0]; } } GLSL end + + def self.bindings(inputs, out) + lines = ["layout(set = 0, binding = 0) buffer BufA { float a[]; };"] + lines << "layout(set = 0, binding = 1) buffer BufB { float b[]; };" if inputs == 2 + slot = inputs == 2 ? 2 : 1 + name = out == "partial" ? "BufOut" : "Buf#{out.upcase}" + lines << "layout(set = 0, binding = #{slot}) buffer #{name} { float #{out}[]; };" + lines.join("\n") + end + end + + # Run the block once over proxies and return the GLSL it traces to. + # `inputs` is 1 (|x|) or 2 (|x, y|). + def self.trace(inputs, reduce, &block) + vars = [Expr.new(:var, ["x"])] + vars << Expr.new(:var, ["y"]) if inputs == 2 + tree = Expr.numeric(KernelContext.new.instance_exec(*vars, &block), "a kernel body") + Kernel.source_for(tree, inputs, reduce) end - # GPU.kernel { |x| ... } -> GPU::Kernel + def self.kernel_inputs(arity, block) + inputs = arity + if inputs.nil? + inputs = block.respond_to?(:arity) && block.arity > 1 ? 2 : 1 + end + unless inputs == 1 || inputs == 2 + raise ArgumentError, "a kernel takes 1 or 2 input arrays, got arity: #{inputs.inspect}" + end + inputs + end + + def self.check_reduce(reduce) + return if reduce.nil? || Kernel::REDUCE.key?(reduce) + raise ArgumentError, + "unknown reduction #{reduce.inspect}; use :sum, :min or :max" + end + + # GPU.kernel { |x| ... } -> GPU::Kernel, element-wise + # GPU.kernel { |x, y| ... } -> two inputs (from the block's arity) + # GPU.kernel(arity: 2) { ... } -> two inputs, stated explicitly + # GPU.kernel(reduce: :sum) { |x| ... } -> Kernel#call returns a Float # # Traces the block and returns a reusable, compiled kernel. Identical source # is compiled once and cached, so calling this in a loop is cheap. - def self.kernel(&block) + def self.kernel(arity: nil, reduce: nil, &block) raise ArgumentError, "GPU.kernel requires a block" unless block - - tree = Expr.wrap(KernelContext.new.instance_exec(Expr.new(:var, ["x"]), &block)) - source = Kernel.source_for(tree.to_glsl) + inputs = kernel_inputs(arity, block) + check_reduce(reduce) + source = trace(inputs, reduce, &block) @kernel_cache ||= {} @kernel_cache[source] ||= begin # mruby-gpu-narray brings Vulkan up lazily on its first operation and # keeps that entry point private to its C code, so reach it through a # public call before compiling. Only runs on a cache miss. GPU.info - Kernel.compile(source) + k = Kernel.compile(source, inputs, reduce ? Kernel::REDUCE[reduce][2] : 0) + k.instance_variable_set(:@inputs, inputs) + k.instance_variable_set(:@reduce, reduce) + k end end # Trace the block into GLSL without compiling it -- for tests and for showing # the generated shader on a slide. - def self.kernel_source(&block) - tree = Expr.wrap(KernelContext.new.instance_exec(Expr.new(:var, ["x"]), &block)) - Kernel.source_for(tree.to_glsl) + def self.kernel_source(arity: nil, reduce: nil, &block) + raise ArgumentError, "GPU.kernel_source requires a block" unless block + check_reduce(reduce) + trace(kernel_inputs(arity, block), reduce, &block) end class NArray # na.map { |x| x * 2 + sin(x) } -- the block becomes one GPU kernel, so the # whole expression is a single dispatch no matter how long it is. def map(&block) - GPU.kernel(&block).call(self) + GPU.kernel(arity: 1, &block).call(self) + end + + # na.map2(nb) { |x, y| x * y + 1 } -- two inputs, one dispatch. + def map2(other, &block) + GPU.kernel(arity: 2, &block).call(self, other) + end + + # na.sum { |x| x * 2 + 1 } -- expression and reduction in one dispatch. + # Without a block this is the C implementation from mruby-gpu-narray. + alias sum_without_block sum + def sum(&block) + block ? GPU.kernel(arity: 1, reduce: :sum, &block).call(self) : sum_without_block + end + + # na.dot(nb) -- inner product, one dispatch. The block captures nothing, + # so the kernel is traced once and kept: tracing a block and looking its + # source up in the cache costs a few hundred microseconds of Ruby time, + # which matters in a loop. (a.sum { ... } traces on every call, because + # its block may capture a local that changes.) + def dot(other) + GPU::NArray.dot_kernel.call(self, other) + end + + def self.dot_kernel + @dot_kernel ||= GPU.kernel(arity: 2, reduce: :sum) { |x, y| x * y } end end end diff --git a/src/gpu_kernel.c b/src/gpu_kernel.c index ec8ec08..a31cadc 100644 --- a/src/gpu_kernel.c +++ b/src/gpu_kernel.c @@ -1,4 +1,4 @@ -/* gpu_kernel.c -- element-wise kernels compiled at run time (GPU::Kernel). +/* gpu_kernel.c -- kernels compiled at run time (GPU::Kernel). * * The Ruby side (mrblib/gpu_kernel.rb) traces a block into GLSL source. This * file takes that source, runs it through glslangValidator, builds a compute @@ -6,8 +6,9 @@ * * Everything here rides on mruby-gpu-narray: the Vulkan context (g_ctx), the * descriptor-set layouts, and the buffer helpers all come from its - * gpu_internal.h. A traced kernel is one input buffer and one output buffer, - * which is exactly LAYOUT_2BUF, so no new descriptor plumbing is needed. + * gpu_internal.h. A kernel is one or two input buffers plus one output buffer + * (the result array, or the per-workgroup partials of a reduction), which is + * exactly LAYOUT_2BUF / LAYOUT_3BUF, so no new descriptor plumbing is needed. * * The dispatch below duplicates the body of mruby-gpu-narray's * dispatch_compute(). Upstream's version is keyed to its PipeId enum and so @@ -22,8 +23,13 @@ #include #include +/* Matches GPU::Kernel::REDUCE in mrblib/gpu_kernel.rb. */ +enum { REDUCE_NONE = 0, REDUCE_SUM = 1, REDUCE_MIN = 2, REDUCE_MAX = 3 }; + typedef struct { VkPipeline pipeline; + int inputs; /* 1 or 2 input arrays */ + int reduce; /* REDUCE_NONE for element-wise, else how partials combine */ } GpuKernel; static void kernel_free(mrb_state *mrb, void *p) { @@ -39,6 +45,11 @@ static const struct mrb_data_type gpu_kernel_type = { "GPU::Kernel", kernel_free }; +/* One input -> (a, out) is LAYOUT_2BUF; two inputs -> (a, b, out) is 3BUF. */ +static LayoutId kernel_layout(int inputs) { + return inputs == 2 ? LAYOUT_3BUF : LAYOUT_2BUF; +} + /* ---- file helpers ---- */ static uint8_t *read_file(const char *path, size_t *size) { @@ -60,7 +71,7 @@ static uint8_t *read_file(const char *path, size_t *size) { /* ---- Vulkan ---- */ -static VkPipeline pipeline_from_spv(const char *path) { +static VkPipeline pipeline_from_spv(const char *path, LayoutId lid) { size_t spv_size = 0; uint8_t *spv = read_file(path, &spv_size); if (!spv) return VK_NULL_HANDLE; @@ -85,7 +96,7 @@ static VkPipeline pipeline_from_spv(const char *path) { .module = shader, .pName = "main" }, - .layout = g_ctx.pipe_layouts[LAYOUT_2BUF] + .layout = g_ctx.pipe_layouts[lid] }; VkPipeline pipeline = VK_NULL_HANDLE; if (vkCreateComputePipelines(g_ctx.device, VK_NULL_HANDLE, 1, &cp_info, @@ -96,22 +107,30 @@ static VkPipeline pipeline_from_spv(const char *path) { return pipeline; } -/* One input buffer, one output buffer, { uint n } push constant. */ -static void dispatch_kernel(VkPipeline pipeline, - VkBuffer *buffers, VkDeviceSize *sizes, +/* num_buffers storage buffers bound in order, { uint n } push constant. + * Records, submits and waits; every failure releases what was created and + * raises through gpu_check. */ +static void dispatch_kernel(mrb_state *mrb, VkPipeline pipeline, LayoutId lid, + VkBuffer *buffers, VkDeviceSize *sizes, int num_buffers, uint32_t n, uint32_t groups) { + if (groups > g_ctx.max_workgroups) { + mrb_raisef(mrb, E_ARGUMENT_ERROR, + "array too large for one dispatch: %i workgroups needed, device allows %i", + (mrb_int)groups, (mrb_int)g_ctx.max_workgroups); + } + VkDescriptorSetAllocateInfo dsai = { .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO, .descriptorPool = g_ctx.desc_pool, .descriptorSetCount = 1, - .pSetLayouts = &g_ctx.desc_layouts[LAYOUT_2BUF] + .pSetLayouts = &g_ctx.desc_layouts[lid] }; VkDescriptorSet desc_set; - vkAllocateDescriptorSets(g_ctx.device, &dsai, &desc_set); + VK_CHECK(mrb, vkAllocateDescriptorSets(g_ctx.device, &dsai, &desc_set)); - VkDescriptorBufferInfo buf_infos[2]; - VkWriteDescriptorSet writes[2]; - for (int i = 0; i < 2; i++) { + VkDescriptorBufferInfo buf_infos[3]; + VkWriteDescriptorSet writes[3]; + for (int i = 0; i < num_buffers; i++) { buf_infos[i] = (VkDescriptorBufferInfo){buffers[i], 0, sizes[i]}; writes[i] = (VkWriteDescriptorSet){ .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, @@ -122,7 +141,7 @@ static void dispatch_kernel(VkPipeline pipeline, .pBufferInfo = &buf_infos[i] }; } - vkUpdateDescriptorSets(g_ctx.device, 2, writes, 0, NULL); + vkUpdateDescriptorSets(g_ctx.device, num_buffers, writes, 0, NULL); VkCommandBufferAllocateInfo cbai = { .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO, @@ -131,41 +150,59 @@ static void dispatch_kernel(VkPipeline pipeline, .commandBufferCount = 1 }; VkCommandBuffer cmd; - vkAllocateCommandBuffers(g_ctx.device, &cbai, &cmd); + VkResult r = vkAllocateCommandBuffers(g_ctx.device, &cbai, &cmd); + if (r != VK_SUCCESS) { + vkFreeDescriptorSets(g_ctx.device, g_ctx.desc_pool, 1, &desc_set); + gpu_check(mrb, r, "vkAllocateCommandBuffers"); + } VkCommandBufferBeginInfo begin = { .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO, .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT }; - vkBeginCommandBuffer(cmd, &begin); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - g_ctx.pipe_layouts[LAYOUT_2BUF], 0, 1, &desc_set, 0, NULL); - vkCmdPushConstants(cmd, g_ctx.pipe_layouts[LAYOUT_2BUF], - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &n); - vkCmdDispatch(cmd, groups, 1, 1); - vkEndCommandBuffer(cmd); - VkFenceCreateInfo fi = {.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO}; - VkFence fence; - vkCreateFence(g_ctx.device, &fi, NULL, &fence); - VkSubmitInfo si = { .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO, .commandBufferCount = 1, .pCommandBuffers = &cmd }; - vkQueueSubmit(g_ctx.queue, 1, &si, fence); - vkWaitForFences(g_ctx.device, 1, &fence, VK_TRUE, UINT64_MAX); + VkFence fence = VK_NULL_HANDLE; + const char *step = "vkBeginCommandBuffer"; + + r = vkBeginCommandBuffer(cmd, &begin); + if (r == VK_SUCCESS) { + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, + g_ctx.pipe_layouts[lid], 0, 1, &desc_set, 0, NULL); + vkCmdPushConstants(cmd, g_ctx.pipe_layouts[lid], + VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &n); + vkCmdDispatch(cmd, groups, 1, 1); + step = "vkEndCommandBuffer"; + r = vkEndCommandBuffer(cmd); + } + if (r == VK_SUCCESS) { + step = "vkCreateFence"; + r = vkCreateFence(g_ctx.device, &fi, NULL, &fence); + } + if (r == VK_SUCCESS) { + step = "vkQueueSubmit"; + r = vkQueueSubmit(g_ctx.queue, 1, &si, fence); + } + if (r == VK_SUCCESS) { + step = "vkWaitForFences"; + r = vkWaitForFences(g_ctx.device, 1, &fence, VK_TRUE, UINT64_MAX); + } - vkDestroyFence(g_ctx.device, fence, NULL); + if (fence != VK_NULL_HANDLE) vkDestroyFence(g_ctx.device, fence, NULL); vkFreeCommandBuffers(g_ctx.device, g_ctx.cmd_pool, 1, &cmd); vkFreeDescriptorSets(g_ctx.device, g_ctx.desc_pool, 1, &desc_set); + + gpu_check(mrb, r, step); } /* ---- mruby surface ---- */ -/* GPU::Kernel.compile(glsl_source) -> GPU::Kernel +/* GPU::Kernel.compile(glsl_source, inputs, reduce_code) -> GPU::Kernel * * Raises GPU::CompileError carrying glslang's own output, so a bad block reads * like a Ruby error rather than a silent failure. @@ -176,12 +213,18 @@ static void dispatch_kernel(VkPipeline pipeline, */ static mrb_value kernel_s_compile(mrb_state *mrb, mrb_value self) { const char *src; - mrb_int src_len; - mrb_get_args(mrb, "s", &src, &src_len); + mrb_int src_len, inputs, reduce; + mrb_get_args(mrb, "sii", &src, &src_len, &inputs, &reduce); struct RClass *gpu = mrb_module_get(mrb, "GPU"); struct RClass *cerr = mrb_class_get_under(mrb, gpu, "CompileError"); + if (inputs != 1 && inputs != 2) { + mrb_raisef(mrb, E_ARGUMENT_ERROR, "a kernel takes 1 or 2 inputs, got %i", inputs); + } + if (reduce < REDUCE_NONE || reduce > REDUCE_MAX) { + mrb_raisef(mrb, E_ARGUMENT_ERROR, "unknown reduction code %i", reduce); + } if (!g_ctx.initialized) { mrb_raise(mrb, E_RUNTIME_ERROR, "no Vulkan device yet -- create a GPU::NArray (or call GPU.info) first"); @@ -223,7 +266,7 @@ static mrb_value kernel_s_compile(mrb_state *mrb, mrb_value self) { mrb_exc_raise(mrb, mrb_exc_new_str(mrb, cerr, msg)); } - VkPipeline pipeline = pipeline_from_spv(spv_path); + VkPipeline pipeline = pipeline_from_spv(spv_path, kernel_layout((int)inputs)); remove(comp_path); remove(log_path); remove(spv_path); @@ -235,6 +278,8 @@ static mrb_value kernel_s_compile(mrb_state *mrb, mrb_value self) { GpuKernel *k = (GpuKernel *)mrb_malloc(mrb, sizeof(GpuKernel)); k->pipeline = pipeline; + k->inputs = (int)inputs; + k->reduce = (int)reduce; struct RObject *obj = (struct RObject *)mrb_obj_alloc(mrb, MRB_TT_CDATA, mrb_class_ptr(self)); @@ -246,25 +291,82 @@ static mrb_value kernel_s_compile(mrb_state *mrb, mrb_value self) { return kernel; } -/* kernel.call(narray) -> narray -- run the kernel over every element. */ +/* kernel.call(a [, b]) -> GPU::NArray for an element-wise kernel, or the + * reduced Float for a reduction. */ static mrb_value kernel_call(mrb_state *mrb, mrb_value self) { - mrb_value arg; - mrb_get_args(mrb, "o", &arg); + const mrb_value *argv; + mrb_int argc; + mrb_get_args(mrb, "*", &argv, &argc); GpuKernel *k = (GpuKernel *)DATA_GET_PTR(mrb, self, &gpu_kernel_type, GpuKernel); - if (mrb_type(arg) != MRB_TT_CDATA || DATA_TYPE(arg) != &gpu_buffer_type) { - mrb_raisef(mrb, E_TYPE_ERROR, "GPU::Kernel#call expects a GPU::NArray, got %s", - mrb_obj_classname(mrb, arg)); + if (argc != k->inputs) { + mrb_raisef(mrb, E_ARGUMENT_ERROR, + "this kernel takes %i array(s) (its block has that many parameters), got %i", + (mrb_int)k->inputs, argc); } - GpuBuffer *a = (GpuBuffer *)DATA_PTR(arg); - GpuBuffer *b = create_buffer(mrb, a->n); - VkBuffer bufs[2] = {a->buffer, b->buffer}; - VkDeviceSize sizes[2] = {a->bytes, b->bytes}; - dispatch_kernel(k->pipeline, bufs, sizes, a->n, (a->n + 255) / 256); + /* Copy out of the VM stack before allocating anything. */ + mrb_value args[2] = {mrb_nil_value(), mrb_nil_value()}; + GpuBuffer *in[2] = {NULL, NULL}; + for (mrb_int i = 0; i < argc; i++) { + args[i] = argv[i]; + if (mrb_type(args[i]) != MRB_TT_CDATA || DATA_TYPE(args[i]) != &gpu_buffer_type) { + mrb_raisef(mrb, E_TYPE_ERROR, "GPU::Kernel#call expects a GPU::NArray, got %s", + mrb_obj_classname(mrb, args[i])); + } + in[i] = (GpuBuffer *)DATA_PTR(args[i]); + } + if (k->inputs == 2 && in[0]->n != in[1]->n) { + mrb_raisef(mrb, E_ARGUMENT_ERROR, "shape mismatch: %i vs %i", + (mrb_int)in[0]->n, (mrb_int)in[1]->n); + } - return wrap_buffer(mrb, mrb_obj_class(mrb, arg), b); + uint32_t n = in[0]->n; + uint32_t groups = (n + 255) / 256; + LayoutId lid = kernel_layout(k->inputs); + VkBuffer bufs[3]; + VkDeviceSize sizes[3]; + for (int i = 0; i < k->inputs; i++) { + bufs[i] = in[i]->buffer; + sizes[i] = in[i]->bytes; + } + + if (k->reduce == REDUCE_NONE) { + GpuBuffer *out = create_buffer(mrb, n); + /* Wrapped before dispatching: a Vulkan failure raises, and the GC can only + * free what it owns. */ + mrb_value result = wrap_buffer(mrb, mrb_obj_class(mrb, args[0]), out); + bufs[k->inputs] = out->buffer; + sizes[k->inputs] = out->bytes; + dispatch_kernel(mrb, k->pipeline, lid, bufs, sizes, k->inputs + 1, n, groups); + return result; + } + + /* Reduction: one partial per workgroup, combined here in double precision + * (the same arrangement as mruby-gpu-narray's #sum). */ + if (n == 0) { + if (k->reduce == REDUCE_SUM) return mrb_float_value(mrb, 0.0); + mrb_raise(mrb, E_ARGUMENT_ERROR, "min/max of an empty array"); + } + GpuBuffer *partial = create_buffer(mrb, groups); + wrap_buffer(mrb, mrb_obj_class(mrb, args[0]), partial); /* GC-owned scratch */ + bufs[k->inputs] = partial->buffer; + sizes[k->inputs] = partial->bytes; + dispatch_kernel(mrb, k->pipeline, lid, bufs, sizes, k->inputs + 1, n, groups); + + float *pm = map_buffer(mrb, partial); + double acc = (k->reduce == REDUCE_SUM) ? 0.0 : (double)pm[0]; + for (uint32_t i = 0; i < groups; i++) { + double v = (double)pm[i]; + switch (k->reduce) { + case REDUCE_SUM: acc += v; break; + case REDUCE_MIN: if (v < acc) acc = v; break; + case REDUCE_MAX: if (v > acc) acc = v; break; + } + } + unmap_buffer(partial); + return mrb_float_value(mrb, (mrb_float)acc); } void mrb_mruby_gpu_kernel_gem_init(mrb_state *mrb) { @@ -273,8 +375,8 @@ void mrb_mruby_gpu_kernel_gem_init(mrb_state *mrb) { struct RClass *kernel = mrb_define_class_under(mrb, gpu, "Kernel", mrb->object_class); MRB_SET_INSTANCE_TT(kernel, MRB_TT_CDATA); - mrb_define_class_method(mrb, kernel, "compile", kernel_s_compile, MRB_ARGS_REQ(1)); - mrb_define_method(mrb, kernel, "call", kernel_call, MRB_ARGS_REQ(1)); + mrb_define_class_method(mrb, kernel, "compile", kernel_s_compile, MRB_ARGS_REQ(3)); + mrb_define_method(mrb, kernel, "call", kernel_call, MRB_ARGS_ANY()); } void mrb_mruby_gpu_kernel_gem_final(mrb_state *mrb) { diff --git a/test/kernel_test.rb b/test/kernel_test.rb index 93bf3b0..b3088a8 100644 --- a/test/kernel_test.rb +++ b/test/kernel_test.rb @@ -115,6 +115,106 @@ def assert_raise(label, klass) GPU.kernel end +# ---- two inputs ---- +assert_ary("map2 { x * y + 1 }", [11.0, 41.0, 91.0], + GPU::SFloat[1, 2, 3].map2(GPU::SFloat[10, 20, 30]) { |x, y| x * y + 1 }.to_a) +assert_ary("kernel { |x, y| } takes two arrays (arity from the block)", [-9.0, -18.0, -27.0], + GPU.kernel { |x, y| x - y }.call(GPU::SFloat[1, 2, 3], GPU::SFloat[10, 20, 30]).to_a) +assert_ary("kernel(arity: 2) stated explicitly", [11.0, 22.0], + GPU.kernel(arity: 2) { |x, y| x + y }.call(GPU::SFloat[1, 2], GPU::SFloat[10, 20]).to_a) +assert_true("two-input source loads y from the second buffer", + GPU.kernel_source { |x, y| x + y }.include?("float y = b[idx];") && + GPU.kernel_source { |x, y| x + y }.include?("c[idx] = (x + y);")) +assert_raise("map2 with mismatched sizes -> ArgumentError", ArgumentError) do + GPU::SFloat[1, 2, 3].map2(GPU::SFloat[1, 2]) { |x, y| x + y } +end +assert_raise("two-input kernel called with one array -> ArgumentError", ArgumentError) do + GPU.kernel { |x, y| x + y }.call(GPU::SFloat[1, 2]) +end +assert_raise("one-input kernel called with two arrays -> ArgumentError", ArgumentError) do + GPU.kernel { |x| x }.call(GPU::SFloat[1, 2], GPU::SFloat[1, 2]) +end + +# ---- reductions: the expression and the sum in one dispatch ---- +assert_near("sum { x * 2 + 1 } over 1000 elements", 1_000_000.0, + GPU::SFloat.new(1000).seq.sum { |x| x * 2 + 1 }) +assert_near("sum without a block still works", 3.0, GPU::SFloat[1, 2].sum) +assert_near("dot product", 32.0, GPU::SFloat[1, 2, 3].dot(GPU::SFloat[4, 5, 6])) +assert_near("kernel(reduce: :sum) { |x, y| } is the dot product", 32.0, + GPU.kernel(reduce: :sum) { |x, y| x * y }.call(GPU::SFloat[1, 2, 3], GPU::SFloat[4, 5, 6])) +assert_near("reduce: :max", 5.0, + GPU.kernel(reduce: :max) { |x| abs(x) }.call(GPU::SFloat[-5, 3, 2])) +assert_near("reduce: :min", -5.0, + GPU.kernel(reduce: :min) { |x| x }.call(GPU::SFloat[-5, 3, 2])) +# 1000 elements = 4 workgroups, so the host really combines several partials. +assert_near("reduce: :max across workgroups", 999.0, + GPU.kernel(reduce: :max) { |x| x }.call(GPU::SFloat.new(1000).seq)) +assert_near("sum of an empty array is 0", 0.0, GPU::SFloat.new(0).sum { |x| x }) +assert_true("reduce source has the shared-memory tree", + GPU.kernel_source(reduce: :sum) { |x| x }.include?("shared float sdata[256];")) +assert_raise("unknown reduction -> ArgumentError", ArgumentError) do + GPU.kernel(reduce: :product) { |x| x } +end + +# ---- control flow: where, & | !, iterate ---- +assert_ary("where(x > 0, x, 0) is relu", [0.0, 0.0, 0.0, 1.0, 2.0], + GPU::SFloat[-2, -1, 0, 1, 2].map { |x| where(x > 0, x, 0) }.to_a) +assert_ary("select is where", [0.0, 1.0], + GPU::SFloat[-1, 1].map { |x| select(x > 0, 1, 0) }.to_a) +assert_ary("(a) & (b)", [0.0, 1.0, 1.0, 1.0, 0.0], + GPU::SFloat[0, 1, 2, 3, 4].map { |x| where((x >= 1) & (x <= 3), 1, 0) }.to_a) +assert_ary("(a) | (b) and !", [0.0, 1.0, 2.0, -1.0, 4.0], + GPU::SFloat[0, 1, 2, 3, 4].map { |x| where(!(x > 2) | (x == 4), x, -1) }.to_a) +assert_ary("== and !=", [100.0, 2.0, 100.0], + GPU::SFloat[1, 2, 3].map { |x| where(x != 2, 100, x) }.to_a) +assert_true("where becomes GLSL's ternary", + GPU.kernel_source { |x| where(x > 0, x, 0) }.include?("((x > 0.0) ? x : 0.0)")) +assert_ary("iterate: Newton's method matches sqrt(x)", + GPU::SFloat[1, 4, 9, 2, 100].map { |x| sqrt(x) }.to_a, + GPU::SFloat[1, 4, 9, 2, 100].map { |x| iterate(x, 20) { |g, i| (g + x / g) * 0.5 } }.to_a) +assert_ary("iterate: i counts from 0", [0.0 + 1 + 2 + 3] * 2, + GPU::SFloat[0, 0].map { |x| iterate(0, 4) { |acc, i| acc + i } }.to_a) +assert_ary("nested iterate", [18.0, 18.0], + GPU::SFloat[0, 0].map { |x| + iterate(0, 3) { |acc, i| acc + iterate(0, 4) { |acc2, j| acc2 + i * j } } + }.to_a) +assert_ary("a loop result used twice is one loop", [16.0, 32.0], + GPU::SFloat[1, 2].map { |x| s = iterate(x, 3) { |a2, i| a2 * 2 }; s + s }.to_a) +assert_true("iterate emits a real GLSL for loop", + GPU.kernel_source { |x| iterate(x, 5) { |acc, i| acc * 2 } }.include?("for (int i2 = 0; i2 < 5; i2++) {"), + GPU.kernel_source { |x| iterate(x, 5) { |acc, i| acc * 2 } }) +assert_ary("a Ruby loop with a known count is unrolled", [8.0, 16.0], + GPU::SFloat[1, 2].map { |x| y = x; 3.times { y = y * 2 }; y }.to_a) +assert_true("...and leaves no loop in the shader", + !GPU.kernel_source { |x| y = x; 3.times { y = y * 2 }; y }.include?("for (")) +assert_true("tracing shows the loop node", + GPU::KernelContext.new.instance_exec(GPU::Expr.new(:var, ["x"])) { |x| + iterate(x, 2) { |acc, i| acc + 1 } + }.inspect == "(iterate x 2 (+ acc 1.0))") + +# ---- control-flow errors are Ruby errors, not glslang's ---- +assert_raise("a condition in arithmetic -> TypeError", TypeError) do + GPU::SFloat[1, 2].map { |x| (x > 1) + 1 } +end +assert_raise("where with a number as the condition -> TypeError", TypeError) do + GPU::SFloat[1, 2].map { |x| where(x, 1, 0) } +end +assert_raise("& on numbers -> TypeError", TypeError) do + GPU::SFloat[1, 2].map { |x| where(x & x, 1, 0) } +end +assert_raise("iterate count must be an Integer (got Float) -> TypeError", TypeError) do + GPU::SFloat[1, 2].map { |x| iterate(x, 2.5) { |acc, i| acc } } +end +assert_raise("iterate count must be known at trace time (got x) -> TypeError", TypeError) do + GPU::SFloat[1, 2].map { |x| iterate(0, x) { |acc, i| acc } } +end +assert_raise("negative iterate count -> ArgumentError", ArgumentError) do + GPU::SFloat[1, 2].map { |x| iterate(0, -1) { |acc, i| acc } } +end +assert_raise("loop variable escaping its block -> ArgumentError", ArgumentError) do + GPU::SFloat[1, 2].map { |x| leaked = nil; iterate(0, 2) { |acc, i| leaked = i; acc }; leaked } +end + # ---- summary ---- puts puts "#{$pass + $fail} tests, #{$pass} passed, #{$fail} failed" From b0edf89d291d2bc5497e605dd95dcf8a222349a5 Mon Sep 17 00:00:00 2001 From: yujiteshima Date: Sat, 12 Sep 2026 23:41:43 +0900 Subject: [PATCH 2/2] Dispatch through mruby-gpu-narray's dispatch_pipeline() Drops the private copy of the dispatch routine. A traced kernel -- one or two inputs, element-wise or reduction -- is now recorded into the base gem's deferred batch like any built-in operator, so a.map { ... } queues one dispatch and nothing is submitted until a result is read. For reductions, mapping the partials is what flushes the batch. examples/kernel_dsl.rb times the operator chain three ways (a wait per operator, one batch, one traced kernel), because under deferred submission the old loop measured recording rather than execution. Needs a mruby-gpu-narray that exports dispatch_pipeline() (its deferred-submit branch, PR #4). Stacked on dsl-control-flow (PR #1). Co-Authored-By: Claude Fable 5.1 --- README.md | 25 ++++++--- examples/kernel_dsl.rb | 19 ++++--- src/gpu_kernel.c | 120 ++++------------------------------------- 3 files changed, 43 insertions(+), 121 deletions(-) diff --git a/README.md b/README.md index 507d78e..3b00e2b 100644 --- a/README.md +++ b/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 @@ -187,17 +198,19 @@ it before this gem in `build_config.rb`. 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 (2 or 3 buffers) +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 diff --git a/examples/kernel_dsl.rb b/examples/kernel_dsl.rb index ccec1af..f8c12d1 100644 --- a/examples/kernel_dsl.rb +++ b/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" diff --git a/src/gpu_kernel.c b/src/gpu_kernel.c index a31cadc..b590e3c 100644 --- a/src/gpu_kernel.c +++ b/src/gpu_kernel.c @@ -10,10 +10,9 @@ * (the result array, or the per-workgroup partials of a reduction), which is * exactly LAYOUT_2BUF / LAYOUT_3BUF, so no new descriptor plumbing is needed. * - * The dispatch below duplicates the body of mruby-gpu-narray's - * dispatch_compute(). Upstream's version is keyed to its PipeId enum and so - * cannot bind a pipeline we built ourselves; when it grows a pipeline-taking - * variant, this copy should go away. + * Dispatch goes through mruby-gpu-narray's dispatch_pipeline(), so a traced + * kernel joins the same deferred batch as the built-in operators: nothing is + * submitted until a result is read. * * We shell out to glslangValidator rather than linking glslang, so the link * dependency stays libvulkan only -- and the compiler is already required to @@ -107,99 +106,6 @@ static VkPipeline pipeline_from_spv(const char *path, LayoutId lid) { return pipeline; } -/* num_buffers storage buffers bound in order, { uint n } push constant. - * Records, submits and waits; every failure releases what was created and - * raises through gpu_check. */ -static void dispatch_kernel(mrb_state *mrb, VkPipeline pipeline, LayoutId lid, - VkBuffer *buffers, VkDeviceSize *sizes, int num_buffers, - uint32_t n, uint32_t groups) { - if (groups > g_ctx.max_workgroups) { - mrb_raisef(mrb, E_ARGUMENT_ERROR, - "array too large for one dispatch: %i workgroups needed, device allows %i", - (mrb_int)groups, (mrb_int)g_ctx.max_workgroups); - } - - VkDescriptorSetAllocateInfo dsai = { - .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO, - .descriptorPool = g_ctx.desc_pool, - .descriptorSetCount = 1, - .pSetLayouts = &g_ctx.desc_layouts[lid] - }; - VkDescriptorSet desc_set; - VK_CHECK(mrb, vkAllocateDescriptorSets(g_ctx.device, &dsai, &desc_set)); - - VkDescriptorBufferInfo buf_infos[3]; - VkWriteDescriptorSet writes[3]; - for (int i = 0; i < num_buffers; i++) { - buf_infos[i] = (VkDescriptorBufferInfo){buffers[i], 0, sizes[i]}; - writes[i] = (VkWriteDescriptorSet){ - .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, - .dstSet = desc_set, - .dstBinding = i, - .descriptorCount = 1, - .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, - .pBufferInfo = &buf_infos[i] - }; - } - vkUpdateDescriptorSets(g_ctx.device, num_buffers, writes, 0, NULL); - - VkCommandBufferAllocateInfo cbai = { - .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO, - .commandPool = g_ctx.cmd_pool, - .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY, - .commandBufferCount = 1 - }; - VkCommandBuffer cmd; - VkResult r = vkAllocateCommandBuffers(g_ctx.device, &cbai, &cmd); - if (r != VK_SUCCESS) { - vkFreeDescriptorSets(g_ctx.device, g_ctx.desc_pool, 1, &desc_set); - gpu_check(mrb, r, "vkAllocateCommandBuffers"); - } - - VkCommandBufferBeginInfo begin = { - .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO, - .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT - }; - VkFenceCreateInfo fi = {.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO}; - VkSubmitInfo si = { - .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO, - .commandBufferCount = 1, - .pCommandBuffers = &cmd - }; - VkFence fence = VK_NULL_HANDLE; - const char *step = "vkBeginCommandBuffer"; - - r = vkBeginCommandBuffer(cmd, &begin); - if (r == VK_SUCCESS) { - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - g_ctx.pipe_layouts[lid], 0, 1, &desc_set, 0, NULL); - vkCmdPushConstants(cmd, g_ctx.pipe_layouts[lid], - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &n); - vkCmdDispatch(cmd, groups, 1, 1); - step = "vkEndCommandBuffer"; - r = vkEndCommandBuffer(cmd); - } - if (r == VK_SUCCESS) { - step = "vkCreateFence"; - r = vkCreateFence(g_ctx.device, &fi, NULL, &fence); - } - if (r == VK_SUCCESS) { - step = "vkQueueSubmit"; - r = vkQueueSubmit(g_ctx.queue, 1, &si, fence); - } - if (r == VK_SUCCESS) { - step = "vkWaitForFences"; - r = vkWaitForFences(g_ctx.device, 1, &fence, VK_TRUE, UINT64_MAX); - } - - if (fence != VK_NULL_HANDLE) vkDestroyFence(g_ctx.device, fence, NULL); - vkFreeCommandBuffers(g_ctx.device, g_ctx.cmd_pool, 1, &cmd); - vkFreeDescriptorSets(g_ctx.device, g_ctx.desc_pool, 1, &desc_set); - - gpu_check(mrb, r, step); -} - /* ---- mruby surface ---- */ /* GPU::Kernel.compile(glsl_source, inputs, reduce_code) -> GPU::Kernel @@ -325,21 +231,16 @@ static mrb_value kernel_call(mrb_state *mrb, mrb_value self) { uint32_t n = in[0]->n; uint32_t groups = (n + 255) / 256; LayoutId lid = kernel_layout(k->inputs); - VkBuffer bufs[3]; - VkDeviceSize sizes[3]; - for (int i = 0; i < k->inputs; i++) { - bufs[i] = in[i]->buffer; - sizes[i] = in[i]->bytes; - } + GpuBuffer *bufs[3] = {in[0], in[1], NULL}; /* inputs, then the output */ if (k->reduce == REDUCE_NONE) { GpuBuffer *out = create_buffer(mrb, n); /* Wrapped before dispatching: a Vulkan failure raises, and the GC can only * free what it owns. */ mrb_value result = wrap_buffer(mrb, mrb_obj_class(mrb, args[0]), out); - bufs[k->inputs] = out->buffer; - sizes[k->inputs] = out->bytes; - dispatch_kernel(mrb, k->pipeline, lid, bufs, sizes, k->inputs + 1, n, groups); + bufs[k->inputs] = out; + dispatch_pipeline(mrb, k->pipeline, lid, bufs, k->inputs + 1, &n, sizeof(n), + groups, 1, 1); return result; } @@ -351,10 +252,11 @@ static mrb_value kernel_call(mrb_state *mrb, mrb_value self) { } GpuBuffer *partial = create_buffer(mrb, groups); wrap_buffer(mrb, mrb_obj_class(mrb, args[0]), partial); /* GC-owned scratch */ - bufs[k->inputs] = partial->buffer; - sizes[k->inputs] = partial->bytes; - dispatch_kernel(mrb, k->pipeline, lid, bufs, sizes, k->inputs + 1, n, groups); + bufs[k->inputs] = partial; + dispatch_pipeline(mrb, k->pipeline, lid, bufs, k->inputs + 1, &n, sizeof(n), + groups, 1, 1); + /* map_buffer flushes the batch: the partials are its last dispatch. */ float *pm = map_buffer(mrb, partial); double acc = (k->reduce == REDUCE_SUM) ? 0.0 : (double)pm[0]; for (uint32_t i = 0; i < groups; i++) {