diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 45f6df640..0408980dd 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -9,6 +9,7 @@ using Strided: StridedViews using MatrixAlgebraKit, Adapt using TensorKit using TensorKit.TensorOperations: linearize, DefaultAllocator +using TensorKit.VectorInterface: One using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check @@ -295,19 +296,25 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREEBRAIDER_CACHE") => GLOBAL_DIAGONAL_TREEBRAIDER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE") => GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE, ) ) # device-caches should probably have smaller sizes: resize!(GLOBAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_DIAGONAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) return nothing end # Allocations must outlive any GPUArrays allocation-cache scope around the operation. -function device_transformer(A, transformer::UniqueTreeTransformer, p) - packed = DeviceUniqueTreeTransformer(transformer, p) +device_transformer(A, transformer::UniqueTreeTransformer, p) = + device_transformer(A, DeviceUniqueTreeTransformer(transformer, p)) + +function device_transformer(A, packed::DeviceUniqueTreeTransformer) return GPUArrays.@uncached begin proto = A(undef, 0) # Metadata and coefficients have different element types; preserve each on upload. @@ -453,6 +460,91 @@ TensorKit._transform_subblocks( transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer} ) = (tdst.data, tsrc.data) +# Diagonal sources +# ---------------- +# A `DiagonalTensorMap` only stores the diagonal of each of its subblocks, and an index +# manipulation maps the entry `(i, i)` of a source subblock onto the entry `(i, i)` of a destination +# subblock. Every block is therefore a one-dimensional strided copy, which the unique kernel +# handles without needing a dense copy of the source or destination. + +const GPUDiagonalTensorMap = DiagonalTensorMap{<:Any, <:Any, <:AnyGPUVector} +const GPUTensorMap = TensorMap{<:Any, <:Any, <:Any, <:Any, <:AnyGPUVector} + +# 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)) + +_diagonal_entries(transformer::UniqueTreeTransformer) = transformer.data +# `V ← V` and any rearrangement of its two indices have a single tree for every subblock +_diagonal_entries(transformer::GenericTreeTransformer) = + ((only(U), only(inds_dst), only(inds_src)) for (U, inds_dst, inds_src) in transformer.data) + +function diagonal_device_transformer(A, transformer, dstdiagonal::Bool) + (; structure_dst, structure_src) = transformer + offsets_src = _diagonal_offsets(structure_src) + offsets_dst = dstdiagonal ? _diagonal_offsets(structure_dst) : nothing + blocks = map(_diagonal_entries(transformer)) do (coeff, idst, isrc) + (sz, strides_dst, offset_dst) = structure_dst[idst] + n = first(sz) + return UniqueTransformerBlock{typeof(coeff), 1}( + coeff, (n,), (1,), dstdiagonal ? (1,) : (sum(strides_dst),), + dstdiagonal ? offsets_dst[idst] : offset_dst, (1,), offsets_src[isrc] + ) + end + work_offsets, nwork = _work_offsets(blk.sz[1] for blk in blocks) + return device_transformer(A, DeviceUniqueTreeTransformer(blocks, work_offsets, nwork)) +end + +@cached function diagonal_treebraider( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, levels::IndexTuple, dstdiagonal::Bool + ) where {TA <: AnyGPUVector} + host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) + return diagonal_device_transformer(A, host, dstdiagonal) +end + +@cached function diagonal_treetransposer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, dstdiagonal::Bool + ) where {TA <: AnyGPUVector} + host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) + return diagonal_device_transformer(A, host, dstdiagonal) +end + +function _add_diagonal!(tdst::GPUDiagonalTensorMap, tsrc, conjsrc, transformer, α, β) + _launch_unique!(tdst.data, tsrc.data, conjsrc ? conj : identity, transformer, α, β, Val(1)) + return tdst +end +# the kernel only touches the diagonal. be careful to apply `β` to the entire tensor. +function _add_diagonal!(tdst::GPUTensorMap, tsrc, conjsrc, transformer, α, β) + scale!(tdst, β) + _launch_unique!(tdst.data, tsrc.data, conjsrc ? conj : identity, transformer, α, One(), Val(1)) + return tdst +end + +Base.@propagate_inbounds function TensorKit._braid!( + tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, + p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator + ) + Vdst, Vsrc = space(tdst), space(tsrc) + @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + transformer = diagonal_treebraider( + storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels, tdst isa DiagonalTensorMap + ) + return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) +end + +Base.@propagate_inbounds function TensorKit._transpose!( + tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, + p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator + ) + Vdst, Vsrc = space(tdst), space(tsrc) + @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + transformer = diagonal_treetransposer( + storagetype(tdst), Vdst, Vsrc, p, conjsrc, tdst isa DiagonalTensorMap + ) + return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) +end + function TensorKit.add_transform_kernel!( dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, device::DeviceUniqueTreeTransformer, α, β, backend, allocator, ntasks::Int diff --git a/test/amd/tensors.jl b/test/amd/tensors.jl index 93c026367..9bf1d1997 100644 --- a/test/amd/tensors.jl +++ b/test/amd/tensors.jl @@ -2,6 +2,7 @@ using Adapt, AMDGPU using Test, TestExtras using Random using TensorKit, Combinatorics +using TensorOperations: TensorOperations as TO ad = adapt(Array) const AMDGPUExt = Base.get_extension(TensorKit, :TensorKitAMDGPUExt) @assert !isnothing(AMDGPUExt) @@ -526,3 +527,76 @@ end end end end + +@timedtestset "DiagonalTensorMap index manipulations with domain $V: test via CPU" for V in ( + (ℂ^4)', + Vect[Z2Irrep](0 => 2, 1 => 3), + Vect[U1Irrep](0 => 2, 1 => 2, -1 => 1), + Vect[FermionParity](0 => 2, 1 => 3), + Vect[SU2Irrep](0 => 2, 1 => 1)', + ) + function test_tensoradd(dst, d, p, conjd, α, β) + cpu = TO.tensoradd!(adapt(Array, dst), adapt(Array, d), p, conjd, α, β, TO.DefaultBackend(), TO.DefaultAllocator()) + gpu = TO.tensoradd!(copy(dst), d, p, conjd, α, β, TO.DefaultBackend(), TO.DefaultAllocator()) + @test gpu isa typeof(dst) + @test adapt(Array, gpu) ≈ cpu + return nothing + end + diagonal(T, W) = DiagonalTensorMap(ROCVector{T}(rand(T, reduceddim(W))), W) + + @testset "eltype $T" for T in (Float64, ComplexF64) + d = diagonal(T, V) + dcpu = adapt(Array, d) + scalars = ((TensorKit.One(), TensorKit.Zero()), (T(2), T(3))) + + 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)) + for (α, β) in scalars + test_tensoradd(dst, d, p, conjd, α, β) + end + end + for p in (((1,), (2,)), ((2,), (1,))), conjd in (false, true) + W = domain(TO.tensoradd_structure(dcpu, p, conjd))[1] + for (α, β) in scalars + test_tensoradd(diagonal(T, W), d, p, conjd, α, β) + end + end + for p in (((), (1, 2)), ((1, 2), ()), ((2, 1), ()), ((2,), (1,))) + @test adapt(Array, permute(d, p)) ≈ permute(dcpu, p) + @test adapt(Array, permute(d', p)) ≈ permute(dcpu', p) + end + for k in (0, 2) + @test adapt(Array, repartition(d, k)) ≈ repartition(dcpu, k) + @test adapt(Array, repartition(d', k)) ≈ repartition(dcpu', k) + end + for levels in ((1, 2), (2, 1)) + @test adapt(Array, braid(d, ((1, 2), ()), levels)) ≈ braid(dcpu, ((1, 2), ()), levels) + end + end + + @testset "real source, complex destination" begin + d = diagonal(Float64, V) + p = ((1,), (2,)) + W = domain(TO.tensoradd_structure(adapt(Array, d), p, true))[1] + test_tensoradd(diagonal(ComplexF64, W), d, p, true, TensorKit.One(), TensorKit.Zero()) + test_tensoradd(diagonal(ComplexF64, W), d, p, true, ComplexF64(2), ComplexF64(3)) + end + + @testset "contractions" begin + d = diagonal(Float64, V) + t = adapt(ROCVector{ComplexF64}, randn(ComplexF64, V ⊗ V ← V)) + s = adapt(ROCVector{Float64}, randn(Float64, V ← V)) + dcpu, tcpu, scpu = adapt(Array, d), adapt(Array, t), adapt(Array, s) + + @tensor gpu[-1 -2; -3] := t[-1 -2; 1] * d[1; -3] + @tensor cpu[-1 -2; -3] := tcpu[-1 -2; 1] * dcpu[1; -3] + @test adapt(Array, gpu) ≈ cpu + + @tensor gpu[-1 -2; -3] := d[-1; 1] * t[1 -2; -3] + @tensor cpu[-1 -2; -3] := dcpu[-1; 1] * tcpu[1 -2; -3] + @test adapt(Array, gpu) ≈ cpu + + @test (@tensor s[1; 2] * d[2; 1]) ≈ (@tensor scpu[1; 2] * dcpu[2; 1]) + end +end diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 59196d52c..68e3bfa8f 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -2,6 +2,7 @@ using Adapt, CUDA, CUDA.cuRAND, cuTENSOR using Test, TestExtras using Random using TensorKit, Combinatorics +using TensorOperations: TensorOperations as TO ad = adapt(Array) const CUDAExt = Base.get_extension(TensorKit, :TensorKitCUDAExt) @assert !isnothing(CUDAExt) @@ -576,3 +577,76 @@ end end end end + +@timedtestset "DiagonalTensorMap index manipulations with domain $V: test via CPU" for V in ( + (ℂ^4)', + Vect[Z2Irrep](0 => 2, 1 => 3), + Vect[U1Irrep](0 => 2, 1 => 2, -1 => 1), + Vect[FermionParity](0 => 2, 1 => 3), + Vect[SU2Irrep](0 => 2, 1 => 1)', + ) + function test_tensoradd(dst, d, p, conjd, α, β) + cpu = TO.tensoradd!(adapt(Array, dst), adapt(Array, d), p, conjd, α, β, TO.DefaultBackend(), TO.DefaultAllocator()) + gpu = TO.tensoradd!(copy(dst), d, p, conjd, α, β, TO.DefaultBackend(), TO.DefaultAllocator()) + @test gpu isa typeof(dst) + @test adapt(Array, gpu) ≈ cpu + return nothing + end + diagonal(T, W) = DiagonalTensorMap(CuVector{T}(rand(T, reduceddim(W))), W) + + @testset "eltype $T" for T in (Float64, ComplexF64) + d = diagonal(T, V) + dcpu = adapt(Array, d) + scalars = ((TensorKit.One(), TensorKit.Zero()), (T(2), T(3))) + + 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)) + for (α, β) in scalars + test_tensoradd(dst, d, p, conjd, α, β) + end + end + for p in (((1,), (2,)), ((2,), (1,))), conjd in (false, true) + W = domain(TO.tensoradd_structure(dcpu, p, conjd))[1] + for (α, β) in scalars + test_tensoradd(diagonal(T, W), d, p, conjd, α, β) + end + end + for p in (((), (1, 2)), ((1, 2), ()), ((2, 1), ()), ((2,), (1,))) + @test adapt(Array, permute(d, p)) ≈ permute(dcpu, p) + @test adapt(Array, permute(d', p)) ≈ permute(dcpu', p) + end + for k in (0, 2) + @test adapt(Array, repartition(d, k)) ≈ repartition(dcpu, k) + @test adapt(Array, repartition(d', k)) ≈ repartition(dcpu', k) + end + for levels in ((1, 2), (2, 1)) + @test adapt(Array, braid(d, ((1, 2), ()), levels)) ≈ braid(dcpu, ((1, 2), ()), levels) + end + end + + @testset "real source, complex destination" begin + d = diagonal(Float64, V) + p = ((1,), (2,)) + W = domain(TO.tensoradd_structure(adapt(Array, d), p, true))[1] + test_tensoradd(diagonal(ComplexF64, W), d, p, true, TensorKit.One(), TensorKit.Zero()) + test_tensoradd(diagonal(ComplexF64, W), d, p, true, ComplexF64(2), ComplexF64(3)) + end + + @testset "contractions" begin + d = diagonal(Float64, V) + t = adapt(CuVector{ComplexF64}, randn(ComplexF64, V ⊗ V ← V)) + s = adapt(CuVector{Float64}, randn(Float64, V ← V)) + dcpu, tcpu, scpu = adapt(Array, d), adapt(Array, t), adapt(Array, s) + + @tensor gpu[-1 -2; -3] := t[-1 -2; 1] * d[1; -3] + @tensor cpu[-1 -2; -3] := tcpu[-1 -2; 1] * dcpu[1; -3] + @test adapt(Array, gpu) ≈ cpu + + @tensor gpu[-1 -2; -3] := d[-1; 1] * t[1 -2; -3] + @tensor cpu[-1 -2; -3] := dcpu[-1; 1] * tcpu[1 -2; -3] + @test adapt(Array, gpu) ≈ cpu + + @test (@tensor s[1; 2] * d[2; 1]) ≈ (@tensor scpu[1; 2] * dcpu[2; 1]) + end +end