change blocktype of TensorMap to StridedView

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

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