Reactant tracing for AnelasticFormulation

[bischtob]

When we try to trace and compile a simulation like the rising bubble with Reactant.jl, then the issue that we run into is that Oceananigans.jl currently requires a planned FFT that isn't supported for traced arrays from Reactant.jl.

[giordano]

Reportedly, XLA does the planning internally when doing FFTs, so we'll only need to teach Reactant to make the plan no-op. I can have a look later in the week.

[glwagner]

@bischtob do you have a short MWE that produces the error?

[giordano]

Yes, that'd help too!

[glwagner]

Probably something like this:

https://github.com/CliMA/Oceananigans.jl/blob/0217789f82583f942b77d9a105ebfb24aa2b8ee6/ext/OceananigansCUDAExt.jl#L111

but just return nothing eh?

[bischtob]

Will get to it in a bit.

[bischtob]

MWE

using Breeze
using Enzyme
using Oceananigans.Architectures: ReactantState
using Oceananigans.Units
using Reactant
using Statistics

function check_breeze_derivative()
        function make_grid(architecture)
            grid = RectilinearGrid(architecture; size = (8, 8), halo = (5, 5),
                                   x = (-10e3, 10e3), z = (0, 10e3),
                                   topology = (Periodic, Flat, Bounded))
            return grid
        end

        function make_model(grid)
            model = AtmosphereModel(grid; advection=WENO(order=9))
            return model
        end

        function make_initial_condition(model)
            r₀ = 2e3
            Δθ = 10 # K
            N² = 1e-6
            θ₀ = model.formulation.reference_state.potential_temperature
            g = model.thermodynamics.gravitational_acceleration

            @inline function θ_init_function(x, z;
                        x₀ = mean(xnodes(model.grid, Center())),
                        z₀ = 0.3*model.grid.Lz,
                        N² = N²)

                θ̄ = θ₀ * exp(N² * z / g)
                r = sqrt((x - x₀)^2 + (z - z₀)^2)
                θ′ = Δθ * max(0, 1 - r / r₀)

                return θ̄ + θ′
            end

            θ = Field{Center, Nothing, Center}(model.grid)
            set!(θ, θ_init_function)

            return θ
        end

        function run_model!(model, θᵢ, nsteps)
            set!(model, θ = θᵢ)        

            model.clock.iteration = 0
            model.clock.time = 0

            Δt = model.clock.last_Δt
            @trace mincut = true checkpointing = true track_numbers = false for i = 1:nsteps
                time_step!(model, Δt)
            end

            return nothing
        end

        function loss_function(model, θᵢ, nsteps)
            run_model!(model, θᵢ, nsteps)        
            return mean(model.temperature.^2)
        end
        
        function grad_loss_function(model, θᵢ, dθᵢ, nsteps)
            mode = set_strong_zero(Enzyme.ReverseWithPrimal)
            _, loss = autodiff(mode,
                               loss_function, 
                               Active,
                               Const(model),
                               Duplicated(θᵢ, dθᵢ),
                               Const(nsteps))

            return loss
        end

        # setup
        architecture = ReactantState()
        grid = make_grid(architecture)
        model = make_model(grid)
        θᵢ = make_initial_condition(model)

        # differentiate the loss function
        nsteps = 3
        dθᵢ = Field{Center, Nothing, Center}(model.grid)
        r_dloss_dθᵢ = @compile raise_first=true raise=true sync=true grad_loss_function(model, θᵢ, dθᵢ, nsteps)
        loss_value = r_dloss_dθᵢ(model, θᵢ, dθᵢ, nsteps)
end

Error:

┌ Error: Error in check_breeze_derivative.
│   exception =
│    MethodError: no method matching plan_forward_transform(::ConcreteIFRTArray{ComplexF64, 3, Nothing}, ::Periodic, ::Tuple{Int64, Int64}, ::UInt32)
│    The function `plan_forward_transform` exists, but no method is defined for this combination of argument types.
│    
│    Closest candidates are:
│      plan_forward_transform(::CuArray, ::Union{Bounded, Periodic}, ::Any, ::Any)
│       @ OceananigansCUDAExt ~/.julia/packages/Oceananigans/xGwrF/ext/OceananigansCUDAExt.jl:111
│      plan_forward_transform(::AbstractArray, ::Flat, ::Any...)
│       @ Oceananigans ~/.julia/packages/Oceananigans/xGwrF/src/Solvers/plan_transforms.jl:36
│      plan_forward_transform(::Array, ::Periodic, ::Any, ::Any)
│       @ Oceananigans ~/.julia/packages/Oceananigans/xGwrF/src/Solvers/plan_transforms.jl:16
│      ...
│    
│    Stacktrace:
|    ...

[bischtob]

However, this simple example works:

using Reactant
using AbstractFFTs

#backend = "cpu"
backend = "gpu"
Reactant.set_default_backend(backend)

x_host = rand(Float32, 16, 16)
x_r = Reactant.to_rarray(x_host)

fft_compiled = @compile fft(x_r)
y_r = fft_compiled(x_r)

which probably means that there is a simple workaround.

[bischtob]

If I understand correctly, then we should have a no-op version for this in OceananigansReactantExt.jl, e.g.,

import Oceananigans.Grids as GD
import Oceananigans.Solvers as SO

function SO.plan_forward_transform(A::ConcreteIFRTArray, ::Union{GD.Bounded, GD.Periodic}, dims, planner_flag)
    return nothing
end

function SO.plan_backward_transform(A::ConcreteIFRTArray, ::Union{GD.Bounded, GD.Periodic}, dims, planner_flag)
    return nothing
end

However, in this case I hit

┌ Error: Error in check_breeze_derivative.
│   exception =
│    TypeError: non-boolean (Reactant.TracedRNumber{Bool}) used in boolean context
│    Stacktrace:
|    ...

I haven't verified that this would actually work for the FFT issue, so take it with a grain of salt.

[giordano]

A lot simpler reproducer, without Oceananigans entirely:

using Reactant, FFTW

x_host = rand(ComplexF32, 16, 16)
x_r = Reactant.to_rarray(x_host)

### Non-planned, out-of-place FFT

fft_compiled = @compile fft(x_r)
y_r = fft_compiled(x_r)

# Check Reactant FFT works
@assert y_r ≈ fft(x_host)

### Planned, in-place FFT

copied_x = copy(x_host)
copied_x_r = Reactant.to_rarray(copied_x)

function planned_fft(x)
    plan = plan_fft!(x)
    return plan * x
end

planned_fft_compiled = @compile planned_fft(x_r)

# Just a sanity check the non-Reactant version works
@assert planned_fft(x_host)       ≈ fft(copied_x)
# Check Reactant FFT works
@assert compiled_planned_fft(x_r) ≈ fft(copied_x)

Errors with

ERROR: LoadError: MethodError: no method matching plan_fft!(::Reactant.TracedRArray{ComplexF32, 2}, ::UnitRange{Int64})
The function `plan_fft!` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  plan_fft!(::StridedArray{T, N}, ::Any; flags, timelimit, num_threads) where {T<:Union{ComplexF64, ComplexF32}, N}
   @ FFTW ~/.julia/packages/FFTW/uNgRy/src/fft.jl:788
  plan_fft!(::AbstractArray; kws...)
   @ AbstractFFTs ~/.julia/packages/AbstractFFTs/4iQz5/src/definitions.jl:68
  plan_fft!(::StridedArray{<:Union{ComplexF64, ComplexF32}}; kws...)
   @ FFTW ~/.julia/packages/FFTW/uNgRy/src/fft.jl:802

We discussed with Greg options to fix this both in Reactant and Oceananigans.

[bischtob]

Nice job! I think this is exactly right.