kumi 0.0.17 → 0.0.19

data/lib/kumi/core/lir/analyze.rb

@@ -0,0 +1,64 @@
+# lib/kumi/core/lir/analyze.rb
+# Single-pass context builder. Derives loop tables and per-register lineage.
+module Kumi
+  module Core
+    module LIR
+      module Analyze
+        module_function
+
+        # Returns:
+        # - loop_table:  { loop_id => { id:, axis:, input:, start:, end: } }
+        # - active_axes: Array<Array<Symbol>> per instruction index
+        # - active_ids:  Array<Array<Symbol>> per instruction index
+        # - definition_table: { reg => index }
+        # - use_table:   { reg => [indices...] }
+        # - register_axes: { reg => [:axis...] } lineage at def site
+        # - register_ids:  { reg => [:L1,:L2,...] } loop ids at def site
+        def context(instructions, ids: nil)
+          frames = []
+          loop_table = {}
+          active_axes = Array.new(instructions.length) { [] }
+          active_ids  = Array.new(instructions.length) { [] }
+          definition_table = {}
+          use_table = Hash.new { |h, k| h[k] = [] }
+
+          instructions.each_with_index do |instruction, idx|
+            definition_table[instruction.result_register] = idx if instruction.result_register
+            instruction.inputs&.each { |r| use_table[r] << idx }
+
+            case instruction.opcode
+            when :LoopStart
+              id = (instruction.attributes[:id] ||= (ids || Ids.new).generate_loop_id)
+              frames << { id:, axis: instruction.attributes[:axis], input: instruction.inputs&.first, start: idx }
+            when :LoopEnd
+              frame = frames.pop or raise Error, "LoopEnd without frame at #{idx}"
+              loop_table[frame[:id]] = frame.merge(end: idx)
+            end
+
+            active_axes[idx] = frames.map { _1[:axis] }
+            active_ids[idx]  = frames.map { _1[:id] }
+          end
+
+          raise Error, "unclosed loops" unless frames.empty?
+
+          register_axes = {}
+          register_ids  = {}
+          definition_table.each do |reg, def_idx|
+            register_axes[reg] = active_axes[def_idx]
+            register_ids[reg]  = active_ids[def_idx]
+          end
+
+          {
+            loop_table:,
+            active_axes:,
+            active_ids:,
+            definition_table:,
+            use_table:,
+            register_axes:,
+            register_ids:
+          }
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/build.rb

@@ -0,0 +1,363 @@
+# lib/kumi/core/lir/build.rb
+# Builders for each opcode. Keep params clear. Carry optional location.
+module Kumi
+  module Core
+    module LIR
+      module Build
+        module_function
+
+        # Constant
+        # Params:
+        #   value:     literal value
+        #   dtype:     dtype of the literal
+        #   as:        result register symbol/name
+        #   location:  optional Location
+        # Result: produces(result_register, stamp(dtype))
+        def constant(value:, dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :Constant,
+            result_register: as,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [],
+            immediates: [Literal.new(value: value, dtype: dtype)],
+            attributes: {},
+            location:
+          )
+        end
+
+        # LoadInput
+        # Params:
+        #   key:       String|Symbol top-level input key
+        #   dtype:     dtype of the loaded field
+        #   as:        result register
+        #   location:  optional Location
+        # Result: produces
+        def load_input(key:, dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :LoadInput,
+            result_register: as,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [],
+            immediates: [Literal.new(value: key.to_s, dtype: :string)],
+            attributes: {},
+            location:
+          )
+        end
+
+        # LoadDeclaration
+        # Params:
+        #   name:      declaration name
+        #   dtype:     dtype of referenced declaration result
+        #   axes:      axes of referenced declaration
+        #   as:        result register
+        #   location:  optional Location
+        # Result: produces
+        def load_declaration(name:, dtype:, axes:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :LoadDeclaration,
+            result_register: as,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [],
+            immediates: [Literal.new(value: name.to_s, dtype: :string)],
+            attributes: { axes: Array(axes).map!(&:to_sym) },
+            location:
+          )
+        end
+
+        # LoadField
+        # Params:
+        #   object_register: register holding a Hash-like object
+        #   key:             field key to access
+        #   dtype:           dtype of the field
+        #   as:              result register
+        #   location:        optional Location
+        # Result: produces
+        def load_field(object_register:, key:, dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :LoadField,
+            result_register: as,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [object_register],
+            immediates: [Literal.new(value: key.to_s, dtype: :string)],
+            attributes: {},
+            location:
+          )
+        end
+
+        # LoopStart
+        # Params:
+        #   collection_register: register holding the array to iterate
+        #   axis:                Symbol axis name
+        #   as_element:          register name for the element within loop body
+        #   as_index:            register name for the numeric index
+        #   id:                  optional loop id; generated if nil
+        #   location:            optional Location
+        # Result: does not produce
+        def loop_start(collection_register:, axis:, as_element:, as_index:, id: nil, ids: nil, location: nil)
+          Instruction.new(
+            opcode: :LoopStart,
+            result_register: nil,
+            stamp: nil,
+            inputs: [collection_register],
+            immediates: [],
+            attributes: { axis: axis.to_sym, as_element:, as_index:, id: id || ids.generate_loop_id },
+            location:
+          )
+        end
+
+        # LoopEnd
+        # Params:
+        #   location: optional Location
+        # Result: does not produce
+        def loop_end(location: nil)
+          Instruction.new(
+            opcode: :LoopEnd,
+            result_register: nil,
+            stamp: nil,
+            inputs: [],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        # KernelCall
+        # Params:
+        #   function:   String function id (e.g., "core.mul")
+        #   args:       Array of register names
+        #   out_dtype:  dtype of the result
+        #   as:         result register
+        #   location:   optional Location
+        # Result: produces
+        def kernel_call(function:, args:, out_dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :KernelCall,
+            result_register: as,
+            stamp: Stamp.new(dtype: out_dtype),
+            inputs: args,
+            immediates: [],
+            attributes: { fn: function.to_s },
+            location:
+          )
+        end
+
+        # Fold
+        # Params:
+        #   function:   String function id (e.g., "core.mul")
+        #   arg:        Collection being folded register
+        #   out_dtype:  dtype of the result
+        #   as:         result register
+        #   location:   optional Location
+        # Result: produces
+        def fold(function:, arg:, out_dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :Fold,
+            result_register: as,
+            stamp: Stamp.new(dtype: out_dtype),
+            inputs: Array(arg),
+            immediates: [],
+            attributes: { fn: function.to_s },
+            location:
+          )
+        end
+
+        # Select
+        # Params:
+        #   cond:       register name (boolean)
+        #   on_true:    register name
+        #   on_false:   register name
+        #   out_dtype:  dtype of the result
+        #   as:         result register
+        #   location:   optional Location
+        # Result: produces
+        def select(cond:, on_true:, on_false:, out_dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :Select,
+            result_register: as,
+            stamp: Stamp.new(dtype: out_dtype),
+            inputs: [cond, on_true, on_false],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        # DeclareAccumulator
+        # Params:
+        #   name:     Symbol accumulator name
+        #   initial:  Literal identity value
+        #   location: optional Location
+        # Result: does not produce
+        def declare_accumulator(dtype:, location: nil, ids: nil, name: nil)
+          name ||= ids.generate_acc
+          Instruction.new(
+            opcode: :DeclareAccumulator,
+            result_register: name.to_sym,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        # Accumulate
+        # Params:
+        #   accumulator: Symbol accumulator name
+        #   function:    String function id (e.g., "core.add")
+        #   value_register: register providing the value to accumulate
+        #   location:      optional Location
+        # Result: does not produce
+        def accumulate(accumulator:, function:, value_register:, dtype:, location: nil)
+          Instruction.new(
+            opcode: :Accumulate,
+            result_register: accumulator.to_sym,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [value_register],
+            immediates: [],
+            attributes: { fn: function.to_s },
+            location:
+          )
+        end
+
+        # LoadAccumulator
+        # Params:
+        #   name:     Symbol accumulator name
+        #   dtype:    dtype of the accumulator
+        #   as:       result register
+        #   location: optional Location
+        # Result: produces
+        def load_accumulator(accumulator:, dtype:, ids:, location: nil)
+          as = ids.generate_temp
+          Instruction.new(
+            opcode: :LoadAccumulator,
+            result_register: as,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [accumulator.to_sym],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        # MakeTuple
+        # Params:
+        #   elements: Array of registers
+        #   out_dtype: tuple dtype
+        #   as: result register
+        # Result: produces tuple
+        def make_tuple(elements:, out_dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :MakeTuple,
+            result_register: as,
+            stamp: Stamp.new(dtype: out_dtype),
+            inputs: elements,
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        # MakeObject
+        # Params:
+        #   keys:   Array<String|Symbol> in the same order as values
+        #   values: Array of registers
+        #   out_dtype: object dtype (or :object)
+        #   as: result
+        def make_object(keys:, values:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :MakeObject,
+            result_register: as,
+            stamp: Stamp.new(dtype: :object),
+            inputs: values,
+            immediates: keys.map { |k| Literal.new(value: k.to_s, dtype: :string) },
+            attributes: {},
+            location: location
+          )
+        end
+
+        # TupleGet (optional)
+        # Params:
+        #   tuple: register
+        #   index: Integer
+        #   out_dtype: element dtype
+        #   as: result
+        def tuple_get(tuple:, index:, out_dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :TupleGet,
+            result_register: as,
+            stamp: Stamp.new(dtype: out_dtype),
+            inputs: [tuple],
+            immediates: [Literal.new(value: Integer(index), dtype: :integer)],
+            attributes: {},
+            location:
+          )
+        end
+
+        def length(collection_register:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :Length,
+            result_register: as,
+            stamp: Stamp.new(dtype: :integer),
+            inputs: [collection_register],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        def gather(collection_register:, index_register:, dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :Gather,
+            result_register: as,
+            stamp: Stamp.new(dtype: dtype),
+            inputs: [collection_register, index_register],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+
+        # Yield
+        # Opcode: Yield
+        # Semantics:
+        # - Exactly one per declaration.
+        # - Must be the last instruction in the declaration.
+        # - The declaration's result axes are Γ, the active loop stack at the Yield site
+        #   (i.e., the sequence of surrounding LoopStart frames).
+        # - The yielded register's stamp.dtype is the declaration's result dtype.
+        # Codegen rule:
+        # - If Γ == [], return the yielded scalar.
+        # - If Γ != [], materialize a container shaped by Γ and write the yielded value
+        #   at each iteration of the surrounding loops.
+        # Params:
+        #   result_register: register that holds the final value
+        #   location:        optional Location
+        # Result: does not produce
+        def yield(result_register:, location: nil)
+          Instruction.new(
+            opcode: :Yield,
+            result_register: nil,
+            stamp: nil,
+            inputs: [result_register],
+            immediates: [],
+            attributes: {},
+            location:
+          )
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/emit.rb

@@ -0,0 +1,62 @@
+module Kumi
+  module Core
+    module LIR
+      class Emit
+        def initialize(registry:, ids:, ops:)
+          (@registry = registry
+           @ids = ids
+           @ops = ops)
+        end
+
+        def iconst(v) = push(Build.constant(value: Integer(v), dtype: :integer, ids: @ids))
+
+        # Kernel bridge with explicit function ids
+        def k_id(fn_id, args, out:)
+          ins = Build.kernel_call(function: @registry.resolve_function(fn_id),
+                                  args:, out_dtype: out, ids: @ids)
+          push(ins)
+        end
+
+        # Integer arithmetic via canonical ids
+        def add_i(a, b) = k_id("core.add", [a, b], out: :integer)
+        def sub_i(a, b) = k_id("core.sub", [a, b], out: :integer)
+        def mod_i(a, b) = k_id("core.mod", [a, b], out: :integer)
+        def lt(a, b)    = k_id("core.lt",  [a, b], out: :boolean)
+        def gt(a, b)    = k_id("core.gt",  [a, b], out: :boolean)
+        def le(a, b)   = k_id("core.lte", [a, b], out: :boolean)
+        def ge(a, b)   = k_id("core.gte", [a, b], out: :boolean)
+        def and_(a, b) = k_id("core.and", [a, b], out: :boolean)
+
+        # Clamp via canonical id
+        def clamp(x, lo, hi, out:) = k_id("core.clamp", [x, lo, hi], out:)
+
+        def length(coll) = push(Build.length(collection_register: coll, ids: @ids))
+        def gather(arr, i, dt) = push(Build.gather(collection_register: arr, index_register: i, dtype: dt, ids: @ids))
+        def select(c, t, f, dt) = push(Build.select(cond: c, on_true: t, on_false: f, out_dtype: dt, ids: @ids))
+
+        def const(value, dtype)
+          push(Build.constant(value: value, dtype: dtype, ids: @ids))
+        end
+
+        def wrap_index(i, off, nlen)
+          i_plus = add_i(i, iconst(off))
+          m1     = mod_i(i_plus, nlen)
+          p2     = add_i(m1, nlen)
+          mod_i(p2, nlen)
+        end
+
+        def clamp_index(j, nlen)
+          hi = add_i(nlen, iconst(-1))
+          clamp(j, iconst(0), hi, out: :integer)
+        end
+
+        private
+
+        def push(ins)
+          (@ops << ins
+           ins.result_register)
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/structs/instruction.rb

@@ -0,0 +1,44 @@
+# lib/kumi/core/lir/structs/instruction.rb
+# Fields:
+# - opcode, result_register, stamp, inputs, immediates, attributes, location
+module Kumi
+  module Core
+    module LIR
+      module Structs
+        Instruction = Struct.new(
+          :opcode, :result_register, :stamp, :inputs, :immediates, :attributes, :location,
+          keyword_init: true
+        ) do
+          SIDE_EFFECT_OPCODES = Set.new(%i[
+                                          LoopStart
+                                          LoopEnd
+                                          DeclareAccumulator
+                                          Accumulate LoadAccumulator
+                                          Yield
+                                        ]).freeze
+
+          def produces? = !result_register.nil?
+
+          def side_effect?
+            SIDE_EFFECT_OPCODES.include?(opcode)
+          end
+
+          def pure?
+            !side_effect?
+          end
+
+          def to_h
+            h = { op: opcode }
+            h[:result]     = result_register if result_register
+            h[:stamp]      = stamp.to_h if stamp
+            h[:inputs]     = inputs unless inputs.nil? || inputs.empty?
+            h[:immediates] = immediates&.map { |x| x.respond_to?(:to_h) ? x.to_h : x } unless immediates.nil? || immediates.empty?
+            h[:attrs]      = attributes unless attributes.nil? || attributes.empty?
+            h[:loc] = { file: location.file, line: location.line, column: location.column } if location
+            h
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/structs/literal.rb

@@ -0,0 +1,14 @@
+# lib/kumi/core/lir/structs/literal.rb
+module Kumi
+  module Core
+    module LIR
+      module Structs
+        Literal = Struct.new(:value, :dtype, keyword_init: true) do
+          # include Typed
+          def value? = true
+          def to_h = { value: value, dtype: dtype }
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/structs/stamp.rb

@@ -0,0 +1,13 @@
+# lib/kumi/core/lir/structs/stamp.rb
+module Kumi
+  module Core
+    module LIR
+      module Structs
+        Stamp = Struct.new(:dtype, keyword_init: true) do
+          # include Typed   # enable if you add typed.rb
+          def to_h = { dtype: dtype }
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/support/error.rb

@@ -0,0 +1,9 @@
+module Kumi
+  module Core
+    module LIR
+      module Support
+        class Error < StandardError; end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/support/ids.rb

@@ -0,0 +1,34 @@
+# lib/kumi/core/lir/ids.rb
+# Ids generates deterministic temp and loop ids. Each compilation context gets its own instance.
+module Kumi
+  module Core
+    module LIR
+      module Support
+        class Ids
+          def initialize
+            reset!
+          end
+
+          def reset!
+            @t = 0
+            @l = 0
+          end
+
+          def generate_temp(prefix: :t)
+            @t += 1
+            :"#{prefix}#{@t}"
+          end
+
+          def generate_acc
+            generate_temp(prefix: :acc)
+          end
+
+          def generate_loop_id
+            @l += 1
+            :"L#{@l}"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/lir/validate.rb

@@ -0,0 +1,74 @@
+# lib/kumi/core/lir/validate.rb
+# Per-declaration and program-level structural checks.
+module Kumi
+  module Core
+    module LIR
+      module Validate
+        module_function
+
+        # Validate one declaration's instruction list
+        def declaration!(instructions)
+          ensure_opcodes!(instructions)
+
+          defs = {}
+          depth = 0
+          yield_seen = false
+
+          instructions.each_with_index do |ins, i|
+            # producer stamps
+            if ins.result_register
+              raise Error, "missing stamp at #{i}" unless ins.stamp.is_a?(Stamp)
+
+              r = ins.result_register
+              raise Error, "redefinition of #{r} at #{i}" if defs.key?(r)
+
+              defs[r] = i
+            end
+
+            case ins.opcode
+            when :LoopStart
+              attrs = ins.attributes || {}
+              raise Error, "LoopStart missing :axis at #{i}" unless attrs[:axis]
+
+              depth += 1
+            when :LoopEnd
+              raise Error, "LoopEnd without LoopStart at #{i}" if depth.zero?
+
+              depth -= 1
+            when :LoadDeclaration
+              raise Error, "LoadDeclaration missing :axes at #{i}" unless Array(ins.attributes[:axes]).is_a?(Array)
+              raise Error, "LoadDeclaration missing stamp at #{i}" unless ins.stamp.is_a?(Stamp)
+            when :Yield
+              raise Error, "multiple Yield (at #{i})" if yield_seen
+
+              yield_seen = true
+              raise Error, "instructions after Yield (at #{i})" unless i == instructions.length - 1
+            end
+          end
+
+          raise Error, "unclosed loops" unless depth.zero?
+          raise Error, "missing Yield" unless yield_seen
+
+          true
+        end
+
+        # Validate whole program: { name => { operations: [...] } }
+        def program!(ops_by_decl)
+          ops_by_decl.each do |_name, h|
+            declaration!(Array(h[:operations]))
+          end
+          true
+        end
+
+        def ensure_opcodes!(instructions)
+          instructions.each_with_index do |ins, i|
+            next if OPCODES.include?(ins.opcode)
+
+            raise Error, "unknown opcode #{ins.opcode.inspect} at #{i}"
+          end
+        end
+        private_class_method :ensure_opcodes!
+      end
+    end
+  end
+end