From 8d5d02f730adba8cb6d5c9524201daa160817c09 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 11:42:07 -0400 Subject: [PATCH 1/4] Precompute planar trace kernels in a `TraceTransformer` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `planartrace!` now precomputes the subblock positions and recoupling coefficients in a cached `TraceTransformer`, and executes them with a sector-independent kernel, analogous to the index manipulations. This replaces the per-tree loop that built ad-hoc single-pair `FusionTreeBlock`s. - `planar_trace` is defined for `Union{FusionTreePair, FusionTreeBlock}`, with unique-fusion pairs returning `dst => coeff` and blocks `dst => U`, and throws when the traced sectors are not dual to each other - the transformer skips such (sub)blocks upfront and groups the data by destination, such that every destination subblock is written once and `β` is handled inside the kernel - blocks with several destination trees are traced (and permuted) into a buffer, recoupled with a single matmul and added into the destination - `FusionTreeBlock`s hash their uncoupled sectors and duality flags Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 1 + docs/src/lib/fusiontrees.md | 1 + src/fusiontrees/duality_manipulations.jl | 115 ++++++++++++++--------- src/fusiontrees/fusiontrees.jl | 9 +- src/planar/planaroperations.jl | 83 +++++++++++++--- src/tensors/treetransformers.jl | 100 ++++++++++++++++++++ test/symmetries/doubletree.jl | 24 ++--- 7 files changed, 258 insertions(+), 75 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index a6ab14c3b..5ea2651e7 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -30,6 +30,7 @@ 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) +- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s (internal) ### Deprecated diff --git a/docs/src/lib/fusiontrees.md b/docs/src/lib/fusiontrees.md index 57033eca6..759eeec12 100644 --- a/docs/src/lib/fusiontrees.md +++ b/docs/src/lib/fusiontrees.md @@ -48,4 +48,5 @@ repartition(src::Union{FusionTreePair, FusionTreeBlock}, N::Int) Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, levels::Index2Tuple) permute(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) +planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, q::Index2Tuple) ``` diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index cae44aca3..9e613d50e 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -584,58 +584,85 @@ CacheStyle(::typeof(fstranspose), k::FSBTransposeKey{I}) where {I} = # -> composite manipulations that depend on the duality (rigidity) and pivotal structure # -> planar manipulations that do not require braiding, everything is in Fsymbol (A/Bsymbol) -function planar_trace((f₁, f₂)::FusionTreePair, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) +# leg positions of a planar trace of a tree pair with `N₁` outgoing and `N₂` incoming legs, after +# repartitioning all legs to the outgoing side: the traced pairs `q′`, the remaining legs `r` in +# increasing order and the permutation `p′` that is applied after removing the traced legs +function _planar_trace_indices(::Val{N₁}, ::Val{N₂}, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) where {N₁, N₂} length(q₁) == length(q₂) || throw(ArgumentError(lazy"trace index tuples q₁ and q₂ must have equal length, got $(length(q₁)) and $(length(q₂))")) - I = sectortype(f₁) N = length(p₁) + length(p₂) + 2 * length(q₁) - length(f₁) + length(f₂) == N || - throw(ArgumentError(lazy"fusion tree pair has $(length(f₁) + length(f₂)) indices, but permutation expects $N = $(length(p₁)) + $(length(p₂)) + 2×$(length(q₁))")) - if isempty(q₁) - return transpose((f₁, f₂), (p₁, p₂)) - end - - linearindex = ( - ntuple(identity, Val(length(f₁)))..., - reverse(length(f₁) .+ ntuple(identity, Val(length(f₂))))..., - ) - + N₁ + N₂ == N || + throw(ArgumentError(lazy"fusion tree pair has $(N₁ + N₂) indices, but permutation expects $N = $(length(p₁)) + $(length(p₂)) + 2×$(length(q₁))")) + linearindex = (ntuple(identity, Val(N₁))..., reverse(N₁ .+ ntuple(identity, Val(N₂)))...) q₁′ = TupleTools.getindices(linearindex, q₁) q₂′ = TupleTools.getindices(linearindex, q₂) - p₁′, p₂′ = let q′ = (q₁′..., q₂′...) - ( - map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₁)), - map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₂)), - ) - end + q′ = (q₁′..., q₂′...) + p₁′ = map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₁)) + p₂′ = map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₂)) + r = TupleTools.sort(TupleTools.getindices(linearindex, (p₁..., p₂...))) + return (p₁′, p₂′), (q₁′, q₂′), r +end - T = fusionscalartype(I) - F₁ = fusiontreetype(I, length(p₁)) - F₂ = fusiontreetype(I, length(p₂)) - newtrees = FusionTreeDict{Tuple{F₁, F₂}, T}() - if FusionStyle(I) isa UniqueFusion - (f₁′, f₂′), coeff′ = repartition((f₁, f₂), N) - for (f₁′′, coeff′′) in planar_trace(f₁′, (q₁′, q₂′)) - (f12′′′, coeff′′′) = transpose((f₁′′, f₂′), (p₁′, p₂′)) - coeff = coeff′ * coeff′′ * coeff′′′ - iszero(coeff) || (newtrees[f12′′′] = get(newtrees, f12′′′, zero(coeff)) + coeff) - end +# the sectors of the traced legs have to be dual to each other for a nonzero trace, where the +# sector of an incoming leg is the dual of the corresponding uncoupled sector of the fusion tree +function trace_sectors_match(src::Union{FusionTreePair, FusionTreeBlock}, (q₁, q₂)::Index2Tuple) + f₁, f₂ = src isa FusionTreePair ? src : first(fusiontrees(src)) + N₁ = length(f₁) + legsector(i) = i <= N₁ ? f₁.uncoupled[i] : dual(f₂.uncoupled[i - N₁]) + return all(map((i, j) -> legsector(i) == dual(legsector(j)), q₁, q₂)) +end + +""" + planar_trace((f₁, f₂)::FusionTreePair, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) -> (f₁′, f₂′) => coeff + planar_trace(src::FusionTreeBlock, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) -> dst => U + +Perform a planar trace of the legs `q₁[i]` with `q₂[i]` of a fusion tree pair, followed by a +planar permutation `(p₁, p₂)` of the remaining legs. The sectors of the traced legs have to be +dual to each other. + +For `FusionStyle(I) isa UniqueFusion`, the input is a single fusion tree pair, which is mapped +to a single output tree pair with a coefficient. For generic `FusionStyle`s, the input and +output consist of `FusionTreeBlock`s, where `src` and `dst` contain all tree pairs of the +traced and remaining legs respectively, and `U[j, i]` is the coefficient of the `j`th tree pair +of `dst` in the trace of the `i`th tree pair of `src`. +""" +function planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) + src isa FusionTreePair && @assert FusionStyle(src) === UniqueFusion() + p′, q′, r = _planar_trace_indices(Val(numout(src)), Val(numin(src)), (p₁, p₂), (q₁, q₂)) + trace_sectors_match(src, (q₁, q₂)) || + throw(ArgumentError(lazy"sectors of the traced legs $q₁ and $q₂ are not dual to each other")) + isempty(q₁) && return transpose(src, (p₁, p₂)) + + # bend all legs to the outgoing side, where the trace only acts on the splitting trees + src′, U₁ = repartition(src, numind(src)) + traced, U₂ = _planar_trace_splitting(src′, q′, r) + # for multiple fusion channels, the remaining legs need not couple to the unit + if traced isa FusionTreeBlock && isempty(traced) + I = sectortype(src) + F = fusiontreetype(I, length(p₁), length(p₂)) + return FusionTreeBlock(F[]) => zeros(fusionscalartype(I), 0, length(src)) else - # TODO: this is a bit of a hack to fix the traces for now - src = FusionTreeBlock([(f₁, f₂)]) - dst, U = repartition(src, N) - for ((f₁′, f₂′), coeff′) in zip(fusiontrees(dst), U) - for (f₁′′, coeff′′) in planar_trace(f₁′, (q₁′, q₂′)) - src′ = FusionTreeBlock([(f₁′′, f₂′)]) - dst′, U′ = transpose(src′, (p₁′, p₂′)) - for (f12′′′, coeff′′′) in zip(fusiontrees(dst′), U′) - coeff = coeff′ * coeff′′ * coeff′′′ - iszero(coeff) || (newtrees[f12′′′] = get(newtrees, f12′′′, zero(coeff)) + coeff) - end - end - end + dst, U₃ = transpose(traced, p′) + return dst => U₃ * U₂ * U₁ end - return newtrees +end + +function _planar_trace_splitting((f, f₀)::FusionTreePair, q′, r) + f′, coeff = only(planar_trace(f, q′)) + return (f′, f₀) => coeff +end +function _planar_trace_splitting(src::FusionTreeBlock, q′, r) + I = sectortype(src) + f₀, = first(fusiontrees(src)) + uncoupled = TupleTools.getindices(f₀.uncoupled, r) + isdual = TupleTools.getindices(f₀.isdual, r) + dst = FusionTreeBlock{I}((uncoupled, ()), (isdual, ())) + U = zeros(fusionscalartype(I), length(dst), length(src)) + positions = Dict(f => i for (i, (f, _)) in enumerate(fusiontrees(dst))) + for (j, (f, _)) in enumerate(fusiontrees(src)), (f′, coeff) in planar_trace(f, q′) + U[positions[f′], j] += coeff + end + return dst => U end """ diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index ea8334462..4ea3ab915 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -238,9 +238,12 @@ function Base.hash(f::FusionTree{I}, h::UInt) where {I} end return h end -# Hashing the first tree pair covers the block's defining data `(uncoupled, isdual)`, while -# also disambiguating the ad-hoc single-pair blocks constructed in `planar_trace` -Base.hash(b::FusionTreeBlock, h::UInt) = isempty(b) ? h : hash(first(fusiontrees(b)), h) +# a block is defined by the uncoupled sectors and duality flags shared by all its tree pairs +function Base.hash(b::FusionTreeBlock, h::UInt) + isempty(b) && return h + f₁, f₂ = first(fusiontrees(b)) + return hash(f₁.uncoupled, hash(f₁.isdual, hash(f₂.uncoupled, hash(f₂.isdual, h)))) +end function Base.:(==)(f₁::FusionTree{I, N}, f₂::FusionTree{I, N}) where {I <: Sector, N} f₁.coupled == f₂.coupled || return false diff --git a/src/planar/planaroperations.jl b/src/planar/planaroperations.jl index 7c24e1399..5cae4ea8f 100644 --- a/src/planar/planaroperations.jl +++ b/src/planar/planaroperations.jl @@ -92,24 +92,79 @@ function planartrace!( end @timeit_debug GLOBAL_TIMER "planartrace!" begin - if iszero(β) - fill!(C, β) - elseif !isone(β) - rmul!(C, β) + transformer = treetracer(storagetype(C), space(C), space(A), (p₁, p₂), (q₁, q₂)) + dst, src = _transform_subblocks(C, A, transformer) + ntasks = get_num_transformer_threads(A) + trace_transform_kernel!(dst, src, (p₁, p₂), (q₁, q₂), transformer, α, β, backend, allocator, ntasks) + end + return C +end + +# kernel for computing `tdst = β * tdst + α * tr(tsrc)`, with the planar trace encoded in `transformer` +function trace_transform_kernel!( + dst::TransformSubblocks, src::TransformSubblocks, p::Index2Tuple, q::Index2Tuple, + transformer::TraceTransformer, α, β, backend, allocator, ntasks::Int + ) + bufsize = buffersize(transformer) + if bufsize == 0 # every block has a single destination subblock + taskforeach(transformer.data, ntasks) do (U, inds_dst, inds_src) + _trace_transform_block!(dst, src, p, q, U, inds_dst, inds_src, nothing, α, β, backend, allocator) + end + else + # see `add_transform_kernel!` for the workspace strategy + 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 + _trace_transform_block!(dst, src, p, q, U, inds_dst, inds_src, buffer, α, β, backend, allocator) + end + foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) + TO.allocator_reset!(allocator, cp) + end + return nothing +end + +function _trace_transform_block!( + dst::TransformSubblocks, src::TransformSubblocks, p::Index2Tuple, q::Index2Tuple, + U, inds_dst, inds_src, buffer, α, β, backend, allocator + ) + if size(U, 1) == 1 # single destination: trace every source directly into it + idst = only(inds_dst) + if isempty(inds_src) + @inbounds iszero(β) ? zerovector!(dst[idst]) : scale!(dst[idst], β) + return nothing end - β′ = One() - for (f₁, f₂) in fusiontrees(A) - for ((f₁′, f₂′), coeff) in planar_trace((f₁, f₂), (p₁, p₂), (q₁, q₂)) - @timeit_debug GLOBAL_TIMER "dense: trace" TO.tensortrace!( - C[f₁′, f₂′], - A[f₁, f₂], (p₁, p₂), (q₁, q₂), false, - α * coeff, β′, - backend, allocator - ) + @timeit_debug GLOBAL_TIMER "dense: trace" @inbounds begin + TO.tensortrace!(dst[idst], src[first(inds_src)], p, q, false, α * U[1, 1], β, backend, allocator) + for i in 2:length(inds_src) + TO.tensortrace!(dst[idst], src[inds_src[i]], p, q, false, α * U[1, i], One(), backend, allocator) end end + else # trace → recoupling matmul → unpack + rows, cols = size(U) + sz = size(@inbounds(dst[first(inds_dst)])) + ptriv = (ntuple(identity, length(sz)), ()) + blocksize = prod(sz) + buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) + buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) + + # the permutation is applied while tracing, such that unpacking is a plain addition + @timeit_debug GLOBAL_TIMER "dense: trace" @inbounds for (i, isrc) in enumerate(inds_src) + TO.tensortrace!( + sreshape(view(buffer_src, :, i), sz), src[isrc], p, q, false, + One(), Zero(), backend, allocator + ) + end + @timeit_debug GLOBAL_TIMER "dense: recouple mul!" _recouple!(buffer, buffer_dst, buffer_src, U) + @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (j, idst) in enumerate(inds_dst) + TO.tensoradd!( + dst[idst], sreshape(view(buffer_dst, :, j), sz), ptriv, false, α, β, backend, allocator + ) + end end - return C + return nothing end # insert default backend diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index b274c6f8c..73d910775 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,4 +202,104 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end + +""" + TraceTransformer{T, N₁, N₂} + +Planar trace of a tensor with `N₂` indices into a tensor with `N₁` indices, stored as +`(U, inds_dst, inds_src)` where `U[j, i]` is the coefficient with which the trace of source +subblock `inds_src[i]` contributes to destination subblock `inds_dst[j]`. The data is grouped by +destination, such that every destination subblock appears in exactly one entry, where those +without contributions have no source subblocks. 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. +""" +struct TraceTransformer{T, N₁, N₂} + data::Vector{GenericTransformerData{T}} + structure_dst::Vector{StridedStructure{N₁}} + structure_src::Vector{StridedStructure{N₂}} +end + +function TraceTransformer(::Type{A}, Vdst::HomSpace, Vsrc::HomSpace, p::Index2Tuple, q::Index2Tuple) where {A} + t₀ = Base.time() + src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) + + # destination positions => (columns of the recoupling matrix, source positions) + groups = Dict{Vector{Int}, Tuple{Vector{Vector{T}}, Vector{Int}}}() + function addcolumn!(inds_dst, u, isrc) + any(!iszero, u) || return nothing + us, inds_src = get!(() -> (Vector{T}[], Int[]), groups, inds_dst) + push!(us, u) + push!(inds_src, isrc) + return nothing + end + @timeit_debug GLOBAL_TIMER "symmetry: trace coefficients" if FusionStyle(sectortype(Vdst)) isa UniqueFusion + for (isrc, f) in enumerate(src_trees) + trace_sectors_match(f, q) || continue + f′, coeff = planar_trace(f, p, q) + _, (_, idst) = gettoken(dst_trees, f′) + addcolumn!([idst], [convert(T, coeff)], isrc) + end + else + for fs_src in fusionblocks(Vsrc) + trace_sectors_match(fs_src, q) || continue + fs_dst, U = planar_trace(fs_src, p, q) + inds_dst = map(fusiontrees(fs_dst)) do f + _, (_, idst) = gettoken(dst_trees, f) + return idst + end + for (i, f) in enumerate(fusiontrees(fs_src)) + _, (_, isrc) = gettoken(src_trees, f) + addcolumn!(inds_dst, convert(Vector{T}, U[:, i]), isrc) + end + end + end + data = GenericTransformerData{T}[(reduce(hcat, us), inds_dst, inds_src) for (inds_dst, (us, inds_src)) in groups] + # destination subblocks without contributions are only scaled + covered = falses(length(dst_trees)) + foreach(inds_dst -> covered[inds_dst] .= true, keys(groups)) + for idst in findall(!, covered) + push!(data, (zeros(T, 1, 0), [idst], Int[])) + end + + structure_dst = degeneracystructure(Vdst).subblockstructure + structure_src = degeneracystructure(Vsrc).subblockstructure + # sort by (approximate) weight to facilitate multi-threading strategies + weights = map(data) do (U, inds_dst, _) + return length(U) * prod(structure_dst[first(inds_dst)][1]) + end + Base.permute!(data, sortperm(weights; rev = true)) + transformer = TraceTransformer(data, structure_dst, structure_src) + + Δt = Base.time() - t₀ + @debug("TraceTransformer for $Vsrc to $Vdst via $p and $q", nblocks = length(data), Δt) + + return transformer +end + +""" + buffersize(transformer::TraceTransformer) -> Int + +Compute the workspace size required to trace, recouple and unpack the largest block with +multiple destination subblocks, i.e. `prod(sz_dst) * (rows + cols)` where `(rows, cols) = size(U)`. +""" +function buffersize(transformer::TraceTransformer) + structure_dst = transformer.structure_dst + return maximum(transformer.data; init = 0) do (U, inds_dst, _) + return size(U, 1) == 1 ? 0 : prod(structure_dst[first(inds_dst)][1]) * sum(size(U)) + end +end + +function tracetransformertype(::Type{A}, Vdst, Vsrc) where {A} + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) + return TraceTransformer{T, numind(Vdst), numind(Vsrc)} +end + +@cached function treetracer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple + )::tracetransformertype(A, Vdst, Vsrc) where {TA} + return TraceTransformer(A, Vdst, Vsrc, p, q) +end + # default cachestyle is GlobalLRUCache diff --git a/test/symmetries/doubletree.jl b/test/symmetries/doubletree.jl index 38d0aaea0..0af5db19c 100644 --- a/test/symmetries/doubletree.jl +++ b/test/symmetries/doubletree.jl @@ -249,28 +249,24 @@ using TensorKitSectors end @testset "Double fusion tree: planar trace" begin + p1 = ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,)) + p2, q2 = ((2:N...,), (1, ((2N):-1:(N + 3))...)), ((N + 1,), (N + 2,)) if FusionStyle(I) isa UniqueFusion f1, f1 = src - dst, U = transpose((f1, f1), ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,))) - d1 = zip((dst,), (U,)) + dst, U = transpose((f1, f1), p1) + @test TK.trace_sectors_match(dst, q2) + dst2, U2 = @constinferred TK.planar_trace(dst, p2, q2) + d2 = Dict(dst2 => U2 * U) else f1, f1 = first(fusiontrees(src)) src′ = FusionTreeBlock{I}((f1.uncoupled, f1.uncoupled), (f1.isdual, f1.isdual)) - dst, U = transpose(src′, ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,))) - d1 = zip(fusiontrees(dst), U[:, 1]) + dst, U = transpose(src′, p1) + @test TK.trace_sectors_match(dst, q2) + dst2, U2 = TK.planar_trace(dst, p2, q2) + d2 = Dict(zip(fusiontrees(dst2), U2 * U[:, 1])) end f1front, = TK.split(f1, N - 1) - T = sectorscalartype(I) - d2 = Dict{typeof((f1front, f1front)), T}() - for ((f1′, f2′), coeff′) in d1 - for ((f1′′, f2′′), coeff′′) in TK.planar_trace( - (f1′, f2′), ((2:N...,), (1, ((2N):-1:(N + 3))...)), ((N + 1,), (N + 2,)) - ) - coeff = coeff′ * coeff′′ - d2[(f1′′, f2′′)] = get(d2, (f1′′, f2′′), zero(coeff)) + coeff - end - end for ((f1_, f2_), coeff) in d2 if (f1_, f2_) == (f1front, f1front) @test coeff ≈ dim(f1.coupled) / dim(f1front.coupled) From dc3cd36eadc777ffb3f97323bcb32fa071a65da6 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 12:06:16 -0400 Subject: [PATCH 2/4] Add a GPU kernel for planar traces GPU tensors pack the `TraceTransformer` into `isbits` descriptors, cached per device storagetype, and trace with one thread per destination element, fusing the sum over the traced entries and the recoupling into a single kernel, analogous to the generic tree transformers. Other tensor types, such as adjoints, are first converted to `TensorMap`s: tracing subblock views with cuTENSOR fails for data offsets that are not aligned (e.g. `Float64` data at odd offsets). `_linear_index` is now written as an explicit loop, since `sum` of an empty tuple, as occurs for a full trace, does not compile for GPUs. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 2 +- ext/TensorKitGPUArraysExt.jl | 167 ++++++++++++++++++++++++++++++++++- test/cuda/planar.jl | 18 ++++ 3 files changed, 182 insertions(+), 5 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 5ea2651e7..02f0d5035 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -30,7 +30,7 @@ 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) -- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s (internal) +- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s, and GPU tensors use a dedicated kernel (internal) ### Deprecated diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 45f6df640..183b31948 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -12,8 +12,9 @@ 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 -import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel! +using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, TraceTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple +using TensorKit: SubblockIterator +import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel!, trace_transform_kernel! function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:AnyGPUMatrix{T}, <:StridedViews.StridedView{T, 4, <:AnyGPUArray{T}}}} # COV_EXCL_START @@ -271,6 +272,75 @@ function DeviceGenericTreeTransformer( ) end +""" + TraceTransformerBlock{N} + +Descriptor for a destination block of a `TraceTransformer`, indexing into the flat +`coeffs`/`structs_dst`/`structs_src` vectors of a `DeviceTraceTransformer`. +""" +struct TraceTransformerBlock{N} + sz::NTuple{N, Int} + dense_strides::NTuple{N, Int} + rows::Int + cols::Int + u_offset::Int + dst_offset::Int + src_offset::Int +end + +# (permuted strides of the remaining legs, sizes, dense strides and summed strides of the traced +# pairs, offset): the traced entries of a source subblock are found through a single stride per pair +const TraceSourceStructure{N, M} = Tuple{NTuple{N, Int}, NTuple{M, Int}, NTuple{M, Int}, NTuple{M, Int}, Int} + +struct DeviceTraceTransformer{VB <: AbstractVector{<:TraceTransformerBlock}, VO <: AbstractVector{Int}, VC <: AbstractVector{<:Number}, VD <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}, VS <: AbstractVector{<:TraceSourceStructure}} + blocks::VB + work_offsets::VO + nwork::Int + coeffs::VC # every `U`, concatenated in column-major order + structs_dst::VD + structs_src::VS +end + +function DeviceTraceTransformer(transformer::TraceTransformer{T, N₁, N₂}, p, (q₁, q₂)) where {T, N₁, N₂} + lp = linearize(p) + M = length(q₁) + blocks = TraceTransformerBlock{N₁}[] + coeffs = T[] + structs_dst = TreeStructure{N₁}[] + structs_src = TraceSourceStructure{N₁, M}[] + + (; structure_dst, structure_src) = transformer + for (U, inds_dst, inds_src) in transformer.data + # all destination trees in a block share the same subblock size + size_dst = first(structure_dst[first(inds_dst)]) + push!( + blocks, TraceTransformerBlock{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 + size_src, strides_src, offset_src = structure_src[isrc] + size_traced = TupleTools.getindices(size_src, q₁) + strides_traced = TupleTools.getindices(strides_src, q₁) .+ TupleTools.getindices(strides_src, q₂) + push!( + structs_src, ( + TupleTools.getindices(strides_src, lp), size_traced, + _dense_strides(size_traced), strides_traced, offset_src, + ) + ) + end + end + + work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) + return DeviceTraceTransformer(blocks, work_offsets, nwork, coeffs, structs_dst, structs_src) +end + # Cache the device transformer directly, constructing the cached host version only on a miss. @cached function TensorKit.treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, @@ -287,6 +357,13 @@ end return device_transformer(A, host, linearize(p)) end +@cached function TensorKit.treetracer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple + ) where {TA <: AnyGPUVector} + host = TensorKit.treetracer(Vector{scalartype(A)}, Vdst, Vsrc, p, q) + return device_transformer(A, host, p, q) +end + const DEFAULT_DEVICECACHE_SIZE = 1000 function __init__() @@ -295,12 +372,14 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_CACHE, ) ) # device-caches should probably have smaller sizes: resize!(GLOBAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) return nothing end @@ -332,6 +411,18 @@ function device_transformer(A, transformer::GenericTreeTransformer, p) end end +function device_transformer(A, transformer::TraceTransformer, p, q) + packed = DeviceTraceTransformer(transformer, p, q) + return GPUArrays.@uncached begin + proto = A(undef, 0) + upload(x) = copy!(similar(proto, eltype(x), size(x)), x) + DeviceTraceTransformer( + upload(packed.blocks), upload(packed.work_offsets), packed.nwork, + upload(packed.coeffs), upload(packed.structs_dst), upload(packed.structs_src) + ) + end +end + # COV_EXCL_START # kernels are not reachable by coverage @@ -360,7 +451,12 @@ end # finds the overall linear index in the output and input arrays corresponding to the **sublock** # coordinates currently being worked on @inline function _linear_index(coords::NTuple{N, Int}, strides::NTuple{N, Int}, offset) where {N} - return offset + sum(ntuple(n -> coords[n] * strides[n], Val(N))) + 1 + # an explicit loop, since `sum` of an empty tuple (e.g. a full trace) does not compile for GPUs + i = offset + 1 + for n in 1:N + i += coords[n] * strides[n] + end + return i end # One thread per destination element in `data_dst`. `op` is `identity` or `conj`, and is @@ -426,6 +522,44 @@ end end end +# COV_EXCL_START +# One thread per destination element, as for the generic transformers, where every source +# additionally contributes the sum over its traced entries. +@kernel function batched_trace_kernel!( + data_dst, data_src, blocks, work_offsets, coeffs, structs_dst, structs_src, α, β, nwork + ) + # w = @index(Global, Linear) - 1 does not work due to JLArrays bug + w = @index(Global, Linear) + w -= 1 + if w < nwork + b = _searchblock(work_offsets, w) + blk = @inbounds blocks[b] + 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.dense_strides) + + st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1] + i_dst = _linear_index(coords, st_dst, offs_dst) + + acc = zero(promote_type(eltype(data_src), eltype(coeffs))) + @inbounds for j in 1:blk.cols + coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] + iszero(coeff) && continue + pst_src, sz_tr, dense_tr, st_tr, offs_src = structs_src[blk.src_offset + j] + i_src = _linear_index(coords, pst_src, offs_src) + tr = zero(eltype(data_src)) + for t in 0:(prod(sz_tr) - 1) + tr += data_src[i_src + _linear_index(_coordinates(t, sz_tr, dense_tr), st_tr, -1)] + end + acc += coeff * tr + end + # destinations without contributions are overwritten for `β == 0` + @inbounds data_dst[i_dst] = iszero(β) ? α * acc : α * acc + β * data_dst[i_dst] + end +end +# COV_EXCL_STOP + function _launch_unique!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} nwork = transformer.nwork nwork == 0 && return nothing @@ -450,7 +584,7 @@ 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} + transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer, DeviceTraceTransformer} ) = (tdst.data, tsrc.data) function TensorKit.add_transform_kernel!( @@ -476,4 +610,29 @@ function TensorKit.add_transform_kernel!( return nothing end +function TensorKit.trace_transform_kernel!( + dst::AnyGPUVector, src::AnyGPUVector, p::Index2Tuple, q::Index2Tuple, + device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int + ) + device.nwork == 0 && return nothing + batched_trace_kernel!(get_backend(dst))( + dst, src, device.blocks, device.work_offsets, device.coeffs, + device.structs_dst, device.structs_src, α, β, device.nwork; ndrange = device.nwork + ) + return nothing +end + +# other tensor types are first converted to `TensorMap`s: tracing subblock by subblock with +# `tensortrace!` is not an option, since cuTENSOR requires aligned data for the subblock views +function TensorKit.trace_transform_kernel!( + dst::SubblockIterator, src::SubblockIterator, p::Index2Tuple, q::Index2Tuple, + device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int + ) + _tensormap(t) = t isa TensorMap ? t : copy!(similar(t, scalartype(t), space(t)), t) + tdst, tsrc = _tensormap(dst.t), _tensormap(src.t) + trace_transform_kernel!(tdst.data, tsrc.data, p, q, device, α, β, backend, allocator, ntasks) + tdst === dst.t || copy!(dst.t, tdst) + return nothing +end + end diff --git a/test/cuda/planar.jl b/test/cuda/planar.jl index 2d992023b..8cec768db 100644 --- a/test/cuda/planar.jl +++ b/test/cuda/planar.jl @@ -74,6 +74,24 @@ for V in spacelist end end +for V in spacelist + I = sectortype(first(V)) + Istr = type_repr(I) + @timedtestset "planartrace + CUDA with symmetry: $Istr" begin + # (space, p, q) for a trace of neighbouring legs and a nested trace + cases = ( + (V[1] ⊗ V[2] ⊗ V[3] ← V[4] ⊗ V[3], ((1, 2), (4,)), ((3,), (5,))), + (V[1] ⊗ V[2] ⊗ V[3] ⊗ V[4] ← V[3] ⊗ V[4], ((1,), (2,)), ((3, 4), (5, 6))), + ) + for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) + A = randn(ComplexF64, W) + C = randn(ComplexF64, TensorKit.select(W, p)) + Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) + @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) + end + end +end + @testset "planar methods" verbose = true begin @testset "planaradd" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^6 ⊗ ℂ^5 ⊗ ℂ^4) From f4fd43ccdcc31b167956cb8071e2b278499e9a45 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 30 Oct 2026 12:05:19 -0400 Subject: [PATCH 3/4] Split the GPU planar trace kernel into partial traces and recoupling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first stage sums chunks of at most `max(256, √ntraced)` traced entries per thread, so full reductions and traces with few remaining entries are still spread over many threads. The second stage recouples these partials with one thread per destination element, such that traced entries are summed once per source rather than once per destination tree. Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitGPUArraysExt.jl | 133 +++++++++++++++++++++++++++-------- test/cuda/planar.jl | 18 +++++ 2 files changed, 121 insertions(+), 30 deletions(-) diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 183b31948..9974ddf35 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -9,6 +9,7 @@ using Strided: StridedViews using MatrixAlgebraKit, Adapt using TensorKit using TensorKit.TensorOperations: linearize, DefaultAllocator +const TO = TensorKit.TensorOperations using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check @@ -272,11 +273,17 @@ function DeviceGenericTreeTransformer( ) end +# Planar traces run in two stages. The first computes partial traces of every source subblock: +# each thread sums a chunk of at most `chunksize` traced entries for a single remaining entry, +# such that large traces (e.g. full reductions) are still spread over many threads. The second +# recouples these partials with one thread per destination element, as for the generic transformers, +# such that the traced entries of a source are summed only once, rather than once per destination tree. + """ TraceTransformerBlock{N} Descriptor for a destination block of a `TraceTransformer`, indexing into the flat -`coeffs`/`structs_dst`/`structs_src` vectors of a `DeviceTraceTransformer`. +`coeffs`/`structs_dst`/`sources` vectors of a `DeviceTraceTransformer`. """ struct TraceTransformerBlock{N} sz::NTuple{N, Int} @@ -288,35 +295,60 @@ struct TraceTransformerBlock{N} src_offset::Int end -# (permuted strides of the remaining legs, sizes, dense strides and summed strides of the traced -# pairs, offset): the traced entries of a source subblock are found through a single stride per pair -const TraceSourceStructure{N, M} = Tuple{NTuple{N, Int}, NTuple{M, Int}, NTuple{M, Int}, NTuple{M, Int}, Int} +""" + TraceSourceBlock{N, M} -struct DeviceTraceTransformer{VB <: AbstractVector{<:TraceTransformerBlock}, VO <: AbstractVector{Int}, VC <: AbstractVector{<:Number}, VD <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}, VS <: AbstractVector{<:TraceSourceStructure}} +Descriptor for the partial traces of a source subblock of a `TraceTransformer`, with `N` remaining +legs (permuted to the destination order) and `M` traced pairs, whose entries are found through a +single stride per pair. The partials are stored chunk-major, as `nchunks` dense arrays of shape `sz`. +""" +struct TraceSourceBlock{N, M} + sz::NTuple{N, Int} + dense_strides::NTuple{N, Int} + permuted_strides::NTuple{N, Int} + offset::Int + sz_traced::NTuple{M, Int} + dense_strides_traced::NTuple{M, Int} + strides_traced::NTuple{M, Int} # strides[q₁] .+ strides[q₂] + nremaining::Int + ntraced::Int + chunksize::Int + nchunks::Int +end + +struct DeviceTraceTransformer{VB <: AbstractVector{<:TraceTransformerBlock}, VO <: AbstractVector{Int}, VC <: AbstractVector{<:Number}, VD <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}, VS <: AbstractVector{<:TraceSourceBlock}} blocks::VB work_offsets::VO nwork::Int coeffs::VC # every `U`, concatenated in column-major order structs_dst::VD - structs_src::VS + sources::VS + partial_offsets::VO # also the work offsets of the first stage + npartials::Int end +# traces with fewer entries are not split, larger ones are split into at most `√ntraced` chunks, +# such that no thread in either stage sums more than `max(256, √ntraced)` entries +const MIN_TRACE_CHUNKSIZE = 256 +_trace_chunksize(ntraced::Int) = max(MIN_TRACE_CHUNKSIZE, isqrt(ntraced)) + function DeviceTraceTransformer(transformer::TraceTransformer{T, N₁, N₂}, p, (q₁, q₂)) where {T, N₁, N₂} lp = linearize(p) M = length(q₁) blocks = TraceTransformerBlock{N₁}[] coeffs = T[] structs_dst = TreeStructure{N₁}[] - structs_src = TraceSourceStructure{N₁, M}[] + sources = TraceSourceBlock{N₁, M}[] (; structure_dst, structure_src) = transformer for (U, inds_dst, inds_src) in transformer.data # all destination trees in a block share the same subblock size size_dst = first(structure_dst[first(inds_dst)]) + dense_strides_dst = _dense_strides(size_dst) push!( blocks, TraceTransformerBlock{N₁}( - size_dst, _dense_strides(size_dst), size(U, 1), size(U, 2), - length(coeffs), length(structs_dst), length(structs_src) + size_dst, dense_strides_dst, size(U, 1), size(U, 2), + length(coeffs), length(structs_dst), length(sources) ) ) append!(coeffs, U) @@ -328,17 +360,23 @@ function DeviceTraceTransformer(transformer::TraceTransformer{T, N₁, N₂}, p, size_src, strides_src, offset_src = structure_src[isrc] size_traced = TupleTools.getindices(size_src, q₁) strides_traced = TupleTools.getindices(strides_src, q₁) .+ TupleTools.getindices(strides_src, q₂) + ntraced = prod(size_traced) + chunksize = _trace_chunksize(ntraced) push!( - structs_src, ( - TupleTools.getindices(strides_src, lp), size_traced, - _dense_strides(size_traced), strides_traced, offset_src, + sources, TraceSourceBlock{N₁, M}( + size_dst, dense_strides_dst, TupleTools.getindices(strides_src, lp), offset_src, + size_traced, _dense_strides(size_traced), strides_traced, + prod(size_dst), ntraced, chunksize, cld(ntraced, chunksize) ) ) end end work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) - return DeviceTraceTransformer(blocks, work_offsets, nwork, coeffs, structs_dst, structs_src) + partial_offsets, npartials = _work_offsets(src.nchunks * src.nremaining for src in sources) + return DeviceTraceTransformer( + blocks, work_offsets, nwork, coeffs, structs_dst, sources, partial_offsets, npartials + ) end # Cache the device transformer directly, constructing the cached host version only on a miss. @@ -418,7 +456,8 @@ function device_transformer(A, transformer::TraceTransformer, p, q) upload(x) = copy!(similar(proto, eltype(x), size(x)), x) DeviceTraceTransformer( upload(packed.blocks), upload(packed.work_offsets), packed.nwork, - upload(packed.coeffs), upload(packed.structs_dst), upload(packed.structs_src) + upload(packed.coeffs), upload(packed.structs_dst), upload(packed.sources), + upload(packed.partial_offsets), packed.npartials ) end end @@ -523,34 +562,57 @@ end end # COV_EXCL_START -# One thread per destination element, as for the generic transformers, where every source -# additionally contributes the sum over its traced entries. -@kernel function batched_trace_kernel!( - data_dst, data_src, blocks, work_offsets, coeffs, structs_dst, structs_src, α, β, nwork - ) +# First stage: one thread per (chunk, remaining entry) of every source subblock, writing its +# partial trace to `partials[w + 1]`. Consecutive threads handle consecutive remaining entries. +@kernel function partial_trace_kernel!(partials, data_src, sources, partial_offsets, npartials) # w = @index(Global, Linear) - 1 does not work due to JLArrays bug w = @index(Global, Linear) w -= 1 + if w < npartials + j = _searchblock(partial_offsets, w) + src = @inbounds sources[j] + k, e = divrem(w - (@inbounds partial_offsets[j]), src.nremaining) + coords = _coordinates(e, src.sz, src.dense_strides) + i_src = _linear_index(coords, src.permuted_strides, src.offset) + + acc = zero(eltype(partials)) + start = k * src.chunksize + @inbounds for t in start:(min(start + src.chunksize, src.ntraced) - 1) + coords_traced = _coordinates(t, src.sz_traced, src.dense_strides_traced) + acc += data_src[i_src + _linear_index(coords_traced, src.strides_traced, -1)] + end + @inbounds partials[w + 1] = acc + end +end + +# Second stage: one thread per destination element, as for the generic transformers, where +# every source contributes the sum of its partial traces. +@kernel function trace_recouple_kernel!( + data_dst, partials, blocks, work_offsets, coeffs, structs_dst, sources, partial_offsets, + α, β, nwork + ) + w = @index(Global, Linear) + w -= 1 if w < nwork b = _searchblock(work_offsets, w) blk = @inbounds blocks[b] 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.dense_strides) + i, e = divrem(local_w, blocksize) # 0-based destination tree and entry + coords = _coordinates(e, blk.sz, blk.dense_strides) st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1] i_dst = _linear_index(coords, st_dst, offs_dst) - acc = zero(promote_type(eltype(data_src), eltype(coeffs))) + acc = zero(promote_type(eltype(partials), eltype(coeffs))) @inbounds for j in 1:blk.cols coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] iszero(coeff) && continue - pst_src, sz_tr, dense_tr, st_tr, offs_src = structs_src[blk.src_offset + j] - i_src = _linear_index(coords, pst_src, offs_src) - tr = zero(eltype(data_src)) - for t in 0:(prod(sz_tr) - 1) - tr += data_src[i_src + _linear_index(_coordinates(t, sz_tr, dense_tr), st_tr, -1)] + js = blk.src_offset + j + offset = partial_offsets[js] + e + 1 + tr = zero(eltype(partials)) + for k in 0:(sources[js].nchunks - 1) + tr += partials[offset + k * blocksize] end acc += coeff * tr end @@ -615,10 +677,21 @@ function TensorKit.trace_transform_kernel!( device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int ) device.nwork == 0 && return nothing - batched_trace_kernel!(get_backend(dst))( - dst, src, device.blocks, device.work_offsets, device.coeffs, - device.structs_dst, device.structs_src, α, β, device.nwork; ndrange = device.nwork + kernel_backend = get_backend(dst) + cp = TO.allocator_checkpoint!(allocator) + partials = TO.tensoralloc(similarstoragetype(src), device.npartials, Val(true), allocator) + if device.npartials > 0 + partial_trace_kernel!(kernel_backend)( + partials, src, device.sources, device.partial_offsets, device.npartials; + ndrange = device.npartials + ) + end + trace_recouple_kernel!(kernel_backend)( + dst, partials, device.blocks, device.work_offsets, device.coeffs, device.structs_dst, + device.sources, device.partial_offsets, α, β, device.nwork; ndrange = device.nwork ) + TO.tensorfree!(partials, allocator) + TO.allocator_reset!(allocator, cp) return nothing end diff --git a/test/cuda/planar.jl b/test/cuda/planar.jl index 8cec768db..c2d46d48c 100644 --- a/test/cuda/planar.jl +++ b/test/cuda/planar.jl @@ -92,6 +92,24 @@ for V in spacelist end end +@timedtestset "planartrace + CUDA with large traces: $I" for (I, V, Vr) in ( + (FermionParity, Vect[FermionParity](0 => 40, 1 => 40), Vect[FermionParity](0 => 2, 1 => 3)), + (FibonacciAnyon, Vect[FibonacciAnyon](:I => 30, :τ => 30), Vect[FibonacciAnyon](:I => 2, :τ => 1)), + ) + # (space, p, q) for a full reduction and for a trace with few remaining entries, which are + # both split into several chunks of traced entries + cases = ( + (V ⊗ V ← V ⊗ V, ((), ()), ((1, 2), (3, 4))), + (Vr ⊗ V ⊗ V ← V ⊗ V, ((1,), ()), ((2, 3), (4, 5))), + ) + for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) + A = randn(ComplexF64, W) + C = randn(ComplexF64, TensorKit.select(W, p)) + Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) + @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) + end +end + @testset "planar methods" verbose = true begin @testset "planaradd" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^6 ⊗ ℂ^5 ⊗ ℂ^4) From 388c35e40adca4811b46df5315fcefd48da6c013 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 30 Oct 2026 15:56:12 -0400 Subject: [PATCH 4/4] Move the large GPU trace tests into the planar methods testset Co-Authored-By: Claude Opus 5.5 --- test/cuda/planar.jl | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/test/cuda/planar.jl b/test/cuda/planar.jl index c2d46d48c..e92faef8e 100644 --- a/test/cuda/planar.jl +++ b/test/cuda/planar.jl @@ -92,24 +92,6 @@ for V in spacelist end end -@timedtestset "planartrace + CUDA with large traces: $I" for (I, V, Vr) in ( - (FermionParity, Vect[FermionParity](0 => 40, 1 => 40), Vect[FermionParity](0 => 2, 1 => 3)), - (FibonacciAnyon, Vect[FibonacciAnyon](:I => 30, :τ => 30), Vect[FibonacciAnyon](:I => 2, :τ => 1)), - ) - # (space, p, q) for a full reduction and for a trace with few remaining entries, which are - # both split into several chunks of traced entries - cases = ( - (V ⊗ V ← V ⊗ V, ((), ()), ((1, 2), (3, 4))), - (Vr ⊗ V ⊗ V ← V ⊗ V, ((1,), ()), ((2, 3), (4, 5))), - ) - for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) - A = randn(ComplexF64, W) - C = randn(ComplexF64, TensorKit.select(W, p)) - Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) - @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) - end -end - @testset "planar methods" verbose = true begin @testset "planaradd" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^6 ⊗ ℂ^5 ⊗ ℂ^4) @@ -134,6 +116,24 @@ end planartrace!(C′, A′, p, q, true, true) end + @testset "planartrace with large traces: $I" for (I, V, Vr) in ( + (FermionParity, Vect[FermionParity](0 => 40, 1 => 40), Vect[FermionParity](0 => 2, 1 => 3)), + (FibonacciAnyon, Vect[FibonacciAnyon](:I => 30, :τ => 30), Vect[FibonacciAnyon](:I => 2, :τ => 1)), + ) + # (space, p, q) for a full reduction and for a trace with few remaining entries, which are + # both split into several chunks of traced entries + cases = ( + (V ⊗ V ← V ⊗ V, ((), ()), ((1, 2), (3, 4))), + (Vr ⊗ V ⊗ V ← V ⊗ V, ((1,), ()), ((2, 3), (4, 5))), + ) + for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) + A = randn(ComplexF64, W) + C = randn(ComplexF64, TensorKit.select(W, p)) + Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) + @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) + end + end + @testset "planarcontract" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^2 ⊗ ℂ^5 ⊗ ℂ^4) B = CUDA.randn(ℂ^2 ⊗ ℂ^4 ← ℂ^4 ⊗ ℂ^3)