Attempt to fix GPU implementation

Author: lkdvos

ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl

@@ -4,7 +4,7 @@ using MatrixAlgebraKit
 using MatrixAlgebraKit: @algdef, Algorithm, check_input
 using MatrixAlgebraKit: one!, zero!, uppertriangular!, lowertriangular!
 using MatrixAlgebraKit: diagview, sign_safe
-using MatrixAlgebraKit: LQViaTransposedQR
+using MatrixAlgebraKit: LQViaTransposedQR, TruncationByValue
 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!
@@ -40,4 +40,17 @@ _gpu_heevj!(A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwar
 _gpu_heevd!(A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwargs...) = YArocSOLVER.heevd!(A, Dd, V; kwargs...)
 _gpu_heev!(A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwargs...)  = YArocSOLVER.heev!(A, Dd, V; kwargs...)
 _gpu_heevx!(A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwargs...) = YArocSOLVER.heevx!(A, Dd, V; kwargs...)
+
+function MatrixAlgebraKit.findtruncated_sorted(values::StridedROCVector, strategy::TruncationByValue)
+    atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p))
+    @assert strategy.by === abs || strategy.by === real "sorting strategy incompatible with implementation"
+    if strategy.rev
+        i = @something findfirst(<(atol) ∘ strategy.by, values) lastindex(values) + 1
+        return i:length(values)
+    else
+        i = @something findlast(>(atol) ∘ strategy.by, values) firstindex(values) - 1
+        return 1:i
+    end
+end
+
 end

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -4,7 +4,7 @@ using MatrixAlgebraKit
 using MatrixAlgebraKit: @algdef, Algorithm, check_input
 using MatrixAlgebraKit: one!, zero!, uppertriangular!, lowertriangular!
 using MatrixAlgebraKit: diagview, sign_safe
-using MatrixAlgebraKit: LQViaTransposedQR
+using MatrixAlgebraKit: LQViaTransposedQR, TruncationByValue
 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!
@@ -44,4 +44,16 @@ _gpu_gesvdj!(A::StridedCuMatrix, S::StridedCuVector, U::StridedCuMatrix, Vᴴ::S
 _gpu_heevj!(A::StridedCuMatrix, Dd::StridedCuVector, V::StridedCuMatrix; kwargs...) = YACUSOLVER.heevj!(A, Dd, V; kwargs...)
 _gpu_heevd!(A::StridedCuMatrix, Dd::StridedCuVector, V::StridedCuMatrix; kwargs...) = YACUSOLVER.heevd!(A, Dd, V; kwargs...)
 
+function MatrixAlgebraKit.findtruncated_sorted(values::StridedCuVector, strategy::TruncationByValue)
+    atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p))
+    @assert strategy.by === abs || strategy.by === real "sorting strategy incompatible with implementation"
+    if strategy.rev
+        i = @something findfirst(<(atol) ∘ strategy.by, values) lastindex(values) + 1
+        return i:length(values)
+    else
+        i = @something findlast(>(atol) ∘ strategy.by, values) firstindex(values) - 1
+        return 1:i
+    end
+end
+
 end