diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 2a853124a..3d4049fd9 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -30,11 +30,14 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Require TensorKitSectors v0.3.10 - `@cached` can be used from other modules, such as package extensions, which then own the cache of the methods they add; these caches are shown with their module in `global_cache_info` after registration in `__init__` (internal) - The `TreeTransformer`s used in index manipulations are now constructed and cached per storagetype of the destination tensor, which is also a dispatch point for storage-specific transformers with their own cache (internal) +- The CPU and GPU index manipulations share a single `TreeTransformer` structure per fusion style, which stores the subblock locations alongside the recoupling coefficients and is moved to the GPU with `Adapt`; the kernels operate on the flat data of a `TensorMap` directly (internal) ([#578](https://github.com/QuantumKitHub/TensorKit.jl/issues/578)) ### Deprecated ### Removed +- `StridedSubblocks`, which is superseded by the subblock locations stored in the `TreeTransformer`s (internal) + ### Fixed - Documentation build failing on TensorKitSectors doctests @@ -42,6 +45,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Performance - Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. +- Index manipulations with many large subblocks visit the subblocks in the order of the destination tensor, which is up to ~15% faster. ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 948f9fa75..f3d20b232 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -4,7 +4,6 @@ using GPUArrays using GPUArrays: @allowscalar using GPUArrays.KernelAbstractions: @kernel, @index, get_backend using Adapt -using TensorKit.TupleTools using Strided: StridedViews using MatrixAlgebraKit, Adapt using TensorKit @@ -12,7 +11,7 @@ using TensorKit.TensorOperations: linearize, DefaultAllocator using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check -using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple +using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, StridedLayout, @cached, TensorMapSpace, Index2Tuple, IndexTuple import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel! function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:AnyGPUMatrix{T}, <:StridedViews.StridedView{T, 4, <:AnyGPUArray{T}}}} @@ -126,165 +125,34 @@ end # Device-side tree transformers # ----------------------------- -# A `TreeTransformer` stores the mapping between subblock positions (plus recoupling -# coefficients) and the `StridedStructure`s of the source and destination spaces. -# We resolve the positions into sizes/strides/offsets and pack all the information on the -# CPU side into dense vectors of numbers, plus some accounting information so we know how -# to unpack in the kernel. Also, we can permute the source strides "in advance" on the CPU side. -# We also precompute the strides of the subblock each kernel index will work on, -# so that the GPU thread can recover the Cartesian coordinates it will need for input/ouput, -# and a running `work_offsets` count of destination elements, so that a kernel can run one -# thread per output element and recover which input that element belongs with. +# The tree transformers with `StridedLayout`s only consist of vectors of `isbits` elements, with the +# source strides permuted "in advance" on the CPU side, and are moved to the device as is. The +# running `work_offsets` count of destination elements allows a kernel to run one thread per output +# element and recover which input that element belongs with. # Some possible TODO here: # - Try cuTILE as this is a classic tile programming problem # - Use shared memory to coalesce the reads # - Use a 2D grid for the Generic case -const TreeStructure{N} = Tuple{NTuple{N, Int}, Int} - -""" - UniqueTransformerBlock{T, N} - -`isbits` descriptor for a subblock that is a *single* scaled permutation: -an entry of a `UniqueTreeTransformer` or a unique (one-tree) block of a -`GenericTreeTransformer`. -""" -struct UniqueTransformerBlock{T, N} - coeff::T - sz::NTuple{N, Int} - dense_strides::NTuple{N, Int} - strides_dst::NTuple{N, Int} - offsets_dst::Int - permuted_strides_src::NTuple{N, Int} # source strides, permuted by `p` - offsets_src::Int -end - -# device-side simple struct that GPU kernels can use -struct DeviceUniqueTreeTransformer{VB <: AbstractVector{<:UniqueTransformerBlock}, VO <: AbstractVector{Int}} - blocks::VB - work_offsets::VO - nwork::Int -end - -""" - GenericTransformerBlock{N} - -Descriptor for a recoupling block of a `GenericTreeTransformer`, indexing into the -flat `coeffs`/`structs_dst`/`structs_src` vectors of a `DeviceGenericTreeTransformer`. -""" -struct GenericTransformerBlock{N} - sz::NTuple{N, Int} - densestrides::NTuple{N, Int} - rows::Int - cols::Int - u_offset::Int # location in the flattened U vector to find this block's U - dst_offset::Int - src_offset::Int -end - -# force all the type signatures here to make sure doing something wrong fails -# before the kernel launch. Kernel error dumps are awful and hard to interpret. -struct DeviceGenericTreeTransformer{VO <: AbstractVector{Int}, DA <: DeviceUniqueTreeTransformer{<:Any, VO}, VB <: AbstractVector{<:GenericTransformerBlock}, VC <: AbstractVector{<:Number}, VS <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}} - unique_blocks::DA # length(U) = 1 blocks, can be handled by unique kernel - blocks::VB - work_offsets::VO - nwork::Int - coeffs::VC # every `U`, concatenated in column-major order - structs_dst::VS - structs_src::VS -end - # strides of a dense array of shape `sz` -_dense_strides(size::Dims) = (1, Base.front(cumprod(size))...) - -# `permute(Vsrc, p) == Vdst` is enforced when the transformer is built, so the permuted -# source shape always matches `sz_dst` and the two views share Cartesian inds. -function _unique_block( - coeff::T, (size_dst, strides_dst, offsets_dst), (_, strides_src, offsets_src), p - ) where {T} - return UniqueTransformerBlock{T, length(size_dst)}( - coeff, size_dst, _dense_strides(size_dst), strides_dst, offsets_dst, - TupleTools.getindices(strides_src, p), offsets_src - ) -end - -function _work_offsets(work) - offsets = cumsum(work) - pushfirst!(offsets, 0) - total = pop!(offsets) - return offsets, total -end - -function DeviceUniqueTreeTransformer(transformer::UniqueTreeTransformer{T, N}, p) where {T, N} - (; structure_dst, structure_src) = transformer - blocks = UniqueTransformerBlock{T, N}[ - _unique_block(coeff, structure_dst[idst], structure_src[isrc], p) - for (coeff, idst, isrc) in transformer.data - ] - work_offsets, nwork = _work_offsets(prod(blk.sz) for blk in blocks) - return DeviceUniqueTreeTransformer(blocks, work_offsets, nwork) -end - -function DeviceGenericTreeTransformer( - transformer::GenericTreeTransformer{T, N}, p - ) where {T, N} - unique_blocks = UniqueTransformerBlock{T, N}[] - blocks = GenericTransformerBlock{N}[] - coeffs = T[] - structs_dst = TreeStructure{N}[] - structs_src = TreeStructure{N}[] - - (; structure_dst, structure_src) = transformer - for (U, inds_dst, inds_src) in transformer.data - if length(U) == 1 # same as the unique (Abelian) case - push!( - unique_blocks, _unique_block( - only(U), structure_dst[only(inds_dst)], structure_src[only(inds_src)], p - ) - ) - else - # all trees in a block share the same subblock size - size_dst = first(structure_dst[first(inds_dst)]) - push!( - blocks, GenericTransformerBlock{N}( - size_dst, _dense_strides(size_dst), size(U, 1), size(U, 2), - length(coeffs), length(structs_dst), length(structs_src) - ) - ) - append!(coeffs, U) - for idst in inds_dst - _, strides_dst, offset_dst = structure_dst[idst] - push!(structs_dst, (strides_dst, offset_dst)) - end - for isrc in inds_src - _, strides_src, offset_src = structure_src[isrc] - push!(structs_src, (TupleTools.getindices(strides_src, p), offset_src)) - end - end - end - - unique_offsets, unique_nwork = _work_offsets(prod(blk.sz) for blk in unique_blocks) - work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) - return DeviceGenericTreeTransformer( - DeviceUniqueTreeTransformer(unique_blocks, unique_offsets, unique_nwork), - blocks, work_offsets, nwork, coeffs, structs_dst, structs_src - ) -end +_dense_strides(size::Dims) = Base.front((1, cumprod(size)...)) # Cache the device transformer directly, constructing the cached host version only on a miss. +# Other layouts are transformed on the host, through the positional subblocks. @cached function TensorKit.treebraider( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, levels::IndexTuple + A::Type{TA}, layout_dst::StridedLayout, layout_src::StridedLayout, + Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple ) where {TA <: AnyGPUVector} - host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) - return device_transformer(A, host, linearize(p)) + host = TensorKit.treebraider(Vector{scalartype(A)}, layout_dst, layout_src, Vdst, Vsrc, p, conjsrc, levels) + return device_transformer(A, host) end @cached function TensorKit.treetransposer( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool + A::Type{TA}, layout_dst::StridedLayout, layout_src::StridedLayout, + Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool ) where {TA <: AnyGPUVector} - host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) - return device_transformer(A, host, linearize(p)) + host = TensorKit.treetransposer(Vector{scalartype(A)}, layout_dst, layout_src, Vdst, Vsrc, p, conjsrc) + return device_transformer(A, host) end const DEFAULT_DEVICECACHE_SIZE = 1000 @@ -305,38 +173,21 @@ function __init__() return nothing end -# Allocations must outlive any GPUArrays allocation-cache scope around the operation. -function device_transformer(A, transformer::UniqueTreeTransformer, p) - packed = DeviceUniqueTreeTransformer(transformer, p) - return GPUArrays.@uncached begin - proto = A(undef, 0) - # Metadata and coefficients have different element types; preserve each on upload. - upload(x) = copy!(similar(proto, eltype(x), size(x)), x) - DeviceUniqueTreeTransformer(upload(packed.blocks), upload(packed.work_offsets), packed.nwork) - end +# Adaptor that uploads arrays to the storage of `proto`, preserving their element types, since +# metadata and coefficients have different element types. +struct DeviceUpload{P <: AnyGPUVector} + proto::P end +Adapt.adapt_storage(to::DeviceUpload, x::AbstractArray) = copy!(similar(to.proto, eltype(x), size(x)), x) -function device_transformer(A, transformer::GenericTreeTransformer, p) - packed = DeviceGenericTreeTransformer(transformer, p) - return GPUArrays.@uncached begin - proto = A(undef, 0) - upload(x) = copy!(similar(proto, eltype(x), size(x)), x) - unique = packed.unique_blocks - device_unique = DeviceUniqueTreeTransformer( - upload(unique.blocks), upload(unique.work_offsets), unique.nwork - ) - DeviceGenericTreeTransformer( - device_unique, upload(packed.blocks), upload(packed.work_offsets), packed.nwork, - upload(packed.coeffs), upload(packed.structs_dst), upload(packed.structs_src) - ) - end -end +# Allocations must outlive any GPUArrays allocation-cache scope around the operation. +device_transformer(A, transformer) = GPUArrays.@uncached adapt(DeviceUpload(A(undef, 0)), transformer) # COV_EXCL_START # kernels are not reachable by coverage -# largest `i` with `offsets[i] <= w`. This corresponds to the -# block which this kernel thread will work on. Since this is +# largest `i` with `offsets[i] <= w`. For `0 <= w < offsets[end]`, this corresponds to the +# (nonempty) block which this kernel thread will work on. Since this is # used inside a GPU kernel, searchsortedlast/searchsortedfirst # won't work. @inline function _searchblock(offsets, w) @@ -353,7 +204,7 @@ end end # Cartesian coordinates of the `w`-th (0-based) entry of a dense subblock of shape `sz`. -@inline function _coordinates(w, size::NTuple{N, Int}, dense_strides::NTuple{N, Int}) where {N} +@inline function _coordinates(w, size::NTuple{N, Int}, dense_strides::NTuple{N, Int} = _dense_strides(size)) where {N} return ntuple(n -> (w ÷ dense_strides[n]) % size[n], Val(N)) end @@ -366,23 +217,22 @@ end # One thread per destination element in `data_dst`. `op` is `identity` or `conj`, and is # applied to the source data only (not to the coefficients). @kernel function unique_batched_permute_kernel!( - data_dst, data_src, op, blocks, work_offsets, α, β, nwork, ::Val{N} + data_dst, data_src, op, coeffs, structs_dst, structs_src, work_offsets, α, β, ::Val{N} ) where {N} # w = @index(Global, Linear) - 1 does not work due to JLArrays bug w = @index(Global, Linear) w -= 1 - if w < nwork + if w < @inbounds work_offsets[end] b = _searchblock(work_offsets, w) - blk = @inbounds blocks[b] - coords = _coordinates(w - (@inbounds work_offsets[b]), blk.sz, blk.dense_strides) - i_dst = _linear_index(coords, blk.strides_dst, blk.offsets_dst) - i_src = _linear_index(coords, blk.permuted_strides_src, blk.offsets_src) - @inbounds data_dst[i_dst] = α * blk.coeff * op(data_src[i_src]) + β * data_dst[i_dst] + sz, strides_dst, offset_dst = @inbounds structs_dst[b] + strides_src, offset_src = @inbounds structs_src[b] + coords = _coordinates(w - (@inbounds work_offsets[b]), sz) + i_dst = _linear_index(coords, strides_dst, offset_dst) + i_src = _linear_index(coords, strides_src, offset_src) + @inbounds data_dst[i_dst] = α * coeffs[b] * op(data_src[i_src]) + β * data_dst[i_dst] end end -# COV_EXCL_STOP - # One thread per destination element in `data_dst`. This makes much better use of the # GPU "massive parallelism" as compared to the one-thread-per-subtransformer approach. # It also more evenly divides the work among threads so the work profile is less @@ -392,22 +242,25 @@ end # TODO: what about symmetries like SU(3), where the column by column approach is not # optimal? @kernel function generic_batched_permute_kernel!( - data_dst, data_src, op, blocks, work_offsets, coeffs, structs_dst, structs_src, - α, β, nwork, ::Val{N} + data_dst, data_src, op, coeffs, structs_dst, structs_src, tree_offsets, coeff_offsets, work_offsets, + α, β, ::Val{N} ) where {N} # w = @index(Global, Linear) - 1 does not work due to JLArrays bug w = @index(Global, Linear) w -= 1 - if w < nwork + if w < @inbounds work_offsets[end] # bookkeeping to figure out where to read from and write to b = _searchblock(work_offsets, w) - blk = @inbounds blocks[b] + t₀ = @inbounds tree_offsets[b] + L = (@inbounds tree_offsets[b + 1]) - t₀ + u₀ = @inbounds coeff_offsets[b] + sz = first(@inbounds structs_dst[t₀ + 1]) + blocksize = prod(sz) local_w = w - (@inbounds work_offsets[b]) - blocksize = prod(blk.sz) i = local_w ÷ blocksize # 0-based destination tree - coords = _coordinates(local_w % blocksize, blk.sz, blk.densestrides) + coords = _coordinates(local_w % blocksize, sz) - st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1] + _, st_dst, offs_dst = @inbounds structs_dst[t₀ + i + 1] i_dst = _linear_index(coords, st_dst, offs_dst) # dst_i = β * dst_i + α * Σ_j U[i, j] * permute(src_j, p): each output tree is a @@ -415,47 +268,46 @@ end # The permutation of src_j was already done by permuting its strides before the # kernel launched. acc = zero(promote_type(eltype(data_src), eltype(coeffs))) - @inbounds for j in 1:blk.cols + @inbounds for j in 1:L # TODO is there a more efficient way to do this read? - coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] + coeff = coeffs[u₀ + 1 + i + (j - 1) * L] iszero(coeff) && continue - pst_src, offs_src = structs_src[blk.src_offset + j] + pst_src, offs_src = structs_src[t₀ + j] acc += coeff * op(data_src[_linear_index(coords, pst_src, offs_src)]) end @inbounds data_dst[i_dst] = α * acc + β * data_dst[i_dst] end end -function _launch_unique!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} - nwork = transformer.nwork - nwork == 0 && return nothing +# COV_EXCL_STOP + +function _launch_unique!(data_dst, data_src, op, transformer::UniqueTreeTransformer, α, β, ::Val{N}) where {N} + (; coeffs, structs_dst, structs_src, work_offsets) = transformer + # the subblocks are disjoint, so the total work is bounded by the destination length + ndrange = isempty(coeffs) ? 0 : length(data_dst) + ndrange == 0 && return nothing unique_batched_permute_kernel!(get_backend(data_dst))( - data_dst, data_src, op, transformer.blocks, transformer.work_offsets, α, β, nwork, - Val(N); ndrange = nwork + data_dst, data_src, op, coeffs, structs_dst, structs_src, work_offsets, α, β, + Val(N); ndrange ) return nothing end -function _launch_generic!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} - nwork = transformer.nwork - nwork == 0 && return nothing +function _launch_generic!(data_dst, data_src, op, transformer::GenericTreeTransformer, α, β, ::Val{N}) where {N} + (; coeffs, structs_dst, structs_src, tree_offsets, coeff_offsets, work_offsets) = transformer + # the subblocks are disjoint, so the total work is bounded by the destination length + ndrange = length(tree_offsets) == 1 ? 0 : length(data_dst) + ndrange == 0 && return nothing generic_batched_permute_kernel!(get_backend(data_dst))( - data_dst, data_src, op, transformer.blocks, transformer.work_offsets, - transformer.coeffs, transformer.structs_dst, transformer.structs_src, α, β, nwork, - Val(N); ndrange = nwork + data_dst, data_src, op, coeffs, structs_dst, structs_src, tree_offsets, coeff_offsets, work_offsets, + α, β, Val(N); ndrange ) return nothing end -# we don't want/need the `StridedSubblocks` here because that information is already in the transformer -TensorKit._transform_subblocks( - tdst::TensorMap, tsrc::TensorMap, - transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer} -) = (tdst.data, tsrc.data) - function TensorKit.add_transform_kernel!( dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, - device::DeviceUniqueTreeTransformer, α, β, backend, allocator, ntasks::Int + device::UniqueTreeTransformer, α, β, backend, allocator, ntasks::Int ) N = length(linearize(p)) op = conjsrc ? conj : identity @@ -465,13 +317,13 @@ end function TensorKit.add_transform_kernel!( dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, - device::DeviceGenericTreeTransformer, α, β, backend, allocator, ntasks::Int + device::GenericTreeTransformer, α, β, backend, allocator, ntasks::Int ) N = length(linearize(p)) op = conjsrc ? conj : identity # one-tree blocks are a scaled permutation, which the unique kernel already handles; the # two kernels touch disjoint subblocks so the launch order does not matter - _launch_unique!(dst, src, op, device.unique_blocks, α, β, Val(N)) + _launch_unique!(dst, src, op, device.unique, α, β, Val(N)) _launch_generic!(dst, src, op, device, α, β, Val(N)) return nothing end diff --git a/src/tensors/blockiterators.jl b/src/tensors/blockiterators.jl index 893ad5ee7..0782184d1 100644 --- a/src/tensors/blockiterators.jl +++ b/src/tensors/blockiterators.jl @@ -188,41 +188,3 @@ 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 diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 61a4a6a12..c31974874 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -631,92 +631,99 @@ function add_transform!( else p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels ntasks = get_num_transformer_threads(tdst) - dst, src = _transform_subblocks(tdst, tsrc, transformer) + dst, src = subblockdata(tdst), subblockdata(tsrc) 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) = - StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src) -_transform_subblocks(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, transformer) = - subblocks(tdst), subblocks(tsrc) - # 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. -const TransformSubblocks = Union{StridedSubblocks, SubblockIterator} +# The kernels only depend on the storage, `numind`, `eltype` and the transformer data, not on the sectortype. +function add_transform_kernel!( + dst, src, p, conjsrc::Bool, transformer::UniqueTreeTransformer, α, β, backend, allocator, ntasks::Int + ) + (; coeffs, structs_dst, structs_src) = transformer + taskforeach(eachindex(coeffs), ntasks) do i + dst_i = _dstblock(dst, structs_dst[i]) + src_i, p_i = _srcblock(src, structs_src[i], size(dst_i), p) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( + dst_i, src_i, p_i, conjsrc, α * coeffs[i], β, backend, allocator + ) + end + return nothing +end + +# subblock accessors for the `SubblockLayout`s, where `p` is already applied to strided sources +@inline _dstblock(data::DenseVector, (sz, str, offset)::StridedStructure) = StridedView(data, sz, str, offset) +@inline _dstblock(iter::SubblockIterator, i::Int) = iter[i] +@inline function _srcblock( + data::DenseVector, (str, offset)::Tuple{NTuple{N, Int}, Int}, sz::NTuple{N, Int}, p::Index2Tuple + ) where {N} + return StridedView(data, sz, str, offset), (ntuple(identity, Val(N)), ()) +end +@inline _srcblock(iter::SubblockIterator, i::Int, sz, p::Index2Tuple) = iter[i], p + +# storage type of the workspace for the kernels +_bufferstoragetype(data::DenseVector) = typeof(data) +_bufferstoragetype(iter::SubblockIterator) = storagetype(iter) + function add_transform_kernel!( - dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, - transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int + dst, src, p, conjsrc::Bool, transformer::GenericTreeTransformer, α, β, backend, allocator, ntasks::Int ) + add_transform_kernel!(dst, src, p, conjsrc, transformer.unique, α, β, backend, allocator, ntasks) + nblocks = length(transformer.tree_offsets) - 1 + nblocks == 0 && return nothing + # 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. 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, 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 - # 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:min(length(transformer.data), ntasks) - ] - taskforeach(transformer.data, buffers) do (U, inds_dst, inds_src), buffer - _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) + cp = TO.allocator_checkpoint!(allocator) + @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ + TO.tensoralloc(_bufferstoragetype(dst), bufsize, Val(true), allocator) for _ in 1:min(nblocks, ntasks) + ] + taskforeach(1:nblocks, buffers) do b, buffer + _add_transform_block!(dst, src, p, conjsrc, transformer, b, buffer, α, β, backend, allocator) end + foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) + TO.allocator_reset!(allocator, cp) return nothing end -# `U` is either a scalar coefficient with integer positions (unique fusion), or a recoupling -# matrix with vectors of positions (generic). +# pack → recoupling matmul → unpack for the multi-tree block `b` function _add_transform_block!( - dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, U, inds_dst, inds_src, buffer, - α, β, backend, allocator + dst, src, p, conjsrc::Bool, transformer::GenericTreeTransformer, b::Int, 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, conjsrc, α * 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) - 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, isrc) in enumerate(inds_src) - TO.tensoradd!( - sreshape(view(buffer_src, :, i), sz_src), src[isrc], - ptriv, conjsrc, One(), Zero(), backend, allocator - ) - end + (; coeffs, structs_dst, structs_src, tree_offsets, coeff_offsets) = transformer + trees = (tree_offsets[b] + 1):tree_offsets[b + 1] + L = length(trees) + U = reshape(view(coeffs, (coeff_offsets[b] + 1):coeff_offsets[b + 1]), L, L) + sz = size(_dstblock(dst, structs_dst[first(trees)])) + blocksize = prod(sz) + ptriv = (ntuple(identity, length(sz)), ()) + buffer_dst = StridedView(buffer, (blocksize, L), (1, blocksize), 0) + buffer_src = StridedView(buffer, (blocksize, L), (1, blocksize), blocksize * L) + + # 1. Extract: copy each source block into column i of buffer_src as a flat vector, + # applying the permutation such that the layout is that of the destination. + @timeit_debug GLOBAL_TIMER "dense: pack" for (i, k) in enumerate(trees) + src_k, p_k = _srcblock(src, structs_src[k], sz, p) + TO.tensoradd!(sreshape(view(buffer_src, :, i), sz), src_k, p_k, conjsrc, One(), Zero(), 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!" _recouple!(buffer, buffer_dst, buffer_src, U) - - # 3. Insert: scatter column j of buffer_dst into the destination, applying the - # actual index permutation p and the scaling α in the same tensoradd! call. - @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (j, idst) in enumerate(inds_dst) - TO.tensoradd!( - dst[idst], sreshape(view(buffer_dst, :, j), sz_src), - p, false, α, β, 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!" _recouple!(buffer, buffer_dst, buffer_src, U) + + # 3. Insert: add column j of buffer_dst to the destination, applying the scaling α. + @timeit_debug GLOBAL_TIMER "dense: unpack" for (j, k) in enumerate(trees) + TO.tensoradd!( + _dstblock(dst, structs_dst[k]), sreshape(view(buffer_dst, :, j), sz), + ptriv, false, α, β, backend, allocator + ) end return nothing end @@ -734,7 +741,7 @@ function _recouple!( end # on the CPU, views of reinterpreted arrays are not `StridedMatrix`, so `mul!` would not use BLAS function _recouple!( - buffer::CPUStorage{Complex{R}}, buffer_dst::StridedView, buffer_src::StridedView, U::Matrix{R} + buffer::CPUStorage{Complex{R}}, buffer_dst::StridedView, buffer_src::StridedView, U::StridedMatrix{R} ) where {R <: LinearAlgebra.BlasReal} rbuffer = reinterpret(R, buffer) LinearAlgebra.BLAS.gemm!('N', 'T', one(R), _realview(rbuffer, buffer_src), U, zero(R), _realview(rbuffer, buffer_dst)) diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index 1aeb0070c..3de8c9009 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -496,9 +496,6 @@ end # Getting and setting the data at the subblock level # -------------------------------------------------- -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))) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 990e50d2d..2a1956422 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -3,8 +3,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`](@ref)), together with the recoupling +The transformers only store how subblocks map onto each other, together with the recoupling coefficients, and are therefore independent of the sectortype once constructed. The 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 @@ -12,7 +11,6 @@ of `tsrc` itself; when `conjsrc` is `true` the fusion trees that are transformed """ abstract type TreeTransformer end -# `StridedSubblocks` report their storage as the `StridedView` parent type, which is `Memory` there @static if isdefined(Core, :Memory) const CPUStorage{T} = Union{Array{T}, Memory{T}} else @@ -33,67 +31,133 @@ function recoupling_scalartype(::Type{A}, ::Type{Tₛ}) where {A, Tₛ <: Number return Tₛ <: Real ? real(T) : complex(T) end -# (coefficient, destination position, source position) -const UniqueTransformerData{T} = Tuple{T, Int, Int} - """ - UniqueTreeTransformer{T, N} <: TreeTransformer + SubblockLayout -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. +How a transformer locates the subblocks of a tensor. [`StridedLayout`](@ref) addresses the flat +data vector of a `TensorMap` through the `(sizes, strides, offset)` of each subblock, while +[`PositionalLayout`](@ref) indexes the [`subblocks`](@ref) of any other `AbstractTensorMap` by their +position in the canonical order of [`fusiontrees`](@ref). """ -struct UniqueTreeTransformer{T, N} <: TreeTransformer - data::Vector{UniqueTransformerData{T}} - structure_dst::Vector{StridedStructure{N}} - structure_src::Vector{StridedStructure{N}} -end +abstract type SubblockLayout end +struct StridedLayout <: SubblockLayout end +struct PositionalLayout <: SubblockLayout end + +subblocklayout(t::AbstractTensorMap) = subblocklayout(typeof(t)) +subblocklayout(::Type{<:AbstractTensorMap}) = PositionalLayout() +subblocklayout(::Type{<:TensorMap}) = StridedLayout() -# (recoupling matrix, destination positions, source positions) -const GenericTransformerData{T} = Tuple{Matrix{T}, Vector{Int}, Vector{Int}} +# the object through which the kernels access the subblocks, as located by `subblocklayout` +subblockdata(t::TensorMap) = t.data +subblockdata(t::AbstractTensorMap) = subblocks(t) """ - GenericTreeTransformer{T, N} <: TreeTransformer + UniqueTreeTransformer{T, VC, VD, VS, VO} <: 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)`, where `U[j, i]` maps source `i` onto destination `j`. -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. +Tree transformation for `UniqueFusion` sectors, where every source subblock maps onto a single +destination subblock with a scalar coefficient. Entry `i` describes the transformation +`dst_i = β * dst_i + α * coeffs[i] * permute(src_i, p)` through the locators `structs_dst[i]` and +`structs_src[i]`, whose type depends on the [`SubblockLayout`](@ref) of either tensor: + +- `StridedLayout`: `structs_dst[i]` is the `(sizes, strides, offset)` of the destination subblock, + and `structs_src[i]` the `(strides, offset)` of the source subblock with its strides already + permuted by `p`, such that both share the destination sizes. +- `PositionalLayout`: the subblock positions, and `p` is applied when transforming. + +The `n + 1` `work_offsets` are the number of destination elements preceding every entry, followed by +their total, such that the work can be distributed per element. All fields are vectors of `isbits` +elements, such that the same structure can be moved to other devices with `Adapt.adapt`. """ -struct GenericTreeTransformer{T, N} <: TreeTransformer - data::Vector{GenericTransformerData{T}} - structure_dst::Vector{StridedStructure{N}} - structure_src::Vector{StridedStructure{N}} +struct UniqueTreeTransformer{ + T, VC <: AbstractVector{T}, VD <: AbstractVector, VS <: AbstractVector, VO <: AbstractVector{Int}, + } <: TreeTransformer + coeffs::VC + structs_dst::VD + structs_src::VS + work_offsets::VO end -function UniqueTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A} +Adapt.@adapt_structure UniqueTreeTransformer + +function UniqueTreeTransformer( + ::Type{A}, layout_dst::SubblockLayout, layout_src::SubblockLayout, transform, p, Vdst, Vsrc, conjsrc::Bool + ) where {A} t₀ = Base.time() spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) - data = Vector{UniqueTransformerData{T}}(undef, length(src_trees)) + # the transformation is a bijection between the trees, so we use the canonical destination order + coeffs = Vector{T}(undef, length(src_trees)) + inds_src = Vector{Int}(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₂)) _, (_, idst) = gettoken(dst_trees, f_dst) - data[isrc] = (convert(T, coeff), idst, isrc) + coeffs[idst] = convert(T, coeff) + inds_src[idst] = isrc end structure_dst = degeneracystructure(Vdst).subblockstructure structure_src = degeneracystructure(Vsrc).subblockstructure - transformer = UniqueTreeTransformer(data, structure_dst, structure_src) + structs_dst = layout_dst isa StridedLayout ? structure_dst : Base.OneTo(length(structure_dst)) + p′ = linearize(p) + structs_src = map(isrc -> _src_locator(layout_src, structure_src, isrc, p′), inds_src) + work_offsets = pushfirst!(cumsum(map(s -> prod(first(s)), structure_dst)), 0) + transformer = UniqueTreeTransformer(coeffs, structs_dst, structs_src, work_offsets) Δt = Base.time() - t₀ - @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = length(data), Δt) + @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = length(coeffs), Δt) return transformer end -function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A} +# locators of subblock `i` for the different `SubblockLayout`s +_dst_locator(::StridedLayout, structure, i::Int) = structure[i] +_dst_locator(::PositionalLayout, structure, i::Int) = i +function _src_locator(::StridedLayout, structure, i::Int, p::IndexTuple) + _, strides, offset = structure[i] + return (TupleTools.getindices(strides, p), offset) +end +_src_locator(::PositionalLayout, structure, i::Int, p::IndexTuple) = i + +dstlocatortype(::StridedLayout, N) = StridedStructure{N} +dstlocatortype(::PositionalLayout, N) = Int +srclocatortype(::StridedLayout, N) = Tuple{NTuple{N, Int}, Int} +srclocatortype(::PositionalLayout, N) = Int + +""" + GenericTreeTransformer{T, U, VC, VD, VS, VO} <: 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 square +recoupling matrix. Blocks that consist of a single tree are stored as the +[`UniqueTreeTransformer`](@ref) `unique`, while for the others block `b` consists of the entries +`tree_offsets[b] + 1:tree_offsets[b + 1]` of the locators `structs_dst` and `structs_src`, with +the same conventions as for `UniqueTreeTransformer`, and its recoupling matrix `U` is stored in +column-major order as `coeffs[coeff_offsets[b] + 1:coeff_offsets[b + 1]]`, where `U[j, i]` maps +source `i` onto destination `j`. The `work_offsets` count the number of destination elements +preceding every block, followed by their total, and the blocks are sorted by decreasing work. +""" +struct GenericTreeTransformer{ + T, U <: UniqueTreeTransformer{T}, VC <: AbstractVector{T}, VD <: AbstractVector, VS <: AbstractVector, + VO <: AbstractVector{Int}, + } <: TreeTransformer + unique::U + coeffs::VC + structs_dst::VD + structs_src::VS + tree_offsets::VO + coeff_offsets::VO + work_offsets::VO +end + +Adapt.@adapt_structure GenericTreeTransformer + +function GenericTreeTransformer( + ::Type{A}, layout_dst::SubblockLayout, layout_src::SubblockLayout, transform, p, Vdst, Vsrc, conjsrc::Bool + ) where {A} t₀ = Base.time() spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) # the fusion blocks that are transformed are those of the adjoint space for a conjugated source @@ -105,7 +169,7 @@ function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bo fblocks = @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" fusionblocks(Vsrc′) nblocks = length(fblocks) - data = Vector{GenericTransformerData{T}}(undef, nblocks) + data = Vector{Tuple{Matrix{T}, Vector{Int}, Vector{Int}}}(undef, nblocks) weights = Vector{Int}(undef, nblocks) nthreads = get_num_manipulation_threads() @@ -113,6 +177,7 @@ function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bo taskforeach(1:nblocks, nthreads) do i fs_src = fblocks[i] fs_dst, U = transform(fs_src) + LinearAlgebra.checksquare(U) @timeit_debug GLOBAL_TIMER "bookkeeping: subblock positions" begin # the token into the fusion tree `Indices` is the subblock position inds_src = map(fusiontrees(fs_src)) do (f₁, f₂) @@ -137,7 +202,36 @@ function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bo # sort by (approximate) weight to facilitate multi-threading strategies @timeit_debug GLOBAL_TIMER "bookkeeping: sort" Base.permute!(data, sortperm(weights; rev = true)) - transformer = GenericTreeTransformer{T, numind(Vdst)}(data, structure_dst, structure_src) + + N = numind(Vdst) + p′ = linearize(p) + unique_coeffs = T[] + unique_dst = dstlocatortype(layout_dst, N)[] + unique_src = srclocatortype(layout_src, N)[] + unique_offsets = [0] + coeffs = T[] + structs_dst = dstlocatortype(layout_dst, N)[] + structs_src = srclocatortype(layout_src, N)[] + tree_offsets, coeff_offsets, work_offsets = [0], [0], [0] + for (U, inds_dst, inds_src) in data + if length(U) == 1 + push!(unique_coeffs, only(U)) + push!(unique_dst, _dst_locator(layout_dst, structure_dst, only(inds_dst))) + push!(unique_src, _src_locator(layout_src, structure_src, only(inds_src), p′)) + push!(unique_offsets, last(unique_offsets) + prod(first(structure_dst[only(inds_dst)]))) + else + append!(coeffs, U) + append!(structs_dst, (_dst_locator(layout_dst, structure_dst, i) for i in inds_dst)) + append!(structs_src, (_src_locator(layout_src, structure_src, i, p′) for i in inds_src)) + push!(tree_offsets, length(structs_dst)) + push!(coeff_offsets, length(coeffs)) + push!(work_offsets, last(work_offsets) + size(U, 1) * prod(structure_dst[first(inds_dst)][1])) + end + end + unique = UniqueTreeTransformer(unique_coeffs, unique_dst, unique_src, unique_offsets) + transformer = GenericTreeTransformer( + unique, coeffs, structs_dst, structs_src, tree_offsets, coeff_offsets, work_offsets + ) Δt = Base.time() - t₀ @debug( @@ -152,62 +246,68 @@ function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bo end """ - buffersize(transformer::TreeTransformer) -> Int + buffersize(transformer::GenericTreeTransformer) -> 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. +block, i.e. twice the number of destination elements of that block. """ -buffersize(::UniqueTreeTransformer) = 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 -end +buffersize(transformer::GenericTreeTransformer) = 2 * maximum(diff(transformer.work_offsets); init = 0) -function treetransformertype(::Type{A}, Vdst, Vsrc) where {A} +function treetransformertype(::Type{A}, layout_dst, layout_src, Vdst, Vsrc) where {A} I = sectortype(Vdst) T = recoupling_scalartype(A, sectorscalartype(I)) N = numind(Vdst) - return FusionStyle(I) == UniqueFusion() ? UniqueTreeTransformer{T, N} : GenericTreeTransformer{T, N} + D, S = dstlocatortype(layout_dst, N), srclocatortype(layout_src, N) + if FusionStyle(I) == UniqueFusion() + Dᵤ = layout_dst isa StridedLayout ? Vector{D} : Base.OneTo{Int} + return UniqueTreeTransformer{T, Vector{T}, Dᵤ, Vector{S}, Vector{Int}} + end + U = UniqueTreeTransformer{T, Vector{T}, Vector{D}, Vector{S}, Vector{Int}} + return GenericTreeTransformer{T, U, Vector{T}, Vector{D}, Vector{S}, Vector{Int}} end function TreeTransformer( - ::Type{A}, transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S}, conjsrc::Bool + ::Type{A}, layout_dst::SubblockLayout, layout_src::SubblockLayout, transform::Function, p, + Vdst::HomSpace{S}, Vsrc::HomSpace{S}, conjsrc::Bool ) where {A, S} I = sectortype(Vdst) return FusionStyle(I) == UniqueFusion() ? - UniqueTreeTransformer(A, transform, p, Vdst, Vsrc, conjsrc) : - GenericTreeTransformer(A, transform, p, Vdst, Vsrc, conjsrc) + UniqueTreeTransformer(A, layout_dst, layout_src, transform, p, Vdst, Vsrc, conjsrc) : + GenericTreeTransformer(A, layout_dst, layout_src, transform, p, Vdst, Vsrc, conjsrc) end # braid is special because it has levels function treebraider( tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple ) - return treebraider(storagetype(tdst), space(tdst), space(tsrc), p, conjsrc, levels) + return treebraider( + storagetype(tdst), subblocklayout(tdst), subblocklayout(tsrc), space(tdst), space(tsrc), p, conjsrc, levels + ) end @cached function treebraider( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple - )::treetransformertype(A, Vdst, Vsrc) where {TA} + A::Type{TA}, layout_dst::SubblockLayout, layout_src::SubblockLayout, + Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple + )::treetransformertype(A, layout_dst, layout_src, Vdst, Vsrc) where {TA} 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′))) fusiontreebraider(f) = braid(f, p′, levels″) - return TreeTransformer(A, fusiontreebraider, p, Vdst, Vsrc, conjsrc) + return TreeTransformer(A, layout_dst, layout_src, fusiontreebraider, p, Vdst, Vsrc, conjsrc) end function treetransposer(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool) - return treetransposer(storagetype(tdst), space(tdst), space(tsrc), p, conjsrc) + return treetransposer( + storagetype(tdst), subblocklayout(tdst), subblocklayout(tsrc), space(tdst), space(tsrc), p, conjsrc + ) end @cached function treetransposer( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool - )::treetransformertype(A, Vdst, Vsrc) where {TA} + A::Type{TA}, layout_dst::SubblockLayout, layout_src::SubblockLayout, + Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool + )::treetransformertype(A, layout_dst, layout_src, Vdst, Vsrc) where {TA} p′ = conjsrc ? adjointtensorindices(Vsrc, p) : p fusiontreetransform(f) = transpose(f, p′) - return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) + return TreeTransformer(A, layout_dst, layout_src, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end # default cachestyle is GlobalLRUCache