change blocktype of TensorMap to StridedView

Project.toml

@@ -66,3 +66,7 @@ TensorOperations = "5.5.2, 5.6"
 TupleTools = "1.5"
 VectorInterface = "0.6"
 julia = "1.10"
+
+[sources]
+Strided = {url = "https://github.com/QuantumKitHub/Strided.jl", rev = "ksh/big"}
+MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "ksh/stridedviews"}

ext/TensorKitEnzymeExt/linalg.jl

@@ -129,15 +129,15 @@ function EnzymeRules.reverse(
 end
 function EnzymeRules.forward(
         config::EnzymeRules.FwdConfigWidth{1},
-        ::Type{RT},
         func::Const{typeof(tr)},
+        ::Type{RT},
         A::Annotation{<:AbstractTensorMap},
     ) where {RT}
     y = EnzymeRules.needs_primal(config) ? tr(A.val) : nothing
     Δy = if EnzymeRules.needs_shadow(config) && !isa(A, Const)
         tr(A.dval)
     elseif EnzymeRules.needs_shadow(config)
-        zero(eltype(A.dval))
+        zero(eltype(A.val))
     else
         nothing
     end

ext/TensorKitFiniteDifferencesExt.jl

@@ -4,6 +4,13 @@ using TensorKit
 using TensorKit: sqrtdim, invsqrtdim, SectorVector
 using VectorInterface: scale!
 using FiniteDifferences
+using Strided: StridedView
+
+function FiniteDifferences.to_vec(x::StridedView)
+    x_vec, from_vec = FiniteDifferences.to_vec(Array(x))
+    StridedView_from_vec(x_vec) = StridedView(from_vec(x_vec))
+    return x_vec, StridedView_from_vec
+end
 
 function FiniteDifferences.to_vec(t::AbstractTensorMap)
     # convert to vector of vectors to make use of existing functionality

src/spaces/homspace.jl

@@ -187,20 +187,6 @@ function fusionblocks(W::HomSpace)
     return fblocks
 end
 
-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
-
 # Operations on HomSpaces
 # -----------------------
 """
@@ -353,3 +339,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/spaces/structure.jl

@@ -187,12 +187,4 @@ See also [`sectorstructure`](@ref), [`blockstructure`](@ref), [`subblockstructur
     return DegeneracyStructure(blockoffset, blockvalues, structurevalues)
 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
-
 CacheStyle(::typeof(degeneracystructure), ::HomSpace) = GlobalLRUCache()

src/tensors/tensor.jl

@@ -467,31 +467,39 @@ 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 = Memory{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, ((d₁, d₂), r)), newstate = next
+    return c => StridedView(iter.t.data, (d₁, d₂), (1, d₁), first(r) - 1), 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₂))
+        offset = first(r) - 1
     else
         # if 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₂))
+        offset = 0
     end
+    return StridedView(iter.t.data, (d₁, d₂), (1, d₁), offset)
 end
 
 # Getting and setting the data at the subblock level

src/tensors/vectorinterface.jl

@@ -129,9 +129,9 @@ function VectorInterface.inner(tx::TensorMap, ty::TensorMap)
     else
         T = VectorInterface.promote_inner(tx, ty)
         s = zero(T)
-        for c in blocksectors(tx)
-            bx = parent(block(tx, c)) # matrix structure (reshape) does not matter
-            by = parent(block(ty, c)) # but does lead to slower path in inner
+        for (c, (_, r)) in pairs(blockstructure(space(tx)))
+            bx = view(tx.data, r) # matrix structure (reshape) does not matter
+            by = view(ty.data, r) # but does lead to slower path in inner
             s += convert(T, dim(c)) * inner(bx, by)
         end
     end

test/enzyme-linalg/tr.jl

@@ -22,18 +22,12 @@ TDs = is_ci ? (Duplicated,) : (Const, Duplicated)
         @testset "tr reverse: RT $RT, TD $TD" for RT in rRTs, TD in TDs
             EnzymeTestUtils.test_reverse(tr, RT, (D1, TD); atol, rtol)
             EnzymeTestUtils.test_reverse(tr, RT, (D2, TD); atol, rtol)
-            # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457
-            @static if VERSION ≥ v"1.11.0-rc"
-                EnzymeTestUtils.test_reverse(tr, RT, (D3, TD); atol, rtol)
-            end
+            EnzymeTestUtils.test_reverse(tr, RT, (D3, TD); atol, rtol)
         end
         @testset "tr forward: RT $RT, TD $TD" for RT in fRTs, TD in TDs
             EnzymeTestUtils.test_forward(tr, RT, (D1, TD); atol, rtol)
             EnzymeTestUtils.test_forward(tr, RT, (D2, TD); atol, rtol)
-            # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457
-            @static if VERSION ≥ v"1.11.0-rc"
-                EnzymeTestUtils.test_forward(tr, RT, (D3, TD); atol, rtol)
-            end
+            EnzymeTestUtils.test_forward(tr, RT, (D3, TD); atol, rtol)
         end
     end
 end

test/tensors/diagonal.jl

@@ -151,8 +151,8 @@ diagspacelist = (
     @timedtestset "Trace, Multiplication and inverse" begin
         t1 = DiagonalTensorMap(rand(Float64, reduceddim(V)), V)
         t2 = DiagonalTensorMap(rand(ComplexF64, reduceddim(V)), V)
-        @test tr(TensorMap(t1)) == @constinferred tr(t1)
-        @test tr(TensorMap(t2)) == @constinferred tr(t2)
+        @test tr(TensorMap(t1)) ≈ @constinferred tr(t1)
+        @test tr(TensorMap(t2)) ≈ @constinferred tr(t2)
         @test TensorMap(@constinferred t1 * t2) ≈ TensorMap(t1) * TensorMap(t2)
         @test TensorMap(@constinferred t1 \ t2) ≈ TensorMap(t1) \ TensorMap(t2)
         @test TensorMap(@constinferred t1 / t2) ≈ TensorMap(t1) / TensorMap(t2)