From 6375089832cf9ad9e740704adaa647f283ffae21 Mon Sep 17 00:00:00 2001 From: Juan Ignacio Polanco Date: Fri, 14 Jun 2024 17:48:14 +0200 Subject: [PATCH 1/2] Add AMDGPU extension --- Project.toml | 3 +++ ext/PencilArraysAMDGPUExt.jl | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+) create mode 100644 ext/PencilArraysAMDGPUExt.jl diff --git a/Project.toml b/Project.toml index aa294af7..09f4dda3 100644 --- a/Project.toml +++ b/Project.toml @@ -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] PencilArraysDiffEqExt = ["DiffEqBase"] PencilArraysHDF5Ext = ["HDF5"] +PencilArraysAMDGPUExt = ["AMDGPU"] [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 From 7d7a26579764eecf93c6a2957d4384197d96070a Mon Sep 17 00:00:00 2001 From: Juan Ignacio Polanco Date: Fri, 14 Jun 2024 18:08:40 +0200 Subject: [PATCH 2/2] Update Project.toml [skip ci] --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 09f4dda3..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" @@ -24,9 +24,9 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" [extensions] +PencilArraysAMDGPUExt = ["AMDGPU"] PencilArraysDiffEqExt = ["DiffEqBase"] PencilArraysHDF5Ext = ["HDF5"] -PencilArraysAMDGPUExt = ["AMDGPU"] [compat] Adapt = "3, 4"