GPU-friendly truncation implementations

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,26 @@
+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
+
+# 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

src/factorizations/truncation.jl

@@ -29,31 +29,36 @@ end
 # ---------
 _blocklength(d::Integer, ind) = _blocklength(Base.OneTo(d), ind)
 _blocklength(ax, ind) = length(ax[ind])
+_blocklength(ax::Base.OneTo, ind::AbstractVector{<:Integer}) = length(ind)
+_blocklength(ax::Base.OneTo, ind::AbstractVector{Bool}) = count(ind)
+
 function truncate_space(V::ElementarySpace, inds)
-    return spacetype(V)(c => _blocklength(dim(V, c), ind) for (c, ind) in inds)
+    return spacetype(V)(c => _blocklength(dim(V, c), ind) for (c, ind) in pairs(inds))
 end
 
 function truncate_domain!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds)
     for (c, b) in blocks(tdst)
         I = get(inds, c, nothing)
         @assert !isnothing(I)
-        copy!(b, view(block(tsrc, c), :, I))
+        b′ = block(tsrc, c)
+        b .= b′[:, I]
     end
     return tdst
 end
 function truncate_codomain!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds)
     for (c, b) in blocks(tdst)
         I = get(inds, c, nothing)
         @assert !isnothing(I)
-        copy!(b, view(block(tsrc, c), I, :))
+        b′ = block(tsrc, c)
+        b .= b′[I, :]
     end
     return tdst
 end
 function truncate_diagonal!(Ddst::DiagonalTensorMap, Dsrc::DiagonalTensorMap, inds)
     for (c, b) in blocks(Ddst)
         I = get(inds, c, nothing)
         @assert !isnothing(I)
-        copy!(diagview(b), view(diagview(block(Dsrc, c)), I))
+        diagview(b) .= @view diagview(block(Dsrc, c))[I]
     end
     return Ddst
 end
@@ -78,7 +83,7 @@ end
 function MAK.truncate(
         ::typeof(left_null!), (U, S)::NTuple{2, AbstractTensorMap}, strategy::TruncationStrategy
     )
-    extended_S = zerovector!(SectorVector{eltype(S)}(undef, fuse(codomain(U))))
+    extended_S = zerovector!(SectorVector{eltype(S), sectortype(S), storagetype(S)}(undef, fuse(codomain(U))))
     for (c, b) in blocks(S)
         copyto!(extended_S[c], diagview(b)) # copyto! since `b` might be shorter
     end
@@ -91,7 +96,7 @@ end
 function MAK.truncate(
         ::typeof(right_null!), (S, Vᴴ)::NTuple{2, AbstractTensorMap}, strategy::TruncationStrategy
     )
-    extended_S = zerovector!(SectorVector{eltype(S)}(undef, fuse(domain(Vᴴ))))
+    extended_S = zerovector!(SectorVector{eltype(S), sectortype(S), storagetype(S)}(undef, fuse(domain(Vᴴ))))
     for (c, b) in blocks(S)
         copyto!(extended_S[c], diagview(b)) # copyto! since `b` might be shorter
     end
@@ -142,57 +147,11 @@ for f! in (:eig_trunc!, :eigh_trunc!)
     end
 end
 
-# Find truncation
-# ---------------
+# findtruncated
+# -------------
 # auxiliary functions
 rtol_to_atol(S, p, atol, rtol) = rtol == 0 ? atol : max(atol, norm(S, p) * rtol)
 
-function _compute_truncerr(Σdata, truncdim, p = 2)
-    I = keytype(Σdata)
-    S = scalartype(valtype(Σdata))
-    return TensorKit._norm(
-        (c => @view(v[(get(truncdim, c, 0) + 1):end]) for (c, v) in Σdata),
-        p, zero(S)
-    )
-end
-
-function _findnexttruncvalue(
-        S, truncdim::SectorDict{I, Int}; by = identity, rev::Bool = true
-    ) where {I <: Sector}
-    # early return
-    (isempty(S) || all(iszero, values(truncdim))) && return nothing
-    if rev
-        σmin, imin = findmin(keys(truncdim)) do c
-            d = truncdim[c]
-            return by(S[c][d])
-        end
-        return σmin, keys(truncdim)[imin]
-    else
-        σmax, imax = findmax(keys(truncdim)) do c
-            d = truncdim[c]
-            return by(S[c][d])
-        end
-        return σmax, keys(truncdim)[imax]
-    end
-end
-
-function _sort_and_perm(values::SectorVector; by = identity, rev::Bool = false)
-    values_sorted = similar(values)
-    perms = SectorDict(
-        (
-                begin
-                    p = sortperm(v; by, rev)
-                    vs = values_sorted[c]
-                    vs .= view(v, p)
-                    c => p
-                end
-            ) for (c, v) in pairs(values)
-    )
-    return values_sorted, perms
-end
-
-# findtruncated
-# -------------
 # Generic fallback
 function MAK.findtruncated_svd(values::SectorVector, strategy::TruncationStrategy)
     return MAK.findtruncated(values, strategy)
@@ -203,24 +162,38 @@ function MAK.findtruncated(values::SectorVector, ::NoTruncation)
 end
 
 function MAK.findtruncated(values::SectorVector, strategy::TruncationByOrder)
-    values_sorted, perms = _sort_and_perm(values; strategy.by, strategy.rev)
-    inds = MAK.findtruncated_svd(values_sorted, truncrank(strategy.howmany))
-    return SectorDict(c => perms[c][I] for (c, I) in inds)
-end
-function MAK.findtruncated_svd(values::SectorVector, strategy::TruncationByOrder)
-    I = keytype(values)
-    truncdim = SectorDict{I, Int}(c => length(d) for (c, d) in pairs(values))
-    totaldim = sum(dim(c) * d for (c, d) in truncdim; init = 0)
-    while totaldim > strategy.howmany
-        next = _findnexttruncvalue(values, truncdim; strategy.by, strategy.rev)
-        isnothing(next) && break
-        _, cmin = next
-        truncdim[cmin] -= 1
-        totaldim -= dim(cmin)
-        truncdim[cmin] == 0 && delete!(truncdim, cmin)
+    I = sectortype(values)
+
+
+    if FusionStyle(I) isa UniqueFusion # dimensions are all 1
+        perm = partialsortperm(parent(values), 1:strategy.howmany; strategy.by, strategy.rev)
+        result = similar(values, Bool)
+        fill!(parent(result), false)
+        parent(result)[perm] .= true
+        return result
+    end
+
+    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)
+    result = similar(values, Bool)
+    fill!(parent(result), false)
+
+    totaldim = 0
+    for i in perm
+        totaldim += dims[i]
+        totaldim > strategy.howmany && break
+        result[i] = true
     end
-    return SectorDict(c => Base.OneTo(d) for (c, d) in truncdim)
+
+    return result
 end
+# disambiguate
+MAK.findtruncated_svd(values::SectorVector, strategy::TruncationByOrder) =
+    MAK.findtruncated(values, strategy)
 
 function MAK.findtruncated(values::SectorVector, strategy::TruncationByFilter)
     return SectorDict(c => findall(strategy.filter, d) for (c, d) in pairs(values))
@@ -237,28 +210,29 @@ function MAK.findtruncated_svd(values::SectorVector, strategy::TruncationByValue
     return SectorDict(c => MAK.findtruncated_svd(d, strategy′) for (c, d) in pairs(values))
 end
 
-function MAK.findtruncated(values::SectorVector, strategy::TruncationByError)
-    values_sorted, perms = _sort_and_perm(values; strategy.by, strategy.rev)
-    inds = MAK.findtruncated_svd(values_sorted, truncrank(strategy.howmany))
-    return SectorDict(c => perms[c][I] for (c, I) in inds)
-end
-function MAK.findtruncated_svd(values::SectorVector, strategy::TruncationByError)
-    I = keytype(values)
-    truncdim = SectorDict{I, Int}(c => length(d) for (c, d) in pairs(values))
-    by(c, v) = abs(v)^strategy.p * dim(c)
-    Nᵖ = sum(((c, v),) -> sum(Base.Fix1(by, c), v), pairs(values))
-    ϵᵖ = max(strategy.atol^strategy.p, strategy.rtol^strategy.p * Nᵖ)
-    truncerrᵖ = zero(real(scalartype(valtype(values))))
-    next = _findnexttruncvalue(values, truncdim)
-    while !isnothing(next)
-        σmin, cmin = next
-        truncerrᵖ += by(cmin, σmin)
-        truncerrᵖ >= ϵᵖ && break
-        (truncdim[cmin] -= 1) == 0 && delete!(truncdim, cmin)
-        next = _findnexttruncvalue(values, truncdim)
+function MAK.findtruncated(values::SectorVector, strategy::MAK.TruncationByError)
+    ϵᵖmax = max(strategy.atol^strategy.p, strategy.rtol^strategy.p * norm(values, strategy.p))
+    ϵᵖ = similar(values, typeof(ϵᵖmax))
+
+    if FusionStyle(sectortype(values)) isa UniqueFusion # dimensions are all 1
+        parent(ϵᵖ) .= abs.(parent(values)) .^ strategy.p
+    else
+        for (c, v) in pairs(values)
+            v′ = ϵᵖ[c]
+            v′ .= abs.(v) .^ strategy.p .* dim(c)
+        end
     end
-    return SectorDict{I, Base.OneTo{Int}}(c => Base.OneTo(d) for (c, d) in truncdim)
+
+    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
+# disambiguate
+MAK.findtruncated_svd(values::SectorVector, strategy::TruncationByError) =
+    MAK.findtruncated(values, strategy)
 
 function MAK.findtruncated(values::SectorVector, strategy::TruncationSpace)
     blockstrategy(c) = truncrank(dim(strategy.space, c); strategy.by, strategy.rev)
@@ -273,17 +247,15 @@ function MAK.findtruncated(values::SectorVector, strategy::TruncationIntersectio
     inds = map(Base.Fix1(MAK.findtruncated, values), strategy.components)
     return SectorDict(
         c => mapreduce(
-                Base.Fix2(getindex, c), MatrixAlgebraKit._ind_intersect, inds;
-                init = trues(length(values[c]))
+                Base.Fix2(getindex, c), MatrixAlgebraKit._ind_intersect, inds
             ) for c in intersect(map(keys, inds)...)
     )
 end
 function MAK.findtruncated_svd(values::SectorVector, strategy::TruncationIntersection)
     inds = map(Base.Fix1(MAK.findtruncated_svd, values), strategy.components)
     return SectorDict(
         c => mapreduce(
-                Base.Fix2(getindex, c), MatrixAlgebraKit._ind_intersect, inds;
-                init = trues(length(values[c]))
+                Base.Fix2(getindex, c), MatrixAlgebraKit._ind_intersect, inds
             ) for c in intersect(map(keys, inds)...)
     )
 end
@@ -293,7 +265,7 @@ end
 MAK.truncation_error(values::SectorVector, ind) = MAK.truncation_error!(copy(values), ind)
 
 function MAK.truncation_error!(values::SectorVector, ind)
-    for (c, ind_c) in ind
+    for (c, ind_c) in pairs(ind)
         v = values[c]
         v[ind_c] .= zero(eltype(v))
     end