Split reads and writes in aliasing optimization

[Pangoraw]

the optimization that bypasses XLA was reading from args, but if the
value of the argument is over-written from being update the computation
then using the previous value would read the wrong value (already updated).

Closes #2870.

[wsmoses]

@Pangoraw can you explai the difference in behavior here, in terms of the code_typed of the generated thunk code? [and/or change to the mlir ]

[Pangoraw]

if you take a look at the reproducer:

mutable struct State{A}
    u::A
    u⁰::A
end

function step!(state)
    state.u⁰ .= state.u
    state.u = state.u .+ 1.0
    return nothing
end

the previous pseudo code_lowered would look something like this:

# u⁰ is donated
linearized_results = XLA.execute_sharded(thunk.exec, thunk.device, (state.u, state.u⁰))

# Unpack result buffers
result_buffer_1 = linearized_results[1]

result = nothing

# Build a cache mapping traced arrays/numbers -> concrete PJRT arrays/numbers
cache_dict = IdDict{Union{TracedRArray, TracedRNumber}, Union{ConcretePJRTArray, ConcretePJRTNumber}}()

# Write result buffer back into state.u
traced_setfield_buffer!(Val{:PJRT}(), cache_dict, result_buffer_1, args[1], 1, (1,))

# Snapshot field 1's data AFTER the result has been written back
argpath_value1 = map(copy, state.u.data)

# Propagate that snapshot into field 2 of args[1] — twice
# once because of (:args, 1, 2) <- (:args, 1, 1)  
traced_setfield!(state.u⁰, :data, argpath_value1, (2,))

# once because of (:resargs, 1, 2) <- (:args, 1, 1)
traced_setfield!(state.u⁰, :data, argpath_value1, (2,))

# therefore state.u⁰ is equal to state.u POST execution

with this pr this looks something like this:

# u⁰ is donated
linearized_results = XLA.execute_sharded(thunk.exec, thunk.device, (state.u, state.u⁰))

result_buffer_1 = Base.getindex(linearized_results, 1)

# read result 1 value before overwriting
argpath_value1 = copy(state.u.data)

# Write result buffer back into state.u
Reactant.Compiler.traced_setfield_buffer!(
  Val{:PJRT}(), cache_dict, result_buffer_1, args[1], 1, (1,))

# (:resargs, 1, 2) <- (:args, 1, 1)  
traced_setfield!(state.u⁰, :data, argpath_value1, (2,))

# noop since already written before (:args, 1, 2) <- (:args, 1, 1)

# therefore state.u⁰ is equal to state.u PRE execution

[wsmoses]

If we need two different results corresponding to the same logical value I think my preference here would for execute_sharded to return that number of the results [aka mlir should undo the redundant return

[wsmoses]

thinking more, I think I'd be fine disbaling [or perhaps changing to a flag, disabled by default] the redundant return elimination opt we have