Improve truncate and CheckSVDDone

ITensor · Oct 23, 2023 · f56a3b3 · f56a3b3
1 parent 28b9e5a
commit f56a3b3
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 21 deletions.
diff --git a/NDTensors/src/blocksparse/linearalgebra.jl b/NDTensors/src/blocksparse/linearalgebra.jl
@@ -76,7 +76,7 @@ function svd(T::BlockSparseMatrix{ElT}; kwargs...) where {ElT}
 
   dropblocks = Int[]
   if truncate
-    truncerr, docut = truncate!(d; kwargs...)
+    d, truncerr, docut = truncate!!(d; kwargs...)
     for n in 1:nnzblocks(T)
       blockdim = _truncated_blockdim(
         Ss[n], docut; min_blockdim, singular_values=true, truncate
@@ -237,7 +237,7 @@ function eigen(
   sort!(d; rev=true, by=abs)
 
   if truncate
-    truncerr, docut = truncate!(d; kwargs...)
+    d, truncerr, docut = truncate!!(d; kwargs...)
     for n in 1:nnzblocks(T)
       blockdim = _truncated_blockdim(Ds[n], docut)
       if blockdim == 0

diff --git a/NDTensors/src/linearalgebra/linearalgebra.jl b/NDTensors/src/linearalgebra/linearalgebra.jl
@@ -169,11 +169,9 @@ function svd(T::DenseTensor{ElT,2,IndsT}; kwargs...) where {ElT,IndsT}
 
   P = MS .^ 2
   if truncate
-    P_cpu = NDTensors.cpu(P)
-    truncerr, _ = truncate!(
-      P_cpu; mindim, maxdim, cutoff, use_absolute_cutoff, use_relative_cutoff, kwargs...
+    P, truncerr, _ = truncate!!(
+      P; mindim, maxdim, cutoff, use_absolute_cutoff, use_relative_cutoff, kwargs...
     )
-    P = adapt(typeof(P), P_cpu)
   else
     truncerr = 0.0
   end
@@ -240,11 +238,9 @@ function eigen(
   VM = VM[:, p]
 
   if truncate
-    DM_cpu = NDTensors.cpu(DM)
-    truncerr, _ = truncate!(
-      DM_cpu; mindim, maxdim, cutoff, use_absolute_cutoff, use_relative_cutoff, kwargs...
+    DM, truncerr, _ = truncate!!(
+      DM; mindim, maxdim, cutoff, use_absolute_cutoff, use_relative_cutoff, kwargs...
     )
-    DM = adapt(typeof(DM), DM_cpu)
     dD = length(DM)
     if dD < size(VM, 2)
       VM = VM[:, 1:dD]
@@ -359,11 +355,9 @@ function eigen(
   #VM = VM[:,p]
 
   if truncate
-    DM_cpu = NDTensors.cpu(DM)
-    truncerr, _ = truncate!(
-      DM_cpu; maxdim, cutoff, use_absolute_cutoff, use_relative_cutoff, kwargs...
+    DM, truncerr, _ = truncate!!(
+      DM; maxdim, cutoff, use_absolute_cutoff, use_relative_cutoff, kwargs...
     )
-    DM = adapt(typeof(DM), DM_cpu)
     dD = length(DM)
     if dD < size(VM, 2)
       VM = VM[:, 1:dD]

diff --git a/NDTensors/src/linearalgebra/svd.jl b/NDTensors/src/linearalgebra/svd.jl
@@ -1,8 +1,9 @@
-
 function checkSVDDone(S::AbstractArray, thresh::Float64)
-  # Convert to CPU to avoid slow scalar indexing
-  # on GPU.
-  S = cpu(S)
+  return checkSVDDone(leaf_parenttype(S), S, thresh)
+end
+
+# CPU version.
+function checkSVDDone(::Type{<:Array}, S::AbstractArray, thresh::Float64)
   N = length(S)
   (N <= 1 || thresh < 0.0) && return (true, 1)
   S1t = S[1] * thresh
@@ -17,6 +18,12 @@ function checkSVDDone(S::AbstractArray, thresh::Float64)
   return (false, start)
 end
 
+# Convert to CPU to avoid slow scalar indexing
+# on GPU.
+function checkSVDDone(::Type{<:AbstractArray}, S::AbstractArray, thresh::Float64)
+  return checkSVDDone(Array, cpu(S), thresh)
+end
+
 function svd_recursive(M::AbstractMatrix; thresh::Float64=1E-3, north_pass::Int=2)
   Mr, Mc = size(M)
   if Mr > Mc

diff --git a/NDTensors/src/truncate.jl b/NDTensors/src/truncate.jl
@@ -1,7 +1,25 @@
 export truncate!
 
-function truncate!(P::AbstractVector{ElT}; kwargs...)::Tuple{ElT,ElT} where {ElT}
-  cutoff::Union{Nothing,ElT} = get(kwargs, :cutoff, zero(ElT))
+function truncate!!(P::AbstractArray; kwargs...)
+  return truncate!!(leaf_parenttype(P), P; kwargs...)
+end
+
+# CPU version.
+function truncate!!(::Type{<:Array}, P::AbstractArray; kwargs...)
+  P, truncerr, docut = truncate!(P; kwargs...)
+  return P, truncerr, docut
+end
+
+# GPU fallback version, convert to CPU.
+function truncate!!(::Type{<:AbstractArray}, P::AbstractArray; kwargs...)
+  P_cpu = cpu(P)
+  P_cpu, truncerr, docut = truncate!(P_cpu; kwargs...)
+  P = adapt(leaf_parenttype(P), P_cpu)
+  return P, truncerr, docut
+end
+
+# CPU implementation.
+function truncate!(P::AbstractVector{ElT}; cutoff=zero(eltype(P)), kwargs...) where {ElT}
   if isnothing(cutoff)
     cutoff = typemin(ElT)
   end
@@ -92,5 +110,5 @@ function truncate!(P::AbstractVector{ElT}; kwargs...)::Tuple{ElT,ElT} where {ElT
   s < 0 && (P .*= s)
   resize!(P, n)
 
-  return truncerr, docut
+  return P, truncerr, docut
 end