Update CI to Julia version to 1.12.0 (#4836)

* Update CI to Julia version to 1.12.0

* try increasing reactant compact

* try without reactant

* just run the docs

* bugfix

* Change grid bottom function to GridFittedBottom

* Update warning for Julia version compatibility

* take up the changes in #3836

* try without distributed

* imports alphabetically

* use a local accumulator and reinstate tests

* reinstate all tests

* white line

* reinstate distributed examples

* Update KernelAbstractions version to 0.9.39

* Disable doctest temporarily

* Update pipeline.yml

* separate out reactant and enzyme tests

* try like this

* don't need to make the new directory

* fixing

* use a new key

* need to create the new folder

* mv manifest to new version

* Update .buildkite/pipeline.yml

Co-authored-by: Mosè Giordano <[email protected]>

* [CI] Automatically create version-specific manifests

* enzyme on 1.10

* whyt would / not compile? meybe because of / 0?

* try removing internal closures

* this should do the trick

* up to six operands

* Update make.jl

* Apply suggestion from @simone-silvestri

* Apply suggestion from @simone-silvestri

* Apply suggestion from @simone-silvestri

* another bugfix

---------

Co-authored-by: Simone Silvestri <[email protected]>
Co-authored-by: Gregory L. Wagner <[email protected]>
Co-authored-by: Mosè Giordano <[email protected]>

Author: 4 people

.buildkite/pipeline.yml

@@ -1,5 +1,6 @@
 env:
-  JULIA_VERSION: "1.10.10"
+  JULIA_VERSION: "1.12.1"
+  JULIA_VERSION_ENZYME: "1.10.10"
   BUILDKITE_HOME: "/var/lib/buildkite-agent"
   JULIA_DEPOT_PATH: "$BUILDKITE_HOME/.julia-oceananigans"
   JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
@@ -12,7 +13,7 @@ with:
   coverage: false
 
 steps:
-  - label: "🏕️ initialize environment"
+  - label: "🏕️ initialize general environment"
     key: "init"
     agents:
       queue: "Oceananigans-nautilus"
@@ -26,6 +27,29 @@ steps:
       cd "$OCEANANIGANS_DIR"
 
       julia +$JULIA_VERSION -O0 --color=yes --project -e 'using Pkg; Pkg.test()'
+      # After we have instantiated the general environment, move manifest to the
+      # version-specific filename, not to clash with the manifest for Enzyme.
+      mv -v Manifest.toml Manifest-v$(echo "${JULIA_VERSION}" | cut -d. -f1-2).toml
+    retry:
+      automatic:
+        - exit_status: 1
+          limit: 1
+
+  - label: "🏕️ initialize enzyme environment"
+    key: "init_enzyme"
+    depends_on: "init"
+    agents:
+      queue: "Oceananigans-nautilus"
+    env:
+      TEST_GROUP: "init"
+    command: |
+      juliaup add $JULIA_VERSION_ENZYME
+      cd "$OCEANANIGANS_DIR"
+
+      julia +$JULIA_VERSION_ENZYME -O0 --color=yes --project -e 'using Pkg; Pkg.test()'
+      # After we have instantiated the Enzyme environment, move manifest to the
+      # version-specific filename, not to clash with the general manifest.
+      # mv -v Manifest.toml Manifest-v$(echo "${JULIA_VERSION_ENZYME}" | cut -d. -f1-2).toml
     retry:
       automatic:
         - exit_status: 1
@@ -58,12 +82,6 @@ steps:
         export TEST_ARCHITECTURE="GPU"
       fi
 
-      # Force IFRT runtime for Reactant (this can be moved to Project.toml on 1.11)
-      touch LocalPreferences.toml
-      echo "[Reactant]" >> LocalPreferences.toml
-      echo "xla_runtime = \"IFRT\"" >> LocalPreferences.toml
-      cat LocalPreferences.toml
-
       # Strip emoji for environment variable
       group="{{ matrix.group }}"
       export TEST_GROUP="\${group#* }"
@@ -80,7 +98,6 @@ steps:
           - "GPU"
         group:
           - "🦖 poisson_solvers_2"
-          - "🎭 reactant_2"
           - "🐇 unit"
           - "👻 abstract_operations"
           - "🕊 poisson_solvers_1"
@@ -99,14 +116,51 @@ steps:
           - "🍂 lagrangian_particles"
           - "🧅 multi_region"
           - "🦧 scripts"
-          - "👺 enzyme"
           - "🍱 xesmf"
-          - "👹 reactant_1"
+
     retry:
       automatic:
         - exit_status: 1
           limit: 1
 
+  - label: "{{ matrix.architecture }} - {{ matrix.group }} tests"
+    key: "tests_enzyme_reactant"
+    depends_on: "init_enzyme"
+    agents:
+      queue: "Oceananigans-nautilus"
+    command: |
+      # Add matrix-specific environment variables
+      if [[ "{{ matrix.architecture }}" == "CPU" ]]; then
+        export CUDA_VISIBLE_DEVICES="-1"
+        export TEST_ARCHITECTURE="CPU"
+      else
+        export CUDA_VISIBLE_DEVICES="0"
+        export TEST_ARCHITECTURE="GPU"
+      fi
+
+      # Strip emoji for environment variable
+      group="{{ matrix.group }}"
+      export TEST_GROUP="\${group#* }"
+      echo $TEST_GROUP
+
+      # Run tests
+      cd "$OCEANANIGANS_DIR"
+      julia +$JULIA_VERSION_ENZYME -O0 --color=yes --project -e 'using Pkg; Pkg.test()'
+
+    matrix:
+      setup:
+        architecture:
+          - "CPU"
+          - "GPU"
+        group:
+          - "🎭 reactant_2"
+          - "👺 enzyme"
+          - "👹 reactant_1"
+
+    retry:
+      automatic:
+        - exit_status: 1
+          limit: 1
   #####
   ##### Blocking wait on tests and documentation
   #####

Project.toml

@@ -77,7 +77,7 @@ GPUArraysCore = "0.2"
 Glob = "1.3"
 InteractiveUtils = "1.9"
 JLD2 = "0.4, 0.5, 0.6"
-KernelAbstractions = "0.9.21"
+KernelAbstractions = "0.9.39"
 Krylov = "0.10.0"
 KrylovPreconditioners = "0.3.2"
 LinearAlgebra = "1.9"
@@ -92,7 +92,7 @@ OffsetArrays = "1.4"
 OrderedCollections = "1.1"
 Printf = "1.9"
 Random = "1.9"
-Reactant = "0.2.162"
+Reactant = "0.2.169"
 ReactantCore = "0.1"
 Rotations = "1.0"
 SeawaterPolynomials = "0.3.9"

src/AbstractOperations/multiary_operations.jl

@@ -13,9 +13,44 @@ struct MultiaryOperation{LX, LY, LZ, N, O, A, IN, G, T} <: AbstractOperation{LX,
     end
 end
 
+# NOTE!!! The GPU compiler has trouble inferring this operation which might lead to dynamic evaluations.
 @inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, N}, i, j, k)  where {LX, LY, LZ, N} =
     Π.op(ntuple(γ -> Π.▶[γ](i, j, k, Π.grid, Π.args[γ]), Val(N))...)
 
+# Try to improve inferrability by hardcoding methods for small N
+@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 1}, i, j, k)  where {LX, LY, LZ} =
+    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]))
+
+@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 2}, i, j, k)  where {LX, LY, LZ} =
+    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]), 
+         Π.▶[2](i, j, k, Π.grid, Π.args[2]))
+
+@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 3}, i, j, k)  where {LX, LY, LZ} =
+    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]), 
+         Π.▶[2](i, j, k, Π.grid, Π.args[2]), 
+         Π.▶[3](i, j, k, Π.grid, Π.args[3]))
+
+@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 4}, i, j, k)  where {LX, LY, LZ} =
+    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]), 
+         Π.▶[2](i, j, k, Π.grid, Π.args[2]), 
+         Π.▶[3](i, j, k, Π.grid, Π.args[3]), 
+         Π.▶[4](i, j, k, Π.grid, Π.args[4]))
+
+@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 5}, i, j, k)  where {LX, LY, LZ} =
+    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]), 
+         Π.▶[2](i, j, k, Π.grid, Π.args[2]), 
+         Π.▶[3](i, j, k, Π.grid, Π.args[3]), 
+         Π.▶[4](i, j, k, Π.grid, Π.args[4]), 
+         Π.▶[5](i, j, k, Π.grid, Π.args[5]))
+
+@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 6}, i, j, k)  where {LX, LY, LZ} =
+    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]), 
+         Π.▶[2](i, j, k, Π.grid, Π.args[2]), 
+         Π.▶[3](i, j, k, Π.grid, Π.args[3]), 
+         Π.▶[4](i, j, k, Π.grid, Π.args[4]), 
+         Π.▶[5](i, j, k, Π.grid, Π.args[5]), 
+         Π.▶[6](i, j, k, Π.grid, Π.args[6]))
+
 #####
 ##### MultiaryOperation construction
 #####

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -61,9 +61,9 @@ function mask_immersed_model_fields!(model, grid)
     η = displacement(model.free_surface)
     fields_to_mask = merge(model.auxiliary_fields, prognostic_fields(model))
 
-    foreach(fields_to_mask) do field
-        if field !== η
-            mask_immersed_field!(field)
+    foreach(keys(fields_to_mask)) do key
+        if key != :η
+            @inbounds mask_immersed_field!(fields_to_mask[key])
         end
     end
     mask_immersed_field_xy!(η, k=size(grid, 3)+1)

src/Models/NonhydrostaticModels/update_hydrostatic_pressure.jl

@@ -12,10 +12,12 @@ the `buoyancy_perturbationᶜᶜᶜ` downwards:
 @kernel function _update_hydrostatic_pressure!(pHY′, grid, buoyancy, C)
     i, j = @index(Global, NTuple)
 
-    @inbounds pHY′[i, j, grid.Nz] = - z_dot_g_bᶜᶜᶠ(i, j, grid.Nz+1, grid, buoyancy, C) * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)
+    pᵏ = - z_dot_g_bᶜᶜᶠ(i, j, grid.Nz+1, grid, buoyancy, C) * Δzᶜᶜᶠ(i, j, grid.Nz+1, grid)
+    @inbounds pHY′[i, j, grid.Nz] = pᵏ
 
     for k in grid.Nz-1 : -1 : 1
-        @inbounds pHY′[i, j, k] = pHY′[i, j, k+1] - z_dot_g_bᶜᶜᶠ(i, j, k+1, grid, buoyancy, C) * Δzᶜᶜᶠ(i, j, k+1, grid)
+        pᵏ -= z_dot_g_bᶜᶜᶠ(i, j, k+1, grid, buoyancy, C) * Δzᶜᶜᶠ(i, j, k+1, grid)
+        @inbounds pHY′[i, j, k] = pᵏ
     end
 end
 

src/MultiRegion/multi_region_field.jl

@@ -6,7 +6,7 @@ using Oceananigans.OutputWriters: output_indices
 
 using Base: @propagate_inbounds
 
-import Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
+import Oceananigans.BoundaryConditions: regularize_field_boundary_conditions, FieldBoundaryConditions
 import Oceananigans.Diagnostics: hasnan
 import Oceananigans.DistributedComputations: reconstruct_global_field, CommunicationBuffers
 import Oceananigans.Fields: set!, compute!, compute_at!, interior, communication_buffers,

src/Oceananigans.jl

@@ -119,10 +119,10 @@ using DocStringExtensions
 using FFTW
 
 function __init__()
-    if VERSION >= v"1.11.0"
-        @warn """You are using Julia v1.11 or later!"
-                 Oceananigans is currently tested on Julia v1.10."
-                 If you find issues with Julia v1.11 or later,"
+    if VERSION >= v"1.13.0"
+        @warn """You are using Julia v1.13 or later!"
+                 Oceananigans is currently tested on Julia v1.12."
+                 If you find issues with Julia v1.13 or later,"
                  please report at https://github.com/CliMA/Oceananigans.jl/issues/new"""
 
     end

test/test_init.jl

@@ -1,4 +1,3 @@
-using Enzyme
 using CUDA
 using Metal
 using AMDGPU