From 73cd3f842595585960ae7d4d22ba349a2cabcf2e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 6 Nov 2026 10:49:07 -0500 Subject: [PATCH] Tiled shared-memory kernel for GPU index manipulations Replace the per-element unique/generic permute kernels with a single tiled KernelAbstractions kernel: each workgroup gathers one power-of-two tile of a destination tree, recoupled over the source trees, into shared memory in source order and writes it out in destination order, so reads and writes are both coalesced. Unique and recoupling blocks share one launch, and the diagonal paths use the same kernel with one-dimensional blocks. - tile shapes are chosen per transformer with a sector-based cost model - the tile length is branched on statically, so `DeviceTreeTransformer` has a concrete type and the device caches declare their return types - tiles of blocks with more source trees are scheduled first - split the GPUArrays extension into several files Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitGPUArraysExt.jl | 823 ------------------ .../TensorKitGPUArraysExt.jl | 85 ++ ext/TensorKitGPUArraysExt/diagonal.jl | 96 ++ .../indexmanipulations.jl | 298 +++++++ ext/TensorKitGPUArraysExt/trace.jl | 245 ++++++ ext/TensorKitGPUArraysExt/truncation.jl | 77 ++ ext/TensorKitGPUArraysExt/utility.jl | 30 + test/cuda/tensors.jl | 2 +- 8 files changed, 832 insertions(+), 824 deletions(-) delete mode 100644 ext/TensorKitGPUArraysExt.jl create mode 100644 ext/TensorKitGPUArraysExt/TensorKitGPUArraysExt.jl create mode 100644 ext/TensorKitGPUArraysExt/diagonal.jl create mode 100644 ext/TensorKitGPUArraysExt/indexmanipulations.jl create mode 100644 ext/TensorKitGPUArraysExt/trace.jl create mode 100644 ext/TensorKitGPUArraysExt/truncation.jl create mode 100644 ext/TensorKitGPUArraysExt/utility.jl diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl deleted file mode 100644 index 5b8e43a97..000000000 --- a/ext/TensorKitGPUArraysExt.jl +++ /dev/null @@ -1,823 +0,0 @@ -module TensorKitGPUArraysExt - -using GPUArrays -using GPUArrays: @allowscalar -using GPUArrays.KernelAbstractions: @kernel, @index, get_backend -using Adapt -using TensorKit.TupleTools -using Strided: StridedViews -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 -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}}}} - @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 - kernel = fill_subblock_kernel!(get_backend(data)) - kernel(data, val; ndrange = size(data)) - return data -end - -const GPUSectorVector{T, I} = TensorKit.SectorVector{T, I, <:AnyGPUVector{T}} - -function MatrixAlgebraKit.findtruncated( - values::GPUSectorVector, strategy::MatrixAlgebraKit.TruncationByOrder - ) - I = sectortype(values) - - dims = similar(values, Base.promote_op(dim, I)) - for (c, v) in pairs(dims) - fill!(v, dim(c)) - end - - isempty(parent(values)) && return similar(values, Bool) - - perm = sortperm(parent(values); strategy.by, strategy.rev) - cumulative_dim = cumsum(Base.permute!(parent(dims), perm)) - - result = similar(values, Bool) - parent(result)[perm] .= cumulative_dim .<= strategy.howmany - return result -end - -function MatrixAlgebraKit.findtruncated( - values::GPUSectorVector, strategy::MatrixAlgebraKit.TruncationByError - ) - (isfinite(strategy.p) && strategy.p > 0) || - throw(ArgumentError(lazy"p-norm with p = $(strategy.p) is currently not supported.")) - ϵᵖmax = max(strategy.atol^strategy.p, strategy.rtol^strategy.p * norm(values, strategy.p)) - ϵᵖ = similar(values, typeof(ϵᵖmax)) - - # dimensions are all 1 so no need to account for weight - if FusionStyle(sectortype(values)) isa UniqueFusion - parent(ϵᵖ) .= abs.(parent(values)) .^ strategy.p - else - for (c, v) in pairs(values) - v′ = ϵᵖ[c] - v′ .= abs.(v) .^ strategy.p .* dim(c) - end - end - - isempty(parent(values)) && return similar(values, Bool) - - perm = sortperm(parent(values); by = abs, rev = false) - cumulative_err = cumsum(Base.permute!(parent(ϵᵖ), perm)) - - result = similar(values, Bool) - parent(result)[perm] .= cumulative_err .> ϵᵖmax - return result -end - -function MatrixAlgebraKit.findtruncated_svd(values::GPUSectorVector, strategy::S) where {S <: MatrixAlgebraKit.TruncationStrategy} - # returning a GPUSectorVector wrecks things in truncate_{co}domain - # because of scalar indexing - return Adapt.adapt(Vector, MatrixAlgebraKit.findtruncated(values, strategy)) -end - -for strat in (:(MatrixAlgebraKit.TruncationByOrder), :(MatrixAlgebraKit.TruncationByError), :(MatrixAlgebraKit.TruncationIntersection), :(MatrixAlgebraKit.TruncationMultiplets), :(TensorKit.Factorizations.TruncationSpace)) - @eval function MatrixAlgebraKit.findtruncated_svd(values::GPUSectorVector, strategy::$strat) - # returning a GPUSectorVector wrecks things in truncate_{co}domain - # because of scalar indexing - return Adapt.adapt(Vector, MatrixAlgebraKit.findtruncated(values, strategy)) - end -end - -function MatrixAlgebraKit.findtruncated_svd(values::GPUSectorVector, strategy::MatrixAlgebraKit.TruncationByValue) - atol = TensorKit.Factorizations.rtol_to_atol(values, strategy.p, strategy.atol, strategy.rtol) - strategy′ = trunctol(; atol, strategy.by, strategy.keep_below) - return SectorDict(c => Adapt.adapt(Vector, MatrixAlgebraKit.findtruncated_svd(d, strategy′)) for (c, d) in pairs(values)) -end - -function MatrixAlgebraKit.truncation_error!(values::GPUSectorVector, ind::AbstractVector{Bool}) - for (c, ind_c) in pairs(ind) - sector_vals = values[c] - @. sector_vals *= !ind_c - end - return norm(values) -end - -# project_symmetric! doesn't yet work for GPU types, so do this on the host, then copy -function TensorKit.project_symmetric_and_check(::Type{T}, ::Type{A}, data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data)))))) where {T, A <: AnyGPUVector{T}} - h_t = TensorKit.TensorMapWithStorage{T, Vector{T}}(undef, V) - h_t = TensorKit.project_symmetric!(h_t, Array(data)) - # verify result - isapprox(Array(reshape(data, dims(h_t))), convert(Array, h_t); atol = tol) || - throw(ArgumentError("Data has non-zero elements at incompatible positions")) - return TensorKit.TensorMapWithStorage{T, A}(A(h_t.data), V) -end - -# Scalar implementation -#----------------------- -function TensorKit.scalar(t::TensorMap{T, S, 0, 0, <:AnyGPUArray}) where {T, S} - inds = findall(!iszero, t.data) - return isempty(inds) ? zero(scalartype(t)) : @allowscalar @inbounds t.data[only(inds)] -end - -# Device-side tree transformers -# ----------------------------- -# A `TreeTransformer` stores the mapping between subblock positions (plus recoupling -# coefficients) and the `StridedStructure`s of the source and destination spaces. -# We resolve the positions into sizes/strides/offsets and pack all the information on the -# CPU side into dense vectors of numbers, plus some accounting information so we know how -# to unpack in the kernel. Also, we can permute the source strides "in advance" on the CPU side. -# We also precompute the strides of the subblock each kernel index will work on, -# so that the GPU thread can recover the Cartesian coordinates it will need for input/ouput, -# and a running `work_offsets` count of destination elements, so that a kernel can run one -# thread per output element and recover which input that element belongs with. -# Some possible TODO here: -# - Try cuTILE as this is a classic tile programming problem -# - Use shared memory to coalesce the reads -# - Use a 2D grid for the Generic case - -const TreeStructure{N} = Tuple{NTuple{N, Int}, Int} - -""" - UniqueTransformerBlock{T, N} - -`isbits` descriptor for a subblock that is a *single* scaled permutation: -an entry of a `UniqueTreeTransformer` or a unique (one-tree) block of a -`GenericTreeTransformer`. -""" -struct UniqueTransformerBlock{T, N} - coeff::T - sz::NTuple{N, Int} - dense_strides::NTuple{N, Int} - strides_dst::NTuple{N, Int} - offsets_dst::Int - permuted_strides_src::NTuple{N, Int} # source strides, permuted by `p` - offsets_src::Int -end - -# device-side simple struct that GPU kernels can use -struct DeviceUniqueTreeTransformer{VB <: AbstractVector{<:UniqueTransformerBlock}, VO <: AbstractVector{Int}} - blocks::VB - work_offsets::VO - nwork::Int -end - -""" - GenericTransformerBlock{N} - -Descriptor for a recoupling block of a `GenericTreeTransformer`, indexing into the -flat `coeffs`/`structs_dst`/`structs_src` vectors of a `DeviceGenericTreeTransformer`. -""" -struct GenericTransformerBlock{N} - sz::NTuple{N, Int} - densestrides::NTuple{N, Int} - rows::Int - cols::Int - u_offset::Int # location in the flattened U vector to find this block's U - dst_offset::Int - src_offset::Int -end - -# force all the type signatures here to make sure doing something wrong fails -# before the kernel launch. Kernel error dumps are awful and hard to interpret. -struct DeviceGenericTreeTransformer{VO <: AbstractVector{Int}, DA <: DeviceUniqueTreeTransformer{<:Any, VO}, VB <: AbstractVector{<:GenericTransformerBlock}, VC <: AbstractVector{<:Number}, VS <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}} - unique_blocks::DA # length(U) = 1 blocks, can be handled by unique kernel - blocks::VB - work_offsets::VO - nwork::Int - coeffs::VC # every `U`, concatenated in column-major order - structs_dst::VS - structs_src::VS -end - -# strides of a dense array of shape `sz` -_dense_strides(size::Dims) = Base.front((1, cumprod(size)...)) - -# `permute(Vsrc, p) == Vdst` is enforced when the transformer is built, so the permuted -# source shape always matches `sz_dst` and the two views share Cartesian inds. -function _unique_block( - coeff::T, (size_dst, strides_dst, offsets_dst), (_, strides_src, offsets_src), p - ) where {T} - return UniqueTransformerBlock{T, length(size_dst)}( - coeff, size_dst, _dense_strides(size_dst), strides_dst, offsets_dst, - TupleTools.getindices(strides_src, p), offsets_src - ) -end - -function _work_offsets(work) - offsets = cumsum(work) - pushfirst!(offsets, 0) - total = pop!(offsets) - return offsets, total -end - -function DeviceUniqueTreeTransformer(transformer::UniqueTreeTransformer{T, N}, p) where {T, N} - (; structure_dst, structure_src) = transformer - blocks = UniqueTransformerBlock{T, N}[ - _unique_block(coeff, structure_dst[idst], structure_src[isrc], p) - for (coeff, idst, isrc) in transformer.data - ] - work_offsets, nwork = _work_offsets(prod(blk.sz) for blk in blocks) - return DeviceUniqueTreeTransformer(blocks, work_offsets, nwork) -end - -function DeviceGenericTreeTransformer( - transformer::GenericTreeTransformer{T, N}, p - ) where {T, N} - unique_blocks = UniqueTransformerBlock{T, N}[] - blocks = GenericTransformerBlock{N}[] - coeffs = T[] - structs_dst = TreeStructure{N}[] - structs_src = TreeStructure{N}[] - - (; structure_dst, structure_src) = transformer - for (U, inds_dst, inds_src) in transformer.data - if length(U) == 1 # same as the unique (Abelian) case - push!( - unique_blocks, _unique_block( - only(U), structure_dst[only(inds_dst)], structure_src[only(inds_src)], p - ) - ) - else - # all trees in a block share the same subblock size - size_dst = first(structure_dst[first(inds_dst)]) - push!( - blocks, GenericTransformerBlock{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 - _, strides_src, offset_src = structure_src[isrc] - push!(structs_src, (TupleTools.getindices(strides_src, p), offset_src)) - end - end - end - - unique_offsets, unique_nwork = _work_offsets(prod(blk.sz) for blk in unique_blocks) - work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) - return DeviceGenericTreeTransformer( - DeviceUniqueTreeTransformer(unique_blocks, unique_offsets, unique_nwork), - blocks, work_offsets, nwork, coeffs, structs_dst, structs_src - ) -end - -""" - 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, - conjsrc::Bool, levels::IndexTuple - ) where {TA <: AnyGPUVector} - host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) - return device_transformer(A, host, linearize(p)) -end - -@cached function TensorKit.treetransposer( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool - ) where {TA <: AnyGPUVector} - host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) - 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 - -@cached function TensorKit.treetwister( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool - ) where {TA <: AnyGPUVector, N} - host = TensorKit.treetwister(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask, inv) - return device_transformer(A, host, linearize(p)) -end - -const DEFAULT_DEVICECACHE_SIZE = 1000 - -function __init__() - # extensions can't modify GLOBAL_CACHES during precompilation, so register explicitly - append!( - TensorKit.GLOBAL_CACHES, ( - Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, - 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_DIAGONAL_TREETWISTER_CACHE") => GLOBAL_DIAGONAL_TREETWISTER_CACHE, - Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_CACHE, - Symbol("TensorKitGPUArraysExt.GLOBAL_TREETWISTER_CACHE") => GLOBAL_TREETWISTER_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_DIAGONAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - resize!(GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - resize!(GLOBAL_DIAGONAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - resize!(GLOBAL_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - resize!(GLOBAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - - return nothing -end - -# Allocations must outlive any GPUArrays allocation-cache scope around the operation. -device_transformer(A, transformer::UniqueTreeTransformer, p) = - device_transformer(A, DeviceUniqueTreeTransformer(transformer, p)) - -function device_transformer(A, packed::DeviceUniqueTreeTransformer) - return GPUArrays.@uncached begin - proto = A(undef, 0) - # Metadata and coefficients have different element types; preserve each on upload. - upload(x) = copy!(similar(proto, eltype(x), size(x)), x) - DeviceUniqueTreeTransformer(upload(packed.blocks), upload(packed.work_offsets), packed.nwork) - end -end - -function device_transformer(A, transformer::GenericTreeTransformer, p) - packed = DeviceGenericTreeTransformer(transformer, p) - return GPUArrays.@uncached begin - proto = A(undef, 0) - upload(x) = copy!(similar(proto, eltype(x), size(x)), x) - unique = packed.unique_blocks - device_unique = DeviceUniqueTreeTransformer( - upload(unique.blocks), upload(unique.work_offsets), unique.nwork - ) - DeviceGenericTreeTransformer( - device_unique, upload(packed.blocks), upload(packed.work_offsets), packed.nwork, - upload(packed.coeffs), upload(packed.structs_dst), upload(packed.structs_src) - ) - 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 - -# 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 -# won't work. -@inline function _searchblock(offsets, w) - lo, hi = 1, length(offsets) - while lo < hi - mid = (lo + hi + 1) >>> 1 - if @inbounds offsets[mid] <= w - lo = mid - else - hi = mid - 1 - end - end - return lo -end - -# Cartesian coordinates of the `w`-th (0-based) entry of a dense subblock of shape `sz`. -@inline function _coordinates(w, size::NTuple{N, Int}, dense_strides::NTuple{N, Int}) where {N} - return ntuple(n -> (w ÷ dense_strides[n]) % size[n], Val(N)) -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} - # 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 -# applied to the source data only (not to the coefficients). -@kernel function unique_batched_permute_kernel!( - data_dst, data_src, op, blocks, work_offsets, α, β, nwork, ::Val{N} - ) where {N} - w = @index(Global, Linear) - 1 - if w < nwork - b = _searchblock(work_offsets, w) - blk = @inbounds blocks[b] - coords = _coordinates(w - (@inbounds work_offsets[b]), blk.sz, blk.dense_strides) - i_dst = _linear_index(coords, blk.strides_dst, blk.offsets_dst) - i_src = _linear_index(coords, blk.permuted_strides_src, blk.offsets_src) - @inbounds data_dst[i_dst] = α * blk.coeff * op(data_src[i_src]) + β * data_dst[i_dst] - end -end - -# 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 -# jagged. Unlike the CPU implementation, there is no extract → recouple → insert process: -# BLAS is not generally reachable from inside a kernel, and fusing the recoupling into -# the strided gather lets us remove the buffer entirely. -# TODO: what about symmetries like SU(3), where the column by column approach is not -# optimal? -@kernel function generic_batched_permute_kernel!( - data_dst, data_src, op, blocks, work_offsets, coeffs, structs_dst, structs_src, - α, β, nwork, ::Val{N} - ) where {N} - w = @index(Global, Linear) - 1 - if w < nwork - # bookkeeping to figure out where to read from and write to - 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.densestrides) - - st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1] - i_dst = _linear_index(coords, st_dst, offs_dst) - - # dst_i = β * dst_i + α * Σ_j U[i, j] * permute(op(src_j), p): each output tree is a - # linear combination of the input trees weighted by the recoupling coefficients. - # The permutation of src_j was already done by permuting its strides before the - # kernel launched. - acc = zero(promote_type(eltype(data_src), eltype(coeffs))) - @inbounds for j in 1:blk.cols - # TODO is there a more efficient way to do this read? - coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] - iszero(coeff) && continue - pst_src, offs_src = structs_src[blk.src_offset + j] - acc += coeff * op(data_src[_linear_index(coords, pst_src, offs_src)]) - end - @inbounds data_dst[i_dst] = α * acc + β * data_dst[i_dst] - end -end - -# 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 - 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 - -# 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 - ) - w = @index(Global, Linear) - 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 - -function _launch_unique!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} - nwork = transformer.nwork - nwork == 0 && return nothing - unique_batched_permute_kernel!(get_backend(data_dst))( - data_dst, data_src, op, transformer.blocks, transformer.work_offsets, α, β, nwork, - Val(N); ndrange = nwork - ) - return nothing -end - -function _launch_generic!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} - nwork = transformer.nwork - nwork == 0 && return nothing - generic_batched_permute_kernel!(get_backend(data_dst))( - data_dst, data_src, op, transformer.blocks, transformer.work_offsets, - transformer.coeffs, transformer.structs_dst, transformer.structs_src, α, β, nwork, - Val(N); ndrange = nwork - ) - return nothing -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, DeviceTraceTransformer} -) = (tdst.data, tsrc.data) - -# Diagonal sources -# ---------------- -# A `DiagonalTensorMap` only stores the diagonal of each of its subblocks, and an index -# manipulation maps the entry `(i, i)` of a source subblock onto the entry `(i, i)` of a destination -# subblock. Every block is therefore a one-dimensional strided copy, which the unique kernel -# handles without needing a dense copy of the source or destination. - -const GPUDiagonalTensorMap = DiagonalTensorMap{<:Any, <:Any, <:AnyGPUVector} -const GPUTensorMap = TensorMap{<:Any, <:Any, <:Any, <:Any, <:AnyGPUVector} - -# position of each subblock in the data of a `DiagonalTensorMap` in the same space -_diagonal_offsets(structure) = first(_work_offsets(first(first(s)) for s in structure)) - -_diagonal_entries(transformer::UniqueTreeTransformer) = transformer.data -# `V ← V` and any rearrangement of its two indices have a single tree for every subblock -_diagonal_entries(transformer::GenericTreeTransformer) = - ((only(U), only(inds_dst), only(inds_src)) for (U, inds_dst, inds_src) in transformer.data) - -function diagonal_device_transformer(A, transformer, dstdiagonal::Bool) - (; structure_dst, structure_src) = transformer - offsets_src = _diagonal_offsets(structure_src) - offsets_dst = dstdiagonal ? _diagonal_offsets(structure_dst) : nothing - blocks = map(_diagonal_entries(transformer)) do (coeff, idst, isrc) - (sz, strides_dst, offset_dst) = structure_dst[idst] - n = first(sz) - return UniqueTransformerBlock{typeof(coeff), 1}( - coeff, (n,), (1,), dstdiagonal ? (1,) : (sum(strides_dst),), - dstdiagonal ? offsets_dst[idst] : offset_dst, (1,), offsets_src[isrc] - ) - end - work_offsets, nwork = _work_offsets(blk.sz[1] for blk in blocks) - return device_transformer(A, DeviceUniqueTreeTransformer(blocks, work_offsets, nwork)) -end - -@cached function diagonal_treebraider( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, levels::IndexTuple, dstdiagonal::Bool - ) where {TA <: AnyGPUVector} - host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) - return diagonal_device_transformer(A, host, dstdiagonal) -end - -@cached function diagonal_treetransposer( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, dstdiagonal::Bool - ) where {TA <: AnyGPUVector} - host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) - return diagonal_device_transformer(A, host, dstdiagonal) -end - -@cached function diagonal_treetwister( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool, dstdiagonal::Bool - ) where {TA <: AnyGPUVector, N} - host = TensorKit.treetwister(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask, inv) - return diagonal_device_transformer(A, host, dstdiagonal) -end - -function _add_diagonal!(tdst::GPUDiagonalTensorMap, tsrc, conjsrc, transformer, α, β) - _launch_unique!(tdst.data, tsrc.data, conjsrc ? conj : identity, transformer, α, β, Val(1)) - return tdst -end -# the kernel only touches the diagonal. be careful to apply `β` to the entire tensor. -function _add_diagonal!(tdst::GPUTensorMap, tsrc, conjsrc, transformer, α, β) - scale!(tdst, β) - _launch_unique!(tdst.data, tsrc.data, conjsrc ? conj : identity, transformer, α, One(), Val(1)) - return tdst -end - -Base.@propagate_inbounds function TensorKit._braid!( - tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, - p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator - ) - Vdst, Vsrc = space(tdst), space(tsrc) - @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) - transformer = diagonal_treebraider( - storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels, tdst isa DiagonalTensorMap - ) - return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) -end - -Base.@propagate_inbounds function TensorKit._transpose!( - tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, - p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator - ) - Vdst, Vsrc = space(tdst), space(tsrc) - @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) - transformer = diagonal_treetransposer( - storagetype(tdst), Vdst, Vsrc, p, conjsrc, tdst isa DiagonalTensorMap - ) - return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) -end - -Base.@propagate_inbounds function TensorKit._twist!( - tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, - p::Index2Tuple, conjsrc::Bool, twistmask::Tuple{Vararg{Bool}}, inv::Bool, α, β, backend, allocator - ) - Vdst, Vsrc = space(tdst), space(tsrc) - @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) - transformer = diagonal_treetwister( - storagetype(tdst), Vdst, Vsrc, p, conjsrc, twistmask, inv, tdst isa DiagonalTensorMap - ) - return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) -end - -function TensorKit.add_transform_kernel!( - dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, - device::DeviceUniqueTreeTransformer, α, β, backend, allocator, ntasks::Int - ) - N = length(linearize(p)) - op = conjsrc ? conj : identity - _launch_unique!(dst, src, op, device, α, β, Val(N)) - return nothing -end - -function TensorKit.add_transform_kernel!( - dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, - device::DeviceGenericTreeTransformer, α, β, backend, allocator, ntasks::Int - ) - N = length(linearize(p)) - op = conjsrc ? conj : identity - # one-tree blocks are a scaled permutation, which the unique kernel already handles; the - # two kernels touch disjoint subblocks so the launch order does not matter - _launch_unique!(dst, src, op, device.unique_blocks, α, β, Val(N)) - _launch_generic!(dst, src, op, device, α, β, Val(N)) - 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 diff --git a/ext/TensorKitGPUArraysExt/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt/TensorKitGPUArraysExt.jl new file mode 100644 index 000000000..1dbbd1153 --- /dev/null +++ b/ext/TensorKitGPUArraysExt/TensorKitGPUArraysExt.jl @@ -0,0 +1,85 @@ +module TensorKitGPUArraysExt + +using GPUArrays +using GPUArrays: @allowscalar +using GPUArrays.KernelAbstractions: @kernel, @index, @localmem, @synchronize, @Const, get_backend +using Adapt +using TensorKit.TupleTools +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, TraceTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple +using TensorKit: SubblockIterator +import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel!, trace_transform_kernel! + +const GPUDiagonalTensorMap = DiagonalTensorMap{<:Any, <:Any, <:AnyGPUVector} +const GPUTensorMap = TensorMap{<:Any, <:Any, <:Any, <:Any, <:AnyGPUVector} + +function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:AnyGPUMatrix{T}, <:StridedViews.StridedView{T, 4, <:AnyGPUArray{T}}}} + @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 + kernel = fill_subblock_kernel!(get_backend(data)) + kernel(data, val; ndrange = size(data)) + return data +end + +# project_symmetric! doesn't yet work for GPU types, so do this on the host, then copy +function TensorKit.project_symmetric_and_check(::Type{T}, ::Type{A}, data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data)))))) where {T, A <: AnyGPUVector{T}} + h_t = TensorKit.TensorMapWithStorage{T, Vector{T}}(undef, V) + h_t = TensorKit.project_symmetric!(h_t, Array(data)) + # verify result + isapprox(Array(reshape(data, dims(h_t))), convert(Array, h_t); atol = tol) || + throw(ArgumentError("Data has non-zero elements at incompatible positions")) + return TensorKit.TensorMapWithStorage{T, A}(A(h_t.data), V) +end + +# Scalar implementation +#----------------------- +function TensorKit.scalar(t::TensorMap{T, S, 0, 0, <:AnyGPUArray}) where {T, S} + inds = findall(!iszero, t.data) + return isempty(inds) ? zero(scalartype(t)) : @allowscalar @inbounds t.data[only(inds)] +end + +include("truncation.jl") +include("utility.jl") +include("indexmanipulations.jl") +include("trace.jl") +include("diagonal.jl") + +const DEFAULT_DEVICECACHE_SIZE = 1000 + +function __init__() + # extensions can't modify GLOBAL_CACHES during precompilation, so register explicitly + append!( + TensorKit.GLOBAL_CACHES, ( + Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, + 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_DIAGONAL_TREETWISTER_CACHE") => GLOBAL_DIAGONAL_TREETWISTER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETWISTER_CACHE") => GLOBAL_TREETWISTER_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_DIAGONAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_DIAGONAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + + return nothing +end + +end diff --git a/ext/TensorKitGPUArraysExt/diagonal.jl b/ext/TensorKitGPUArraysExt/diagonal.jl new file mode 100644 index 000000000..eb9d56945 --- /dev/null +++ b/ext/TensorKitGPUArraysExt/diagonal.jl @@ -0,0 +1,96 @@ +# Diagonal sources +# ---------------- +# An index manipulation of a `DiagonalTensorMap` maps the diagonal of every source subblock onto +# the diagonal of a destination subblock, i.e. every block is a one-dimensional strided copy. + +# position of each subblock in the data of a `DiagonalTensorMap` in the same space +_diagonal_offsets(structure) = first(_work_offsets(first(first(s)) for s in structure)) + +_diagonal_entries(transformer::UniqueTreeTransformer) = transformer.data +# `V ← V` and any rearrangement of its two indices have a single tree for every subblock +_diagonal_entries(transformer::GenericTreeTransformer) = + ((only(U), only(inds_dst), only(inds_src)) for (U, inds_dst, inds_src) in transformer.data) + +function diagonal_device_transformer(A, transformer, dstdiagonal::Bool) + (; structure_dst, structure_src) = transformer + offsets_src = _diagonal_offsets(structure_src) + offsets_dst = dstdiagonal ? _diagonal_offsets(structure_dst) : nothing + entries = map(_diagonal_entries(transformer)) do (coeff, idst, isrc) + (sz, strides_dst, offset_dst) = structure_dst[idst] + n = first(sz) + s_dst = dstdiagonal ? ((n,), (1,), offsets_dst[idst]) : ((n,), (sum(strides_dst),), offset_dst) + return (fill(coeff, 1, 1), (s_dst,), (((n,), (1,), offsets_src[isrc]),)) + end + return device_transformer(A, DeviceTreeTransformer(scalartype(transformer), entries, (1,), sizeof(scalartype(A)))) +end + +@cached function diagonal_treebraider( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, levels::IndexTuple, dstdiagonal::Bool + )::devicetransformertype(A, sectortype(Vdst), Val(1)) where {TA <: AnyGPUVector} + host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) + return diagonal_device_transformer(A, host, dstdiagonal) +end + +@cached function diagonal_treetransposer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, dstdiagonal::Bool + )::devicetransformertype(A, sectortype(Vdst), Val(1)) where {TA <: AnyGPUVector} + host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) + return diagonal_device_transformer(A, host, dstdiagonal) +end + +@cached function diagonal_treetwister( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool, dstdiagonal::Bool + )::devicetransformertype(A, sectortype(Vdst), Val(1)) where {TA <: AnyGPUVector, N} + host = TensorKit.treetwister(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask, inv) + return diagonal_device_transformer(A, host, dstdiagonal) +end + +function _add_diagonal!(tdst::GPUDiagonalTensorMap, tsrc, conjsrc, transformer, α, β) + _launch_tiled!(tdst.data, tsrc.data, conjsrc, transformer, α, β) + return tdst +end +# the kernel only touches the diagonal. be careful to apply `β` to the entire tensor. +function _add_diagonal!(tdst::GPUTensorMap, tsrc, conjsrc, transformer, α, β) + scale!(tdst, β) + _launch_tiled!(tdst.data, tsrc.data, conjsrc, transformer, α, one(scalartype(tdst))) + return tdst +end + +Base.@propagate_inbounds function TensorKit._braid!( + tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, + p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator + ) + Vdst, Vsrc = space(tdst), space(tsrc) + @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + transformer = diagonal_treebraider( + storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels, tdst isa DiagonalTensorMap + ) + return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) +end + +Base.@propagate_inbounds function TensorKit._transpose!( + tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, + p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator + ) + Vdst, Vsrc = space(tdst), space(tsrc) + @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + transformer = diagonal_treetransposer( + storagetype(tdst), Vdst, Vsrc, p, conjsrc, tdst isa DiagonalTensorMap + ) + return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) +end + +Base.@propagate_inbounds function TensorKit._twist!( + tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, + p::Index2Tuple, conjsrc::Bool, twistmask::Tuple{Vararg{Bool}}, inv::Bool, α, β, backend, allocator + ) + Vdst, Vsrc = space(tdst), space(tsrc) + @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + transformer = diagonal_treetwister( + storagetype(tdst), Vdst, Vsrc, p, conjsrc, twistmask, inv, tdst isa DiagonalTensorMap + ) + return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) +end diff --git a/ext/TensorKitGPUArraysExt/indexmanipulations.jl b/ext/TensorKitGPUArraysExt/indexmanipulations.jl new file mode 100644 index 000000000..df716064a --- /dev/null +++ b/ext/TensorKitGPUArraysExt/indexmanipulations.jl @@ -0,0 +1,298 @@ +# Tiled index manipulations +# ------------------------- +# One workgroup of `G` threads per tile of `L` entries: the tile is gathered (recoupled) in source +# order into shared memory, then written out in destination order. The `L ÷ G` entries of a +# thread are unrolled through tuples, such that all its loads are in flight at once. + +""" + TransformerBlock{N} + +Descriptor for a block of a `DeviceTreeTransformer`, indexing into its flat +`coeffs`/`structs_dst`/`structs_src` vectors. A unique (one-tree) block has `rows = cols = 1`. +""" +struct TransformerBlock{N} + sz::NTuple{N, Int32} # destination subblock size + rows::Int32 + cols::Int32 + u_offset::Int32 # location of this block's `U` in the flattened `coeffs` vector + dst_offset::Int32 + src_offset::Int32 +end + +""" + TransformerTile{N} + +Work item of a single workgroup: a tile of the destination tree `tree` (0-based) of `block`. +""" +struct TransformerTile{N} + block::Int32 + tree::Int32 + origin::NTuple{N, Int32} # 0-based coordinates of the first tile entry +end + +""" + TileShape{N} + +A tile with power-of-two extents `2 .^ shift`, in the index order of the destination. +Tile entries are numbered both in destination order and in source order, and the coordinates +of an entry `e` are `(e >> shift[n]) & mask[n]` with the corresponding `shift`. +""" +struct TileShape{N} + dshift::NTuple{N, Int32} + sshift::NTuple{N, Int32} + mask::NTuple{N, Int32} + len::Int32 + padshift::Int32 # shared memory has one padding entry per `2^padshift` entries +end + +function TileShape(sz::NTuple{N, Int}, p::NTuple{N, Int}) where {N} + lg = map(trailing_zeros, sz) + dshift = ntuple(n -> sum(lg[1:(n - 1)]; init = 0), Val(N)) + # destination index `n` is source index `p[n]` + q = invperm(p) + spos = ntuple(k -> sum(lg[q[i]] for i in 1:(k - 1); init = 0), Val(N)) + sshift = ntuple(n -> spos[p[n]], Val(N)) + padshift = N == 0 ? 0 : lg[1] + return TileShape{N}(Int32.(dshift), Int32.(sshift), Int32.(sz .- 1), Int32(prod(sz)), Int32(padshift)) +end + +struct DeviceTreeTransformer{N, VB <: AbstractVector{TransformerBlock{N}}, VT <: AbstractVector{TransformerTile{N}}, VC <: AbstractVector{<:Number}, VS <: AbstractVector{TreeStructure{N}}} + blocks::VB + tiles::VT + coeffs::VC # every `U`, concatenated in column-major order + structs_dst::VS + structs_src::VS # with source strides permuted to the destination order + tile::TileShape{N} +end + +# the type of the `DeviceTreeTransformer`s for the space `V` on storage `A`, for type-stable caches +devicetransformertype(A, ::TensorMapSpace{S, N₁, N₂}) where {S, N₁, N₂} = + devicetransformertype(A, sectortype(S), Val(N₁ + N₂)) +function devicetransformertype(::Type{A}, ::Type{I}, ::Val{N}) where {A, I, N} + T = TensorKit.recoupling_scalartype(A, sectorscalartype(I)) + vec(::Type{X}) where {X} = Base.promote_op(similar, A, Type{X}, Dims{1}) + return DeviceTreeTransformer{N, vec(TransformerBlock{N}), vec(TransformerTile{N}), vec(T), vec(TreeStructure{N})} +end + +Adapt.adapt_structure(to, t::DeviceTreeTransformer) = DeviceTreeTransformer( + adapt(to, t.blocks), adapt(to, t.tiles), adapt(to, t.coeffs), adapt(to, t.structs_dst), + adapt(to, t.structs_src), t.tile +) + +# Tile extents are doubled greedily while the estimated cost decreases: 32-byte sectors read +# (contiguous along `q₁`) and written (along index 1), one extra per contiguous run, a fixed +# overhead per tile and the idle threads of partially filled tiles. +const MAX_TILE_LENGTH = 256 +const MIN_TILE_LENGTH = 32 # tiles should at least fill a warp +const MAX_GROUPSIZE = 128 + +function _tile_cost(tile, q₁, sizes, elsize) + sectors(len) = cld(len * elsize, 32) + 1 + cost = 0.0 + for sz in sizes + ext = min.(tile, sz) + n = prod(ext) + reads = sectors(ext[q₁]) * (n ÷ ext[q₁]) + writes = sectors(ext[1]) * (n ÷ ext[1]) + cost += prod(cld.(sz, tile)) * (64 + reads + writes + prod(tile) / 32) + end + return cost +end + +function _choose_tile(q₁, sizes, elsize, ::Val{N}) where {N} + N == 0 && return () + # at least two entries along the first index bounds the shared memory padding + tile = ntuple(n -> n == 1 ? 2 : 1, Val(N)) + cost = _tile_cost(tile, q₁, sizes, elsize) + while prod(tile) < MAX_TILE_LENGTH + candidates = ntuple(d -> Base.setindex(tile, 2 * tile[d], d), Val(N)) + costs = map(t -> _tile_cost(t, q₁, sizes, elsize), candidates) + newcost, d = findmin(costs) + newcost < cost || prod(tile) < MIN_TILE_LENGTH || break + tile, cost = candidates[d], newcost + end + return tile +end + +# `entries` iterates `(U, structures_dst, structures_src)`, with `structures_src` in source order +function DeviceTreeTransformer(::Type{T}, entries, p::NTuple{N, Int}, elsize::Int) where {T, N} + blocks = TransformerBlock{N}[] + coeffs = T[] + structs_dst = TreeStructure{N}[] + structs_src = TreeStructure{N}[] + sizes = NTuple{N, Int}[] # one per source tile that is read + for (U, structures_dst, structures_src) in entries + # all trees in a block share the same subblock size + sz = first(first(structures_dst)) + push!(blocks, TransformerBlock{N}(sz, size(U)..., length(coeffs), length(structs_dst), length(structs_src))) + append!(coeffs, U) + for (_, strides, offset) in structures_dst + push!(structs_dst, (strides, offset)) + end + for (_, strides, offset) in structures_src + push!(structs_src, (TupleTools.getindices(strides, p), offset)) + end + append!(sizes, Iterators.repeated(sz, length(U))) + end + q₁ = N == 0 ? 0 : findfirst(==(1), p) + tile = _choose_tile(q₁, sizes, elsize, Val(N)) + tiles = TransformerTile{N}[] + for (b, blk) in enumerate(blocks), i in 0:(blk.rows - 1), I in CartesianIndices(cld.(blk.sz, tile)) + push!(tiles, TransformerTile{N}(b, i, (Tuple(I) .- 1) .* tile)) + end + # the work of a tile scales with the number of source trees, so schedule the heaviest first + sort!(tiles; by = t -> blocks[t.block].cols, rev = true) + return DeviceTreeTransformer(blocks, tiles, coeffs, structs_dst, structs_src, TileShape(tile, p)) +end + +function _tree_entries(transformer::UniqueTreeTransformer) + (; structure_dst, structure_src) = transformer + return ( + (fill(coeff, 1, 1), (structure_dst[idst],), (structure_src[isrc],)) + for (coeff, idst, isrc) in transformer.data + ) +end +function _tree_entries(transformer::GenericTreeTransformer) + (; structure_dst, structure_src) = transformer + return ( + (U, view(structure_dst, inds_dst), view(structure_src, inds_src)) + for (U, inds_dst, inds_src) in transformer.data + ) +end + +function device_transformer(A, transformer::Union{UniqueTreeTransformer, GenericTreeTransformer}, p) + entries = _tree_entries(transformer) + packed = DeviceTreeTransformer(scalartype(transformer), entries, Tuple(p), sizeof(scalartype(A))) + return device_transformer(A, packed) +end + +function device_transformer(A, packed::DeviceTreeTransformer) + return GPUArrays.@uncached begin + upload = _uploader(A) + DeviceTreeTransformer( + upload(packed.blocks), upload(packed.tiles), upload(packed.coeffs), + upload(packed.structs_dst), upload(packed.structs_src), packed.tile + ) + end +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, + conjsrc::Bool, levels::IndexTuple + )::devicetransformertype(A, Vdst) where {TA <: AnyGPUVector} + host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) + return device_transformer(A, host, linearize(p)) +end + +@cached function TensorKit.treetransposer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool + )::devicetransformertype(A, Vdst) where {TA <: AnyGPUVector} + host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) + return device_transformer(A, host, linearize(p)) +end + +@cached function TensorKit.treetwister( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool + )::devicetransformertype(A, Vdst) where {TA <: AnyGPUVector, N} + host = TensorKit.treetwister(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask, inv) + return device_transformer(A, host, linearize(p)) +end + +# we don't want/need the `StridedSubblocks` here because that information is already in the transformer +TensorKit._transform_subblocks(tdst::TensorMap, tsrc::TensorMap, ::DeviceTreeTransformer) = (tdst.data, tsrc.data) + +@inline _tilecoords(e, shift, mask, origin) = map((s, m, o) -> ((e >> s) & m) + o, shift, mask, origin) +@inline _isinside(coords, sz) = all(map(<, coords, sz)) +@inline function _tileindex(coords, origin, shift) + e = Int32(0) + for n in 1:length(coords) + e += (coords[n] - origin[n]) << shift[n] + end + return e +end +@inline _shmindex(e, padshift) = e + (e >> padshift) + Int32(1) + +@inline _gather(src, op, coords, inside, (strides, offset)) = + inside ? op(@inbounds src[_linear_index(coords, strides, offset)]) : zero(eltype(src)) + +# destination indices of the entries of this thread, and `β` times their current values +@inline function _dstplan(dst, (strides, offset), es, tile, origin, sz, β, ::Val{B}) where {B} + coords = map(e -> _tilecoords(e, tile.dshift, tile.mask, origin), es) + inside = map(c -> _isinside(c, sz), coords) + inds = map(c -> _linear_index(c, strides, offset), coords) + olds = map(inside, inds) do b, i + return B && b ? β * @inbounds(dst[i]) : zero(eltype(dst)) + end + return inside, inds, olds +end + +@kernel unsafe_indices = true function tiled_transform_kernel!( + dst, @Const(src), op, transformer, α, β, ::Val{L}, ::Val{G}, hasβ::Val{B} + ) where {L, G, B} + (; blocks, tiles, coeffs, structs_dst, structs_src, tile) = transformer + shm = @localmem eltype(dst) (L + L ÷ 2,) + t = Int32(@index(Local, Linear) - 1) + job = @inbounds tiles[@index(Group, Linear)] + blk = @inbounds blocks[job.block] + i = job.tree + es = ntuple(r -> t + Int32((r - 1) * G), Val(L ÷ G)) + struct_dst = @inbounds structs_dst[blk.dst_offset + i + Int32(1)] + # with `β != 0`, the destination is read early so that this overlaps with the gather + B && (plan = _dstplan(dst, struct_dst, es, tile, job.origin, blk.sz, β, hasβ)) + + # dst_i = β * dst_i + α * Σ_j U[i, j] * permute(op(src_j), p), gathered in source order + coords = map(e -> _tilecoords(e, tile.sshift, tile.mask, job.origin), es) + inside = map(c -> _isinside(c, blk.sz), coords) + acc = map(_ -> zero(eltype(dst)), es) + for j in Int32(1):blk.cols + coeff = @inbounds coeffs[blk.u_offset + i + (j - Int32(1)) * blk.rows + Int32(1)] + struct_src = @inbounds structs_src[blk.src_offset + j] + acc = map((a, c, b) -> a + coeff * _gather(src, op, c, b, struct_src), acc, coords, inside) + end + map(acc, coords) do a, c + @inbounds shm[_shmindex(_tileindex(c, job.origin, tile.dshift), tile.padshift)] = a + nothing + end + @synchronize + + # write out in destination order + B || (plan = _dstplan(dst, struct_dst, es, tile, job.origin, blk.sz, β, hasβ)) + map(es, plan...) do e, b, id, old + b && @inbounds(dst[id] = α * shm[_shmindex(e, tile.padshift)] + old) + nothing + end +end + +# all tile lengths are powers of two in `MIN_TILE_LENGTH:MAX_TILE_LENGTH`, or 1 without indices, +# and are branched on statically, as constructing a `Val` at runtime costs a dynamic dispatch +function _launch_tiled!(dst, src, conjsrc::Bool, transformer::DeviceTreeTransformer, α, β) + isempty(transformer.tiles) && return nothing + op = conjsrc ? conj : identity + return iszero(β) ? _launch_tiled!(dst, src, op, transformer, α, β, Val(false)) : + _launch_tiled!(dst, src, op, transformer, α, β, Val(true)) +end +function _launch_tiled!(dst, src, op, transformer, α, β, hasβ::Val) + L = transformer.tile.len + L == 32 && return _launch_tiled!(dst, src, op, transformer, α, β, Val(32), hasβ) + L == 64 && return _launch_tiled!(dst, src, op, transformer, α, β, Val(64), hasβ) + L == 128 && return _launch_tiled!(dst, src, op, transformer, α, β, Val(128), hasβ) + L == 256 && return _launch_tiled!(dst, src, op, transformer, α, β, Val(256), hasβ) + L == 1 && return _launch_tiled!(dst, src, op, transformer, α, β, Val(1), hasβ) + error("unexpected tile length $L") +end +function _launch_tiled!(dst, src, op, transformer, α, β, ::Val{L}, hasβ::Val) where {L} + G = min(L, MAX_GROUPSIZE) + kernel = tiled_transform_kernel!(get_backend(dst), G) + kernel(dst, src, op, transformer, α, β, Val(L), Val(G), hasβ; ndrange = length(transformer.tiles) * G) + return nothing +end + +function TensorKit.add_transform_kernel!( + dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, + device::DeviceTreeTransformer, α, β, backend, allocator, ntasks::Int + ) + _launch_tiled!(dst, src, conjsrc, device, α, β) + return nothing +end diff --git a/ext/TensorKitGPUArraysExt/trace.jl b/ext/TensorKitGPUArraysExt/trace.jl new file mode 100644 index 000000000..e4900265d --- /dev/null +++ b/ext/TensorKitGPUArraysExt/trace.jl @@ -0,0 +1,245 @@ +# 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. + +""" + 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 + +function device_transformer(A, transformer::TraceTransformer, p, q) + packed = DeviceTraceTransformer(transformer, p, q) + return GPUArrays.@uncached begin + upload = _uploader(A) + 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 + +@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 + +TensorKit._transform_subblocks(tdst::TensorMap, tsrc::TensorMap, ::DeviceTraceTransformer) = (tdst.data, tsrc.data) + +# 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 +# won't work. +@inline function _searchblock(offsets, w) + lo, hi = 1, length(offsets) + while lo < hi + mid = (lo + hi + 1) >>> 1 + if @inbounds offsets[mid] <= w + lo = mid + else + hi = mid - 1 + end + end + return lo +end + +# Cartesian coordinates of the `w`-th (0-based) entry of a dense subblock of shape `sz`. +@inline function _coordinates(w, size::NTuple{N, Int}, dense_strides::NTuple{N, Int}) where {N} + return ntuple(n -> (w ÷ dense_strides[n]) % size[n], Val(N)) +end + +# `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 + 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 + +# 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 + ) + w = @index(Global, Linear) - 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 + +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 diff --git a/ext/TensorKitGPUArraysExt/truncation.jl b/ext/TensorKitGPUArraysExt/truncation.jl new file mode 100644 index 000000000..9396fcd24 --- /dev/null +++ b/ext/TensorKitGPUArraysExt/truncation.jl @@ -0,0 +1,77 @@ +const GPUSectorVector{T, I} = TensorKit.SectorVector{T, I, <:AnyGPUVector{T}} + +function MatrixAlgebraKit.findtruncated( + values::GPUSectorVector, strategy::MatrixAlgebraKit.TruncationByOrder + ) + I = sectortype(values) + + dims = similar(values, Base.promote_op(dim, I)) + for (c, v) in pairs(dims) + fill!(v, dim(c)) + end + + isempty(parent(values)) && return similar(values, Bool) + + perm = sortperm(parent(values); strategy.by, strategy.rev) + cumulative_dim = cumsum(Base.permute!(parent(dims), perm)) + + result = similar(values, Bool) + parent(result)[perm] .= cumulative_dim .<= strategy.howmany + return result +end + +function MatrixAlgebraKit.findtruncated( + values::GPUSectorVector, strategy::MatrixAlgebraKit.TruncationByError + ) + (isfinite(strategy.p) && strategy.p > 0) || + throw(ArgumentError(lazy"p-norm with p = $(strategy.p) is currently not supported.")) + ϵᵖmax = max(strategy.atol^strategy.p, strategy.rtol^strategy.p * norm(values, strategy.p)) + ϵᵖ = similar(values, typeof(ϵᵖmax)) + + # dimensions are all 1 so no need to account for weight + if FusionStyle(sectortype(values)) isa UniqueFusion + parent(ϵᵖ) .= abs.(parent(values)) .^ strategy.p + else + for (c, v) in pairs(values) + v′ = ϵᵖ[c] + v′ .= abs.(v) .^ strategy.p .* dim(c) + end + end + + isempty(parent(values)) && return similar(values, Bool) + + perm = sortperm(parent(values); by = abs, rev = false) + cumulative_err = cumsum(Base.permute!(parent(ϵᵖ), perm)) + + result = similar(values, Bool) + parent(result)[perm] .= cumulative_err .> ϵᵖmax + return result +end + +function MatrixAlgebraKit.findtruncated_svd(values::GPUSectorVector, strategy::S) where {S <: MatrixAlgebraKit.TruncationStrategy} + # returning a GPUSectorVector wrecks things in truncate_{co}domain + # because of scalar indexing + return Adapt.adapt(Vector, MatrixAlgebraKit.findtruncated(values, strategy)) +end + +for strat in (:(MatrixAlgebraKit.TruncationByOrder), :(MatrixAlgebraKit.TruncationByError), :(MatrixAlgebraKit.TruncationIntersection), :(MatrixAlgebraKit.TruncationMultiplets), :(TensorKit.Factorizations.TruncationSpace)) + @eval function MatrixAlgebraKit.findtruncated_svd(values::GPUSectorVector, strategy::$strat) + # returning a GPUSectorVector wrecks things in truncate_{co}domain + # because of scalar indexing + return Adapt.adapt(Vector, MatrixAlgebraKit.findtruncated(values, strategy)) + end +end + +function MatrixAlgebraKit.findtruncated_svd(values::GPUSectorVector, strategy::MatrixAlgebraKit.TruncationByValue) + atol = TensorKit.Factorizations.rtol_to_atol(values, strategy.p, strategy.atol, strategy.rtol) + strategy′ = trunctol(; atol, strategy.by, strategy.keep_below) + return SectorDict(c => Adapt.adapt(Vector, MatrixAlgebraKit.findtruncated_svd(d, strategy′)) for (c, d) in pairs(values)) +end + +function MatrixAlgebraKit.truncation_error!(values::GPUSectorVector, ind::AbstractVector{Bool}) + for (c, ind_c) in pairs(ind) + sector_vals = values[c] + @. sector_vals *= !ind_c + end + return norm(values) +end diff --git a/ext/TensorKitGPUArraysExt/utility.jl b/ext/TensorKitGPUArraysExt/utility.jl new file mode 100644 index 000000000..007e2ec1c --- /dev/null +++ b/ext/TensorKitGPUArraysExt/utility.jl @@ -0,0 +1,30 @@ +# Device-side tree transformers +# ----------------------------- +# The positions and recoupling coefficients of a `TreeTransformer` are resolved into +# sizes/strides/offsets and packed into flat device vectors. + +const TreeStructure{N} = Tuple{NTuple{N, Int}, Int} + +# strides of a dense array of shape `sz` +_dense_strides(size::Dims) = Base.front((1, cumprod(size)...)) + +function _work_offsets(work) + offsets = cumsum(work) + pushfirst!(offsets, 0) + total = pop!(offsets) + return offsets, total +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, Integer}, strides::NTuple{N, Int}, offset) where {N} + # 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 + +# uploads preserve the element type, and must outlive any GPUArrays allocation-cache scope +_uploader(A) = (proto = A(undef, 0); x -> copy!(similar(proto, eltype(x), size(x)), x)) diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 750d7384a..f0ac2b212 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -345,7 +345,7 @@ for V in spacelist twistmask = (true, false, true, false, true) p = (codomainind(t), domainind(t)) twister = TensorKit.treetwister(storagetype(t), space(t), space(t), p, false, twistmask, false) - @test twister isa getglobal(GPUArraysExt, :DeviceUniqueTreeTransformer) + @test twister isa getglobal(GPUArraysExt, :DeviceTreeTransformer) @test twister === TensorKit.treetwister(storagetype(t), space(t), space(t), p, false, twistmask, false) end end