Bad XLA codegen for exact TopK

[x66ccff]

Hi, I'm not sure if I'm using Reactant correctly. I've found that in my code, the current performance bottleneck seems to be the operation of selecting certain positions from a huge tensor to the CPU. Specifically:

val_R, idx_R = model.top_k_compiled(-mean_errors_R, model.PSRN_topk)

@info "idx_R:"
@time @info idx_R

I find that printing idx_R here is very slow. I understand this slowness may be due to the synchronization needed from GPU to CPU. But I can't understand why this process is so much slower than PyTorch? 🤔

Julia Implementation with Reactant (~1.8 seconds)

using Reactant
n = 1_000_000_000  # 1 billion elements
k = 100
x = rand(Float32, n)
xr = Reactant.to_rarray(x)
partialsortperm_compiled = @compile partialsortperm(xr, 1:k)

GC.gc()

# First benchmark: compute but don't access the result data
@time begin
    res = partialsortperm_compiled(xr, 1:k)
    @info 1
end

[ Info: 1
  0.870469 seconds (1.86 M allocations: 93.889 MiB, 99.52% compilation time)

GC.gc()

# Second benchmark: compute and access the result data (forcing GPU->CPU transfer)
@time begin
    res = partialsortperm_compiled(xr, 1:k)
    @info res
end

[ Info: SubArray{Int64, 1, ConcretePJRTArray{Int64, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, Tuple{UnitRange{Int64}}, true}([955852229, 932404690, 906993663, 899721285, 894668984, 877960574, 871792407, 796790415, 792612165, ..337289477, 332445067, 267677250, 222104083, 214054217, 194744078, 193590981, 193251762, 185722187, 114722941])
 16.426357 seconds (93.05 M allocations: 4.559 GiB, 3.10% gc time, 99.32% compilation time)

When running this a second time after compilation, the timing becomes much better (around 1.8 seconds), but it's still significantly slower than the PyTorch equivalent.

Python Implementation with PyTorch for Comparison (0.001586 seconds)

For comparison, I ran the equivalent operation in Python with PyTorch, and it's significantly faster:

import torch
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.rand(1_000_000_000, dtype=torch.float32, device=device)  # 1 billion elements

start_time = time.time()
indices = torch.topk(-x, 100)[1]
print(f"Time: {time.time() - start_time:.6f}s, {indices.cpu().tolist()}")

Time: 0.001586s

As you can see, the Python/PyTorch implementation is about 1000× faster for the same operation, even when explicitly calling .cpu().tolist() to transfer data back to CPU.

Is this a known issue with Reactant's GPU-to-CPU data transfer? Even after accounting for Julia's compilation time, the actual data transfer seems much slower than PyTorch. Are there any recommended workarounds or optimizations I could implement to improve this performance bottleneck?

[avik-pal]

@time begin
    res = partialsortperm_compiled(xr, 1:k)
    @info res
end

this will do a device to host transfer to display the array

[x66ccff]

yeah, i understand this, but even this is slower than torch code

@time begin
    res = partialsortperm_compiled(xr, 1:k)
    @info 1
end

which cost around 0.8 seconds. Not to mention that pytorch code also display the values

[wsmoses]

What if you compile the following
Function foo(x,rng)
copy(x[rng])
End

Instead

[x66ccff]

@wsmoses i didn't quite understand. could you explain it please? like this?

julia> copyker=function foo(x,rng)
              copy(x[rng])
              end
foo (generic function with 1 method)

julia> fooker=@compile foo(xr,3)
ERROR: Scalar indexing is disallowed.

[wsmoses]

Oh sorry use partialsortperm not getindex

[x66ccff]

using Reactant
n = 1_000_000_000  # 1 billion elements
k = 100
x = rand(Float32, n)
xr = Reactant.to_rarray(x)
partialsortperm_compiled = @compile copy(partialsortperm(xr, 1:k))

GC.gc()

# First benchmark: compute but don't access the result data
@time begin
    res = partialsortperm_compiled(xr, 1:k)
    @info 1
end

this? 😂 add a copy here?

julia> partialsortperm_compiled = @compile copy(partialsortperm(xr, 1:k))
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::ConcretePJRTArray, ::Vararg{Int, N}) resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

[wsmoses]

the compile macro only compiles the outermost function call. you need to wrap it in an outer function

[x66ccff]

yeah this is faster No there is no difference. if i do not need to access (display) the value when run the @info 1 block 2nd time (cost 0.001s)

But still cost 1.8 seconds to display the value (@info res)

using Reactant
n = 1_000_000_000  # 1 billion elements
k = 100
x = rand(Float32, n)
xr = Reactant.to_rarray(x)



function foo(xr,k)
copy(partialsortperm(xr, 1:k))
end
partialsortperm_compiled = @compile foo(xr,k)

GC.gc()

# First benchmark: compute but don't access the result data
@time begin
    res = partialsortperm_compiled(xr, k)
    @info 1
end

GC.gc()

# Second benchmark: compute and access the result data (forcing GPU->CPU transfer)
@time begin
    res = partialsortperm_compiled(xr, k)
    @info res
end

[avik-pal]

julia> @code_hlo partialsortperm(xr, 1:100)
module @reactant_partial... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<1000000000xf32>) -> tensor<1000000000xi64> {
    %c = stablehlo.constant dense<1> : tensor<1000000000xi32>
    %values, %indices = chlo.top_k(%arg0, k = 1000000000) : tensor<1000000000xf32> -> (tensor<1000000000xf32>, tensor<1000000000xi32>)
    %0 = stablehlo.add %indices, %c : tensor<1000000000xi32>
    %1 = stablehlo.convert %0 : (tensor<1000000000xi32>) -> tensor<1000000000xi64>
    %2 = stablehlo.reverse %1, dims = [0] : tensor<1000000000xi64>
    return %2 : tensor<1000000000xi64>
  }
}

our codegen for this isn't really good. We need to do a partialsortperm(negate x, 1:k; rev=true) if k is a "small" number

[avik-pal]

#1198 (comment)

turns out XLA doesn't have exact topk defined nicely. You can use https://github.com/EnzymeAD/Reactant.jl/pull/1198/files#diff-8a00fd5468cbab4c89477359a95159628ab1a84e095b2c1008c218a58e46309aR217 to lower to ApproxTopK but that works only on TPUs