CPU kernel: array indexing with traced loop variable returns TracedRArray{0} instead of TracedRNumber

[KookiesNKareem]

When indexing an array with a traced loop variable inside a KernelAbstractions kernel on the CPU backend, the result is a TracedRArray{Float32,0} instead of TracedRNumber{Float32}. Arithmetic between these types is not defined. Julia v1.10, KernelAbstractions v0.9.39, Reactant v0.2.212 ~/.julia/dev/Reactant, ReactantCore v0.1.16, Reactant_jll v0.0.315+0.

MWE:

using Reactant
using Reactant: @compile
using KernelAbstractions
using KernelAbstractions: @kernel, @index, @Const
import CUDA

Reactant.set_default_backend("cpu")
Reactant.allowscalar(true)

@kernel function failing_kernel!(y, @Const(x), @Const(vals), N::Int32)
    i = @index(Global, Linear)
    s_idx = Int32(1)
    accum = x[i]
    while s_idx <= N
        v = vals[s_idx] # Returns TracedRArray{Float32,0}
        accum = accum + v  # MethodError: + not defined
        s_idx += Int32(1)
    end
    y[i] = accum
end

@kernel function working_kernel!(y, @Const(x), @Const(vals), N::Int32)
    i = @index(Global, Linear)
    s_idx = Int32(1)
    accum = x[i]
    while s_idx <= N
        v = vals[s_idx][] # [] extracts TracedRNumber from 0-dim array
        accum = accum + v
        s_idx += Int32(1)
    end
    y[i] = accum
end

function test_kernel(name, kernel_fn)
    N, Nt = 10, Int32(4)

    function loss(y, x, vals, N)
        backend = KernelAbstractions.CPU()
        k = kernel_fn(backend, 32)
        k(y, x, vals, N; ndrange=length(x))
        KernelAbstractions.synchronize(backend)
        return sum(y)
    end

    println("\n$name")
    y_ra = Reactant.to_rarray(zeros(Float32, N))
    x_ra = Reactant.to_rarray(ones(Float32, N))
    vals_ra = Reactant.to_rarray(ones(Float32, Nt) .* 0.1f0)
    try
        compiled = @compile raise=true raise_first=true loss(y_ra, x_ra, vals_ra, Nt)
        result = compiled(y_ra, x_ra, vals_ra, Nt)
        println("Success: $result")
    catch e
        showerror(stdout, e)
        println()
    end
end

test_kernel("Failing: vals[s_idx]", failing_kernel!)
test_kernel("Working: vals[s_idx][]", working_kernel!)