`getbc` fails with Reactant due to 2D indexing
This came up while working on Reactant for Terrarium (https://github.com/NumericalEarth/Terrarium.jl/pull/157). (Text partially LLM generated)
## Problem description
Oceananigans evaluates array-valued boundary conditions inside kernels through `getbc`
(`src/BoundaryConditions/boundary_condition.jl`):
```julia
@inline getbc(condition::AbstractArray, i::Integer, j::Integer, grid::AbstractGrid, args...) =
@inbounds condition[i, j]
```
The condition is indexed with **two** indices, but every Oceananigans field-like object is
three-dimensional. That works today only by accident, and only for one concrete type:
- **`Field`** defines a catch-all
```julia
@propagate_inbounds Base.getindex(f::Field, inds...) = getindex(f.data, inds...)
```
(`src/Fields/field.jl:424`), which forwards any number of indices straight to the underlying
`OffsetArray`, where a 2D index into a 3D array is padded cheaply and statically.
- **`AbstractOperation`** (e.g. the `UnaryOperation` produced by `-some_field`) defines only the
three-index form. `op[i, j]` therefore falls back to `Base`'s generic `AbstractArray` path:
```
_getindex -> _to_subscript_indices -> axes(::AbstractField) -> size(::Field) -> size(::AbstractGrid, loc, indices)
```
`axes(::Abstract3DField)` (`src/Fields/abstract_field.jl:62`) destructures `indices(f)`, and inside a
compiled kernel that whole chain becomes a dynamic dispatch.
On the CPU the fallback is harmless. Under Reactant the kernel fails to compile:
```
InvalidIRError: compiling MethodInstance for Oceananigans.BoundaryConditions.gpu__compute_z_bcs!(...)
Reason: unsupported call to an unknown function (call to jl_f_throw_methoderror)
Stacktrace:
[1] indexed_iterate @ ./tuple.jl:165
[2] axes @ Oceananigans/src/Fields/abstract_field.jl:62
[3] _to_subscript_indices @ ./abstractarray.jl:1400
[4] _to_subscript_indices @ ./abstractarray.jl:1398
[5] _getindex @ ./abstractarray.jl:1383
[6] getindex @ ./abstractarray.jl:1342
[7] getbc @ Oceananigans/src/BoundaryConditions/boundary_condition.jl:182
[8] getbc @ Oceananigans/src/BoundaryConditions/boundary_condition.jl:174
[9] compute_z_top_bc! @ Oceananigans/src/BoundaryConditions/compute_flux_bcs.jl:161
```
## Solution
Either we could just define
```julia
@inline getbc(condition::AbstractArray, i::Integer, j::Integer, grid::AbstractGrid, args...) =
@inbounds condition[i, j, 1]
```
or we give `AbstractOperation` a catch-all `getindex(op, inds...)` that supports 2D indexing like for `Field`. I am not sure if either options has drawbacks for the rest of the model. What do you prefer? Happy to do a PR.
## Full MWE with Reactant:
```Julia
using Oceananigans
using Oceananigans.Architectures: ReactantState, architecture
using Oceananigans.BoundaryConditions: FieldBoundaryConditions, FluxBoundaryCondition, compute_z_bcs!
using Oceananigans.Fields: Field, set!
using Oceananigans.Grids: Center
using Reactant
using CUDA
const NF = Float32
grid = RectilinearGrid(
ReactantState(), NF;
size = (1, 4), x = (0, 1), z = (-1, 0),
topology = (Periodic, Flat, Bounded)
)
# What Terrarium's timestepper does for every prognostic variable: add the flux boundary conditions of
# `c` to its tendency `Gc`.
apply_bcs!(Gc, c) = (compute_z_bcs!(Gc, c, architecture(c.grid)); return nothing)
function attempt(label, condition)
bcs = FieldBoundaryConditions(grid, (Center(), Center(), Center()); top = FluxBoundaryCondition(condition))
c = CenterField(grid; boundary_conditions = bcs)
Gc = CenterField(grid)
println(label)
try
@compile raise = true raise_first = true sync = true apply_bcs!(Gc, c)
println(" OK")
catch err
println(" FAILED:")
for line in first(split(sprint(showerror, err), "\n"), 12)
println(" ", line[1:min(end, 150)])
end
end
return nothing
end
surface_flux = Field{Center, Center, Nothing}(grid)
set!(surface_flux, NF(1))
# (1) BC condition is a `Field` -- compiles.
attempt("(1) Field-valued top flux BC", surface_flux)
# (2) BC condition is a lazy `AbstractOperation` -- fails. `-surface_flux` computes nothing; it builds
# a `UnaryOperation` that is evaluated when indexed.
attempt("(2) AbstractOperation-valued top flux BC (`-surface_flux`)", -surface_flux)
```
关闭于 13 天前 2 条评论