Add extensions for AMDGPU.jl and oneAPI.jl (#4381)

* Add an extension for AMDGPU.jl

* Add an extension for oneAPI.jl

* Setup CI for the extensions of AMDGPU.jl and oneAPI.jl

Author: amontoison

.buildkite/pipeline.yml

@@ -65,6 +65,36 @@ steps:
           limit: 1
     depends_on: "init"
 
+  # - label: "AMDGPU.jl unit tests"
+  #   env:
+  #     TEST_GROUP: "amdgpu"
+  #   plugins:
+  #     - JuliaCI/julia#v1:
+  #         version: "1.10"
+  #   agents:
+  #     queue: "juliagpu"
+  #     rocm: "*"
+  #     rocmgpu: "*"
+  #   env:
+  #     JULIA_NUM_THREADS: 4
+  #     JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+  #     JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+  #     JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
+  #   command: |
+  #     julia --color=yes --project -e 'using Pkg; Pkg.test()'"
+
+  # - label: "oneAPI.jl unit tests"
+  #   env:
+  #     TEST_GROUP: "oneapi"
+  #   plugins:
+  #     - JuliaCI/julia#v1:
+  #         version: "1.10"
+  #   agents:
+  #     queue: "juliagpu"
+  #     intel: "*"
+  #   command: |
+  #     julia --color=yes --project -e 'using Pkg; Pkg.test()'"
+
 #####
 ##### Solver tests
 #####

Project.toml

@@ -37,22 +37,27 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
-Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]
+OceananigansAMDGPUExt = "AMDGPU"
 OceananigansEnzymeExt = "Enzyme"
 OceananigansMakieExt = ["MakieCore", "Makie"]
 OceananigansMetalExt = "Metal"
 OceananigansNCDatasetsExt = "NCDatasets"
+OceananigansOneAPIExt = "oneAPI"
 OceananigansReactantExt = ["Reactant", "KernelAbstractions", "ConstructionBase"]
 
 [compat]
+AMDGPU = "1.2.7"
 Adapt = "4.1.1"
 CUDA = "5.7"
 ConstructionBase = "1"
@@ -81,8 +86,8 @@ OffsetArrays = "1.4"
 OrderedCollections = "1.1"
 Printf = "1.9"
 Random = "1.9"
-ReactantCore = "0.1"
 Reactant = "0.2.63"
+ReactantCore = "0.1"
 Rotations = "1.0"
 SeawaterPolynomials = "0.3.9"
 SparseArrays = "1.9"
@@ -91,17 +96,20 @@ Statistics = "1.9"
 StructArrays = "0.4, 0.5, 0.6, 0.7"
 TimesDates = "0.3"
 julia = "1.9"
+oneAPI = "2.0.1"
 
 [extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
 
 [targets]
-test = ["DataDeps", "SafeTestsets", "Test", "Enzyme", "Reactant", "Metal", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates", "NCDatasets"]
+test = ["AMDGPU", "oneAPI", "DataDeps", "SafeTestsets", "Test", "Enzyme", "Reactant", "Metal", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates", "NCDatasets"]

docs/src/grids.md

@@ -65,9 +65,8 @@ grid = RectilinearGrid(architecture,
 └── Bounded  z ∈ [0.0, 10.0] variably spaced with min(Δz)=1.0, max(Δz)=4.0
 ```
 
-!!! note "GPU architecture requires a CUDA- or Metal-enabled device"
-    To run the above example and create a grid on the GPU, either an Nvidia or Metal GPU has to be available.
-    For more information about CUDA, see the [`CUDA.jl` documentation](https://cuda.juliagpu.org/stable/).
+!!! note "GPU architecture requires a CUDA-, ROC- or Metal-enabled device"
+    Running this example and creating a grid on the GPU requires a device that supports CUDA, ROC, or Metal. In other words, you must have an Nvidia, AMD, or Apple-compatible GPU available. For more information about CUDA, see the [`CUDA.jl` documentation](https://cuda.juliagpu.org/stable/).
 
 The ``y``-dimension is "missing" because it's marked `Flat` in `topology = (Periodic, Flat, Bounded)`.
 So nothing varies in ``y``: `y`-derivatives are 0.

ext/OceananigansAMDGPUExt.jl

@@ -0,0 +1,30 @@
+module OceananigansAMDGPUExt
+
+using AMDGPU
+using Oceananigans
+
+import Oceananigans.Architectures:
+    architecture,
+    convert_to_device,
+    on_architecture
+
+const ROCGPU = GPU{<:AMDGPU.ROCBackend}
+ROCGPU() = GPU(Metal.AMDGPU.ROCBackend())
+
+architecture(::ROCArray) = ROCGPU()
+Base.summary(::ROCGPU) = "ROCGPU"
+
+on_architecture(::ROCGPU, a::Number) = a
+on_architecture(::ROCGPU, a::Array) = ROCArray(a)
+on_architecture(::ROCGPU, a::BitArray) = ROCArray(a)
+on_architecture(::ROCGPU, a::SubArray{<:Any, <:Any, <:Array}) = ROCArray(a)
+on_architecture(::CPU, a::ROCArray) = Array(a)
+on_architecture(::CPU, a::SubArray{<:Any, <:Any, <:ROCArray}) = Array(a)
+on_architecture(::ROCGPU, a::ROCArray) = a
+on_architecture(::ROCGPU, a::SubArray{<:Any, <:Any, <:ROCArray}) = a
+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)
+
+end # module

ext/OceananigansMetalExt.jl

@@ -10,6 +10,7 @@ import Oceananigans.Architectures:
 
 const MetalGPU = GPU{<:Metal.MetalBackend}
 MetalGPU() = GPU(Metal.MetalBackend())
+Base.summary(::MetalGPU) = "MetalGPU"
 
 architecture(::MtlArray) = MetalGPU()
 

ext/OceananigansOneAPIExt.jl

@@ -0,0 +1,30 @@
+module OceananigansOneAPIExt
+
+using oneAPI
+using Oceananigans
+
+import Oceananigans.Architectures:
+    architecture,
+    convert_to_device,
+    on_architecture
+
+const ONEGPU = GPU{<:oneAPI.oneAPIBackend}
+ONEGPU() = GPU(oneAPI.oneAPIBackend())
+
+architecture(::oneArray) = ONEGPU()
+Base.summary(::ONEGPU) = "ONEGPU"
+
+on_architecture(::ONEGPU, a::Number) = a
+on_architecture(::ONEGPU, a::Array) = oneArray(a)
+on_architecture(::ONEGPU, a::BitArray) = oneArray(a)
+on_architecture(::ONEGPU, a::SubArray{<:Any, <:Any, <:Array}) = oneArray(a)
+on_architecture(::CPU, a::oneArray) = Array(a)
+on_architecture(::CPU, a::SubArray{<:Any, <:Any, <:oneArray}) = Array(a)
+on_architecture(::ONEGPU, a::oneArray) = a
+on_architecture(::ONEGPU, a::SubArray{<:Any, <:Any, <:oneArray}) = a
+on_architecture(::ONEGPU, a::StepRangeLen) = a
+
+@inline convert_to_device(::ONEGPU, args) = oneAPI.kernel_convert(args)
+@inline convert_to_device(::ONEGPU, args::Tuple) = map(oneAPI.kernel_convert, args)
+
+end # module

test/runtests.jl

@@ -221,7 +221,7 @@ CUDA.allowscalar() do
         end
     end
 
-    # Tests for Metal extension
+    # Tests for MPI extension
     if group == :mpi_tripolar || group == :all
         @testset "Distributed tripolar tests" begin
             include("test_mpi_tripolar.jl")
@@ -257,12 +257,25 @@ CUDA.allowscalar() do
     end
 
     # Tests for Metal extension
-    if group == :metal|| group == :all
+    if group == :metal || group == :all
         @testset "Metal extension tests" begin
             include("test_metal.jl")
         end
     end
 
+    # Tests for AMDGPU extension
+    if group == :amdgpu || group == :all
+        @testset "AMDGPU extension tests" begin
+            include("test_amdgpu.jl")
+        end
+    end
+    # Tests for oneAPI extension
+    if group == :oneapi || group == :all
+        @testset "oneAPI extension tests" begin
+            include("test_oneapi.jl")
+        end
+    end
+
     if group == :convergence
         include("test_convergence.jl")
     end

test/test_amdgpu.jl

@@ -0,0 +1,33 @@
+include("dependencies_for_runtests.jl")
+
+using AMDGPU
+
+@testset "AMDGPU extension" begin
+    roc = AMDGPU.ROCBackend()
+    arch = GPU(roc)
+    grid = RectilinearGrid(arch, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16))
+
+    @test parent(grid.xᶠᵃᵃ) isa ROCArray
+    @test parent(grid.xᶜᵃᵃ) isa ROCArray
+    @test eltype(grid) == Float64
+    @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))
+
+    @test parent(model.velocities.u) isa ROCArray
+    @test parent(model.velocities.v) isa ROCArray
+    @test parent(model.velocities.w) isa ROCArray
+    @test parent(model.tracers.b) isa ROCArray
+
+    simulation = Simulation(model, Δt=1minute, stop_iteration=3)
+    run!(simulation)
+
+    @test iteration(simulation) == 3
+    @test time(simulation) == 3minutes
+end

test/test_init.jl

@@ -1,6 +1,8 @@
 using Reactant
 using Enzyme
 using Metal
+using AMDGPU
+using oneAPI
 
 Pkg.instantiate(; verbose=true)
 Pkg.precompile(; strict=true)

test/test_oneapi.jl

@@ -0,0 +1,33 @@
+include("dependencies_for_runtests.jl")
+
+using oneAPI
+
+@testset "oneAPI extension" begin
+    oneapi = oneAPI.oneAPIBackend()
+    arch = GPU(oneapi)
+    grid = RectilinearGrid(arch, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16))
+
+    @test parent(grid.xᶠᵃᵃ) isa oneArray
+    @test parent(grid.xᶜᵃᵃ) isa oneArray
+    @test eltype(grid) == Float64
+    @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))
+
+    @test parent(model.velocities.u) isa oneArray
+    @test parent(model.velocities.v) isa oneArray
+    @test parent(model.velocities.w) isa oneArray
+    @test parent(model.tracers.b) isa oneArray
+
+    simulation = Simulation(model, Δt=1minute, stop_iteration=3)
+    run!(simulation)
+
+    @test iteration(simulation) == 3
+    @test time(simulation) == 3minutes
+end