Working GPU col permute

Author: kshyatt

ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl

@@ -206,4 +206,11 @@ function MatrixAlgebraKit._mul_herm!(C::StridedROCMatrix{T}, A::StridedROCMatrix
     return C
 end
 
+function MatrixAlgebraKit.permute_V_cols!(V, I::ROCVector{Int})
+    I_ixs = ROCArray(collect(1:size(V, 1)))
+    c_ixs = map(CartesianIndex, I, I_ixs)
+    V[c_ixs] .= one(eltype(V))
+    return V
+end
+
 end

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -191,4 +191,11 @@ function MatrixAlgebraKit._mul_herm!(C::StridedCuMatrix{T}, A::StridedCuMatrix{T
     return C
 end
 
+function MatrixAlgebraKit.permute_V_cols!(V, I::CuVector{Int})
+    I_ixs = CuArray(collect(1:size(V, 1)))
+    c_ixs = map(CartesianIndex, I, I_ixs)
+    V[c_ixs] .= one(eltype(V))
+    return V
+end
+
 end

src/implementations/eigh.jl

@@ -141,6 +141,8 @@ function eigh_trunc_no_error!(A, DV, alg::TruncatedAlgorithm)
     return DVtrunc
 end
 
+permute_V_cols!(V, I::Vector{Int}) = Base.permutecols!!(V, I)
+
 # Diagonal logic
 # --------------
 function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm)
@@ -153,8 +155,7 @@ function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm)
         diagview(D) .= real.(diagview(A))[I]
     end
     zero!(V)
-    Is = [CartesianIndex(ix, I[ix]) for ix in 1:size(A, 1)]
-    V[Is] .= one(eltype(A))
+    V = permute_V_cols!(V, I)
     return D, V
 end