From a1df9bbc72bc051d4944a2efaa6ff05127e31081 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 11:42:07 -0400 Subject: [PATCH 1/9] 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 | 2 + 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, 259 insertions(+), 75 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 3fc79606b..b6fbd8914 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -24,6 +24,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed +- `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 ### Removed 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 96dfa3cb7bf9fd1c590d841a25f9af9b3ac9b1f8 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 12:06:16 -0400 Subject: [PATCH 2/9] 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 b6fbd8914..c4f157b1e 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -24,7 +24,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed -- `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 0408980dd..b3adb2a67 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -13,8 +13,9 @@ using TensorKit.VectorInterface: One 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 @@ -272,6 +273,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, @@ -288,6 +358,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__() @@ -298,6 +375,7 @@ function __init__() Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREEBRAIDER_CACHE") => GLOBAL_DIAGONAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE") => GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_CACHE, ) ) @@ -306,6 +384,7 @@ function __init__() resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) resize!(GLOBAL_DIAGONAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) resize!(GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) return nothing end @@ -339,6 +418,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 @@ -367,7 +458,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 @@ -433,6 +529,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 @@ -457,7 +591,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) # Diagonal sources @@ -568,4 +702,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 b70d162d60c9ad9243e6cbafbbb0e657b71a3189 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 30 Oct 2026 12:05:19 -0400 Subject: [PATCH 3/9] 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 b3adb2a67..cf5dbb967 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -10,6 +10,7 @@ using MatrixAlgebraKit, Adapt using TensorKit using TensorKit.TensorOperations: linearize, DefaultAllocator using TensorKit.VectorInterface: One +const TO = TensorKit.TensorOperations using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check @@ -273,11 +274,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} @@ -289,35 +296,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) @@ -329,17 +361,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. @@ -425,7 +463,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 @@ -530,34 +569,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 @@ -707,10 +769,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 18adc659bf199549ee2869c866fa0401c263b613 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 30 Oct 2026 15:56:12 -0400 Subject: [PATCH 4/9] 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) From 607fe0c6205a8a94edfaf1adbfafbe27aa9fcd84 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 2 Nov 2026 16:26:27 -0500 Subject: [PATCH 5/9] Add Enzyme rules for `planartrace!` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enzyme cannot differentiate through the threaded `taskforeach` in the new trace kernel, so `planartrace!` now has forward and reverse rules, as for the index manipulations. The pullback `planartrace_pullback_ΔA!` reuses the cached `TraceTransformer`: it adds `ΔC` to the traced diagonal of every source subblock with coefficients `conj(α * U[i, j]) * dim(cᵢ) / dim(cⱼ)`, going through the subblocks such that adjoint tensors are supported, and falls back to `trace_permute_pullback_ΔA!` for bosonic sectors. `treetracer` is marked inactive. The Enzyme `planartrace!` test skipped every draw without a remaining leg or traced pair, which in practice skipped (almost) all of them. It now always draws at least one of each, retries draws whose multifusion spaces have mismatched units, and also compares the pullback of adjoint tensors. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 2 + ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl | 1 + ext/TensorKitEnzymeExt/planaroperations.jl | 93 ++++++++++++++++++++ ext/TensorKitEnzymeExt/utility.jl | 1 + src/pullbacks/tensoroperations.jl | 56 ++++++++++++ test/enzyme/trace.jl | 54 ++++++++---- 6 files changed, 192 insertions(+), 15 deletions(-) create mode 100644 ext/TensorKitEnzymeExt/planaroperations.jl diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index c4f157b1e..bc2dc34ad 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -22,6 +22,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Added +- Enzyme rules for `planartrace!`, with a pullback that also supports adjoint tensors + ### Changed - `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) diff --git a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl index 98bc70c40..4de124b5b 100644 --- a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl +++ b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl @@ -30,5 +30,6 @@ include("utility.jl") include("linalg.jl") include("indexmanipulations.jl") include("tensoroperations.jl") +include("planaroperations.jl") end diff --git a/ext/TensorKitEnzymeExt/planaroperations.jl b/ext/TensorKitEnzymeExt/planaroperations.jl new file mode 100644 index 000000000..9ff0a00bd --- /dev/null +++ b/ext/TensorKitEnzymeExt/planaroperations.jl @@ -0,0 +1,93 @@ +# planartrace! +# ------------ +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof(TK.planartrace!)}, + ::Type{RT}, + C::Annotation{<:AbstractTensorMap}, + A::Annotation{<:AbstractTensorMap}, + p::Const{<:Index2Tuple}, + q::Const{<:Index2Tuple}, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + backend::Const, + allocator::Const + ) where {RT} + C_cache = !isa(β, Const) ? copy(C.val) : nothing + # if we need to compute Δα, store the trace of A instead of recomputing it in the pullback + At = if !isa(α, Const) + At = TK.planartrace!(zerovector(C.val), A.val, p.val, q.val, One(), Zero(), backend.val, allocator.val) + add!(C.val, At, α.val, β.val) + At + else + TK.planartrace!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val, allocator.val) + nothing + end + cache = (C_cache, At) + primal = EnzymeRules.needs_primal(config) ? C.val : nothing + shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing + return EnzymeRules.AugmentedReturn(primal, shadow, cache) +end + +function EnzymeRules.reverse( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof(TK.planartrace!)}, + ::Type{RT}, + cache, + C::Annotation{<:AbstractTensorMap}, + A::Annotation{<:AbstractTensorMap}, + p::Const{<:Index2Tuple}, + q::Const{<:Index2Tuple}, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + backend::Const, + allocator::Const + ) where {RT} + C_cache, At = cache + Cval = something(C_cache, C.val) + # the pullback of A is linear, so only the spaces of A are needed + if !is_inactive(config, A) && !is_inactive(config, C) + TK.planartrace_pullback_ΔA!(A.dval, C.dval, A.val, p.val, q.val, α.val, backend.val, allocator.val) + end + Δα = pullback_dα(α, C, At) + Δβ = pullback_dβ(β, C, Cval) + !is_inactive(config, C) && pullback_dC!(C.dval, β.val) + return nothing, nothing, nothing, nothing, Δα, Δβ, nothing, nothing +end + +function EnzymeRules.forward( + config::EnzymeRules.FwdConfigWidth{1}, + func::Const{typeof(TK.planartrace!)}, + ::Type{RT}, + C::Annotation{<:AbstractTensorMap}, + A::Annotation{<:AbstractTensorMap}, + p::Annotation{<:Index2Tuple}, + q::Annotation{<:Index2Tuple}, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + backend::Const, + allocator::Const + ) where {RT} + # ΔC′ = ΔC*β + C*Δβ + tr(A)*Δα + tr(ΔA)*α + if !isa(C, Const) + if isa(β, Const) + scale!(C.dval, β.val) + else + add!(C.dval, C.val, β.dval, β.val) + end + !isa(α, Const) && TK.planartrace!(C.dval, A.val, p.val, q.val, α.dval, One(), backend.val, allocator.val) + !isa(A, Const) && TK.planartrace!(C.dval, A.dval, p.val, q.val, α.val, One(), backend.val, allocator.val) + end + TK.planartrace!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val, allocator.val) + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + return C + elseif EnzymeRules.needs_primal(config) + return C.val + elseif EnzymeRules.needs_shadow(config) && !isa(C, Const) + return C.dval + elseif EnzymeRules.needs_shadow(config) && isa(C, Const) + return Enzyme.make_zero(C.val) + else + return nothing + end +end diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index bb9a21db0..33b38da17 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -150,6 +150,7 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.product_transform), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.treetracer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing diff --git a/src/pullbacks/tensoroperations.jl b/src/pullbacks/tensoroperations.jl index 4e784af5b..5b972a3e3 100644 --- a/src/pullbacks/tensoroperations.jl +++ b/src/pullbacks/tensoroperations.jl @@ -64,3 +64,59 @@ function trace_permute_pullback_ΔA!( ) return nothing end + +# The planar trace adds `α * U[i, j]` times the dense trace of every source subblock `j` to +# destination subblock `i`. Since `inner` weighs every subblock with the dimension of its coupled +# sector, the adjoint adds `ΔC` to the traced diagonal of every source subblock, with coefficients +# `conj(α * U[i, j]) * dim(cᵢ) / dim(cⱼ)`. +function planartrace_pullback_ΔA!( + ΔA, ΔC, A::AbstractTensorMap, p, q, α, backend, allocator + ) + BraidingStyle(sectortype(A)) == Bosonic() && + return trace_permute_pullback_ΔA!(ΔA, ΔC, A, p, q, α, backend) + transformer = treetracer(storagetype(ΔC), space(ΔC), space(A), p, q) + T = promote_type(scalartype(ΔC), scalartype(α), _coefficienttype(transformer)) + if T <: Real || !(scalartype(ΔA) <: Real) + _planartrace_pullback_ΔA!(ΔA, ΔC, transformer, p, q, α, backend, allocator) + else + ΔAc = zerovector(ΔA, T) + _planartrace_pullback_ΔA!(ΔAc, ΔC, transformer, p, q, α, backend, allocator) + add!(ΔA, real(ΔAc)) + end + return nothing +end + +_coefficienttype(::TraceTransformer{T}) where {T} = T + +function _planartrace_pullback_ΔA!( + ΔA, ΔC, transformer::TraceTransformer, (p₁, p₂), (q₁, q₂), α, backend, allocator + ) + lp = (p₁..., p₂...) + dims_dst = [dim(f₁.coupled) for (f₁, _) in fusiontrees(ΔC)] + dims_src = [dim(f₁.coupled) for (f₁, _) in fusiontrees(ΔA)] + dst, src = _transform_subblocks(ΔC, ΔA, transformer) + for (U, inds_dst, inds_src) in transformer.data, (j, isrc) in enumerate(inds_src) + # the traced diagonal of the source subblock, with the remaining legs permuted by `p` + a = StridedView(src[isrc]) + sz, str = size(a), strides(a) + sz_traced = TupleTools.getindices(sz, q₁) + str_traced = TupleTools.getindices(str, q₁) .+ TupleTools.getindices(str, q₂) + diag = StridedView( + parent(a), (TupleTools.getindices(sz, lp)..., sz_traced...), + (TupleTools.getindices(str, lp)..., str_traced...), a.offset, a.op + ) + pdiag = (ntuple(identity, ndims(diag)), ()) + for (i, idst) in enumerate(inds_dst) + iszero(U[i, j]) && continue + # the destination subblock, constant along the traced legs + c = StridedView(dst[idst]) + ΔCi = StridedView( + parent(c), (size(c)..., sz_traced...), (strides(c)..., map(zero, sz_traced)...), + c.offset, c.op + ) + coeff = conj(α * U[i, j]) * (dims_dst[idst] / dims_src[isrc]) + TO.tensoradd!(diag, ΔCi, pdiag, false, coeff, One(), backend, allocator) + end + end + return ΔA +end diff --git a/test/enzyme/trace.jl b/test/enzyme/trace.jl index 82f1f9053..6ccfca2df 100644 --- a/test/enzyme/trace.jl +++ b/test/enzyme/trace.jl @@ -57,9 +57,11 @@ TAs = is_ci ? (Duplicated,) : (Const, Duplicated) end end @timedtestset "planartrace!" begin - for _ in 1:5 - k1 = rand(0:2) - k2 = rand(0:1) + # a random planar trace with at least one remaining leg and one traced pair, or `nothing` + # for multifusion sectors whose random spaces have mismatched units, or empty output + function randplanartrace() + k1 = rand(1:2) + k2 = 1 V1 = map(v -> rand(Bool) ? v' : v, rand(V, k1)) V2 = map(v -> rand(Bool) ? v' : v, rand(V, k2)) V3 = prod(x -> x ⊗ x', V2[1:k2]; init = one(V[1])) @@ -69,19 +71,41 @@ TAs = is_ci ? (Duplicated,) : (Const, Duplicated) (_p, _q) = randcircshift(k′, k1 + 2k2 - k′, k1) p = _repartition(_p, rand(0:k1)) q = (tuple(_q[1:2:end]...), tuple(_q[2:2:end]...)) - if !all(isempty, p) && !all(isempty, q) - α = randn(T) - β = randn(T) - ip = _repartition(invperm(linearize((_p, _q))), k′) - A = randn(T, permute(prod(V1) ⊗ V3 ← V4, ip)) - C = randn!(TensorOperations.tensoralloc_add(T, A, p, false, Val(false))) - for Tα in (Const, Active), Tβ in (Const, Active) - EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace reverse Tα $Tα Tβ $Tβ") - end - for Tα in (Const, Duplicated), Tβ in (Const, Duplicated) - EnzymeTestUtils.test_forward(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace forward Tα $Tα Tβ $Tβ") - end + ip = _repartition(invperm(linearize((_p, _q))), k′) + W = try + permute(prod(V1) ⊗ V3 ← V4, ip) + catch e + e isa SpaceMismatch || rethrow() + return nothing end + return dim(TensorKit.select(W, p)) > 0 ? (W, p, q) : nothing + end + for _ in 1:3 + trace = nothing + for _ in 1:50 + trace = randplanartrace() + isnothing(trace) || break + end + isnothing(trace) && continue + W, p, q = trace + A = randn(T, W) + α = randn(T) + β = randn(T) + C = randn!(TensorOperations.tensoralloc_add(T, A, p, false, Val(false))) + for Tα in rTαs, Tβ in rTβs + EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace reverse Tα $Tα Tβ $Tβ") + end + for Tα in fTαs, Tβ in fTβs + EnzymeTestUtils.test_forward(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace forward Tα $Tα Tβ $Tβ") + end + # the pullback of adjoint tensors goes through their subblocks + ba = (TensorOperations.DefaultBackend(), TensorOperations.DefaultAllocator()) + ΔC = randn!(similar(C)) + ΔA = zerovector(A) + TensorKit.planartrace_pullback_ΔA!(ΔA, ΔC, A, p, q, α, ba...) + ΔA′ = zerovector(copy(A'))' + TensorKit.planartrace_pullback_ΔA!(ΔA′, copy(ΔC')', copy(A')', p, q, α, ba...) + @test copy(ΔA′) ≈ ΔA end end end From b5b6991e186c56dbca11e2b72b0389714c5eaac4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 3 Nov 2026 10:39:35 -0500 Subject: [PATCH 6/9] Drop the coverage exclusions around GPU kernels and the `@index` workaround The workaround for nested `@index` in the CPU backend of KernelAbstractions is no longer needed since KernelAbstractions v0.9.45 (JuliaGPU/KernelAbstractions.jl#877). Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitGPUArraysExt.jl | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index cf5dbb967..53a96742d 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -19,14 +19,11 @@ 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 - # kernels are not reachable by coverage @kernel function fill_subblock_kernel!(subblock, val) idx = @index(Global, Cartesian) idx_val = idx[1] == idx[4] && idx[2] == idx[3] ? val : zero(val) @inbounds subblock[idx] = idx_val end - # COV_EXCL_STOP kernel = fill_subblock_kernel!(get_backend(data)) kernel(data, val; ndrange = size(data)) return data @@ -469,9 +466,6 @@ function device_transformer(A, transformer::TraceTransformer, p, q) end end -# COV_EXCL_START -# kernels are not reachable by coverage - # largest `i` with `offsets[i] <= w`. This corresponds to the # block which this kernel thread will work on. Since this is # used inside a GPU kernel, searchsortedlast/searchsortedfirst @@ -510,9 +504,7 @@ end @kernel function unique_batched_permute_kernel!( data_dst, data_src, op, blocks, work_offsets, α, β, nwork, ::Val{N} ) where {N} - # w = @index(Global, Linear) - 1 does not work due to JLArrays bug - w = @index(Global, Linear) - w -= 1 + w = @index(Global, Linear) - 1 if w < nwork b = _searchblock(work_offsets, w) blk = @inbounds blocks[b] @@ -523,8 +515,6 @@ end end end -# COV_EXCL_STOP - # One thread per destination element in `data_dst`. This makes much better use of the # GPU "massive parallelism" as compared to the one-thread-per-subtransformer approach. # It also more evenly divides the work among threads so the work profile is less @@ -537,9 +527,7 @@ end data_dst, data_src, op, blocks, work_offsets, coeffs, structs_dst, structs_src, α, β, nwork, ::Val{N} ) where {N} - # w = @index(Global, Linear) - 1 does not work due to JLArrays bug - w = @index(Global, Linear) - w -= 1 + w = @index(Global, Linear) - 1 if w < nwork # bookkeeping to figure out where to read from and write to b = _searchblock(work_offsets, w) @@ -568,13 +556,10 @@ end end end -# COV_EXCL_START # 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 + w = @index(Global, Linear) - 1 if w < npartials j = _searchblock(partial_offsets, w) src = @inbounds sources[j] @@ -598,8 +583,7 @@ end data_dst, partials, blocks, work_offsets, coeffs, structs_dst, sources, partial_offsets, α, β, nwork ) - w = @index(Global, Linear) - w -= 1 + w = @index(Global, Linear) - 1 if w < nwork b = _searchblock(work_offsets, w) blk = @inbounds blocks[b] @@ -627,7 +611,6 @@ end @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 From 76da813815c8d9697618d4882a187033cdb22b56 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 3 Nov 2026 11:19:38 -0500 Subject: [PATCH 7/9] Explain the two-stage GPU trace kernels next to the kernels Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitGPUArraysExt.jl | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 53a96742d..379223531 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -271,12 +271,6 @@ 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} @@ -556,8 +550,19 @@ end end end -# 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. +# Planar traces run in two kernels, such that the work per thread stays bounded even when a few +# destination elements receive a huge number of traced entries (e.g. a full trace): +# 1. `partial_trace_kernel!` splits the traced entries of every source subblock into `nchunks` +# chunks of `chunksize`, with one thread summing one chunk for one remaining entry. +# 2. `trace_recouple_kernel!` has one thread per destination element, which only combines the +# `nchunks` partials of each of the `cols` sources of its block with the coefficients in `U`. +# Since `chunksize` and `nchunks` are both about `√ntraced`, no thread in either kernel sums more +# than `max(256, √ntraced)` entries per source, and each traced entry is read only once rather +# than once per destination tree. A full trace of `n` entries uses `√n` threads in the first +# kernel and a single thread per destination tree in the second. + +# `partials[w + 1]` is the sum of chunk `k` of the traced entries of source `j` at its remaining +# entry `e`; consecutive threads share `k` and have consecutive `e`. @kernel function partial_trace_kernel!(partials, data_src, sources, partial_offsets, npartials) w = @index(Global, Linear) - 1 if w < npartials @@ -577,8 +582,8 @@ end end end -# Second stage: one thread per destination element, as for the generic transformers, where -# every source contributes the sum of its partial traces. +# The traced entries were already summed by `partial_trace_kernel!`, so each thread only adds up +# the partials of the sources of its destination tree. @kernel function trace_recouple_kernel!( data_dst, partials, blocks, work_offsets, coeffs, structs_dst, sources, partial_offsets, α, β, nwork From 834183c305b41787a8bfeec1114abf0cf2da2985 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 4 Nov 2026 09:32:04 -0500 Subject: [PATCH 8/9] Test a single random planar trace per sector in the Enzyme tests Co-Authored-By: Claude Opus 5.5 --- test/enzyme/trace.jl | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/test/enzyme/trace.jl b/test/enzyme/trace.jl index 6ccfca2df..d36d35520 100644 --- a/test/enzyme/trace.jl +++ b/test/enzyme/trace.jl @@ -80,13 +80,12 @@ TAs = is_ci ? (Duplicated,) : (Const, Duplicated) end return dim(TensorKit.select(W, p)) > 0 ? (W, p, q) : nothing end - for _ in 1:3 - trace = nothing - for _ in 1:50 - trace = randplanartrace() - isnothing(trace) || break - end - isnothing(trace) && continue + trace = nothing + for _ in 1:50 + trace = randplanartrace() + isnothing(trace) || break + end + if !isnothing(trace) W, p, q = trace A = randn(T, W) α = randn(T) From 4a0352eb33a6e4ea1dcaa33e4ecc3af749bff17e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 4 Nov 2026 09:49:46 -0500 Subject: [PATCH 9/9] Define `scalartype` for tree transformers Co-Authored-By: Claude Opus 5.5 --- src/pullbacks/tensoroperations.jl | 4 +--- src/tensors/treetransformers.jl | 4 ++++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/pullbacks/tensoroperations.jl b/src/pullbacks/tensoroperations.jl index 5b972a3e3..c157b0987 100644 --- a/src/pullbacks/tensoroperations.jl +++ b/src/pullbacks/tensoroperations.jl @@ -75,7 +75,7 @@ function planartrace_pullback_ΔA!( BraidingStyle(sectortype(A)) == Bosonic() && return trace_permute_pullback_ΔA!(ΔA, ΔC, A, p, q, α, backend) transformer = treetracer(storagetype(ΔC), space(ΔC), space(A), p, q) - T = promote_type(scalartype(ΔC), scalartype(α), _coefficienttype(transformer)) + T = promote_type(scalartype(ΔC), scalartype(α), scalartype(transformer)) if T <: Real || !(scalartype(ΔA) <: Real) _planartrace_pullback_ΔA!(ΔA, ΔC, transformer, p, q, α, backend, allocator) else @@ -86,8 +86,6 @@ function planartrace_pullback_ΔA!( return nothing end -_coefficienttype(::TraceTransformer{T}) where {T} = T - function _planartrace_pullback_ΔA!( ΔA, ΔC, transformer::TraceTransformer, (p₁, p₂), (q₁, q₂), α, backend, allocator ) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 73d910775..ba3e9df6b 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -220,6 +220,10 @@ struct TraceTransformer{T, N₁, N₂} structure_src::Vector{StridedStructure{N₂}} end +VectorInterface.scalartype( + ::Type{<:Union{UniqueTreeTransformer{T}, GenericTreeTransformer{T}, TraceTransformer{T}}} +) where {T} = T + 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)