From fb14a4109ce139d3f4a4ac3eff814e5602a32704 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Mon, 5 Jan 2026 17:03:23 +0100 Subject: [PATCH] change blocktype of TensorMap to `StridedView` --- src/spaces/homspace.jl | 122 +++++++++++++++++++++++++++++++++++++++++ src/tensors/tensor.jl | 33 ++++++----- 2 files changed, 142 insertions(+), 13 deletions(-) diff --git a/src/spaces/homspace.jl b/src/spaces/homspace.jl index 8460f7578..103074c7c 100644 --- a/src/spaces/homspace.jl +++ b/src/spaces/homspace.jl @@ -353,3 +353,125 @@ function removeunit(P::HomSpace, ::Val{i}) where {i} return codomain(P) ← removeunit(domain(P), Val(i - numout(P))) end end + +# Block and fusion tree ranges: structure information for building tensors +#-------------------------------------------------------------------------- + +# sizes, strides, offset +const StridedStructure{N} = Tuple{NTuple{N, Int}, NTuple{N, Int}, Int} + +struct FusionBlockStructure{I, N, F₁, F₂} + totaldim::Int + blockstructure::SectorDict{I, StridedStructure{2}} + fusiontreelist::Vector{Tuple{F₁, F₂}} + fusiontreestructure::Vector{StridedStructure{N}} + fusiontreeindices::FusionTreeDict{Tuple{F₁, F₂}, Int} +end + +function fusionblockstructuretype(W::HomSpace) + N₁ = length(codomain(W)) + N₂ = length(domain(W)) + N = N₁ + N₂ + I = sectortype(W) + F₁ = fusiontreetype(I, N₁) + F₂ = fusiontreetype(I, N₂) + return FusionBlockStructure{I, N, F₁, F₂} +end + +@cached function fusionblockstructure(W::HomSpace)::fusionblockstructuretype(W) + codom = codomain(W) + dom = domain(W) + N₁ = length(codom) + N₂ = length(dom) + I = sectortype(W) + F₁ = fusiontreetype(I, N₁) + F₂ = fusiontreetype(I, N₂) + + # output structure + blockstructure = SectorDict{I, StridedStructure{2}}() # size, strides, offset + fusiontreelist = Vector{Tuple{F₁, F₂}}() + fusiontreestructure = Vector{StridedStructure{N₁ + N₂}}() # size, strides, offset + + # temporary data structures + splittingtrees = Vector{F₁}() + splittingstructure = Vector{Tuple{Int, Int}}() + + # main computational routine + blockoffset = 0 + for c in blocksectors(W) + empty!(splittingtrees) + empty!(splittingstructure) + + offset₁ = 0 + for f₁ in fusiontrees(codom, c) + push!(splittingtrees, f₁) + d₁ = dim(codom, f₁.uncoupled) + push!(splittingstructure, (offset₁, d₁)) + offset₁ += d₁ + end + blockdim₁ = offset₁ + strides = (1, blockdim₁) + + offset₂ = 0 + for f₂ in fusiontrees(dom, c) + s₂ = f₂.uncoupled + d₂ = dim(dom, s₂) + for (f₁, (offset₁, d₁)) in zip(splittingtrees, splittingstructure) + push!(fusiontreelist, (f₁, f₂)) + totaloffset = blockoffset + offset₂ * blockdim₁ + offset₁ + subsz = (dims(codom, f₁.uncoupled)..., dims(dom, f₂.uncoupled)...) + @assert !any(isequal(0), subsz) + substr = _subblock_strides(subsz, (d₁, d₂), strides) + push!(fusiontreestructure, (subsz, substr, totaloffset)) + end + offset₂ += d₂ + end + blockdim₂ = offset₂ + blocksize = (blockdim₁, blockdim₂) + blocklength = blockdim₁ * blockdim₂ + blockrange = (blockoffset + 1):(blockoffset + blocklength) + blockstructure[c] = (blocksize, strides, blockoffset) + blockoffset = last(blockrange) + end + + fusiontreeindices = sizehint!( + FusionTreeDict{Tuple{F₁, F₂}, Int}(), length(fusiontreelist) + ) + for (i, f₁₂) in enumerate(fusiontreelist) + fusiontreeindices[f₁₂] = i + end + totaldim = blockoffset + structure = FusionBlockStructure( + totaldim, blockstructure, fusiontreelist, fusiontreestructure, fusiontreeindices + ) + return structure +end + +function _subblock_strides(subsz, sz, str) + sz_simplify = Strided.StridedViews._simplifydims(sz, str) + strides = Strided.StridedViews._computereshapestrides(subsz, sz_simplify...) + isnothing(strides) && + throw(ArgumentError("unexpected error in computing subblock strides")) + return strides +end + +function CacheStyle(::typeof(fusionblockstructure), W::HomSpace) + return GlobalLRUCache() +end + +# Diagonal ranges +#---------------- +# TODO: is this something we want to cache? +function diagonalblockstructure(W::HomSpace) + ((numin(W) == numout(W) == 1) && domain(W) == codomain(W)) || + throw(SpaceMismatch("Diagonal only support on V←V with a single space V")) + structure = SectorDict{sectortype(W), UnitRange{Int}}() # range + offset = 0 + dom = domain(W)[1] + for c in blocksectors(W) + d = dim(dom, c) + structure[c] = offset .+ (1:d) + offset += d + end + return structure +end diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index ed5682367..47bdfbd46 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -467,31 +467,38 @@ block(t::TensorMap, c::Sector) = blocks(t)[c] blocks(t::TensorMap) = BlockIterator(t, blockstructure(space(t))) -function blocktype(::Type{TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: Vector{T}} - return Base.ReshapedArray{T, 2, SubArray{T, 1, A, Tuple{UnitRange{Int}}, true}, Tuple{}} +function blocktype(::Type{TT}) where {TT <: TensorMap} + A = storagetype(TT) + T = eltype(A) + @static if isdefined(Core, :Memory) # StridedViews normalizes parent types! + if A <: Vector{T} + A = GenericMemory{T} + end + end + return StridedView{T, 2, A, typeof(identity)} end function Base.iterate(iter::BlockIterator{<:TensorMap}, state...) next = iterate(pairs(iter.structure), state...) isnothing(next) && return next - (c, (sz, r)), newstate = next - return c => reshape(view(iter.t.data, r), sz), newstate + (c, (sz, str, offset)), newstate = next + return c => StridedView(iter.t.data, sz, str, offset), newstate end function Base.getindex(iter::BlockIterator{<:TensorMap}, c::Sector) sectortype(iter.t) === typeof(c) || throw(SectorMismatch()) - found, token = gettoken(iter.structure, c) - if found - (d₁, d₂), r = gettokenvalue(iter.structure, token) - return reshape(view(iter.t.data, r), (d₁, d₂)) - else - # if c is not a key, at least one of the two dimensions will be zero: + (d₁, d₂), (s₁, s₂), offset = get(iter.structure, c) do + # is c is not a key, at least one of the two dimensions will be zero: # it then does not matter where exactly we construct a view in `t.data`, # as it will have length zero anyway - d₁ = blockdim(codomain(iter.t), c) - d₂ = blockdim(domain(iter.t), c) - return reshape(view(iter.t.data, 1:(d₁ * d₂)), (d₁, d₂)) + d₁′ = blockdim(codomain(iter.t), c) + d₂′ = blockdim(domain(iter.t), c) + s₁ = 1 + s₂ = 0 + offset = 0 + return (d₁′, d₂′), (s₁, s₂), offset end + return StridedView(iter.t.data, (d₁, d₂), (s₁, s₂), offset) end # Getting and setting the data at the subblock level