Constant propagation and identity fold for the rewrite framework (#160)

Replace the ad-hoc _is_transparent tracing in PCall pattern matching with
two proper mechanisms:

1. Identity fold rewrite rules: eliminate identity broadcasts and reshapes
   (same shape in/out) that are no-ops left behind by the broadcast system.
   FMA patterns now match directly without transparent-op tracing.

2. Constant propagation analysis: propagate_constants() builds a
   Dict{SSAValue, Array} mapping SSA values to their known constant tile
   contents (e.g., broadcast(1, (16,)) → fill(Int32(1), 16)). This feeds
   into the rewriter via a constants kwarg.

3. PLiteral pattern node: $() on the LHS matches literal values via O(1)
   constants dict lookup instead of recursive _traces_to_const tracing.

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: maleadt

src/compiler/passes/pipeline.jl

@@ -50,15 +50,117 @@ const ALGEBRA_RULES = RewriteRule[
 
 algebra_pass!(sci::StructuredIRCode) = rewrite_patterns!(sci, ALGEBRA_RULES)
 
+#=============================================================================
+ Identity Fold (rewrite)
+=============================================================================#
+
+# Eliminate identity broadcasts and reshapes (same shape in/out). These are
+# no-ops left behind by the broadcast system after scalar elimination.
+
+function is_identity_op(match, driver)
+    x = match.bindings[:x]
+    val = first(match.matched_ssas)
+    entry = driver.defs[val]
+    in_t = value_type(entry.block, x)
+    out_t = value_type(entry.block, val)
+    in_t === nothing && return false
+    out_t === nothing && return false
+    in_T = CC.widenconst(in_t)
+    out_T = CC.widenconst(out_t)
+    in_T <: Tile && out_T <: Tile || return false
+    return size(in_T) == size(out_T)
+end
+
+const IDENTITY_RULES = RewriteRule[
+    @rewrite(Intrinsics.broadcast(~x, ~shape) => ~x, is_identity_op)
+    @rewrite(Intrinsics.reshape(~x, ~shape) => ~x, is_identity_op)
+]
+
 #=============================================================================
  Combined Rule Set
 =============================================================================#
 
 const OPTIMIZATION_RULES = RewriteRule[
+    IDENTITY_RULES...,
     ALGEBRA_RULES...,
     FMA_RULES...,
 ]
 
+#=============================================================================
+ Constant Propagation (analysis)
+=============================================================================#
+
+# Tracks which SSA values have known constant values. Constants are represented
+# as Julia Arrays matching the tile's element type and shape. This enables O(1)
+# constant lookups in the rewrite pattern matcher (PLiteral).
+
+"""
+    propagate_constants(sci) -> Dict{SSAValue, Any}
+
+Build a map from SSA values to their known constant values. Walks all blocks
+in program order so transitive constants (e.g. reshape of a broadcast of a
+literal) resolve correctly.
+"""
+function propagate_constants(sci::StructuredIRCode)
+    constants = Dict{SSAValue, Any}()
+    propagate_constants!(constants, sci.entry)
+    return constants
+end
+
+function propagate_constants!(constants::Dict{SSAValue, Any}, block::Block)
+    # Recurse into nested control flow first
+    for inst in instructions(block)
+        s = stmt(inst)
+        if s isa ForOp
+            propagate_constants!(constants, s.body)
+        elseif s isa IfOp
+            propagate_constants!(constants, s.then_region)
+            propagate_constants!(constants, s.else_region)
+        elseif s isa WhileOp
+            propagate_constants!(constants, s.before)
+            propagate_constants!(constants, s.after)
+        elseif s isa LoopOp
+            propagate_constants!(constants, s.body)
+        end
+    end
+
+    for inst in instructions(block)
+        call = resolve_call(block, inst)
+        call === nothing && continue
+        func, ops = call
+
+        # Transparent ops (broadcast, reshape) propagate constants from operand
+        if (func === Intrinsics.broadcast || func === Intrinsics.reshape) &&
+                length(ops) >= 1
+            scalar = const_value(constants, ops[1])
+            scalar === nothing && continue
+            vt = value_type(block, SSAValue(inst))
+            vt === nothing && continue
+            T = CC.widenconst(vt)
+            T <: Tile || continue
+            S = size(T)
+            constants[SSAValue(inst)] = fill(convert(eltype(T), scalar), S)
+        end
+    end
+end
+
+"""Resolve an operand to its scalar constant value, or `nothing`."""
+function const_value(constants::Dict{SSAValue, Any}, @nospecialize(op))
+    if op isa Number
+        return op
+    elseif op isa QuoteNode && op.value isa Number
+        return op.value
+    elseif op isa SSAValue
+        c = get(constants, op, nothing)
+        c isa AbstractArray || return nothing
+        isempty(c) && return nothing
+        v = first(c)
+        all(==(v), c) && return v
+        return nothing
+    end
+    return nothing
+end
+
 #=============================================================================
  Pass Pipeline
 =============================================================================#
@@ -72,7 +174,8 @@ and subprogram compilation.
 function run_passes!(sci::StructuredIRCode)
     canonicalize!(sci)
 
-    rewrite_patterns!(sci, OPTIMIZATION_RULES)
+    constants = propagate_constants(sci)
+    rewrite_patterns!(sci, OPTIMIZATION_RULES; constants)
 
     alias_result = alias_analysis_pass!(sci)
     token_order_pass!(sci, alias_result)

src/compiler/passes/rewrite.jl

@@ -26,6 +26,7 @@ struct PCall <: PatternNode; func::Any; operands::Vector{PatternNode}; end
 struct PBind <: PatternNode; name::Symbol; end
 struct PTypedBind <: PatternNode; name::Symbol; type::Type; end
 struct POneUse <: PatternNode; inner::PatternNode; end
+struct PLiteral <: PatternNode; val::Any; end
 
 abstract type RewriteNode end
 struct RCall <: RewriteNode; func::Any; operands::Vector{RewriteNode}; end
@@ -86,6 +87,10 @@ macro rewriter(ex)
 end
 
 function _compile_lhs(ex)
+    # $(expr) on the LHS: match a literal value
+    if ex isa Expr && ex.head === :$
+        return :(PLiteral($(ex.args[1])))
+    end
     ex isa Expr && ex.head === :call || error("@rewrite LHS: expected call, got $ex")
     f = ex.args[1]
     if f === :~
@@ -172,17 +177,14 @@ mutable struct RewriteDriver
     defs::Dict{SSAValue, DefEntry}
     dispatch::Dict{Any, Vector{RewriteRule}}
     worklist::Worklist
+    constants::Dict{SSAValue, Any}   # SSA → constant value (from propagate_constants)
     max_rewrites::Int
 end
 
 """Compute fresh use count for an SSA value."""
 _use_count(driver::RewriteDriver, val::SSAValue) =
     length(uses(driver.sci.entry, val))
 
-# Codegen no-ops that pattern matching traces through transparently.
-_is_transparent(func) = func === Intrinsics.broadcast ||
-                         func === Intrinsics.reshape
-
 #=============================================================================
  Notifications
 =============================================================================#
@@ -266,29 +268,6 @@ function pattern_match(driver::RewriteDriver, @nospecialize(val), pat::PCall,
         end
     end
 
-    # Trace through single-use transparent ops to find the underlying operation
-    if _is_transparent(entry.func)
-        _use_count(driver, val) == 1 || return nothing
-        ops = _def_operands(entry)
-        isempty(ops) && return nothing
-        if entry.func === Intrinsics.broadcast
-            inner = ops[1]
-            if inner isa SSAValue
-                inner_entry = get(driver.defs, inner, nothing)
-                if inner_entry !== nothing
-                    it = value_type(entry.block, inner)
-                    ot = value_type(entry.block, val)
-                    it !== nothing && ot !== nothing || return nothing
-                    CC.widenconst(it) <: Tile && CC.widenconst(ot) <: Tile || return nothing
-                    size(CC.widenconst(it)) == size(CC.widenconst(ot)) || return nothing
-                end
-            end
-        end
-        result = pattern_match(driver, ops[1], pat, entry.block)
-        result === nothing && return nothing
-        push!(result.matched_ssas, val)
-        return result
-    end
     return nothing
 end
 
@@ -309,6 +288,23 @@ function pattern_match(driver::RewriteDriver, @nospecialize(val), pat::POneUse,
     pattern_match(driver, val, pat.inner, block)
 end
 
+# PLiteral: match if the operand equals the given value.
+# For non-SSA operands (enum constants, predicates): checks ===.
+# For SSA operands: O(1) lookup in the constants map built by propagate_constants.
+function pattern_match(driver::RewriteDriver, @nospecialize(val), pat::PLiteral,
+                       block::Block=driver.sci.entry)
+    val === pat.val && return MatchResult(Dict{Symbol,Any}(), SSAValue[])
+    if val isa SSAValue
+        c = get(driver.constants, val, nothing)
+        if c isa AbstractArray
+            all(==(pat.val), c) && return MatchResult(Dict{Symbol,Any}(), SSAValue[])
+        elseif c !== nothing
+            c == pat.val && return MatchResult(Dict{Symbol,Any}(), SSAValue[])
+        end
+    end
+    return nothing
+end
+
 #=============================================================================
  Rewrite Application
 =============================================================================#
@@ -373,7 +369,8 @@ Rules are tried until no more matches fire or `max_rewrites` is reached.
 Dead code left behind is cleaned up by the pipeline's `dce_pass!`.
 """
 function rewrite_patterns!(sci::StructuredIRCode, rules::Vector{RewriteRule};
-                           max_rewrites::Int=10_000)
+                           max_rewrites::Int=10_000,
+                           constants::Dict{SSAValue, Any}=Dict{SSAValue, Any}())
     # Build dispatch table
     dispatch = Dict{Any, Vector{RewriteRule}}()
     for rule in rules
@@ -401,7 +398,7 @@ function rewrite_patterns!(sci::StructuredIRCode, rules::Vector{RewriteRule};
         end
     end
 
-    driver = RewriteDriver(sci, defs, dispatch, wl, max_rewrites)
+    driver = RewriteDriver(sci, defs, dispatch, wl, constants, max_rewrites)
 
     num_rewrites = 0
     while !isempty(driver.worklist) && num_rewrites < driver.max_rewrites