API to initialize directly on device

[alexandergunnarson]

Love Reactant.jl — an extremely exciting piece of work! 🎉

I was wondering — is there a way to wrap an existing GPU array? Or even just create an empty Reactant array with particular dimensions?

Doing this causes an OOM on my L40S GPU:

Reactant.set_default_backend("gpu")
vectors_gpu = CUDA.zeros(Float32, 768, 10_000_000)
vectors_rx = Reactant.to_rarray(vectors_gpu)

vectors_gpu takes up around 75% of my GPU, which is intended, but vectors_rx allocates on both CPU and GPU. It just so happens that I have lots of RAM, so I can do:

Reactant.set_default_backend("gpu")
vectors_cpu = zeros(Float32, 768, 10_000_000)
vectors_rx = Reactant.to_rarray(vectors_cpu)

However, notably, Reactant.to_rarray copies the array. The problem is mitigated, of course, by not maintaining a reference to the source CPU array:

Reactant.set_default_backend("gpu")
vectors_rx = zeros(Float32, 768, 10_000_000) |> Reactant.to_rarray

However, in either case, it's wasteful to have a mirror of the GPU array on CPU when I really don't need it in my application.

How can I mitigate this problem?

[avik-pal]

We might be able to use some sort of DLPack API for this (xref #1362). Though I don't think XLA exposes any API for transferring reference to an existing device array to PJRT

[alexandergunnarson]

Appreciate your response!

Hmm. Is it considered normal in XLA, then, to force the same amount of RAM in use as VRAM? That seems wasteful, no? I’d be so surprised if this is the case. Many GPU machines have very low RAM because they simply don’t need it.

[alexandergunnarson]

Perhaps it’s possible to, in XLA, express the shape of a tensor and then instantiate it only on GPU (not CPU) when the time comes. I don’t necessarily need to wrap the GPU array I suppose; I primarily need to just not double up on memory.

[avik-pal]

I see. Directly initialization on device memory is definitely possible. We should implement an API for this, but for the time being the following should work:

@jit Reactant.Ops.fill(3.0f0, (3, 4, 5))

[avik-pal]

Similarly for random numbers

julia> rng = Random.default_rng()
TaskLocalRNG()

julia> function rand_on_device(rng)
           x = Reactant.Ops.fill(0.0f0, (3, 4, 5))
           rand!(rng, x)
           return x
       end
rand_on_device (generic function with 2 methods)

julia> @code_hlo rand_on_device(rng)
module @reactant_rand_on... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main() -> tensor<5x4x3xf32> {
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<3x4x5xf32>
    %c = stablehlo.constant dense<1065353216> : tensor<3x4x5xui32>
    %c_0 = stablehlo.constant dense<9> : tensor<3x4x5xui32>
    %c_1 = stablehlo.constant dense<[13088689619725983284, 4141439666682415456]> : tensor<2xui64>
    %output_state, %output = stablehlo.rng_bit_generator %c_1, algorithm =  DEFAULT : (tensor<2xui64>) -> (tensor<2xui64>, tensor<3x4x5xui32>)
    %0 = stablehlo.shift_right_logical %output, %c_0 : tensor<3x4x5xui32>
    %1 = stablehlo.or %0, %c : tensor<3x4x5xui32>
    %2 = stablehlo.bitcast_convert %1 : (tensor<3x4x5xui32>) -> tensor<3x4x5xf32>
    %3 = stablehlo.subtract %2, %cst : tensor<3x4x5xf32>
    %4 = stablehlo.transpose %3, dims = [2, 1, 0] : (tensor<3x4x5xf32>) -> tensor<5x4x3xf32>
    return %4 : tensor<5x4x3xf32>
  }
}

[alexandergunnarson]

Amazing, thank you!! Looks to be just what I needed. Will try it out.

Any way to grab a CuArray reference to the underlying data? I ask because I’m currently using CUDA.jl to achieve atomic writes. However perhaps that’s not necessary and XLA will handle atomicity. I haven’t looked at the lowered code yet.

[avik-pal]

Any way to grab a CuArray reference to the underlying data?

No there isn't a nice way to make them interact (atleast that I know of)

[alexandergunnarson]

All right, thank you!