From 736d5f34aac8b5b22e75f7e30fd738d35563a277 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 3 Sep 2026 11:18:10 -0400 Subject: [PATCH 01/16] Refactor index manipulation kernels around position-indexed subblocks Index manipulations now run through a single kernel that operates on subblocks addressed by position: `StridedSubblocks` (sector-independent views into the flat data of a `TensorMap`) or `TreeSubblocks` (any `AbstractTensorMap`, through `subblock`), both carrying an optional lazy conjugation. `TreeTransformer`s store only the mapping between subblock positions and recoupling coefficients, alongside the subblock structures, and are cached for every tensor type. Adjoint sources and destinations, as well as `conj` in `tensoradd!`, are folded into a conjugation flag, relabeled permutation and levels, and conjugated scalars, so that `AdjointTensorMap` wrappers no longer force the uncached generic path (fixes #516, supersedes #519 and #520). Co-Authored-By: Claude Fable 5.1 --- .../IndexManipulationBenchmarks.jl | 13 +- .../indexmanipulations/benchparams.toml | 2 + docs/src/Changelog.md | 6 + ext/TensorKitEnzymeExt/utility.jl | 1 + src/spaces/homspace.jl | 14 + src/tensors/blockiterator.jl | 70 +++++ src/tensors/braidingtensor.jl | 12 +- src/tensors/indexmanipulations.jl | 275 ++++++++---------- src/tensors/tensor.jl | 20 ++ src/tensors/tensoroperations.jl | 9 +- src/tensors/treetransformers.jl | 243 ++++++++-------- test/tensors/indexmanipulations.jl | 82 ++++++ 12 files changed, 446 insertions(+), 301 deletions(-) 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 282799375..827f755ca 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -25,6 +25,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed - For sector types with `GenericUnit` such that colorings are not unique, `GradedSpace`, `ProductSpace` and `HomSpace` now check for this compatibility. In particular, this prevents the construction of `TensorMap`s with incompatible colorings, which previously either errored or produced empty tensors inconsistently. ([#515](https://github.com/QuantumKitHub/TensorKit.jl/pull/515)) +- 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 - The type alias `ZNSpace{N}` is deprecated in favour of `Vect[ZNIrrep{N}]` or `Rep[ℤ{N}]`: a type alias cannot compute the storage type from `N`, so the two only agree for small `N`. ([#511](https://github.com/QuantumKitHub/TensorKit.jl/pull/511)) @@ -32,9 +34,13 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### 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 - `GradedSpace` operations (`dim`, `flip`, `⊕`, `⊖`, `fuse`, `infimum`, `supremum`, truncation) are now specialised on the storage type of the degeneracy dimensions, and tuple storage is used only for sector types with at most `TensorKit._NTUPLE_STORAGE_THRESHOLD` sectors so that sector types with many sectors no longer burden the compiler. ([#511](https://github.com/QuantumKitHub/TensorKit.jl/pull/511)) +- 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 7f3418944..c8431c531 100644 --- a/src/spaces/homspace.jl +++ b/src/spaces/homspace.jl @@ -201,6 +201,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 7da5f84dd..3209e69e0 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) From ee2b8ae57ac1ad200f63455834f5eaf075c28fdc Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 9 Sep 2026 16:18:20 -0400 Subject: [PATCH 02/16] blockiterator->blockiterators --- src/TensorKit.jl | 2 +- src/tensors/{blockiterator.jl => blockiterators.jl} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename src/tensors/{blockiterator.jl => blockiterators.jl} (100%) diff --git a/src/TensorKit.jl b/src/TensorKit.jl index 41a7d8ae1..11a49993e 100644 --- a/src/TensorKit.jl +++ b/src/TensorKit.jl @@ -256,7 +256,7 @@ end # general definitions include("tensors/abstracttensor.jl") include("tensors/backends.jl") -include("tensors/blockiterator.jl") +include("tensors/blockiterators.jl") include("tensors/sectorvector.jl") include("tensors/tensor.jl") include("tensors/adjoint.jl") diff --git a/src/tensors/blockiterator.jl b/src/tensors/blockiterators.jl similarity index 100% rename from src/tensors/blockiterator.jl rename to src/tensors/blockiterators.jl From 3dbaf456aae89e334c6f4b0459391e534081bf2d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 9 Sep 2026 16:26:52 -0400 Subject: [PATCH 03/16] clean up trivial symmetry bypassing overhead --- src/tensors/indexmanipulations.jl | 92 ++++++++++++++----------------- src/tensors/tensoroperations.jl | 5 +- 2 files changed, 42 insertions(+), 55 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index cdfc01c0c..6c9efff75 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -221,12 +221,6 @@ See also [`permute`](@ref) for creating a new tensor. ) @boundscheck spacecheck_transform(permute, tdst, tsrc, p) @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 @@ -317,12 +311,6 @@ See also [`braid`](@ref) for creating a new tensor. ) @boundscheck spacecheck_transform(braid, tdst, tsrc, p, levels) @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′, levels′, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, levels, false, α, β) @inbounds _braid!(tdst′, tsrc′, p′, conjsrc, levels′, α′, β′, backend, allocator) end @@ -396,15 +384,8 @@ end ) @boundscheck spacecheck_transform(transpose, tdst, tsrc, p) @timeit_debug GLOBAL_TIMER "transpose!" 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, α, β) - transformer = treetransposer(tdst′, tsrc′, p′, conjsrc) - @inbounds add_transform!(tdst′, tsrc′, p′, conjsrc, transformer, α′, β′, backend, allocator) + @inbounds _transpose!(tdst′, tsrc′, p′, conjsrc, α′, β′, backend, allocator) end return tdst end @@ -604,19 +585,38 @@ function unwrap_adjoints(tdst, tsrc, p::Index2Tuple, levels, conjsrc::Bool, α, return (tdst′, tsrc′, p″, levels′, conjsrc″, α′, β′) end -# shared by `permute!`, `braid!` and `TO.tensoradd!` after the adjoints have been unwrapped +# dense transform that bypasses overhead +function _dense_transform!(tdst, tsrc, p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator) + p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( + tdst[], tsrc[], p2, conjsrc, α, β, backend, allocator + ) + return tdst +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 + @propagate_inbounds function _braid!( tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator ) @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) + has_array_view(tdst, tsrc) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) 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′) +# counterpart of `_braid!` for `transpose!`; the cyclicity of `p` is checked by the caller +@propagate_inbounds function _transpose!( + tdst, tsrc, p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator + ) + @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) + has_array_view(tdst, tsrc) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) + transformer = treetransposer(tdst, tsrc, p, conjsrc) + return @inbounds add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end """ @@ -635,20 +635,14 @@ of `tsrc`, using the fusion tree transformation encoded in `transformer` (see [` add!(tdst, tsrc, α, β) else 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, conjsrc, α, β, backend, allocator - ) + ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 + # 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 - ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 - # 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 - dst, src = _transform_subblocks(tdst, tsrc, transformer, identity) - add_transform_kernel!(dst, src, p2, transformer, α, β, backend, allocator, ntasks) - end + dst, src = _transform_subblocks(tdst, tsrc, transformer, identity) + add_transform_kernel!(dst, src, p2, transformer, α, β, backend, allocator, ntasks) end end @@ -656,19 +650,15 @@ of `tsrc`, using the fusion tree transformation encoded in `transformer` (see [` 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 -function use_threaded_transform(t::AbstractTensorMap, transformer) - return get_num_transformer_threads() > 1 && dim(space(t)) > Strided.MINTHREADLENGTH -end +_transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer, op) = + StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src, op) +_transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer, op) = + TreeSubblocks(tdst), TreeSubblocks(tsrc, op) + +use_threaded_transform(t::TensorMap, transformer) = + get_num_transformer_threads() > 1 && length(t.data) > Strided.MINTHREADLENGTH +use_threaded_transform(t::AbstractTensorMap, transformer) = + get_num_transformer_threads() > 1 && dim(space(t)) > Strided.MINTHREADLENGTH # 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. diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 3209e69e0..f6f6eeefd 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -39,6 +39,7 @@ has_array_view(t) = has_array_view(typeof(t)) has_array_view(::Type) = false has_array_view(::Type{T}) where {T <: TensorMap} = sectortype(T) === Trivial has_array_view(::Type{T}) where {T <: AdjointTensorMap} = has_array_view(parenttype(T)) +has_array_view(t, ts...) = has_array_view(t) && has_array_view(ts...) # tensoradd! function TO.tensoradd!( @@ -47,10 +48,6 @@ function TO.tensoradd!( α::Number, β::Number, backend, allocator ) - if has_array_view(C) && has_array_view(A) - TO.tensoradd!(C[], A[], pA, conjA, α, β, backend, allocator) - return C - 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 From 25503d4e6cd6db1c7b714cba84853ab43988705f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 15 Sep 2026 11:45:19 -0400 Subject: [PATCH 04/16] Address review: normalize subblock parent, prime conjsrc `StridedSubblocks` now stores its data the way `StridedView` parents it (an `Array` becomes its underlying `Memory` on Julia >= 1.11), asking `StridedView` itself rather than reproducing that rule. This makes the view type a direct function of the type parameters, so `eltype` can be written out instead of going through `Core.Compiler.return_type`. As a consequence `storagetype` reports the normalized type, which is not an `Array`, so the CPU branch of `_adapt_recoupling` is keyed on a `CPUStorage` alias to keep the recoupling matrices off the `Adapt` path. Co-Authored-By: Claude Opus 5 (1M context) --- src/tensors/blockiterators.jl | 9 ++++++++- src/tensors/indexmanipulations.jl | 12 ++++++------ src/tensors/treetransformers.jl | 4 +++- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index a762ca35a..9a3204434 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -190,11 +190,18 @@ struct StridedSubblocks{A <: DenseVector, N, F <: SubblockOp} data::A structure::Vector{StridedStructure{N}} op::F + # store the data as `StridedView` parents it, so that `A` is also the parent type of the views + function StridedSubblocks( + data::DenseVector, structure::Vector{StridedStructure{N}}, op::F = identity + ) where {N, F <: SubblockOp} + data′ = parent(StridedView(data)) + return new{typeof(data′), N, F}(data′, structure, op) + end 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.eltype(::Type{StridedSubblocks{A, N, F}}) where {A, N, F} = StridedView{eltype(A), N, A, F} Base.@propagate_inbounds function Base.getindex(s::StridedSubblocks, i::Int) sz, str, offset = s.structure[i] diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 6c9efff75..d8f42f446 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -221,8 +221,8 @@ See also [`permute`](@ref) for creating a new tensor. ) @boundscheck spacecheck_transform(permute, tdst, tsrc, p) @timeit_debug GLOBAL_TIMER "permute!/braid!" begin - tdst′, tsrc′, p′, _, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β) - @inbounds _braid!(tdst′, tsrc′, p′, conjsrc, allind(tsrc′), α′, β′, backend, allocator) + 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 @@ -311,8 +311,8 @@ See also [`braid`](@ref) for creating a new tensor. ) @boundscheck spacecheck_transform(braid, tdst, tsrc, p, levels) @timeit_debug GLOBAL_TIMER "permute!/braid!" begin - tdst′, tsrc′, p′, levels′, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, levels, false, α, β) - @inbounds _braid!(tdst′, tsrc′, p′, conjsrc, levels′, α′, β′, 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 @@ -384,8 +384,8 @@ end ) @boundscheck spacecheck_transform(transpose, tdst, tsrc, p) @timeit_debug GLOBAL_TIMER "transpose!" begin - tdst′, tsrc′, p′, _, conjsrc, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β) - @inbounds _transpose!(tdst′, tsrc′, p′, conjsrc, α′, β′, backend, allocator) + tdst′, tsrc′, p′, _, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β) + @inbounds _transpose!(tdst′, tsrc′, p′, conjsrc′, α′, β′, backend, allocator) end return tdst end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index d25349ad0..f50e631bc 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -197,5 +197,7 @@ end # Adapt.adapt would make (which additionally dispatches dynamically). Other storage # types (e.g. GPU arrays) do require the conversion. # TODO: transformers with dedicated storagetypes -_adapt_recoupling(::Type{<:Array}, U::Matrix) = StridedView(U) +# `StridedSubblocks` report their storage as the `StridedView` parent type, which is `Memory` there +const CPUStorage = @static isdefined(Core, :Memory) ? Union{Array, Memory} : Array +_adapt_recoupling(::Type{<:CPUStorage}, U::Matrix) = StridedView(U) _adapt_recoupling(::Type{A}, U::Matrix) where {A} = Adapt.adapt(A, StridedView(U)) From a222e0f736a7ef5d330f0186c801bd51b65f1c9c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 15 Sep 2026 15:58:44 -0400 Subject: [PATCH 05/16] Test fix: pick a multifusion-compatible space for the isometry test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `V1 ⊗ V2 ← V3 ⊗ V4` does not close the unit cycle that `GenericUnit` sectors require, so constructing it threw a `SpaceMismatch` for the `IsingBimodule` space lists. The sibling `Permutations: adjoint operands` testset never hit this because it sits behind `symmetricbraiding`, which is false for multifusion; this one builds its space unconditionally. `V1 ⊗ V5 ← V2 ⊗ V4` does close the cycle, and is equally general for the other sectors, so use that rather than skipping multifusion: the transpose half of the test then covers them too, while the braid half stays behind `hasbraiding` (multifusion is `NoBraiding`). Only Windows and macOS saw this, since `default_spacelist` hands out different space lists per OS on CI and only those two include the multifusion entries. Co-Authored-By: Claude Opus 5 (1M context) --- test/tensors/indexmanipulations.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/tensors/indexmanipulations.jl b/test/tensors/indexmanipulations.jl index 025b58010..0cc276c6d 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -183,7 +183,9 @@ for V in spacelist @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 + # this leg selection also closes the unit cycle for multifusion sectors, where + # `V1 ⊗ V2 ← V3 ⊗ V4` does not + W = V1 ⊗ V5 ← V2 ⊗ V4 x = rand(ComplexF64, W) y = rand(ComplexF64, W') if hasbraiding From 19f516c347caa21d58aa3a6938fe38d8080a9c9c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 16 Sep 2026 10:39:56 -0400 Subject: [PATCH 06/16] reorganize to fix docstring --- src/tensors/blockiterators.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index 9a3204434..b648f88fc 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -172,6 +172,8 @@ function Base.show(io::IO, mime::MIME"text/plain", iter::SubblockIterator) return nothing end +const SubblockOp = Union{typeof(identity), typeof(conj)} + """ struct StridedSubblocks{A <: DenseVector, N, F} StridedSubblocks(t::TensorMap, [op = identity]) @@ -185,7 +187,6 @@ 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}} @@ -198,6 +199,9 @@ struct StridedSubblocks{A <: DenseVector, N, F <: SubblockOp} return new{typeof(data′), N, F}(data′, structure, op) end end + +storagetype(::Type{StridedSubblocks{A, N, F}}) where {A, N, F} = A + Base.length(s::StridedSubblocks) = length(s.structure) Base.firstindex(s::StridedSubblocks) = 1 Base.lastindex(s::StridedSubblocks) = length(s) @@ -213,7 +217,6 @@ function Base.iterate(s::StridedSubblocks, i::Int = 1) return @inbounds(s[i]), i + 1 end -storagetype(::Type{StridedSubblocks{A, N, F}}) where {A, N, F} = A """ struct TreeSubblocks{TT <: AbstractTensorMap, I, F} From 9f2cd595235cc4b0246f25b39fc89e921db213ca Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 16 Sep 2026 10:40:10 -0400 Subject: [PATCH 07/16] simplify getting number of transformer_threads --- src/tensors/indexmanipulations.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index d8f42f446..30c8c86d2 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -635,7 +635,7 @@ of `tsrc`, using the fusion tree transformation encoded in `transformer` (see [` add!(tdst, tsrc, α, β) else p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels - ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 + ntasks = get_num_transformer_threads(tdst) # 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) @@ -655,10 +655,9 @@ _transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer, op) = _transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer, op) = TreeSubblocks(tdst), TreeSubblocks(tsrc, op) -use_threaded_transform(t::TensorMap, transformer) = - get_num_transformer_threads() > 1 && length(t.data) > Strided.MINTHREADLENGTH -use_threaded_transform(t::AbstractTensorMap, transformer) = - get_num_transformer_threads() > 1 && dim(space(t)) > Strided.MINTHREADLENGTH +# Don't thread if overhead is not worth it +get_num_transformer_threads(t::AbstractTensorMap) = + dim(t) <= Strided.MINTHREADLENGTH ? 1 : get_num_transformer_threads() # 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. From 9da1592f6fe93778e3155e6a92e59c6d653bc936 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 16 Sep 2026 10:59:50 -0400 Subject: [PATCH 08/16] mark `add_transform` as non-public --- src/tensors/indexmanipulations.jl | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 30c8c86d2..f9b594b5e 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -606,7 +606,7 @@ end @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) has_array_view(tdst, tsrc) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) transformer = treebraider(tdst, tsrc, p, conjsrc, levels) - return @inbounds add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) + return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end # counterpart of `_braid!` for `transpose!`; the cyclicity of `p` is checked by the caller @@ -616,21 +616,15 @@ end @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) has_array_view(tdst, tsrc) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) transformer = treetransposer(tdst, tsrc, p, conjsrc) - return @inbounds add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) + return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) 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!( +# kernel for computing `tdst = β * tdst + α * permutedims(conjsrc ? conj(tsrc) : tsrc, p)` +# with fusion tree transformation encoded in `transformer`. +function add_transform!( tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, transformer, α::Number, β::Number, backend, allocator ) - @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) - if !conjsrc && p[1] === codomainind(tsrc) && p[2] === domainind(tsrc) add!(tdst, tsrc, α, β) else From e8b083ce37add69d09b2856de92cb3d518faf306 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 16 Sep 2026 14:44:45 -0400 Subject: [PATCH 09/16] fix (unrelated) docstring sentence --- src/tensors/blockiterators.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index b648f88fc..09e4d1da6 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -110,7 +110,7 @@ end struct SubblockIterator{T <: AbstractTensorMap, S} Iterator over the subblocks of a tensor of type `T`, possibly holding some pre-computed data of type `S`. -This is typically constructed through of [`subblocks`](@ref). +This is typically constructed through [`subblocks`](@ref). """ struct SubblockIterator{T <: AbstractTensorMap, S} t::T From 197729393c780eae6815782c524864c90dfe4963 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 16 Sep 2026 14:51:47 -0400 Subject: [PATCH 10/16] all has_array_view --- src/tensors/indexmanipulations.jl | 4 ++-- src/tensors/tensoroperations.jl | 5 ++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index f9b594b5e..22d90e169 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -604,7 +604,7 @@ end tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator ) @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) - has_array_view(tdst, tsrc) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) + all(has_array_view, (tdst, tsrc)) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) transformer = treebraider(tdst, tsrc, p, conjsrc, levels) return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end @@ -614,7 +614,7 @@ end tdst, tsrc, p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator ) @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc) - has_array_view(tdst, tsrc) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) + all(has_array_view, (tdst, tsrc)) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) transformer = treetransposer(tdst, tsrc, p, conjsrc) return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index f6f6eeefd..91fa672bb 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -39,7 +39,6 @@ has_array_view(t) = has_array_view(typeof(t)) has_array_view(::Type) = false has_array_view(::Type{T}) where {T <: TensorMap} = sectortype(T) === Trivial has_array_view(::Type{T}) where {T <: AdjointTensorMap} = has_array_view(parenttype(T)) -has_array_view(t, ts...) = has_array_view(t) && has_array_view(ts...) # tensoradd! function TO.tensoradd!( @@ -128,7 +127,7 @@ function TO.tensorcontract!( ) pAB′ = _canonicalize(pAB, C) @boundscheck spacecheck_contract(C, A, pA, conjA, B, pB, conjB, pAB′) - if has_array_view(C) && has_array_view(A) && has_array_view(B) + if all(has_array_view, (C, A, B)) TO.tensorcontract!(C[], A[], pA, conjA, B[], pB, conjB, pAB′, α, β, backend, allocator) return C end @@ -227,7 +226,7 @@ function trace_permute!( end @timeit_debug GLOBAL_TIMER "trace_permute!" begin - if has_array_view(tdst) && has_array_view(tsrc) + if all(has_array_view, (tdst, tsrc)) TO.tensortrace!(tdst[], tsrc[], (p₁, p₂), (q₁, q₂), false, α, β, backend) else _trace_permute!(FusionStyle(I), tdst, tsrc, (p₁, p₂), (q₁, q₂), α, β, backend) From 2366bb21873408e0cf6d30e4c59e0c9ba8980bb1 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 17 Sep 2026 11:07:36 -0400 Subject: [PATCH 11/16] clamp -> min --- src/tensors/indexmanipulations.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 22d90e169..fceaba98a 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -673,7 +673,7 @@ function add_transform_kernel!( 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) + for _ in 1:min(length(transformer.data), 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) From e0fddbf5125d620ed8690e6c9eea6123f8487bf2 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 17 Sep 2026 11:10:12 -0400 Subject: [PATCH 12/16] docstring improvements --- src/tensors/blockiterators.jl | 14 ++++++-------- src/tensors/treetransformers.jl | 2 +- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index 09e4d1da6..0d831e852 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -178,14 +178,12 @@ const SubblockOp = Union{typeof(identity), typeof(conj)} 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`. +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`. +It only contains type parameters `N` (number of indices) and `F` (element-wise operation). """ struct StridedSubblocks{A <: DenseVector, N, F <: SubblockOp} data::A diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index f50e631bc..6e6326637 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -4,7 +4,7 @@ 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 +[`StridedSubblocks`](@ref) (the canonical order of [`fusiontrees`](@ref)), 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 From 11998d14a2b4036bcb2eb564861786830b5da1b3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 17 Sep 2026 13:45:11 -0400 Subject: [PATCH 13/16] simplify treetransformer implementations --- src/spaces/homspace.jl | 14 --------- src/tensors/blockiterators.jl | 2 +- src/tensors/indexmanipulations.jl | 13 ++++---- src/tensors/treetransformers.jl | 51 +++++++++++++++---------------- 4 files changed, 33 insertions(+), 47 deletions(-) diff --git a/src/spaces/homspace.jl b/src/spaces/homspace.jl index c8431c531..7f3418944 100644 --- a/src/spaces/homspace.jl +++ b/src/spaces/homspace.jl @@ -201,20 +201,6 @@ 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/blockiterators.jl b/src/tensors/blockiterators.jl index 0d831e852..d8a13d774 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -179,7 +179,7 @@ const SubblockOp = Union{typeof(identity), typeof(conj)} 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). +Subblock `i` corresponds to the `i`th fusion tree pair in the canonical order of `fusiontrees(space(t))`. 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`. diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index fceaba98a..34c018e86 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -521,6 +521,13 @@ spacecheck_transform(f, tdst::AbstractTensorMap, tsrc::AbstractTensorMap, args.. ) return nothing end +# a conjugated source is transformed through its adjoint space, whose legs `p` is relabeled onto +@noinline function spacecheck_transform( + f, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool + ) + Vsrc′, p′ = conjsrc ? (Vsrc', adjointtensorindices(Vsrc, p)) : (Vsrc, p) + return spacecheck_transform(f, Vdst, Vsrc′, p′) +end @noinline function spacecheck_transform(::typeof(braid), Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, levels::IndexTuple) check_spacetype(Vdst, Vsrc) braid(Vsrc, p, levels) == Vdst || @@ -594,12 +601,6 @@ function _dense_transform!(tdst, tsrc, p::Index2Tuple, conjsrc::Bool, α, β, ba return tdst 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 - @propagate_inbounds function _braid!( tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator ) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 6e6326637..c33891106 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -47,27 +47,19 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer structure_src::Vector{StridedStructure{N}} end -# 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() - 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) + + spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + + src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) T = sectorscalartype(sectortype(Vdst)) - data = Vector{AbelianTransformerData{T}}(undef, L) + data = Vector{AbelianTransformerData{T}}(undef, length(src_trees)) - @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (i, f_src) in enumerate(trees_src) - f_dst, coeff = transform(f_src) - data[i] = (coeff, dstindices[f_dst], srcindices[source_tree(f_src, conjsrc)]) + @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (isrc, (f₁, f₂)) in enumerate(src_trees) + f_dst, coeff = transform(conjsrc ? (f₂, f₁) : (f₁, f₂)) + _, (_, idst) = gettoken(dst_trees, f_dst) + data[isrc] = (coeff, idst, isrc) end structure_dst = degeneracystructure(Vdst).subblockstructure @@ -75,17 +67,17 @@ function AbelianTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) transformer = AbelianTreeTransformer(data, structure_dst, structure_src) Δt = Base.time() - t₀ - @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = L, Δt) + @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = length(data), Δt) return transformer end function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) t₀ = Base.time() - Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) - spacecheck_transform(permute, Vdst, Vsrc′, p′) - srcindices = fusiontreeindices(Vsrc) - dstindices = fusiontreeindices(Vdst) + spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + # the fusion blocks that are transformed are those of the adjoint space for a conjugated source + Vsrc′ = conjsrc ? Vsrc' : Vsrc + src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) structure_dst = degeneracystructure(Vdst).subblockstructure structure_src = degeneracystructure(Vsrc).subblockstructure T = sectorscalartype(sectortype(Vsrc)) @@ -101,8 +93,15 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) fs_src = fblocks[i] fs_dst, U = transform(fs_src) @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)) + # the token into the fusion tree `Indices` is the subblock position + inds_src = map(fusiontrees(fs_src)) do (f₁, f₂) + _, (_, isrc) = gettoken(src_trees, conjsrc ? (f₂, f₁) : (f₁, f₂)) + return isrc + end + inds_dst = map(fusiontrees(fs_dst)) do f + _, (_, idst) = gettoken(dst_trees, f) + return idst + end end data[i] = (U, inds_dst, inds_src) # cost model: L input blocks each going to L output blocks of a given length @@ -171,7 +170,7 @@ end @cached function treebraider( Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple )::treetransformertype(Vdst, Vsrc) - Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) + Vsrc′, p′ = conjsrc ? (Vsrc', adjointtensorindices(Vsrc, p)) : (Vsrc, p) # 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′))) @@ -185,7 +184,7 @@ end @cached function treetransposer( Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool )::treetransformertype(Vdst, Vsrc) - Vsrc′, p′ = transform_source(Vsrc, p, conjsrc) + p′ = conjsrc ? adjointtensorindices(Vsrc, p) : p fusiontreetransform(f) = transpose(f, p′) return TreeTransformer(fusiontreetransform, p, Vdst, Vsrc, conjsrc) end From cf67a5ec46f89314358122fd86e9e35637d7c035 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 19 Sep 2026 08:16:02 -0400 Subject: [PATCH 14/16] Rename `AbelianTreeTransformer` to `UniqueTreeTransformer` "Abelian" is ambiguous for sectors: it can refer either to the fusion of two sectors having a unique result, or to the commutativity of the fusion rules. The transformer is selected on `FusionStyle(I) == UniqueFusion()`, so name it after that. Co-Authored-By: Claude Opus 5 (1M context) --- ext/TensorKitEnzymeExt/utility.jl | 2 +- src/tensors/indexmanipulations.jl | 8 ++++---- src/tensors/treetransformers.jl | 20 ++++++++++---------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 4280a25e0..35ce295fd 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -24,7 +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.UniqueTreeTransformer}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.LRU}) = true diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 34c018e86..6e11ab0d6 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -658,8 +658,8 @@ get_num_transformer_threads(t::AbstractTensorMap) = # depends on `numind`, `eltype` and the transformer data, not on the sectortype. const TransformSubblocks = Union{StridedSubblocks, TreeSubblocks} function add_transform_kernel!( - dst::TransformSubblocks, src::TransformSubblocks, p, - transformer::Union{AbelianTreeTransformer, GenericTreeTransformer}, + dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, + transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int ) bufsize = buffersize(transformer) @@ -685,8 +685,8 @@ function add_transform_kernel!( return nothing end -# `U` is either a scalar coefficient with integer positions (abelian), or a recoupling matrix -# with vectors of positions (generic). +# `U` is either a scalar coefficient with integer positions (unique fusion), or a recoupling +# matrix with vectors of positions (generic). function _add_transform_block!( dst::TransformSubblocks, src::TransformSubblocks, p, U, inds_dst, inds_src, buffer, α, β, backend, allocator diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index c33891106..44a200cb8 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -13,18 +13,18 @@ those of `adjoint(space(tsrc))`, which read the subblocks of `tsrc` with the tre abstract type TreeTransformer end # (coefficient, destination position, source position) -const AbelianTransformerData{T} = Tuple{T, Int, Int} +const UniqueTransformerData{T} = Tuple{T, Int, Int} """ - AbelianTreeTransformer{T, N} <: TreeTransformer + UniqueTreeTransformer{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}} +struct UniqueTreeTransformer{T, N} <: TreeTransformer + data::Vector{UniqueTransformerData{T}} structure_dst::Vector{StridedStructure{N}} structure_src::Vector{StridedStructure{N}} end @@ -47,14 +47,14 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer structure_src::Vector{StridedStructure{N}} end -function AbelianTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) +function UniqueTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) t₀ = Base.time() spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) T = sectorscalartype(sectortype(Vdst)) - data = Vector{AbelianTransformerData{T}}(undef, length(src_trees)) + data = Vector{UniqueTransformerData{T}}(undef, length(src_trees)) @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (isrc, (f₁, f₂)) in enumerate(src_trees) f_dst, coeff = transform(conjsrc ? (f₂, f₁) : (f₁, f₂)) @@ -64,7 +64,7 @@ function AbelianTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) structure_dst = degeneracystructure(Vdst).subblockstructure structure_src = degeneracystructure(Vsrc).subblockstructure - transformer = AbelianTreeTransformer(data, structure_dst, structure_src) + transformer = UniqueTreeTransformer(data, structure_dst, structure_src) Δt = Base.time() - t₀ @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = length(data), Δt) @@ -137,7 +137,7 @@ Compute the workspace size required to pack, recouple and unpack the largest mul block, i.e. `prod(sz_src) * (rows + cols)` where `(rows, cols) = size(U)` is the size of the recoupling matrix. """ -buffersize(::AbelianTreeTransformer) = 0 +buffersize(::UniqueTreeTransformer) = 0 function buffersize(transformer::GenericTreeTransformer) structure_src = transformer.structure_src return maximum(transformer.data; init = 0) do (U, _, inds_src) @@ -149,7 +149,7 @@ function treetransformertype(Vdst, Vsrc) I = sectortype(Vdst) T = sectorscalartype(I) N = numind(Vdst) - return FusionStyle(I) == UniqueFusion() ? AbelianTreeTransformer{T, N} : GenericTreeTransformer{T, N} + return FusionStyle(I) == UniqueFusion() ? UniqueTreeTransformer{T, N} : GenericTreeTransformer{T, N} end function TreeTransformer( @@ -157,7 +157,7 @@ function TreeTransformer( ) where {S} I = sectortype(Vdst) return FusionStyle(I) == UniqueFusion() ? - AbelianTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) : + UniqueTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) : GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) end From 40e53aa8da0f021a4ba6b0d73318001e0aff09c4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 19 Sep 2026 08:22:38 -0400 Subject: [PATCH 15/16] Pass conjugation as a flag instead of a subblock view op `StridedSubblocks` and `TreeSubblocks` no longer apply `identity`/`conj` to every view. Instead `conjsrc` is threaded through `add_transform_kernel!` into `_add_transform_block!`, where it is handed to `TO.tensoradd!` as its `conjA` argument, at the single-tree call and when packing a multi-tree block. This drops a type parameter from both collections, so the kernel compiles to one instance per (storage, numind) rather than one per conjugation. The runtime flag is free: `flag2op` is union-split, and `conj` of a real-eltype `StridedView` is a type-level no-op, which also makes the previous `scalartype(t) <: Real` guard redundant. Co-Authored-By: Claude Opus 5 (1M context) --- docs/src/Changelog.md | 2 +- src/tensors/blockiterators.jl | 47 +++++++++++++------------------ src/tensors/indexmanipulations.jl | 31 ++++++++------------ src/tensors/tensor.jl | 8 +++--- src/tensors/treetransformers.jl | 6 ++-- 5 files changed, 39 insertions(+), 55 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 827f755ca..23d1e33c8 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -25,7 +25,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed - For sector types with `GenericUnit` such that colorings are not unique, `GradedSpace`, `ProductSpace` and `HomSpace` now check for this compatibility. In particular, this prevents the construction of `TensorMap`s with incompatible colorings, which previously either errored or produced empty tensors inconsistently. ([#515](https://github.com/QuantumKitHub/TensorKit.jl/pull/515)) -- 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)) +- 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`). 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 - The type alias `ZNSpace{N}` is deprecated in favour of `Vect[ZNIrrep{N}]` or `Rep[ℤ{N}]`: a type alias cannot compute the storage type from `N`, so the two only agree for small `N`. ([#511](https://github.com/QuantumKitHub/TensorKit.jl/pull/511)) diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index d8a13d774..6f2384fac 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -172,42 +172,36 @@ function Base.show(io::IO, mime::MIME"text/plain", iter::SubblockIterator) return nothing end -const SubblockOp = Union{typeof(identity), typeof(conj)} - """ - struct StridedSubblocks{A <: DenseVector, N, F} - StridedSubblocks(t::TensorMap, [op = identity]) + struct StridedSubblocks{A <: DenseVector, N} + StridedSubblocks(t::TensorMap) 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))`. -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`. -It only contains type parameters `N` (number of indices) and `F` (element-wise operation). +This is the data structure consumed by the index manipulation kernels, whose type does not depend on the sectortype of `t`: +it only carries the storage type `A` of the flat data vector and the number of indices `N` of the subblocks. """ -struct StridedSubblocks{A <: DenseVector, N, F <: SubblockOp} +struct StridedSubblocks{A <: DenseVector, N} data::A structure::Vector{StridedStructure{N}} - op::F # store the data as `StridedView` parents it, so that `A` is also the parent type of the views - function StridedSubblocks( - data::DenseVector, structure::Vector{StridedStructure{N}}, op::F = identity - ) where {N, F <: SubblockOp} + function StridedSubblocks(data::DenseVector, structure::Vector{StridedStructure{N}}) where {N} data′ = parent(StridedView(data)) - return new{typeof(data′), N, F}(data′, structure, op) + return new{typeof(data′), N}(data′, structure) end end -storagetype(::Type{StridedSubblocks{A, N, F}}) where {A, N, F} = A +storagetype(::Type{StridedSubblocks{A, N}}) where {A, N} = A Base.length(s::StridedSubblocks) = length(s.structure) Base.firstindex(s::StridedSubblocks) = 1 Base.lastindex(s::StridedSubblocks) = length(s) -Base.eltype(::Type{StridedSubblocks{A, N, F}}) where {A, N, F} = StridedView{eltype(A), N, A, F} +Base.eltype(::Type{StridedSubblocks{A, N}}) where {A, N} = StridedView{eltype(A), N, A, typeof(identity)} 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) + return StridedView(s.data, sz, str, offset) end function Base.iterate(s::StridedSubblocks, i::Int = 1) @@ -217,32 +211,29 @@ end """ - struct TreeSubblocks{TT <: AbstractTensorMap, I, F} - TreeSubblocks(t::AbstractTensorMap, [op = identity]) + struct TreeSubblocks{TT <: AbstractTensorMap, I} + TreeSubblocks(t::AbstractTensorMap) 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. +[`subblock`](@ref). 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} +struct TreeSubblocks{TT <: AbstractTensorMap, I} 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 +TreeSubblocks(t::AbstractTensorMap) = TreeSubblocks(t, fusiontrees(t)) -storagetype(::Type{TreeSubblocks{TT, I, F}}) where {TT, I, F} = storagetype(TT) +storagetype(::Type{TreeSubblocks{TT, I}}) where {TT, I} = 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.eltype(::Type{TreeSubblocks{TT, I}}) where {TT, I} = subblocktype(TT) Base.@propagate_inbounds function Base.getindex(s::TreeSubblocks, i::Int) - return s.op(subblock(s.t, gettokenvalue(s.trees, i))) + return subblock(s.t, gettokenvalue(s.trees, i)) end function Base.iterate(s::TreeSubblocks, i::Int = 1) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 6e11ab0d6..bec9c67ab 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -631,24 +631,18 @@ function add_transform!( else p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels ntasks = get_num_transformer_threads(tdst) - # 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 - dst, src = _transform_subblocks(tdst, tsrc, transformer, identity) - add_transform_kernel!(dst, src, p2, transformer, α, β, backend, allocator, ntasks) - end + dst, src = _transform_subblocks(tdst, tsrc, transformer) + add_transform_kernel!(dst, src, p2, conjsrc, transformer, α, β, backend, allocator, ntasks) end return tdst end # TensorMaps address their flat data directly, other tensor types go through `subblock` -_transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer, op) = - StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src, op) -_transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer, op) = - TreeSubblocks(tdst), TreeSubblocks(tsrc, op) +_transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer) = + StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src) +_transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer) = + TreeSubblocks(tdst), TreeSubblocks(tsrc) # Don't thread if overhead is not worth it get_num_transformer_threads(t::AbstractTensorMap) = @@ -659,13 +653,12 @@ get_num_transformer_threads(t::AbstractTensorMap) = const TransformSubblocks = Union{StridedSubblocks, TreeSubblocks} function add_transform_kernel!( dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, - transformer::TreeTransformer, - α, β, backend, allocator, ntasks::Int + transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int ) bufsize = buffersize(transformer) 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) + _add_transform_block!(dst, src, p, conjsrc, U, inds_dst, inds_src, nothing, α, β, backend, allocator) end else # One max-sized workspace per task (a single one that is reused by all blocks when @@ -677,7 +670,7 @@ function add_transform_kernel!( for _ in 1:min(length(transformer.data), 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) + _add_transform_block!(dst, src, p, conjsrc, U, inds_dst, inds_src, buffer, α, β, backend, allocator) end foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) TO.allocator_reset!(allocator, cp) @@ -688,12 +681,12 @@ end # `U` is either a scalar coefficient with integer positions (unique fusion), or a recoupling # matrix with vectors of positions (generic). function _add_transform_block!( - dst::TransformSubblocks, src::TransformSubblocks, p, U, inds_dst, inds_src, buffer, + dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, U, inds_dst, inds_src, buffer, α, β, 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 + dst[only(inds_dst)], src[only(inds_src)], p, conjsrc, α * only(U), β, backend, allocator ) else # Multi-tree block: pack → recoupling matmul → unpack. rows, cols = size(U) @@ -708,7 +701,7 @@ function _add_transform_block!( @timeit_debug GLOBAL_TIMER "dense: pack" @inbounds for (i, isrc) in enumerate(inds_src) TO.tensoradd!( sreshape(view(buffer_src, :, i), sz_src), src[isrc], - ptriv, false, One(), Zero(), backend, allocator + ptriv, conjsrc, One(), Zero(), backend, allocator ) end diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index d42dd320c..162aebf48 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -496,11 +496,11 @@ 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) +function StridedSubblocks(t::TensorMap) + return StridedSubblocks(t, degeneracystructure(space(t)).subblockstructure) end -function StridedSubblocks(t::TensorMap, structure::Vector{<:StridedStructure}, op::SubblockOp = identity) - return StridedSubblocks(t.data, structure, scalartype(t) <: Real ? identity : op) +function StridedSubblocks(t::TensorMap, structure::Vector{<:StridedStructure}) + return StridedSubblocks(t.data, structure) end # iterate the subblock views in canonical order alongside the fusion trees, without hashing diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 44a200cb8..231171da8 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -6,9 +6,9 @@ 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`](@ref)), 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. +transformation is that of `permutedims(conjsrc ? conj(tsrc) : tsrc, p)` where `p` indexes the legs +of `tsrc` itself; when `conjsrc` is `true` 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 From 6b2dd79fd89bca760a1763c72bf19074f7949135 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 19 Sep 2026 16:42:59 -0400 Subject: [PATCH 16/16] remove TreeSubblocks and route through StridedSubblocks --- docs/src/Changelog.md | 4 +- src/tensors/abstracttensor.jl | 2 +- src/tensors/blockiterators.jl | 74 +++++++++++++------------------ src/tensors/indexmanipulations.jl | 4 +- src/tensors/tensor.jl | 40 +++++++++-------- 5 files changed, 57 insertions(+), 67 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 23d1e33c8..9c00e6ae8 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -25,7 +25,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed - For sector types with `GenericUnit` such that colorings are not unique, `GradedSpace`, `ProductSpace` and `HomSpace` now check for this compatibility. In particular, this prevents the construction of `TensorMap`s with incompatible colorings, which previously either errored or produced empty tensors inconsistently. ([#515](https://github.com/QuantumKitHub/TensorKit.jl/pull/515)) -- 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`). 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)) +- Index manipulations use a single kernel operating on subblocks addressed by position: `StridedSubblocks` (sector-independent views into the flat data of a `TensorMap`) or `SubblockIterator` (any `AbstractTensorMap`, through `subblock`). 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 - The type alias `ZNSpace{N}` is deprecated in favour of `Vect[ZNIrrep{N}]` or `Rep[ℤ{N}]`: a type alias cannot compute the storage type from `N`, so the two only agree for small `N`. ([#511](https://github.com/QuantumKitHub/TensorKit.jl/pull/511)) @@ -39,7 +39,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Performance - `GradedSpace` operations (`dim`, `flip`, `⊕`, `⊖`, `fuse`, `infimum`, `supremum`, truncation) are now specialised on the storage type of the degeneracy dimensions, and tuple storage is used only for sector types with at most `TensorKit._NTUPLE_STORAGE_THRESHOLD` sectors so that sector types with many sectors no longer burden the compiler. ([#511](https://github.com/QuantumKitHub/TensorKit.jl/pull/511)) -- 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)) +- 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` iterates without repeating any space-level lookup ([#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 diff --git a/src/tensors/abstracttensor.jl b/src/tensors/abstracttensor.jl index 842cf1074..b701dd1df 100644 --- a/src/tensors/abstracttensor.jl +++ b/src/tensors/abstracttensor.jl @@ -412,7 +412,7 @@ corresponding tensor subblocks. See also [`subblock`](@ref) and [`fusiontrees`](@ref). """ -subblocks(t::AbstractTensorMap) = SubblockIterator(t, fusiontrees(t)) +subblocks(t::AbstractTensorMap) = SubblockIterator(t) const _doc_subblock = """ Return a view into the data of `t` corresponding to the splitting - fusion tree pair diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index 6f2384fac..893ad5ee7 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -106,32 +106,49 @@ function Base.show(io::IO, mime::MIME"text/plain", b::BlockIterator) return nothing end +# Positional subblock collections +# ------------------------------- +# These address subblocks by their position in the canonical order of `fusiontrees(space(t))`, +# hoisting the space-level lookups out of the individual accesses. + """ struct SubblockIterator{T <: AbstractTensorMap, S} + SubblockIterator(t::AbstractTensorMap) + +Collection of the subblocks of a tensor of type `T`, indexable both by fusion tree pair and by +token, i.e. by the position in the canonical order of [`fusiontrees`](@ref), and iterating over +`(f₁, f₂) => subblock` pairs. This is what [`subblocks`](@ref) returns. -Iterator over the subblocks of a tensor of type `T`, possibly holding some pre-computed data of type `S`. -This is typically constructed through [`subblocks`](@ref). +This object additionally has a `structure` field which can be used to precompute data that facilitates +fast indexing/iteration. By default this only holds the fusion tree pairs, but it can be any dictionary +mapping those onto the data needed to address the relevant subblocks, as `TensorMap` does. """ struct SubblockIterator{T <: AbstractTensorMap, S} t::T structure::S end +# default just holds the set of fusiontrees for fast iteration and mapping index to fusiontree +SubblockIterator(t::AbstractTensorMap) = SubblockIterator(t, fusiontrees(t)) + +storagetype(::Type{<:SubblockIterator{T}}) where {T} = storagetype(T) + Base.IteratorSize(::SubblockIterator) = Base.HasLength() Base.IteratorEltype(::SubblockIterator) = Base.HasEltype() Base.eltype(::Type{<:SubblockIterator{T}}) where {T} = Pair{fusiontreetype(T), subblocktype(T)} Base.length(iter::SubblockIterator) = length(iter.structure) -Base.isdone(iter::SubblockIterator, state...) = Base.isdone(iter.structure, state...) +Base.firstindex(::SubblockIterator) = 1 +Base.lastindex(iter::SubblockIterator) = length(iter) +Base.isdone(iter::SubblockIterator, i::Int = 1) = i > length(iter) -# default implementation assumes `structure = fusiontrees(t)` -function Base.iterate(iter::SubblockIterator, state...) - next = Base.iterate(iter.structure, state...) - isnothing(next) && return nothing - (f₁, f₂), state = next - @inbounds data = subblock(iter.t, (f₁, f₂)) - return (f₁, f₂) => data, state -end +@propagate_inbounds Base.getindex(iter::SubblockIterator, i::Int) = + subblock(iter.t, gettokenvalue(keys(iter.structure), i)) +@propagate_inbounds Base.getindex(iter::SubblockIterator, f::FusionTreePair) = subblock(iter.t, f) +function Base.iterate(iter::SubblockIterator, i::Int = 1) + i > length(iter) && return nothing + @inbounds return gettokenvalue(keys(iter.structure), i) => iter[i], i + 1 +end function Base.showarg(io::IO, iter::SubblockIterator, toplevel::Bool) print(io, "subblocks(") @@ -177,10 +194,11 @@ end StridedSubblocks(t::TensorMap) 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))`. +Subblock `i` corresponds to the `i`th fusion tree pair in the canonical order of [`fusiontrees`](@ref). This is the data structure consumed by the index manipulation kernels, whose type does not depend on the sectortype of `t`: it only carries the storage type `A` of the flat data vector and the number of indices `N` of the subblocks. +As a result, the kernels do not have to be recompiled for each new symmetry type. """ struct StridedSubblocks{A <: DenseVector, N} data::A @@ -208,35 +226,3 @@ function Base.iterate(s::StridedSubblocks, i::Int = 1) i > length(s) && return nothing return @inbounds(s[i]), i + 1 end - - -""" - struct TreeSubblocks{TT <: AbstractTensorMap, I} - TreeSubblocks(t::AbstractTensorMap) - -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). 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} - t::TT - trees::I -end -TreeSubblocks(t::AbstractTensorMap) = TreeSubblocks(t, fusiontrees(t)) - -storagetype(::Type{TreeSubblocks{TT, I}}) where {TT, I} = storagetype(TT) - -Base.length(s::TreeSubblocks) = length(s.trees) -Base.firstindex(s::TreeSubblocks) = 1 -Base.lastindex(s::TreeSubblocks) = length(s) -Base.eltype(::Type{TreeSubblocks{TT, I}}) where {TT, I} = subblocktype(TT) - -Base.@propagate_inbounds function Base.getindex(s::TreeSubblocks, i::Int) - return 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/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index bec9c67ab..106f020a4 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -642,7 +642,7 @@ end _transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer) = StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src) _transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer) = - TreeSubblocks(tdst), TreeSubblocks(tsrc) + subblocks(tdst), subblocks(tsrc) # Don't thread if overhead is not worth it get_num_transformer_threads(t::AbstractTensorMap) = @@ -650,7 +650,7 @@ get_num_transformer_threads(t::AbstractTensorMap) = # 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} +const TransformSubblocks = Union{StridedSubblocks, SubblockIterator} function add_transform_kernel!( dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index 162aebf48..0c4697469 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -496,32 +496,36 @@ end # Getting and setting the data at the subblock level # -------------------------------------------------- -function StridedSubblocks(t::TensorMap) - return StridedSubblocks(t, degeneracystructure(space(t)).subblockstructure) -end -function StridedSubblocks(t::TensorMap, structure::Vector{<:StridedStructure}) - return StridedSubblocks(t.data, structure) -end +StridedSubblocks(t::TensorMap) = StridedSubblocks(t.data, degeneracystructure(space(t)).subblockstructure) +StridedSubblocks(t::TensorMap, structure::Vector{<:StridedStructure}) = StridedSubblocks(t.data, structure) + +# the subblocks of a `TensorMap` are strided views into its flat data, which the subblock +# structure of its space describes directly +SubblockIterator(t::TensorMap) = SubblockIterator(t, subblockstructure(space(t))) -# 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)) +# the throws are kept out of line: mentioning `iter` in the error would make it escape, +# which costs the caller scalar replacement of the subblock views +@noinline _throw_subblock_bounds(iter, i) = throw(BoundsError(iter, i)) +@noinline _throw_subblock_missing(f) = throw(SectorMismatch(lazy"fusion tree pair $f is not present")) + +@propagate_inbounds function Base.getindex(iter::SubblockIterator{<:TensorMap}, i::Int) + @boundscheck 0 < i <= length(iter.structure) || _throw_subblock_bounds(iter, i) + @inbounds sz, str, offset = gettokenvalue(iter.structure, i) + return StridedView(iter.t.data, sz, str, offset) 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 +@propagate_inbounds function Base.getindex(iter::SubblockIterator{<:TensorMap}, f::FusionTreePair) + found, token = gettoken(iter.structure, f) + @boundscheck found || _throw_subblock_missing(f) + @inbounds sz, str, offset = gettokenvalue(iter.structure, token) + return StridedView(iter.t.data, sz, str, offset) end -function subblock( +@propagate_inbounds function subblock( t::TensorMap{T, S, N₁, N₂}, (f₁, f₂)::Tuple{FusionTree{I, N₁}, FusionTree{I, N₂}} ) where {T, S, N₁, N₂, I <: Sector} fts = subblockstructure(space(t)) found, token = gettoken(fts, (f₁, f₂)) - @boundscheck found || throw(SectorMismatch(lazy"fusion tree pair ($(f₁, f₂)) is not present")) + @boundscheck found || _throw_subblock_missing((f₁, f₂)) @inbounds begin sz, str, offset = gettokenvalue(fts, token) return StridedView(t.data, sz, str, offset)