diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index a6ab14c3b..02f0d5035 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -30,6 +30,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Require TensorKitSectors v0.3.10 - `@cached` can be used from other modules, such as package extensions, which then own the cache of the methods they add; these caches are shown with their module in `global_cache_info` after registration in `__init__` (internal) - The `TreeTransformer`s used in index manipulations are now constructed and cached per storagetype of the destination tensor, which is also a dispatch point for storage-specific transformers with their own cache (internal) +- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s, and GPU tensors use a dedicated kernel (internal) ### Deprecated diff --git a/docs/src/lib/fusiontrees.md b/docs/src/lib/fusiontrees.md index 57033eca6..759eeec12 100644 --- a/docs/src/lib/fusiontrees.md +++ b/docs/src/lib/fusiontrees.md @@ -48,4 +48,5 @@ repartition(src::Union{FusionTreePair, FusionTreeBlock}, N::Int) Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, levels::Index2Tuple) permute(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) +planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, q::Index2Tuple) ``` diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 45f6df640..9974ddf35 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -9,11 +9,13 @@ using Strided: StridedViews using MatrixAlgebraKit, Adapt using TensorKit using TensorKit.TensorOperations: linearize, DefaultAllocator +const TO = TensorKit.TensorOperations using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check -using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple -import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel! +using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, TraceTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple +using TensorKit: SubblockIterator +import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel!, trace_transform_kernel! function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:AnyGPUMatrix{T}, <:StridedViews.StridedView{T, 4, <:AnyGPUArray{T}}}} # COV_EXCL_START @@ -271,6 +273,112 @@ function DeviceGenericTreeTransformer( ) end +# Planar traces run in two stages. The first computes partial traces of every source subblock: +# each thread sums a chunk of at most `chunksize` traced entries for a single remaining entry, +# such that large traces (e.g. full reductions) are still spread over many threads. The second +# recouples these partials with one thread per destination element, as for the generic transformers, +# such that the traced entries of a source are summed only once, rather than once per destination tree. + +""" + TraceTransformerBlock{N} + +Descriptor for a destination block of a `TraceTransformer`, indexing into the flat +`coeffs`/`structs_dst`/`sources` vectors of a `DeviceTraceTransformer`. +""" +struct TraceTransformerBlock{N} + sz::NTuple{N, Int} + dense_strides::NTuple{N, Int} + rows::Int + cols::Int + u_offset::Int + dst_offset::Int + src_offset::Int +end + +""" + TraceSourceBlock{N, M} + +Descriptor for the partial traces of a source subblock of a `TraceTransformer`, with `N` remaining +legs (permuted to the destination order) and `M` traced pairs, whose entries are found through a +single stride per pair. The partials are stored chunk-major, as `nchunks` dense arrays of shape `sz`. +""" +struct TraceSourceBlock{N, M} + sz::NTuple{N, Int} + dense_strides::NTuple{N, Int} + permuted_strides::NTuple{N, Int} + offset::Int + sz_traced::NTuple{M, Int} + dense_strides_traced::NTuple{M, Int} + strides_traced::NTuple{M, Int} # strides[q₁] .+ strides[q₂] + nremaining::Int + ntraced::Int + chunksize::Int + nchunks::Int +end + +struct DeviceTraceTransformer{VB <: AbstractVector{<:TraceTransformerBlock}, VO <: AbstractVector{Int}, VC <: AbstractVector{<:Number}, VD <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}, VS <: AbstractVector{<:TraceSourceBlock}} + blocks::VB + work_offsets::VO + nwork::Int + coeffs::VC # every `U`, concatenated in column-major order + structs_dst::VD + sources::VS + partial_offsets::VO # also the work offsets of the first stage + npartials::Int +end + +# traces with fewer entries are not split, larger ones are split into at most `√ntraced` chunks, +# such that no thread in either stage sums more than `max(256, √ntraced)` entries +const MIN_TRACE_CHUNKSIZE = 256 +_trace_chunksize(ntraced::Int) = max(MIN_TRACE_CHUNKSIZE, isqrt(ntraced)) + +function DeviceTraceTransformer(transformer::TraceTransformer{T, N₁, N₂}, p, (q₁, q₂)) where {T, N₁, N₂} + lp = linearize(p) + M = length(q₁) + blocks = TraceTransformerBlock{N₁}[] + coeffs = T[] + structs_dst = TreeStructure{N₁}[] + sources = TraceSourceBlock{N₁, M}[] + + (; structure_dst, structure_src) = transformer + for (U, inds_dst, inds_src) in transformer.data + # all destination trees in a block share the same subblock size + size_dst = first(structure_dst[first(inds_dst)]) + dense_strides_dst = _dense_strides(size_dst) + push!( + blocks, TraceTransformerBlock{N₁}( + size_dst, dense_strides_dst, size(U, 1), size(U, 2), + length(coeffs), length(structs_dst), length(sources) + ) + ) + append!(coeffs, U) + for idst in inds_dst + _, strides_dst, offset_dst = structure_dst[idst] + push!(structs_dst, (strides_dst, offset_dst)) + end + for isrc in inds_src + size_src, strides_src, offset_src = structure_src[isrc] + size_traced = TupleTools.getindices(size_src, q₁) + strides_traced = TupleTools.getindices(strides_src, q₁) .+ TupleTools.getindices(strides_src, q₂) + ntraced = prod(size_traced) + chunksize = _trace_chunksize(ntraced) + push!( + sources, TraceSourceBlock{N₁, M}( + size_dst, dense_strides_dst, TupleTools.getindices(strides_src, lp), offset_src, + size_traced, _dense_strides(size_traced), strides_traced, + prod(size_dst), ntraced, chunksize, cld(ntraced, chunksize) + ) + ) + end + end + + work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) + partial_offsets, npartials = _work_offsets(src.nchunks * src.nremaining for src in sources) + return DeviceTraceTransformer( + blocks, work_offsets, nwork, coeffs, structs_dst, sources, partial_offsets, npartials + ) +end + # Cache the device transformer directly, constructing the cached host version only on a miss. @cached function TensorKit.treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, @@ -287,6 +395,13 @@ end return device_transformer(A, host, linearize(p)) end +@cached function TensorKit.treetracer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple + ) where {TA <: AnyGPUVector} + host = TensorKit.treetracer(Vector{scalartype(A)}, Vdst, Vsrc, p, q) + return device_transformer(A, host, p, q) +end + const DEFAULT_DEVICECACHE_SIZE = 1000 function __init__() @@ -295,12 +410,14 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_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_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) return nothing end @@ -332,6 +449,19 @@ function device_transformer(A, transformer::GenericTreeTransformer, p) end end +function device_transformer(A, transformer::TraceTransformer, p, q) + packed = DeviceTraceTransformer(transformer, p, q) + return GPUArrays.@uncached begin + proto = A(undef, 0) + upload(x) = copy!(similar(proto, eltype(x), size(x)), x) + DeviceTraceTransformer( + upload(packed.blocks), upload(packed.work_offsets), packed.nwork, + upload(packed.coeffs), upload(packed.structs_dst), upload(packed.sources), + upload(packed.partial_offsets), packed.npartials + ) + end +end + # COV_EXCL_START # kernels are not reachable by coverage @@ -360,7 +490,12 @@ end # finds the overall linear index in the output and input arrays corresponding to the **sublock** # coordinates currently being worked on @inline function _linear_index(coords::NTuple{N, Int}, strides::NTuple{N, Int}, offset) where {N} - return offset + sum(ntuple(n -> coords[n] * strides[n], Val(N))) + 1 + # an explicit loop, since `sum` of an empty tuple (e.g. a full trace) does not compile for GPUs + i = offset + 1 + for n in 1:N + i += coords[n] * strides[n] + end + return i end # One thread per destination element in `data_dst`. `op` is `identity` or `conj`, and is @@ -426,6 +561,67 @@ end end end +# COV_EXCL_START +# First stage: one thread per (chunk, remaining entry) of every source subblock, writing its +# partial trace to `partials[w + 1]`. Consecutive threads handle consecutive remaining entries. +@kernel function partial_trace_kernel!(partials, data_src, sources, partial_offsets, npartials) + # w = @index(Global, Linear) - 1 does not work due to JLArrays bug + w = @index(Global, Linear) + w -= 1 + if w < npartials + j = _searchblock(partial_offsets, w) + src = @inbounds sources[j] + k, e = divrem(w - (@inbounds partial_offsets[j]), src.nremaining) + coords = _coordinates(e, src.sz, src.dense_strides) + i_src = _linear_index(coords, src.permuted_strides, src.offset) + + acc = zero(eltype(partials)) + start = k * src.chunksize + @inbounds for t in start:(min(start + src.chunksize, src.ntraced) - 1) + coords_traced = _coordinates(t, src.sz_traced, src.dense_strides_traced) + acc += data_src[i_src + _linear_index(coords_traced, src.strides_traced, -1)] + end + @inbounds partials[w + 1] = acc + end +end + +# Second stage: one thread per destination element, as for the generic transformers, where +# every source contributes the sum of its partial traces. +@kernel function trace_recouple_kernel!( + data_dst, partials, blocks, work_offsets, coeffs, structs_dst, sources, partial_offsets, + α, β, nwork + ) + w = @index(Global, Linear) + w -= 1 + if w < nwork + b = _searchblock(work_offsets, w) + blk = @inbounds blocks[b] + local_w = w - (@inbounds work_offsets[b]) + blocksize = prod(blk.sz) + i, e = divrem(local_w, blocksize) # 0-based destination tree and entry + coords = _coordinates(e, blk.sz, blk.dense_strides) + + st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1] + i_dst = _linear_index(coords, st_dst, offs_dst) + + acc = zero(promote_type(eltype(partials), eltype(coeffs))) + @inbounds for j in 1:blk.cols + coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] + iszero(coeff) && continue + js = blk.src_offset + j + offset = partial_offsets[js] + e + 1 + tr = zero(eltype(partials)) + for k in 0:(sources[js].nchunks - 1) + tr += partials[offset + k * blocksize] + end + acc += coeff * tr + end + # destinations without contributions are overwritten for `β == 0` + @inbounds data_dst[i_dst] = iszero(β) ? α * acc : α * acc + β * data_dst[i_dst] + end +end +# COV_EXCL_STOP + function _launch_unique!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} nwork = transformer.nwork nwork == 0 && return nothing @@ -450,7 +646,7 @@ end # we don't want/need the `StridedSubblocks` here because that information is already in the transformer TensorKit._transform_subblocks( tdst::TensorMap, tsrc::TensorMap, - transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer} + transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer, DeviceTraceTransformer} ) = (tdst.data, tsrc.data) function TensorKit.add_transform_kernel!( @@ -476,4 +672,40 @@ function TensorKit.add_transform_kernel!( return nothing end +function TensorKit.trace_transform_kernel!( + dst::AnyGPUVector, src::AnyGPUVector, p::Index2Tuple, q::Index2Tuple, + device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int + ) + device.nwork == 0 && return nothing + kernel_backend = get_backend(dst) + cp = TO.allocator_checkpoint!(allocator) + partials = TO.tensoralloc(similarstoragetype(src), device.npartials, Val(true), allocator) + if device.npartials > 0 + partial_trace_kernel!(kernel_backend)( + partials, src, device.sources, device.partial_offsets, device.npartials; + ndrange = device.npartials + ) + end + trace_recouple_kernel!(kernel_backend)( + dst, partials, device.blocks, device.work_offsets, device.coeffs, device.structs_dst, + device.sources, device.partial_offsets, α, β, device.nwork; ndrange = device.nwork + ) + TO.tensorfree!(partials, allocator) + TO.allocator_reset!(allocator, cp) + return nothing +end + +# other tensor types are first converted to `TensorMap`s: tracing subblock by subblock with +# `tensortrace!` is not an option, since cuTENSOR requires aligned data for the subblock views +function TensorKit.trace_transform_kernel!( + dst::SubblockIterator, src::SubblockIterator, p::Index2Tuple, q::Index2Tuple, + device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int + ) + _tensormap(t) = t isa TensorMap ? t : copy!(similar(t, scalartype(t), space(t)), t) + tdst, tsrc = _tensormap(dst.t), _tensormap(src.t) + trace_transform_kernel!(tdst.data, tsrc.data, p, q, device, α, β, backend, allocator, ntasks) + tdst === dst.t || copy!(dst.t, tdst) + return nothing +end + end diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index cae44aca3..9e613d50e 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -584,58 +584,85 @@ CacheStyle(::typeof(fstranspose), k::FSBTransposeKey{I}) where {I} = # -> composite manipulations that depend on the duality (rigidity) and pivotal structure # -> planar manipulations that do not require braiding, everything is in Fsymbol (A/Bsymbol) -function planar_trace((f₁, f₂)::FusionTreePair, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) +# leg positions of a planar trace of a tree pair with `N₁` outgoing and `N₂` incoming legs, after +# repartitioning all legs to the outgoing side: the traced pairs `q′`, the remaining legs `r` in +# increasing order and the permutation `p′` that is applied after removing the traced legs +function _planar_trace_indices(::Val{N₁}, ::Val{N₂}, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) where {N₁, N₂} length(q₁) == length(q₂) || throw(ArgumentError(lazy"trace index tuples q₁ and q₂ must have equal length, got $(length(q₁)) and $(length(q₂))")) - I = sectortype(f₁) N = length(p₁) + length(p₂) + 2 * length(q₁) - length(f₁) + length(f₂) == N || - throw(ArgumentError(lazy"fusion tree pair has $(length(f₁) + length(f₂)) indices, but permutation expects $N = $(length(p₁)) + $(length(p₂)) + 2×$(length(q₁))")) - if isempty(q₁) - return transpose((f₁, f₂), (p₁, p₂)) - end - - linearindex = ( - ntuple(identity, Val(length(f₁)))..., - reverse(length(f₁) .+ ntuple(identity, Val(length(f₂))))..., - ) - + N₁ + N₂ == N || + throw(ArgumentError(lazy"fusion tree pair has $(N₁ + N₂) indices, but permutation expects $N = $(length(p₁)) + $(length(p₂)) + 2×$(length(q₁))")) + linearindex = (ntuple(identity, Val(N₁))..., reverse(N₁ .+ ntuple(identity, Val(N₂)))...) q₁′ = TupleTools.getindices(linearindex, q₁) q₂′ = TupleTools.getindices(linearindex, q₂) - p₁′, p₂′ = let q′ = (q₁′..., q₂′...) - ( - map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₁)), - map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₂)), - ) - end + q′ = (q₁′..., q₂′...) + p₁′ = map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₁)) + p₂′ = map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₂)) + r = TupleTools.sort(TupleTools.getindices(linearindex, (p₁..., p₂...))) + return (p₁′, p₂′), (q₁′, q₂′), r +end - T = fusionscalartype(I) - F₁ = fusiontreetype(I, length(p₁)) - F₂ = fusiontreetype(I, length(p₂)) - newtrees = FusionTreeDict{Tuple{F₁, F₂}, T}() - if FusionStyle(I) isa UniqueFusion - (f₁′, f₂′), coeff′ = repartition((f₁, f₂), N) - for (f₁′′, coeff′′) in planar_trace(f₁′, (q₁′, q₂′)) - (f12′′′, coeff′′′) = transpose((f₁′′, f₂′), (p₁′, p₂′)) - coeff = coeff′ * coeff′′ * coeff′′′ - iszero(coeff) || (newtrees[f12′′′] = get(newtrees, f12′′′, zero(coeff)) + coeff) - end +# the sectors of the traced legs have to be dual to each other for a nonzero trace, where the +# sector of an incoming leg is the dual of the corresponding uncoupled sector of the fusion tree +function trace_sectors_match(src::Union{FusionTreePair, FusionTreeBlock}, (q₁, q₂)::Index2Tuple) + f₁, f₂ = src isa FusionTreePair ? src : first(fusiontrees(src)) + N₁ = length(f₁) + legsector(i) = i <= N₁ ? f₁.uncoupled[i] : dual(f₂.uncoupled[i - N₁]) + return all(map((i, j) -> legsector(i) == dual(legsector(j)), q₁, q₂)) +end + +""" + planar_trace((f₁, f₂)::FusionTreePair, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) -> (f₁′, f₂′) => coeff + planar_trace(src::FusionTreeBlock, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) -> dst => U + +Perform a planar trace of the legs `q₁[i]` with `q₂[i]` of a fusion tree pair, followed by a +planar permutation `(p₁, p₂)` of the remaining legs. The sectors of the traced legs have to be +dual to each other. + +For `FusionStyle(I) isa UniqueFusion`, the input is a single fusion tree pair, which is mapped +to a single output tree pair with a coefficient. For generic `FusionStyle`s, the input and +output consist of `FusionTreeBlock`s, where `src` and `dst` contain all tree pairs of the +traced and remaining legs respectively, and `U[j, i]` is the coefficient of the `j`th tree pair +of `dst` in the trace of the `i`th tree pair of `src`. +""" +function planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) + src isa FusionTreePair && @assert FusionStyle(src) === UniqueFusion() + p′, q′, r = _planar_trace_indices(Val(numout(src)), Val(numin(src)), (p₁, p₂), (q₁, q₂)) + trace_sectors_match(src, (q₁, q₂)) || + throw(ArgumentError(lazy"sectors of the traced legs $q₁ and $q₂ are not dual to each other")) + isempty(q₁) && return transpose(src, (p₁, p₂)) + + # bend all legs to the outgoing side, where the trace only acts on the splitting trees + src′, U₁ = repartition(src, numind(src)) + traced, U₂ = _planar_trace_splitting(src′, q′, r) + # for multiple fusion channels, the remaining legs need not couple to the unit + if traced isa FusionTreeBlock && isempty(traced) + I = sectortype(src) + F = fusiontreetype(I, length(p₁), length(p₂)) + return FusionTreeBlock(F[]) => zeros(fusionscalartype(I), 0, length(src)) else - # TODO: this is a bit of a hack to fix the traces for now - src = FusionTreeBlock([(f₁, f₂)]) - dst, U = repartition(src, N) - for ((f₁′, f₂′), coeff′) in zip(fusiontrees(dst), U) - for (f₁′′, coeff′′) in planar_trace(f₁′, (q₁′, q₂′)) - src′ = FusionTreeBlock([(f₁′′, f₂′)]) - dst′, U′ = transpose(src′, (p₁′, p₂′)) - for (f12′′′, coeff′′′) in zip(fusiontrees(dst′), U′) - coeff = coeff′ * coeff′′ * coeff′′′ - iszero(coeff) || (newtrees[f12′′′] = get(newtrees, f12′′′, zero(coeff)) + coeff) - end - end - end + dst, U₃ = transpose(traced, p′) + return dst => U₃ * U₂ * U₁ end - return newtrees +end + +function _planar_trace_splitting((f, f₀)::FusionTreePair, q′, r) + f′, coeff = only(planar_trace(f, q′)) + return (f′, f₀) => coeff +end +function _planar_trace_splitting(src::FusionTreeBlock, q′, r) + I = sectortype(src) + f₀, = first(fusiontrees(src)) + uncoupled = TupleTools.getindices(f₀.uncoupled, r) + isdual = TupleTools.getindices(f₀.isdual, r) + dst = FusionTreeBlock{I}((uncoupled, ()), (isdual, ())) + U = zeros(fusionscalartype(I), length(dst), length(src)) + positions = Dict(f => i for (i, (f, _)) in enumerate(fusiontrees(dst))) + for (j, (f, _)) in enumerate(fusiontrees(src)), (f′, coeff) in planar_trace(f, q′) + U[positions[f′], j] += coeff + end + return dst => U end """ diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index ea8334462..4ea3ab915 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -238,9 +238,12 @@ function Base.hash(f::FusionTree{I}, h::UInt) where {I} end return h end -# Hashing the first tree pair covers the block's defining data `(uncoupled, isdual)`, while -# also disambiguating the ad-hoc single-pair blocks constructed in `planar_trace` -Base.hash(b::FusionTreeBlock, h::UInt) = isempty(b) ? h : hash(first(fusiontrees(b)), h) +# a block is defined by the uncoupled sectors and duality flags shared by all its tree pairs +function Base.hash(b::FusionTreeBlock, h::UInt) + isempty(b) && return h + f₁, f₂ = first(fusiontrees(b)) + return hash(f₁.uncoupled, hash(f₁.isdual, hash(f₂.uncoupled, hash(f₂.isdual, h)))) +end function Base.:(==)(f₁::FusionTree{I, N}, f₂::FusionTree{I, N}) where {I <: Sector, N} f₁.coupled == f₂.coupled || return false diff --git a/src/planar/planaroperations.jl b/src/planar/planaroperations.jl index 7c24e1399..5cae4ea8f 100644 --- a/src/planar/planaroperations.jl +++ b/src/planar/planaroperations.jl @@ -92,24 +92,79 @@ function planartrace!( end @timeit_debug GLOBAL_TIMER "planartrace!" begin - if iszero(β) - fill!(C, β) - elseif !isone(β) - rmul!(C, β) + transformer = treetracer(storagetype(C), space(C), space(A), (p₁, p₂), (q₁, q₂)) + dst, src = _transform_subblocks(C, A, transformer) + ntasks = get_num_transformer_threads(A) + trace_transform_kernel!(dst, src, (p₁, p₂), (q₁, q₂), transformer, α, β, backend, allocator, ntasks) + end + return C +end + +# kernel for computing `tdst = β * tdst + α * tr(tsrc)`, with the planar trace encoded in `transformer` +function trace_transform_kernel!( + dst::TransformSubblocks, src::TransformSubblocks, p::Index2Tuple, q::Index2Tuple, + transformer::TraceTransformer, α, β, backend, allocator, ntasks::Int + ) + bufsize = buffersize(transformer) + if bufsize == 0 # every block has a single destination subblock + taskforeach(transformer.data, ntasks) do (U, inds_dst, inds_src) + _trace_transform_block!(dst, src, p, q, U, inds_dst, inds_src, nothing, α, β, backend, allocator) + end + else + # see `add_transform_kernel!` for the workspace strategy + cp = TO.allocator_checkpoint!(allocator) + @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ + TO.tensoralloc(storagetype(dst), bufsize, Val(true), allocator) + for _ in 1:min(length(transformer.data), ntasks) + ] + taskforeach(transformer.data, buffers) do (U, inds_dst, inds_src), buffer + _trace_transform_block!(dst, src, p, q, U, inds_dst, inds_src, buffer, α, β, backend, allocator) + end + foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) + TO.allocator_reset!(allocator, cp) + end + return nothing +end + +function _trace_transform_block!( + dst::TransformSubblocks, src::TransformSubblocks, p::Index2Tuple, q::Index2Tuple, + U, inds_dst, inds_src, buffer, α, β, backend, allocator + ) + if size(U, 1) == 1 # single destination: trace every source directly into it + idst = only(inds_dst) + if isempty(inds_src) + @inbounds iszero(β) ? zerovector!(dst[idst]) : scale!(dst[idst], β) + return nothing end - β′ = One() - for (f₁, f₂) in fusiontrees(A) - for ((f₁′, f₂′), coeff) in planar_trace((f₁, f₂), (p₁, p₂), (q₁, q₂)) - @timeit_debug GLOBAL_TIMER "dense: trace" TO.tensortrace!( - C[f₁′, f₂′], - A[f₁, f₂], (p₁, p₂), (q₁, q₂), false, - α * coeff, β′, - backend, allocator - ) + @timeit_debug GLOBAL_TIMER "dense: trace" @inbounds begin + TO.tensortrace!(dst[idst], src[first(inds_src)], p, q, false, α * U[1, 1], β, backend, allocator) + for i in 2:length(inds_src) + TO.tensortrace!(dst[idst], src[inds_src[i]], p, q, false, α * U[1, i], One(), backend, allocator) end end + else # trace → recoupling matmul → unpack + rows, cols = size(U) + sz = size(@inbounds(dst[first(inds_dst)])) + ptriv = (ntuple(identity, length(sz)), ()) + blocksize = prod(sz) + buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) + buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) + + # the permutation is applied while tracing, such that unpacking is a plain addition + @timeit_debug GLOBAL_TIMER "dense: trace" @inbounds for (i, isrc) in enumerate(inds_src) + TO.tensortrace!( + sreshape(view(buffer_src, :, i), sz), src[isrc], p, q, false, + One(), Zero(), backend, allocator + ) + end + @timeit_debug GLOBAL_TIMER "dense: recouple mul!" _recouple!(buffer, buffer_dst, buffer_src, U) + @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (j, idst) in enumerate(inds_dst) + TO.tensoradd!( + dst[idst], sreshape(view(buffer_dst, :, j), sz), ptriv, false, α, β, backend, allocator + ) + end end - return C + return nothing end # insert default backend diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index b274c6f8c..73d910775 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,4 +202,104 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end + +""" + TraceTransformer{T, N₁, N₂} + +Planar trace of a tensor with `N₂` indices into a tensor with `N₁` indices, stored as +`(U, inds_dst, inds_src)` where `U[j, i]` is the coefficient with which the trace of source +subblock `inds_src[i]` contributes to destination subblock `inds_dst[j]`. The data is grouped by +destination, such that every destination subblock appears in exactly one entry, where those +without contributions have no source subblocks. The subblock structures of the destination and source +spaces are kept alongside, such that the [`StridedSubblocks`](@ref) of both tensors can be +created without further lookups. +""" +struct TraceTransformer{T, N₁, N₂} + data::Vector{GenericTransformerData{T}} + structure_dst::Vector{StridedStructure{N₁}} + structure_src::Vector{StridedStructure{N₂}} +end + +function TraceTransformer(::Type{A}, Vdst::HomSpace, Vsrc::HomSpace, p::Index2Tuple, q::Index2Tuple) where {A} + t₀ = Base.time() + src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) + + # destination positions => (columns of the recoupling matrix, source positions) + groups = Dict{Vector{Int}, Tuple{Vector{Vector{T}}, Vector{Int}}}() + function addcolumn!(inds_dst, u, isrc) + any(!iszero, u) || return nothing + us, inds_src = get!(() -> (Vector{T}[], Int[]), groups, inds_dst) + push!(us, u) + push!(inds_src, isrc) + return nothing + end + @timeit_debug GLOBAL_TIMER "symmetry: trace coefficients" if FusionStyle(sectortype(Vdst)) isa UniqueFusion + for (isrc, f) in enumerate(src_trees) + trace_sectors_match(f, q) || continue + f′, coeff = planar_trace(f, p, q) + _, (_, idst) = gettoken(dst_trees, f′) + addcolumn!([idst], [convert(T, coeff)], isrc) + end + else + for fs_src in fusionblocks(Vsrc) + trace_sectors_match(fs_src, q) || continue + fs_dst, U = planar_trace(fs_src, p, q) + inds_dst = map(fusiontrees(fs_dst)) do f + _, (_, idst) = gettoken(dst_trees, f) + return idst + end + for (i, f) in enumerate(fusiontrees(fs_src)) + _, (_, isrc) = gettoken(src_trees, f) + addcolumn!(inds_dst, convert(Vector{T}, U[:, i]), isrc) + end + end + end + data = GenericTransformerData{T}[(reduce(hcat, us), inds_dst, inds_src) for (inds_dst, (us, inds_src)) in groups] + # destination subblocks without contributions are only scaled + covered = falses(length(dst_trees)) + foreach(inds_dst -> covered[inds_dst] .= true, keys(groups)) + for idst in findall(!, covered) + push!(data, (zeros(T, 1, 0), [idst], Int[])) + end + + structure_dst = degeneracystructure(Vdst).subblockstructure + structure_src = degeneracystructure(Vsrc).subblockstructure + # sort by (approximate) weight to facilitate multi-threading strategies + weights = map(data) do (U, inds_dst, _) + return length(U) * prod(structure_dst[first(inds_dst)][1]) + end + Base.permute!(data, sortperm(weights; rev = true)) + transformer = TraceTransformer(data, structure_dst, structure_src) + + Δt = Base.time() - t₀ + @debug("TraceTransformer for $Vsrc to $Vdst via $p and $q", nblocks = length(data), Δt) + + return transformer +end + +""" + buffersize(transformer::TraceTransformer) -> Int + +Compute the workspace size required to trace, recouple and unpack the largest block with +multiple destination subblocks, i.e. `prod(sz_dst) * (rows + cols)` where `(rows, cols) = size(U)`. +""" +function buffersize(transformer::TraceTransformer) + structure_dst = transformer.structure_dst + return maximum(transformer.data; init = 0) do (U, inds_dst, _) + return size(U, 1) == 1 ? 0 : prod(structure_dst[first(inds_dst)][1]) * sum(size(U)) + end +end + +function tracetransformertype(::Type{A}, Vdst, Vsrc) where {A} + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) + return TraceTransformer{T, numind(Vdst), numind(Vsrc)} +end + +@cached function treetracer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple + )::tracetransformertype(A, Vdst, Vsrc) where {TA} + return TraceTransformer(A, Vdst, Vsrc, p, q) +end + # default cachestyle is GlobalLRUCache diff --git a/test/cuda/planar.jl b/test/cuda/planar.jl index 2d992023b..e92faef8e 100644 --- a/test/cuda/planar.jl +++ b/test/cuda/planar.jl @@ -74,6 +74,24 @@ for V in spacelist end end +for V in spacelist + I = sectortype(first(V)) + Istr = type_repr(I) + @timedtestset "planartrace + CUDA with symmetry: $Istr" begin + # (space, p, q) for a trace of neighbouring legs and a nested trace + cases = ( + (V[1] ⊗ V[2] ⊗ V[3] ← V[4] ⊗ V[3], ((1, 2), (4,)), ((3,), (5,))), + (V[1] ⊗ V[2] ⊗ V[3] ⊗ V[4] ← V[3] ⊗ V[4], ((1,), (2,)), ((3, 4), (5, 6))), + ) + for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) + A = randn(ComplexF64, W) + C = randn(ComplexF64, TensorKit.select(W, p)) + Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) + @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) + end + end +end + @testset "planar methods" verbose = true begin @testset "planaradd" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^6 ⊗ ℂ^5 ⊗ ℂ^4) @@ -98,6 +116,24 @@ end planartrace!(C′, A′, p, q, true, true) end + @testset "planartrace with large traces: $I" for (I, V, Vr) in ( + (FermionParity, Vect[FermionParity](0 => 40, 1 => 40), Vect[FermionParity](0 => 2, 1 => 3)), + (FibonacciAnyon, Vect[FibonacciAnyon](:I => 30, :τ => 30), Vect[FibonacciAnyon](:I => 2, :τ => 1)), + ) + # (space, p, q) for a full reduction and for a trace with few remaining entries, which are + # both split into several chunks of traced entries + cases = ( + (V ⊗ V ← V ⊗ V, ((), ()), ((1, 2), (3, 4))), + (Vr ⊗ V ⊗ V ← V ⊗ V, ((1,), ()), ((2, 3), (4, 5))), + ) + for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) + A = randn(ComplexF64, W) + C = randn(ComplexF64, TensorKit.select(W, p)) + Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) + @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) + end + end + @testset "planarcontract" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^2 ⊗ ℂ^5 ⊗ ℂ^4) B = CUDA.randn(ℂ^2 ⊗ ℂ^4 ← ℂ^4 ⊗ ℂ^3) diff --git a/test/symmetries/doubletree.jl b/test/symmetries/doubletree.jl index 38d0aaea0..0af5db19c 100644 --- a/test/symmetries/doubletree.jl +++ b/test/symmetries/doubletree.jl @@ -249,28 +249,24 @@ using TensorKitSectors end @testset "Double fusion tree: planar trace" begin + p1 = ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,)) + p2, q2 = ((2:N...,), (1, ((2N):-1:(N + 3))...)), ((N + 1,), (N + 2,)) if FusionStyle(I) isa UniqueFusion f1, f1 = src - dst, U = transpose((f1, f1), ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,))) - d1 = zip((dst,), (U,)) + dst, U = transpose((f1, f1), p1) + @test TK.trace_sectors_match(dst, q2) + dst2, U2 = @constinferred TK.planar_trace(dst, p2, q2) + d2 = Dict(dst2 => U2 * U) else f1, f1 = first(fusiontrees(src)) src′ = FusionTreeBlock{I}((f1.uncoupled, f1.uncoupled), (f1.isdual, f1.isdual)) - dst, U = transpose(src′, ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,))) - d1 = zip(fusiontrees(dst), U[:, 1]) + dst, U = transpose(src′, p1) + @test TK.trace_sectors_match(dst, q2) + dst2, U2 = TK.planar_trace(dst, p2, q2) + d2 = Dict(zip(fusiontrees(dst2), U2 * U[:, 1])) end f1front, = TK.split(f1, N - 1) - T = sectorscalartype(I) - d2 = Dict{typeof((f1front, f1front)), T}() - for ((f1′, f2′), coeff′) in d1 - for ((f1′′, f2′′), coeff′′) in TK.planar_trace( - (f1′, f2′), ((2:N...,), (1, ((2N):-1:(N + 3))...)), ((N + 1,), (N + 2,)) - ) - coeff = coeff′ * coeff′′ - d2[(f1′′, f2′′)] = get(d2, (f1′′, f2′′), zero(coeff)) + coeff - end - end for ((f1_, f2_), coeff) in d2 if (f1_, f2_) == (f1front, f1front) @test coeff ≈ dim(f1.coupled) / dim(f1front.coupled)