diff --git a/benchmark/TensorKitBenchmarks/indexmanipulations/IndexManipulationBenchmarks.jl b/benchmark/TensorKitBenchmarks/indexmanipulations/IndexManipulationBenchmarks.jl index 1529db44e..735738306 100644 --- a/benchmark/TensorKitBenchmarks/indexmanipulations/IndexManipulationBenchmarks.jl +++ b/benchmark/TensorKitBenchmarks/indexmanipulations/IndexManipulationBenchmarks.jl @@ -12,9 +12,9 @@ const all_parameters = TOML.parsefile(joinpath(@__DIR__, "benchparams.toml")) # permute! # -------- -function init_permute_tensors(T, W, p) +function init_permute_tensors(T, W, p, adjoint::Bool) C = randn(T, permute(W, p)) - A = randn(T, W) + A = adjoint ? randn(T, W')' : randn(T, W) return C, A end function benchmark_permute!(benchgroup, params::Dict) @@ -25,7 +25,9 @@ function benchmark_permute!(benchgroup, params::Dict) end return nothing end -function benchmark_permute!(bench; sigmas = nothing, T = "Float64", I = "Trivial", dims, p) +function benchmark_permute!( + bench; sigmas = nothing, T = "Float64", I = "Trivial", dims, p, adjoint = false + ) T_ = parse_type(T) I_ = parse_type(I) @@ -34,9 +36,10 @@ function benchmark_permute!(bench; sigmas = nothing, T = "Float64", I = "Trivial codomain = mapreduce(Base.Fix1(getindex, Vs), ⊗, p_[1]; init = one(eltype(Vs))) domain = mapreduce(Base.Fix1(getindex, Vs), ⊗, p_[2]; init = one(eltype(Vs))) - init() = init_permute_tensors(T_, codomain ← domain, p_) + init() = init_permute_tensors(T_, codomain ← domain, p_, adjoint) - bench[T, I, dims, sigmas, p] = @benchmarkable permute!(C, A, $p_) setup = ((C, A) = $init()) + key = adjoint ? (T, I, dims, sigmas, p, "adjoint") : (T, I, dims, sigmas, p) + bench[key...] = @benchmarkable permute!(C, A, $p_) setup = ((C, A) = $init()) return nothing end diff --git a/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml b/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml index 549072d23..05a12c590 100644 --- a/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml +++ b/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml @@ -10,6 +10,7 @@ I = "Z2Irrep" p = [[[2, 1], []]] dims = [[7264, 7264], [43408, 1216]] sigmas = [[0.5, 0.5]] +adjoint = [false, true] [[permute]] T = ["Float64"] @@ -24,3 +25,4 @@ I = "SU2Irrep" p = [[[1, 3], [2, 4]], [[4, 2, 3], [1]]] dims = [[48, 48, 48, 48]] sigmas = [[1.0, 1.0, 1.0, 1.0]] +adjoint = [false, true] diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 74a238fc2..36d5ea90d 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -24,14 +24,20 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed +- Index manipulations use a single kernel operating on subblocks addressed by position: `StridedSubblocks` (sector-independent views into the flat data of a `TensorMap`) or `TreeSubblocks` (any `AbstractTensorMap`, through `subblock`), both able to carry a lazy conjugation. The `TreeTransformer`s store the mapping between subblock positions and recoupling coefficients and are cached for every tensor type; conjugated and adjoint operands are handled through this mechanism instead of through `AdjointTensorMap` wrappers (internal) ([#516](https://github.com/QuantumKitHub/TensorKit.jl/issues/516)) + ### Deprecated ### Removed ### Fixed +- `braid!`, `permute!` and `transpose!` with a `BraidingTensor` source now use the cached fusion tree transformers ([#516](https://github.com/QuantumKitHub/TensorKit.jl/issues/516)) + ### Performance +- In-place `permute!`, `braid!` and `transpose!` with `AdjointTensorMap` sources or destinations, as well as `@tensor` expressions with `conj`, now use the same cached and sector-independent kernel as plain `TensorMap`s; other tensor types (e.g. `DiagonalTensorMap`) also use the cached fusion tree transformers, and `subblocks(::TensorMap)` iterates without hashing fusion trees ([#516](https://github.com/QuantumKitHub/TensorKit.jl/issues/516), [#519](https://github.com/QuantumKitHub/TensorKit.jl/pull/519), [#520](https://github.com/QuantumKitHub/TensorKit.jl/pull/520)) + ## [0.17.1](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.0...v0.17.1) - 2026-07-13 ### Added diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index c0f178cf0..4280a25e0 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -24,6 +24,7 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β)) @inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTreeDict}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTreeBlock}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true +@inline EnzymeRules.inactive_type(::Type{<:TensorKit.AbelianTreeTransformer}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.LRU}) = true diff --git a/src/spaces/homspace.jl b/src/spaces/homspace.jl index 8460f7578..f7e600427 100644 --- a/src/spaces/homspace.jl +++ b/src/spaces/homspace.jl @@ -169,6 +169,20 @@ See also [`degeneracystructure`](@ref), [`blockstructure`](@ref). """ subblockstructure(W::HomSpace) = Dictionary(fusiontrees(W), degeneracystructure(W).subblockstructure) +""" + fusiontreeindices(W::HomSpace) -> Dictionary + +Return a `Dictionary` mapping each fusion tree pair `(f₁, f₂)` to its position in +[`fusiontrees`](@ref)`(W)`, which coincides with its position in +[`subblockstructure`](@ref)`(W)` and in the subblocks of a `TensorMap` on `W`. + +See also [`fusiontrees`](@ref), [`subblockstructure`](@ref). +""" +function fusiontreeindices(W::HomSpace) + trees = fusiontrees(W) + return Dictionary(trees, 1:length(trees)) +end + """ fusionblocks(W::HomSpace) diff --git a/src/tensors/blockiterator.jl b/src/tensors/blockiterator.jl index 1fbe63b5e..a762ca35a 100644 --- a/src/tensors/blockiterator.jl +++ b/src/tensors/blockiterator.jl @@ -171,3 +171,73 @@ function Base.show(io::IO, mime::MIME"text/plain", iter::SubblockIterator) show_subblocks(io, mime, iter) return nothing end + +""" + struct StridedSubblocks{A <: DenseVector, N, F} + StridedSubblocks(t::TensorMap, [op = identity]) + +Sector-independent, integer-indexable collection of the subblocks of a `TensorMap`, as +`StridedView`s into its flat data vector. Subblock `i` corresponds to the `i`th fusion tree pair +in the canonical order of `fusiontrees(space(t))`, see also [`fusiontreeindices`](@ref). +The operation `op` (`identity` or `conj`) is applied lazily to every view, which allows +representing the subblocks of a conjugated tensor without materializing it. + +This is the data structure consumed by the index manipulation kernels, whose type does not +depend on the sectortype of `t`. +""" +const SubblockOp = Union{typeof(identity), typeof(conj)} +struct StridedSubblocks{A <: DenseVector, N, F <: SubblockOp} + data::A + structure::Vector{StridedStructure{N}} + op::F +end +Base.length(s::StridedSubblocks) = length(s.structure) +Base.firstindex(s::StridedSubblocks) = 1 +Base.lastindex(s::StridedSubblocks) = length(s) +Base.eltype(::Type{S}) where {S <: StridedSubblocks} = Core.Compiler.return_type(getindex, Tuple{S, Int}) + +Base.@propagate_inbounds function Base.getindex(s::StridedSubblocks, i::Int) + sz, str, offset = s.structure[i] + return StridedView(s.data, sz, str, offset, s.op) +end + +function Base.iterate(s::StridedSubblocks, i::Int = 1) + i > length(s) && return nothing + return @inbounds(s[i]), i + 1 +end + +storagetype(::Type{StridedSubblocks{A, N, F}}) where {A, N, F} = A + +""" + struct TreeSubblocks{TT <: AbstractTensorMap, I, F} + TreeSubblocks(t::AbstractTensorMap, [op = identity]) + +Integer-indexable collection of the subblocks of an arbitrary tensor `t`, where position `i` +refers to the `i`th fusion tree pair of `fusiontrees(space(t))` and the data is retrieved through +[`subblock`](@ref), with `op` (`identity` or `conj`) applied. This is the generic counterpart of +[`StridedSubblocks`](@ref) for tensor types that do not store their data in a flat vector. +""" +struct TreeSubblocks{TT <: AbstractTensorMap, I, F <: SubblockOp} + t::TT + trees::I + op::F +end +function TreeSubblocks(t::AbstractTensorMap, op::SubblockOp = identity) + return TreeSubblocks(t, fusiontrees(t), scalartype(t) <: Real ? identity : op) +end + +storagetype(::Type{TreeSubblocks{TT, I, F}}) where {TT, I, F} = storagetype(TT) + +Base.length(s::TreeSubblocks) = length(s.trees) +Base.firstindex(s::TreeSubblocks) = 1 +Base.lastindex(s::TreeSubblocks) = length(s) +Base.eltype(::Type{S}) where {S <: TreeSubblocks} = Core.Compiler.return_type(getindex, Tuple{S, Int}) + +Base.@propagate_inbounds function Base.getindex(s::TreeSubblocks, i::Int) + return s.op(subblock(s.t, gettokenvalue(s.trees, i))) +end + +function Base.iterate(s::TreeSubblocks, i::Int = 1) + i > length(s) && return nothing + return @inbounds(s[i]), i + 1 +end diff --git a/src/tensors/braidingtensor.jl b/src/tensors/braidingtensor.jl index add0df48d..cdf2f887b 100644 --- a/src/tensors/braidingtensor.jl +++ b/src/tensors/braidingtensor.jl @@ -195,16 +195,8 @@ end # Index manipulations # ------------------- has_shared_permute(t::BraidingTensor, ::Index2Tuple) = false -function add_transform!( - tdst::AbstractTensorMap, - tsrc::BraidingTensor, (p₁, p₂)::Index2Tuple, - fusiontreetransform, - α::Number, β::Number, backend::AbstractBackend... - ) - return add_transform!( - tdst, TensorMap(tsrc), (p₁, p₂), fusiontreetransform, α, β, - backend... - ) +function unwrap_adjoints(tdst, tsrc::BraidingTensor, p::Index2Tuple, levels, conjsrc::Bool, α, β) + return unwrap_adjoints(tdst, TensorMap(tsrc), p, levels, conjsrc, α, β) end function planarcontract!( diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index da1a62d84..cdfc01c0c 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -220,8 +220,17 @@ See also [`permute`](@ref) for creating a new tensor. backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator() ) @boundscheck spacecheck_transform(permute, tdst, tsrc, p) - levels = ntuple(identity, numind(tsrc)) - return @inbounds braid!(tdst, tsrc, p, levels, α, β, backend, allocator) + @timeit_debug GLOBAL_TIMER "permute!/braid!" begin + if has_array_view(tdst) && has_array_view(tsrc) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( + tdst[], tsrc[], p, false, α, β, backend, allocator + ) + return tdst + end + tdst′, tsrc′, p′, _, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β) + @inbounds _braid!(tdst′, tsrc′, p′, conjsrc, allind(tsrc′), α′, β′, backend, allocator) + end + return tdst end """ @@ -314,11 +323,10 @@ See also [`braid`](@ref) for creating a new tensor. ) return tdst end - levels1 = TupleTools.getindices(levels, codomainind(tsrc)) - levels2 = TupleTools.getindices(levels, domainind(tsrc)) - transformer = treebraider(tdst, tsrc, p, (levels1, levels2)) - return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) + tdst′, tsrc′, p′, levels′, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, levels, false, α, β) + @inbounds _braid!(tdst′, tsrc′, p′, conjsrc, levels′, α′, β′, backend, allocator) end + return tdst end """ @@ -394,9 +402,11 @@ end ) return tdst end - transformer = treetransposer(tdst, tsrc, p) - return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) + tdst′, tsrc′, p′, _, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β) + transformer = treetransposer(tdst′, tsrc′, p′, conjsrc) + @inbounds add_transform!(tdst′, tsrc′, p′, conjsrc, transformer, α′, β′, backend, allocator) end + return tdst end """ @@ -562,26 +572,82 @@ Base.@deprecate( # Kernel implementation # --------------------- +# Index manipulations are expressed as `tdst = β * tdst + α * permutedims(conjsrc ? conj(tsrc) : tsrc, p)`, +# with `p` indexing the legs of `tsrc`. Adjoint wrappers are absorbed into `conjsrc`, `p`, `levels` +# and the scalars, such that the actual work always happens on the subblocks of the parent tensors. + +# levels are attached to the legs: relabel them from the legs of `t'` to those of `t` +_adjoint_levels(t, ::Nothing) = nothing +_adjoint_levels(t, levels::IndexTuple) = TupleTools.getindices(levels, adjointtensorindices(t, allind(t))) + +_unwrap_source(tsrc::AbstractTensorMap, p, levels, conjsrc::Bool) = (tsrc, p, levels, conjsrc) +function _unwrap_source(tsrc::AdjointTensorMap, p, levels, conjsrc::Bool) + tp = parent(tsrc) + return (tp, adjointtensorindices(tsrc, p), _adjoint_levels(tp, levels), !conjsrc) +end + +_unwrap_destination(tdst::AbstractTensorMap, p, conjsrc::Bool, α, β) = (tdst, p, conjsrc, α, β) +function _unwrap_destination(tdst::AdjointTensorMap, p, conjsrc::Bool, α, β) + return (parent(tdst), (p[2], p[1]), !conjsrc, conj(α), conj(β)) +end + +""" + unwrap_adjoints(tdst, tsrc, p, levels, conjsrc::Bool, α, β) -> (tdst′, tsrc′, p′, levels′, conjsrc′, α′, β′) + +Rewrite the operation `tdst = β * tdst + α * braid(conjsrc ? conj(tsrc) : tsrc, p, levels)` such that +neither `tdst′` nor `tsrc′` is an `AdjointTensorMap`, by absorbing the adjoints into the conjugation +flag, the permutation, the `levels` (which may be `nothing`) and the scalars. +""" +function unwrap_adjoints(tdst, tsrc, p::Index2Tuple, levels, conjsrc::Bool, α, β) + tsrc′, p′, levels′, conjsrc′ = _unwrap_source(tsrc, p, levels, conjsrc) + tdst′, p″, conjsrc″, α′, β′ = _unwrap_destination(tdst, p′, conjsrc′, α, β) + return (tdst′, tsrc′, p″, levels′, conjsrc″, α′, β′) +end + +# shared by `permute!`, `braid!` and `TO.tensoradd!` after the adjoints have been unwrapped +@propagate_inbounds function _braid!( + tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator + ) + @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) + transformer = treebraider(tdst, tsrc, p, conjsrc, levels) + return @inbounds add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) +end + +# space check for `tdst = permutedims(conjsrc ? conj(tsrc) : tsrc, p)` +function spacecheck_transform(f, tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool) + Vsrc′, p′ = transform_source(space(tsrc), p, conjsrc) + return spacecheck_transform(f, space(tdst), Vsrc′, p′) +end + +""" + add_transform!(tdst, tsrc, p, conjsrc::Bool, transformer, α, β, backend, allocator) -> tdst + +Compute `tdst = β * tdst + α * permutedims(conjsrc ? conj(tsrc) : tsrc, p)`, where `p` indexes the legs +of `tsrc`, using the fusion tree transformation encoded in `transformer` (see [`TreeTransformer`](@ref)). +""" @propagate_inbounds function add_transform!( - tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, transformer, + tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, transformer, α::Number, β::Number, backend, allocator ) - @boundscheck spacecheck_transform(permute, tdst, tsrc, p) + @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) - if p[1] === codomainind(tsrc) && p[2] === domainind(tsrc) + if !conjsrc && p[1] === codomainind(tsrc) && p[2] === domainind(tsrc) add!(tdst, tsrc, α, β) else - p2 = (linearize(p), ()) + p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels if has_array_view(tdst) && has_array_view(tsrc) @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( - tdst[], tsrc[], p2, false, α, β, backend, allocator + tdst[], tsrc[], p2, conjsrc, α, β, backend, allocator ) else ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 - if tdst isa TensorMap && tsrc isa TensorMap # unpack data fields to avoid specializing - add_transform_kernel!(tdst.data, tsrc.data, p2, transformer, α, β, backend, allocator, ntasks) + # resolve the conjugation flag into the view type here, with a statically typed call per branch + if conjsrc + dst, src = _transform_subblocks(tdst, tsrc, transformer, conj) + add_transform_kernel!(dst, src, p2, transformer, α, β, backend, allocator, ntasks) else - add_transform_kernel!(tdst, tsrc, p2, transformer, α, β, backend, allocator, ntasks) + dst, src = _transform_subblocks(tdst, tsrc, transformer, identity) + add_transform_kernel!(dst, src, p2, transformer, α, β, backend, allocator, ntasks) end end end @@ -589,6 +655,14 @@ Base.@deprecate( return tdst end +# TensorMaps address their flat data directly, other tensor types go through `subblock` +function _transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer, op) + return StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src, op) +end +function _transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer, op) + return TreeSubblocks(tdst), TreeSubblocks(tsrc, op) +end + function use_threaded_transform(t::TensorMap, transformer) return get_num_transformer_threads() > 1 && length(t.data) > Strided.MINTHREADLENGTH end @@ -596,149 +670,50 @@ function use_threaded_transform(t::AbstractTensorMap, transformer) return get_num_transformer_threads() > 1 && dim(space(t)) > Strided.MINTHREADLENGTH end +# The kernel operates on the subblocks addressed by position, so that for `TensorMap`s this only +# depends on `numind`, `eltype` and the transformer data, not on the sectortype. +const TransformSubblocks = Union{StridedSubblocks, TreeSubblocks} function add_transform_kernel!( - tdst, tsrc, p, transformer, α, β, backend, allocator, ntasks::Int - ) - I = sectortype(tdst) - if FusionStyle(I) === UniqueFusion() - @timeit_debug GLOBAL_TIMER "dense: tensoradd" begin - taskforeach(fusiontrees(tsrc), ntasks) do (f₁, f₂) - (f₁′, f₂′), coeff = transformer((f₁, f₂)) - @inbounds TO.tensoradd!( - tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β, backend, allocator - ) - end - end - return nothing - end - - @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" begin - fblocks = fusionblocks(tsrc) - bufsize = buffersize(tsrc, fblocks) - end - - # One max-sized workspace per task (a single one that is reused by all blocks when - # serial), allocated on the calling thread before any task spawns, so that also - # allocators that are not thread-safe can be used. - cp = TO.allocator_checkpoint!(allocator) - @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ - TO.tensoralloc(storagetype(tdst), bufsize, Val(true), allocator) - for _ in 1:clamp(length(fblocks), 1, ntasks) - ] - taskforeach(fblocks, buffers) do src, buffer - _add_transform_block!( - tdst, tsrc, p, src, transformer, buffer, α, β, backend, allocator - ) - end - foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) - TO.allocator_reset!(allocator, cp) - return nothing -end - -# TensorMap specializations: operate directly on the flat data vector to avoid -# repeated specialization -- this only depends on `numind` and `eltype`. -function add_transform_kernel!( - data_dst::DenseVector, data_src::DenseVector, p, transformer::AbelianTreeTransformer, - α, β, backend, allocator, ntasks::Int - ) - @timeit_debug GLOBAL_TIMER "dense: tensoradd" begin - taskforeach(transformer.data, ntasks) do (coeff, struct_dst, struct_src) - TO.tensoradd!( - StridedView(data_dst, struct_dst...), StridedView(data_src, struct_src...), - p, false, α * coeff, β, backend, allocator - ) - end - end - return nothing -end -function add_transform_kernel!( - data_dst::DenseVector, data_src::DenseVector, p, transformer::GenericTreeTransformer, + dst::TransformSubblocks, src::TransformSubblocks, p, + transformer::Union{AbelianTreeTransformer, GenericTreeTransformer}, α, β, backend, allocator, ntasks::Int ) bufsize = buffersize(transformer) - - # One max-sized workspace per task (a single one that is reused by all blocks when - # serial), allocated on the calling thread before any task spawns, so that also - # allocators that are not thread-safe can be used. - cp = TO.allocator_checkpoint!(allocator) - @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ - TO.tensoralloc(typeof(data_dst), bufsize, Val(true), allocator) - for _ in 1:clamp(length(transformer.data), 1, ntasks) - ] - taskforeach(transformer.data, buffers) do subtransformer, buffer - _add_transform_block!( - data_dst, data_src, p, subtransformer, buffer, α, β, backend, allocator - ) - end - foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) - TO.allocator_reset!(allocator, cp) - return nothing -end - -function _add_transform_block!( - tdst, tsrc, p, src::FusionTreeBlock, transformer, buffer, - α, β, backend, allocator - ) - dst, U = transformer(src) - - if length(src) == 1 # Degenerate block with a single tree: no matmul needed. - (f₁, f₂) = only(fusiontrees(src)) - (f₁′, f₂′) = only(fusiontrees(dst)) - @timeit_debug GLOBAL_TIMER "dense: tensoradd" @inbounds TO.tensoradd!( - tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * only(U), β, backend, allocator - ) - else # Multi-tree block: pack → recoupling matmul → unpack. - rows, cols = size(U) - sz_src = size(tsrc[first(fusiontrees(src))...]) - blocksize = prod(sz_src) - # the buffer was sized assuming a square recoupling matrix - rows == cols || throw(DimensionMismatch(lazy"recoupling matrix is not square: $(size(U))")) - ptriv = (ntuple(identity, length(sz_src)), ()) - buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) - buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) - - # 1. Extract: copy each source block into column i of buffer_src as a flat vector, - # using a trivial permutation so the layout is canonical before the matmul. - @timeit_debug GLOBAL_TIMER "dense: pack" @inbounds for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) - TO.tensoradd!( - sreshape(view(buffer_src, :, i), sz_src), tsrc[f₁, f₂], - ptriv, false, One(), Zero(), backend, allocator - ) + if bufsize == 0 # no recoupling needed: every block consists of a single tree + taskforeach(transformer.data, ntasks) do (U, inds_dst, inds_src) + _add_transform_block!(dst, src, p, U, inds_dst, inds_src, nothing, α, β, backend, allocator) end - - # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear - # combination of input trees weighted by the recoupling coefficients). - @timeit_debug GLOBAL_TIMER "dense: recouple mul!" begin - U′ = _adapt_recoupling(storagetype(tdst), U) - mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) - end - - # 3. Insert: scatter column i of buffer_dst into the destination, applying the - # actual index permutation p in the same tensoradd! call. - @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (i, (f₃, f₄)) in enumerate(fusiontrees(dst)) - TO.tensoradd!( - tdst[f₃, f₄], sreshape(view(buffer_dst, :, i), sz_src), - p, false, One(), β, backend, allocator - ) + else + # One max-sized workspace per task (a single one that is reused by all blocks when + # serial), allocated on the calling thread before any task spawns, so that also + # allocators that are not thread-safe can be used. + cp = TO.allocator_checkpoint!(allocator) + @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ + TO.tensoralloc(storagetype(dst), bufsize, Val(true), allocator) + for _ in 1:clamp(length(transformer.data), 1, ntasks) + ] + taskforeach(transformer.data, buffers) do (U, inds_dst, inds_src), buffer + _add_transform_block!(dst, src, p, 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 +# `U` is either a scalar coefficient with integer positions (abelian), or a recoupling matrix +# with vectors of positions (generic). function _add_transform_block!( - data_dst::DenseVector, data_src::DenseVector, p, - ((U, (sz_dst, structs_dst), (sz_src, structs_src)))::GenericTransformerData, - buffer, α, β, backend, allocator + dst::TransformSubblocks, src::TransformSubblocks, p, U, inds_dst, inds_src, buffer, + α, β, backend, allocator ) - if length(U) == 1 # Degenerate block with a single tree: no matmul needed. - coeff = only(U) - @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( - StridedView(data_dst, sz_dst, only(structs_dst)...), - StridedView(data_src, sz_src, only(structs_src)...), - p, false, α * coeff, β, backend, allocator + if length(U) == 1 # single tree: no matmul needed + @timeit_debug GLOBAL_TIMER "dense: tensoradd" @inbounds TO.tensoradd!( + dst[only(inds_dst)], src[only(inds_src)], p, false, α * only(U), β, backend, allocator ) else # Multi-tree block: pack → recoupling matmul → unpack. rows, cols = size(U) + sz_src = size(@inbounds(src[first(inds_src)])) blocksize = prod(sz_src) ptriv = (ntuple(identity, length(sz_src)), ()) buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) @@ -746,9 +721,9 @@ function _add_transform_block!( # 1. Extract: copy each source block into column i of buffer_src as a flat vector, # using a trivial permutation so the layout is canonical before the matmul. - @timeit_debug GLOBAL_TIMER "dense: pack" @inbounds for (i, struct_src_i) in enumerate(structs_src) + @timeit_debug GLOBAL_TIMER "dense: pack" @inbounds for (i, isrc) in enumerate(inds_src) TO.tensoradd!( - sreshape(view(buffer_src, :, i), sz_src), StridedView(data_src, sz_src, struct_src_i...), + sreshape(view(buffer_src, :, i), sz_src), src[isrc], ptriv, false, One(), Zero(), backend, allocator ) end @@ -756,15 +731,15 @@ function _add_transform_block!( # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear # combination of input trees weighted by the recoupling coefficients). @timeit_debug GLOBAL_TIMER "dense: recouple mul!" begin - U′ = _adapt_recoupling(typeof(data_dst), U) + U′ = _adapt_recoupling(storagetype(dst), U) mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) end - # 3. Insert: scatter column i of buffer_dst into the destination, applying the + # 3. Insert: scatter column j of buffer_dst into the destination, applying the # actual index permutation p in the same tensoradd! call. - @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (i, struct_dst_i) in enumerate(structs_dst) + @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (j, idst) in enumerate(inds_dst) TO.tensoradd!( - StridedView(data_dst, sz_dst, struct_dst_i...), sreshape(view(buffer_dst, :, i), sz_src), + dst[idst], sreshape(view(buffer_dst, :, j), sz_src), p, false, One(), β, backend, allocator ) end diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index ed5682367..d42dd320c 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -496,6 +496,26 @@ end # Getting and setting the data at the subblock level # -------------------------------------------------- +function StridedSubblocks(t::TensorMap, op::SubblockOp = identity) + return StridedSubblocks(t, degeneracystructure(space(t)).subblockstructure, op) +end +function StridedSubblocks(t::TensorMap, structure::Vector{<:StridedStructure}, op::SubblockOp = identity) + return StridedSubblocks(t.data, structure, scalartype(t) <: Real ? identity : op) +end + +# iterate the subblock views in canonical order alongside the fusion trees, without hashing +function subblocks(t::TensorMap) + return sectortype(t) === Trivial ? SubblockIterator(t, fusiontrees(t)) : + SubblockIterator(t, StridedSubblocks(t)) +end +function Base.iterate(iter::SubblockIterator{<:TensorMap, <:StridedSubblocks}, i::Int = 1) + i > length(iter.structure) && return nothing + @inbounds begin + f = gettokenvalue(fusiontrees(iter.t), i) + return f => iter.structure[i], i + 1 + end +end + function subblock( t::TensorMap{T, S, N₁, N₂}, (f₁, f₂)::Tuple{FusionTree{I, N₁}, FusionTree{I, N₂}} ) where {T, S, N₁, N₂, I <: Sector} diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 4cb274cf3..6654cdc9c 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -51,13 +51,8 @@ function TO.tensoradd!( TO.tensoradd!(C[], A[], pA, conjA, α, β, backend, allocator) return C end - if conjA - A′ = adjoint(A) - pA′ = adjointtensorindices(A, _canonicalize(pA, C)) - permute!(C, A′, pA′, α, β, backend, allocator) - else - permute!(C, A, _canonicalize(pA, C), α, β, backend, allocator) - end + tdst, tsrc, p, _, conjA′, α′, β′ = unwrap_adjoints(C, A, _canonicalize(pA, C), nothing, conjA, α, β) + _braid!(tdst, tsrc, p, conjA′, allind(tsrc), α′, β′, backend, allocator) return C end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 00c0ec24b..d25349ad0 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -2,101 +2,111 @@ TreeTransformer Supertype for structures containing the data for a tree transformation. + +The transformers only store how subblocks map onto each other in terms of their positions in +[`StridedSubblocks`](@ref) (the canonical order of `fusiontrees`), together with the recoupling +coefficients, and are therefore independent of the sectortype once constructed. The +transformation is that of `permutedims(op(tsrc), p)` where `p` indexes the legs of `tsrc` itself and +`op` is either `identity` or `conj`; in the latter case the fusion trees that are transformed are +those of `adjoint(space(tsrc))`, which read the subblocks of `tsrc` with the tree pair swapped. """ abstract type TreeTransformer end -struct TrivialTreeTransformer <: TreeTransformer end +# (coefficient, destination position, source position) +const AbelianTransformerData{T} = Tuple{T, Int, Int} -const AbelianTransformerData{T, N} = Tuple{T, StridedStructure{N}, StridedStructure{N}} +""" + AbelianTreeTransformer{T, N} <: TreeTransformer +Tree transformation for `UniqueFusion` sectors, where every source subblock maps onto a single +destination subblock with a scalar coefficient, stored as `(coeff, idst, isrc)`. 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 AbelianTreeTransformer{T, N} <: TreeTransformer - data::Vector{AbelianTransformerData{T, N}} + data::Vector{AbelianTransformerData{T}} + structure_dst::Vector{StridedStructure{N}} + structure_src::Vector{StridedStructure{N}} +end + +# (recoupling matrix, destination positions, source positions): U[j, i] maps source i onto destination j +const GenericTransformerData{T} = Tuple{Matrix{T}, Vector{Int}, Vector{Int}} + +""" + GenericTreeTransformer{T, N} <: TreeTransformer + +Tree transformation for sectors with multiple fusion channels, where the subblocks of a +[`FusionTreeBlock`](@ref) map onto the subblocks of the transformed block through a recoupling +matrix, stored as `(U, inds_dst, inds_src)`. 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 GenericTreeTransformer{T, N} <: TreeTransformer + data::Vector{GenericTransformerData{T}} + structure_dst::Vector{StridedStructure{N}} + structure_src::Vector{StridedStructure{N}} end -function AbelianTreeTransformer(transform, p, Vdst, Vsrc) +# Space and permutation as seen by the tree manipulations: for a conjugated source these are +# those of the adjoint space, whose tree pairs `(f₁, f₂)` address the subblock `(f₂, f₁)` of the source. +function transform_source(Vsrc::HomSpace, p::Index2Tuple, conjsrc::Bool) + return conjsrc ? (adjoint(Vsrc), adjointtensorindices(Vsrc, p)) : (Vsrc, p) +end +source_tree((f₁, f₂)::FusionTreePair, conjsrc::Bool) = conjsrc ? (f₂, f₁) : (f₁, f₂) + +function AbelianTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) t₀ = Base.time() - permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting.")) - fts_src = subblockstructure(Vsrc) - fts_dst = subblockstructure(Vdst) - L = length(fts_src) + Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) + spacecheck_transform(permute, Vdst, Vsrc′, p′) + srcindices = fusiontreeindices(Vsrc) + dstindices = fusiontreeindices(Vdst) + trees_src = fusiontrees(Vsrc′) + L = length(trees_src) T = sectorscalartype(sectortype(Vdst)) - N = numind(Vsrc) - data = Vector{Tuple{T, StridedStructure{N}, StridedStructure{N}}}(undef, L) + data = Vector{AbelianTransformerData{T}}(undef, L) - @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (i, (f_src, stridestructure_src)) in enumerate(pairs(fts_src)) + @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (i, f_src) in enumerate(trees_src) f_dst, coeff = transform(f_src) - stridestructure_dst = fts_dst[f_dst] - data[i] = (coeff, stridestructure_dst, stridestructure_src) + data[i] = (coeff, dstindices[f_dst], srcindices[source_tree(f_src, conjsrc)]) end - transformer = AbelianTreeTransformer(data) - - # sort by (approximate) weight to facilitate multi-threading strategies - # sort!(transformer) + structure_dst = degeneracystructure(Vdst).subblockstructure + structure_src = degeneracystructure(Vsrc).subblockstructure + transformer = AbelianTreeTransformer(data, structure_dst, structure_src) Δt = Base.time() - t₀ - - @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", nblocks = L, Δt) + @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = L, Δt) return transformer end -const GenericTransformerData{T, N} = Tuple{ - Matrix{T}, - Tuple{NTuple{N, Int}, Vector{Tuple{NTuple{N, Int}, Int}}}, - Tuple{NTuple{N, Int}, Vector{Tuple{NTuple{N, Int}, Int}}}, -} - -struct GenericTreeTransformer{T, N} <: TreeTransformer - data::Vector{GenericTransformerData{T, N}} -end - -""" - buffersize(transformer::GenericTreeTransformer) -> Int - buffersize(t::AbstractTensorMap, fblocks) -> Int - -Compute the workspace size required to pack, recouple and unpack the largest multi-tree -block, i.e. `prod(sz_src) * (rows + cols)` where `(rows, cols) = size(U)` is the size of -the recoupling matrix. The second form computes this from the fusion-tree blocks of a -tensor instead, using the fact that the recoupling matrix is square. -""" -function buffersize(transformer::GenericTreeTransformer) - return maximum(transformer.data; init = 0) do (U, _, (sz_src, _)) - return length(U) == 1 ? 0 : prod(sz_src) * sum(size(U)) - end -end -function buffersize(t::AbstractTensorMap, fblocks) - return maximum(fblocks; init = 0) do src - n = length(src) - return n == 1 ? 0 : 2 * n * length(t[first(fusiontrees(src))...]) - end -end - -function GenericTreeTransformer(transform, p, Vdst, Vsrc) +function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) t₀ = Base.time() - permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting.")) - fusionstructure_dst = subblockstructure(Vdst) - fusionstructure_src = subblockstructure(Vsrc) - I = sectortype(Vsrc) - T = sectorscalartype(I) - N = numind(Vdst) - N₁ = numout(Vsrc) - N₂ = numin(Vsrc) - - fblocks = @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" fusionblocks(Vsrc) + Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) + spacecheck_transform(permute, Vdst, Vsrc′, p′) + srcindices = fusiontreeindices(Vsrc) + dstindices = fusiontreeindices(Vdst) + structure_dst = degeneracystructure(Vdst).subblockstructure + structure_src = degeneracystructure(Vsrc).subblockstructure + T = sectorscalartype(sectortype(Vsrc)) + + fblocks = @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" fusionblocks(Vsrc′) nblocks = length(fblocks) - data = Vector{GenericTransformerData{T, N}}(undef, nblocks) + data = Vector{GenericTransformerData{T}}(undef, nblocks) + weights = Vector{Int}(undef, nblocks) nthreads = get_num_manipulation_threads() @timeit_debug GLOBAL_TIMER "symmetry: recoupling matrices" begin taskforeach(1:nblocks, nthreads) do i fs_src = fblocks[i] fs_dst, U = transform(fs_src) - @timeit_debug GLOBAL_TIMER "bookkeeping: repack" begin - sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, fusiontrees(fs_src)) - sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, fusiontrees(fs_dst)) + @timeit_debug GLOBAL_TIMER "bookkeeping: subblock positions" begin + inds_src = map(f -> srcindices[source_tree(f, conjsrc)], fusiontrees(fs_src)) + inds_dst = map(f -> dstindices[f], fusiontrees(fs_dst)) end - data[i] = U, (sz_dst, newstructs_dst), (sz_src, newstructs_src) + data[i] = (U, inds_dst, inds_src) + # cost model: L input blocks each going to L output blocks of a given length + weights[i] = length(U) * prod(structure_dst[first(inds_dst)][1]) @debug( lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", @@ -104,96 +114,84 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc) ) end end - transformer = GenericTreeTransformer{T, N}(data) # sort by (approximate) weight to facilitate multi-threading strategies - @timeit_debug GLOBAL_TIMER "bookkeeping: sort" sort!(transformer) + @timeit_debug GLOBAL_TIMER "bookkeeping: sort" Base.permute!(data, sortperm(weights; rev = true)) + transformer = GenericTreeTransformer(data, structure_dst, structure_src) Δt = Base.time() - t₀ - @debug( - lazy"TreeTransformer for $Vsrc to $Vdst via $p", - nblocks = length(transformer.data), - sz_median = size(transformer.data[cld(end, 2)][1], 1), - sz_max = size(transformer.data[1][1], 1), + lazy"TreeTransformer for $Vsrc to $Vdst via $p", conjsrc, + nblocks = nblocks, + sz_median = nblocks > 0 ? size(data[cld(end, 2)][1], 1) : 0, + sz_max = nblocks > 0 ? size(data[1][1], 1) : 0, Δt ) return transformer end -function repack_transformer_structure(structures::Dictionary, trees) - sz = structures[first(trees)][1] - strides_offsets = map(trees) do f - _, stride, offset = structures[f] - return stride, offset +""" + buffersize(transformer::TreeTransformer) -> Int + +Compute the workspace size required to pack, recouple and unpack the largest multi-tree +block, i.e. `prod(sz_src) * (rows + cols)` where `(rows, cols) = size(U)` is the size of +the recoupling matrix. +""" +buffersize(::AbelianTreeTransformer) = 0 +function buffersize(transformer::GenericTreeTransformer) + structure_src = transformer.structure_src + return maximum(transformer.data; init = 0) do (U, _, inds_src) + return length(U) == 1 ? 0 : prod(structure_src[first(inds_src)][1]) * sum(size(U)) end - return sz, strides_offsets end - function treetransformertype(Vdst, Vsrc) I = sectortype(Vdst) - I === Trivial && return TrivialTreeTransformer - T = sectorscalartype(I) N = numind(Vdst) return FusionStyle(I) == UniqueFusion() ? AbelianTreeTransformer{T, N} : GenericTreeTransformer{T, N} end function TreeTransformer( - transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S} + transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S}, conjsrc::Bool ) where {S} - permute(Vsrc, p) == Vdst || - throw(SpaceMismatch("Incompatible spaces for permuting")) - I = sectortype(Vdst) - I === Trivial && return TrivialTreeTransformer() - return FusionStyle(I) == UniqueFusion() ? - AbelianTreeTransformer(transform, p, Vdst, Vsrc) : - GenericTreeTransformer(transform, p, Vdst, Vsrc) + AbelianTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) : + GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) end # braid is special because it has levels -function treebraider(::AbstractTensorMap, ::AbstractTensorMap, p::Index2Tuple, levels) - return fusiontreetransform(f) = braid(f, p, levels) -end -function treebraider(tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple, levels) - return treebraider(space(tdst), space(tsrc), p, levels) +function treebraider( + tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple + ) + return treebraider(space(tdst), space(tsrc), p, conjsrc, levels) end @cached function treebraider( - Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, levels + Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple )::treetransformertype(Vdst, Vsrc) - fusiontreebraider(f) = braid(f, p, levels) - return TreeTransformer(fusiontreebraider, p, Vdst, Vsrc) + Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) + # levels are attached to the legs, so they follow the same relabeling as the permutation + levels′ = conjsrc ? TupleTools.getindices(levels, adjointtensorindices(Vsrc′, allind(Vsrc′))) : levels + levels″ = (TupleTools.getindices(levels′, codomainind(Vsrc′)), TupleTools.getindices(levels′, domainind(Vsrc′))) + fusiontreebraider(f) = braid(f, p′, levels″) + return TreeTransformer(fusiontreebraider, p, Vdst, Vsrc, conjsrc) end -function treetransposer(::AbstractTensorMap, ::AbstractTensorMap, p::Index2Tuple) - return fusiontreetransform(f) = transpose(f, p) -end -function treetransposer(tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple) - return treetransposer(space(tdst), space(tsrc), p) +function treetransposer(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool) + return treetransposer(space(tdst), space(tsrc), p, conjsrc) end @cached function treetransposer( - Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple + Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool )::treetransformertype(Vdst, Vsrc) - fusiontreetransform(f) = transpose(f, p) - return TreeTransformer(fusiontreetransform, p, Vdst, Vsrc) + Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) + fusiontreetransform(f) = transpose(f, p′) + return TreeTransformer(fusiontreetransform, p, Vdst, Vsrc, conjsrc) end # default cachestyle is GlobalLRUCache -# Sorting based on cost model -# --------------------------- -function Base.sort!( - transformer::Union{AbelianTreeTransformer, GenericTreeTransformer}; - by = _transformer_weight, rev::Bool = true - ) - sort!(transformer.data; by, rev) - return transformer -end - # For CPU arrays the recoupling matrix can be used as is, also when the scalar types # do not match, since Strided handles mixed-eltype mul! without the copy that # Adapt.adapt would make (which additionally dispatches dynamically). Other storage @@ -201,16 +199,3 @@ end # TODO: transformers with dedicated storagetypes _adapt_recoupling(::Type{<:Array}, U::Matrix) = StridedView(U) _adapt_recoupling(::Type{A}, U::Matrix) where {A} = Adapt.adapt(A, StridedView(U)) - -function _transformer_weight((coeff, struct_dst, struct_src)::AbelianTransformerData) - return prod(struct_dst[1]) -end - -# Cost model for transforming a set of subblocks with fixed uncoupled sectors: -# L x L x length(subblock) where L is the number of subblocks -# this is L input blocks each going to L output blocks of given length -# Note that it might be the case that the permutations are dominant, in which case the -# actual cost model would scale like L x length(subblock) -function _transformer_weight((mat, structs_dst, structs_src)::GenericTransformerData) - return length(mat) * prod(structs_dst[1]) -end diff --git a/test/tensors/indexmanipulations.jl b/test/tensors/indexmanipulations.jl index c4b0d924d..025b58010 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -129,6 +129,76 @@ for V in spacelist @tensor tb[a, b] := flip(t1, (1, 3))[x, y, a, z] * flip(t2, (2, 4))[y, b, z, x] @test flip(ta, (1, 2)) ≈ tb end + symmetricbraiding && @timedtestset "Permutations: adjoint operands" begin + W = V1 ⊗ V2 ← V3 ⊗ V4 + for T in (Float64, ComplexF64) + t = rand(T, W) + tref = copy(t') # genuine TensorMap on W' + α = T <: Complex ? T(2.1 + 0.3im) : T(2.1) + β = T <: Complex ? T(-0.7im) : T(-0.7) + for p in (((1, 2), (3, 4)), ((2, 3), (4, 1)), ((3,), (1, 2, 4)), ((), (4, 3, 2, 1))) + tdst = rand(T, permute(space(tref), p)) + ref = permute!(copy(tdst), tref, p, α, β) + # adjoint source + @test permute!(copy(tdst), t', p, α, β) ≈ ref + @test @constinferred(permute!(copy(tdst), t', p)) ≈ permute(tref, p) + # adjoint destination + D = copy(tdst') + permute!(D', tref, p, α, β) + @test D' ≈ ref + # adjoint source and destination + D = copy(tdst') + permute!(D', t', p, α, β) + @test D' ≈ ref + end + p = ((2, 4), (1, 3)) # cyclic + tdst = rand(T, transpose(space(tref), p)) + ref = transpose!(copy(tdst), tref, p, α, β) + @test transpose!(copy(tdst), t', p, α, β) ≈ ref + D = copy(tdst') + transpose!(D', t', p, α, β) + @test D' ≈ ref + + # conjugation through TensorOperations + A = rand(T, W) + Aref = copy(A') + @tensor C[a, b; c, d] := conj(A[c, a; d, b]) + @test C ≈ permute(Aref, ((4, 2), (3, 1))) + B = rand(T, W) + @tensor C2[a, b; c, d] := conj(A[x, y; a, b]) * B[x, y; c, d] + @tensor C2ref[a, b; c, d] := Aref[a, b; x, y] * B[x, y; c, d] + @test C2 ≈ C2ref + if BraidingStyle(I) isa Bosonic && hasfusiontensor(I) + @test convert(Array, C) ≈ permutedims(conj(convert(Array, A)), (2, 4, 1, 3)) + end + end + end + symmetricbraiding && @timedtestset "Permutations: BraidingTensor source" begin + τ = BraidingTensor(V1, V2) + p = ((2, 1), (3, 4)) + levels = (1, 2, 3, 4) + tdst = rand(ComplexF64, braid(space(τ), p, levels)) + @test braid!(copy(tdst), τ, p, levels) ≈ braid!(copy(tdst), TensorMap(τ), p, levels) + end + @timedtestset "Adjoint operands: isometry" begin + # independent of the adjoint convention: a wrongly conjugated recoupling matrix U + # would break dot(U * x', U * y) == dot(x', y) whenever U is genuinely complex + W = V1 ⊗ V2 ← V3 ⊗ V4 + x = rand(ComplexF64, W) + y = rand(ComplexF64, W') + if hasbraiding + p = ((2,), (1, 3, 4)) + levels = (1, 3, 2, 4) + bx = braid(x', p, levels) + by = braid(y, p, levels) + @test dot(bx, by) ≈ dot(x', y) + D = similar(y, braid(space(y), p, levels)') + braid!(D', x', p, levels) + @test dot(D', by) ≈ dot(x', y) + end + pc = ((2, 4), (1, 3)) + @test dot(transpose(x', pc), transpose(y, pc)) ≈ dot(x', y) + end hasbraiding && !symmetricbraiding && @timedtestset "Braid AdjointTensorMap: adjoint identity" begin t = rand(ComplexF64, V1 ⊗ V2 ← V3) p = ((2,), (1, 3)) @@ -136,6 +206,18 @@ for V in spacelist t1 = copy(braid(t', p, levels)) t2 = braid(copy(t'), p, levels) @test t1 ≈ t2 + + tref = copy(t') + α, β = 1.5im, 0.3 + tdst = rand(ComplexF64, braid(space(tref), p, levels)) + ref = braid!(copy(tdst), tref, p, levels, α, β) + @test braid!(copy(tdst), t', p, levels, α, β) ≈ ref + D = copy(tdst') + braid!(D', tref, p, levels, α, β) + @test D' ≈ ref + D = copy(tdst') + braid!(D', t', p, levels, α, β) + @test D' ≈ ref end hasbraiding && !symmetricbraiding && @timedtestset "Braid: invalid levels" begin t = rand(ComplexF64, V1 ⊗ V2 ← V3)