GPU + ClimaCore compilation error (but only sometimes)

[trontrytel]

I have added a GPU + ClimaCore test in the CI. It works when the below code block is inside a function (i.e. the way it is currently written). However if instead you just try to execute

FT = Float32                                                                    
                                                                                
Ch2022 = CMP.Chen2022VelType(FT)                                                
liquid = CMP.CloudLiquid(FT)                                                    
ice = CMP.CloudIce(FT)                                                          
rain = CMP.Rain(FT)                                                             
snow = CMP.Snow(FT)                                                             
                                                                                
params = (; liquid, ice, Ch2022)                                                
                                                                                
space_3d_ρq = make_extruded_sphere(FT)                                          
space_3d_ρ = make_extruded_sphere(FT)                                           
space_3d_w = make_extruded_sphere(FT)                                           
                                                                                
ρq = CC.Fields.ones(space_3d_ρq) .* FT(1e-3)                                    
ρ = CC.Fields.ones(space_3d_ρ)                                                  
w = CC.Fields.zeros(space_3d_w)                                                 
                                                                                
Y = (; ρq, ρ)                                                                   
p = (; w, params)                                                               
                                                                                
t = 1                                                                           
                                                                                
set_sedimentation_precomputed_quantities(Y, p, t)   

it fails with

ERROR: LoadError: InvalidIRError: compiling MethodInstance for ClimaCoreCUDAExt.knl_copyto!(::ClimaCore.DataLayouts.VIJFH{…}, ::Base.Broadcast.Broadcasted{…}, ::ClimaCore.DataLayouts.UniversalSize{…}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to _broadcast_getindex_evalf(f::Tf, args::Vararg{Any, N}) where {Tf, N} @ Base.Broadcast broadcast.jl:709)

Here is the full test: https://github.com/CliMA/CloudMicrophysics.jl/blob/main/test/gpu_clima_core_test.jl

This looks like the same error I'm getting in ClimaAtmos when trying to add the cloud sedimentation CliMA/ClimaAtmos.jl#3442. Maybe it will be easier to debug here?

[trontrytel]

cc: @sriharshakandala @charleskawczynski

[trontrytel]

Does it mean I could try putting the whole Atmos driver into a function?....

[charleskawczynski]

Sorry I'm arriving a bit late to this.

The issue here is with this line:

    @. w = CMN.terminal_velocity(
        params.liquid,
        params.Ch2022.rain,
        Y.ρ,
        max(0, Y.ρq / Y.ρ),
    )

Julia's compiler is not able to infer through Base's broadcast recursive getindex.

This is because we create a tuple of tuple of objects (((params.liquid, ), (params.Ch2022.rain, ))), and the compiler struggles to infer unwrapping this object.

This limitation may be fixed in later versions of Julia (e.g., this PR), so I'd rather not suggest we change the way terminal_velocity is defined, that should work just fine.

The reason there is a difference between the posted code working and not working inside and outside a function (respectively) is likely due to inference being a heuristic-- other things can impact inference (and on the GPU that translates to the function working or not).

One simple (albeit a bit inconvenient) solution is to require users to define a wrapper function. We can change:

@. w = CMN.terminal_velocity(
        params.liquid,
        params.Ch2022.rain,
        Y.ρ,
        max(0, Y.ρq / Y.ρ),
    )

to

wrapper_terminal_velocity(p,ρ,x) = CMN.terminal_velocity(p.a,p.b,ρ,x)

struct TwoParams{A, B}
    a::A
    b::B
end
Base.Broadcast.broadcastable(x::TwoParams) = tuple(x)

p = TwoParams(params.liquid, params.Ch2022.rain)
@. w = wrapper_terminal_velocity(
        p,
        Y.ρ,
        max(0, Y.ρq / Y.ρ),
    )