In-place mutation allocates despite `assert_nonallocating=true`

[alexandergunnarson]

I'm trying to simply write a vector to a column-oriented matrix. What I have works for small sizes, but it allocates despite assert_nonallocating=true.

# NOTE — There's probably a simpler way to do this that's not so low-level. But maybe not, because
#        XLA is immutable by default.
function write_row!(xs_mut, i, v)
    v_col = Reactant.Ops.reshape(v, [length(v), 1])
    updated_xs = Reactant.Ops.dynamic_update_slice(xs_mut, v_col, [1, i])

    # Mutatively update xs with the new data
    Reactant.TracedUtils.set_mlir_data!(xs_mut, Reactant.TracedUtils.get_mlir_data(updated_xs))

    nothing
end

# Same issue with this
function write_row_simple!(xs_mut, i, v)
    xs_mut[:, i] = v
    nothing
end

I compile and run like so

T = Float32
dim = 768
max_n = 10_000_0
vectors_gpu = Reactant.@jit Reactant.Ops.fill(T(0), (dim, max_n))
i_default = 1 # just to have a value to pass into `write_row_rx!`
vector_gpu = Reactant.@jit Reactant.Ops.fill(T(0), (dim,))

write_row_rx! = Reactant.@compile assert_nonallocating=true sync=true write_row!(
    vectors_gpu, i_default, vector_gpu
)

write_row_rx!(vectors_gpu, 1, vector_gpu)

And get

julia> write_row_rx!(vectors_gpu, 1, vector_gpu)
2025-07-16 22:08:55.628804: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 28.61GiB (rounded to 30720000000)requested by op
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation.
Current allocation summary follows.
Current allocation summary follows.
2025-07-16 22:08:55.628872: I external/xla/xla/tsl/framework/bfc_allocator.cc:1049] BFCAllocator dump for GPU_0_bfc
2025-07-16 22:08:55.628888: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (256):  Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.628899: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (512):  Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.628909: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (1024):         Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.628947: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (2048):         Total Chunks: 1, Chunks in use: 1. 3.0KiB allocated for chunks. 3.0KiB in use in bin. 3.0KiB client-requested in use in bin.
2025-07-16 22:08:55.628968: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (4096):         Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.628992: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (8192):         Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629005: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (16384):        Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629034: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (32768):        Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629053: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (65536):        Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629074: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (131072):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629087: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (262144):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629132: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (524288):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629142: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (1048576):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629152: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (2097152):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629162: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (4194304):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629190: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (8388608):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629210: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (16777216):     Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629224: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (33554432):     Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629233: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (67108864):     Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629247: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (134217728):    Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-07-16 22:08:55.629263: I external/xla/xla/tsl/framework/bfc_allocator.cc:1056] Bin (268435456):    Total Chunks: 2, Chunks in use: 1. 32.00GiB allocated for chunks. 28.61GiB in use in bin. 28.61GiB client-requested in use in bin.
2025-07-16 22:08:55.629295: I external/xla/xla/tsl/framework/bfc_allocator.cc:1072] Bin for 28.61GiB was 256.00MiB, Chunk State:
2025-07-16 22:08:55.629320: I external/xla/xla/tsl/framework/bfc_allocator.cc:1078]   Size: 3.39GiB | Requested Size: 0B | in_use: 0 | bin_num: 20, prev:   Size: 3.0KiB | Requested Size: 3.0KiB | in_use: 1 | bin_num: -1
2025-07-16 22:08:55.629331: I external/xla/xla/tsl/framework/bfc_allocator.cc:1085] Next region of size 34359738368
2025-07-16 22:08:55.629342: I external/xla/xla/tsl/framework/bfc_allocator.cc:1105] InUse at 72fe14000000 of size 30720000000 next 1
2025-07-16 22:08:55.629356: I external/xla/xla/tsl/framework/bfc_allocator.cc:1105] InUse at 73053b0e0000 of size 3072 next 2
2025-07-16 22:08:55.629417: I external/xla/xla/tsl/framework/bfc_allocator.cc:1105] Free  at 73053b0e0c00 of size 3639735296 next 18446744073709551615
2025-07-16 22:08:55.629436: I external/xla/xla/tsl/framework/bfc_allocator.cc:1110]      Summary of in-use Chunks by size:
2025-07-16 22:08:55.629445: I external/xla/xla/tsl/framework/bfc_allocator.cc:1113] 1 Chunks of size 3072 totalling 3.0KiB
2025-07-16 22:08:55.629454: I external/xla/xla/tsl/framework/bfc_allocator.cc:1113] 1 Chunks of size 30720000000 totalling 28.61GiB
2025-07-16 22:08:55.629467: I external/xla/xla/tsl/framework/bfc_allocator.cc:1117] Sum Total of in-use chunks: 28.61GiB
2025-07-16 22:08:55.629479: I external/xla/xla/tsl/framework/bfc_allocator.cc:1119] Total bytes in pool: 34359738368 memory_limit_: 42909460070 available bytes: 8549721702 curr_region_allocation_bytes_: 34359738368
2025-07-16 22:08:55.629491: I external/xla/xla/tsl/framework/bfc_allocator.cc:1124] Stats:
Limit:                     42909460070
InUse:                     30720003072
MaxInUse:                  30720003072
NumAllocs:                           2
MaxAllocSize:              30720000000
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0

2025-07-16 22:08:55.629517: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] ******************************************************************************************__________
E0000 00:00:1752703735.629561 2610542 pjrt_stream_executor_client.cc:2939] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 30720000000 bytes. [tf-allocator-allocation-error='']
ERROR: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 30720000000 bytes. [tf-allocator-allocation-error='']

Stacktrace:
 [1] reactant_err(msg::Cstring)
   @ Reactant.XLA ~/.julia/packages/Reactant/bK8OR/src/xla/Utils.jl:12
 [2] macro expansion
   @ ~/.julia/packages/Reactant/bK8OR/src/xla/PJRT/LoadedExecutable.jl:195 [inlined]
 [3] execute_sharded
   @ ~/.julia/packages/Reactant/bK8OR/src/xla/PJRT/LoadedExecutable.jl:164 [inlined]
 [4] macro expansion
   @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3197 [inlined]
 [5] (::Reactant.Compiler.Thunk{…})(::Reactant.ConcretePJRTArray{…}, ::Int64, ::Reactant.ConcretePJRTArray{…})
   @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3644
 [6] top-level scope
   @ REPL[10]:1
Some type information was truncated. Use `show(err)` to see complete types.

Then if I try to run it again:

julia> write_row_rx!(vectors_gpu, 1, vector_gpu)
ERROR: Reactant.ConcretePJRTArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}} has already been donated!
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] getproperty(x::Reactant.ConcretePJRTArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, f::Symbol)
   @ Reactant ~/.julia/packages/Reactant/bK8OR/src/Types.jl:10
 [3] __resolve_device_and_client(client::Nothing, seen_args::Reactant.OrderedIdDict{Any, Any}, linear_args::Vector{Union{…}}, is_sharded::Bool)
   @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3343
 [4] compile_xla(f::Function, args::Tuple{…}; client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
   @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3408
 [5] compile_xla
   @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3377 [inlined]
 [6] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
   @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3462
 [7] top-level scope
   @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:2557
Some type information was truncated. Use `show(err)` to see complete types.

Am I missing something? In CUDA.jl this kind of mutation is easy, but it doesn't optimize like Reactant.jl does (~4x slower for brute-force similarity search). I'm considering getting a pointer to the data and passing it to CUDA.jl, but that seems dangerous. Any suggestions?

Also curious how the concurrency/atomicity model works. What are the visibility guarantees of Reactant.jl operations? I'd want writes to be linearizably reflected by subsequent reads to the same location.

[alexandergunnarson]

It's interesting — it only allocates once.

I tried decreasing max_n to e.g. 2M instead of 10M so the OOM wouldn't happen. The first allocation you see below is for vectors_gpu. The second allocation is when I call write_row_rx! the first time. But subsequent times are indeed allocation-free. Not sure how to accommodate for this. Perhaps some scratch space is being used?

Also, the operation takes 38 ms on my L40S where max_n is 2M, and 19 ms where max_n is 1M. Basically it's writing the whole matrix every time rather than updating only the tiny part that needs updating. Not only is this highly suboptimal, but I wonder whether there are concurrency implications (writes blocking reads).

[wsmoses]

can you paste both the result of @ code_hlo and @ code_xla [insert this instead of @ compile]

[wsmoses]

also I added a PJRT API UninitPJRTBuffer (

Reactant.jl/deps/ReactantExtra/API.cpp

Line 802 in 17130b7

extern "C" PjRtBuffer *UninitPJRTBuffer(PjRtClient *client, PjRtDevice *device,

) that we should use to replace similar [and thus wont allocate host side to copy into for an initial vector, instead of your JIT fill, assuming you just wanted some data]

[alexandergunnarson]

Both @code_hlo and @code_xla, with both write_row_rx! and write_row_simple_rx!, yield this stacktrace:

julia> Reactant.@code_hlo write_row_rx!(vectors_gpu, 1, vector_gpu)
"Inconsistent guaranteed error IR   1 ─      nothing::Nothing\n3 2 ─ %2 = Core.getfield(_5, 1)::Reactant.TracedRArray{Float32, 2}\n  │   %3 = Core.getfield(_5, 2)::Any\n  │   %4 = Core.getfield(_5, 3)::Any\n  │   %5 = (_4)(%2, %3, %4)::Any\n  └──      return %5\n  "
ERROR:
The Reactant-compiled function `Reactant.Compiler.Thunk{typeof(write_row!), Symbol("##write_row!_reactant#347"), Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, Int64, Reactant.ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}, false, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}}` exists, but no method is defined for this combination of argument types.
You passed in arguments with types
        (Reactant.TracedRArray{Float32, 2}, Int64, Reactant.TracedRArray{Float32, 1})
However the method you are calling was compiled for arguments with types
        (Reactant.ConcretePJRTArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, Int64, Reactant.ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}})
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3652 [inlined]
  [2] Thunk
    @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3644 [inlined]
  [3] (::Nothing)(none::Reactant.Compiler.Thunk{…}, none::Tuple{…})
    @ Reactant ./<missing>:0
  [4] macro expansion
    @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3652 [inlined]
  [5] Thunk
    @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:3644 [inlined]
  [6] call_with_reactant(::Reactant.Compiler.Thunk{…}, ::Reactant.TracedRArray{…}, ::Int64, ::Reactant.TracedRArray{…})
    @ Reactant ~/.julia/packages/Reactant/bK8OR/src/utils.jl:0
  [7] #apply#54
    @ ~/.julia/packages/Reactant/bK8OR/src/utils.jl:3 [inlined]
  [8] (::Nothing)(none::Reactant.var"##apply#54", none::@Kwargs{}, none::typeof(Reactant.apply), none::Reactant.Compiler.Thunk{…}, none::Tuple{…})
    @ Reactant ./<missing>:0
  [9] call_with_reactant(::Reactant.MustThrowError, ::Reactant.var"##apply#54", ::@Kwargs{}, ::typeof(Reactant.apply), ::Reactant.Compiler.Thunk{…}, ::Reactant.TracedRArray{…}, ::Any, ::Any)
    @ Reactant ~/.julia/packages/Reactant/bK8OR/src/utils.jl:877
 [10] apply
    @ ~/.julia/packages/Reactant/bK8OR/src/utils.jl:2 [inlined]
 [11] (::Nothing)(none::typeof(Reactant.apply), none::Reactant.Compiler.Thunk{…}, none::Tuple{…})
    @ Reactant ./<missing>:0
 [12] call_with_reactant(::typeof(Reactant.apply), ::Reactant.Compiler.Thunk{…}, ::Reactant.TracedRArray{…}, ::Int64, ::Reactant.TracedRArray{…})
    @ Reactant ~/.julia/packages/Reactant/bK8OR/src/utils.jl:2
 [13] make_mlir_fn(f::typeof(Reactant.apply), args::Tuple{…}, kwargs::@NamedTuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/bK8OR/src/TracedUtils.jl:330
 [14] make_mlir_fn
    @ ~/.julia/packages/Reactant/bK8OR/src/TracedUtils.jl:260 [inlined]
 [15] #make_mlir_fn#6
    @ ~/.julia/packages/Reactant/bK8OR/src/TracedUtils.jl:282 [inlined]
 [16] make_mlir_fn
    @ ~/.julia/packages/Reactant/bK8OR/src/TracedUtils.jl:260 [inlined]
 [17] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Reactant.Compiler.Thunk{…}, args::Tuple{…}, compile_options::Reactant.CompileOptions, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:1510
 [18] compile_mlir!
    @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:1478 [inlined]
 [19] (::Reactant.Compiler.var"#9#10"{Nothing, Reactant.Compiler.Thunk{…}, Tuple{…}})(ctx::Reactant.MLIR.IR.Context)
    @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:1335
 [20] with_context(f::Reactant.Compiler.var"#9#10"{Nothing, Reactant.Compiler.Thunk{…}, Tuple{…}}; allow_use_existing::Bool)
    @ Reactant.MLIR.IR ~/.julia/packages/Reactant/bK8OR/src/mlir/IR/Context.jl:99
 [21] with_context
    @ ~/.julia/packages/Reactant/bK8OR/src/mlir/IR/Context.jl:82 [inlined]
 [22] compile_mlir(f::Reactant.Compiler.Thunk{…}, args::Tuple{…}; client::Nothing, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:1331
 [23] top-level scope
    @ ~/.julia/packages/Reactant/bK8OR/src/Compiler.jl:2557
Some type information was truncated. Use `show(err)` to see complete types.

This is (perhaps unsurprisingly) the case whether or not I do

vectors_gpu = Reactant.@jit Reactant.Ops.fill(Float32(0), (dim, max_n))
vector_gpu = Reactant.@jit Reactant.Ops.fill(Float32(0), (dim,))

or

vectors_gpu = fill(0.0f0, (dim, max_n)) |> Reactant.to_rarray
vector_gpu = fill(0.0f0, (dim,)) |> Reactant.to_rarray

or

vectors_gpu = fill(0.0f0, (dim, max_n)) |> Reactant.ConcreteRArray
vector_gpu = fill(0.0f0, (dim,)) |> Reactant.ConcreteRArray

Pretty sure this didn't happen in the previous release I was using (which is why I settled on the Reactant.Ops.dynamic_update_slice approach), but I could be wrong.

[alexandergunnarson]

Interesting that write_row_rx!(vectors_gpu, 1, vector_gpu) doesn't error out, whereas @code_* does.

[wsmoses]

Ah sorry, I meant:

@code_hlo write_row!(
    vectors_gpu, i_default, vector_gpu
)

in place of

@compile ...

[alexandergunnarson]

Ah oops, thanks for that clarification. My brain was not on.

Interestingly, I see the same @code_hlo output for these two. In Reactant 0.2.126, I'm pretty sure I remember seeing dynamic_update_slice in the @code_hlo for write_row!. (Otherwise I wouldn't have written different versions.) It makes a little more sense now why there's the (one-time) allocation going on. There's still the mystery of 1) why the behavior changed and 2) why it's allocating despite assert_nonallocating=true.

julia> Reactant.@code_hlo write_row!(vectors_gpu, i_default, vector_gpu)
module @"reactant_write_row!" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<1000000x768xf32> {tf.aliasing_output = 0 : i32}, %arg1: tensor<768xf32>) -> tensor<1000000x768xf32> {
    %0 = stablehlo.reshape %arg1 : (tensor<768xf32>) -> tensor<1x768xf32>
    %1 = stablehlo.slice %arg0 [1:1000000, 0:768] : (tensor<1000000x768xf32>) -> tensor<999999x768xf32>
    %2 = stablehlo.concatenate %0, %1, dim = 0 : (tensor<1x768xf32>, tensor<999999x768xf32>) -> tensor<1000000x768xf32>
    return %2 : tensor<1000000x768xf32>
  }
}

julia> Reactant.@code_hlo write_row_simple!(vectors_gpu, i_default, vector_gpu)
module @reactant_write_r... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<1000000x768xf32> {tf.aliasing_output = 0 : i32}, %arg1: tensor<768xf32>) -> tensor<1000000x768xf32> {
    %0 = stablehlo.reshape %arg1 : (tensor<768xf32>) -> tensor<1x768xf32>
    %1 = stablehlo.slice %arg0 [1:1000000, 0:768] : (tensor<1000000x768xf32>) -> tensor<999999x768xf32>
    %2 = stablehlo.concatenate %0, %1, dim = 0 : (tensor<1x768xf32>, tensor<999999x768xf32>) -> tensor<1000000x768xf32>
    return %2 : tensor<1000000x768xf32>
  }
}

julia> Reactant.@code_xla write_row!(vectors_gpu, i_default, vector_gpu)
HloModule reactant_write_row_, is_scheduled=true, input_output_alias={ {}: (0, {}, may-alias) }, entry_computation_layout={(f32[1000000,768]{1,0}, f32[768]{0})->f32[1000000,768]{1,0}}, frontend_attributes={fingerprint_before_lhs="b4fc583b08d25c9b2e5912fdf4ec026a"}

%fused_concatenate (param_0.2: f32[768], param_1.2: f32[1000000,768]) -> f32[1000000,768] {
  %param_0.2 = f32[768]{0} parameter(0)
  %bitcast.2 = f32[1,768]{1,0} bitcast(%param_0.2), metadata={op_name="arg3 (path=(:args, 3))"}
  %param_1.2 = f32[1000000,768]{1,0} parameter(1)
  %slice.4.1 = f32[999999,768]{1,0} slice(%param_1.2), slice={[1:1000000], [0:768]}
  ROOT %concatenate.5.1 = f32[1000000,768]{1,0} concatenate(%bitcast.2, %slice.4.1), dimensions={0}
}

ENTRY %main.6 (Arg_0.1: f32[1000000,768], Arg_1.2: f32[768]) -> f32[1000000,768] {
  %Arg_1.2 = f32[768]{0} parameter(1), metadata={op_name="arg3 (path=(:args, 3))"}
  %Arg_0.1 = f32[1000000,768]{1,0} parameter(0), metadata={op_name="arg1 (path=(:args, 1))"}
  %input_concatenate_fusion = f32[1000000,768]{1,0} fusion(%Arg_1.2, %Arg_0.1), kind=kInput, calls=%fused_concatenate
  ROOT %copy.1 = f32[1000000,768]{1,0} copy(%input_concatenate_fusion)
}

julia> Reactant.@code_xla write_row_simple!(vectors_gpu, i_default, vector_gpu)
HloModule reactant_write_r..., is_scheduled=true, input_output_alias={ {}: (0, {}, may-alias) }, entry_computation_layout={(f32[1000000,768]{1,0}, f32[768]{0})->f32[1000000,768]{1,0}}, frontend_attributes={fingerprint_before_lhs="39717f9a792e1022750e39ddf4cf8bda"}

%fused_concatenate (param_0.2: f32[768], param_1.2: f32[1000000,768]) -> f32[1000000,768] {
  %param_0.2 = f32[768]{0} parameter(0)
  %bitcast.2 = f32[1,768]{1,0} bitcast(%param_0.2), metadata={op_name="arg3 (path=(:args, 3))"}
  %param_1.2 = f32[1000000,768]{1,0} parameter(1)
  %slice.4.1 = f32[999999,768]{1,0} slice(%param_1.2), slice={[1:1000000], [0:768]}
  ROOT %concatenate.5.1 = f32[1000000,768]{1,0} concatenate(%bitcast.2, %slice.4.1), dimensions={0}
}

ENTRY %main.6 (Arg_0.1: f32[1000000,768], Arg_1.2: f32[768]) -> f32[1000000,768] {
  %Arg_1.2 = f32[768]{0} parameter(1), metadata={op_name="arg3 (path=(:args, 3))"}
  %Arg_0.1 = f32[1000000,768]{1,0} parameter(0), metadata={op_name="arg1 (path=(:args, 1))"}
  %input_concatenate_fusion = f32[1000000,768]{1,0} fusion(%Arg_1.2, %Arg_0.1), kind=kInput, calls=%fused_concatenate
  ROOT %copy.1 = f32[1000000,768]{1,0} copy(%input_concatenate_fusion)
}

[wsmoses]

it still should be nonallocating here? the output aliases the input (so no allocation for return), and XLA should be able to handle this concat without a new alloc (cc @avik-pal )

[alexandergunnarson]

Problem 1:

• I manually allocate Reactant.@jit Reactant.Ops.fill(T(0), (dim, max_n)). Memory allocated is M.
• I do Reactant.@compile
• I run write_row_rx!(vectors_gpu, 1, vector_gpu). It allocates an M-sized chunk of memory.
• I run write_row_rx! again, arbitrarily many times. It allocates no memory.

Problem 2:
Targeted dynamic_update_slice gets "washed out" into a suboptimal whole-matrix-update approach. I believe this is a regression.

[avik-pal]

BufferAssignment:
allocation 0: size 307200000, parameter 0, shape |f32[100000,768]| at ShapeIndex {}, output shape is |f32[100000,768]|, maybe-live-out:
 value: <5 Arg_0.1 @0> (size=307200000,offset=0): f32[100000,768]{1,0}
 value: <7 copy.1 @0> (size=307200000,offset=0): f32[100000,768]{1,0}
allocation 1: size 3072, parameter 1, shape |f32[768]| at ShapeIndex {}:
 value: <4 Arg_1.2 @0> (size=3072,offset=0): f32[768]{0}
allocation 2: size 307200000, preallocated-temp:
 value: <6 input_concatenate_fusion @0> (size=307200000,offset=0): f32[100000,768]{1,0}

Total bytes used: 614403072 (585.94MiB)

Used values:
<4 Arg_1.2 @0>
 positions:
  Arg_1.2
 uses:
  input_concatenate_fusion, operand 0
 from instruction: %Arg_1.2 = f32[768]{0} parameter(1), metadata={op_name="arg3 (path=(:args, 3))"}
<5 Arg_0.1 @0>
 positions:
  Arg_0.1
 uses:
  input_concatenate_fusion, operand 1
 from instruction: %Arg_0.1 = f32[100000,768]{1,0} parameter(0), metadata={op_name="arg1 (path=(:args, 1))"}
<6 input_concatenate_fusion @0>
 positions:
  input_concatenate_fusion
 uses:
  copy.1, operand 0
 from instruction: %input_concatenate_fusion = f32[100000,768]{1,0} fusion(%Arg_1.2, %Arg_0.1), kind=kInput, calls=%fused_concatenate
<7 copy.1 @0>
 positions:
  copy.1
 uses:
 from instruction: %copy.1 = f32[100000,768]{1,0} copy(%input_concatenate_fusion)


HloLiveRange (max 4):
  InstructionSequence:
    0:Arg_1.2
    1:Arg_0.1
    2:input_concatenate_fusion
    3:copy.1
  BufferLiveRange:
    Arg_1.2{}:0-4
    Arg_0.1{}:0-2
    input_concatenate_fusion{}:2-3
    copy.1{}:3-4
  Live ranges at 3 (peak):
    copy.1{}: 307200000 bytes (cumulative: 307200000 bytes)
    input_concatenate_fusion{}: 307200000 bytes (cumulative: 614400000 bytes)
    Arg_1.2{}: 3072 bytes (cumulative: 614403072 bytes)

Total bytes used: 614403072 (585.94MiB)

Allocations sorted by size:

cumulative_size; total_size - cumulative_size; allocation
------------------------------------------------------------------------------
 292.97MiB( 50%);  292.97MiB; allocation 0: size 292.97MiB, parameter 0, shape |f32[100000,768]| at ShapeIndex {}, output shape is |f32[100000,768]|, maybe-live-out:
 585.94MiB(100%);     3.0KiB; allocation 2: size 292.97MiB, preallocated-temp:
 585.94MiB(100%);         0B; allocation 1: size 3.0KiB, parameter 1, shape |f32[768]| at ShapeIndex {}:


Allocations sorted by size with their values:
allocation 0: size 292.97MiB, parameter 0, shape |f32[100000,768]| at ShapeIndex {}, output shape is |f32[100000,768]|, maybe-live-out:

	cumulative_size;       size;       offset; used_by_n_values; shapes_list
	------------------------------------------------------------
	292.97MiB(100%);  292.97MiB;            0;                2; 2×f32[100000,768]

allocation 2: size 292.97MiB, preallocated-temp:
allocation 1: size 3.0KiB, parameter 1, shape |f32[768]| at ShapeIndex {}:

[alexandergunnarson]

(I should mention I'm running with XLA_REACTANT_GPU_PREALLOCATE=false.)

@avik-pal your point is to show that it indeed does allocate? Trying to grok the output you posted

[avik-pal]

Yeah it does allocate.

I was just dumping the logs here. Didn't mean to respond to anything yet 😅 (still analyzing)

[wsmoses]

there's a concat_to_dus pas we should run at the end [since it otherwise may interfere with the dus_to_concat pass], that we could run to remedy.

[avik-pal]

HloModule reactant_write_row_, is_scheduled=true, input_output_alias={ {}: (0, {}, may-alias) }, entry_computation_layout={(f32[100000,768]{1,0}, f32[768]{0})->f32[100000,768]{1,0}}, frontend_attributes={fingerprint_before_lhs="001d76dd1a5a232d00c7a03f0464e990"}

%fused_concatenate (param_0.2: f32[768], param_1.2: f32[100000,768]) -> f32[100000,768] {
  %param_0.2 = f32[768]{0} parameter(0)
  %bitcast.2 = f32[1,768]{1,0} bitcast(%param_0.2), metadata={op_name="arg3 (path=(:args, 3))"}
  %param_1.2 = f32[100000,768]{1,0} parameter(1)
  %slice.4.1 = f32[99999,768]{1,0} slice(%param_1.2), slice={[1:100000], [0:768]}
  ROOT %concatenate.5.1 = f32[100000,768]{1,0} concatenate(%bitcast.2, %slice.4.1), dimensions={0}
}

ENTRY %main.6 (Arg_0.1: f32[100000,768], Arg_1.2: f32[768]) -> f32[100000,768] {
  %Arg_1.2 = f32[768]{0} parameter(1), metadata={op_name="arg3 (path=(:args, 3))"}
  %Arg_0.1 = f32[100000,768]{1,0} parameter(0), metadata={op_name="arg1 (path=(:args, 1))"}
  %input_concatenate_fusion = f32[100000,768]{1,0} fusion(%Arg_1.2, %Arg_0.1), kind=kInput, calls=%fused_concatenate
  ROOT %copy.1 = f32[100000,768]{1,0} copy(%input_concatenate_fusion)
}

Can't XLA handle aliasing between fusion inputs (param_1.2: f32[100000,768]) and outputs (%input_concatenate_fusion)? %input_concatenate_fusion is the allocated temp (probably)

But yeah a concat to dus should fix this

[alexandergunnarson]

Curious — will concat to dus make it so the matrix is written to in a much more targeted way (single column) as opposed to a whole-matrix update, or does this not fix that?