Oceananigans.jl: Complex AbstractOperations cannot be computed on GPU

[tomchor]

I have the following output writer set-up in my simulation:

import Oceananigans.Fields: ComputedField
u, v, w = model.velocities.u, model.velocities.v, model.velocities.w
b, pHY′, pNHS = model.tracers.b, model.pressures.pHY′, model.pressures.pNHS
ν, νₑ = model.closure.ν, model.diffusivities.νₑ
#-----

#-----
import Oceananigans.AbstractOperations: ∂x, ∂y, ∂z

ddx = ∂x(u)^2 + ∂x(v)^2 + ∂x(w)^2
ddy = ∂y(u)^2 + ∂y(v)^2 + ∂y(w)^2
ddz = ∂z(u)^2 + ∂z(v)^2 + ∂z(w)^2
#avg_ε = AveragedField(ComputedField(ddx + ddy + ddz), dims=(1,))
avg_ε = AveragedField(ν*(ddx + ddy + ddz), dims=(1,))

outputs = (ε = avg_ε,)
simulation.output_writers[:avg_field_writer] =
    NetCDFOutputWriter(model, outputs,
                       filepath = "avg.jd15_3dbounded.nc",
                       schedule = TimeInterval(2minutes),
                       mode = "c")

This works successfully on CPUs, but running on GPUs I get a huge amount of error lines with some

 [16] compute! at /glade/u/home/tomasc/.julia/packages/Oceananigans/6JcUu/src/Fields/averaged_field.jl:86 [inlined]

Running the simulation without that output works for both GPUs and CPUs.

Am I doing something wrong here?

[glwagner]

Ah, you may be running into some hidden limitations of AbstractOperations --- they seem to fail when the operations are too complex.

It's a shifty problem, because it appears to depend on the julia compiler. It's good that you opened this issue because it would be nice to document our efforts for solving this tricky problem.

I don't think this is a problem of AveragedFields specifically. Rather it's an issue with ComputedFields, are more specifically, evaluating complex AbstractOperations on the GPU.

Over at LESbrary, we are circumventing this issue by hand-writing particularly important complicated kernels. I think ViscousDissipation may be the object you're looking for:

https://github.com/CliMA/LESbrary.jl/blob/master/src/TurbulenceStatistics/viscous_dissipation.jl

[tomchor]

Thanks! Yeah, that's what I was trying to do. I see you already have some SGS stuff there as well.

Do you have an example of how to use LESbrary? I'm not that familiar with Julia and just from reading the README I can' quite understand how it's supposed to work.

Thanks!

[glwagner]

There's some LESbrary examples but the package definitely needs some love (wink wink):

https://github.com/CliMA/LESbrary.jl/blob/6cfe60cf44031f9357936058f08eed0f1bb5f9fc/examples/langmuir_turbulence.jl#L240

To get LESbrary.jl into your environment you use the package manager, eg:

using Pkg
Pkg.add("https://github.com/CliMA/LESbrary.jl.git")

using LESbrary.TurbulenceStatistics: ViscousDissipation

[glwagner]

One strategy that might help with complex AbstractOperations is to define intermediate ComputedFields, and then reuse these ComputedFields in higher-order AbstractOperations. This incurs extra computation but might be a stopgap if a low-level kernel like the one that's been written for ViscousDissipation is not available.

[glwagner]

This is probably related to some of the struggles in #870

[tomchor]

One strategy that might help with complex AbstractOperations is to define intermediate ComputedFields, and then reuse these ComputedFields in higher-order AbstractOperations. This incurs extra computation but might be a stopgap if a low-level kernel like the one that's been written for ViscousDissipation is not available.

I've tried a version of that already and it didn't work, but there are probably more things I could try. For now it's best not to reinvent the wheel and use LESbrary! But I'll keep this in mind when doing stuff like this in the future.

[glwagner]

Ultimately @tomchor we would definitely rather use AbstractOperations to diagnose ViscousDissipation rather than hard-coding it, as we had to resort to in LESbrary.jl. But this is tough issue to solve I think so yes, I think LESbrary.jl stuff is good for now. It'd be nice also to add other diagnostics that cannot be specified by AbstractOperations to LESbrary.jl. And we need documentation and tests, etc, etc...

[glwagner]

Another issue is that I believe one cannot embed AveragedFields into AbstractOperations on the GPU, despite that we have developed the abstraction and machinery so that everything does work as expected on the CPU. Embedding AveragedFields in AbstractOperations is needed to calculate things like turbulent kinetic energy and so is important for LES. Currently there's a custom field TurbulentKineticEnergy in the LESbrary.jl for this purpose:

https://github.com/CliMA/LESbrary.jl/blob/master/src/TurbulenceStatistics/turbulent_kinetic_energy.jl

It might be good to come up with a list of AbstractOperations that we would like to work on the GPU eventually, and then write tests for them and add them in a PR. We can then @test_skip them until either we figure out how to solve the issue or, perhaps, the julia compiler changes in a way that makes the tests pass... :-D

[francispoulin]

I have not used GPUs and don't appreciate the difficulty here at all but would be happy to discuss this sometime if people wanted to have a brainstorming session. Certainly starting simple is what I would recommend.

[ali-ramadhan]

Yeah I think it's just an unknown GPU compilation problem/failure. Unclear to me whether it's the fault of Oceananigans.AbstractOperations or something in CUDA.jl (or even KernelAbstractions.jl).

Certainly starting simple is what I would recommend.

Maybe it makes sense to try and condense some abstract operations into an isolated minimal working example which we can use to open an issue on CUDA.jl if that turns out to be the problem? Might help isolate the problem, and would certainly be much easier to debug.

I recall we encounter a limitation of the GPU compiler when trying to construct a GPU model with too many arbitrary tracers. I think in this case the type information was too large to even fit into the argument of a CUDA kernel or something so maybe this was a hard GPU limitation? Would be unfortunate if something similar is happening for complex abstract operations.

[glwagner]

Here's a starting list:

• ViscousDissipation, which @tomchor attempts to compute above, also being computed here (and note that there's an open issue about the expression).

• TurbulentKineticEnergy, which requires embedding AveragedFields in an operation

• ShearProduction, similar to above.

@ali-ramadhan not sure if this is what you mean but when trying to adapt Field to the GPU we did encounter an error that was something like "ptx arguments consume too much parameter space". I believe this is an issue compiling very large objects (Field is large because it has boundary conditions). I don't think that this is the error we get for too-complex AbstractOperations though. I think for abstract operations its a type-inference issue. We should dump the whole error into this issue so that we know. I also think the error might change as the julia compiler evolves, or when we upgrade CUDA.jl.

[glwagner]

Here's the specific error we got when we tried to get Field, including all its glorious boundary conditions, to compile on the GPU:

Entry function 'ptxcall_calculate_Gu__66' uses too much parameter space (0x16c8 bytes, 0x1100 max)

dredged up from #746 . Some workarounds were suggested there, but I think our solution is actually better / simpler (adapt fields by unwrapping the underlying data and throwing away boundary conditions, rather than wrestling to get all the field info onto the poor GPU).

[francispoulin]

Does someone have a minimal example that reproduces this error? I'm just curious to learn more about the problem.

[glwagner]

using Oceananigans
using Oceananigans.AbstractOperations
using Oceananigans.Fields

grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1)) 

model = IncompressibleModel(architecture=GPU(), grid=grid)

u, v, w = model.velocities

Σˣˣ = ∂x(u)
Σʸʸ = ∂y(v)
Σᶻᶻ = ∂z(w)

Σˣʸ = (∂y(u) + ∂x(v)) / 2 
Σˣᶻ = (∂z(u) + ∂x(w)) / 2 
Σʸᶻ = (∂z(v) + ∂y(w)) / 2 

ϵ = model.closure.ν * 2 * (Σˣˣ^2 + Σʸʸ^2 + Σᶻᶻ^2 + 2 * (Σˣʸ^2 + Σˣᶻ^2 + Σʸᶻ^2))

ϵ_field = ComputedField(ϵ)

compute!(ϵ_field)

produces

julia> include("complex_output.jl")
[ Info: Precompiling Oceananigans [9e8cae18-63c1-5223-a75c-80ca9d6e9a09]
[ Info: Oceananigans will use 24 threads
ERROR: LoadError: InvalidIRError: compiling kernel gpu__compute!(Cassette.Context{nametype(CUDACtx),KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},KernelAbstractions.NDIteration.DynamicCheck,Nothing,Nothing,KernelAbstractions.NDIteration.NDRange{3,KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},Nothing,Nothing}},Nothing,KernelAbstractions.var"##PassType#253",Nothing,Cassette.DisableHooks}, typeof(Oceananigans.Fields.gpu__compute!), OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}}, Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(*),Float64,Oceananigans.AbstractOperations.MultiaryOperation{Cell,Cell,Cell,4,typeof(+),Tuple{Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂xᶜᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂yᵃᶜᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂zᵃᵃᶜ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(*),Int64,Oceananigans.AbstractOperations.MultiaryOperation{Face,Face,Cell,3,typeof(+),Tuple{Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(+),Oceananigans.AbstractOperations.Derivative{Face,Face,Cell,typeof(Oceananigans.Operators.∂yᵃᶠᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Face,Face,Cell,typeof(Oceananigans.Operators.∂xᶠᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(+),Oceananigans.AbstractOperations.Derivative{Face,Cell,Face,typeof(Oceananigans.Operators.∂zᵃᵃᶠ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Face,Cell,Face,typeof(Oceananigans.Operators.∂xᶠᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(+),Oceananigans.AbstractOperations.Derivative{Cell,Face,Face,typeof(Oceananigans.Operators.∂zᵃᵃᶠ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Cell,Face,Face,typeof(Oceananigans.Operators.∂yᵃᶠᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}},Tuple{typeof(identity),typeof(Oceananigans.Operators.ℑyzᵃᶠᶜ),typeof(Oceananigans.Operators.ℑxzᶠᵃᶜ)},RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}},Tuple{typeof(identity),typeof(identity),typeof(identity),typeof(Oceananigans.Operators.ℑxyᶜᶜᵃ)},RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to overdub(overdub_context::Cassette.Context, overdub_arguments...) in Cassette at /home/glwagner/.julia/packages/Cassette/158rp/src/overdub.jl:586)
Stacktrace:
 [1] getindex at /home/glwagner/.julia/packages/Oceananigans/cLFd3/src/AbstractOperations/binary_operations.jl:34
 [2] macro expansion at /home/glwagner/.julia/packages/Oceananigans/cLFd3/src/Fields/computed_field.jl:86
 [3] gpu__compute! at /home/glwagner/.julia/packages/KernelAbstractions/jAutM/src/macros.jl:80
 [4] overdub at /home/glwagner/.julia/packages/Cassette/158rp/src/overdub.jl:0
Reason: unsupported dynamic function invocation (call to overdub)
Stacktrace:
 [1] macro expansion at /home/glwagner/.julia/packages/Oceananigans/cLFd3/src/Fields/computed_field.jl:86
 [2] gpu__compute! at /home/glwagner/.julia/packages/KernelAbstractions/jAutM/src/macros.jl:80
 [3] overdub at /home/glwagner/.julia/packages/Cassette/158rp/src/overdub.jl:0
Stacktrace:
 [1] check_ir(::GPUCompiler.CompilerJob{GPUCompiler.PTXCompilerTarget,CUDA.CUDACompilerParams}, ::LLVM.Module) at /home/glwagner/.julia/packages/GPUCompiler/uTpNx/src/validation.jl:123
 [2] macro expansion at /home/glwagner/.julia/packages/GPUCompiler/uTpNx/src/driver.jl:239 [inlined]
 [3] macro expansion at /home/glwagner/.julia/packages/TimerOutputs/ZmKD7/src/TimerOutput.jl:206 [inlined]
 [4] codegen(::Symbol, ::GPUCompiler.CompilerJob; libraries::Bool, deferred_codegen::Bool, optimize::Bool, strip::Bool, validate::Bool, only_entry::Bool) at /home/glwagner/.julia/packages/GPUCompiler/uTpNx/src/driver.jl:237
 [5] compile(::Symbol, ::GPUCompiler.CompilerJob; libraries::Bool, deferred_codegen::Bool, optimize::Bool, strip::Bool, validate::Bool, only_entry::Bool) at /home/glwagner/.julia/packages/GPUCompiler/uTpNx/src/driver.jl:39
 [6] compile at /home/glwagner/.julia/packages/GPUCompiler/uTpNx/src/driver.jl:35 [inlined]
 [7] cufunction_compile(::GPUCompiler.FunctionSpec; kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxthreads,),Tuple{Int64}}}) at /home/glwagner/.julia/packages/CUDA/YeS8q/src/compiler/execution.jl:310
 [8] check_cache(::Dict{UInt64,Any}, ::Any, ::Any, ::GPUCompiler.FunctionSpec{typeof(Cassette.overdub),Tuple{Cassette.Context{nametype(CUDACtx),KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},KernelAbstractions.NDIteration.DynamicCheck,Nothing,Nothing,KernelAbstractions.NDIteration.NDRange{3,KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},Nothing,Nothing}},Nothing,KernelAbstractions.var"##PassType#253",Nothing,Cassette.DisableHooks},typeof(Oceananigans.Fields.gpu__compute!),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(*),Float64,Oceananigans.AbstractOperations.MultiaryOperation{Cell,Cell,Cell,4,typeof(+),Tuple{Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂xᶜᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂yᵃᶜᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂zᵃᵃᶜ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(*),Int64,Oceananigans.AbstractOperations.MultiaryOperation{Face,Face,Cell,3,typeof(+),Tuple{Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(+),Oceananigans.AbstractOperations.Derivative{Face,Face,Cell,typeof(Oceananigans.Operators.∂yᵃᶠᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Face,Face,Cell,typeof(Oceananigans.Operators.∂xᶠᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(+),Oceananigans.AbstractOperations.Derivative{Face,Cell,Face,typeof(Oceananigans.Operators.∂zᵃᵃᶠ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Face,Cell,Face,typeof(Oceananigans.Operators.∂xᶠᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(+),Oceananigans.AbstractOperations.Derivative{Cell,Face,Face,typeof(Oceananigans.Operators.∂zᵃᵃᶠ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Cell,Face,Face,typeof(Oceananigans.Operators.∂yᵃᶠᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}},Tuple{typeof(identity),typeof(Oceananigans.Operators.ℑyzᵃᶠᶜ),typeof(Oceananigans.Operators.ℑxzᶠᵃᶜ)},RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}},Tuple{typeof(identity),typeof(identity),typeof(identity),typeof(Oceananigans.Operators.ℑxyᶜᶜᵃ)},RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}}}, ::UInt64; kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxthreads,),Tuple{Int64}}}) at /home/glwagner/.julia/packages/GPUCompiler/uTpNx/src/cache.jl:40
 [9] gpu__compute! at ./array.jl:0 [inlined]
 [10] cufunction(::typeof(Cassette.overdub), ::Type{Tuple{Cassette.Context{nametype(CUDACtx),KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},KernelAbstractions.NDIteration.DynamicCheck,Nothing,Nothing,KernelAbstractions.NDIteration.NDRange{3,KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},Nothing,Nothing}},Nothing,KernelAbstractions.var"##PassType#253",Nothing,Cassette.DisableHooks},typeof(Oceananigans.Fields.gpu__compute!),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(*),Float64,Oceananigans.AbstractOperations.MultiaryOperation{Cell,Cell,Cell,4,typeof(+),Tuple{Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂xᶜᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂yᵃᶜᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Cell,Cell,typeof(^),Oceananigans.AbstractOperations.Derivative{Cell,Cell,Cell,typeof(Oceananigans.Operators.∂zᵃᵃᶜ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(*),Int64,Oceananigans.AbstractOperations.MultiaryOperation{Face,Face,Cell,3,typeof(+),Tuple{Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Face,Face,Cell,typeof(+),Oceananigans.AbstractOperations.Derivative{Face,Face,Cell,typeof(Oceananigans.Operators.∂yᵃᶠᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Face,Face,Cell,typeof(Oceananigans.Operators.∂xᶠᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Face,Cell,Face,typeof(+),Oceananigans.AbstractOperations.Derivative{Face,Cell,Face,typeof(Oceananigans.Operators.∂zᵃᵃᶠ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Face,Cell,Face,typeof(Oceananigans.Operators.∂xᶠᵃᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(^),Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(/),Oceananigans.AbstractOperations.BinaryOperation{Cell,Face,Face,typeof(+),Oceananigans.AbstractOperations.Derivative{Cell,Face,Face,typeof(Oceananigans.Operators.∂zᵃᵃᶠ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Oceananigans.AbstractOperations.Derivative{Cell,Face,Face,typeof(Oceananigans.Operators.∂yᵃᶠᵃ),OffsetArrays.OffsetArray{Float64,3,CUDA.CuDeviceArray{Float64,3,1}},typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},Int64,typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}},Tuple{typeof(identity),typeof(Oceananigans.Operators.ℑyzᵃᶠᶜ),typeof(Oceananigans.Operators.ℑxzᶠᵃᶜ)},RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}},Tuple{typeof(identity),typeof(identity),typeof(identity),typeof(Oceananigans.Operators.ℑxyᶜᶜᵃ)},RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}},typeof(identity),typeof(identity),typeof(identity),RegularCartesianGrid{Float64,Periodic,Periodic,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}}}; name::String, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxthreads,),Tuple{Int64}}}) at /home/glwagner/.julia/packages/CUDA/YeS8q/src/compiler/execution.jl:297
 [11] macro expansion at /home/glwagner/.julia/packages/CUDA/YeS8q/src/compiler/execution.jl:109 [inlined]
 [12] (::KernelAbstractions.Kernel{KernelAbstractions.CUDADevice,KernelAbstractions.NDIteration.StaticSize{(1, 1)},KernelAbstractions.NDIteration.StaticSize{(1, 1, 1)},typeof(Oceananigans.Fields.gpu__compute!)})(::OffsetArrays.OffsetArray{Float64,3,CUDA.CuArray{Float64,3}}, ::Vararg{Any,N} where N; ndrange::Nothing, dependencies::KernelAbstractions.CudaEvent, workgroupsize::Nothing, progress::Function) at /home/glwagner/.julia/packages/KernelAbstractions/jAutM/src/backends/cuda.jl:185
 [13] compute!(::ComputedField{Cell,Cell,Cell

[ali-ramadhan]

Hmmm, looks like a Cassette issue. Maybe it's just being overly sensitive to the contents of the kernel (or kernel arguments) like was the case with #828?

[glwagner]

A "dynamic function invocation" means that the compiler thinks a function is being called whose scope can change "dynamically" (I think). This is the error one gets when a function depends on a global variable that is not const (for example). In this case, the error tells us that the types of the objects involved in calling getindex defined on BinaryOperation:

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/AbstractOperations/binary_operations.jl#L34

are not correctly inferred.

The way getindex comes into play is in the kernel function _compute! that evaluates the AbstractOperation:

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/Fields/computed_field.jl#L84-L87

calling getindex(operand, i, j, k) (or equivalently operand[i, j, k]) triggers recursive getindex calls that traverse the AbstractOperation tree. It seems that when the tree is too large, this traversal cannot be entirely compiled.

Perhaps there are tricks we might use to help the compiler parse this kind of operation, like putting some type annotations / hints into getindex(b::BinaryOperation, ...). Not sure. Another possibility is to figure out how to simplify the object BinaryOperation, MultiaryOperation, so that the compiler is less stressed trying to compile them... ?

[ali-ramadhan]

Hmmm, if it's indeed incapable of compiling the entire operation might be helpful to be able to find out exactly at which size it fails.

I'm still thinking that it might be something to be fixed/improved in CUDA.jl.

I can try to do some test as I update #870.

[francispoulin]

I tried the code above and do get the same error, so I can confirm that.

I tried to trim it down and found that the following, slightly more minimal example, produces the same error. It seems that squaring and multiplying togther is too much. But if you remove one or the other, it seems to work fine.

using Oceananigans
using Oceananigans.AbstractOperations
using Oceananigans.Fields

grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1)) 

model = IncompressibleModel(architecture=GPU(), grid=grid)

u, v, w = model.velocities

ϵ = 2 * (∂x(u)^2)

ϵ_field = ComputedField(ϵ)

compute!(ϵ_field)

[glwagner]

Hmm that's good to know the limit of complexity. I can confirm that behavior, eg:

julia> compute!(ComputedField(2 * Σˣˣ))

julia> compute!(ComputedField(Σˣˣ^2))

julia> compute!(ComputedField(2 * Σˣˣ^2))
ERROR: InvalidIRError: compiling kernel gpu__compute!...

Here's the tree for 2 * Σˣˣ^2:

julia> 2 * Σˣˣ^2
BinaryOperation at (Cell, Cell, Cell)
├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Cell, Cell, Cell) via identity
    ├── 2
    └── ^ at (Cell, Cell, Cell) via identity
        ├── ∂xᶜᵃᵃ at (Cell, Cell, Cell) via identity
        │   └── Field located at (Face, Cell, Cell)
        └── 2

[glwagner]

Also

julia> compute!(ComputedField(Σˣˣ^2 + Σʸʸ^2))
ERROR: InvalidIRError: compiling kernel gpu__compute!...

[glwagner]

It's not _strictly_ a nesting issue, since:

julia> compute!(ComputedField(∂x(∂x(u))))

julia> compute!(ComputedField(∂x(∂x(∂x(u)))))

julia> compute!(ComputedField(∂x(∂x(∂x(∂x(u))))))

is fine.

[glwagner]

Whoa. Here's a hint.

julia> compute!(ComputedField(u^2 + v^2 + w^2))

julia> compute!(ComputedField(u^2 + v^2))
ERROR: InvalidIRError: compiling kernel gpu__compute!

!!

Note:

julia> u^2 + v^2
BinaryOperation at (Face, Cell, Cell)
├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    + at (Face, Cell, Cell) via identity
    ├── ^ at (Face, Cell, Cell) via identity
    │   ├── Field located at (Face, Cell, Cell)
    │   └── 2
    └── ^ at (Cell, Face, Cell) via identity
        ├── Field located at (Cell, Face, Cell)
        └── 2

julia> u^2 + v^2 + w^2
MultiaryOperation at (Face, Cell, Cell)
├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    + at (Face, Cell, Cell)
    ├── ^ at (Face, Cell, Cell) via identity
    │   ├── Field located at (Face, Cell, Cell)
    │   └── 2
    ├── ^ at (Cell, Face, Cell) via identity
    │   ├── Field located at (Cell, Face, Cell)
    │   └── 2
    └── ^ at (Cell, Cell, Face) via identity
        ├── Field located at (Cell, Cell, Face)
        └── 2

Surprisingly, MultiaryOperations are better behaved.

[ali-ramadhan]

I can confirm that I was able to get away with √(u^2 + v^2 + w^2) on GPUs in one of the particle tracking tests (PR #1091).

[glwagner]

While getindex for a MultiaryOperation is complicated (unrolls a loop!):

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/AbstractOperations/multiary_operations.jl#L14-L15

the MultiaryOperation object is _simpler_ than a BinaryOperation because it stores _fewer_ interpolation functions (we only support the case that every member of a multiary operation is _first_ interpolated to a common location before being op'd on):

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/AbstractOperations/multiary_operations.jl#L3-L12

compare with

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/AbstractOperations/binary_operations.jl#L8-L32

[glwagner]

The boundary might be two nested BinaryOperations. @francispoulin gave one example; here's another:

julia> compute!(ComputedField(u + v))

julia> compute!(ComputedField(u + v - w))
ERROR: InvalidIRError: compiling kernel gpu__compute!

u + v - w nests the operation u + v inside (u + v) - w:

julia> u + v - w
BinaryOperation at (Face, Cell, Cell)
├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    - at (Face, Cell, Cell) via identity
    ├── + at (Face, Cell, Cell) via identity
    │   ├── Field located at (Face, Cell, Cell)
    │   └── Field located at (Cell, Face, Cell)
    └── Field located at (Cell, Cell, Face)

and for good measure and also good hilarity:

julia> u + v + (-1*w)
MultiaryOperation at (Face, Cell, Cell)
├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    + at (Face, Cell, Cell)
    ├── Field located at (Face, Cell, Cell)
    ├── Field located at (Cell, Face, Cell)
    └── * at (Cell, Cell, Face) via identity
        ├── -1
        └── Field located at (Cell, Cell, Face)

julia> compute!(ComputedField(u + v + (-1*w)))

is perfectly good.

[francispoulin]

I'm trying to wrap my head around this (and figure out interpolation in ShallowWater) and am getting bits and pieces. It is surprising to me, a novice, that MultiaryOperations are better behaved, as I would expect that they would be more complicated. But good news that they are. If that's the case, should we rethink how BinaryOperation is put together, and maybe redo it or just use Multiary?

Sorry if this is naive but wanted to through it out there.

[glwagner]

@francispoulin you shouldn't have to worry about this when developing ShallowWaterModel. This purely concerns AbstractOperations, output, and diagnostics, which doesn't affect the time-stepping / physics kernels. You still care about resolving this of course, because you're using Fields and therefore can use AbstractOperations for specifying output...

I think you're right that right now you can hack your way to success by avoiding nested Binary's (when possible). The above example is one case...

[francispoulin]

I don't have to but I am interested, hence trying to follow.

And I believe I almost have 2D ShallowWater compiling, then to see whether it spits out garbage. But that is for another issue ;)

[glwagner]

Wow this is an amazing black hole of errors. So now that we know Multiary is better, we can actually compute a bit more of the viscous dissipation than we thought, eg:

julia> compute!(ComputedField(∂x(u)^2 + ∂y(v)^2 + ∂z(w)^2))

But there be dragons...

julia> compute!(ComputedField(∂x(u)^2 + ∂y(v)^2 + ∂z(w)^2 + ∂x(w)^2))
ERROR: CUDA error: device kernel image is invalid (code 200, ERROR_INVALID_IMAGE)

amazingly, the error continues to morph

julia> compute!(ComputedField(∂x(u)^2 + ∂y(v)^2 + ∂z(w)^2 + ∂x(w)^2 + ∂y(w)^2))
ERROR: CUDA error: a PTX JIT compilation failed (code 218, ERROR_INVALID_PTX)
ptxas application ptx input, line 803; error   : Entry function '_Z19julia_gpu__compute_7ContextI14__CUDACtx_Name16CompilerMetadataI10StaticSizeI9_1__1__1_E12DynamicCheckvv7NDRangeILi3ES2_I9_1__1__1_ES2_I9_1__1__1_EvvEEv14__PassType_253v12DisableHooksE14_gpu__compute_11OffsetArrayI7Float64Li3E13CuDeviceArrayIS9_Li3ELi1EEE17MultiaryOperationI4CellS12_S12_Li5E2__5TupleI15BinaryOperationIS12_S12_S12_S13_10DerivativeIS12_S12_S12_6__x___S8_IS9_Li3ES10_IS9_Li3ELi1EEE9_identity20RegularCartesianGridIS9_8PeriodicS20_7BoundedS8_IS9_Li1E12StepRangeLenIS9_14TwicePrecisionIS9_ES23_IS9_EEEEE5Int64S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_IS12_S12_S12_S13_S16_IS12_S12_S12_6__y___S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_IS12_S12_S12_S13_S16_IS12_S12_S12_6__z___S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_I4FaceS12_S27_S13_S16_IS27_S12_S27_S17_S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_IS12_S27_S27_S13_S16_IS12_S27_S27_S25_S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEEES14_IS18_S18_S18_7__xz___7__yz___ES19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEE' uses too much parameter space (0x1408 bytes, 0x1100 max).
ptxas fatal   : Ptx assembly aborted due to errors

[glwagner]

@vchuravy we need you. And yes of course we need to use Cthulhu.

[glwagner]

In summary, we can encounter at least 3 different errors:

1. Dynamic function invocation

julia> compute!(ComputedField(u + v - w))
ERROR: LoadError: InvalidIRError: compiling kernel gpu__compute!(Cassette.Context{nametype(CUDACtx),...
Reason: unsupported dynamic function invocation...

Solution: probably the compiler isn't inferring types correctly. I think we can use Cthulhu to fully diagnose this error (though we may still need to be creative to solve the problem).

2. device kernel image is invalid?!?!

julia> compute!(ComputedField(∂x(u)^2 + ∂y(v)^2 + ∂z(w)^2 + ∂x(w)^2))
ERROR: CUDA error: device kernel image is invalid (code 200, ERROR_INVALID_IMAGE)

Solution: ??????

3. "Entry function uses too much parameter space"

julia> compute!(ComputedField(∂x(u)^2 + ∂y(v)^2 + ∂z(w)^2 + ∂x(w)^2 + ∂y(w)^2))
ERROR: CUDA error: a PTX JIT compilation failed (code 218, ERROR_INVALID_PTX)
ptxas application ptx input, line 803; error   : Entry function '_Z19julia_gpu__compute_7ContextI14__CUDACtx_Name16CompilerMetadataI10StaticSizeI9_1__1__1_E12DynamicCheckvv7NDRangeILi3ES2_I9_1__1__1_ES2_I9_1__1__1_EvvEEv14__PassType_253v12DisableHooksE14_gpu__compute_11OffsetArrayI7Float64Li3E13CuDeviceArrayIS9_Li3ELi1EEE17MultiaryOperationI4CellS12_S12_Li5E2__5TupleI15BinaryOperationIS12_S12_S12_S13_10DerivativeIS12_S12_S12_6__x___S8_IS9_Li3ES10_IS9_Li3ELi1EEE9_identity20RegularCartesianGridIS9_8PeriodicS20_7BoundedS8_IS9_Li1E12StepRangeLenIS9_14TwicePrecisionIS9_ES23_IS9_EEEEE5Int64S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_IS12_S12_S12_S13_S16_IS12_S12_S12_6__y___S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_IS12_S12_S12_S13_S16_IS12_S12_S12_6__z___S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_I4FaceS12_S27_S13_S16_IS27_S12_S27_S17_S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES15_IS12_S27_S27_S13_S16_IS12_S27_S27_S25_S8_IS9_Li3ES10_IS9_Li3ELi1EEES18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEES24_S18_S18_S18_S19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEEES14_IS18_S18_S18_7__xz___7__yz___ES19_IS9_S20_S20_S21_S8_IS9_Li1ES22_IS9_S23_IS9_ES23_IS9_EEEEE' uses too much parameter space (0x1408 bytes, 0x1100 max).
ptxas fatal   : Ptx assembly aborted due to errors

Solution: probably some of the suggestions on #746 are useful; we might need to submit a PR to CUDA.jl for this. I think this might be the toughest of all (more argument for also resolving #1246)

[ali-ramadhan]

Haha that last error might actually be worth framing on a wall!

RIP PTX assembler

[glwagner]

I always suspected that entry function _Z19julia_gpu__compute_7ContextI14... was problematic.

[francispoulin]

Great job @glwagner with sorting this out. Even though I don't follow all the details it seems like we have a much better understanding of the problem and how to fix some aspects of it. Well done!

[vchuravy]

I summon @maleadt for 2. & 3.

[glwagner]

On problem 1: meeting with @vchuravy we think there is a "recursive call cycle with levels of indirection". In other words, calling getindex on a BinaryOperation:

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/AbstractOperations/binary_operations.jl#L34

calls β.▶op = identity:

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/Operators/interpolation_utils.jl#L11

which calls op:

https://github.com/CliMA/Oceananigans.jl/blob/c3b688f9ef125faeac3aa9d7fd57f0dd2d392380/src/AbstractOperations/binary_operations.jl#L59-L60

which may invoke another call to either ▶a=identity or ▶b=identity (which might subsequently go back to getindex...) Due to some aspect of this process, the compiler throws up its hands because of something like "you wouldn't want unbound recursion in your compiler, right?"

A possible solution, which is also a hilarious hack, is to define multiple identity functions and use an internal counter to use the different yet identical copies of this function when constructing BinaryOperations. Then we wouldn't be recursing over identity (from the compiler's point of view), since we'd be calling identity1, identity2, etc. If we also need different flavors of BinaryOperation, we can do that too...

[vchuravy]

Alternatively we could think about linearizing the recursion before we go off and compile it.

[tomchor]

I'm not sure if it matters too much, but I'm going over some examples posted here and I can't reproduce some of the examples by @glwagner. I can compute some stuff that he can't, while the opposite is true for others

julia> compute!(ComputedField(∂x(∂x(u)))) # works as expected

julia> compute!(ComputedField(∂x(∂x(∂x(u))))) # works as expected

julia> compute!(ComputedField(∂x(∂x(∂x(∂x(u)))))) # doesn't work (not expected)

julia> compute!(ComputedField(u^2 + v^2)) # works (unexpected)

julia> compute!(ComputedField(u^2 + v^2 + w^2)) # works (expected)

Hopefully that'll give you guys more hints on what's going on. I'm running this on a Tesla V100 at NCAR's Casper cluster btw.

• Julia version 1.5.2
• Oceananigans version 0.45.2.
• For some reason CUDA doesn't appear when I type ]st, but I think it's 2.2.1 (sorry for the unfamiliarity with Julia)

EDIT:

The results above were obtained initializing a minimum model as posted above. I tried the same examples I just posted initializing a model differently and many more things are failing now (for example, compute!(ComputedField(u^2 + v^2)) fails).

[glwagner]

Can you please post the errors that you obtain in each case?

[tomchor]

Hmm, I can't reproduce the same results exactly. All I did before was honestly open a Julia session and just paste the examples you guys posted one by one.

Here's a pastebin with my whole session testing the commands I got in the previous post. (The comments of course don't reflect the outcome anymore.)

[glwagner]

Awesome, thank you, that's helpful. Stochastic errors are troubling.