search

tomchor / Oceanostics.jl

Diagnostics for Oceananigans
https://tomchor.github.io/Oceanostics.jl/
MIT License
24 stars 8 forks source link

Error calculating `TracerVarianceDissipationRate` for tuple closures on GPU #111

Closed tomchor closed 1 year ago

tomchor commented 1 year ago

When calculating TracerVarianceDissipationRate on a GPU using a tuple closure I get:

ERROR: LoadError: InvalidIRError: compiling kernel #gpu__compute!(KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.StaticSize{(12, 2, 12)}, KernelAbstractions.NDIteration.DynamicCheck, Nothing, Nothing, KernelAbstractions.NDIteration.NDRange{3, KernelAbstractions.NDIteration.StaticSize{(1, 1, 12)}, KernelAbstractions.NDIteration.StaticSize{(16, 16, 1)}, Nothing, Nothing}}, OffsetArrays.OffsetArray{Float64, 3, SubArray{Float64, 3, CuDeviceArray{Float64, 3, 1}, Tuple{UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}}, false}}, Oceananigans.AbstractOperations.BinaryOperation{Center, Center, Center, typeof(/), KernelFunctionOperation{Center, Center, Center, NamedTuple{(:closure, :clock, :buoyancy, :id), Tuple{Tuple{ScalarDiffusivity{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Oceananigans.TurbulenceClosures.ThreeDimensionalFormulation, Float64, NamedTuple{(:b,), Tuple{Float64}}}, SmagorinskyLilly{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Float64, NamedTuple{(:b,), Tuple{Float64}}}}, NamedTuple{(:time, :iteration, :stage), Tuple{Float64, Int64, Int64}}, Buoyancy{BuoyancyTracer, Oceananigans.Grids.ZDirection}, Val{1}}}, RectilinearGrid{Float64, Bounded, Periodic, Bounded, Float64, Float64, OffsetArrays.OffsetVector{Float64, CuDeviceVector{Float64, 1}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, CuDeviceVector{Float64, 1}}, Nothing}, Float64, typeof(Oceanostics.FlowDiagnostics.tracer_variance_dissipation_rate_ccc), Tuple{Tuple{Nothing, NamedTuple{(:νₑ,), Tuple{OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}}}}, OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, NamedTuple{(:u, :v, :w, :b), NTuple{4, OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}}}}}, Float64, typeof(Oceananigans.Operators.identity1), typeof(Oceananigans.Operators.identity2), RectilinearGrid{Float64, Bounded, Periodic, Bounded, Float64, Float64, OffsetArrays.OffsetVector{Float64, CuDeviceVector{Float64, 1}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, CuDeviceVector{Float64, 1}}, Nothing}, Float64}, Tuple{UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to mapreduce_empty)
Stacktrace:
  [1] reduce_empty
    @ ./reduce.jl:356
  [2] reduce_empty_iter
    @ ./reduce.jl:379
  [3] reduce_empty_iter
    @ ./reduce.jl:378
  [4] foldl_impl
    @ ./reduce.jl:49
  [5] mapfoldl_impl
    @ ./reduce.jl:44
  [6] #mapfoldl#259
    @ ./reduce.jl:170
  [7] mapfoldl
    @ ./reduce.jl:170
  [8] #mapreduce#263
    @ ./reduce.jl:302
  [9] mapreduce
    @ ./reduce.jl:302
 [10] #sum#266
    @ ./reduce.jl:528
 [11] sum
    @ ./reduce.jl:528
 [12] #sum#267
    @ ./reduce.jl:557
 [13] sum
    @ ./reduce.jl:557
 [14] diffusive_flux_x
    @ /glade/work/tomasc/.julia/packages/Oceanostics/d5inD/src/FlowDiagnostics.jl:354
 [15] χxᶠᶜᶜ
    @ /glade/work/tomasc/.julia/packages/Oceanostics/d5inD/src/FlowDiagnostics.jl:359
 [16] ℑxᶜᵃᵃ
    @ /glade/work/tomasc/.julia/packages/Oceananigans/Zg9Bd/src/Operators/interpolation_operators.jl:33
 [17] tracer_variance_dissipation_rate_ccc
    @ /glade/work/tomasc/.julia/packages/Oceanostics/d5inD/src/FlowDiagnostics.jl:371

The issue seem to be coming from these lines: https://github.com/tomchor/Oceanostics.jl/blob/376c25e90397dfa3e13b305fed3bb16aa7dbed6a/src/FlowDiagnostics.jl#L353-L356

Although it's not clear to me why that's the case.

It's also a shame we don't have an available GPU for testing for this repo. It would have caught that.

tomchor commented 1 year ago

Here's a MWE:

using Oceananigans
using Oceanostics.FlowDiagnostics: TracerVarianceDissipationRate

grid = RectilinearGrid(GPU(), size=(4, 4, 4), extent=(1,1,1))
model = NonhydrostaticModel(; grid, closure=(VerticalScalarDiffusivity(ν=1), HorizontalScalarDiffusivity(ν=2)), tracers=:b)
χ = Field(TracerVarianceDissipationRate(model, :b))
compute!(χ)
tomchor commented 1 year ago

Apparently the reason why this is failing is because you can't dynamically allocate an Array inside a GPU Kernel. So things like

diff_flux(i, j, k, grid, closure_tuple::Tuple, diffusivity_fields, args...) = 
        sum(diff_flux(i, j, k, grid, closure, diffusivities, args...) for (closure, diffusivities) in zip(closure_tuple, diffusivity_fields))

Have to replaced by stuff like

diff_flux(i, j, k, grid, closure_tuple::Tuple, diffusivity_fields, args...) = 
        diff_flux(i, j, k, grid, closure_tuple[1], diffusivity_fields[1], args...) + 
        diff_flux(i, j, k, grid, closure_tuple[2:end], diffusivity_fields[2:end], args...)