diff --git a/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl b/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl index a54d2490d..6f0f3edad 100644 --- a/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl +++ b/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 diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 5b8e43a97..368e5abe2 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/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)) diff --git a/src/tensors/linalg.jl b/src/tensors/linalg.jl index 549595596..3ea59a153 100644 --- a/src/tensors/linalg.jl +++ b/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. diff --git a/test/amd/tensors.jl b/test/amd/tensors.jl index d0c9ca014..14b4dc62d 100644 --- a/test/amd/tensors.jl +++ b/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)) diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 750d7384a..0d2b866f9 100644 --- a/test/cuda/tensors.jl +++ b/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))