[alexreinking]

CodeGen_LLVM::codegen_atomic_rmw lowers an atomic() update that isn't a
single atomic read-modify-write to a compare-and-swap loop. The loop compared
against a phi of the atomic load and cmpxchg result, but computed the new value
with codegen(op->value). In that value, the self-load was a plain reload of
the location. The cmpxchg could therefore succeed with a value computed from a
different (stale) value of the location. That is a data race in the LLVM IR,
and on CUDA it loses updates: the weak ld.global can read a stale L1 line.

Fix

Inside the existing loop, each Load of the stored buffer in the value is
replaced with the phi (reinterpreted or cast when the buffer is stored as a
different type, see upgrade_type_for_storage) before codegen:

• If simplify(load_index == store_index) is true, the load is replaced with
  the phi.
• If it is false, the load is of another location and is kept as it is.
• Otherwise it becomes select(load_index == store_index, phi, load). This
  case matters in practice. After lowering, the self-load index and the store
  index of a += can be different expressions with the same value, e.g.
  out[t1832] = out[t2214 + 2] + a where both are lets of separately CSE'd
  expressions. That is also why such a += misses the atomic-add path. LLVM
  usually proves the comparison true and folds the select and the load away. In
  the case above, the PTX loop has no ld.global left.

Vector stores are still scalarized one lane at a time. For each lane the value
and the store index are both reduced with extract_lane, so a vector self-load
at the store's index becomes a scalar load at that lane's index, the same
expression as that lane's store index, and is replaced exactly. A vector load
of the buffer that remains, e.g. one at a broadcast index, is compared
lane-wise.

AddAtomicMutex doesn't rule out a self-load at a different index for stores
it leaves to the CAS path. With atomic(true) such a load is the user's
responsibility, so it is left alone rather than asserted against.

The atomic-add path is unchanged.

Tests

• New correctness_atomics_cas_loop: compiles CAS-path updates (int32 min,
  float max, uint32 *, a vectorized int32 min, and float16 min, which
  goes through the storage-type reinterpret) to LLVM assembly. It checks that no
  basic block containing a cmpxchg also contains a non-atomic load, and that
  at least one such block exists. It fails on main for all five cases and passes
  with this change.
• correctness_atomics: test_parallel_cas_update now uses few enough updates
  that no bin reaches the min cap, so lost updates change the result. Before,
  6 of 7 bins saturated at 100. A new GPU test, test_gpu_cas_tiles, runs on
  CUDA/OpenCL: 200 blocks each min-update 16 tiles in a staggered order with
  exact integer deltas. It is a best-effort stress test. On an RTX 3090 it fails
  on main in all 3 of 3 runs and passes with this change.

Results on an RTX 3090 (HL_JIT_TARGET=host-cuda): correctness_atomics,
correctness_atomic_tuples, correctness_atomics_gpu_8_bit,
correctness_atomics_gpu_mutex and correctness_atomics_cas_loop all pass. On
the original GEMM-shaped reproducer, main gives wrong results in 16 of 18 runs
and this branch in 0 of 27.

Performance: in the CAS loop, the reload is replaced with a value already held
in a register. When a select is needed and LLVM can't fold it, there is one
extra compare and select per iteration. The atomic-add path is unchanged.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PWSVRwgKGuCHpMzU4aSRWi

[codecov]

Codecov Report

❌ Patch coverage is 81.25000% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.89%. Comparing base (f8d34e4) to head (cd902d5).

Files with missing lines	Patch %	Lines
src/CodeGen_LLVM.cpp	81.25%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #9700      +/-   ##
==========================================
- Coverage   72.89%   72.89%   -0.01%     
==========================================
  Files         266      266              
  Lines       83478    83508      +30     
  Branches    20497    20503       +6     
==========================================
+ Hits        60855    60872      +17     
- Misses      16173    16197      +24     
+ Partials     6450     6439      -11     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.