Heap2Local bug on cmpxchng

[kripken]

(module
 (rec
  (type $0 (sub (shared (array v128))))
  (type $1 (sub (shared (struct (field (mut (ref null $1)))))))
  (type $2 (shared (func (param (ref $5)) (result i32))))
  (type $3 (sub (func (param f64) (result contref))))
  (type $4 (sub (shared (struct (field i32) (field i8) (field (ref null $8)) (field (mut (ref (shared any)))) (field i32)))))
  (type $5 (func (param v128 f64) (result (ref null $1))))
  (type $6 (sub $0 (shared (array v128))))
  (type $7 (sub (shared (array v128))))
  (type $8 (sub $4 (shared (struct (field i32) (field i8) (field (ref null $8)) (field (mut (ref (shared any)))) (field i32)))))
  (type $9 (sub $4 (shared (struct (field i32) (field i8) (field (ref $8)) (field (mut (ref (shared any)))) (field i32) (field (mut (ref null (shared extern))))))))
  (type $10 (shared (func (param (ref null $5) f64 (ref null $3) i64 (ref $2) nullcontref (ref null $1)) (result (ref null $4) i32 (ref null $4) (ref $9) (ref null $11)))))
  (type $11 (sub $9 (shared (struct (field i32) (field i8) (field (ref $8)) (field (mut (ref (shared any)))) (field i32) (field (mut (ref null (shared extern))))))))
 )
 (export "func_482" (func $0))
 (@binaryen.js.called)
 (func $0 (type $5) (param $0 v128) (param $1 f64) (result (ref null $1))
  (local $2 (ref $1))
  (struct.atomic.rmw.xchg acqrel acqrel $1 0
   (local.tee $2
    (struct.new_default $1)
   )
   (struct.new $1
    (struct.atomic.rmw.cmpxchg acqrel acqrel $1 0
     (local.get $2)
     (struct.new_default $1)
     (struct.new_default $1)
    )
   )
  )
 )
)

$ bin/wasm-opt -all --fuzz-exec w.wat --heap2local
[fuzz-exec] export func_482
[fuzz-exec] note result: func_482 => null
[fuzz-exec] export func_482
[fuzz-exec] note result: func_482 => object(null)
[fuzz-exec] comparing func_482
values not identical! shared [ref (type $struct.0 (sub (shared (struct (field (mut (ref null $struct.0))))))) shared nullref] != shared nullref
[fuzz-exec] optimization passes changed results

This looks related to a worry we had back here:

#8491 (comment)

@tlively perhaps worth revisiting that, what do you think?

[tlively]

Simplified:

(module
 (type $struct (shared (struct (field (mut (ref null (shared eq)))))))
 (export "func" (func $func))
 (func $func (result i32)
  (local $local (ref null $struct))
  (drop
   (struct.atomic.rmw.cmpxchg acqrel acqrel $struct 0
    (local.tee $local
     (struct.new_default $struct)
    )
    (struct.new_default $struct)
    (struct.new_default $struct)
   )
  )
  (ref.is_null
   (struct.get $struct 0 (local.get $local))
  )
 )
)

Will investigate further, but the urge is strong to just rewrite Heap2Local to stop interleaving analysis and optimization :)

[kripken]

I hope that instead of a full rewrite, we just need to keep those maps fully updated all the time, as I suggested back in that other issue? Though there could be other issues I'm missing, or other benefits to a rewrite, ofc

[tlively]

The problem is that we don't optimize the ref.eq that's created when we optimize the cmpxchg. If we did, we would see that an allocation flows into one side but not the other, and we would correctly optimize it to a constant 0 instead of letting incorrectly compare two separately optimized values. From most surgical (but hacky) to most general, I can think of the following fixes:

1. Update analyzer.parents and analyzer.scratchInfo enough to let the new ref.eq get optimized correctly when we optimize the "expected" allocation. (Heap2Local: fix stale data cmpxchg bug #8856)
2. Simply clear and recompute analyzer.parents and analyzer.localGraph after each optimized allocation to ensure they are not stale. (Heap2Local: reset data structures after optimizing #8857)
3. Rewrite the pass to analyze all allocations before optimizing, allowing the optimizer to see all allocations flowing into the cmpxchg and letting it optimize the whole thing down to a single local.get. (WIP)

[ArkadySkv]

The problem is that we don't optimize the ref.eq that's created when we optimize the cmpxchg. If we did, we would see that an allocation flows into one side but not the other, and we would correctly optimize it to a constant 0 instead of letting incorrectly compare two separately optimized values. From most surgical (but hacky) to most general, I can think of the following fixes:

1. Update analyzer.parents and analyzer.scratchInfo enough to let the new ref.eq get optimized correctly when we optimize the "expected" allocation. (Heap2Local: fix stale data cmpxchg bug #8856)
2. Simply clear and recompute analyzer.parents and analyzer.localGraph after each optimized allocation to ensure they are not stale. (Heap2Local: reset data structures after optimizing #8857)
3. Rewrite the pass to analyze all allocations before optimizing, allowing the optimizer to see all allocations flowing into the cmpxchg and letting it optimize the whole thing down to a single local.get. (WIP)

Thanks for the detailed analysis and the three options you laid out. After looking through the discussion and the PRs (#8856, #8857), here's my take on a path forward that builds on what you've already started.

Your simplified test case helps isolate the core issue. The ref.eq not being optimized after the cmpxchg transformation is indeed the culprit, and your three proposed fixes cover the space well.

Option 1 (PR #8856): Updating analyzer.parents and analyzer.scratchInfo surgically is the most targeted approach. It keeps the existing optimization pipeline intact and only fixes the specific stale-data bug. The downside, as you noted, is that it's hacky — we'd be playing whack-a-mole with data structure staleness, and future changes to the pass could easily reintroduce similar bugs.

Option 2 (PR #8857): Resetting the data structures after each optimized allocation is a much more robust solution. It eliminates the entire class of stale-data bugs in one go. The overhead concern is valid, but as @kripken pointed out, optimizations are not super-common in practice, and the cost of rebuilding parents and localGraph per allocation is likely acceptable. The helper function you added for the reset logic keeps it clean.

Option 3 (Full rewrite): This is the most general and elegant solution — analyze all allocations first, then optimize in a single pass. It would completely avoid the incremental staleness problem and likely enable more aggressive optimizations (e.g., seeing all allocations flowing into a cmpxchg and folding the whole thing to a single local.get). However, it's a significant undertaking and would need careful design to avoid regressions.

My recommendation: Option 2 (PR #8857) as the immediate fix, with a path toward Option 3 as a future improvement.

Here's why:

1. PR Heap2Local: reset data structures after optimizing #8857 is already merged, and it solves the immediate issue (Heap2Local bug on cmpxchng #8850) correctly and cleanly. The reset approach is battle-tested in other Binaryen passes (e.g., the pop pattern @kripken mentioned) and gives us a solid foundation.
2. The overhead is acceptable for the reasons above. If we later find it's too slow on large modules, we can optimize it incrementally — e.g., only reset when new code is generated that could affect subsequent analyses, rather than after every allocation.
3. Option 3 is still worth pursuing, but as a separate, longer-term effort. The current pass interleaves analysis and optimization in a way that makes correctness hard to reason about. A rewrite that separates these phases would:
   · Eliminate the need for ScratchInfo and manual parent tracking entirely
   · Allow the optimizer to see the full data flow graph before making any transformations
   · Enable more aggressive folding (e.g., the ref.eq case you mentioned)
   · Make the pass easier to maintain and extend
4. In the meantime, we should add a test case for the nested RMW pattern from the original issue description (the xchg with a cmpxchg in its replacement operand) to ensure the reset approach handles it correctly. The simplified test you provided covers the ref.eq case, but the nested RMW case is a different data flow pattern that might still expose gaps.

[ArkadySkv]

The problem is that we don't optimize the ref.eq that's created when we optimize the cmpxchg. If we did, we would see that an allocation flows into one side but not the other, and we would correctly optimize it to a constant 0 instead of letting incorrectly compare two separately optimized values. From most surgical (but hacky) to most general, I can think of the following fixes:

1. Update analyzer.parents and analyzer.scratchInfo enough to let the new ref.eq get optimized correctly when we optimize the "expected" allocation. (Heap2Local: fix stale data cmpxchg bug #8856)
2. Simply clear and recompute analyzer.parents and analyzer.localGraph after each optimized allocation to ensure they are not stale. (Heap2Local: reset data structures after optimizing #8857)
3. Rewrite the pass to analyze all allocations before optimizing, allowing the optimizer to see all allocations flowing into the cmpxchg and letting it optimize the whole thing down to a single local.get. (WIP)

Thanks for the detailed analysis and the three options you laid out. After looking through the discussion and the PRs (#8856, #8857), here's my take on a path forward that builds on what you've already started.

Your simplified test case helps isolate the core issue. The ref.eq not being optimized after the cmpxchg transformation is indeed the culprit, and your three proposed fixes cover the space well.

Option 1 (PR #8856): Updating analyzer.parents and analyzer.scratchInfo surgically is the most targeted approach. It keeps the existing optimization pipeline intact and only fixes the specific stale-data bug. The downside, as you noted, is that it's hacky — we'd be playing whack-a-mole with data structure staleness, and future changes to the pass could easily reintroduce similar bugs.

Option 2 (PR #8857): Resetting the data structures after each optimized allocation is a much more robust solution. It eliminates the entire class of stale-data bugs in one go. The overhead concern is valid, but as @kripken pointed out, optimizations are not super-common in practice, and the cost of rebuilding parents and localGraph per allocation is likely acceptable. The helper function you added for the reset logic keeps it clean.

Option 3 (Full rewrite): This is the most general and elegant solution — analyze all allocations first, then optimize in a single pass. It would completely avoid the incremental staleness problem and likely enable more aggressive optimizations (e.g., seeing all allocations flowing into a cmpxchg and folding the whole thing to a single local.get). However, it's a significant undertaking and would need careful design to avoid regressions.

My recommendation: Option 2 (PR #8857) as the immediate fix, with a path toward Option 3 as a future improvement.

Here's why:

1. PR Heap2Local: reset data structures after optimizing #8857 is already merged, and it solves the immediate issue (Heap2Local bug on cmpxchng #8850) correctly and cleanly. The reset approach is battle-tested in other Binaryen passes (e.g., the pop pattern @kripken mentioned) and gives us a solid foundation.
2. The overhead is acceptable for the reasons above. If we later find it's too slow on large modules, we can optimize it incrementally — e.g., only reset when new code is generated that could affect subsequent analyses, rather than after every allocation.
3. Option 3 is still worth pursuing, but as a separate, longer-term effort. The current pass interleaves analysis and optimization in a way that makes correctness hard to reason about. A rewrite that separates these phases would:
   · Eliminate the need for ScratchInfo and manual parent tracking entirely
   · Allow the optimizer to see the full data flow graph before making any transformations
   · Enable more aggressive folding (e.g., the ref.eq case you mentioned)
   · Make the pass easier to maintain and extend
4. In the meantime, we should add a test case for the nested RMW pattern from the original issue description (the xchg with a cmpxchg in its replacement operand) to ensure the reset approach handles it correctly. The simplified test you provided covers the ref.eq case, but the nested RMW case is a different data flow pattern that might still expose gaps.

Proposed Implementation

1. src/passes/Heap2Local.cpp – Add RMW result tracking infrastructure

// In Heap2Local struct definition, after ScratchInfo scratchInfo;

// Tracks the scratch local used to hold the result of an RMW operation that
// returns a reference type, so that we can follow data flow through it.
struct RMWResultInfo {
  Expression* originalRMW;  // the RMW being replaced
  LocalSet* resultSet;      // local.set that writes the result
  LocalGet* resultGet;      // local.get that reads the result
};
using RMWResultInfoMap = std::unordered_map<Expression*, RMWResultInfo>;
RMWResultInfoMap rmwResultInfo;

Reason: This extends the ScratchInfo pattern to track RMW results specifically. As discussed in our RFC, when an RMW returns a reference type that is itself an allocation, the result should be tracked as a flow through subsequent operations.

---

2. src/passes/Heap2Local.cpp – Helper to detect RMW result reads

// Helper to check if a local.get reads from an RMW result scratch local.
bool isRMWResultGet(LocalGet* get) {
  return std::any_of(rmwResultInfo.begin(), rmwResultInfo.end(),
                     [&](auto& pair) { return pair.second.resultGet == get; });
}

Reason: This allows the pass to detect when an operand is reading from a previous RMW result, which is the key insight needed to handle the nested case from #8850.

---

3. src/passes/Heap2Local.cpp – Extend EscapeAnalyzer to follow RMW result flows

// In EscapeAnalyzer class, add reference to rmwResultInfo
RMWResultInfoMap& rmwResultInfo;

EscapeAnalyzer(const LazyLocalGraph& localGraph,
               ScratchInfo& scratchInfo,
               Parents& parents,
               RMWResultInfoMap& rmwResultInfo)
  : localGraph(localGraph), scratchInfo(scratchInfo),
    parents(parents), rmwResultInfo(rmwResultInfo) {}

// In checkFlows() method, after the scratchInfo lookup:
// If this set is writing an RMW result, then the flow continues
// through the corresponding get.
if (auto it = rmwResultInfo.find(parent); it != rmwResultInfo.end()) {
  flows.push({it->second.resultGet, parents.getParent(it->second.resultGet)});
  continue;
}

Reason: This allows the escape analysis to follow data flow through RMW results, which is essential for correctly identifying when allocations flow through nested operations.

---

4. src/passes/Heap2Local.cpp – Modify visitStructRMW to handle RMW result flows

void visitStructRMW(StructRMW* curr) {
  // ... existing early returns ...
  
  // Check if the replacement operand is a LocalGet that reads from a previous
  // RMW result. If so, we must not fully optimize away this RMW; we need to
  // keep the operation but replace its ref operand with locals.
  bool replacementIsRMWResult =
    curr->replacement->is<LocalGet>() &&
    isRMWResultGet(curr->replacement->cast<LocalGet>());
  
  // If the replacement is an RMW result, we cannot replace the whole RMW
  // with a local.get; we must keep the RMW operation and just localize
  // the ref operand.
  if (replacementIsRMWResult) {
    auto refType = curr->ref->type;
    auto scratch = builder.addVar(func, refType);
    auto* set = builder.makeLocalSet(scratch, curr->ref);
    auto* get = builder.makeLocalGet(scratch, refType);
    curr->ref = get;
    builder.replaceCurrent(builder.makeSequence(set, curr));
    return;
  }
  
  // If the RMW result is a reference type, store it in a scratch local
  // so that downstream uses can see the flow.
  if (curr->type.isRef() && !curr->type.isNull()) {
    auto resultType = curr->type;
    auto scratch = builder.addVar(func, resultType);
    auto* setScratch = builder.makeLocalSet(scratch, replacementExpr);
    auto* getScratch = builder.makeLocalGet(scratch, resultType);
    rmwResultInfo[curr] = {curr, setScratch, getScratch};
    builder.replaceCurrent(builder.makeBlock({setScratch, getScratch}));
    return;
  }
}

Reason: This handles the core case from #8850 where a cmpxchg result flows into the replacement operand of an outer xchg. The first check preserves the RMW when its replacement is an RMW result. The second check stores reference-typed RMW results in scratch locals for downstream tracking.

---

5. src/passes/Heap2Local.cpp – Modify visitStructCmpxchg similarly

void visitStructCmpxchg(StructCmpxchg* curr) {
  
  // Check if the expected or replacement operands are RMW results.
  bool expectedIsRMWResult =
    curr->expected->is<LocalGet>() &&
    isRMWResultGet(curr->expected->cast<LocalGet>());
  bool replacementIsRMWResult =
    curr->replacement->is<LocalGet>() &&
    isRMWResultGet(curr->replacement->cast<LocalGet>());

  if (expectedIsRMWResult || replacementIsRMWResult) {
    // If expected or replacement is an RMW result, we cannot optimize the
    // cmpxchg away; we just localize the ref operand.
    // (This case is rare but we handle it conservatively.)
  }
  
  // If the cmpxchg result is a reference type, store it in a scratch local
  // for downstream uses.
  if (curr->type.isRef() && !curr->type.isNull()) {
    auto resultType = curr->type;
    auto scratch = builder.addVar(func, resultType);
    auto* setScratch = builder.makeLocalSet(scratch, newCurr);
    auto* getScratch = builder.makeLocalGet(scratch, resultType);
    rmwResultInfo[curr] = {curr, setScratch, getScratch};
    builder.replaceCurrent(builder.makeBlock({setScratch, getScratch}));
    return;
  }
}

Reason: Symmetric handling for cmpxchg operations. The detection of RMW results in expected or replacement operands prevents incorrect optimization, while storing cmpxchg results in scratch locals enables proper flow tracking.

---

6. test/lit/passes/heap2local-rmw.wast – Add test case for nested RMW pattern

;; Test nested RMW where inner cmpxchg result flows into outer xchg replacement.
;; This ensures that we correctly track the flow of RMW results.

(module
  (type $s (struct (field (mut i32))))
  (memory 1)

  ;;(func $nested-rmw (type $0) (result (ref null $s))
  ;;   (local $0 (ref null $s))
  ;;   (local $1 (ref null $s))
  ;;   (local $2 (ref null $s))
  ;;   (block
  ;;    (local.set $0
  ;;     (struct.new_default $s)
  ;;    )
  ;;    (local.set $1
  ;;     (struct.atomic.rmw.cmpxchg $s 0
  ;;      (local.get $0)
  ;;      (struct.new_default $s)
  ;;      (struct.new_default $s)
  ;;     )
  ;;    )
  ;;    (local.set $2
  ;;     (struct.atomic.rmw.xchg $s 0
  ;;     (local.get $0)
  ;;     (local.get $1)
  ;;     )
  ;;    )
  ;;   (local.get $2)
  ;;  )
  ;; )
  (func $nested-rmw (result (ref null $s))
    (struct.atomic.rmw.xchg $s 0
      (struct.new_default $s)
      (struct.atomic.rmw.cmpxchg $s 0
        (struct.new_default $s)
        (struct.new_default $s)
        (struct.new_default $s)
      )
    )
  )
)

Reason: This test case reproduces the exact pattern from #8850 and verifies that the optimization correctly preserves the data flow through nested RMW operations. The CHECK lines confirm that all allocations are localized and the RMW operations are preserved with proper local references.

[tlively]

Hi @ArkadySkv, thanks for your interest in the issue. You're welcome to open a PR or discuss bugs and improvements on the issues, but large posts of raw LLM output aren't very helpful. Please only post new information or questions next time.

[ArkadySkv]

Hi @ArkadySkv, thanks for your interest in the issue. You're welcome to open a PR or discuss bugs and improvements on the issues, but large posts of raw LLM output aren't very helpful. Please only post new information or questions next time.

Sorry about the wording, typed it in a rush from my phone while out of town and used an LLM to clean up the formatting. Just wanted to point out that Heap2Local misses StructAtomicRMW. I'll try to provide a clean minimal testcase/patch once I am back at my main machine.