From 009bff7cf7d033595043a816b4cd396a0721b53e Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 10 Jun 2025 16:51:35 +1000 Subject: [PATCH 01/12] extend plan transforms for Flat --- ext/OceananigansAMDGPUExt.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ext/OceananigansAMDGPUExt.jl b/ext/OceananigansAMDGPUExt.jl index e89791da44..38d3d7c379 100644 --- a/ext/OceananigansAMDGPUExt.jl +++ b/ext/OceananigansAMDGPUExt.jl @@ -4,12 +4,16 @@ using AMDGPU using Oceananigans using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata using KernelAbstractions: __dynamic_checkbounds, __iterspace + import KernelAbstractions: __validindex import Oceananigans.Architectures: architecture, convert_to_device, on_architecture +import Oceananigans.Solvers: + plan_backward_transform, + plan_forward_transform const ROCGPU = GPU{<:AMDGPU.ROCBackend} ROCGPU() = GPU(AMDGPU.ROCBackend()) @@ -30,6 +34,9 @@ on_architecture(::ROCGPU, a::StepRangeLen) = a @inline convert_to_device(::ROCGPU, args) = AMDGPU.rocconvert(args) @inline convert_to_device(::ROCGPU, args::Tuple) = map(AMDGPU.rocconvert, args) +plan_backward_transform(A::ROCArray, ::Flat, args...) = nothing +plan_forward_transform(A::ROCArray, ::Flat, args...) = nothing + AMDGPU.Device.@device_override @inline function __validindex(ctx::MappedCompilerMetadata) if __dynamic_checkbounds(ctx) I = @inbounds linear_expand(__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x) From 8e2621b83255c9ccf0a8ed1063c7c48e1e7aa6dd Mon Sep 17 00:00:00 2001 From: Navid Constantinou Date: Tue, 10 Jun 2025 18:43:27 +0800 Subject: [PATCH 02/12] add methods for plan_fwd/bwd_transform --- ext/OceananigansAMDGPUExt.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ext/OceananigansAMDGPUExt.jl b/ext/OceananigansAMDGPUExt.jl index 38d3d7c379..83ad0624f7 100644 --- a/ext/OceananigansAMDGPUExt.jl +++ b/ext/OceananigansAMDGPUExt.jl @@ -1,6 +1,7 @@ module OceananigansAMDGPUExt using AMDGPU +using AMDGPU.rocFFT using Oceananigans using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata using KernelAbstractions: __dynamic_checkbounds, __iterspace @@ -34,6 +35,16 @@ on_architecture(::ROCGPU, a::StepRangeLen) = a @inline convert_to_device(::ROCGPU, args) = AMDGPU.rocconvert(args) @inline convert_to_device(::ROCGPU, args::Tuple) = map(AMDGPU.rocconvert, args) +function plan_forward_transform(A::ROCArray, ::Union{Bounded, Periodic}, dims, planner_flag) + length(dims) == 0 && return nothing + return AMDGPU.rocFFT.plan_fft!(A, dims) +end + +function plan_backward_transform(A::ROCArray, ::Union{Bounded, Periodic}, dims, planner_flag) + length(dims) == 0 && return nothing + return AMDGPU.rocFFT.plan_bfft!(A, dims) +end + plan_backward_transform(A::ROCArray, ::Flat, args...) = nothing plan_forward_transform(A::ROCArray, ::Flat, args...) = nothing From 661887f6cdafe3e931499297ef98eb5bae495848 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 10 Jun 2025 19:01:01 +0800 Subject: [PATCH 03/12] add tests for NonhydrostaticModel --- test/test_amdgpu.jl | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl index 86e86ff4b6..01516157e2 100644 --- a/test/test_amdgpu.jl +++ b/test/test_amdgpu.jl @@ -7,7 +7,7 @@ using AMDGPU arch = GPU(roc) for FT in float_types - @info " Testing on $arch with $FT" + @info " Testing HydrostaticFreeSurfaceModel on $arch with $FT" grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16)) @@ -33,5 +33,30 @@ using AMDGPU @test iteration(simulation) == 3 @test time(simulation) == 3minutes + + @info " Testing NonhydrostaticModel on $arch with $FT" + + pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid, maxiter=10; reltol=1e-7, abstol=1e-7, preconditioner=nothing), + Oceananigans.Solvers.FFTBasedPoissonSolver(grid)) + + for pressure_solver in pressure_solvers + + model = NonhydrostaticModel(; grid, pressure_solver + coriolis = FPlane(latitude=45), + buoyancy = BuoyancyTracer(), + tracers = :b, + momentum_advection = WENO(order=5), + tracer_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 end From 0c9783504e9d5d9953ab1644aa936a2975b40887 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 10 Jun 2025 19:09:33 +0800 Subject: [PATCH 04/12] compactify tests --- test/test_amdgpu.jl | 63 ++++++++++++++++++++++----------------------- 1 file changed, 31 insertions(+), 32 deletions(-) diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl index 01516157e2..dc7ff309fc 100644 --- a/test/test_amdgpu.jl +++ b/test/test_amdgpu.jl @@ -2,6 +2,20 @@ include("dependencies_for_runtests.jl") using AMDGPU +function timestep_simulation(model) + 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 + + return nothing +end + @testset "AMDGPU extension" begin roc = AMDGPU.ROCBackend() arch = GPU(roc) @@ -16,47 +30,32 @@ using AMDGPU @test eltype(grid) == FT @test architecture(grid) isa GPU - model = HydrostaticFreeSurfaceModel(; grid, - coriolis = FPlane(latitude=45), - buoyancy = BuoyancyTracer(), - tracers = :b, - momentum_advection = WENO(order=5), - tracer_advection = WENO(order=5), - free_surface = SplitExplicitFreeSurface(grid; substeps=60)) + coriolis = FPlane(latitude=45) + buoyancy = BuoyancyTracer() + tracers = :b + advection = WENO(order=5) + momentum_advection = tracer_advection = advection - for field in merge(model.velocities, model.tracers) - @test parent(field) isa ROCArray - end + pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid, maxiter=10; reltol=1e-7, abstol=1e-7, preconditioner=nothing), + Oceananigans.Solvers.FFTBasedPoissonSolver(grid)) + + free_surface = SplitExplicitFreeSurface(grid; substeps=60) - simulation = Simulation(model, Δt=1minute, stop_iteration=3) - run!(simulation) + model = HydrostaticFreeSurfaceModel(; grid, free_surface, + coriolis, buoyancy, tracers, + momentum_advection, tracer_advection) - @test iteration(simulation) == 3 - @test time(simulation) == 3minutes + timestep_simulation(model) @info " Testing NonhydrostaticModel on $arch with $FT" - pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid, maxiter=10; reltol=1e-7, abstol=1e-7, preconditioner=nothing), - Oceananigans.Solvers.FFTBasedPoissonSolver(grid)) for pressure_solver in pressure_solvers + model = NonhydrostaticModel(; grid, pressure_solver, + coriolis, buoyancy, + tracers, advection) - model = NonhydrostaticModel(; grid, pressure_solver - coriolis = FPlane(latitude=45), - buoyancy = BuoyancyTracer(), - tracers = :b, - momentum_advection = WENO(order=5), - tracer_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 + timestep_simulation(model) end end end From 30eec3a80a57a9e73405ba80e7aecb3db57560cd Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 10 Jun 2025 21:44:18 +0800 Subject: [PATCH 05/12] reorder --- ext/OceananigansAMDGPUExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/OceananigansAMDGPUExt.jl b/ext/OceananigansAMDGPUExt.jl index 83ad0624f7..dd0c8ea143 100644 --- a/ext/OceananigansAMDGPUExt.jl +++ b/ext/OceananigansAMDGPUExt.jl @@ -45,8 +45,8 @@ function plan_backward_transform(A::ROCArray, ::Union{Bounded, Periodic}, dims, return AMDGPU.rocFFT.plan_bfft!(A, dims) end -plan_backward_transform(A::ROCArray, ::Flat, args...) = nothing plan_forward_transform(A::ROCArray, ::Flat, args...) = nothing +plan_backward_transform(A::ROCArray, ::Flat, args...) = nothing AMDGPU.Device.@device_override @inline function __validindex(ctx::MappedCompilerMetadata) if __dynamic_checkbounds(ctx) From 495be4819b91d21881ace64f485f786b7cf218f1 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 10 Jun 2025 21:44:53 +0800 Subject: [PATCH 06/12] test few grids --- test/test_amdgpu.jl | 64 +++++++++++++++++++++++++++++---------------- 1 file changed, 41 insertions(+), 23 deletions(-) diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl index dc7ff309fc..4dea3d112d 100644 --- a/test/test_amdgpu.jl +++ b/test/test_amdgpu.jl @@ -2,16 +2,18 @@ include("dependencies_for_runtests.jl") using AMDGPU -function timestep_simulation(model) +function build_and_timestep_simulation(model) + FT = eltype(model) + for field in merge(model.velocities, model.tracers) @test parent(field) isa ROCArray end - simulation = Simulation(model, Δt=1minute, stop_iteration=3) + simulation = Simulation(model, Δt=1minute, stop_iteration=3, verbose=false) run!(simulation) @test iteration(simulation) == 3 - @test time(simulation) == 3minutes + @test time(simulation) ≈ FT(3minutes) return nothing end @@ -21,41 +23,57 @@ end arch = GPU(roc) for FT in float_types - @info " Testing HydrostaticFreeSurfaceModel on $arch with $FT" + @info "Testing grids on $arch with $FT..." + + regular_grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=(0, 4), y=(0, 1), z=(0, 16)) + horizontally_stretched_grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16)) + vertically_stretched_grid = RectilinearGrid(arch, FT, size=(16, 8, 4), x=(0, 16), y=(0, 1), z=[0, 1, 2, 3, 4]) + + @test parent(horizontally_stretched_grid.xᶠᵃᵃ) isa ROCArray + @test parent(horizontally_stretched_grid.xᶜᵃᵃ) isa ROCArray - grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16)) + @test parent(vertically_stretched_grid.z.cᵃᵃᶠ) isa ROCArray + @test parent(vertically_stretched_grid.z.cᵃᵃᶜ) isa ROCArray + @test parent(vertically_stretched_grid.z.Δᵃᵃᶠ) isa ROCArray + @test parent(vertically_stretched_grid.z.Δᵃᵃᶜ) isa ROCArray - @test parent(grid.xᶠᵃᵃ) isa ROCArray - @test parent(grid.xᶜᵃᵃ) isa ROCArray - @test eltype(grid) == FT - @test architecture(grid) isa GPU + for grid in (regular_grid, horizontally_stretched_grid, vertically_stretched_grid) + @test eltype(grid) == FT + @test architecture(grid) isa GPU + end + + @info "Testing HydrostaticFreeSurfaceModel on $arch with $FT..." coriolis = FPlane(latitude=45) buoyancy = BuoyancyTracer() tracers = :b advection = WENO(order=5) - momentum_advection = tracer_advection = advection - pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid, maxiter=10; reltol=1e-7, abstol=1e-7, preconditioner=nothing), - Oceananigans.Solvers.FFTBasedPoissonSolver(grid)) + for grid in (regular_grid, horizontally_stretched_grid, vertically_stretched_grid) + momentum_advection = tracer_advection = advection - free_surface = SplitExplicitFreeSurface(grid; substeps=60) + free_surface = SplitExplicitFreeSurface(grid; substeps=60) - model = HydrostaticFreeSurfaceModel(; grid, free_surface, - coriolis, buoyancy, tracers, - momentum_advection, tracer_advection) + model = HydrostaticFreeSurfaceModel(; grid, free_surface, + coriolis, buoyancy, tracers, + momentum_advection, tracer_advection) - timestep_simulation(model) + build_and_timestep_simulation(model) + end - @info " Testing NonhydrostaticModel on $arch with $FT" + @info "Testing NonhydrostaticModel on $arch with $FT..." + for grid in (regular_grid, vertically_stretched_grid) + pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid, maxiter=10; reltol=1e-7, abstol=1e-7, preconditioner=nothing), + Oceananigans.Solvers.FFTBasedPoissonSolver(grid)) - for pressure_solver in pressure_solvers - model = NonhydrostaticModel(; grid, pressure_solver, - coriolis, buoyancy, - tracers, advection) + for pressure_solver in pressure_solvers + model = NonhydrostaticModel(; grid, pressure_solver, + coriolis, buoyancy, + tracers, advection) - timestep_simulation(model) + build_and_timestep_simulation(model) + end end end end From 2f54d0ea350bca5efc425120e13378d77f9ab082 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 17 Jul 2025 07:37:13 +1000 Subject: [PATCH 07/12] merge main --- ext/OceananigansAMDGPUExt.jl | 40 ++++-------------------------------- 1 file changed, 4 insertions(+), 36 deletions(-) diff --git a/ext/OceananigansAMDGPUExt.jl b/ext/OceananigansAMDGPUExt.jl index c1c94325ab..e5b23d064b 100644 --- a/ext/OceananigansAMDGPUExt.jl +++ b/ext/OceananigansAMDGPUExt.jl @@ -1,27 +1,11 @@ module OceananigansAMDGPUExt -<<<<<<< HEAD -using AMDGPU -using AMDGPU.rocFFT -======= ->>>>>>> main using Oceananigans -using InteractiveUtils using AMDGPU, AMDGPU.rocSPARSE, AMDGPU.rocFFT + using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata +using InteractiveUtils using KernelAbstractions: __dynamic_checkbounds, __iterspace -<<<<<<< HEAD - -import KernelAbstractions: __validindex - -import Oceananigans.Architectures: - architecture, - convert_to_device, - on_architecture -import Oceananigans.Solvers: - plan_backward_transform, - plan_forward_transform -======= using KernelAbstractions import Oceananigans.Architectures as AC @@ -31,11 +15,10 @@ import Oceananigans.Fields as FD import Oceananigans.Grids as GD import Oceananigans.Solvers as SO import Oceananigans.Utils as UT -import SparseArrays: SparseMatrixCSC +import Oceananigans.DistributedComputations: Distributed import KernelAbstractions: __iterspace, __groupindex, __dynamic_checkbounds, __validindex, CompilerMetadata -import Oceananigans.DistributedComputations: Distributed ->>>>>>> main +import SparseArrays: SparseMatrixCSC const GPUVar = Union{ROCArray, Ptr} @@ -97,30 +80,16 @@ end @inline convert_to_device(::ROCGPU, args) = AMDGPU.rocconvert(args) @inline convert_to_device(::ROCGPU, args::Tuple) = map(AMDGPU.rocconvert, args) -<<<<<<< HEAD -function plan_forward_transform(A::ROCArray, ::Union{Bounded, Periodic}, dims, planner_flag) -======= BC.validate_boundary_condition_architecture(::ROCArray, ::AC.GPU, bc, side) = nothing BC.validate_boundary_condition_architecture(::ROCArray, ::AC.CPU, bc, side) = throw(ArgumentError("$side $bc must use `Array` rather than `ROCArray` on CPU architectures!")) function SO.plan_forward_transform(A::ROCArray, ::Union{GD.Bounded, GD.Periodic}, dims, planner_flag) ->>>>>>> main length(dims) == 0 && return nothing return AMDGPU.rocFFT.plan_fft!(A, dims) end -<<<<<<< HEAD -function plan_backward_transform(A::ROCArray, ::Union{Bounded, Periodic}, dims, planner_flag) - length(dims) == 0 && return nothing - return AMDGPU.rocFFT.plan_bfft!(A, dims) -end - -plan_forward_transform(A::ROCArray, ::Flat, args...) = nothing -plan_backward_transform(A::ROCArray, ::Flat, args...) = nothing - -======= FD.set!(v::Field, a::ROCArray) = FD.set_to_array!(v, a) DC.set!(v::DC.DistributedField, a::ROCArray) = DC.set_to_array!(v, a) @@ -129,7 +98,6 @@ function SO.plan_backward_transform(A::ROCArray, ::Union{GD.Bounded, GD.Periodic return AMDGPU.rocFFT.plan_ifft!(A, dims) end ->>>>>>> main AMDGPU.Device.@device_override @inline function __validindex(ctx::MappedCompilerMetadata) if __dynamic_checkbounds(ctx) I = @inbounds linear_expand(__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x) From f835c0dc88cbed8d11bb2f5c68d2cccb67495963 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Fri, 12 Dec 2025 16:00:51 -0700 Subject: [PATCH 08/12] Add tests for AMDGPU with HydrostaticFreeSurfaceModel --- test/test_amdgpu.jl | 43 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl index 4dea3d112d..b49b3f3bc8 100644 --- a/test/test_amdgpu.jl +++ b/test/test_amdgpu.jl @@ -1,6 +1,7 @@ include("dependencies_for_runtests.jl") using AMDGPU +using SeawaterPolynomials.TEOS10: TEOS10EquationOfState function build_and_timestep_simulation(model) FT = eltype(model) @@ -18,7 +19,7 @@ function build_and_timestep_simulation(model) return nothing end -@testset "AMDGPU extension" begin +@testset "AMDGPU on RectilinearGrids" begin roc = AMDGPU.ROCBackend() arch = GPU(roc) @@ -64,7 +65,8 @@ end @info "Testing NonhydrostaticModel on $arch with $FT..." for grid in (regular_grid, vertically_stretched_grid) - pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid, maxiter=10; reltol=1e-7, abstol=1e-7, preconditioner=nothing), + cg_kw = (maxiter=10, reltol=1e-7, abstol=1e-7, preconditioner=nothing) + pressure_solvers = (Oceananigans.Solvers.ConjugateGradientPoissonSolver(grid; cg_kw...), Oceananigans.Solvers.FFTBasedPoissonSolver(grid)) for pressure_solver in pressure_solvers @@ -77,3 +79,40 @@ end end end end + +@testset "AMDGPU on LatitudeLongitudeGrid with HydrostaticFreeSurfaceModel" begin + roc = AMDGPU.ROCBackend() + arch = GPU(roc) + + for FT in float_types + @info " Testing on $arch with $FT" + + 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.Δ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 + + equation_of_state = TEOS10EquationOfState() + buoyancy = SeawaterBuoyancy(; equation_of_state) + + model = HydrostaticFreeSurfaceModel(; grid, buoyancy, + coriolis = FPlane(latitude=45), + tracers = (:T, :S), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), + free_surface = SplitExplicitFreeSurface(grid; substeps=60)) + + build_and_time_step_simulation(model) + end +end + From b07b3c061b4f7c9453559cdf9fc28629112acc1f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mos=C3=A8=20Giordano?= Date: Sun, 4 Jan 2026 18:49:12 +0100 Subject: [PATCH 09/12] [AMDGPUExt] Use FFT functions from `AbstractFFTs` instead non-existing ones in `rocFFT` --- Project.toml | 3 ++- ext/OceananigansAMDGPUExt.jl | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index 524d24c87e..021bf71341 100755 --- a/Project.toml +++ b/Project.toml @@ -36,6 +36,7 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" @@ -49,7 +50,7 @@ XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] -OceananigansAMDGPUExt = "AMDGPU" +OceananigansAMDGPUExt = ["AMDGPU", "AbstractFFTs"] OceananigansCUDAExt = ["CUDA", "GPUArrays", "GPUArraysCore"] OceananigansConservativeRegriddingExt = "ConservativeRegridding" OceananigansEnzymeExt = "Enzyme" diff --git a/ext/OceananigansAMDGPUExt.jl b/ext/OceananigansAMDGPUExt.jl index 06db66ffef..473e53ed38 100644 --- a/ext/OceananigansAMDGPUExt.jl +++ b/ext/OceananigansAMDGPUExt.jl @@ -1,7 +1,8 @@ module OceananigansAMDGPUExt using Oceananigans -using AMDGPU, AMDGPU.rocSPARSE, AMDGPU.rocFFT +using AMDGPU, AMDGPU.rocSPARSE +using AbstractFFTs: plan_fft!, plan_ifft! using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata using InteractiveUtils @@ -90,7 +91,7 @@ BC.validate_boundary_condition_architecture(::ROCArray, ::AC.CPU, bc, side) = function SO.plan_forward_transform(A::ROCArray, ::Union{GD.Bounded, GD.Periodic}, dims, planner_flag) length(dims) == 0 && return nothing - return AMDGPU.rocFFT.plan_fft!(A, dims) + return plan_fft!(A, dims) end FD.set!(v::Field, a::ROCArray) = FD.set_to_array!(v, a) @@ -98,7 +99,7 @@ DC.set!(v::DC.DistributedField, a::ROCArray) = DC.set_to_array!(v, a) function SO.plan_backward_transform(A::ROCArray, ::Union{GD.Bounded, GD.Periodic}, dims, planner_flag) length(dims) == 0 && return nothing - return AMDGPU.rocFFT.plan_ifft!(A, dims) + return plan_ifft!(A, dims) end AMDGPU.Device.@device_override @inline function __validindex(ctx::MappedCompilerMetadata) From de4e9e96ab6d577e82c595acf871b2299c410f12 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Sun, 4 Jan 2026 11:19:10 -0700 Subject: [PATCH 10/12] Update test/test_amdgpu.jl --- test/test_amdgpu.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl index 1480005b10..91d667da96 100644 --- a/test/test_amdgpu.jl +++ b/test/test_amdgpu.jl @@ -93,8 +93,6 @@ end @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 From fa511bf2cd0b961cfdb15b80cf1b3a596e41ae3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mos=C3=A8=20Giordano?= Date: Mon, 5 Jan 2026 00:02:41 +0100 Subject: [PATCH 11/12] Add compat bound for `AbstractFFTs` --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 021bf71341..264c027097 100755 --- a/Project.toml +++ b/Project.toml @@ -63,6 +63,7 @@ OceananigansXESMFExt = ["XESMF"] [compat] AMDGPU = "1.3.6, 2" +AbstractFFTs = "1.5" Adapt = "4.1.1" Aqua = "0.8" CUDA = "=5.8.5, 5.9.1" From f3da27f11c8c4e7b945b708c293e3b6baa7c01d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mos=C3=A8=20Giordano?= <765740+giordano@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:58:27 +0100 Subject: [PATCH 12/12] Remove trailing blank line --- test/test_amdgpu.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl index 91d667da96..7a26b3b61a 100644 --- a/test/test_amdgpu.jl +++ b/test/test_amdgpu.jl @@ -113,4 +113,3 @@ end build_and_time_step_simulation(model) end end -