KernelAbstractions + Reactant: UndefVarError: pop not defined

[glwagner]

Trying to use Reactant with KernelAbstractions (on CPU for now). This code:

using KernelAbstractions
using Random
using Reactant
using CUDA

const backend = CPU() #CUDA.CUDABackend(; always_inline=true)

Reactant.allowscalar(true)

# Simple kernel for matrix multiplication
@kernel function matmul_kernel!(output, a, b)
    i, j = @index(Global, NTuple)

    # creating a temporary sum variable for matrix multiplication
    tmp_sum = zero(eltype(output))
    @trace for k in 1:size(a)[2]
        tmp_sum += @inbounds a[i, k] * b[k, j]
    end

    @inbounds output[i, j] = tmp_sum
end

# Creating a wrapper kernel for launching with error checks
function matmul!(output, a, b)
    if size(a)[2] != size(b)[1]
        println("Matrix size mismatch!")
        return nothing
    end
    #backend = KernelAbstractions.get_backend(a)
    kernel! = matmul_kernel!(backend)
    kernel!(output, a, b, ndrange = size(output))
end

a = rand!(allocate(backend, Float32, 256, 123))
b = rand!(allocate(backend, Float32, 123, 45))
output = KernelAbstractions.zeros(backend, Float32, 256, 45)

matmul!(output, a, b)
KernelAbstractions.synchronize(backend)

@show isapprox(output, a * b)

ra = Reactant.to_rarray(a)
rb = Reactant.to_rarray(b)
ro = Reactant.to_rarray(output)

if CUDA.functional()
    @jit matmul!(ro, ra, rb)
    @show all(Array(ro) .≈ output)
else
    @code_hlo optimize = :before_kernel matmul!(ro, ra, rb)
end

produces

julia> include("reactant_kernels.jl")
ERROR: LoadError: UndefVarError: `pop` not defined
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/ReactantCore/Lh7r5/src/ReactantCore.jl:185 [inlined]
  [2] macro expansion
    @ ~/Projects/Enzymantics/Enzymanigans.jl/reactant_kernels.jl:16 [inlined]
  [3] cpu_matmul_kernel!
    @ ~/.julia/packages/KernelAbstractions/0r40T/src/macros.jl:293 [inlined]
  [4] cpu_matmul_kernel!(__ctx__::KernelAbstractions.CompilerMetadata{…}, output::Matrix{…}, a::Matrix{…}, b::Matrix{…})
    @ Main ./none:0
  [5] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Tuple{…}, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/0r40T/src/cpu.jl:142
  [6] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Tuple{…}, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck, static_threads::Bool)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/0r40T/src/cpu.jl:109
  [7] (::KernelAbstractions.Kernel{…})(::Matrix{…}, ::Vararg{…}; ndrange::Tuple{…}, workgroupsize::Nothing)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/0r40T/src/cpu.jl:44
  [8] Kernel
    @ ~/.julia/packages/KernelAbstractions/0r40T/src/cpu.jl:37 [inlined]
  [9] matmul!(output::Matrix{Float32}, a::Matrix{Float32}, b::Matrix{Float32})
    @ Main ~/Projects/Enzymantics/Enzymanigans.jl/reactant_kernels.jl:31
 [10] top-level scope
    @ ~/Projects/Enzymantics/Enzymanigans.jl/reactant_kernels.jl:38
 [11] include(fname::String)
    @ Base.MainInclude ./client.jl:494
 [12] top-level scope
    @ REPL[1]:1
in expression starting at /Users/gregorywagner/Projects/Enzymantics/Enzymanigans.jl/reactant_kernels.jl:38
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

@Pangoraw so this issue popped up when using traced for (regular worked, but emitted ~3000000 instructions). Can you take a look?

[Pangoraw]

I'll take a look. Just to be clear, on the CPU KA.jl uses the Reactant tracing infrastructure?

[wsmoses]

Yup!

[Pangoraw]

This particular issue should be fixed by updating ExpressionExplorer.jl (a dependency of ReactantCore) to 1.1.1.