kumi 0.0.13 → 0.0.15

data/lib/kumi/core/functions/signature_resolver.rb

@@ -0,0 +1,272 @@
+# frozen_string_literal: true
+
+require_relative "errors"
+require_relative "shape"
+require_relative "signature"
+
+module Kumi
+  module Core
+    module Functions
+      # Given a set of signatures and actual argument shapes, pick the best match.
+      # Supports NEP 20 extensions: fixed-size, flexible, and broadcastable dimensions.
+      #
+      # Inputs:
+      #   signatures : Array<Signature> (with Dimension objects)
+      #   arg_shapes : Array<Array<Symbol|Integer>>   e.g., [[:i], [:i]] or [[], [3]] or [[2, :i]]
+      #
+      # Returns:
+      #   { signature:, result_axes:, join_policy:, dropped_axes:, effective_signature: }
+      #
+      # NEP 20 Matching rules:
+      # - Arity must match exactly (before flexible dimension resolution).
+      # - Fixed-size dimensions (integers) must match exactly.
+      # - Flexible dimensions (?) can be omitted if not present in all operands.
+      # - Broadcastable dimensions (|1) can match scalar or size-1 dimensions.
+      # - For each param position, shapes are checked according to NEP 20 rules.
+      # - We prefer exact matches, then flexible matches, then broadcast matches.
+      class SignatureResolver
+        class << self
+          def choose(signatures:, arg_shapes:)
+            # Handle empty arg_shapes for zero-arity functions
+            arg_shapes = [] if arg_shapes.nil?
+            sanity_check_args!(arg_shapes)
+
+            candidates = signatures.map do |sig|
+              score = match_score(sig, arg_shapes)
+              next if score.nil?
+
+              # Convert arg_shapes to normalized Dimension arrays for environment building
+              normalized_args = arg_shapes.map { |shape| normalize_shape(shape) }
+              env = build_dimension_environment(sig, normalized_args)
+              next if env.nil?  # Skip candidates with dimension conflicts
+
+              {
+                signature: sig,
+                score: score,
+                result_axes: sig.out_shape.map(&:name), # Convert Dimension objects to names for backward compatibility
+                join_policy: sig.join_policy,
+                dropped_axes: sig.dropped_axes.map { |name| name.is_a?(Symbol) ? name : name.to_sym }, # Convert to symbols
+                env: env
+              }
+            end.compact
+
+            raise SignatureMatchError, mismatch_message(signatures, arg_shapes) if candidates.empty?
+
+            # Lower score is better: 0 = exact-everywhere, then number of broadcasts
+            best = candidates.min_by { |c| c[:score] }
+            
+            # Add effective signature and environment for analyzer/lowering
+            best[:effective_signature] = {
+              in_shapes: best[:signature].in_shapes.map { |dims| dims.map(&:name) },
+              out_shape: best[:signature].out_shape.map(&:name),
+              join_policy: best[:signature].join_policy
+            }
+            # env is already included from candidate building
+            
+            best
+          end
+
+          private
+
+          def sanity_check_args!(arg_shapes)
+            unless arg_shapes.is_a?(Array) &&
+                   arg_shapes.all? { |s| s.is_a?(Array) && s.all? { |a| a.is_a?(Symbol) || a.is_a?(Integer) } }
+              raise SignatureMatchError, "arg_shapes must be an array of dimension arrays (symbols or integers), got: #{arg_shapes.inspect}"
+            end
+          end
+
+          # Returns an integer "broadcast cost" or nil if not matchable.
+          # Lower score = better match: 0 = exact, then increasing cost for broadcasts/flexibility
+          def match_score(sig, arg_shapes)
+            return nil unless sig.arity == arg_shapes.length
+
+            # Convert arg_shapes to normalized Dimension arrays for comparison
+            normalized_args = arg_shapes.map { |shape| normalize_shape(shape) }
+
+            # Try to match each argument against its expected signature shape
+            cost = 0
+            sig.in_shapes.each_with_index do |expected_dims, idx|
+              got_dims = normalized_args[idx]
+              arg_cost = match_argument_cost(got: got_dims, expected: expected_dims)
+              return nil if arg_cost.nil?
+
+              cost += arg_cost
+            end
+
+            # Additional checks for join_policy constraints
+            return nil unless valid_join_policy?(sig, normalized_args)
+
+            cost
+          end
+
+          private
+
+          # Convert a shape array (symbols/integers) to normalized Dimension array
+          def normalize_shape(shape)
+            shape.map do |dim|
+              case dim
+              when Symbol
+                Dimension.new(dim)
+              when Integer
+                Dimension.new(dim)
+              else
+                raise SignatureMatchError, "Invalid dimension type: #{dim.class}"
+              end
+            end
+          end
+
+          # Calculate cost of matching one argument against expected dimensions
+          def match_argument_cost(got:, expected:)
+            # Handle scalar first
+            if got.empty?
+              return expected.empty? ? 0 : (expected.any?(&:flexible?) ? 10 : 1) # scalar broadcast or flexible-tail
+            end
+
+            # Try strict matching first if no flexible dimensions
+            if !expected.any?(&:flexible?) && got.length == expected.length
+              total = 0
+              got.zip(expected).each do |g, e|
+                c = match_dimension_cost(got: g, expected: e)
+                return nil if c.nil?
+                total += c
+              end
+              return total
+            end
+
+            # Use right-aligned flexible matching
+            right_align_match(got: got, expected: expected)
+          end
+
+          # Right-aligned matching for flexible dimensions (NEP 20 ? modifier)
+          def right_align_match(got:, expected:)
+            gi = got.length - 1
+            ei = expected.length - 1
+            cost = 0
+
+            while ei >= 0
+              exp = expected[ei]
+
+              if exp.flexible? && gi < 0
+                # optional tail dimension that we don't have → ok, consume expected only
+                ei -= 1
+                cost += 10
+                next
+              end
+
+              return nil if gi < 0 # ran out of got dims and exp wasn't flexible
+
+              got_dim = got[gi]
+              dim_cost = match_dimension_cost(got: got_dim, expected: exp)
+              if dim_cost.nil?
+                # if exp is flexible, we can try to drop it
+                if exp.flexible?
+                  ei -= 1
+                  cost += 10
+                  next
+                else
+                  return nil
+                end
+              else
+                cost += dim_cost
+                gi -= 1
+                ei -= 1
+              end
+            end
+
+            # if we still have leftover got dims, argument is longer than expected → not a match
+            return nil if gi >= 0
+
+            cost
+          end
+
+          # Calculate cost of matching one dimension against another
+          def match_dimension_cost(got:, expected:)
+            return 0 if got == expected # Exact match
+
+            # Fixed-size equality
+            if got.fixed_size? && expected.fixed_size?
+              return got.size == expected.size ? 0 : nil
+            end
+
+            # Same symbolic name (ignoring modifiers) → ok unless one is fixed and the other isn't (penalize)
+            if got.named? && expected.named? && got.name == expected.name
+              return (got.fixed_size? || expected.fixed_size?) ? 2 : 0
+            end
+
+            # Broadcastable expected dim accepts scalar or size-1
+            if expected.broadcastable?
+              # scalar at argument level would have been handled in match_argument_cost
+              # so here we check for size-1 fixed dimensions
+              return 3 if got.fixed_size? && got.size == 1
+              # Named dimensions that could be size-1 at runtime also get broadcast cost
+              return 3 if got.named?
+            end
+
+            nil # No match possible
+          end
+
+          # Check if join_policy constraints are satisfied
+          def valid_join_policy?(sig, normalized_args)
+            return true if sig.join_policy # :zip or :product allows different axes
+
+            # nil join_policy: check if dimension names are consistent
+            non_scalar_args = normalized_args.reject { |a| Shape.scalar?(a) }
+            return true if non_scalar_args.empty?
+
+            # For nil join_policy, we allow different dimension names if:
+            # 1. All args have same dimension names (element-wise operations), OR
+            # 2. The constraint solver can validate cross-dimensional consistency (like matmul)
+            first_names = non_scalar_args.first.map(&:name)
+            same_names = non_scalar_args.all? { |arg| arg.map(&:name) == first_names }
+            
+            return true if same_names
+            
+            # If dimension names differ, check if constraint solver can handle it
+            # This allows operations like matmul where dimensions are linked across arguments
+            env = build_dimension_environment(sig, normalized_args)
+            !env.nil?
+          end
+
+          def mismatch_message(signatures, arg_shapes)
+            sigs = signatures.map(&:inspect).join(", ")
+            "no matching signature for shapes #{pp_shapes(arg_shapes)} among [#{sigs}]"
+          end
+
+          def pp_shapes(shapes)
+            shapes.map { |ax| "(#{ax.join(',')})" }.join(", ")
+          end
+
+          # Build dimension environment by checking consistency of named dimensions across arguments
+          def build_dimension_environment(sig, normalized_args)
+            env = {}
+            
+            # Walk all expected dimensions across all arguments
+            sig.in_shapes.each_with_index do |expected_shape, arg_idx|
+              got_shape = normalized_args[arg_idx] || []
+              
+              expected_shape.each_with_index do |exp_dim, dim_idx|
+                next unless exp_dim.named? && dim_idx < got_shape.length
+                
+                got_dim = got_shape[dim_idx]
+                dim_name = exp_dim.name
+                
+                # Check for consistency: same dimension name must map to same concrete value
+                if env.key?(dim_name)
+                  # If we've seen this dimension name before, it must match
+                  if env[dim_name] != got_dim
+                    return nil  # Inconsistent binding - signature doesn't match
+                  end
+                else
+                  # First time seeing this dimension name - record the binding
+                  env[dim_name] = got_dim
+                end
+              end
+            end
+            
+            env
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/ir/execution_engine/interpreter.rb

@@ -9,20 +9,70 @@ module Kumi
           PRODUCES_SLOT = %i[const load_input ref array map reduce lift align_to switch].freeze
           NON_PRODUCERS = %i[guard_push guard_pop assign store].freeze
 
+          def self.build_name_index(ir_module)
+            index = {}
+            ir_module.decls.each do |decl|
+              decl.ops.each do |op|
+                next unless op.tag == :store
+
+                name = op.attrs[:name]
+                index[name] = decl if name
+              end
+            end
+            index
+          end
+
           def self.run(ir_module, ctx, accessors:, registry:)
             # Validate registry is properly initialized
             raise ArgumentError, "Registry cannot be nil" if registry.nil?
             raise ArgumentError, "Registry must be a Hash, got #{registry.class}" unless registry.is_a?(Hash)
 
+            # --- PROFILER: init per run (but not in persistent mode) ---
+            if Profiler.enabled?
+              schema_name = ctx[:schema_name] || "UnknownSchema"
+              if Profiler.persistent?
+                # In persistent mode, just update schema name without full reset
+                Profiler.set_schema_name(schema_name)
+              else
+                # Normal mode: full reset with schema name
+                Profiler.reset!(meta: { decls: ir_module.decls&.size || 0, schema_name: schema_name })
+              end
+            end
+
             outputs = {}
             target = ctx[:target]
             guard_stack = [true]
 
-            ir_module.decls.each do |decl|
+            # Always ensure we have a declaration cache - either from caller or new for this VM run
+            declaration_cache = ctx[:declaration_cache] || {}
+
+            # Build name index for targeting by stored names
+            name_index = ctx[:name_index] || (target ? build_name_index(ir_module) : nil)
+
+            # Choose declarations to execute - prefer explicit schedule if present
+            decls_to_run =
+              if ctx[:decls_to_run]
+                ctx[:decls_to_run] # array of decl objects
+              elsif target
+                # Prefer a decl that STORES the target (covers __vec twins)
+                d = name_index && name_index[target]
+                # Fallback: allow targeting by decl name (legacy behavior)
+                d ||= ir_module.decls.find { |dd| dd.name == target }
+                raise "Unknown target: #{target}" unless d
+
+                [d]
+              else
+                ir_module.decls
+              end
+
+            decls_to_run.each do |decl|
               slots = []
               guard_stack = [true] # reset per decl
 
               decl.ops.each_with_index do |op, op_index|
+                t0 = Profiler.enabled? ? Profiler.t0 : nil
+                cpu_t0 = Profiler.enabled? ? Profiler.cpu_t0 : nil
+                rows_touched = nil
                 if ENV["ASSERT_VM_SLOTS"] == "1"
                   expected = op_index
                   unless slots.length == expected
@@ -47,17 +97,26 @@ module Kumi
                                    false
                                  end
                   slots << nil # keep slot_id == op_index
+                  if t0
+                    Profiler.record!(decl: decl.name, idx: op_index, tag: op.tag, op: op, t0: t0, cpu_t0: cpu_t0, rows: 0,
+                                     note: "enter")
+                  end
                   next
 
                 when :guard_pop
                   guard_stack.pop
                   slots << nil
+                  Profiler.record!(decl: decl.name, idx: op_index, tag: op.tag, op: op, t0: t0, cpu_t0: cpu_t0, rows: 0, note: "exit") if t0
                   next
                 end
 
                 # Skip body when guarded off, but keep indices aligned
                 unless guard_stack.last
                   slots << nil if PRODUCES_SLOT.include?(op.tag) || NON_PRODUCERS.include?(op.tag)
+                  if t0
+                    Profiler.record!(decl: decl.name, idx: op_index, tag: op.tag, op: op, t0: t0, cpu_t0: cpu_t0, rows: 0,
+                                     note: "skipped")
+                  end
                   next
                 end
 
@@ -69,35 +128,74 @@ module Kumi
                   raise "assign: dst/src OOB" if dst >= slots.length || src >= slots.length
 
                   slots[dst] = slots[src]
+                  Profiler.record!(decl: decl.name, idx: op_index, tag: :assign, op: op, t0: t0, cpu_t0: cpu_t0, rows: 1) if t0
 
                 when :const
                   result = Values.scalar(op.attrs[:value])
                   puts "DEBUG Const #{op.attrs[:value].inspect}: result=#{result}" if ENV["DEBUG_VM_ARGS"]
                   slots << result
+                  Profiler.record!(decl: decl.name, idx: op_index, tag: :const, op: op, t0: t0, cpu_t0: cpu_t0, rows: 1) if t0
 
                 when :load_input
                   plan_id = op.attrs[:plan_id]
                   scope = op.attrs[:scope] || []
                   scalar = op.attrs[:is_scalar]
                   indexed = op.attrs[:has_idx]
-                  raw = accessors.fetch(plan_id).call(ctx[:input] || ctx["input"])
 
-                  puts "DEBUG LoadInput plan_id: #{plan_id} raw_values: #{raw.inspect}" if ENV["DEBUG_VM_ARGS"]
+                  # NEW: consult runtime accessor cache
+                  acc_cache = ctx[:accessor_cache] || {}
+                  input_obj = ctx[:input] || ctx["input"]
+                  cache_key = [plan_id, input_obj.object_id]
+
+                  if acc_cache.key?(cache_key)
+                    raw = acc_cache[cache_key]
+                    hit = true
+                  else
+                    raw = accessors.fetch(plan_id).call(input_obj)
+                    acc_cache[cache_key] = raw
+                    hit = false
+                  end
+
+                  puts "DEBUG LoadInput plan_id: #{plan_id} raw_values: #{raw.inspect} cache_hit: #{hit}" if ENV["DEBUG_VM_ARGS"]
                   slots << if scalar
                              Values.scalar(raw)
                            elsif indexed
+                             rows_touched = raw.respond_to?(:size) ? raw.size : raw.count
                              Values.vec(scope, raw.map { |v, idx| { v: v, idx: Array(idx) } }, true)
                            else
+                             rows_touched = raw.respond_to?(:size) ? raw.size : raw.count
                              Values.vec(scope, raw.map { |v| { v: v } }, false)
                            end
+                  rows_touched ||= 1
+                  cache_note = hit ? "hit:#{plan_id}" : "miss:#{plan_id}"
+                  if t0
+                    Profiler.record!(decl: decl.name, idx: op_index, tag: :load_input, op: op, t0: t0, cpu_t0: cpu_t0,
+                                     rows: rows_touched, note: cache_note)
+                  end
 
                 when :ref
                   name = op.attrs[:name]
-                  referenced_value = outputs.fetch(name) { raise "Missing output for reference: #{name}" }
+
+                  if outputs.key?(name)
+                    referenced = outputs[name]
+                    hit = :outputs
+                  elsif declaration_cache.key?(name)
+                    referenced = declaration_cache[name]
+                    hit = :cache
+                  else
+                    raise "unscheduled ref #{name}: producer not executed or dependency analysis failed"
+                  end
+
                   if ENV["DEBUG_VM_ARGS"]
-                    puts "DEBUG Ref #{name}: #{referenced_value[:k] == :scalar ? "scalar(#{referenced_value[:v].inspect})" : "#{referenced_value[:k]}(#{referenced_value[:rows]&.size || 0} rows)"}"
+                    puts "DEBUG Ref #{name}: #{referenced[:k] == :scalar ? "scalar(#{referenced[:v].inspect})" : "#{referenced[:k]}(#{referenced[:rows]&.size || 0} rows)"}"
+                  end
+
+                  slots << referenced
+                  rows_touched = referenced[:k] == :vec ? (referenced[:rows]&.size || 0) : 1
+                  if t0
+                    Profiler.record!(decl: decl.name, idx: op_index, tag: :ref, op: op, t0: t0, cpu_t0: cpu_t0,
+                                     rows: rows_touched, note: hit)
                   end
-                  slots << referenced_value
 
                 when :array
                   # Validate slot indices before accessing
@@ -216,7 +314,10 @@ module Kumi
                     raise "Store operation '#{name}': source slot #{src} is nil (available slots: #{slots.length}, non-nil slots: #{slots.compact.length})"
                   end
 
-                  outputs[name] = slots[src]
+                  result = slots[src]
+                  outputs[name] = result
+                  # Also store in declaration cache for future ref operations
+                  declaration_cache[name] = result
 
                   # keep slot_id == op_index invariant
                   slots << nil
@@ -327,6 +428,8 @@ module Kumi
               end
             end
 
+            # --- end-of-run summary ---
+            Profiler.emit_summary! if Profiler.enabled?
             outputs
           end
         end