diff --git a/src/compiler/transform/pipeline.jl b/src/compiler/transform/pipeline.jl index e6a9d907..773760b4 100644 --- a/src/compiler/transform/pipeline.jl +++ b/src/compiler/transform/pipeline.jl @@ -207,7 +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). + # 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 7ab25b82..0ee03137 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,25 @@ function find_matched_ssa(driver, pat::PCall, bindings) return nothing end +""" +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) + 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. + # 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 ebc03c30..0ad18787 100644 --- a/test/codegen/integration.jl +++ b/test/codegen/integration.jl @@ -565,6 +565,30 @@ 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 in place, + # otherwise the gather sees the wrong index + @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)