Error from `mapfirst!` into a SubArray of TracedArray

[glwagner]

Originates from ClimaOcean.Bathymetry.regrid_bathymetry ultimately

cc @NoraLoose from your global simulation example

julia> include("global_ocean_simulation.jl")
[ Info: Oceananigans will use 16 threads
ERROR: LoadError: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::Vararg{Int, N}) resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex(::Reactant.TracedRArray{Float64, 2}, ::Int64, ::Int64)
    @ Reactant.TracedRArrayOverrides ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:62
  [6] getindex
    @ ~/.julia/packages/Reactant/yO34a/ext/ReactantOffsetArraysExt.jl:52 [inlined]
  [7] Δxᶜᶜᵃ
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/Operators/spacings_and_areas_and_volumes.jl:230 [inlined]
  [8] Δxᶜᶜᶜ
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/Operators/spacings_and_areas_and_volumes.jl:93 [inlined]
  [9] Δx
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/Operators/spacings_and_areas_and_volumes.jl:370 [inlined]
 [10] getindex
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/AbstractOperations/kernel_function_operation.jl:73 [inlined]
 [11] getindex
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/AbstractOperations/conditional_operations.jl:146 [inlined]
 [12] getindex
    @ ./subarray.jl:290 [inlined]
 [13] _getindex
    @ ./abstractarray.jl:1340 [inlined]
 [14] getindex
    @ ./abstractarray.jl:1290 [inlined]
 [15] iterate
    @ ./abstractarray.jl:1216 [inlined]
 [16] iterate
    @ ./abstractarray.jl:1214 [inlined]
 [17] copyto_unaliased!(deststyle::IndexLinear, dest::Array{…}, srcstyle::IndexCartesian, src::SubArray{…})
    @ Base ./abstractarray.jl:1092
 [18] copyto!
    @ ./abstractarray.jl:1067 [inlined]
 [19] _collect_indices
    @ ./array.jl:777 [inlined]
 [20] collect(A::SubArray{Any, 3, Oceananigans.AbstractOperations.ConditionalOperation{…}, Tuple{…}, false})
    @ Base ./array.jl:761
 [21] constant(x::SubArray{Any, 3, Oceananigans.AbstractOperations.ConditionalOperation{…}, Tuple{…}, false}; location::Reactant.MLIR.IR.Location)
    @ Reactant.Ops ~/.julia/packages/Reactant/yO34a/src/Ops.jl:140
 [22] constant(x::SubArray{Any, 3, Oceananigans.AbstractOperations.ConditionalOperation{…}, Tuple{…}, false})
    @ Reactant.Ops ~/.julia/packages/Reactant/yO34a/src/Ops.jl:137
 [23] broadcast_to_size(arg::SubArray{…}, rsize::Base.Generator{…})
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/yO34a/src/TracedUtils.jl:966
 [24] broadcast_to_size(arg::Base.Broadcast.Extruded{SubArray{…}, Tuple{…}, Tuple{…}}, rsize::Tuple{Int64, Int64, Int64})
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/yO34a/src/TracedUtils.jl:1012
 [25] (::Reactant.TracedRArrayOverrides.var"#16#17"{Base.Broadcast.Broadcasted{…}})(a::Base.Broadcast.Extruded{SubArray{…}, Tuple{…}, Tuple{…}})
    @ Reactant.TracedRArrayOverrides ./none:0
 [26] iterate(::Base.Generator{Tuple{Base.Broadcast.Extruded{…}}, Reactant.TracedRArrayOverrides.var"#16#17"{Base.Broadcast.Broadcasted{…}}})
    @ Base ./generator.jl:47
 [27] _copyto!
    @ ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:704 [inlined]
 [28] _copyto!(none::SubArray{…}, none::Base.Broadcast.Broadcasted{…})
    @ Reactant ./<missing>:0
 [29] getproperty
    @ ./Base.jl:37 [inlined]
 [30] axes
    @ ./subarray.jl:491 [inlined]
 [31] _copyto!
    @ ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:697 [inlined]
 [32] call_with_reactant(::typeof(Reactant.TracedRArrayOverrides._copyto!), ::SubArray{…}, ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/yO34a/src/utils.jl:0
 [33] materialize!
    @ ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:668 [inlined]
 [34] materialize!(none::Reactant.TracedRArrayOverrides.AbstractReactantArrayStyle{…}, none::SubArray{…}, none::Base.Broadcast.Broadcasted{…})
    @ Reactant ./<missing>:0
 [35] getproperty
    @ ./Base.jl:37 [inlined]
 [36] materialize!
    @ ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:668 [inlined]
 [37] call_with_reactant(::typeof(Base.Broadcast.materialize!), ::Reactant.TracedRArrayOverrides.AbstractReactantArrayStyle{…}, ::SubArray{…}, ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/yO34a/src/utils.jl:0
 [38] materialize!
    @ ./broadcast.jl:911 [inlined]
 [39] map!
    @ ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:1168 [inlined]
 [40] mymap!
    @ ~/.julia/packages/Reactant/yO34a/src/ConcreteRArray.jl:586 [inlined]
 [41] mymap!(none::typeof(identity), none::SubArray{…}, none::SubArray{…})
    @ Reactant ./<missing>:0
 [42] broadcasted
    @ ./broadcast.jl:1349 [inlined]
 [43] broadcasted
    @ ./broadcast.jl:1341 [inlined]
 [44] map!
    @ ~/.julia/packages/Reactant/yO34a/src/TracedRArray.jl:1168 [inlined]
 [45] mymap!
    @ ~/.julia/packages/Reactant/yO34a/src/ConcreteRArray.jl:586 [inlined]
 [46] call_with_reactant(::typeof(Reactant.mymap!), ::typeof(identity), ::SubArray{…}, ::SubArray{…})
    @ Reactant ~/.julia/packages/Reactant/yO34a/src/utils.jl:0
 [47] make_mlir_fn(f::typeof(Reactant.mymap!), args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/yO34a/src/TracedUtils.jl:409
 [48] make_mlir_fn
    @ ~/.julia/packages/Reactant/yO34a/src/TracedUtils.jl:237 [inlined]
 [49] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(Reactant.mymap!), args::Tuple{…}, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::Tuple{}, optimize::Bool, cudnn_hlo_optimize::Bool, shardy_passes::Symbol, no_nan::Bool, transpose_propagate::Symbol, reshape_propagate::Symbol, optimize_communications::Bool, assert_nonallocating::Bool, backend::String, raise::Bool, raise_first::Bool, donated_args::Symbol, optimize_then_pad::Bool, runtime::Val{…}, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/yO34a/src/Compiler.jl:1150
 [50] compile_mlir!
    @ ~/.julia/packages/Reactant/yO34a/src/Compiler.jl:1104 [inlined]
 [51] compile_xla(f::Function, args::Tuple{typeof(identity), SubArray{…}, SubArray{…}}; client::Nothing, serializable::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/yO34a/src/Compiler.jl:2911
 [52] compile_xla
    @ ~/.julia/packages/Reactant/yO34a/src/Compiler.jl:2893 [inlined]
 [53] compile(f::Function, args::Tuple{typeof(identity), SubArray{…}, SubArray{…}}; sync::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/yO34a/src/Compiler.jl:2964
 [54] compile(f::Function, args::Tuple{typeof(identity), SubArray{…}, SubArray{…}})
    @ Reactant.Compiler ~/.julia/packages/Reactant/yO34a/src/Compiler.jl:2963
 [55] map!(f::Function, R::SubArray{…}, A::SubArray{…})
    @ Reactant ~/.julia/packages/Reactant/yO34a/src/ConcreteRArray.jl:591
 [56] mapfirst!(f::typeof(identity), R::SubArray{…}, A::Oceananigans.AbstractOperations.ConditionalOperation{…})
    @ Base ./reducedim.jl:274
 [57] initialize_reduced_field!(::typeof(minimum!), f::Function, r::Field{…}, c::Oceananigans.AbstractOperations.ConditionalOperation{…})
    @ Oceananigans.Fields ~/.julia/packages/Oceananigans/lJ5jT/src/Fields/field.jl:613
 [58] minimum(f::Function, c::KernelFunctionOperation{…}; condition::Nothing, mask::Float64, dims::Function)
    @ Oceananigans.Fields ~/.julia/packages/Oceananigans/lJ5jT/src/Fields/field.jl:706
 [59] minimum
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/Fields/field.jl:696 [inlined]
 [60] minimum
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/Fields/field.jl:716 [inlined]
 [61] minimum_xspacing
    @ ~/.julia/packages/Oceananigans/lJ5jT/src/Grids/nodes_and_spacings.jl:160 [inlined]
 [62] interpolate_bathymetry_in_passes(native_z::Field{…}, target_grid::OrthogonalSphericalShellGrid{…}; passes::Int64)
    @ ClimaOcean.Bathymetry ~/.julia/packages/ClimaOcean/DFzhM/src/Bathymetry.jl:220
 [63] regrid_bathymetry(target_grid::OrthogonalSphericalShellGrid{…}; height_above_water::Nothing, minimum_depth::Int64, dir::String, url::String, filename::String, interpolation_passes::Int64, major_basins::Int64)
    @ ClimaOcean.Bathymetry ~/.julia/packages/ClimaOcean/DFzhM/src/Bathymetry.jl:188
 [64] top-level scope
    @ ~/Projects/Enzymanigans.jl/global/global_ocean_simulation.jl:32
 [65] include(fname::String)
    @ Base.MainInclude ./client.jl:494
in expression starting at /Users/gregorywagner/Projects/Enzymanigans.jl/global/global_ocean_simulation.jl:32
Some type information was truncated. Use `show(err)` to see complete types.