Skip to content

Remove the in-place rewrite mode #306

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
maleadt merged 1 commit into from
Aug 31, 2026
Merged

Remove the in-place rewrite mode #306

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions src/compiler/transform/pipeline.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,10 +206,7 @@ const COMPARISON_RULES = RewriteRule[
Intrinsics.cmpi(~x, ~y, $(ComparisonPredicate.LessThan), $(Signedness.Signed))

# Nested: cmpi(addi(a, addi(b, 1)), y, <=, signed) → cmpi(addi(a, b), y, <, signed)
# Uses inplace=true to modify the existing addi and cmpi ops' operands rather
# than creating new ones (which would cascade the worklist).
@rewrite(inplace=true,
Intrinsics.cmpi(Intrinsics.addi(~a, Intrinsics.addi(~b, $(1))), ~y,
@rewrite(Intrinsics.cmpi(Intrinsics.addi(~a, Intrinsics.addi(~b, $(1))), ~y,
$(ComparisonPredicate.LessThanOrEqual), $(Signedness.Signed)) =>
Intrinsics.cmpi(Intrinsics.addi(~a, ~b), ~y,
$(ComparisonPredicate.LessThan), $(Signedness.Signed)))
Expand Down
149 changes: 20 additions & 129 deletions src/compiler/transform/rewrite.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,8 @@ struct RewriteRule
lhs::PCall
rhs::RewriteNode
guard::Union{Function, Nothing} # (match, driver) -> Bool, or nothing
inplace::Bool # true = modify matched ops in-place (no new instructions)
end
RewriteRule(lhs::PCall, rhs::RewriteNode) = RewriteRule(lhs, rhs, nothing, false)
RewriteRule(lhs::PCall, rhs::RewriteNode, guard) = RewriteRule(lhs, rhs, guard, false)
RewriteRule(lhs::PCall, rhs::RewriteNode) = RewriteRule(lhs, rhs, nothing)

root_func(rule::RewriteRule) = rule.lhs.func

Expand All @@ -62,8 +60,6 @@ root_func(rule::RewriteRule) = rule.lhs.func
"""
@rewrite lhs => rhs
@rewrite(lhs => rhs, guard)
@rewrite(inplace=true, lhs => rhs)
@rewrite(inplace=true, lhs => rhs, guard)

Compile a declarative rewrite rule. LHS: `func(args...)` matches calls,
`~x` binds (repeated names require equality), `~x::T` binds with type constraint,
Expand All @@ -72,32 +68,12 @@ RHS: `func(args...)` emits calls, `~x` references bindings,
`\$(expr)` injects a literal constant.

Optional `guard` is a function `(match, driver) -> Bool` checked after pattern match.

With `inplace=true`, the RHS describes modifications to the matched ops' operands
rather than creating new ops. The LHS and RHS trees are walked in parallel: where
they share the same function, the existing op is modified in-place; where the RHS
has a different binding or constant, the operand is replaced. This avoids the
worklist cascade that occurs when the standard mode creates new instructions.
"""
macro rewrite(args...)
# Parse keyword arguments
inplace = false
positional = Any[]
for arg in args
if arg isa Expr && arg.head === :(=) && arg.args[1] === :inplace
inplace = arg.args[2]::Bool
else
push!(positional, arg)
end
end
length(positional) >= 1 || error("@rewrite expects: lhs => rhs")
ex = positional[1]
guard = length(positional) >= 2 ? positional[2] : nothing

macro rewrite(ex, guard=nothing)
ex isa Expr && ex.head === :call && ex.args[1] === :(=>) ||
error("@rewrite expects: lhs => rhs")
g = guard === nothing ? :nothing : guard
esc(:(RewriteRule($(compile_lhs(ex.args[2])), $(compile_rhs(ex.args[3])), $g, $inplace)))
esc(:(RewriteRule($(compile_lhs(ex.args[2])), $(compile_rhs(ex.args[3])), $g)))
end

"""
Expand Down Expand Up @@ -255,8 +231,8 @@ function notify_modified!(d::RewriteDriver, block::Block, val::SSAValue,
d.defs[val] = DefEntry(block, val, func)
end
# Worklist cascading is done by the callers (`apply_rewrite!`,
# `apply_inplace_rewrite!`, `commute_arith_transparent`), which re-seed
# `val` and its users themselves.
# `commute_arith_transparent`), which re-seed `val` and its users
# themselves.
end

function notify_erased!(d::RewriteDriver, ::Block, val::SSAValue,
Expand Down Expand Up @@ -378,8 +354,8 @@ end

"""Resolve an RHS operand, inserting sub-calls before `ref` as needed.
`root_typ` is the type of the original matched instruction — used only for the
outermost RCall (which replaces the root in-place). Intermediate RCalls infer
their type from the first SSA operand, since element-wise ops preserve type."""
outermost RCall (whose statement replaces the root). Intermediate RCalls infer
their type from the first value-like operand, since element-wise ops preserve type."""
resolve_rhs(driver, block, ref, op::RBind, bindings, root_typ) = bindings[op.name]
resolve_rhs(driver, block, ref, op::RConst, bindings, root_typ) = op.val
function resolve_rhs(driver::RewriteDriver, block, ref, op::RCall, bindings, root_typ)
Expand All @@ -392,111 +368,26 @@ function resolve_rhs(driver::RewriteDriver, block, ref, op::RCall, bindings, roo
push!(operands, resolve_rhs(driver, block, ref, sub, bindings, root_typ))
end
end
# Infer type from first SSA operand — correct for element-wise ops (addi,
# subi, negf, etc.) whose result type matches their operands. Falls back to
# root_typ when no SSA operand is available.
# Infer type from the first value-like operand (SSA, argument, or block
# argument) — correct for element-wise ops (addi, subi, negf, etc.) whose
# result type matches their operands. Kernel arguments still have their
# source-level scalar type, so promote them to the canonical 0-D tile type.
# Literal operands are skipped: their scalar type is not necessarily the
# tile result type. Falls back to root_typ when no value operand is available.
typ = root_typ
for o in operands
o isa SSAValue || continue
is_trackable_value(o) || continue
t = value_type(block, o)
t === nothing && continue
typ = CC.widenconst(t)
typ = boundary_jltype(CC.widenconst(t))
break
end
inst = insert_before!(driver.rewriter, block, ref, Expr(:call, op.func, operands...), typ;
flag=inferred_flags(op.func))
SSAValue(inst)
end

"""
In-place rewrite: walk LHS/RHS trees in parallel, modifying matched ops' operands.
No new instructions are created — existing ops are modified in-place.
"""
function apply_inplace_rewrite!(driver::RewriteDriver, block, val::SSAValue, rule, match)
haskey(block, val.id) || return false

# Build new operands for the root op from the RHS.
new_operands = Any[resolve_inplace_rhs(driver, match.bindings, op, lhs_op)
for (op, lhs_op) in zip(rule.rhs.operands, rule.lhs.operands)]
new_stmt = Expr(:call, rule.rhs.func, new_operands...)
# Same-func rewrites (most common: only operands change) preserve flag
# via the partial-NamedTuple setindex. Different-func rewrites recompute
# the flag from the new func's declared effects (`efunc` overrides),
# mirroring inference's `flags_for_effects` — the inferred bits on the
# old call describe the OLD op and don't carry over.
flag = rule.rhs.func === driver.defs[val].func ? nothing : inferred_flags(rule.rhs.func)
replace_stmt!(driver.rewriter, block, val, new_stmt; flag)
push!(driver.worklist, val)
add_users_to_worklist!(driver, val)
return true
end

"""Resolve an RHS operand for in-place mode. If the RHS sub-tree is an RCall
matching a PCall at the same position, modify the existing op in-place and
return its SSAValue. Otherwise fall back to bindings/constants."""
resolve_inplace_rhs(driver, bindings, op::RBind, @nospecialize(lhs_op)) = bindings[op.name]
resolve_inplace_rhs(driver, bindings, op::RConst, @nospecialize(lhs_op)) = op.val

function resolve_inplace_rhs(driver, bindings, op::RCall, lhs_op::PCall)
# The LHS matched an op with this function. Find it via the match bindings
# or the defs index and modify it in-place.
op.func === lhs_op.func && length(op.operands) == length(lhs_op.operands) ||
error("inplace rewrite: RHS sub-call $(op.func) doesn't match LHS structure")
matched_ssa = @something find_matched_ssa(driver, lhs_op, bindings) error(
"inplace rewrite: could not find matched SSA for $(lhs_op.func)")
entry = @something get(driver.defs, matched_ssa, nothing) error(
"inplace rewrite: no def entry for $matched_ssa")
haskey(entry.block, matched_ssa.id) ||
error("inplace rewrite: $matched_ssa not found in block")
new_ops = Any[resolve_inplace_rhs(driver, bindings, sub_rhs, sub_lhs)
for (sub_rhs, sub_lhs) in zip(op.operands, lhs_op.operands)]
# Sub-call func is enforced equal to the matched lhs func above
# (`op.func === lhs_op.func`), so the flag is still valid — only operands
# changed.
replace_stmt!(driver.rewriter, entry.block, matched_ssa, Expr(:call, op.func, new_ops...))
push!(driver.worklist, matched_ssa)
return matched_ssa
end

# Fallback for mismatched LHS/RHS structure
function resolve_inplace_rhs(driver, bindings, op::RCall, @nospecialize(lhs_op))
error("inplace rewrite: RHS has RCall but LHS has $(typeof(lhs_op)) at same position")
end

"""Find the SSA value that was matched by a PCall pattern node during matching.
The matched_ssas in MatchResult are ordered root-first, but we need to find
the specific SSA for a sub-pattern. We do this by looking up the first operand's
binding and finding the op that defines it."""
function find_matched_ssa(driver, pat::PCall, bindings)
for sub in pat.operands
if sub isa PBind
bound = get(bindings, sub.name, nothing)
bound isa SSAValue || continue
for u in users_of(driver, bound)
def_entry = get(driver.defs, u, nothing)
def_entry === nothing && continue
def_entry.func === pat.func && return u
end
elseif sub isa PCall
inner_ssa = find_matched_ssa(driver, sub, bindings)
if inner_ssa !== nothing
for u in users_of(driver, inner_ssa)
def_entry = get(driver.defs, u, nothing)
def_entry === nothing && continue
def_entry.func === pat.func && return u
end
end
end
end
return nothing
end

function apply_rewrite!(driver::RewriteDriver, block, val::SSAValue, rule, match)
# In-place mode: modify matched ops' operands without creating new instructions
if rule.inplace
return apply_inplace_rewrite!(driver, block, val, rule, match)
end

entry = driver.defs[val]
if rule.rhs isa RFunc
# Look up live instruction for RFunc interface
Expand All @@ -518,7 +409,7 @@ function apply_rewrite!(driver::RewriteDriver, block, val::SSAValue, rule, match
replace_uses!(driver.rewriter, val, new_val)
erase!(driver.rewriter, entry.block, val)
else
# Substitution: replace root in-place, clean up dead intermediates.
# Substitution: rewrite the root statement, clean up dead intermediates.
# Only delete intermediates with no remaining uses — transparent-op
# tracing may have added multi-use intermediates to matched_ssas.
for dead_val in match.matched_ssas
Expand All @@ -538,10 +429,10 @@ function apply_rewrite!(driver::RewriteDriver, block, val::SSAValue, rule, match
push!(operands, resolve_rhs(driver, block, val, op, match.bindings, typ))
end
end
# The substituted func almost always differs from the matched root
# (otherwise this would be an inplace rule). Inferred IR_FLAG_* bits
# describe the OLD op; recompute from the new func's `efunc` effects
# so downstream gates (CSE, LICM) see fresh, correct information.
# When the substituted func differs from the matched root, the
# inferred IR_FLAG_* bits describe the OLD op; recompute from the new
# func's `efunc` effects so downstream gates (CSE, LICM) see fresh,
# correct information.
new_stmt = Expr(:call, rule.rhs.func, operands...)
flag = rule.rhs.func === driver.defs[val].func ? nothing : inferred_flags(rule.rhs.func)
replace_stmt!(driver.rewriter, block, val, new_stmt; flag)
Expand Down
81 changes: 81 additions & 0 deletions test/codegen/integration.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@

import IRStructurizer

#=============================================================================
Type Support
=============================================================================#
Expand Down Expand Up @@ -61,6 +63,33 @@
end
end

#=============================================================================
Rewrite Pipeline
=============================================================================#

@testset "comparison rewrite preserves the type of kernel arguments" begin
sci, _ = only(IRStructurizer.code_structured(Tuple{Int,Int,Int}) do a::Int, b::Int, y::Int
a + (b + 1) <= y
end)

ct.run_passes!(sci)

addis = [inst for inst in IRStructurizer.instructions(sci.entry)
if let call = IRStructurizer.resolve_call(sci.entry, inst)
call !== nothing && first(call) === ct.Intrinsics.addi
end]
@test length(addis) == 1
@test IRStructurizer.value_type(only(addis)) === ct.Tile{Int,Tuple{}}

cmpis = [inst for inst in IRStructurizer.instructions(sci.entry)
if let call = IRStructurizer.resolve_call(sci.entry, inst)
call !== nothing && first(call) === ct.Intrinsics.cmpi
end]
@test length(cmpis) == 1
_, operands = IRStructurizer.resolve_call(sci.entry, only(cmpis))
@test operands[3] === ct.ComparisonPredicate.LessThan
end

#=============================================================================
Integration Tests
=============================================================================#
Expand Down Expand Up @@ -565,6 +594,58 @@ end
end
end

@testset "masked gather: shared 1-based index survives comparison rewrite" begin
# `c` is used in both the mask comparison and the gather. The rewrite
# rule `x + 1 ≤ y => x < y` must not mutate the shared addi,
# otherwise the gather sees the wrong index
@test @filecheck implicit_check_not="less_than_or_equal" begin
@check_label "entry"
@check "cmpi less_than [[C:%[^,]+]], {{[^,]+}}, signed"
@check_not "subi [[C]],"
@check "load_ptr_tko"
code_tiled(Tuple{ct.TileArray{Float32,1,Int32,spec},
ct.TileArray{Float32,1,Int32,spec},
Int32}) do a, b, lim
pid = ct.bid(1)
off = (pid - Int32(1)) * Int32(16)
c = off .+ ct.arange(16)
mask = c .≤ lim
tile = ct.gather(a, c; mask)
ct.store(b, pid, tile)
return
end
end
end

@testset "masked gather: decoy addi sharing an operand is not clobbered" begin
# `d = s + s` shares operand `s` with the addi the comparison rule
# matches (`c = s + arange`), and precedes it in the use index. A
# rewrite of the mask comparison must leave `d` intact: the gather
# must still index off `addi(s, s)`, not off the rewritten mask
# operand `addi(s, iota)`.
@test @filecheck implicit_check_not="less_than_or_equal" begin
@check_label "entry"
@check "[[S:%[^ ]+]] = addi [[U:%[^,]+]], [[U]] :"
@check "[[D:%[^ ]+]] = addi [[S]], [[S]] :"
@check_dag "cmpi less_than"
@check_dag "subi [[D]],"
@check "load_ptr_tko"
code_tiled(Tuple{ct.TileArray{Float32,1,Int32,spec},
ct.TileArray{Float32,1,Int32,spec},
Int32}) do a, b, lim
pid = ct.bid(1)
u = ct.arange(16)
s = u .+ u
d = s + s
c = s + ct.arange(16)
mask = c .≤ lim
tile = ct.gather(a, d; mask)
ct.store(b, pid, tile)
return
end
end
end

@testset "contiguous-axis stride folds out of 2D gather offset" begin
spec_out = ct.ArraySpec{1}(16, true)

Expand Down