Keep normalizations of GPUTensorMaps from synchronizing

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -19,4 +19,11 @@ using Random
 
 include("cutensormap.jl")
 
+function TensorKit._nrm2!(
+        out::CuVector{R}, x::cuBLAS.StridedCuVecOrDenseMat{T}
+    ) where {R <: Union{Float32, Float64}, T <: Union{R, Complex{R}}}
+    cuBLAS.nrm2(x, out)
+    return out
+end
+
 end

ext/TensorKitGPUArraysExt.jl

@@ -1,6 +1,7 @@
 module TensorKitGPUArraysExt
 
 using GPUArrays
+using LinearAlgebra: LinearAlgebra
 using GPUArrays: @allowscalar
 using GPUArrays.KernelAbstractions: @kernel, @index, get_backend
 using Adapt
@@ -666,6 +667,46 @@ TensorKit._transform_subblocks(
 const GPUDiagonalTensorMap = DiagonalTensorMap{<:Any, <:Any, <:AnyGPUVector}
 const GPUTensorMap = TensorMap{<:Any, <:Any, <:Any, <:Any, <:AnyGPUVector}
 
+# `norm` returns a host scalar by default, which forces the GPU to wait and synchronize.
+# Here we add some methods for various `p` that avoid copying back to the host if we
+# don't need to, e.g. for normalization. For the 2-norm, the norm is instead kept in a
+# 1-element device array and broadcast into the data. This matches the vendor device-side
+# nrm2 versions.
+function TensorKit._nrm2!(out::AnyGPUVector, x::AnyGPUArray)
+    out .= sqrt.(vec(sum(abs2, x; dims = ntuple(identity, ndims(x)))))
+    return out
+end
+TensorKit._nrm2!(out::AnyGPUVector, x::LinearAlgebra.Diagonal) = TensorKit._nrm2!(out, x.diag)
+
+function _device_norm2(t::Union{GPUTensorMap, GPUDiagonalTensorMap})
+    InnerProductStyle(t) === EuclideanInnerProduct() || TensorKit.throw_invalid_innerproduct(:norm)
+    out = similar(t.data, float(real(scalartype(t))), 1)
+    FusionStyle(sectortype(t)) isa UniqueFusion && return TensorKit._nrm2!(out, t.data)
+    fill!(out, zero(eltype(out)))
+    nb = similar(out)
+    for (c, b) in blocks(t)
+        out .+= dim(c) .* abs2.(TensorKit._nrm2!(nb, b))
+    end
+    return out .= sqrt.(out)
+end
+
+function LinearAlgebra.normalize!(t::Union{GPUTensorMap, GPUDiagonalTensorMap}, p::Real = 2)
+    p == 2 || return scale!(t, inv(norm(t, p)))
+    t.data ./= _device_norm2(t)
+    return t
+end
+function LinearAlgebra.normalize(t::Union{GPUTensorMap, GPUDiagonalTensorMap}, p::Real = 2)
+    p == 2 || return scale(t, inv(norm(t, p)))
+    tdst = similar(t, float(scalartype(t)))
+    tdst.data .= t.data ./ _device_norm2(t)
+    return tdst
+end
+function LinearAlgebra.normalize!(t::AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:GPUTensorMap}, p::Real = 2)
+    normalize!(parent(t), p)
+    return t
+end
+LinearAlgebra.normalize(t::AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:GPUTensorMap}, p::Real = 2) = adjoint(normalize(parent(t), p))
+
 # position of each subblock in the data of a `DiagonalTensorMap` in the same space
 _diagonal_offsets(structure) = first(_work_offsets(first(first(s)) for s in structure))
 

src/tensors/linalg.jl

@@ -18,6 +18,14 @@ Base.:\(α::Number, t::AbstractTensorMap) = *(t, one(scalartype(t)) / α)
 LinearAlgebra.normalize!(t::AbstractTensorMap, p::Real = 2) = scale!(t, inv(norm(t, p)))
 LinearAlgebra.normalize(t::AbstractTensorMap, p::Real = 2) = scale(t, inv(norm(t, p)))
 
+"""
+    _nrm2!(out, x)
+
+Write the 2-norm of the array `x` into the 1-element array `out`, without copying it to the host.
+Used to normalize GPU-backed tensors without synchronizing; the GPU extensions provide methods.
+"""
+function _nrm2! end
+
 # destination allocation for matrix multiplication
 # note that we don't fall back to `tensoralloc_contract` since that needs to account for
 # permutations, which might require complex scalartypes even if the inputs are real.

test/amd/tensors.jl

@@ -153,6 +153,17 @@ for V in spacelist
                 p = 3 * rand(Float64)
                 @test norm(t + t, p) ≈ 2 * norm(t, p)
                 @test norm(t) ≈ norm(t')
+                for p in (1, 2, 3 * rand(Float64), Inf)
+                    tn = @constinferred normalize(t, p)
+                    @test storagetype(tn) == storagetype(t)
+                    @test norm(tn, p) ≈ 1
+                    @test tn ≈ t / norm(t, p)
+                    @test normalize!(copy(t), p) ≈ tn
+                    @test normalize(t', p) ≈ tn'
+                    ta = copy(t)'
+                    @test normalize!(ta, p) === ta
+                    @test ta ≈ tn'
+                end
 
                 t2 = @constinferred rand!(similar(t))
                 β = rand(T)
@@ -566,6 +577,13 @@ end
         dcpu = adapt(Array, d)
         scalars = ((TensorKit.One(), TensorKit.Zero()), (T(2), T(3)))
 
+        for p in (1, 2, 3 * rand(Float64), Inf)
+            dn = @constinferred normalize(d, p)
+            @test dn isa typeof(d)
+            @test adapt(Array, dn) ≈ normalize(dcpu, p)
+            @test adapt(Array, normalize!(copy(d), p)) ≈ normalize(dcpu, p)
+        end
+
         for p in (((), (1, 2)), ((1, 2), ()), ((2, 1), ()), ((), (2, 1)), ((1,), (2,)), ((2,), (1,))), conjd in (false, true)
             W = TO.tensoradd_structure(dcpu, p, conjd)
             dst = adapt(ROCVector{T}, randn(T, W))

test/cuda/tensors.jl

@@ -172,6 +172,17 @@ for V in spacelist
                 p = 3 * rand(Float64)
                 @test norm(t + t, p) ≈ 2 * norm(t, p)
                 @test norm(t) ≈ norm(t')
+                for p in (1, 2, 3 * rand(Float64), Inf)
+                    tn = @constinferred normalize(t, p)
+                    @test storagetype(tn) == storagetype(t)
+                    @test norm(tn, p) ≈ 1
+                    @test tn ≈ t / norm(t, p)
+                    @test normalize!(copy(t), p) ≈ tn
+                    @test normalize(t', p) ≈ tn'
+                    ta = copy(t)'
+                    @test normalize!(ta, p) === ta
+                    @test ta ≈ tn'
+                end
 
                 t2 = @constinferred rand!(similar(t))
                 β = rand(T)
@@ -639,6 +650,13 @@ end
         dcpu = adapt(Array, d)
         scalars = ((TensorKit.One(), TensorKit.Zero()), (T(2), T(3)))
 
+        for p in (1, 2, 3 * rand(Float64), Inf)
+            dn = @constinferred normalize(d, p)
+            @test dn isa typeof(d)
+            @test adapt(Array, dn) ≈ normalize(dcpu, p)
+            @test adapt(Array, normalize!(copy(d), p)) ≈ normalize(dcpu, p)
+        end
+
         for p in (((), (1, 2)), ((1, 2), ()), ((2, 1), ()), ((), (2, 1)), ((1,), (2,)), ((2,), (1,))), conjd in (false, true)
             W = TO.tensoradd_structure(dcpu, p, conjd)
             dst = adapt(CuVector{T}, randn(T, W))