Add AMDGPU extension (#95)

* Add AMDGPU extension

* Update Project.toml [skip ci]

Project.toml

@@ -1,7 +1,7 @@
 name = "PencilArrays"
 uuid = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e"
 authors = ["Juan Ignacio Polanco <[email protected]> and contributors"]
-version = "0.19.4"
+version = "0.19.5"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -19,15 +19,18 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 VersionParsing = "81def892-9a0e-5fdd-b105-ffc91e053289"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 
 [extensions]
+PencilArraysAMDGPUExt = ["AMDGPU"]
 PencilArraysDiffEqExt = ["DiffEqBase"]
 PencilArraysHDF5Ext = ["HDF5"]
 
 [compat]
 Adapt = "3, 4"
+AMDGPU = "0.8, 0.9"
 DiffEqBase = "6"
 HDF5 = "0.16, 0.17"
 JSON3 = "1.4"

ext/PencilArraysAMDGPUExt.jl

@@ -0,0 +1,22 @@
+module PencilArraysAMDGPUExt
+
+using PencilArrays: typeof_array, typeof_ptr
+using PencilArrays.Transpositions: Transpositions
+using AMDGPU: ROCVector
+
+# Workaround `unsafe_wrap` not allowing the `own` keyword argument in the AMDGPU
+# implementation.
+# Moreover, one needs to set the `lock = false` argument to indicate that we want to wrap an
+# array which is already in the GPU.
+function Transpositions.unsafe_as_array(::Type{T}, x::ROCVector{UInt8}, dims::Tuple) where {T}
+    p = typeof_ptr(x){T}(pointer(x))
+    unsafe_wrap(typeof_array(x), p, dims; lock = false)
+end
+
+# Workaround `unsafe_wrap` for ROCArrays not providing a definition for dims::Integer.
+# We convert that argument to a tuple, which is accepted by the implementation in AMDGPU.
+function Transpositions.unsafe_as_array(::Type{T}, x::ROCVector{UInt8}, N::Integer) where {T}
+    Transpositions.unsafe_as_array(T, x, (N,))
+end
+
+end