GPU-friendly truncation implementations (#349)

* try to make truncation GPU-friendly

* Temporarily fix StridedViews version

* Revert "Temporarily fix StridedViews version"

This reverts commit 77f0ffa.

* Small update for diagonal pullbacks

* Fix last error

* Reenable truncated CUDA tests

* make truncation run on GPU

* bypass scalar indexing by specializing

* convenience overloads

* gpu-friendly copies

* retain storagetype in extended_S

* avoid GPU issues with truncated adjoint tensormaps

* various utility improvements

* complete rewrite of implementation

* GPU doesn't like `trues`

* remove CUDA specializations and temporarily add missing MatrixAlgebraKit thingies

* better dimension testing

* fix unbound type parameter

* add missing import

* be careful about double method definitions

* disable diagonal test

* bump MatrixAlgebraKit dependency

* Revert "disable diagonal test"

This reverts commit f26cffe.

* remove unnecessary specializations

* specialize CPU implementations

* add explanation TruncationByOrder

* add explanation and specialization TruncationByError

* fix stupidity

* fix views

* enfore positive and finite p-norms

---------

Co-authored-by: Katharine Hyatt <kslimes@gmail.com>

Author: lkdvos

Project.toml

@@ -44,7 +44,7 @@ FiniteDifferences = "0.12"
 GPUArrays = "11.3.1"
 LRUCache = "1.0.2"
 LinearAlgebra = "1"
-MatrixAlgebraKit = "0.6.2"
+MatrixAlgebraKit = "0.6.3"
 Mooncake = "0.4.183"
 OhMyThreads = "0.8.0"
 Printf = "1"

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -17,5 +17,6 @@ using TensorKit: MatrixAlgebraKit
 using Random
 
 include("cutensormap.jl")
+include("truncation.jl")
 
 end

ext/TensorKitCUDAExt/truncation.jl

@@ -0,0 +1,52 @@
+const CuSectorVector{T, I} = TensorKit.SectorVector{T, I, <:CuVector{T}}
+
+function MatrixAlgebraKit.findtruncated(
+        values::CuSectorVector, strategy::MatrixAlgebraKit.TruncationByOrder
+    )
+    I = sectortype(values)
+
+    dims = similar(values, Base.promote_op(dim, I))
+    for (c, v) in pairs(dims)
+        fill!(v, dim(c))
+    end
+
+    perm = sortperm(parent(values); strategy.by, strategy.rev)
+    cumulative_dim = cumsum(Base.permute!(parent(dims), perm))
+
+    result = similar(values, Bool)
+    parent(result)[perm] .= cumulative_dim .<= strategy.howmany
+    return result
+end
+
+function MatrixAlgebraKit.findtruncated(
+        values::CuSectorVector, strategy::MatrixAlgebraKit.TruncationByError
+    )
+    (isfinite(strategy.p) && strategy.p > 0) ||
+        throw(ArgumentError(lazy"p-norm with p = $(strategy.p) is currently not supported."))
+    ϵᵖmax = max(strategy.atol^strategy.p, strategy.rtol^strategy.p * norm(values, strategy.p))
+    ϵᵖ = similar(values, typeof(ϵᵖmax))
+
+    # dimensions are all 1 so no need to account for weight
+    if FusionStyle(sectortype(values)) isa UniqueFusion
+        parent(ϵᵖ) .= abs.(parent(values)) .^ strategy.p
+    else
+        for (c, v) in pairs(values)
+            v′ = ϵᵖ[c]
+            v′ .= abs.(v) .^ strategy.p .* dim(c)
+        end
+    end
+
+    perm = sortperm(parent(values); by = abs, rev = false)
+    cumulative_err = cumsum(Base.permute!(parent(ϵᵖ), perm))
+
+    result = similar(values, Bool)
+    parent(result)[perm] .= cumulative_err .> ϵᵖmax
+    return result
+end
+
+# Needed until MatrixAlgebraKit patch hits...
+function MatrixAlgebraKit._ind_intersect(A::CuVector{Bool}, B::CuVector{Int})
+    result = fill!(similar(A), false)
+    result[B] .= @view A[B]
+    return result
+end

src/factorizations/adjoint.jl

@@ -7,6 +7,7 @@ _adjoint(alg::MAK.LAPACK_HouseholderLQ) = MAK.LAPACK_HouseholderQR(; alg.kwargs.
 _adjoint(alg::MAK.LAPACK_HouseholderQL) = MAK.LAPACK_HouseholderRQ(; alg.kwargs...)
 _adjoint(alg::MAK.LAPACK_HouseholderRQ) = MAK.LAPACK_HouseholderQL(; alg.kwargs...)
 _adjoint(alg::MAK.PolarViaSVD) = MAK.PolarViaSVD(_adjoint(alg.svd_alg))
+_adjoint(alg::TruncatedAlgorithm) = TruncatedAlgorithm(_adjoint(alg.alg), alg.trunc)
 _adjoint(alg::AbstractAlgorithm) = alg
 
 _adjoint(alg::MAK.CUSOLVER_HouseholderQR) = MAK.LQViaTransposedQR(alg)
@@ -81,7 +82,7 @@ for (left_f, right_f) in zip(
 end
 
 # 3-arg functions
-for f in (:svd_full, :svd_compact)
+for f in (:svd_full, :svd_compact, :svd_trunc)
     f! = Symbol(f, :!)
     @eval function MAK.copy_input(::typeof($f), t::AdjointTensorMap)
         return adjoint(MAK.copy_input($f, adjoint(t)))
@@ -93,9 +94,16 @@ for f in (:svd_full, :svd_compact)
         return reverse(adjoint.(MAK.initialize_output($f!, adjoint(t), _adjoint(alg))))
     end
 
-    @eval function MAK.$f!(t::AdjointTensorMap, F, alg::AbstractAlgorithm)
-        F′ = $f!(adjoint(t), reverse(adjoint.(F)), _adjoint(alg))
-        return reverse(adjoint.(F′))
+    if f === :svd_trunc
+        function MAK.svd_trunc!(t::AdjointTensorMap, F, alg::AbstractAlgorithm)
+            U, S, Vᴴ, ϵ = svd_trunc!(adjoint(t), reverse(adjoint.(F)), _adjoint(alg))
+            return Vᴴ', S, U', ϵ
+        end
+    else
+        @eval function MAK.$f!(t::AdjointTensorMap, F, alg::AbstractAlgorithm)
+            F′ = $f!(adjoint(t), reverse(adjoint.(F)), _adjoint(alg))
+            return reverse(adjoint.(F′))
+        end
     end
 
     # disambiguate by prohibition
@@ -111,6 +119,15 @@ function MAK.svd_compact!(t::AdjointTensorMap, F, alg::DiagonalAlgorithm)
     F′ = svd_compact!(adjoint(t), reverse(adjoint.(F)), _adjoint(alg))
     return reverse(adjoint.(F′))
 end
+function MAK.initialize_output(
+        ::typeof(svd_trunc!), t::AdjointTensorMap, alg::TruncatedAlgorithm
+    )
+    return reverse(adjoint.(MAK.initialize_output(svd_trunc!, adjoint(t), _adjoint(alg))))
+end
+function MAK.svd_trunc!(t::AdjointTensorMap, F, alg::TruncatedAlgorithm)
+    U, S, Vᴴ, ϵ = svd_trunc!(adjoint(t), reverse(adjoint.(F)), _adjoint(alg))
+    return Vᴴ', S, U', ϵ
+end
 
 function LinearAlgebra.isposdef(t::AdjointTensorMap)
     return isposdef(adjoint(t))

src/factorizations/diagonal.jl

@@ -13,26 +13,6 @@ for f in (
     @eval MAK.copy_input(::typeof($f), d::DiagonalTensorMap) = copy(d)
 end
 
-for f! in (:eig_full!, :eig_trunc!)
-    @eval function MAK.initialize_output(
-            ::typeof($f!), d::AbstractTensorMap, ::DiagonalAlgorithm
-        )
-        return d, similar(d)
-    end
-end
-
-for f! in (:eigh_full!, :eigh_trunc!)
-    @eval function MAK.initialize_output(
-            ::typeof($f!), d::AbstractTensorMap, ::DiagonalAlgorithm
-        )
-        if scalartype(d) <: Real
-            return d, similar(d, space(d))
-        else
-            return similar(d, real(scalartype(d))), similar(d, space(d))
-        end
-    end
-end
-
 for f! in (:qr_full!, :qr_compact!)
     @eval function MAK.initialize_output(
             ::typeof($f!), d::AbstractTensorMap, ::DiagonalAlgorithm
@@ -93,7 +73,7 @@ end
 # For diagonal inputs we don't have to promote the scalartype since we know they are symmetric
 function MAK.initialize_output(::typeof(eig_vals!), t::AbstractTensorMap, alg::DiagonalAlgorithm)
     V_D = fuse(domain(t))
-    Tc = scalartype(t)
+    Tc = complex(scalartype(t))
     A = similarstoragetype(t, Tc)
     return SectorVector{Tc, sectortype(t), A}(undef, V_D)
 end