From 8ae41c55a1c4451723f319a2e0a5fe8dabef421f Mon Sep 17 00:00:00 2001 From: Simeon David Schaub Date: Sat, 29 Aug 2026 12:26:46 +0200 Subject: [PATCH 1/2] Apply in-place rewrites only to single-use sub-ops An inplace=true rewrite rule modifies the matched ops' operands instead of building new instructions. That mutation is only sound when the match is the single user of every op it mutates; any other user silently sees the rewritten value. The one such rule, the nested comparison strength reduction cmpi(addi(a, addi(b, 1)), y, <=) -> cmpi(addi(a, b), y, <), miscompiles kernels in which the 1-based index tile feeds both a mask comparison and a gather or scatter: the shared addi chain is rewritten to 0-based for the comparison while the gather's lowering still subtracts 1 from it, shifting every gathered element down by one. Whether the bug fires depends on worklist order: when the algebra cancellation subi(addi(a, addi(b, c)), c) -> addi(a, b) consumes the gather's subtraction first, the comparison is left as the chain's only user and the mutation happens to be sound. The driver now checks use counts before applying a rule in-place and falls back to the standard mode -- fresh ops, only the root replaced -- when a sub-op that would be mutated is shared. The rewrite fires in both cases; in-place application remains a worklist-efficiency detail. A FileCheck regression test asserts that the strength-reduced mask operand of a masked gather is never decremented again. Assisted-by: Claude Fable 5 --- src/compiler/transform/pipeline.jl | 5 ++++- src/compiler/transform/rewrite.jl | 29 +++++++++++++++++++++++++++-- test/codegen/integration.jl | 27 +++++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/src/compiler/transform/pipeline.jl b/src/compiler/transform/pipeline.jl index e6a9d907..2266fff3 100644 --- a/src/compiler/transform/pipeline.jl +++ b/src/compiler/transform/pipeline.jl @@ -207,7 +207,10 @@ const COMPARISON_RULES = RewriteRule[ # 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). + # than creating new ones (which would cascade the worklist). When the addi + # chain has users besides the cmpi — e.g. the same 1-based index tile also + # feeds a gather, whose lowering subtracts the 1 itself — the driver applies + # the rule in standard mode instead, leaving the chain intact for them. @rewrite(inplace=true, Intrinsics.cmpi(Intrinsics.addi(~a, Intrinsics.addi(~b, $(1))), ~y, $(ComparisonPredicate.LessThanOrEqual), $(Signedness.Signed)) => diff --git a/src/compiler/transform/rewrite.jl b/src/compiler/transform/rewrite.jl index 7ab25b82..790579b8 100644 --- a/src/compiler/transform/rewrite.jl +++ b/src/compiler/transform/rewrite.jl @@ -78,6 +78,9 @@ rather than creating new ops. The LHS and RHS trees are walked in parallel: wher 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. +When an op that would be mutated has users outside the match, the driver applies +the rule in standard (op-building) mode instead, so shared sub-expressions keep +their value for the other users. """ macro rewrite(args...) # Parse keyword arguments @@ -491,9 +494,31 @@ function find_matched_ssa(driver, pat::PCall, bindings) return nothing end +""" +An in-place rewrite mutates every matched inner op (an LHS `PCall` paired with +an RHS `RCall` below the root). That is only sound when each such op's single +use is the match itself — any other user keeps referring to the mutated op and +silently sees the new value. Returns `false` when an op that would be mutated +is shared (or cannot be located). +""" +function inplace_mutation_sound(driver::RewriteDriver, rhs::RCall, lhs::PCall, bindings) + for (sub_rhs, sub_lhs) in zip(rhs.operands, lhs.operands) + sub_rhs isa RCall && sub_lhs isa PCall || continue + ssa = find_matched_ssa(driver, sub_lhs, bindings) + ssa isa SSAValue && use_count(driver, ssa) == 1 || return false + inplace_mutation_sound(driver, sub_rhs, sub_lhs, bindings) || return false + end + return true +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 + # In-place mode: modify matched ops' operands without creating new + # instructions. Only sound when the match is the single user of every op it + # mutates — e.g. a 1-based index tile feeding both a mask comparison and a + # gather (whose lowering subtracts 1) must keep its value for the gather. + # Shared matches fall through to the standard path below, which builds + # fresh ops and replaces only the root. + if rule.inplace && inplace_mutation_sound(driver, rule.rhs::RCall, rule.lhs, match.bindings) return apply_inplace_rewrite!(driver, block, val, rule, match) end diff --git a/test/codegen/integration.jl b/test/codegen/integration.jl index ebc03c30..512a3035 100644 --- a/test/codegen/integration.jl +++ b/test/codegen/integration.jl @@ -565,6 +565,33 @@ end end end + @testset "masked gather: shared 1-based index survives comparison rewrite" begin + # `c` feeds both the mask comparison and the gather (whose lowering + # subtracts the 1 itself). The comparison strength reduction + # (x+1 ≤ y → x < y) must not mutate the shared addi chain in place: + # the gather would then subtract 1 from an already 0-based index and + # read every lane one element too low. The mask's 0-based operand + # must never itself be decremented again. + @test @filecheck begin + @check_label "entry" + @check_not "less_than_or_equal" + @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 "contiguous-axis stride folds out of 2D gather offset" begin spec_out = ct.ArraySpec{1}(16, true) From de6753e7eb00e0de54df45f54e6e151e5330872e Mon Sep 17 00:00:00 2001 From: Simeon David Schaub Date: Sat, 29 Aug 2026 12:42:29 +0200 Subject: [PATCH 2/2] clean up comments --- src/compiler/transform/pipeline.jl | 6 ++---- src/compiler/transform/rewrite.jl | 18 ++++++------------ test/codegen/integration.jl | 9 +++------ 3 files changed, 11 insertions(+), 22 deletions(-) diff --git a/src/compiler/transform/pipeline.jl b/src/compiler/transform/pipeline.jl index 2266fff3..773760b4 100644 --- a/src/compiler/transform/pipeline.jl +++ b/src/compiler/transform/pipeline.jl @@ -207,10 +207,8 @@ const COMPARISON_RULES = RewriteRule[ # 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). When the addi - # chain has users besides the cmpi — e.g. the same 1-based index tile also - # feeds a gather, whose lowering subtracts the 1 itself — the driver applies - # the rule in standard mode instead, leaving the chain intact for them. + # than creating new ones (which would cascade the worklist), but only if the same + # tile is not used elsewhere. @rewrite(inplace=true, Intrinsics.cmpi(Intrinsics.addi(~a, Intrinsics.addi(~b, $(1))), ~y, $(ComparisonPredicate.LessThanOrEqual), $(Signedness.Signed)) => diff --git a/src/compiler/transform/rewrite.jl b/src/compiler/transform/rewrite.jl index 790579b8..0ee03137 100644 --- a/src/compiler/transform/rewrite.jl +++ b/src/compiler/transform/rewrite.jl @@ -495,11 +495,8 @@ function find_matched_ssa(driver, pat::PCall, bindings) end """ -An in-place rewrite mutates every matched inner op (an LHS `PCall` paired with -an RHS `RCall` below the root). That is only sound when each such op's single -use is the match itself — any other user keeps referring to the mutated op and -silently sees the new value. Returns `false` when an op that would be mutated -is shared (or cannot be located). +Return true if an in-place rewrite is sound for the given match, return false +if the op to be mutated is shared and referenced elsewhere. """ function inplace_mutation_sound(driver::RewriteDriver, rhs::RCall, lhs::PCall, bindings) for (sub_rhs, sub_lhs) in zip(rhs.operands, lhs.operands) @@ -512,13 +509,10 @@ function inplace_mutation_sound(driver::RewriteDriver, rhs::RCall, lhs::PCall, b end function apply_rewrite!(driver::RewriteDriver, block, val::SSAValue, rule, match) - # In-place mode: modify matched ops' operands without creating new - # instructions. Only sound when the match is the single user of every op it - # mutates — e.g. a 1-based index tile feeding both a mask comparison and a - # gather (whose lowering subtracts 1) must keep its value for the gather. - # Shared matches fall through to the standard path below, which builds - # fresh ops and replaces only the root. - if rule.inplace && inplace_mutation_sound(driver, rule.rhs::RCall, rule.lhs, match.bindings) + # In-place mode: modify matched ops' operands without creating new instructions. + # Needs to check `inplace_mutation_sound` to ensure the matched ops are not used + # anywhere else. + if rule.inplace && inplace_mutation_sound(driver, rule.rhs, rule.lhs, match.bindings) return apply_inplace_rewrite!(driver, block, val, rule, match) end diff --git a/test/codegen/integration.jl b/test/codegen/integration.jl index 512a3035..0ad18787 100644 --- a/test/codegen/integration.jl +++ b/test/codegen/integration.jl @@ -566,12 +566,9 @@ end end @testset "masked gather: shared 1-based index survives comparison rewrite" begin - # `c` feeds both the mask comparison and the gather (whose lowering - # subtracts the 1 itself). The comparison strength reduction - # (x+1 ≤ y → x < y) must not mutate the shared addi chain in place: - # the gather would then subtract 1 from an already 0-based index and - # read every lane one element too low. The mask's 0-based operand - # must never itself be decremented again. + # `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 in place, + # otherwise the gather sees the wrong index @test @filecheck begin @check_label "entry" @check_not "less_than_or_equal"