Precompute planar trace kernels in a TraceTransformer

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, and GPU tensors use a dedicated kernel (internal)
 
 ### Deprecated
 

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)
 ```

ext/TensorKitGPUArraysExt.jl

@@ -9,11 +9,13 @@ 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
-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 +273,112 @@ 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`/`sources` 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
+
+"""
+    TraceSourceBlock{N, M}
+
+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
+    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₁}[]
+    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_dst, size(U, 1), size(U, 2),
+                length(coeffs), length(structs_dst), length(sources)
+            )
+        )
+        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₂)
+            ntraced = prod(size_traced)
+            chunksize = _trace_chunksize(ntraced)
+            push!(
+                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)
+    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.
 @cached function TensorKit.treebraider(
         A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple,
@@ -287,6 +395,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 +410,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 +449,19 @@ 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.sources),
+            upload(packed.partial_offsets), packed.npartials
+        )
+    end
+end
+
 # COV_EXCL_START
 # kernels are not reachable by coverage
 
@@ -360,7 +490,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 +561,67 @@ 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
+    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, 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(partials), eltype(coeffs)))
+        @inbounds for j in 1:blk.cols
+            coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows]
+            iszero(coeff) && continue
+            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
+        # 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 +646,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 +672,40 @@ 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
+    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
+
+# 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