From 34930db12e6b1bf369721b760fab72232276d290 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <gregory.leclaire.wagner@gmail.com>
Date: Tue, 9 Dec 2025 17:35:44 -0700
Subject: [PATCH 1/2] Test NonhydrostaticModel on AMDGPU

---
 test/test_amdgpu.jl | 54 +++++++++++++++++++++++++++++++++++++++------
 1 file changed, 47 insertions(+), 7 deletions(-)

diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl
index 86e86ff4b6..60dde42df5 100644
--- a/test/test_amdgpu.jl
+++ b/test/test_amdgpu.jl
@@ -1,25 +1,32 @@
 include("dependencies_for_runtests.jl")
 
 using AMDGPU
+using SeawaterPolynomials.TEOS10: TEOS10EquationOfState
 
-@testset "AMDGPU extension" begin
+@testset "AMDGPU with HydrostaticFreeSurfaceModel" begin
     roc = AMDGPU.ROCBackend()
     arch = GPU(roc)
 
     for FT in float_types
         @info "    Testing on $arch with $FT"
 
-        grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16))
+        grid = LatitudeLongitudeGrid(arch, FT, size=(4, 8, 16), longitude=(-60, 60), latitude=(0, 60), z=(0, 1))
 
-        @test parent(grid.xᶠᵃᵃ) isa ROCArray
-        @test parent(grid.xᶜᵃᵃ) isa ROCArray
+        @test parent(grid.λᶠᵃᵃ) isa ROCArray
+        @test parent(grid.λᶜᵃᵃ) isa ROCArray
+        @test parent(grid.φᵃᶠᵃ) isa ROCArray
+        @test parent(grid.φᵃᶜᵃ) isa ROCArray
+        @test parent(grid.zᵃᵃᶠ) isa ROCArray
+        @test parent(grid.zᵃᵃᶜ) isa ROCArray
         @test eltype(grid) == FT
         @test architecture(grid) isa GPU
 
-        model = HydrostaticFreeSurfaceModel(; grid,
+        equation_of_state = TEOS10EquationOfState()
+        buoyancy = SeawaterBuoyancy(; equation_of_state)
+
+        model = HydrostaticFreeSurfaceModel(; grid, buoyancy,
                                             coriolis = FPlane(latitude=45),
-                                            buoyancy = BuoyancyTracer(),
-                                            tracers = :b,
+                                            tracers = (:T, :S),
                                             momentum_advection = WENO(order=5),
                                             tracer_advection = WENO(order=5),
                                             free_surface = SplitExplicitFreeSurface(grid; substeps=60))
@@ -35,3 +42,36 @@ using AMDGPU
         @test time(simulation) == 3minutes
     end
 end
+
+@testset "AMDGPU with NonhydrostaticModel" begin
+    roc = AMDGPU.ROCBackend()
+    arch = GPU(roc)
+
+    for FT in float_types
+        @info "    Testing on $arch with $FT"
+
+        z = 0:16 |> collect
+        grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=(0, 1), y=(0, 1), z=z)
+
+        @test parent(grid.xᶠᵃᵃ) isa ROCArray
+        @test parent(grid.xᶜᵃᵃ) isa ROCArray
+        @test eltype(grid) == FT
+        @test architecture(grid) isa GPU
+
+        model = NonhydrostaticModel(; grid,
+                                    coriolis = FPlane(latitude=45),
+                                    buoyancy = BuoyancyTracer(),
+                                    tracers = :b,
+                                    advection = WENO(order=5))
+
+        for field in merge(model.velocities, model.tracers)
+            @test parent(field) isa ROCArray
+        end
+
+        simulation = Simulation(model, Δt=1minute, stop_iteration=3)
+        run!(simulation)
+
+        @test iteration(simulation) == 3
+        @test time(simulation) == 3minutes
+    end
+end

From 73fda6175e71c7e83f3b1d65454e23eea1966494 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <gregory.leclaire.wagner@gmail.com>
Date: Wed, 10 Dec 2025 10:47:03 -0700
Subject: [PATCH 2/2] fix amd tests

---
 test/test_amdgpu.jl | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl
index 60dde42df5..cc63063b9f 100644
--- a/test/test_amdgpu.jl
+++ b/test/test_amdgpu.jl
@@ -12,12 +12,16 @@ using SeawaterPolynomials.TEOS10: TEOS10EquationOfState
 
         grid = LatitudeLongitudeGrid(arch, FT, size=(4, 8, 16), longitude=(-60, 60), latitude=(0, 60), z=(0, 1))
 
-        @test parent(grid.λᶠᵃᵃ) isa ROCArray
-        @test parent(grid.λᶜᵃᵃ) isa ROCArray
-        @test parent(grid.φᵃᶠᵃ) isa ROCArray
-        @test parent(grid.φᵃᶜᵃ) isa ROCArray
-        @test parent(grid.zᵃᵃᶠ) isa ROCArray
-        @test parent(grid.zᵃᵃᶜ) isa ROCArray
+        @test parent(grid.Δxᶜᶜᵃ) isa ROCArray
+        @test parent(grid.Δxᶠᶜᵃ) isa ROCArray
+        @test parent(grid.Δxᶜᶠᵃ) isa ROCArray
+        @test parent(grid.Δxᶠᶠᵃ) isa ROCArray
+        @test parent(grid.Δyᶠᶜᵃ) isa ROCArray
+        @test parent(grid.Δyᶜᶠᵃ) isa ROCArray
+        @test parent(grid.Azᶜᶜᵃ) isa ROCArray
+        @test parent(grid.Azᶠᶜᵃ) isa ROCArray
+        @test parent(grid.Azᶜᶠᵃ) isa ROCArray
+        @test parent(grid.Azᶠᶠᵃ) isa ROCArray
         @test eltype(grid) == FT
         @test architecture(grid) isa GPU
 
@@ -53,8 +57,8 @@ end
         z = 0:16 |> collect
         grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=(0, 1), y=(0, 1), z=z)
 
-        @test parent(grid.xᶠᵃᵃ) isa ROCArray
-        @test parent(grid.xᶜᵃᵃ) isa ROCArray
+        @test parent(grid.z.cᵃᵃᶠ) isa ROCArray
+        @test parent(grid.z.cᵃᵃᶜ) isa ROCArray
         @test eltype(grid) == FT
         @test architecture(grid) isa GPU