diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index b91e0e68b..07fdc3bcf 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 @@ -43,6 +44,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. - Constructing a `FusionTreeBlock` only takes the fusion product of the shorter side, and `fusionblocks` groups the cached fusion trees of the space instead of constructing a block for every combination of sectors (internal) +- Fusion tree block manipulations no longer construct and index the canonical destination block, but keep the trees of the result in the order in which they are produced; the buffers of the intermediate blocks of a braid are recycled, and blocks with a single tree are braided as their tree pair (internal) ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 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 5e013865d..610df92ff 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -12,8 +12,9 @@ using TensorKit.TensorOperations: linearize, DefaultAllocator 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 +272,75 @@ function DeviceGenericTreeTransformer( ) end +""" + TraceTransformerBlock{N} + +Descriptor for a destination block of a `TraceTransformer`, indexing into the flat +`coeffs`/`structs_dst`/`structs_src` 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 + +# (permuted strides of the remaining legs, sizes, dense strides and summed strides of the traced +# pairs, offset): the traced entries of a source subblock are found through a single stride per pair +const TraceSourceStructure{N, M} = Tuple{NTuple{N, Int}, NTuple{M, Int}, NTuple{M, Int}, NTuple{M, Int}, Int} + +struct DeviceTraceTransformer{VB <: AbstractVector{<:TraceTransformerBlock}, VO <: AbstractVector{Int}, VC <: AbstractVector{<:Number}, VD <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}, VS <: AbstractVector{<:TraceSourceStructure}} + blocks::VB + work_offsets::VO + nwork::Int + coeffs::VC # every `U`, concatenated in column-major order + structs_dst::VD + structs_src::VS +end + +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₁}[] + structs_src = TraceSourceStructure{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)]) + push!( + blocks, TraceTransformerBlock{N₁}( + size_dst, _dense_strides(size_dst), size(U, 1), size(U, 2), + length(coeffs), length(structs_dst), length(structs_src) + ) + ) + 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₂) + push!( + structs_src, ( + TupleTools.getindices(strides_src, lp), size_traced, + _dense_strides(size_traced), strides_traced, offset_src, + ) + ) + end + end + + work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) + return DeviceTraceTransformer(blocks, work_offsets, nwork, coeffs, structs_dst, structs_src) +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 +357,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 +372,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 +411,18 @@ 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.structs_src) + ) + end +end + # COV_EXCL_START # kernels are not reachable by coverage @@ -360,7 +451,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 +522,44 @@ end end end +# COV_EXCL_START +# One thread per destination element, as for the generic transformers, where every source +# additionally contributes the sum over its traced entries. +@kernel function batched_trace_kernel!( + data_dst, data_src, blocks, work_offsets, coeffs, structs_dst, structs_src, α, β, nwork + ) + # w = @index(Global, Linear) - 1 does not work due to JLArrays bug + 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 = local_w ÷ blocksize # 0-based destination tree + coords = _coordinates(local_w % blocksize, 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(data_src), eltype(coeffs))) + @inbounds for j in 1:blk.cols + coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] + iszero(coeff) && continue + pst_src, sz_tr, dense_tr, st_tr, offs_src = structs_src[blk.src_offset + j] + i_src = _linear_index(coords, pst_src, offs_src) + tr = zero(eltype(data_src)) + for t in 0:(prod(sz_tr) - 1) + tr += data_src[i_src + _linear_index(_coordinates(t, sz_tr, dense_tr), st_tr, -1)] + 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 +584,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 +610,29 @@ 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 + batched_trace_kernel!(get_backend(dst))( + dst, src, device.blocks, device.work_offsets, device.coeffs, + device.structs_dst, device.structs_src, α, β, device.nwork; ndrange = device.nwork + ) + 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/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 368e326a7..ceaa1882b 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -14,10 +14,12 @@ applying `artin_braid(f′, i; inv = true)` to all the outputs `f′` of `artin_braid(f, i; inv = false)` and collecting the results should yield a single fusion tree with non-zero coefficient, namely `f` with coefficient `1`. This keyword has no effect if `BraidingStyle(sectortype(f)) isa SymmetricBraiding`. + +For sectors without `UniqueFusion`, the braid of `f` has to be a single fusion tree, which is the +case if `f` is the only fusion tree with its uncoupled and coupled sectors; this throws otherwise. """ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} 1 <= i < N || throw(ArgumentError(lazy"Cannot swap outputs i=$i and i+1 out of only $N outputs")) - @assert FusionStyle(I) === UniqueFusion() uncoupled = f.uncoupled a, b = uncoupled[i], uncoupled[i + 1] @@ -53,7 +55,7 @@ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} if i == 1 c = N > 2 ? inner[1] : coupled′ - R = oftype(oneT, (inv ? conj(Rsymbol(b, a, c)) : Rsymbol(a, b, c))) + R = oftype(oneT, (inv ? conj(only(Rsymbol(b, a, c))) : only(Rsymbol(a, b, c)))) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner, vertices) return f′ => R end @@ -64,13 +66,13 @@ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} a = inner_extended[i - 1] c = inner_extended[i] e = inner_extended[i + 1] - c′ = only(a ⊗ d) + c′ = _artin_channel(a, b, d, e) coeff = oftype( oneT, if inv - conj(Rsymbol(d, c, e) * Fsymbol(d, a, b, e, c′, c)) * Rsymbol(d, a, c′) + conj(only(Rsymbol(d, c, e)) * only(Fsymbol(d, a, b, e, c′, c))) * only(Rsymbol(d, a, c′)) else - Rsymbol(c, d, e) * conj(Fsymbol(d, a, b, e, c′, c) * Rsymbol(a, d, c′)) + only(Rsymbol(c, d, e)) * conj(only(Fsymbol(d, a, b, e, c′, c)) * only(Rsymbol(a, d, c′))) end ) inner′ = TupleTools.setindex(inner, c′, i - 1) @@ -78,9 +80,26 @@ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} return f′ => coeff end +# the inner lines `c′` of `a ⊗ d → c′ → c′ ⊗ b → e`, i.e. `intersect(a ⊗ d, e ⊗ conj(b))`, +# without constructing the second fusion product +_artin_channels(a, b, d, e) = Iterators.filter(c′ -> !iszero(Nsymbol(c′, b, e)), a ⊗ d) +# the only such line, for a fusion tree that braids into a single one (`only` throws otherwise) +_artin_channel(a, b, d, e) = _artin_channel(FusionStyle(a), a, b, d, e) +_artin_channel(::UniqueFusion, a, b, d, e) = only(a ⊗ d) +_artin_channel(::FusionStyle, a, b, d, e) = only(_artin_channels(a, b, d, e)) + function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where {I, N} + n = length(src) # braiding a complete block keeps its size + return artin_braid!(Matrix{sectorscalartype(I)}(undef, n, n), src, i; inv) # zeroed by `artin_braid!` +end +# `U` is overwritten with the braid matrix of `src`, a complete block; the trees of the result are +# stored in `trees` if given +function artin_braid!( + U::AbstractMatrix, src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, trees = nothing + ) where {I, N} 1 <= i < N || throw(ArgumentError("Cannot swap outputs i=$i and i+1 out of only $N outputs")) - isempty(src) && return src => zeros(sectorscalartype(I), 0, 0) + size(U) == (length(src), length(src)) || throw(DimensionMismatch(lazy"braid matrix of size $(size(U)) for a block of $(length(src)) trees")) + isempty(src) && return src => U uncoupled = src.uncoupled[1] a, b = uncoupled[i], uncoupled[i + 1] @@ -91,12 +110,10 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where isdual = src.isdual[1] isdual′ = TupleTools.setindex(isdual, isdual[i], i + 1) isdual′ = TupleTools.setindex(isdual′, isdual[i + 1], i) - dst = FusionTreeBlock{I}((uncoupled′, ()), (isdual′, ()); sizehint = length(src)) + rows = TreeRows(fusiontreetype(I, N, 0), length(src), trees) oneT = one(sectorscalartype(I)) - - indexmap = treeindex_map(dst) - U = zeros(sectorscalartype(I), length(dst), length(src)) + fill!(U, zero(oneT)) if isunit(a) || isunit(b) # braiding with trivial sector: simple and always possible for (col, (f, f₂)) in enumerate(fusiontrees(src)) @@ -113,10 +130,10 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where vertices′ = TupleTools.setindex(vertices′, vertices[i - 1], i) end f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′, vertices′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = oneT end - return dst => U + return FusionTreeBlock(rows.trees) => U end BraidingStyle(I) isa NoBraiding && @@ -132,7 +149,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where if FusionStyle(I) isa MultiplicityFreeFusion R = oftype(oneT, (inv ? conj(Rsymbol(b, a, c)) : Rsymbol(a, b, c))) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner, vertices) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = R else # GenericFusion μ = vertices[1] @@ -142,7 +159,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where iszero(R) && continue vertices′ = TupleTools.setindex(vertices, ν, 1) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner, vertices′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = R end end @@ -155,38 +172,39 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where c = inner_extended[i] e = inner_extended[i + 1] if FusionStyle(I) isa MultiplicityFreeFusion - for c′ in intersect(a ⊗ d, e ⊗ conj(b)) + R = inv ? Rsymbol(d, c, e) : Rsymbol(c, d, e) + for c′ in _artin_channels(a, b, d, e) coeff = if inv - conj(Rsymbol(d, c, e) * Fsymbol(d, a, b, e, c′, c)) * Rsymbol(d, a, c′) + conj(R * Fsymbol(d, a, b, e, c′, c)) * Rsymbol(d, a, c′) else - Rsymbol(c, d, e) * conj(Fsymbol(d, a, b, e, c′, c) * Rsymbol(a, d, c′)) + R * conj(Fsymbol(d, a, b, e, c′, c) * Rsymbol(a, d, c′)) end iszero(coeff) && continue inner′ = TupleTools.setindex(inner, c′, i - 1) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = coeff end else # GenericFusion - for c′ in intersect(a ⊗ d, e ⊗ conj(b)) - Rmat1 = inv ? Rsymbol(d, c, e)' : Rsymbol(c, d, e) + Rmat1 = inv ? Rsymbol(d, c, e)' : Rsymbol(c, d, e) + μ = vertices[i - 1] + ν = vertices[i] + for c′ in _artin_channels(a, b, d, e) Rmat2 = inv ? Rsymbol(d, a, c′)' : Rsymbol(a, d, c′) Fmat = Fsymbol(d, a, b, e, c′, c) - μ = vertices[i - 1] - ν = vertices[i] - for σ in 1:Nsymbol(a, d, c′) - for λ in 1:Nsymbol(c′, b, e) + inner′ = TupleTools.setindex(inner, c′, i - 1) + for σ in axes(Rmat2, 1) + vertices_σ = TupleTools.setindex(vertices, σ, i - 1) + for λ in axes(Fmat, 2) coeff = zero(oneT) - for ρ in 1:Nsymbol(d, c, e), κ in 1:Nsymbol(d, a, c′) + for ρ in axes(Rmat1, 2), κ in axes(Rmat2, 2) coeff += Rmat1[ν, ρ] * conj(Fmat[κ, λ, μ, ρ]) * conj(Rmat2[σ, κ]) end iszero(coeff) && continue - vertices′ = TupleTools.setindex(vertices, σ, i - 1) - vertices′ = TupleTools.setindex(vertices′, λ, i) - inner′ = TupleTools.setindex(inner, c′, i - 1) + vertices′ = TupleTools.setindex(vertices_σ, λ, i) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′, vertices′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = coeff end end @@ -194,7 +212,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where end end - return dst => U + return FusionTreeBlock(rows.trees) => U end # braid fusion tree @@ -216,9 +234,8 @@ braid(f::FusionTree{I, N}, p::IndexTuple{N}, levels::IndexTuple{N}) where {I, N} braid(f, (p, ()), (levels, ())) function braid(f::FusionTree{I, N}, (p, _)::Index2Tuple{N, 0}, (levels, _)::Index2Tuple{N, 0}) where {I, N} TupleTools.isperm(p) || throw(ArgumentError(lazy"not a valid permutation: $p")) - @assert FusionStyle(I) isa UniqueFusion braid_style = BraidingStyle(I) - if braid_style isa SymmetricBraiding # this assumes Fsymbols are 1! + if FusionStyle(I) isa UniqueFusion && braid_style isa SymmetricBraiding # Fsymbols are 1 coeff = one(sectorscalartype(I)) for i in 1:N for j in 1:(i - 1) @@ -303,42 +320,56 @@ end @cached function fsbraid(key::K)::_fsdicttype(K) where {I, N₁, N₂, K <: Union{FSPBraidKey{I, N₁, N₂}, FSBBraidKey{I, N₁, N₂}}} if K <: FSPBraidKey - ((f₁, f₂), (p1, p2), (l1, l2)) = key - p = linearizepermutation(p1, p2, length(f₁), length(f₂)) - levels = (l1..., reverse(l2)...) - (f, f0), coeff1 = repartition((f₁, f₂), N₁ + N₂) - f′, coeff2 = braid(f, p, levels) - (f₁′, f₂′), coeff3 = repartition((f′, f0), N₁) - return (f₁′, f₂′) => coeff1 * coeff2 * coeff3 + f, p, levels = key + return _braid_single(f::FusionTreePair{I}, p, levels) else src, (p1, p2), (l1, l2) = key + src = src::FusionTreeBlock{I} + + if length(src) == 1 # then so is every intermediate block: the braid is a product of scalars + f′, coeff = _braid_single(only(fusiontrees(src)), (p1, p2), (l1, l2)) + return FusionTreeBlock([f′]) => fill(coeff, 1, 1) + end p = linearizepermutation(p1, p2, numout(src), numin(src)) levels = (l1..., reverse(l2)...) - dst, U = repartition(src, numind(src)) + dst, U = repartition(src, Val(numind(src))) + # the swaps keep the size of the block, so that the tree vector and the matrices of the + # previous step are recycled (except for the trees of `src`, which are part of the key) + trees = nothing + U_tmp = Matrix{sectorscalartype(I)}(undef, length(dst), length(dst)) + U_spare = nothing braid_style = BraidingStyle(I) for s in permutation2swaps(p) _check_levels(braid_style, levels, s) inv = levels[s] > levels[s + 1] - dst, U_tmp = artin_braid(dst, s; inv) - U = U_tmp * U + dst′, _ = artin_braid!(U_tmp, dst, s; inv, trees) + trees = dst.trees === src.trees ? nothing : dst.trees + dst = dst′ + U_spare, U = U, U_spare isa typeof(U_tmp) ? mul!(U_spare, U_tmp, U) : U_tmp * U l = levels[s] levels = TupleTools.setindex(levels, levels[s + 1], s) levels = TupleTools.setindex(levels, l, s + 1) end - if N₂ == 0 - return dst => U - else - dst, U_tmp = repartition(dst::FusionTreeBlock{I}, N₁) - U = U_tmp * U - return dst => U - end + N₂ == 0 && return dst => U + dst, U_tmp = repartition(dst::FusionTreeBlock{I}, Val(N₁)) + return dst => U_tmp * U end end +# also for a tree pair that is the only one with its uncoupled sectors (`only` throws otherwise) +function _braid_single((f₁, f₂)::FusionTreePair, (p1, p2)::Index2Tuple{N₁, N₂}, (l1, l2)::Index2Tuple) where {N₁, N₂} + p = linearizepermutation(p1, p2, length(f₁), length(f₂)) + levels = (l1..., reverse(l2)...) + (f, f0), coeff1 = repartition((f₁, f₂), N₁ + N₂) + f′, coeff2 = braid(f, p, levels) + (f₁′, f₂′), coeff3 = repartition((f′, f0), N₁) + return (f₁′, f₂′) => coeff1 * coeff2 * coeff3 +end + CacheStyle(::typeof(fsbraid), k::Union{FSPBraidKey{I}, FSBBraidKey{I}}) where {I} = FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index e399f87ab..58e542988 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -53,16 +53,15 @@ function _bendright_treepair((f₁, f₂)::FusionTreePair) return (a, b, c), (f₁′, f₂′) end +# also for a final splitting vertex without multiplicity (`only` throws otherwise) function bendright((f₁, f₂)::FusionTreePair) - I = sectortype((f₁, f₂)) N₁ = numout((f₁, f₂)) - @assert FusionStyle(I) === UniqueFusion() (a, b, c), (f₁′, f₂′) = _bendright_treepair((f₁, f₂)) # compute the coefficient coeff₀ = sqrtdim(c) * invsqrtdim(a) f₁.isdual[N₁] && (coeff₀ *= conj(frobenius_schur_phase(dual(b)))) - coeff = coeff₀ * Bsymbol(a, b, c) + coeff = coeff₀ * only(Bsymbol(a, b, c)) return (f₁′, f₂′) => coeff end @@ -71,18 +70,9 @@ function bendright(src::FusionTreeBlock) N₁ = numout(src) N₂ = numin(src) @assert N₁ > 0 - uncoupled_dst = ( - TupleTools.front(src.uncoupled[1]), - (src.uncoupled[2]..., dual(src.uncoupled[1][N₁])), - ) - isdual_dst = ( - TupleTools.front(src.isdual[1]), - (src.isdual[2]..., !(src.isdual[1][N₁])), - ) - - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) - U = zeros(fusionscalartype(I), length(dst), length(src)) + + rows = TreeRows(fusiontreetype(I, N₁ - 1, N₂ + 1), length(src)) + U = zeros(fusionscalartype(I), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) (a, b, c), (f₁′, f₂′) = _bendright_treepair((f₁, f₂)) @@ -92,7 +82,7 @@ function bendright(src::FusionTreeBlock) end if FusionStyle(I) isa MultiplicityFreeFusion coeff = coeff₀ * Bsymbol(a, b, c) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = coeff else Bmat = Bsymbol(a, b, c) @@ -106,12 +96,12 @@ function bendright(src::FusionTreeBlock) iszero(coeff) && continue vertices₂ = N₂ > 0 ? (f₂.vertices..., ν) : () f₂′ = FusionTree(uncoupled₂, coupled₂, isdual₂, inner₂, vertices₂) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = coeff end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end @doc """ @@ -138,7 +128,6 @@ See also [`bendright`](@ref). """ bendleft function bendleft((f₁, f₂)::FusionTreePair) - @assert FusionStyle((f₁, f₂)) === UniqueFusion() (f₂′, f₁′), coeff = bendright((f₂, f₁)) return (f₁′, f₂′) => conj(coeff) end @@ -150,18 +139,9 @@ function bendleft(src::FusionTreeBlock) N₁ = numout(src) N₂ = numin(src) @assert N₂ > 0 - uncoupled_dst = ( - (src.uncoupled[1]..., dual(src.uncoupled[2][N₂])), - TupleTools.front(src.uncoupled[2]), - ) - isdual_dst = ( - (src.isdual[1]..., !(src.isdual[2][N₂])), - TupleTools.front(src.isdual[2]), - ) - - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) - U = zeros(fusionscalartype(I), length(dst), length(src)) + + rows = TreeRows(fusiontreetype(I, N₁ + 1, N₂ - 1), length(src)) + U = zeros(fusionscalartype(I), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) (a, b, c), (f₂′, f₁′) = _bendright_treepair((f₂, f₁)) @@ -171,7 +151,7 @@ function bendleft(src::FusionTreeBlock) end if FusionStyle(I) isa MultiplicityFreeFusion coeff = coeff₀ * Bsymbol(a, b, c) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = conj(coeff) else Bmat = Bsymbol(a, b, c) @@ -185,12 +165,12 @@ function bendleft(src::FusionTreeBlock) iszero(coeff) && continue vertices₁ = N₁ > 0 ? (f₁.vertices..., ν) : () f₁′ = FusionTree(uncoupled₁, coupled₁, isdual₁, inner₁, vertices₁) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = conj(coeff) end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end @doc """ @@ -236,17 +216,11 @@ function foldright((f₁, f₂)::FusionTreePair) end function foldright(src::FusionTreeBlock) - uncoupled_dst = ( - Base.tail(src.uncoupled[1]), - (dual(first(src.uncoupled[1])), src.uncoupled[2]...), - ) - isdual_dst = (Base.tail(src.isdual[1]), (!first(src.isdual[1]), src.isdual[2]...)) I = sectortype(src) N₁ = numout(src) @assert N₁ > 0 - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) + rows = TreeRows(fusiontreetype(I, N₁ - 1, numin(src) + 1), length(src)) f₁, f₂ = first(fusiontrees(src)) a::I = f₁.uncoupled[1] @@ -260,7 +234,7 @@ function foldright(src::FusionTreeBlock) c::I = f₁.coupled cache₃ = Dict((b, c) => Asymbol(a, b, c)) - U = zeros(eltype(coeff₁), length(dst), length(src)) + U = zeros(eltype(coeff₁), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) f₁′s, coeffs₁ = get!(cache₁, f₁) do multi_Fmove(f₁) @@ -280,12 +254,12 @@ function foldright(src::FusionTreeBlock) if isduala coeff *= κₐ end - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] += coeff end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end @@ -319,21 +293,12 @@ function foldleft((f₁, f₂)::FusionTreePair) end function foldleft(src::FusionTreeBlock) - uncoupled_dst = ( - (dual(first(src.uncoupled[2])), src.uncoupled[1]...), - Base.tail(src.uncoupled[2]), - ) - isdual_dst = ( - (!first(src.isdual[2]), src.isdual[1]...), - Base.tail(src.isdual[2]), - ) I = sectortype(src) N₁ = numin(src) N₂ = numout(src) @assert N₁ > 0 - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) + rows = TreeRows(fusiontreetype(I, N₂ + 1, N₁ - 1), length(src)) f₁, f₂ = first(fusiontrees(src)) a::I = f₂.uncoupled[1] @@ -347,7 +312,7 @@ function foldleft(src::FusionTreeBlock) c::I = f₂.coupled cache₃ = Dict((b, c) => Asymbol(a, b, c)) - U = zeros(eltype(coeff₂), length(dst), length(src)) + U = zeros(eltype(coeff₂), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) f₂′s, coeffs₂ = get!(cache₂, f₂) do multi_Fmove(f₂) @@ -367,12 +332,12 @@ function foldleft(src::FusionTreeBlock) if isduala coeff *= conj(κₐ) end - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] += coeff end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end # clockwise cyclic permutation while preserving (N₁, N₂): foldright & bendleft @@ -482,7 +447,7 @@ function _repartition_body(N) if N == 0 ex = quote T = fusionscalartype(sectortype(src)) - if FusionStyle(src) === UniqueFusion() + if src isa FusionTreePair return src => one(T) else U = copyto!(zeros(T, length(src), length(src)), LinearAlgebra.I) @@ -577,58 +542,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 40a2961b6..fef21b76b 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -117,6 +117,10 @@ const FusionTreePair{I, N₁, N₂} = Tuple{FusionTree{I, N₁}, FusionTree{I, N Collection of fusion-splitting tree pairs that share the same uncoupled sectors. Mostly internal structure to speed up non-`UniqueFusion` fusiontree manipulation computations. +Manipulating a block requires it to be complete, i.e. to contain all the fusion-splitting tree +pairs with its uncoupled sectors and duality flags, as `FusionTreeBlock{I}(uncoupled, isdual)` +does. The resulting block holds its trees in the order in which they are produced, rather than in +the canonical order of `FusionTreeBlock{I}(uncoupled, isdual)`. """ struct FusionTreeBlock{I, N₁, N₂, F <: FusionTreePair{I, N₁, N₂}} trees::Vector{F} @@ -231,9 +235,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 @@ -268,9 +275,37 @@ function treeindex_data((f₁, f₂)) return false end end -function treeindex_map(fs::FusionTreeBlock) - I = sectortype(fs) - return fusiontreedict(I)(treeindex_data(f) => ind for (ind, f) in enumerate(fusiontrees(fs))) + +# The trees of the destination block of a block manipulation, in the order in which they are +# produced, and the row of each of them in the transformation matrix. The source block has to be +# complete, such that the destination block has as many trees (`n`); the tree vector of a block +# that is no longer needed can be passed on as `buffer`. +struct TreeRows{F, D} + trees::Vector{F} + index::Union{D, Nothing} # `nothing`: few trees, the rows are found by a linear search + n::Int + # no default outer constructor, of which `D` would be unbound if `index === nothing` + TreeRows{F, D}(trees::Vector{F}, index, n::Int) where {F, D} = new{F, D}(trees, index, n) +end +function TreeRows(::Type{F}, n::Int, buffer = nothing) where {F} + trees = buffer isa Vector{F} ? empty!(buffer) : sizehint!(Vector{F}(), n) + K = Base.promote_op(treeindex_data, F) + index = n > 32 ? Dict{K, Int}() : nothing + return TreeRows{F, Dict{K, Int}}(trees, index, n) +end +function row!(rows::TreeRows, f) + key = treeindex_data(f) + i = if isnothing(rows.index) + findfirst(f′ -> treeindex_data(f′) == key, rows.trees) + else + get(rows.index, key, nothing) + end + isnothing(i) || return i + length(rows.trees) < rows.n || + throw(ArgumentError("manipulating a block requires it to be complete")) + push!(rows.trees, f) + isnothing(rows.index) || (rows.index[key] = length(rows.trees)) + return length(rows.trees) end 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..8cec768db 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) diff --git a/test/setup.jl b/test/setup.jl index 2030228a5..e928fdd11 100644 --- a/test/setup.jl +++ b/test/setup.jl @@ -9,12 +9,12 @@ export sectorlist, fast_sectorlist export default_spacelist, factorization_spacelist, ad_spacelist export VIBM, VIBMRepA4 export test_ad_rrule -export _isunitary, _isone +export _isunitary, _isone, _inorder using Random using Test: @test using TensorKit -using TensorKit: ℙ, PlanarTrivial +using TensorKit: ℙ, PlanarTrivial, FusionTreeBlock using TensorKitSectors using TensorOperations: IndexTuple, Index2Tuple using Base.Iterators: take, product @@ -149,6 +149,15 @@ _isunitary(x::Number; kwargs...) = isapprox(x * x', one(x); kwargs...) _isunitary(x; kwargs...) = isunitary(x; kwargs...) _isone(x; kwargs...) = isapprox(x, one(x); kwargs...) +# the rows of the transformation matrix `U` onto the trees of `dst`, in the order of the trees of +# `ref`, as block manipulations keep their trees in the order in which they are produced +_inorder(U::Number, dst, ref) = (@test dst == ref; U) +function _inorder(U::AbstractMatrix, dst::FusionTreeBlock, ref::FusionTreeBlock) + rows = indexin(fusiontrees(ref), fusiontrees(dst)) + @test length(dst) == length(ref) && !any(isnothing, rows) + return U[something.(rows, 1), :] +end + uniquefusionsectorlist = ( Z2Irrep, Z3Irrep, Z4Irrep, Z3Irrep ⊠ Z4Irrep, U1Irrep, diff --git a/test/symmetries/doubletree.jl b/test/symmetries/doubletree.jl index 180e19bdc..e6fd1ca95 100644 --- a/test/symmetries/doubletree.jl +++ b/test/symmetries/doubletree.jl @@ -57,15 +57,13 @@ using TensorKitSectors # single bend dst, U = @constinferred TK.bendright(src) dst2, U2 = @constinferred TK.bendleft(dst) - @test src == dst2 - @test _isone(U2 * U) + @test _isone(_inorder(U2 * U, dst2, src)) # double bend dst1, U1 = @constinferred TK.bendleft(src) dst2, U2 = @constinferred TK.bendleft(dst1) dst3, U3 = @constinferred TK.bendright(dst2) dst4, U4 = @constinferred TK.bendright(dst3) - @test src == dst4 - @test _isone(U4 * U3 * U2 * U1) + @test _isone(_inorder(U4 * U3 * U2 * U1, dst4, src)) if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) all_inds = (ntuple(identity, numout(src))..., reverse(ntuple(i -> i + numout(src), numin(src)))...) @@ -88,15 +86,13 @@ using TensorKitSectors # single bend dst, U = @constinferred TK.foldleft(src) dst2, U2 = @constinferred TK.foldright(dst) - @test src == dst2 - @test _isone(U2 * U) + @test _isone(_inorder(U2 * U, dst2, src)) # double bend dst1, U1 = @constinferred TK.foldright(src) dst2, U2 = @constinferred TK.foldright(dst1) dst3, U3 = @constinferred TK.foldleft(dst2) dst4, U4 = @constinferred TK.foldleft(dst3) - @test src == dst4 - @test _isone(U4 * U3 * U2 * U1) + @test _isone(_inorder(U4 * U3 * U2 * U1, dst4, src)) if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) all_inds = TupleTools.circshift((ntuple(identity, numout(src))..., reverse(ntuple(i -> i + numout(src), numin(src)))...), -2) @@ -121,7 +117,7 @@ using TensorKitSectors # @test _isunitary(U) dst′, U′ = repartition(dst, N) - @test _isone(U * U′) + @test _isone(U * _inorder(U′, dst′, src)) if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) all_inds = (ntuple(identity, numout(src))..., reverse(ntuple(i -> i + numout(src), numin(src)))...) @@ -152,11 +148,11 @@ using TensorKitSectors dst, U = @constinferred transpose(src, (p1, p2)) dst′, U′ = @constinferred transpose(dst, (ip1, ip2)) - @test _isone(U * U′) + @test _isone(U * _inorder(U′, dst′, src)) if BraidingStyle(I) isa Bosonic dst″, U″ = permute(src, (p1, p2)) - @test U″ ≈ U + @test _inorder(U″, dst″, dst) ≈ U end if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) @@ -217,7 +213,7 @@ using TensorKitSectors else dst′, U′ = @constinferred TensorKit.braid(dst, (ip1, ip2), (il1, il2)) end - @test _isone(U * U′) + @test _isone(U * _inorder(U′, dst′, src)) # check fusiontensor compatibility if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) @@ -235,6 +231,31 @@ using TensorKitSectors end end + FusionStyle(I) isa UniqueFusion || !(BraidingStyle(I) isa HasBraiding) || + @testset "Double fusion tree: braiding a block with a single tree" begin + a, b = random_fusion(I, Val(2)) + cs = filter(c -> Nsymbol(a, b, c) == 1, collect(a ⊗ b)) + if !isempty(cs) + src1 = FusionTreeBlock{I}(((a, b), (rand(cs),)), ((rand(Bool), rand(Bool)), (rand(Bool),))) + @test length(src1) == 1 + for n in 0:3 + p = (randperm(3)...,) + p1, p2 = p[1:n], p[(n + 1):3] + ip = invperm(p) + ilevels = TupleTools.getindices((1, 2, 3), p) + dst, U = @constinferred TensorKit.braid(src1, (p1, p2), ((1, 2), (3,))) + @test length(dst) == 1 && size(U) == (1, 1) + dst′, U′ = @constinferred TensorKit.braid(dst, (ip[1:2], ip[3:3]), (ilevels[1:n], ilevels[(n + 1):3])) + @test dst′ == src1 + @test _isone(U * U′) + if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) + A1 = fusiontensor(only(fusiontrees(src1))) + @test permutedims(A1, (p1..., p2...)) ≈ only(U) * fusiontensor(only(fusiontrees(dst))) + end + end + end + end + UnitStyle(I) isa SimpleUnit && @testset "Double fusion tree: blocks of a HomSpace" begin V = Vect[I](Dict(s => 1 for s in (out..., incoming))) W = V ⊗ V' ← V ⊗ V' # never empty: the unit couples both sides @@ -249,28 +270,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)