Piracy fix

Author: kshyatt

ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl

@@ -7,7 +7,7 @@ using MatrixAlgebraKit: diagview, sign_safe
 using MatrixAlgebraKit: LQViaTransposedQR, TruncationStrategy, NoTruncation, TruncationByValue, AbstractAlgorithm
 using MatrixAlgebraKit: default_qr_algorithm, default_lq_algorithm, default_svd_algorithm, default_eigh_algorithm
 import MatrixAlgebraKit: _gpu_geqrf!, _gpu_ungqr!, _gpu_unmqr!, _gpu_gesvd!, _gpu_Xgesvdp!, _gpu_gesvdj!
-import MatrixAlgebraKit: _gpu_heevj!, _gpu_heevd!, _gpu_heev!, _gpu_heevx!
+import MatrixAlgebraKit: _gpu_heevj!, _gpu_heevd!, _gpu_heev!, _gpu_heevx!, _sylvester
 using AMDGPU
 using LinearAlgebra
 using LinearAlgebra: BlasFloat
@@ -171,4 +171,9 @@ end
 MatrixAlgebraKit._ind_intersect(A::ROCVector{Int}, B::ROCVector{Int}) =
     MatrixAlgebraKit._ind_intersect(collect(A), collect(B))
 
+function _sylvester(A::AnyROCMatrix, B::AnyROCMatrix, C::AnyROCMatrix)
+    hX = sylvester(collect(A), collect(B), collect(C))
+    return ROCArray(hX)
+end
+
 end

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -7,7 +7,7 @@ using MatrixAlgebraKit: diagview, sign_safe, default_pullback_gauge_atol, defaul
 using MatrixAlgebraKit: LQViaTransposedQR, TruncationByValue, AbstractAlgorithm
 using MatrixAlgebraKit: default_qr_algorithm, default_lq_algorithm, default_svd_algorithm, default_eig_algorithm, default_eigh_algorithm
 import MatrixAlgebraKit: _gpu_geqrf!, _gpu_ungqr!, _gpu_unmqr!, _gpu_gesvd!, _gpu_Xgesvdp!, _gpu_Xgesvdr!, _gpu_gesvdj!, _gpu_geev!
-import MatrixAlgebraKit: _gpu_heevj!, _gpu_heevd!
+import MatrixAlgebraKit: _gpu_heevj!, _gpu_heevd!, _sylvester
 using CUDA, CUDA.CUBLAS
 using CUDA: i32
 using LinearAlgebra
@@ -202,14 +202,9 @@ function MatrixAlgebraKit.default_pullback_gauge_atol(A::AnyCuArray, As...)
     return isempty(As′) ? 0 : eps(norm(CuArray.(As′), Inf))^(3 / 4)
 end
 
-function LinearAlgebra.sylvester(A::AnyCuMatrix, B::AnyCuMatrix, C::AnyCuMatrix)
-    #=m = size(A, 1)
-    n = size(B, 2)
-    I_n = fill!(similar(A, n), one(eltype(A)))
-    I_m = fill!(similar(B, m), one(eltype(B)))
-    L   = kron(diagm(I_n), A) + kron(adjoint(B), diagm(I_m))
-    x_vec = L \ -vec(C)
-    X = CuMatrix(reshape(x_vec, m, n))=#
+function _sylvester(A::AnyCuMatrix, B::AnyCuMatrix, C::AnyCuMatrix)
+    # https://github.com/JuliaGPU/CUDA.jl/issues/3021
+    # to add native sylvester to CUDA
     hX = sylvester(collect(A), collect(B), collect(C))
     return CuArray(hX)
 end

src/common/pullbacks.jl

@@ -10,3 +10,6 @@ function iszerotangent end
 
 iszerotangent(::Any) = false
 iszerotangent(::Nothing) = true
+
+# fallback
+_sylvester(A, B, C) = LinearAlgebra.sylvester(A, B, C)

src/pullbacks/eig.jl

@@ -140,7 +140,7 @@ function eig_trunc_pullback!(
     # add contribution from orthogonal complement
     PA = A - (A * V) / V
     Y = mul!(ΔVperp, PA', Z, 1, 1)
-    X = sylvester(PA', -Dmat', Y)
+    X = _sylvester(PA', -Dmat', Y)
     Z .+= X
 
     if eltype(ΔA) <: Real

src/pullbacks/eigh.jl

@@ -142,7 +142,7 @@ function eigh_trunc_pullback!(
         # add contribution from orthogonal complement
         W = qr_null(V)
         WᴴΔV = W' * ΔV
-        X = sylvester(W' * A * W, -Dmat, WᴴΔV)
+        X = _sylvester(W' * A * W, -Dmat, WᴴΔV)
         Z = mul!(Z, W, X, 1, 1)
 
         # put everything together: symmetrize for hermitian case

src/pullbacks/polar.jl

@@ -16,7 +16,7 @@ function left_polar_pullback!(ΔA::AbstractMatrix, A, WP, ΔWP; kwargs...)
     M = zero(P)
     !iszerotangent(ΔW) && mul!(M, W', ΔW, 1, 1)
     !iszerotangent(ΔP) && mul!(M, ΔP, P, -1, 1)
-    C = sylvester(P, P, M' - M)
+    C = _sylvester(P, P, M' - M)
     C .+= ΔP
     ΔA = mul!(ΔA, W, C, 1, 1)
     if !iszerotangent(ΔW)
@@ -52,7 +52,7 @@ function right_polar_pullback!(ΔA::AbstractMatrix, A, PWᴴ, ΔPWᴴ; kwargs...
     M = zero(P)
     !iszerotangent(ΔWᴴ) && mul!(M, ΔWᴴ, Wᴴ', 1, 1)
     !iszerotangent(ΔP) && mul!(M, P, ΔP, -1, 1)
-    C = sylvester(P, P, M' - M)
+    C = _sylvester(P, P, M' - M)
     C .+= ΔP
     ΔA = mul!(ΔA, C, Wᴴ, 1, 1)
     if !iszerotangent(ΔWᴴ)

src/pullbacks/svd.jl

@@ -207,7 +207,7 @@ function svd_trunc_pullback!(
     else
         fill!(view(rhs, m̃ .+ (1:ñ), :), 0)
     end
-    XY = sylvester(ÃÃ, -Smat, rhs)
+    XY = _sylvester(ÃÃ, -Smat, rhs)
     X = view(XY, 1:m̃, :)
     Y = view(XY, m̃ .+ (1:ñ), :)
     ΔA = mul!(ΔA, Ũ, X * Vᴴ, 1, 1)