Can't raise `_compute_source_term!` with Reactant

[dkytezab]

_compute_source_term! with FFTBasedPossionSolver throws an error whenever I try to raise it with Reactant. I think this is because rhs is Complex64 array but active * δ is Real64.

@kernel function _compute_source_term!(rhs, grid, Ũ)
    i, j, k = @index(Global, NTuple)
    active = !inactive_cell(i, j, k, grid)
    u, v, w = Ũ
    δ = divᶜᶜᶜ(i, j, k, grid, u, v, w)
    @inbounds rhs[i, j, k] = active * δ
end

Here is a MWE for this issue

using Oceananigans
using Oceananigans.Architectures: ReactantState
using Oceananigans.Models.NonhydrostaticModels: compute_source_term!
using Reactant
using CUDA

Reactant.set_default_backend("cpu")

grid = RectilinearGrid(ReactantState(); size=(4, 4), extent=(1, 1),
                       halo=(3, 3), topology=(Periodic, Periodic, Flat))

model  = NonhydrostaticModel(grid; timestepper=:QuasiAdamsBashforth2)
solver = model.pressure_solver
Ũ      = model.velocities
Δt     = 1.0

function go!(solver, Ũ)
    compute_source_term!(solver, nothing, Ũ, Δt)
    return nothing
end

Reactant.@compile raise=true sync=true go!(solver, Ũ)

loc(callsite(fused<#llvm.di_subprogram<id = distinct[0]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/LLVM/fEIbx/src/interop/base.jl" in ".">, name = "macro expansion;", linkageName = "macro expansion", file = <"/Users/danielkz/.julia/packages/LLVM/fEIbx/src/interop/base.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/LLVM/fEIbx/src/interop/base.jl":39:0] at callsite(fused<#llvm.di_subprogram<id = distinct[2]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"none" in ".">, name = "macro expansion;", linkageName = "macro expansion", file = <"none" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["none":0:0] at callsite(fused<#llvm.di_subprogram<id = distinct[3]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"none" in ".">, name = "pointerset;", linkageName = "pointerset", file = <"none" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["none":0:0] at callsite(fused<#llvm.di_subprogram<id = distinct[4]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/LLVM/fEIbx/src/interop/pointer.jl" in ".">, name = "unsafe_store!;", linkageName = "unsafe_store!", file = <"/Users/danielkz/.julia/packages/LLVM/fEIbx/src/interop/pointer.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/LLVM/fEIbx/src/interop/pointer.jl":88:0] at callsite(fused<#llvm.di_subprogram<id = distinct[5]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, name = "arrayset_bits;", linkageName = "arrayset_bits", file = <"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl":356:0] at callsite(fused<#llvm.di_subprogram<id = distinct[6]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, name = "#arrayset;", linkageName = "#arrayset", file = <"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl":349:0] at callsite(fused<#llvm.di_subprogram<id = distinct[7]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, name = "setindex!;", linkageName = "setindex!", file = <"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl":391:0] at callsite(fused<#llvm.di_subprogram<id = distinct[7]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, name = "setindex!;", linkageName = "setindex!", file = <"/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/Reactant/vZMtU/ext/ReactantCUDAExt.jl":403:0] at callsite(fused<#llvm.di_subprogram<id = distinct[8]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/Aeolus2/Oceananigans.jl/src/Models/NonhydrostaticModels/solve_for_pressure.jl" in ".">, name = "macro expansion;", linkageName = "macro expansion", file = <"/Users/danielkz/Aeolus2/Oceananigans.jl/src/Models/NonhydrostaticModels/solve_for_pressure.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/Aeolus2/Oceananigans.jl/src/Models/NonhydrostaticModels/solve_for_pressure.jl":17:0] at callsite(fused<#llvm.di_subprogram<id = distinct[9]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, scope = #llvm.di_file<"/Users/danielkz/.julia/packages/KernelAbstractions/X5fk1/src/macros.jl" in ".">, name = "gpu__compute_source_term!;", linkageName = "gpu__compute_source_term!", file = <"/Users/danielkz/.julia/packages/KernelAbstractions/X5fk1/src/macros.jl" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["/Users/danielkz/.julia/packages/KernelAbstractions/X5fk1/src/macros.jl":332:0] at fused<#llvm.di_subprogram<id = distinct[10]<>, compileUnit = <id = distinct[1]<>, sourceLanguage = DW_LANG_Julia, file = <"julia" in ".">, producer = "julia", isOptimized = true, emissionKind = None, nameTableKind = None>, name = "gpu__compute_source_term!", linkageName = "julia_gpu__compute_source_term!_78855", file = <"none" in ".">, subprogramFlags = "Definition|Optimized", type = <>>>["none":0:0]))))))))))): error: 'affine.store' op value to store must have the same type as memref element type

[glwagner]

Perhaps you just need

@inbounds rhs[i, j, k] = Complex{T}(active * δ)

I suggest keeping the eltype agnostic just in case --- eg some devices only support f32.

PS @dkytezab if you paste a link to the source code, it will helpfully display inline, which is better than copy-pasting the same (because we know where the code is, we could even click the link and edit in github editor for a fast 30 s PR). for example:

Oceananigans.jl/src/Models/NonhydrostaticModels/solve_for_pressure.jl

Lines 12 to 18 in ae70538

	@kernel function _compute_source_term!(rhs, grid, Ũ)
	i, j, k = @index(Global, NTuple)
	active = !inactive_cell(i, j, k, grid)
	u, v, w = Ũ
	δ = divᶜᶜᶜ(i, j, k, grid, u, v, w)
	@inbounds rhs[i, j, k] = active * δ
	end

[dkytezab]

Gotcha, will do going forward. I think this still breaks - seems like this kernel try and call this function jl_f_apply_type that is unsupported on GPU-compiled kernels.

There is probably an underlying Reactant bug here generating invalid MLIR. Seems like Reactant can handle Real32 <-> Real64 inside of a kernel but not Real <-> Complex.

using Reactant, KernelAbstractions

Reactant.set_default_backend("cpu")

A = Reactant.to_rarray(zeros(ComplexF64, 4))
B = Reactant.to_rarray(zeros(Float64, 4))

@kernel function write!(A, B)
    i = @index(Global)
    @inbounds A[i] = B[i]  
end

f_kernel!(A, B) = (write!(KernelAbstractions.get_backend(A))(A, B; ndrange=length(A)); nothing)
Reactant.@compile raise=true sync=true f_kernel!(A, B)

[glwagner]

I think this still breaks - seems like this kernel try and call this function jl_f_apply_type that is unsupported on GPU-compiled kernels.

Sometimes this kind of thing can be a symptom of a type inference failure. That would be surprising here though, if it already runs on GPU.

[dkytezab]

I wonder if we just need better support for type promotion / demotion inside of these kernels? See the reverse (complex -> real) problem here: EnzymeAD/Reactant.jl#2200

[glwagner]

Does your MWE work if you write

@inbounds A[i] = ComplexF64(B[i])

?

I would add that element to the MWE just to be sure. Then it seems like a pretty cut and dry issue to raise!

[dkytezab]

Nah still breaks for the same reason :(

[glwagner]

nice find!

[glwagner]

x-ref to this issue when you raise it in Reactant!

[dkytezab]

EnzymeAD/Reactant.jl#2406 !!

[dkytezab]

seems like i created a fully connected graph of issues haha

[wsmoses]

The issue is that complex arguments to kernels are currently not supported

[glwagner]

@dkytezab I think if we can store active * δ inside a scratch field or array which is real-valued, we can then use

rhs .= interior(scratch)

this could be tested by using a temp array / field for scratch.

the gist being that complex-valued arrays are not supported in KA kernels, but they might work with broadcasting.

[dkytezab]

I did find actually that broadcasting generally works towards this end