Replay the stores through a Julia array's data pointer when rematerializing it

[wsmoses]

Fixes the wrong gradient behind the Julia 1.10 Turing failure in EnzymeAD/Enzyme.jl#3734 (demo_dot_assume_observe_index, reverse mode with runtime activity). The bug is older than that; #3325 only exposed it.

Cause

A Julia array allocated in a function differentiated in split mode is rematerialized in the reverse pass: the allocation is redone and the stores into it are replayed. The use walk in computeForwardingProperties stops at loads, so the array's data pointer (an addrspace(13) pointer the allocator sets up, loaded from the object) and everything reached through it were invisible. The fill stores were never recorded as stores to replay, and the data loads never as loads that must be reloadable. The reverse pass re-ran the data pointer load on the fresh object and read its unfilled buffer.

This went unnoticed while the reads were cached by other means. A load through a pointer loaded from memory was cached whenever any write at all followed the call in the caller. Since #3325 only unknown writes count, so the callee reloads. In the Turing model, sqmahal's x .- μ temporary is recreated empty by sqmahal's reverse before invwsumsq's reverse reads it (sqmahal's later writes are all to its sret and allocas).

The same happens without any callee once nothing follows the loads in the function. On Julia 1.10 with Enzyme_jll 0.0.299 this already returns uninitialized memory:

@noinline function mid(x, w)
    t = x .- 1.0
    s = 0.0
    @inbounds for i in eachindex(t); s += t[i] * t[i] / w[i]; end
    return s
end
top(x, w) = (r = mid(x, w); r * r)   # mid is differentiated in split mode
Enzyme.autodiff(Reverse, top, Active, Duplicated(x, dx), Duplicated(w, dw))

Julia 1.11+ is unaffected: its data pointer flows through julia.gc_loaded, a call the walk already follows, so the stores are replayed there.

Change

Walk through an addrspace(13) pointer loaded from the allocation as if it were pointer arithmetic on it, as already happens for julia.gc_loaded. The data stores are then replayed before the reverse pass reads the data, and the existing checks (a store that may execute after a load, writing calls) apply to the data loads and stores.

Also tried: refusing to promote such an allocation when its primal is needed in the reverse pass, so it is cached whole. That fixes the model too, but changes rematprimal.ll and keeps arrays alive longer, so it was not kept.

Tests

• New ReverseMode/julia_array_data_remat.ll, reduced from the Julia 1.10 IR of the reproducer above (the array fill in a split-mode callee of a combined function). Without this change the reverse of julia_mid reallocates the array and re-runs the fill loop without its stores, then reloads the data from the fresh buffer; the CHECK lines fail at the missing load/store. With it the fill loop is replayed.
• lit (LLVM 15): all pass except the pre-existing partial_int_window.ll. lit (LLVM 20): no new failures against a build of main.
• Enzyme.jl test suite on Julia 1.10 and 1.11: passes.
• Enzyme.jl Turing integration suite on Julia 1.10: passes, including the model that failed.

🤖 Generated with Claude Code

https://claude.ai/code/session_01G84mmXoaait2MBqCH5BKNb