Apply in-place rewrites only if there are no other uses

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)) =>

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
 

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)