Merge pull request #479 from vchuravy/add-amdgpu-extension

Add AMDGPU.jl extension

Author: ChrisRackauckas

Project.toml

@@ -7,6 +7,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -21,6 +22,7 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [extensions]
+ArrayInterfaceAMDGPUExt = "AMDGPU"
 ArrayInterfaceBandedMatricesExt = "BandedMatrices"
 ArrayInterfaceBlockBandedMatricesExt = "BlockBandedMatrices"
 ArrayInterfaceCUDAExt = "CUDA"
@@ -35,6 +37,7 @@ ArrayInterfaceStaticArraysCoreExt = "StaticArraysCore"
 ArrayInterfaceTrackerExt = "Tracker"
 
 [compat]
+AMDGPU = "2"
 Adapt = "4"
 BandedMatrices = "1"
 BlockBandedMatrices = "0.13"

ext/ArrayInterfaceAMDGPUExt.jl

@@ -0,0 +1,15 @@
+module ArrayInterfaceAMDGPUExt
+
+using ArrayInterface
+using AMDGPU
+using LinearAlgebra
+
+function ArrayInterface.lu_instance(A::ROCMatrix{T}) where {T}
+    ipiv = ROCVector{Cint}(undef, 0)
+    info = zero(Int)
+    return LinearAlgebra.LU(similar(A, 0, 0), ipiv, info)
+end
+
+ArrayInterface.device(::Type{<:AMDGPU.ROCArray}) = ArrayInterface.GPU()
+
+end # module

test/gpu/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"

test/gpu/amdgpu.jl

@@ -0,0 +1,13 @@
+using AMDGPU
+using ArrayInterface
+using LinearAlgebra
+
+using Test
+
+A = ROCMatrix(Float32[1 0; 0 1])
+
+# Test that lu_instance works with AMDGPU.jl ROC arrays
+@test isa(ArrayInterface.lu_instance(A), LU)
+
+# Test that device returns GPU()
+@test ArrayInterface.device(typeof(A)) == ArrayInterface.GPU()