Calculating cell time scale in triply bounded `ImmersedBoundaryGrid` errors on GPU

[tomchor]

On main, trying to use TimeStepWizard fails due to a GPU-related error when calculating cell_advection_timescale. Here's a MWE:

using Oceananigans
using CUDA

N = 8
underlying_grid = RectilinearGrid(GPU(), topology = (Bounded, Bounded, Bounded), size = (N, N, N), extent = (1, 1, 1))

immersed_bottom = GridFittedBottom(-1/2)
model = NonhydrostaticModel(grid = ImmersedBoundaryGrid(underlying_grid, immersed_bottom))
simulation = Simulation(model; Δt = 1, stop_time = 1)

wizard = TimeStepWizard(cfl=0.2)
wizard(simulation)

This fails with

ERROR: Number of threads per block exceeds kernel limit (384 > 256).
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] diagnose_launch_failure(f::CuFunction, err::CuError; blockdim::CuDim3, threaddim::CuDim3, shmem::Int64)
    @ CUDA ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:126
  [3] launch(::CuFunction, ::CUDA.KernelState, ::CartesianIndices{…}, ::CartesianIndices{…}, ::CuDeviceArray{…}, ::Oceananigans.AbstractOperations.ConditionalOperation{…}; blocks::Int64, threads::Int64, cooperative::Bool, shmem::Int64, stream::CuStream)
    @ CUDA ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:79
  [4] launch
    @ ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:58 [inlined]
  [5] #1007
    @ ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:195 [inlined]
  [6] macro expansion
    @ ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:155 [inlined]
  [7] macro expansion
    @ ./none:0 [inlined]
  [8] convert_arguments
    @ ./none:0 [inlined]
  [9] #cudacall#1006
    @ ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:197 [inlined]
 [10] cudacall
    @ ~/.julia/packages/CUDA/ja0IX/lib/cudadrv/execution.jl:193 [inlined]
 [11] macro expansion
    @ ~/.julia/packages/CUDA/ja0IX/src/compiler/execution.jl:272 [inlined]
 [12] macro expansion
    @ ./none:0 [inlined]
 [13] (::CUDA.HostKernel{…})(::typeof(identity), ::typeof(min), ::Nothing, ::CartesianIndices{…}, ::CartesianIndices{…}, ::Val{…}, ::CuDeviceArray{…}, ::Oceananigans.AbstractOperations.ConditionalOperation{…}; convert::Val{…}, call_kwargs::@Kwargs{…})
    @ CUDA ./none:0
 [14] AbstractKernel
    @ ./none:0 [inlined]
 [15] macro expansion
    @ ~/.julia/packages/CUDA/ja0IX/src/compiler/execution.jl:114 [inlined]
 [16] mapreducedim!(f::typeof(identity), op::typeof(min), R::SubArray{…}, A::Oceananigans.AbstractOperations.ConditionalOperation{…}; init::Nothing)
    @ CUDA ~/.julia/packages/CUDA/ja0IX/src/mapreduce.jl:266
 [17] mapreducedim!(f::typeof(identity), op::typeof(min), R::SubArray{…}, A::Oceananigans.AbstractOperations.ConditionalOperation{…})
    @ CUDA ~/.julia/packages/CUDA/ja0IX/src/mapreduce.jl:169
 [18] mapreducedim!(f::Function, op::Function, R::SubArray{…}, A::Oceananigans.AbstractOperations.ConditionalOperation{…})
    @ GPUArrays ~/.julia/packages/GPUArrays/u6tui/src/host/mapreduce.jl:10
 [19] minimum!(f::Function, r::SubArray{…}, A::Oceananigans.AbstractOperations.ConditionalOperation{…}; init::Bool)
    @ Base ./reducedim.jl:1036
 [20] minimum(f::Function, c::KernelFunctionOperation{…}; condition::Nothing, mask::Float64, dims::Function)
    @ Oceananigans.Fields /mnt/storage1/tomaschor/Oceananigans.jl/src/Fields/field.jl:728
 [21] minimum
    @ /mnt/storage1/tomaschor/Oceananigans.jl/src/Fields/field.jl:717 [inlined]
 [22] minimum
    @ /mnt/storage1/tomaschor/Oceananigans.jl/src/Fields/field.jl:737 [inlined]
 [23] cell_advection_timescale
    @ /mnt/storage1/tomaschor/Oceananigans.jl/src/Advection/cell_advection_timescale.jl:16 [inlined]
 [24] cell_advection_timescale
    @ /mnt/storage1/tomaschor/Oceananigans.jl/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl:81 [inlined]
 [25] new_time_step(old_Δt::Float64, wizard::TimeStepWizard{…}, model::NonhydrostaticModel{…})
    @ Oceananigans.Simulations /mnt/storage1/tomaschor/Oceananigans.jl/src/Simulations/time_step_wizard.jl:103
 [26] (::TimeStepWizard{…})(simulation::Simulation{…})
    @ Oceananigans.Simulations /mnt/storage1/tomaschor/Oceananigans.jl/src/Simulations/time_step_wizard.jl:117
 [27] top-level scope
    @ REPL[10]:1

[tomchor]

Probably related to #4047

Also note, at least for me a workaround is to use a stretched grid and the issue goes away.

[ali-ramadhan]

Use your MWE @tomchor I was able to produce the same error on a V100 32GB. And by fixing field reductions to call interior on both arguments (see diff below) I managed to get rid of the error!

There may still be a CUDA.jl issue but maybe it was always an Oceananigans.jl issue?

I'm working on reproducing the error and adding the fix on this branch: https://github.com/CliMA/Oceananigans.jl/tree/ali/fix-kernel-limit-bug

diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl
index 953700fe7..c19198b3b 100644
--- a/src/AbstractOperations/AbstractOperations.jl
+++ b/src/AbstractOperations/AbstractOperations.jl
@@ -23,7 +23,7 @@ import Adapt

 import Oceananigans.Architectures: architecture, on_architecture
 import Oceananigans.BoundaryConditions: fill_halo_regions!
-import Oceananigans.Fields: compute_at!, indices
+import Oceananigans.Fields: compute_at!, indices, interior

 #####
 ##### Basic functionality
@@ -93,4 +93,10 @@ eval(define_multiary_operator(:*))
 push!(operators, :*)
 push!(multiary_operators, :*)

+# Fix for:
+# https://github.com/CliMA/Oceananigans.jl/issues/4047
+# https://github.com/CliMA/Oceananigans.jl/issues/4663
+
+interior(op::AbstractOperation) = interior(compute!(Field(op)))
+
 end # module
diff --git a/src/Fields/field.jl b/src/Fields/field.jl
index 250cf4e78..9325f64f7 100644
--- a/src/Fields/field.jl
+++ b/src/Fields/field.jl
@@ -750,7 +750,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod)
             loc = reduced_location(instantiated_location(c); dims)
             r = Field(loc, c.grid, T; indices=indices(c))
             initialize_reduced_field!(Base.$(reduction!), identity, r, conditioned_c)
-            Base.$(reduction!)(identity, interior(r), conditioned_c, init=false)
+            Base.$(reduction!)(identity, interior(r), interior(conditioned_c), init=false)

             if dims isa Colon
                 return @allowscalar first(r)

[tomchor]

Nice! Did you test this on other MWEs that reproduce the error? For example, here's a small one: #4694 (comment)

[glwagner]

I don't think we want to use

interior(op::AbstractOperation) = interior(compute!(Field(op)))

it's important to support the reduction of operations directly, otherwise memory allocations can be explode for some cases

[glwagner]

Basically on immersed boundary grid, we have to mask fields/operations before reducing them with ConditionalOperation. We don't want to always precompute the masked field/operation in a temporarily allocated Field, this could cause allocations to explode (eg if you compute diagnostics/monitoring inside progress). So I think its a strict requirement that we can reduce operations directly with no allocation.

[ali-ramadhan]

Nice! Did you test this on other MWEs that reproduce the error? For example, here's a small one: #4694 (comment)

Unfortunately this one works and does not produce the error haha :( Feels quite random and I'm not sure when/why it works and doesn't.

[ali-ramadhan]

I don't think we want to use

interior(op::AbstractOperation) = interior(compute!(Field(op)))

it's important to support the reduction of operations directly, otherwise memory allocations can be explode for some cases

Basically on immersed boundary grid, we have to mask fields/operations before reducing them with ConditionalOperation. We don't want to always precompute the masked field/operation in a temporarily allocated Field, this could cause allocations to explode (eg if you compute diagnostics/monitoring inside progress). So I think its a strict requirement that we can reduce operations directly with no allocation.

Definitely agree. I think interior(compute!(Field(op))) was my first instinct but the main issue I think is that CUDA.jl needs to be made aware that we should be looking at the interior of both r and conditioned_c.

There should be a way of communicating interior(conditioned_c) without computing the whole thing?

[ali-ramadhan]

Side note but the error does show up in CI with Julia 1.11 (on Tartarus with a Titan V) in a few places, with mapreducedim! featuring in the stacktraces:
https://buildkite.com/clima/oceananigans/builds/24879#019882d7-966f-4f03-8f09-d6e952b14e27/28-1042
https://buildkite.com/clima/oceananigans/builds/24879#019882d7-96ce-4b2b-8b2d-66f1098593a9/28-1002
https://buildkite.com/clima/oceananigans/builds/24879#019882d7-979f-47ac-8c64-6bcf69b0954a/28-2982

[glwagner]

interior(f::Field) returns a SubArray which is a view into f.data without the halos. For operations, there's no corresponding concept for interior -- we can't return/conceive of a SubArray of an operation.

[ali-ramadhan]

I can confirm that JuliaGPU/CUDA.jl#2869 fixes the MWE from this issue!