Custom Kernel outputs zeros?

[avik-pal]

using Reactant, CUDA

Reactant.set_default_backend("cpu")

function square_kernel!(x, y)
    i = threadIdx().x
    x[i] = y[i] * y[i]
    # We don't yet auto lower this via polygeist
    # sync_threads()
    return nothing
end

function square!(x, y)
    @cuda blocks = 1 threads = length(x) square_kernel!(x, y)
    return nothing
end

function square(x)
    y = similar(x)
    square!(y, x)
    return y
end

x = Reactant.to_rarray(collect(1:1:64) ./ 64)

@code_hlo square(x)

@jit square(x)

64-element ConcretePJRTArray{Float64, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 ⋮
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

[mofeing]

mmm just to discard synchronicity problems. does it change if you wait a couple of seconds before printing?

[avik-pal]

no it doesnt

[mofeing]

mmm what happens if you initialize y with ones instead of similar? i suspect that side-effects are not properly tracked and it might be optimizing away the kernel