Updates for TensorKit compatibility

Author: kshyatt

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -10,24 +10,42 @@ import MatrixAlgebraKit: _gpu_geqrf!, _gpu_ungqr!, _gpu_unmqr!, _gpu_gesvd!, _gp
 import MatrixAlgebraKit: _gpu_heevj!, _gpu_heevd!
 using CUDA
 using LinearAlgebra
-using LinearAlgebra: BlasFloat
+using LinearAlgebra: BlasFloat, eigvals!
 
 include("yacusolver.jl")
 
-function MatrixAlgebraKit.default_qr_algorithm(::Type{T}; kwargs...) where {T<:StridedCuMatrix}
+function MatrixAlgebraKit.default_qr_algorithm(::Type{T}; kwargs...) where {TT<:BlasFloat, T<:StridedCuMatrix{TT}}
     return CUSOLVER_HouseholderQR(; kwargs...)
 end
-function MatrixAlgebraKit.default_lq_algorithm(::Type{T}; kwargs...) where {T<:StridedCuMatrix}
+function MatrixAlgebraKit.default_lq_algorithm(::Type{T}; kwargs...) where {TT<:BlasFloat, T<:StridedCuMatrix{TT}}
     qr_alg = CUSOLVER_HouseholderQR(; kwargs...)
     return LQViaTransposedQR(qr_alg)
 end
-function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T<:StridedCuMatrix}
+function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {TT<:BlasFloat, T<:StridedCuMatrix{TT}}
     return CUSOLVER_QRIteration(; kwargs...)
 end
-function MatrixAlgebraKit.default_eig_algorithm(::Type{T}; kwargs...) where {T<:StridedCuMatrix}
+function MatrixAlgebraKit.default_eig_algorithm(::Type{T}; kwargs...) where {TT<:BlasFloat, T<:StridedCuMatrix{TT}}
     return CUSOLVER_Simple(; kwargs...)
 end
-function MatrixAlgebraKit.default_eigh_algorithm(::Type{T}; kwargs...) where {T<:StridedCuMatrix}
+function MatrixAlgebraKit.default_eigh_algorithm(::Type{T}; kwargs...) where {TT<:BlasFloat, T<:StridedCuMatrix{TT}}
+    return CUSOLVER_DivideAndConquer(; kwargs...)
+end
+
+# include for block sector support
+function MatrixAlgebraKit.default_qr_algorithm(::Type{Base.ReshapedArray{T,2,SubArray{T,1,A,Tuple{UnitRange{Int}},true},Tuple{}}}; kwargs...) where {T<:BlasFloat, A<:CuVecOrMat{T}}
+    return CUSOLVER_HouseholderQR(; kwargs...)
+end
+function MatrixAlgebraKit.default_lq_algorithm(::Type{Base.ReshapedArray{T,2,SubArray{T,1,A,Tuple{UnitRange{Int}},true},Tuple{}}}; kwargs...) where {T<:BlasFloat, A<:CuVecOrMat{T}}
+    qr_alg = CUSOLVER_HouseholderQR(; kwargs...)
+    return LQViaTransposedQR(qr_alg)
+end
+function MatrixAlgebraKit.default_svd_algorithm(::Type{Base.ReshapedArray{T,2,SubArray{T,1,A,Tuple{UnitRange{Int}},true},Tuple{}}}; kwargs...) where {T<:BlasFloat, A<:CuVecOrMat{T}}
+    return CUSOLVER_QRIteration(; kwargs...)
+end
+function MatrixAlgebraKit.default_eig_algorithm(::Type{Base.ReshapedArray{T,2,SubArray{T,1,A,Tuple{UnitRange{Int}},true},Tuple{}}}; kwargs...) where {T<:BlasFloat, A<:CuVecOrMat{T}}
+    return CUSOLVER_Simple(; kwargs...)
+end
+function MatrixAlgebraKit.default_eigh_algorithm(::Type{Base.ReshapedArray{T,2,SubArray{T,1,A,Tuple{UnitRange{Int}},true},Tuple{}}}; kwargs...) where {T<:BlasFloat, A<:CuVecOrMat{T}}
     return CUSOLVER_DivideAndConquer(; kwargs...)
 end
 

ext/MatrixAlgebraKitCUDAExt/yacusolver.jl

@@ -28,7 +28,7 @@ for (bname, fname, elty, relty) in
         #! format: on
             chkstride1(A, U, Vᴴ, S)
             m, n = size(A)
-            (m < n) && throw(ArgumentError("CUSOLVER's gesvd requires m ≥ n"))
+            (m < n) && throw(ArgumentError("CUSOLVER's gesvd requires m ($m) ≥ n ($n)"))
             minmn = min(m, n)
             if length(U) == 0
                 jobu = 'N'

src/implementations/eig.jl

@@ -102,8 +102,9 @@ function eig_vals!(A::AbstractMatrix, D, alg::GPU_EigAlgorithm)
     check_input(eig_vals!, A, D, alg)
     V  = similar(A, complex(eltype(A)), (size(A, 1), 0))
     if alg isa GPU_Simple
-        isempty(alg.kwargs) ||
-            throw(ArgumentError("LAPACK_Simple (geev) does not accept any keyword arguments"))
+        # TODO filter out nothing kwargs
+        #isempty(alg.kwargs) ||
+        #    throw(ArgumentError("GPU_Simple (geev) does not accept any keyword arguments"))
         _gpu_geev!(A, D, V)
     end
     return D