Debug Global Sensitivity Analysis #58
Description
#56 implemented a simple sensitivity analysis of the column model.
Doing a global analysis with a ColumnGrid currently doesn't work
Here's a script
using Terrarium
using CUDA
using Dates
using Rasters, NCDatasets
using Statistics
using CairoMakie, GeoMakie
using Enzyme, Checkpointing
import RingGrids
import SpeedyWeather
# run on GPU if available
arch = CUDA.functional() ? Terrarium.GPU() : Terrarium.CPU()
# Load land-sea mask at ~1° resolution
land_sea_frac = convert.(Float32, dropdims(Raster("inputs/era5-land_land_sea_mask_N72.nc"), dims=Ti))
land_sea_frac_field = RingGrids.FullGaussianField(Matrix(land_sea_frac), input_as=Matrix)
heatmap(land_sea_frac_field)
# Set up grids
land_mask = land_sea_frac_field .> 0.5 # select only grid points with > 50% land
grid = ColumnRingGrid(arch, Float64, ExponentialSpacing(N=30), land_mask.grid, land_mask)
lon, lat = RingGrids.get_londlatds(grid.rings)
# Initial conditions
initializer = FieldInitializers(
# steady-ish state initial condition for temperature
temperature = (x,z) -> -1 - 0.01*z,
# fully saturated soil pores
saturation_water_ice = 1.0,
)
model = SoilModel(grid; initializer)
# constant surface temperature of 1°C
bcs = PrescribedSurfaceTemperature(:T_ub, 1.0)
integrator = initialize(model, ForwardEuler(), boundary_conditions=bcs)
# spin up a little
@time run!(integrator, period=Day(5), Δt=900.0)
# Enzyme prep
scheme = Revolve(1)
dintegrator = make_zero(integrator)
N_t = 200
autodiff(set_runtime_activity(Enzyme.Reverse), run!, Const, Duplicated(integrator, dintegrator), Const(scheme), Const(N_t))And here's the error
ERROR: EnzymeRuntimeException: Enzyme execution failed within
var"#launch!#55"(::Base.Pairs{…}, ::typeof(Oceananigans.Utils.launch!), ::ColumnRingGrid{…}, ::Symbol, ::typeof(Terrarium.compute_energy_tendency!), ::StateVariables{…}, ::ColumnRingGrid{…}, ::SoilEnergyBalance{…}, ::SoilHydrology{…}, ::HomogeneousSoil{…}, ::ConstantSoilCarbonDensity{…})
@ Terrarium ~/Nextcloud/Terrarium.jl/src/grids/grid_utils.jl:2
Hint: catch this exception as `err` and call `code_typed(err)` to inspect the surrounding code.
Enzyme: jl_call calling convention not implemented in aug_forward for %82 = call nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* %81, {} addrspace(10)* noundef @"ejl_inserted$jl_global_131982$false$4729261040", {} addrspace(10)* @"ejl_inserted$_KernelAbstractions_construct_131983$false$5590564240", {} addrspace(10)* @"ejl_inserted$jl_global_131984$false$5590563920", {} addrspace(10)* @"ejl_inserted$jl_global_131985$false$14749146656", {} addrspace(10)* nonnull %75, {} addrspace(10)* nonnull %value_phi.i, {} addrspace(10)* @"ejl_inserted$jl_global_131986$false$13715176320") #18, !dbg !97
Stacktrace:
[1] construct
@ ~/.julia/packages/KernelAbstractions/X5fk1/src/KernelAbstractions.jl:779 [inlined]
[2] _compute_energy_tendency!
@ ~/.julia/packages/KernelAbstractions/X5fk1/src/macros.jl:61 [inlined]
[3] compute_energy_tendency!
@ ~/.julia/packages/KernelAbstractions/X5fk1/src/macros.jl:70 [inlined]
[4] #configure_kernel#8
@ ~/.julia/packages/Oceananigans/VsQ7B/src/Utils/kernel_launching.jl:284 [inlined]
[5] configure_kernel
@ ~/.julia/packages/Oceananigans/VsQ7B/src/Utils/kernel_launching.jl:278 [inlined]
[6] #_launch!#17
@ ~/.julia/packages/Oceananigans/VsQ7B/src/Utils/kernel_launching.jl:372 [inlined]
[7] _launch!
@ ~/.julia/packages/Oceananigans/VsQ7B/src/Utils/kernel_launching.jl:365 [inlined]
[8] launch!
@ ~/.julia/packages/Oceananigans/VsQ7B/src/Utils/kernel_launching.jl:342 [inlined]
[9] #launch!#55
@ ~/Nextcloud/Terrarium.jl/src/grids/grid_utils.jl:4 [inlined]
[10] augmented_julia__launch__55_131968_inner_1wrap
@ ~/Nextcloud/Terrarium.jl/src/grids/grid_utils.jl:0
[11] macro expansion
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6652 [inlined]
[12] enzyme_call
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6131 [inlined]
[13] AugmentedForwardThunk
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6079 [inlined]
[14] macro expansion
@ ~/.julia/packages/Enzyme/i91IH/src/rules/jitrules.jl:447 [inlined]
[15] runtime_generic_augfwd(::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Terrarium.var"##launch!#55", ::Nothing, ::@Kwargs{}, ::Nothing, ::typeof(Oceananigans.Utils.launch!), ::Nothing, ::ColumnRingGrid{…}, ::Base.RefValue{…}, ::Symbol, ::Nothing, ::typeof(Terrarium.compute_energy_tendency!), ::Nothing, ::StateVariables{…}, ::Base.RefValue{…}, ::ColumnRingGrid{…}, ::Base.RefValue{…}, ::SoilEnergyBalance{…}, ::Base.RefValue{…}, ::SoilHydrology{…}, ::Base.RefValue{…}, ::HomogeneousSoil{…}, ::Base.RefValue{…}, ::ConstantSoilCarbonDensity{…}, ::Base.RefValue{…})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/i91IH/src/rules/jitrules.jl:574
[16] launch!
@ ~/Nextcloud/Terrarium.jl/src/grids/grid_utils.jl:2 [inlined]
[17] compute_tendencies!
@ ~/Nextcloud/Terrarium.jl/src/processes/soil/energy/soil_energy.jl:75 [inlined]
[18] compute_tendencies!
@ ~/Nextcloud/Terrarium.jl/src/models/soil/soil_model.jl:77 [inlined]
[19] #update_state!#91
@ ~/Nextcloud/Terrarium.jl/src/state_variables.jl:100
[20] update_state!
@ ~/Nextcloud/Terrarium.jl/src/state_variables.jl:94 [inlined]
[21] update_state!
@ ~/Nextcloud/Terrarium.jl/src/timesteppers/model_integrator.jl:58 [inlined]
[22] #timestep!#191
@ ~/Nextcloud/Terrarium.jl/src/timesteppers/model_integrator.jl:95 [inlined]
[23] timestep!
@ ~/Nextcloud/Terrarium.jl/src/timesteppers/model_integrator.jl:94 [inlined]
[24] macro expansion
@ ~/Nextcloud/Terrarium.jl/ext/TerrariumCheckpointingExt/TerrariumCheckpointingExt.jl:25 [inlined]
[25] #1
@ ~/.julia/packages/Checkpointing/FXuqe/src/Checkpointing.jl:133 [inlined]
[26] diffejulia__1_127828_inner_1wrap
@ ~/.julia/packages/Checkpointing/FXuqe/src/Checkpointing.jl:0
[27] macro expansion
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6652 [inlined]
[28] enzyme_call
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6131 [inlined]
[29] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6015 [inlined]
[30] autodiff(mode::ReverseMode{…}, f::Duplicated{…}, ::Type{…}, args::Const{…})
@ Enzyme ~/.julia/packages/Enzyme/i91IH/src/Enzyme.jl:521
[31] rev_checkpoint_for(config::EnzymeCore.EnzymeRules.RevConfigWidth{…}, body_input::Function, dbody::Function, alg::Revolve{…}, range::UnitRange{…})
@ Checkpointing ~/.julia/packages/Checkpointing/FXuqe/src/Schemes/Revolve.jl:433
[32] reverse
@ ~/.julia/packages/Checkpointing/FXuqe/src/Rules/EnzymeRules.jl:51 [inlined]
[33] macro expansion
@ ~/.julia/packages/Checkpointing/FXuqe/src/Checkpointing.jl:135 [inlined]
[34] run!
@ ~/Nextcloud/Terrarium.jl/ext/TerrariumCheckpointingExt/TerrariumCheckpointingExt.jl:24 [inlined]
[35] run!
@ ~/Nextcloud/Terrarium.jl/ext/TerrariumCheckpointingExt/TerrariumCheckpointingExt.jl:22 [inlined]
[36] run!
@ ~/Nextcloud/Terrarium.jl/ext/TerrariumCheckpointingExt/TerrariumCheckpointingExt.jl:0 [inlined]
[37] diffejulia_run__127612_inner_2wrap
@ ~/Nextcloud/Terrarium.jl/ext/TerrariumCheckpointingExt/TerrariumCheckpointingExt.jl:0
[38] macro expansion
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6652 [inlined]
[39] enzyme_call
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6131 [inlined]
[40] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/i91IH/src/compiler.jl:6015 [inlined]
[41] autodiff
@ ~/.julia/packages/Enzyme/i91IH/src/Enzyme.jl:521 [inlined]
[42] autodiff(::ReverseMode{…}, ::typeof(run!), ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Const{…})
@ Enzyme ~/.julia/packages/Enzyme/i91IH/src/Enzyme.jl:542
[43] top-level scope
@ REPL[71]:1
Some type information was truncated. Use `show(err)` to see complete types.