kumi 0.0.18 → 0.0.20

data/lib/kumi/core/analyzer/passes/lir/validation_pass.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        module LIR
+          class ValidationPass < PassBase
+            def run(_errors)
+              ops_by_decl = get_state(:lir_module)
+
+              ops_by_decl.each do |decl, payload|
+                validate_local_defs!(Array(payload[:operations]), decl)
+                validate_single_yield!(Array(payload[:operations]), decl)
+              end
+
+              state
+            end
+
+            private
+
+            def validate_local_defs!(ops, decl_name)
+              defs = Set.new
+              loop_depth = 0
+
+              ops.each_with_index do |ins, _idx|
+                # 1) check all inputs are defined
+                Array(ins.inputs).each do |r|
+                  next if r.nil?
+                  # --- ADD THIS LINE ---
+                  next if r == :__immediate_placeholder__
+                  # --- END OF FIX ---
+                  raise "use-before-def #{r} in #{decl_name}" unless defs.include?(r)
+                end
+
+                # 2) special handling per opcode
+                case ins.opcode
+                when :LoopStart
+                  # collection_register must be defined already (checked above)
+                  # define loop-introduced registers *here*
+                  el = ins.attributes&.[](:as_element)
+                  ix = ins.attributes&.[](:as_index)
+                  defs << el if el
+                  defs << ix if ix
+                  loop_depth += 1
+
+                when :LoopEnd
+                  loop_depth -= 1
+                  raise "unbalanced LoopEnd in #{decl_name}" if loop_depth < 0
+
+                when :DeclareAccumulator
+                  # no regs defined (symbolic name only)
+
+                when :Yield
+                  # input already checked
+
+                else
+                  # 3) normal producers: record their result register as defined
+                  defs << ins.result_register if ins.result_register
+                end
+              end
+
+              raise "unbalanced loops (depth=#{loop_depth}) in #{decl_name}" unless loop_depth.zero?
+            end
+
+            def validate_single_yield!(ops, decl_name)
+              yi = ops.index { _1.opcode == :Yield }
+              raise "no Yield in #{decl_name}" unless yi
+
+              # exactly one Yield
+              raise "multiple Yields in #{decl_name}" if ops.each_with_index.any? { |ins, i| ins.opcode == :Yield && i != yi }
+
+              # after Yield, only structural LoopEnd is allowed
+              trailing = ops[(yi + 1)..] || []
+              bad = trailing.reject { _1.opcode == :LoopEnd }
+              raise "instructions after Yield in #{decl_name}" unless bad.empty?
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/analyzer/passes/load_input_cse.rb

@@ -10,11 +10,11 @@ module Kumi
         # were already stored by earlier declarations.
         #
         # OPTIMIZATION STRATEGY:
-        # - Cross-declaration load reuse: If a load_input with the same 
-        #   (plan_id, scope, is_scalar, has_idx) was already stored by an 
-        #   earlier declaration, rewrite later identical loads to ref the 
+        # - Cross-declaration load reuse: If a load_input with the same
+        #   (plan_id, scope, is_scalar, has_idx) was already stored by an
+        #   earlier declaration, rewrite later identical loads to ref the
         #   stored value instead of re-loading.
-        # - Only reuses producers that appear earlier in module order 
+        # - Only reuses producers that appear earlier in module order
         #   (no reordering/hoisting).
         # - Safe because interpreter's outputs persist across declarations
         #   and ref operations resolve previously stored values.
@@ -26,11 +26,11 @@ module Kumi
         # DEBUG:
         # - Set DEBUG_LOAD_CSE=1 to see optimization decisions
         class LoadInputCSE < PassBase
-          def run(errors)
+          def run(_errors)
             ir = get_state(:ir_module, required: true)
             return state unless ir&.decls
 
-            debug = ENV["DEBUG_LOAD_CSE"]
+            debug = ENV.fetch("DEBUG_LOAD_CSE", nil)
 
             # Map: key -> { name:, decl_index: }
             producers = {}
@@ -41,14 +41,14 @@ module Kumi
             ir.decls.each_with_index do |decl, di|
               decl.ops.each_with_index do |op, oi|
                 next unless op.tag == :load_input
-                
+
                 key = load_key(op)
                 # Does this decl store that slot under a name?
                 store_name = name_storing_slot(decl.ops, oi)
                 next unless store_name
-                
+
                 # Keep earliest producer only
-                if !producers.key?(key)
+                unless producers.key?(key)
                   producers[key] = { name: store_name, decl_index: di }
                   puts "LOAD_CSE: Found producer #{store_name} in decl #{di} for key #{key.inspect}" if debug
                 end
@@ -62,10 +62,10 @@ module Kumi
             new_decls = ir.decls.each_with_index.map do |decl, di|
               new_ops = decl.ops.each_with_index.map do |op, oi|
                 next op unless op.tag == :load_input
-                
+
                 key = load_key(op)
                 prod = producers[key]
-                
+
                 # Only rewrite if producer is in an earlier decl
                 if prod && prod[:decl_index] < di
                   optimizations += 1
@@ -75,11 +75,11 @@ module Kumi
                   op
                 end
               end
-              
+
               Kumi::Core::IR::Decl.new(
-                name: decl.name, 
-                kind: decl.kind, 
-                shape: decl.shape, 
+                name: decl.name,
+                kind: decl.kind,
+                shape: decl.shape,
                 ops: new_ops
               )
             end
@@ -108,6 +108,7 @@ module Kumi
           def name_storing_slot(ops, slot_id)
             ops.each do |op|
               next unless op.tag == :store
+
               src = op.args && op.args[0]
               return op.attrs[:name] if src == slot_id
             end
@@ -117,4 +118,4 @@ module Kumi
       end
     end
   end
-end
+end

data/lib/kumi/core/analyzer/passes/lower_to_ir_pass.rb

@@ -294,7 +294,7 @@ module Kumi
             end
           end
 
-          def lower_declaration(name, decl, access_plans, join_reduce_plans, scope_plan)
+          def lower_declaration(name, decl, access_plans, _join_reduce_plans, scope_plan)
             ops = []
 
             plan = @join_reduce_plans[name]
@@ -310,7 +310,7 @@ module Kumi
               end
 
             last_slot = lower_expression(decl.expression, ops, access_plans, scope_plan,
-                                         need_indices = true, req_scope)
+                                         true, req_scope)
 
             # Apply broadcasting for scalar-to-vector join plans
             if plan && plan.respond_to?(:policy) && plan.policy == :broadcast
@@ -523,6 +523,8 @@ module Kumi
                 elem_slots = insert_align_to_if_needed(elem_slots, ops, access_plans, on_missing: :error) unless target.empty?
                 ops << Kumi::Core::IR::Ops.Array(elem_slots.size, *elem_slots)
                 return ops.size - 1
+              when Syntax::HashExpression
+                # Skip
 
               when Syntax::CascadeExpression
                 # Find a base (true) case, if present

data/lib/kumi/core/analyzer/passes/lower_to_irv2_pass.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+require_relative "../../irv2/builder"
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        class LowerToIRV2Pass < PassBase
+          # In:  state[:snast_module], state[:evaluation_order], state[:input_table]
+          # Out: state[:irv2_module]
+          def run(errors)
+            snast = get_state(:snast_module, required: true)
+            order = get_state(:evaluation_order, required: true)
+            @input_table = get_state(:input_table, required: true)
+
+            builder = Kumi::Core::IRV2::Builder.new
+            declarations = {}
+
+            order.each do |decl_name|
+              declarations[decl_name] = lower_one_declaration(decl_name, snast, @input_table, builder, errors)
+            end
+
+            irv2_module = IRV2::Module.new(declarations, {})
+            state.with(:irv2_module, irv2_module)
+          end
+
+          private
+
+          def lower_one_declaration(decl_name, snast, _input_table, builder, errors)
+            decl = snast.decls.fetch(decl_name)
+
+            # Temporarily set instance variables for the context of this declaration.
+            @current_decl_name = decl_name
+            @b = builder # Use the shared builder
+            @input_cache = {}
+            @current_input_table = @input_table
+            @current_decls = snast.decls
+
+            start_op_count = builder.values.length
+            ir_val = lower_expr(decl.body, errors)
+
+            # NOTE: The optimizer and parameter collection now need to operate on a slice of the builder's operations.
+            new_ops = builder.values.slice(start_op_count..-1)
+            parameters = collect_declaration_parameters(new_ops, snast.decls)
+
+            # The declaration now just stores references to the globally unique operations.
+            IRV2::Declaration.new(decl_name, new_ops, ir_val, parameters)
+          ensure
+            # Clean up instance variables
+            @current_decl_name = @b = @input_cache = @current_input_table = @current_decls = nil
+          end
+
+          def fqid(val)
+            "#{@current_decl_name}/#{val.id}"
+          end
+
+          def ser_stamp(snast_stamp)
+            {
+              "dtype" => (snast_stamp["dtype"] || snast_stamp[:dtype]).to_s,
+              "axes" => Array(snast_stamp["axes"] || snast_stamp[:axes]).map(&:to_s)
+            }
+          end
+
+          def lower_expr(node, errors)
+            ir_val =
+              case node
+              when Kumi::Core::NAST::Const
+                stamp = ser_stamp(node.meta[:stamp])
+                @b.const(node.value, stamp: stamp)
+
+              when Kumi::Core::NAST::InputRef
+                plan = @input_table.find { |i| i.path_fqn == node.path_fqn }
+                unless @input_cache[node.path_fqn]
+                  axes  = plan.axes
+                  dtype = plan.dtype
+                  stamp = { "dtype" => dtype.to_s, "axes" => axes.map(&:to_s) }
+                  @input_cache[node.path_fqn] = @b.load_input(node.path_fqn, stamp: stamp)
+                end
+                @input_cache[node.path_fqn]
+
+              when Kumi::Core::NAST::Ref
+                dep_name = node.name
+                stamp = ser_stamp(@current_decls[dep_name].meta[:stamp])
+                @b.load_declaration(dep_name.to_s, stamp: stamp)
+
+              when Kumi::Core::NAST::Tuple
+                node.meta[:plan] or (errors << "Missing plan on Tuple"
+                                     return @b.const(nil))
+
+                stamp = ser_stamp(node.meta[:stamp])
+                lowered = node.elements.map { |e| lower_expr(e, errors) }
+                node.elements.map { |elem| ser_stamp(elem.meta[:stamp]) }
+
+                @b.construct_tuple(*lowered, stamp: stamp)
+
+              when Kumi::Core::NAST::Call
+                plan = node.meta[:plan] or (errors << "Missing plan on Call #{node.fn}"
+                                            return @b.const(nil))
+
+                case plan[:kind]
+                when :elementwise
+                  args  = node.args.map { |a| lower_expr(a, errors) }
+                  stamp = ser_stamp(node.meta[:stamp])
+                  if node.fn == BUILTIN_SELECT
+                    @b.select(*args, stamp: stamp)
+                  else
+                    @b.map(node.fn.to_s, *args, stamp: stamp)
+                  end
+
+                when :reduce
+                  unless node.args.size == 1
+                    errors << "Reducer arity must be 1, got #{node.args.size}"
+                    return @b.const(nil)
+                  end
+                  v = lower_expr(node.args.first, errors)
+                  stamp = ser_stamp(node.meta[:stamp])
+                  @b.reduce(node.fn.to_s, v, plan[:last_axis_token], stamp: stamp)
+
+                else
+                  errors << "Unsupported call kind for lowering: #{plan[:kind].inspect}"
+                  return @b.const(nil)
+                end
+
+              else
+                errors << "Unhandled SNAST node in lowering: #{node.class}"
+                return @b.const(nil)
+              end
+
+            ir_val or (errors << "Missing IR value from lowering"
+                       @b.const(nil))
+          end
+
+          def collect_declaration_parameters(operations, _decls)
+            parameters = []
+
+            operations.select { |op| op.op == :LoadInput }.each do |op|
+              path = op.args.first
+              next if parameters.any? { |p| p[:type] == :input && p[:source] == path }
+
+              parameters << { type: :input, source: path }
+            end
+
+            operations.select { |op| op.op == :LoadDeclaration }.each do |op|
+              dep_name = op.args.first
+              next if parameters.any? { |p| p[:type] == :dependency && p[:source] == dep_name }
+
+              parameters << { type: :dependency, source: dep_name }
+            end
+
+            parameters
+          end
+
+          def optimize_declaration(builder)
+            value_map = {}
+
+            builder.values.each do |val|
+              signature = operation_signature(val)
+              if value_map[signature]
+                original = value_map[signature]
+                replace_references(builder.values, val, original)
+              else
+                value_map[signature] = val
+              end
+            end
+
+            builder.values.reject! do |val|
+              signature = operation_signature(val)
+              value_map[signature] != val
+            end
+          end
+
+          def operation_signature(val)
+            case val.op
+            when :LoadDeclaration, :LoadInput, :Const
+              [val.op, val.args.first]
+            when :AlignTo
+              [val.op, val.args.first.id, Array(val.attrs[:target_axes]).map(&:to_s)]
+            when :Map
+              [val.op, val.attrs[:fn], val.args.map(&:id)]
+            when :Reduce
+              [val.op, val.args.first.id, val.attrs[:axis].to_s, val.attrs[:fn]]
+            else
+              [val.op, val.args.map(&:id), val.attrs]
+            end
+          end
+
+          def replace_references(operations, old_val, new_val)
+            operations.each do |op|
+              op.args.map! { |arg| arg == old_val ? new_val : arg }
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/analyzer/passes/name_indexer.rb

@@ -7,17 +7,20 @@ module Kumi
         # RESPONSIBILITY: Build definitions index and detect duplicate names
         # DEPENDENCIES: None (first pass in pipeline)
         # PRODUCES: :declarations - Hash mapping names to declaration nodes
+        #            - annotates hints to declarations (e.g. inlining)
         # INTERFACE: new(schema, state).run(errors)
         class NameIndexer < PassBase
           def run(errors)
             definitions = {}
+            hints = {}
 
             each_decl do |decl|
               report_error(errors, "duplicated definition `#{decl.name}`", location: decl.loc) if definitions.key?(decl.name)
               definitions[decl.name] = decl
+              hints[decl.name] = decl.hints
             end
 
-            state.with(:declarations, definitions.freeze)
+            state.with(:declarations, definitions.freeze).with(:hints, hints)
           end
         end
       end

data/lib/kumi/core/analyzer/passes/nast_dimensional_analyzer_pass.rb

@@ -0,0 +1,237 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        # Extracts dimensional and type metadata from NAST tree
+        # Uses minimal function specs to resolve Call nodes and propagate types
+        #
+        # Input:  state[:nast_module], state[:input_table] (chain-free; must include :axes and :dtype per path)
+        # Output: state[:metadata_table], state[:declaration_table]
+        class NASTDimensionalAnalyzerPass < PassBase
+          def run(errors)
+            nast_module  = get_state(:nast_module, required: true)
+            @input_table = get_state(:input_table, required: true)
+            @registry = get_state(:registry, required: true)
+
+            @function_specs    = Functions::Loader.load_minimal_functions
+
+            @metadata_table    = {}
+            @declaration_table = {}
+
+            debug "Analyzing NAST module with #{nast_module.decls.size} declarations"
+            debug "Function specs loaded: #{@function_specs.keys.join(', ')}"
+
+            nast_module.decls.each { |name, decl| analyze_declaration(name, decl, errors) }
+
+            debug "Generated metadata_table with #{@metadata_table.size} entries"
+            debug "Generated declaration_table with #{@declaration_table.size} entries"
+
+            state.with(:metadata_table, @metadata_table.freeze)
+                 .with(:declaration_table, @declaration_table.freeze)
+          end
+
+          private
+
+          def analyze_declaration(name, decl, errors)
+            debug "Analyzing #{name}"
+            result_metadata = analyze_expression(decl.body, errors)
+
+            decl_metadata = {
+              kind: decl.kind,
+              result_type: result_metadata[:type],
+              result_scope: result_metadata[:scope],
+              target_name: name
+            }.freeze
+
+            @metadata_table[node_id(decl)] = decl_metadata
+            @declaration_table[name] = decl_metadata
+
+            debug "  #{name}: #{result_metadata[:type]} in scope #{result_metadata[:scope].inspect}"
+          end
+
+          def analyze_expression(expr, errors)
+            case expr
+            when Kumi::Core::NAST::Call         then analyze_call_expression(expr, errors)
+            when Kumi::Core::NAST::Tuple        then analyze_tuple(expr, errors)
+            when Kumi::Core::NAST::InputRef     then analyze_input_ref(expr)
+            when Kumi::Core::NAST::IndexRef     then analyze_index_ref(expr, errors)
+            when Kumi::Core::NAST::Const        then analyze_const(expr)
+            when Kumi::Core::NAST::Pair         then analyze_pair(expr, errors)
+            when Kumi::Core::NAST::Ref          then analyze_declaration_ref(expr)
+            when Kumi::Core::NAST::Hash         then analyze_hash(expr, errors)
+
+            else
+              raise "Unknown NAST node type: #{expr.class}"
+            end
+          end
+
+          def analyze_call_expression(call, errors)
+            function_spec = @registry.function(call.fn.to_s)
+
+            arg_metadata = call.args.map { |arg| analyze_expression(arg, errors) }
+            arg_types    = arg_metadata.map { |m| m[:type] }
+            arg_scopes   = arg_metadata.map { |m| m[:scope] }
+
+            debug "    Call #{call.fn}: arg_scopes=#{arg_scopes.inspect}"
+
+            if ENV["DEBUG_NAST_DIMENSIONAL_ANALYZER"] == "1" && function_spec.param_names.size != arg_types.size
+              puts "[NASTDimensionalAnalyzer] WARNING: #{call.fn} expects #{function_spec.param_names.size} args, got #{arg_types.size}"
+            end
+
+            named_types =
+              if function_spec.params.size == arg_types.size
+                Hash[function_spec.param_names.zip(arg_types)]
+              else
+                { function_spec.param_names.first => arg_types } # variadic TODO: this is a hack
+              end
+
+            begin
+              result_type = function_spec.dtype_rule.call(named_types)
+            rescue StandardError
+              # Maybe we have the wrong function, lets try to see if another function with same name works
+              # TODO: Fix this hack
+
+              raise
+            end
+
+            over_collection = arg_types.size == 1 && Types.collection?(arg_types[0])
+            result_scope = compute_result_scope(function_spec, arg_scopes, over_collection)
+
+            @metadata_table[node_id(call)] = {
+              function: function_spec.id,
+              kind: function_spec.kind,
+              params: function_spec.params,
+              result_type: result_type,
+              result_scope: result_scope,
+              arg_types: arg_types,
+              arg_scopes: arg_scopes,
+              last_axis_token: (function_spec.kind == :reduce ? (arg_scopes.first || []).last : nil) # TODO: REMOVE
+            }.freeze
+
+            debug "    Call #{function_spec.id}: (#{arg_types.join(', ')}) -> #{result_type} in #{result_scope.inspect}"
+            { type: result_type, scope: result_scope }
+          end
+
+          def analyze_tuple(node, errors)
+            elems           = node.args.map { |e| analyze_expression(e, errors) }
+            element_types   = elems.map { |m| m[:type] }
+            element_scopes  = elems.map { |m| m[:scope] }
+            result_scope    = lub_by_prefix(element_scopes)
+
+            result_type = if element_types.uniq.size == 1
+                            "tuple<#{element_types.uniq[0]}>"
+                          else
+                            "tuple<#{element_types.join(', ')}>"
+                          end
+
+            @metadata_table[node_id(node)] = {
+              parameter_names: [],
+              result_type: result_type,
+              result_scope: result_scope,
+              arg_types: element_types,
+              arg_scopes: element_scopes,
+              last_axis_token: nil
+            }.freeze
+
+            debug "    Tuple: (#{element_types.join(', ')}) -> #{result_type} in #{result_scope.inspect}"
+            { type: result_type, scope: result_scope }
+          end
+
+          def analyze_hash(node, errors)
+            fields = node.pairs.map { |e| analyze_expression(e, errors) }
+            fields_scopes = fields.map { |m| m[:scope] }
+            scope = lub_by_prefix(fields_scopes)
+            dtype = :hash
+
+            @metadata_table[node_id(node)] = {
+              type: dtype,
+              scope: scope
+            }.freeze
+          end
+
+          def analyze_pair(node, errors)
+            value_node = analyze_expression(node.value, errors)
+            dtype = :pair
+
+            @metadata_table[node_id(node)] = {
+              type: dtype,
+              scope: value_node[:scope]
+            }.freeze
+          end
+
+          # STRICT: requires entry with :axes and :dtype (no fallbacks)
+          def analyze_input_ref(input_ref)
+            entry = @input_table.find { |imp| imp[:path_fqn] == input_ref.path_fqn }
+            entry or raise KeyError, "Input path not found in input_table: #{input_ref.path_fqn}"
+
+            axes  = entry.axes
+            dtype = entry.dtype
+
+            { type: dtype, scope: axes }
+          end
+
+          def analyze_const(node)
+            type = Types.normalize(node.value.class)
+            meta = { type: type, scope: [] }
+
+            @metadata_table[node_id(node)] = meta
+          end
+
+          def analyze_index_ref(node, _errors)
+            meta = @input_table.find { _1.path_fqn == node.input_fqn } or raise "Index plan found: #{n.name.inspect}"
+            axes = Array(meta[:axes])
+            type = :integer
+
+            debug "    IndexRef #{node.name}: input_fqn=#{node.input_fqn}, axes=#{axes.inspect}"
+
+            @metadata_table[node_id(node)] = { type:, scope: axes }.freeze
+            { type:, scope: axes }
+          end
+
+          def analyze_declaration_ref(ref)
+            meta = @declaration_table.fetch(ref.name)
+            @metadata_table[node_id(ref)] = {
+              kind: :ref,
+              result_type: meta[:result_type],
+              result_scope: meta[:result_scope],
+              referenced_name: ref.name
+            }.freeze
+            { type: meta[:result_type], scope: meta[:result_scope] }
+          end
+
+          def compute_result_scope(function_spec, arg_scopes, over_collection = false)
+            case function_spec.kind
+            when :elementwise, :constructor
+              lub_by_prefix(arg_scopes)
+            when :reduce
+              if over_collection
+                lub_by_prefix(arg_scopes)
+              else
+                child = arg_scopes.first || []
+                child[0...-1]
+              end
+            else
+              []
+            end
+          end
+
+          def lub_by_prefix(list)
+            return [] if list.empty?
+
+            candidate = list.max_by(&:length)
+            list.each do |axes|
+              unless axes.each_with_index.all? { |tok, i| candidate[i] == tok }
+                raise Kumi::Core::Errors::SemanticError, "prefix mismatch: #{axes.inspect} vs #{candidate.inspect}"
+              end
+            end
+            candidate
+          end
+
+          def node_id(node) = "#{node.class}_#{node.id}"
+        end
+      end
+    end
+  end
+end