Refactor index manipulation kernels around position-indexed subblocks

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
 

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]

docs/src/Changelog.md

@@ -25,16 +25,22 @@ 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 `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))
 
 ### Removed
 
 ### Fixed
 
+- `braid!`, `permute!` and `transpose!` with a `BraidingTensor` source now use the cached fusion tree transformers ([#516](https://github.com/QuantumKitHub/TensorKit.jl/issues/516))
+
 ### Performance
 - `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` 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
 
 ### Added

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.UniqueTreeTransformer}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.LRU}) = true
 

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")

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

src/tensors/blockiterator.jl → 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 of [`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(")
@@ -171,3 +188,41 @@ function Base.show(io::IO, mime::MIME"text/plain", iter::SubblockIterator)
     show_subblocks(io, mime, iter)
     return nothing
 end
+
+"""
+    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`](@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
+    structure::Vector{StridedStructure{N}}
+    # 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}}) where {N}
+        data′ = parent(StridedView(data))
+        return new{typeof(data′), N}(data′, structure)
+    end
+end
+
+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}}) 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)
+end
+
+function Base.iterate(s::StridedSubblocks, i::Int = 1)
+    i > length(s) && return nothing
+    return @inbounds(s[i]), i + 1
+end

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!(