diff --git a/Project.toml b/Project.toml index 5edabec65..291c1524f 100644 --- a/Project.toml +++ b/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"} diff --git a/ext/TensorKitEnzymeExt/linalg.jl b/ext/TensorKitEnzymeExt/linalg.jl index 35e3eaaf8..736b46767 100644 --- a/ext/TensorKitEnzymeExt/linalg.jl +++ b/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 diff --git a/ext/TensorKitFiniteDifferencesExt.jl b/ext/TensorKitFiniteDifferencesExt.jl index 8a19ae62a..5d692d48a 100644 --- a/ext/TensorKitFiniteDifferencesExt.jl +++ b/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 diff --git a/src/spaces/homspace.jl b/src/spaces/homspace.jl index 8460f7578..89f6ab584 100644 --- a/src/spaces/homspace.jl +++ b/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 diff --git a/src/spaces/structure.jl b/src/spaces/structure.jl index 9a4ddf5c0..78b6fbbb4 100644 --- a/src/spaces/structure.jl +++ b/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() diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl index ed5682367..fef2c97e3 100644 --- a/src/tensors/tensor.jl +++ b/src/tensors/tensor.jl @@ -467,15 +467,22 @@ 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) @@ -483,15 +490,16 @@ function Base.getindex(iter::BlockIterator{<:TensorMap}, c::Sector) 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 diff --git a/src/tensors/vectorinterface.jl b/src/tensors/vectorinterface.jl index 460000787..a439337b1 100644 --- a/src/tensors/vectorinterface.jl +++ b/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 diff --git a/test/enzyme-linalg/tr.jl b/test/enzyme-linalg/tr.jl index f42e02ee1..1ba0c2df7 100644 --- a/test/enzyme-linalg/tr.jl +++ b/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 diff --git a/test/tensors/diagonal.jl b/test/tensors/diagonal.jl index 8a4c46161..c2dfcd369 100644 --- a/test/tensors/diagonal.jl +++ b/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)