diff --git a/Project.toml b/Project.toml index aa294af7..911f3983 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "PencilArrays" uuid = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e" authors = ["Juan Ignacio Polanco 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" diff --git a/ext/PencilArraysAMDGPUExt.jl b/ext/PencilArraysAMDGPUExt.jl new file mode 100644 index 00000000..eb223143 --- /dev/null +++ b/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