kumi 0.0.17 → 0.0.19

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

@@ -0,0 +1,156 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        module LIR
+          # LoopFusionPass (Simplified Version)
+          # --------------
+          #
+          # This pass runs AFTER InlineDeclarationsPass AND LocalCSEPass.
+          # The CSE pass is crucial as it normalizes the LIR, ensuring that any
+          # redundant collection loads are eliminated. This simplifies loop fusion
+          # from a data-flow analysis problem to a simple check for adjacent loops
+          # iterating over the *exact same register*.
+          #
+          # A loop pair is considered "fusable" if:
+          #   1. They are "semantically adjacent," meaning they are separated only by
+          #      instructions that can be correctly reordered around the fused loop
+          #      (e.g., DeclareAccumulator, LoadAccumulator).
+          #   2. They iterate over the identical collection register.
+          #
+          class LoopFusionPass < PassBase
+            LIR = Kumi::Core::LIR
+            def run(_errors)
+              fused_module = get_state(:lir_module).transform_values do |decl|
+                # The decl[:name] might not exist in older states, so provide a fallback.
+                decl_name = decl.is_a?(Hash) ? decl.fetch(:name, "anonymous") : "anonymous"
+                debug "\n--- Fusing loops in: #{decl_name} ---"
+                { operations: fuse_loops_in_block(Array(decl[:operations])) }
+              end
+
+              state.with(:lir_module, fused_module).with(:lir_04_1_loop_fusion, fused_module)
+            end
+
+            private
+
+            def fuse_loops_in_block(ops)
+              changed = true
+              current_ops = ops
+              current_ops, changed = run_one_fusion_pass(current_ops) while changed
+              current_ops
+            end
+
+            def run_one_fusion_pass(ops)
+              new_ops = []
+              changed = false
+              i = 0
+
+              while i < ops.length
+                ins1 = ops[i]
+
+                if ins1.opcode == :LoopStart
+                  debug "  > Found LoopStart(#{ins1.attributes[:id]}) at index #{i} over register #{ins1.inputs.first}"
+                  end1_idx = find_matching_loop_end(ops, i)
+
+                  # NEW: Intelligently scan and classify instructions between loops.
+                  pre_fusion_ops, post_fusion_ops, next_ins_idx = scan_and_classify_intervening_ops(ops, end1_idx)
+
+                  if next_ins_idx
+                    ins2 = ops[next_ins_idx]
+                    debug "    - Found next candidate LoopStart(#{ins2.attributes[:id]}) at index #{next_ins_idx}"
+
+                    # CONDITION: Loop collections must be the IDENTICAL register (thanks to CSE).
+                    are_collections_identical = ins1.inputs.first == ins2.inputs.first
+                    debug "    - Checking collections '#{ins1.inputs.first}' vs '#{ins2.inputs.first}'... Identical? #{are_collections_identical}"
+
+                    if are_collections_identical
+                      changed = true
+                      debug "    - SUCCESS: Fusing loops #{ins1.attributes[:id]} and #{ins2.attributes[:id]}."
+
+                      body1 = ops[(i + 1)...end1_idx]
+                      end2_idx = find_matching_loop_end(ops, next_ins_idx)
+                      body2 = ops[(next_ins_idx + 1)...end2_idx]
+
+                      remap = {
+                        ins2.attributes[:as_element] => ins1.attributes[:as_element],
+                        ins2.attributes[:as_index] => ins1.attributes[:as_index]
+                      }
+                      remapped_body2 = remap_registers(body2, remap)
+                      fused_body = fuse_loops_in_block(body1 + remapped_body2)
+
+                      # Correctly re-order the intervening instructions around the fused loop.
+                      new_ops.concat(pre_fusion_ops) # Declarations for loop 2 go before.
+                      new_ops << ins1
+                      new_ops.concat(fused_body)
+                      new_ops << ops[end1_idx]
+                      new_ops.concat(post_fusion_ops) # Loads from loop 1 go after.
+
+                      i = end2_idx + 1
+                      next
+                    else
+                      debug "    - INFO: Collections not identical. No fusion."
+                    end
+                  end
+                end
+
+                new_ops << ins1
+                i += 1
+              end
+
+              [new_ops, changed]
+            end
+
+            # --- NEW: Replaced dumb scan with intelligent classification ---
+            # Scans between loops and classifies instructions into two buckets:
+            # - pre_ops: Ops that must execute before the fused loop (e.g., DeclareAccumulator for loop 2)
+            # - post_ops: Ops that must execute after the fused loop (e.g., LoadAccumulator from loop 1)
+            # If any unmovable instruction is found, it aborts.
+            def scan_and_classify_intervening_ops(ops, start_idx)
+              pre_ops = []
+              post_ops = []
+              (start_idx + 1...ops.length).each do |i|
+                ins = ops[i]
+                case ins.opcode
+                when :LoopStart
+                  return [pre_ops, post_ops, i]
+                when :DeclareAccumulator
+                  pre_ops << ins
+                when :LoadAccumulator
+                  post_ops << ins
+                else
+                  # Found an unmovable instruction; fusion is not possible.
+                  return [[], [], nil]
+                end
+              end
+              [[], [], nil] # No next loop found
+            end
+
+            def find_matching_loop_end(ops, start_index)
+              depth = 1
+              (start_index + 1...ops.length).each do |i|
+                op = ops[i].opcode
+                depth += 1 if op == :LoopStart
+                depth -= 1 if op == :LoopEnd
+                return i if depth.zero?
+              end
+              raise "Unbalanced LoopStart at index #{start_index}"
+            end
+
+            def remap_registers(ops, remap)
+              ops.map do |ins|
+                new_inputs = Array(ins.inputs).map { |r| remap.fetch(r, r) }
+                LIR::Instruction.new(
+                  opcode: ins.opcode, result_register: ins.result_register, stamp: ins.stamp,
+                  inputs: new_inputs, immediates: ins.immediates,
+                  attributes: ins.attributes, location: ins.location
+                )
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+require "set"
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        module LIR
+          class LoopInvariantCodeMotionPass < PassBase
+            LIR = Kumi::Core::LIR
+            MAX_PASSES = 10
+
+            def run(_errors)
+              current_ops = get_state(:lir_module)
+
+              MAX_PASSES.times do
+                new_ops, changed = run_one_pass(current_ops)
+                return state.with(:lir_module, new_ops.freeze).with(:lir_04_loop_invcm, new_ops) unless changed
+
+                current_ops = new_ops
+              end
+
+              raise "LICM did not converge after #{MAX_PASSES} passes."
+            end
+
+            private
+
+            def run_one_pass(ops_by_decl)
+              changed = false
+              new_ops_by_decl = {}
+
+              ops_by_decl.each do |name, payload|
+                debug "\n--- LICM: Processing declaration: #{name} ---"
+                original_ops = Array(payload[:operations])
+                optimized_ops = optimize_block(original_ops)
+                new_ops_by_decl[name] = { operations: optimized_ops }
+                changed ||= (original_ops != optimized_ops)
+              end
+
+              [new_ops_by_decl, changed]
+            end
+
+            def optimize_block(ops, depth = 0)
+              new_ops = []
+              i = 0
+              while i < ops.length
+                ins = ops[i]
+                prefix = "  " * depth
+
+                if ins.opcode == :LoopStart
+                  debug "#{prefix}> Found LoopStart(#{ins.attributes[:id]})"
+                  end_index = find_matching_loop_end(ops, i)
+                  loop_body = ops[(i + 1)...end_index]
+
+                  optimized_nested_body = optimize_block(loop_body, depth + 1)
+
+                  all_hoisted_ops = []
+                  current_body = optimized_nested_body
+                  pass_num = 0
+
+                  while true
+                    pass_num += 1
+                    debug "#{prefix}  - Hoisting Pass ##{pass_num} for Loop(#{ins.attributes[:id]})"
+
+                    defs_in_loop = get_defs_in_ops(current_body)
+                    defs_in_loop.add(ins.attributes[:as_element])
+                    defs_in_loop.add(ins.attributes[:as_index])
+                    debug "#{prefix}    - Defs inside loop: #{defs_in_loop.to_a.sort.join(', ')}"
+
+                    hoisted_this_pass, remaining_body = current_body.partition do |body_ins|
+                      is_invariant(body_ins, defs_in_loop, prefix, depth)
+                    end
+
+                    if hoisted_this_pass.empty?
+                      debug "#{prefix}    - Convergence. No more invariants found."
+                      break
+                    end
+
+                    debug "#{prefix}    - Hoisted #{hoisted_this_pass.size} instruction(s) this pass."
+                    all_hoisted_ops.concat(hoisted_this_pass)
+                    current_body = remaining_body
+                  end
+
+                  new_ops.concat(all_hoisted_ops)
+                  new_ops << ins
+                  new_ops.concat(current_body)
+                  new_ops << ops[end_index]
+                  debug "#{prefix}< Finished LoopStart(#{ins.attributes[:id]})"
+
+                  i = end_index + 1
+                else
+                  new_ops << ins
+                  i += 1
+                end
+              end
+              new_ops
+            end
+
+            def is_invariant(ins, defs_in_loop, prefix, depth)
+              prefix_inner = "#{prefix}      "
+              debug "#{prefix_inner}- Checking: #{ins.result_register || '(no result)'} = #{ins.opcode}(#{ins.inputs.join(', ')})"
+
+              unless ins.pure?
+                debug "#{prefix_inner}  - REJECT: Instruction is not pure."
+                return false
+              end
+
+              invariant_inputs = (Set.new(Array(ins.inputs)) & defs_in_loop).empty?
+
+              unless invariant_inputs
+                offending_defs = (Set.new(Array(ins.inputs)) & defs_in_loop).to_a
+                debug "#{prefix_inner}  - REJECT: Depends on loop-variant registers: #{offending_defs.join(', ')}"
+                return false
+              end
+
+              debug "#{prefix_inner}  - ACCEPT: Hoisting instruction."
+              true
+            end
+
+            def get_defs_in_ops(ops)
+              defs = Set.new
+              ops.each do |ins|
+                defs.add(ins.result_register) if ins.result_register
+                if ins.opcode == :LoopStart
+                  defs.add(ins.attributes[:as_element])
+                  defs.add(ins.attributes[:as_index])
+                end
+              end
+              defs
+            end
+
+            def find_matching_loop_end(ops, start_index)
+              depth = 1
+              (start_index + 1).upto(ops.length - 1) do |i|
+                opcode = ops[i].opcode
+                depth += 1 if opcode == :LoopStart
+                depth -= 1 if opcode == :LoopEnd
+                return i if depth.zero?
+              end
+              raise "Unbalanced LoopStart at index #{start_index}"
+            end
+          end
+        end
+      end
+    end
+  end
+end