Reactant 10x slower than CUDA on simple broadcasting

[albertomercurio]

I have this simple example, where I compare the performance of Reactant compared to CUDA. However, it is 10x slower.

using LinearAlgebra
using CUDA
using Reactant
using BenchmarkTools

T = Float32
N = 1000000

# coeffs = rand(T, N-1)
coeffs = sqrt.(T.(1:N-1))
v = rand(Complex{T}, N)
w = similar(v)

coeffs_gpu = CuArray(coeffs)
v_gpu = CuArray(v)
w_gpu = similar(v_gpu)

coeffs_reactant = Reactant.to_rarray(coeffs);
v_reactant = Reactant.to_rarray(v)
w_reactant = Reactant.to_rarray(w)

function myfunc(w, coeffs, v)
    N = length(v)
    @views w[1:(N - 1)] .= coeffs .* v[2:N]

    return w
end

@benchmark myfunc($w, $coeffs, $v) # 325.731 μs
@benchmark myfunc($w_gpu, $coeffs_gpu, $v_gpu) # 5.384 μs

myfunc_compiled = @compile myfunc(w_reactant, coeffs_reactant, v_reactant)

@benchmark myfunc_compiled($w_reactant, $coeffs_reactant, $v_reactant) # 49.838 μs

I'm using the GPU on Reactant I guess

julia> Reactant.devices()
1-element Vector{Reactant.XLA.PJRT.Device}:
 Reactant.XLA.PJRT.Device(Ptr{Nothing}(0x000000002b6aede0), "CUDA:0 NVIDIA GeForce RTX 4090"
zzz

[Pangoraw]

thanks for the report!

looks like reactant could do a better job here:

module @reactant_myfunc attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<1000000xcomplex<f32>> {enzymexla.memory_effects = [], tf.aliasing_output = 0 : i32}, %arg1: tensor<999999xf32> {enzymexla.memory_effects = []}, %arg2: tensor<1000000xcomplex<f32>> {enzymexla.memory_effects = []}) -> tensor<1000000xcomplex<f32>> attributes {enzymexla.memory_effects = []} {
    %0 = stablehlo.reshape %arg0 {enzymexla.complex_is_purely_imaginary = [#enzymexla<guaranteed NOTGUARANTEED>], enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<1000000xcomplex<f32>>) -> tensor<1x1000000xcomplex<f32>>
    %1 = stablehlo.slice %arg2 [1:1000000] {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<1000000xcomplex<f32>>) -> tensor<999999xcomplex<f32>>
    %2 = stablehlo.reshape %1 {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<999999xcomplex<f32>>) -> tensor<1x999999xcomplex<f32>>
    %3 = stablehlo.convert %arg1 {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed GUARANTEED>]} : (tensor<999999xf32>) -> tensor<999999xcomplex<f32>>
    %4 = stablehlo.reshape %3 {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed GUARANTEED>]} : (tensor<999999xcomplex<f32>>) -> tensor<1x999999xcomplex<f32>>
    %5 = stablehlo.multiply %4, %2 {enzymexla.complex_is_purely_imaginary = [#enzymexla<guaranteed NOTGUARANTEED>], enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<1x999999xcomplex<f32>>
    %6 = stablehlo.iota dim = 0 : tensor<999999x1xi64>
    %7 = stablehlo.real %0 : (tensor<1x1000000xcomplex<f32>>) -> tensor<1x1000000xf32>
    %8 = stablehlo.imag %0 : (tensor<1x1000000xcomplex<f32>>) -> tensor<1x1000000xf32>
    %9 = stablehlo.concatenate %7, %8, dim = 0 : (tensor<1x1000000xf32>, tensor<1x1000000xf32>) -> tensor<2x1000000xf32>
    %10 = stablehlo.real %5 : (tensor<1x999999xcomplex<f32>>) -> tensor<1x999999xf32>
    %11 = stablehlo.imag %5 : (tensor<1x999999xcomplex<f32>>) -> tensor<1x999999xf32>
    %12 = stablehlo.concatenate %10, %11, dim = 0 : (tensor<1x999999xf32>, tensor<1x999999xf32>) -> tensor<2x999999xf32>
    %13 = "stablehlo.scatter"(%9, %6, %12) <{scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = true}> ({
    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
      stablehlo.return %arg4 : tensor<f32>
    }) : (tensor<2x1000000xf32>, tensor<999999x1xi64>, tensor<2x999999xf32>) -> tensor<2x1000000xf32>
    %14 = stablehlo.slice %13 [0:1, 0:1000000] : (tensor<2x1000000xf32>) -> tensor<1x1000000xf32>
    %15 = stablehlo.slice %13 [1:2, 0:1000000] : (tensor<2x1000000xf32>) -> tensor<1x1000000xf32>
    %16 = stablehlo.complex %14, %15 : tensor<1x1000000xcomplex<f32>>
    %17 = stablehlo.reshape %16 : (tensor<1x1000000xcomplex<f32>>) -> tensor<1000000xcomplex<f32>>
    return %17 : tensor<1000000xcomplex<f32>>
  }
}

[avik-pal]

We need a scatter to DUS pattern here (there exists a primitive version of that already, but clearly doesn't handle this case)

[albertomercurio]

Do you think it is something that can be "easily" fixed? I wanted to show some fancy stuff using Reactant at a workshop about QuantumToolbox.jl

[Pangoraw]

it should not be too hard to fix, but unlikely that someone fixes before april 15th since we are on a deadline