Add JLArrays, Metal and AMDGPU extensions and S kwarg tests (rebased)

.buildkite/pipeline.yml

@@ -11,5 +11,24 @@ steps:
       using Pkg
       Pkg.add("CUDA")
       Pkg.instantiate()
+      include("test/gpu/test_S_kwarg.jl")
       include("test/gpu/nvidia.jl")'
     timeout_in_minutes: 30
+
+  - label: "AMD GPUs -- LinearOperators.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    command: |
+      julia --color=yes --project -e '
+      using Pkg
+      Pkg.add("AMDGPU")
+      Pkg.instantiate()
+      include("test/gpu/test_S_kwarg.jl")
+      include("test/gpu/amdgpu.jl")'
+    timeout_in_minutes: 30
+

Project.toml

@@ -11,28 +11,40 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 
 [extensions]
+LinearOperatorsAMDGPUExt = "AMDGPU" 
 LinearOperatorsChainRulesCoreExt = "ChainRulesCore"
 LinearOperatorsCUDAExt = "CUDA"
 LinearOperatorsLDLFactorizationsExt = "LDLFactorizations"
+LinearOperatorsJLArraysExt = "JLArrays"
+LinearOperatorsMetalExt = "Metal"
 
 [compat]
+AMDGPU = "0.6, 0.7, 0.8, 0.9"
 ChainRulesCore = "1"
 CUDA = "4, 5"
 FastClosures = "0.2, 0.3"
+JLArrays = "0.1"
 LDLFactorizations = "0.9, 0.10"
 LinearAlgebra = "1"
+Metal = "1.1"
 Printf = "1"
 Requires = "1"
 SparseArrays = "1"
 TimerOutputs = "^0.5"
 julia = "^1.6.0"
 
 [extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"

ext/LinearOperatorsAMDGPUExt.jl

@@ -0,0 +1,8 @@
+module LinearOperatorsAMDGPUExt
+
+using LinearOperators
+isdefined(Base, :get_extension) ? (using AMDGPU) : (using ..AMDGPU)
+
+LinearOperators.storage_type(::ROCArray{T, 2, B}) where {T, B} = ROCArray{T, 1, B}
+
+end # module

ext/LinearOperatorsJLArraysExt.jl

@@ -0,0 +1,8 @@
+module LinearOperatorsJLArraysExt
+
+using LinearOperators
+isdefined(Base, :get_extension) ? (using JLArrays) : (using ..JLArrays)
+
+LinearOperators.storage_type(::JLArray{T, 2}) where {T} = JLArray{T, 1}
+
+end # module

ext/LinearOperatorsMetalExt.jl

@@ -0,0 +1,8 @@
+module LinearOperatorsMetalExt
+
+using LinearOperators
+isdefined(Base, :get_extension) ? (using Metal) : (using ..Metal)
+
+LinearOperators.storage_type(::MtlArray{T, 2, S}) where {T, S} = MtlArray{T, 1, S}
+
+end # module

src/LinearOperators.jl

@@ -34,6 +34,9 @@ end
 
 @static if !isdefined(Base, :get_extension)
   function __init__()
+    Requires.@require AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" begin
+      include("../ext/LinearOperatorsAMDGPUExt.jl")
+    end
     Requires.@require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" begin
       include("../ext/LinearOperatorsChainRulesCoreExt.jl")
     end
@@ -43,6 +46,12 @@ end
     Requires.@require LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" begin
       include("../ext/LinearOperatorsLDLFactorizationsExt.jl")
     end
+    Requires.@require JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" begin
+      include("../ext/LinearOperatorsJLArraysExt.jl")
+    end
+    Requires.@require Metal = "dde4c033-4e86-420c-a63e-0dd931031962" begin
+      include("../ext/LinearOperatorsMetalExt.jl")
+    end
   end
 end
 

src/special-operators.jl

@@ -68,7 +68,7 @@ Change `S` to use LinearOperators on GPU.
 """
 function opEye(T::DataType, nrow::I, ncol::I; S = Vector{T}) where {I <: Integer}
   if nrow == ncol
-    return opEye(T, nrow)
+    return opEye(T, nrow; S = S)
   end
   prod! = @closure (res, v, α, β) -> mulOpEye!(res, v, α, β, min(nrow, ncol))
   return LinearOperator{T}(nrow, ncol, false, false, prod!, prod!, prod!, S = S)

test/Project.toml

@@ -2,8 +2,10 @@
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -16,7 +18,9 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 Arpack = "0.5"
 ChainRulesCore = "1"
 FastClosures = "0.2, 0.3"
+JLArrays = "0.1"
 LDLFactorizations = "0.9, 0.10"
+Metal = "1.1"
 Requires = "1"
 Test = "1"
 TestSetExtensions = "3"

test/gpu/amdgpu.jl

@@ -0,0 +1,15 @@
+using Test, LinearAlgebra, SparseArrays
+using LinearOperators, AMDGPU
+
+@testset "AMDGPU -- AMDGPU.jl" begin
+  A = ROCArray(rand(Float32, 5, 5))
+  B = ROCArray(rand(Float32, 10, 10))
+  C = ROCArray(rand(Float32, 20, 20))
+  M = BlockDiagonalOperator(A, B, C)
+
+  v = ROCArray(rand(Float32, 35))
+  y = M * v
+  @test y isa ROCArray{Float32}
+
+  @testset "AMDGPU S kwarg" test_S_kwarg(arrayType = ROCArray)
+end

test/gpu/jlarrays.jl

@@ -0,0 +1 @@
+test_S_kwarg(arrayType = JLArray)

test/gpu/metal.jl

@@ -0,0 +1,2 @@
+using Metal
+test_S_kwarg(arrayType = MtlArray, notMetal = false)

test/gpu/nvidia.jl

@@ -13,4 +13,5 @@ using LinearOperators, CUDA, CUDA.CUSPARSE, CUDA.CUSOLVER
   v = CUDA.rand(35)
   y = M * v
   @test y isa CuVector{Float32}
+  @testset "Nvidia S kwarg" test_S_kwarg(arrayType = CuArray)
 end

test/gpu/test_S_kwarg.jl

@@ -0,0 +1,42 @@
+using Test, LinearOperators, LinearAlgebra
+
+function test_S_kwarg(; arrayType, notMetal = true)
+  mat = arrayType(rand(Float32, 32, 32))
+  vec = arrayType(rand(Float32, 32))
+  vecT = typeof(vec)
+
+  if notMetal
+    # To test operators which can derive a default storage_type from their arguments
+    vecTother = typeof(arrayType(rand(Float32, 32)))
+  end
+
+  @testset "S Kwarg with arrayType $(arrayType)" begin
+    @test vecT == LinearOperators.storage_type(mat)
+
+    # constructors.jl
+    @test LinearOperators.storage_type(LinearOperator(mat)) == LinearOperators.storage_type(mat) # default
+    notMetal && @test LinearOperators.storage_type(LinearOperator(mat; S = vecTother)) == vecTother
+    @test LinearOperators.storage_type(LinearOperator(Symmetric(mat); S = vecT)) == vecT
+    #notMetal && @test LinearOperators.storage_type(LinearOperator(SymTridiagonal(Symmetric(mat)); S = vecT)) == vecT
+    @test LinearOperators.storage_type(LinearOperator(Hermitian(mat); S = vecT)) == vecT
+    @test LinearOperators.storage_type(LinearOperator(Float32, 32, 32, true, true, () -> 0; S = vecT)) == vecT
+
+    # special-operators.jl
+    @test LinearOperators.storage_type(opEye(Float32, 32; S = vecT)) == vecT
+    @test LinearOperators.storage_type(opEye(Float32, 16, 32; S = vecT)) == vecT
+    @test LinearOperators.storage_type(opEye(Float32, 32, 32; S = vecT)) == vecT
+
+    @test LinearOperators.storage_type(opOnes(Float32, 32, 32; S = vecT)) == vecT
+    @test LinearOperators.storage_type(opZeros(Float32, 32, 32; S = vecT)) == vecT
+
+    @test LinearOperators.storage_type(opDiagonal(vec)) == vecT
+    @test LinearOperators.storage_type(opDiagonal(32, 32, vec)) == vecT
+
+    @test LinearOperators.storage_type(opRestriction([1, 2, 3], 32; S = vecT)) == vecT
+    @test LinearOperators.storage_type(opExtension([1, 2, 3], 32; S = vecT)) == vecT
+
+    @test LinearOperators.storage_type(BlockDiagonalOperator(mat, mat)) == vecT # default
+    notMetal && @test LinearOperators.storage_type(BlockDiagonalOperator(mat, mat; S = vecTother)) == vecTother
+  end
+
+end

test/runtests.jl

@@ -1,5 +1,5 @@
 using Arpack, Test, TestSetExtensions, LinearOperators
-using LinearAlgebra, LDLFactorizations, SparseArrays
+using LinearAlgebra, LDLFactorizations, SparseArrays, JLArrays
 using Zygote
 include("test_aux.jl")
 
@@ -16,3 +16,8 @@ include("test_normest.jl")
 include("test_diag.jl")
 include("test_chainrules.jl")
 include("test_solve_shifted_system.jl")
+include("gpu/test_S_kwarg.jl")
+include("gpu/jlarrays.jl")
+if Sys.isapple() && occursin("arm64", Sys.MACHINE)
+    include("gpu/metal.jl")
+end