specialize ipiv2perm for CPUStorage

JuliaGPU · Mar 11, 2024 · 5ee17c9 · 5ee17c9
1 parent 69fa85d
commit 5ee17c9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lib/mps/linalg.jl b/lib/mps/linalg.jl
@@ -156,7 +156,8 @@ end
 
 # Metal's pivoting sequence needs to be iterated sequentially...
 # TODO: figure out a GPU-compatible way to get the permutation matrix
-LinearAlgebra.ipiv2perm(v::MtlVector{T}, maxi::Integer) where T = LinearAlgebra.ipiv2perm(Array(v), maxi)
+LinearAlgebra.ipiv2perm(v::MtlVector{T,S}, maxi::Integer) where {T,S} = LinearAlgebra.ipiv2perm(Array(v), maxi)
+LinearAlgebra.ipiv2perm(v::MtlVector{T,S}, maxi::Integer) where {T,S<:CPUStorage} = LinearAlgebra.ipiv2perm(unsafe_wrap(Array, v), maxi)
 
 function LinearAlgebra.lu(A::MtlMatrix{T}; check::Bool = true) where {T<:MtlFloat}
     M,N = size(A)