Redo add_transform! for GPU backed Tensors to be much, much faster

ext/TensorKitGPUArraysExt.jl

@@ -3,14 +3,15 @@ module TensorKitGPUArraysExt
 using GPUArrays
 using GPUArrays: @allowscalar
 using GPUArrays.KernelAbstractions: @kernel, @index, get_backend
-
+using Adapt
 using Strided: StridedViews
 using MatrixAlgebraKit, Adapt
 using TensorKit
+using TensorKit.TensorOperations: linearize
 using TensorKit.Factorizations
 using TensorKit.Factorizations: AbstractAlgorithm
 using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check
-import TensorKit: randisometry, rand, randn, fill_braidingsubblock!
+import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel!
 
 function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:AnyGPUMatrix{T}, <:StridedViews.StridedView{T, 4, <:AnyGPUArray{T}}}}
     # COV_EXCL_START
@@ -121,5 +122,339 @@ function TensorKit.scalar(t::TensorMap{T, S, 0, 0, <:AnyGPUArray}) where {T, S}
     return isempty(inds) ? zero(scalartype(t)) : @allowscalar @inbounds t.data[only(inds)]
 end
 
+# Device-side tree transformers
+# -----------------------------
+# `AbelianTransformerData` is `isbits`, but `GenericTransformerData` isn't.
+# `Matrix{T}` and two ragged `Vector`s, so we can't just call `adapt` on it.
+# But we can work around this but packing all the information on the CPU side
+# into dense vectors of numbers, plus some accounting information so we know how
+# to unpack in the generic 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 subblock 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}
+
+"""
+    AbelianTransformerBlock{T, N}
+
+Isbits descriptor for a subblock that is a single scaled permutation: either an entry of
+an `AbelianTreeTransformer`, or a degenerate (one-tree) block of a
+`GenericTreeTransformer`.
+"""
+struct AbelianTransformerBlock{T, N}
+    coeff::T
+    sz::NTuple{N, Int}
+    densestrides::NTuple{N, Int}
+    st_dst::NTuple{N, Int}
+    offs_dst::Int
+    pst_src::NTuple{N, Int}  # source strides, permuted by `p`
+    offs_src::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`.
+** All offsets are 0-based since it makes the arithmetic easier. **
+"""
+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
+
+struct DeviceAbelianTreeTransformer{VB <: AbstractVector{<:AbelianTransformerBlock}, VO <: AbstractVector{Int}}
+    blocks::VB
+    work_offsets::VO
+    nwork::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 <: DeviceAbelianTreeTransformer{<:Any, VO}, VB <: AbstractVector{<:GenericTransformerBlock}, VC <: AbstractVector{<:Number}, VS <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}}
+    degenerate::DA  # length(U) = 1 blocks, can be handled by Abelian 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`
+_densestrides(sz::NTuple{N, Int}) where {N} = ntuple(n -> prod(sz[1:(n - 1)]; init = 1), Val(N))
+_permutestrides(st::NTuple{N, Int}, p) where {N} = ntuple(n -> st[p[n]], Val(N))
+
+# `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 _abelian_block(
+        coeff::T, (sz_dst, st_dst, offs_dst), (_, st_src, offs_src), p
+    ) where {T}
+    return AbelianTransformerBlock{T, length(sz_dst)}(
+        coeff, sz_dst, _densestrides(sz_dst), st_dst, offs_dst,
+        _permutestrides(st_src, p), offs_src
+    )
+end
+
+function _work_offsets(work)
+    offsets = cumsum(work)
+    pushfirst!(offsets, 0)
+    total = pop!(offsets)
+    return offsets, total
+end
+
+function DeviceAbelianTreeTransformer(
+        transformer::TensorKit.AbelianTreeTransformer{T, N}, p
+    ) where {T, N}
+    blocks = AbelianTransformerBlock{T, N}[_abelian_block(entry..., p) for entry in transformer.data]
+    work_offsets, nwork = _work_offsets(prod(blk.sz) for blk in blocks)
+    return DeviceAbelianTreeTransformer(blocks, work_offsets, nwork)
+end
+
+function DeviceGenericTreeTransformer(
+        transformer::TensorKit.GenericTreeTransformer{T, N}, p
+    ) where {T, N}
+    degenerate = AbelianTransformerBlock{T, N}[]
+    blocks = GenericTransformerBlock{N}[]
+    coeffs = T[]
+    structs_dst = TreeStructure{N}[]
+    structs_src = TreeStructure{N}[]
+
+    for (U, (sz_dst, sts_dst), (sz_src, sts_src)) in transformer.data
+        if length(U) == 1 # same as the Abelian case
+            push!(
+                degenerate, _abelian_block(
+                    only(U), (sz_dst, only(sts_dst)...), (sz_src, only(sts_src)...), p
+                )
+            )
+        else
+            push!(
+                blocks, GenericTransformerBlock{N}(
+                    sz_dst, _densestrides(sz_dst), size(U, 1), size(U, 2),
+                    length(coeffs), length(structs_dst), length(structs_src)
+                )
+            )
+            append!(coeffs, U)
+            append!(structs_dst, sts_dst)
+            for (st_src, offs_src) in sts_src
+                push!(structs_src, (_permutestrides(st_src, p), offs_src))
+            end
+        end
+    end
+
+    deg_offsets, deg_nwork = _work_offsets(prod(blk.sz) for blk in degenerate)
+    work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks)
+    return DeviceGenericTreeTransformer(
+        DeviceAbelianTreeTransformer(degenerate, deg_offsets, deg_nwork),
+        blocks, work_offsets, nwork, coeffs, structs_dst, structs_src
+    )
+end
+
+"""
+    StorageAdaptor(proto)
+
+`Adapt` adaptor moving arrays onto the same device and array type as `proto`, preserving
+their element type. `adapt(CuVector{Float64}, ::Vector{Int})` would force-convert the Int
+to Float64, while `similar(proto, Int, n)` doesn't.
+"""
+struct StorageAdaptor{A <: AbstractArray}
+    proto::A
+end
+function Adapt.adapt_storage(a::StorageAdaptor, x::AbstractArray)
+    dst = similar(a.proto, eltype(x), size(x))
+    isempty(x) && return dst
+    return copyto!(dst, x)
+end
+
+function Adapt.adapt_structure(to, t::DeviceAbelianTreeTransformer)
+    return DeviceAbelianTreeTransformer(
+        Adapt.adapt(to, t.blocks), Adapt.adapt(to, t.work_offsets), t.nwork
+    )
+end
+function Adapt.adapt_structure(to, t::DeviceGenericTreeTransformer)
+    return DeviceGenericTreeTransformer(
+        Adapt.adapt(to, t.degenerate), Adapt.adapt(to, t.blocks),
+        Adapt.adapt(to, t.work_offsets), t.nwork, Adapt.adapt(to, t.coeffs),
+        Adapt.adapt(to, t.structs_dst), Adapt.adapt(to, t.structs_src)
+    )
+end
+
+# Copying a transformer to GPU is more expensive than running it, and transformers are
+# themselves cached (and thus long-lived) by `treebraider`/`treetransposer`, so we cache the
+# device copy for as long as the CPU original is "alive". The key is:
+# - `transformer.data`
+# - the storage type
+# - `p`, which is baked into the permuted source strides.
+# Using `objectid` avoids walking every recoupling matrix on every lookup.
+# TODO: should this live in the main package?
+const DEVICE_TRANSFORMER_CACHE = Dict{UInt, Tuple{WeakRef, Dict{Any, Any}}}()
+const DEVICE_TRANSFORMER_LOCK = ReentrantLock()
+
+# We have this complicated setup because a naive `adapt` doesn't work.
+# Rather we copy everything to GPU-native arrays and have kernels that can work
+# with that.
+function device_transformer(proto::AbstractArray, transformer, p)
+    key = transformer.data
+    return Base.@lock DEVICE_TRANSFORMER_LOCK begin
+        entry = get(DEVICE_TRANSFORMER_CACHE, objectid(key), nothing)
+        if isnothing(entry) || entry[1].value !== key
+            # misses are rare (once per transformer), so pay for a sweep of the entries
+            # whose host transformer has been collected and free their device copies
+            filter!(kv -> !isnothing(last(kv)[1].value), DEVICE_TRANSFORMER_CACHE)
+            entry = (WeakRef(key), Dict{Any, Any}())
+            DEVICE_TRANSFORMER_CACHE[objectid(key)] = entry
+        end
+        get!(last(entry), (typeof(proto), p)) do
+            Adapt.adapt(StorageAdaptor(proto), _device_transformer(transformer, p))
+        end
+    end
+end
+
+_device_transformer(t::TensorKit.AbelianTreeTransformer, p) = DeviceAbelianTreeTransformer(t, p)
+_device_transformer(t::TensorKit.GenericTreeTransformer, p) = DeviceGenericTreeTransformer(t, p)
+
+# 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.
+@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`.
+# Computed once per thread and then reused for every strided view of that subblock.
+# This avoids `StridedView` redoing these divisions on every single element access.
+# Integer division on GPU is usually pretty slow.
+@inline function _coordinates(w, sz::NTuple{N, Int}, densestrides::NTuple{N, Int}) where {N}
+    return ntuple(n -> (w ÷ densestrides[n]) % sz[n], Val(N))
+end
+
+# finds the overall offset in the output and input arrays corresponding to the **sublock**
+# coordinates currently being worked on
+@inline function _offset(coords::NTuple{N, Int}, st::NTuple{N, Int}, offs) where {N}
+    return offs + sum(ntuple(n -> coords[n] * st[n], Val(N))) + 1
+end
+
+# One thread per destination element in `data_dst`.
+@kernel function abelian_batched_permute!(
+        data_dst, data_src, 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.densestrides)
+        i_dst = _offset(coords, blk.st_dst, blk.offs_dst)
+        i_src = _offset(coords, blk.pst_src, blk.offs_src)
+        @inbounds data_dst[i_dst] = α * blk.coeff * 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.
+@kernel function generic_batched_permute!(
+        data_dst, data_src, 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 = _offset(coords, st_dst, offs_dst)
+
+        # dst_i = β * dst_i + α * Σ_j U[i, j] * permute(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)))
+        for j in 1:blk.cols
+            # TODO is there a more efficient way to do this read?
+            coeff = @inbounds coeffs[blk.u_offset + (j - 1) * blk.rows + i + 1]
+            iszero(coeff) && continue
+            pst_src, offs_src = @inbounds structs_src[blk.src_offset + j]
+            acc += coeff * @inbounds data_src[_offset(coords, pst_src, offs_src)]
+        end
+        @inbounds data_dst[i_dst] = α * acc + β * data_dst[i_dst]
+    end
+end
+# COV_EXCL_STOP
+
+function _launch_abelian!(data_dst, data_src, transformer, α, β, ::Val{N}) where {N}
+    nwork = transformer.nwork
+    nwork == 0 && return nothing
+    abelian_batched_permute!(get_backend(data_dst))(
+        data_dst, data_src, transformer.blocks, transformer.work_offsets, α, β, nwork,
+        Val(N); ndrange = nwork
+    )
+    return nothing
+end
+
+function _launch_generic!(data_dst, data_src, transformer, α, β, ::Val{N}) where {N}
+    nwork = transformer.nwork
+    nwork == 0 && return nothing
+    generic_batched_permute!(get_backend(data_dst))(
+        data_dst, data_src, transformer.blocks, transformer.work_offsets,
+        transformer.coeffs, transformer.structs_dst, transformer.structs_src, α, β, nwork,
+        Val(N); ndrange = nwork
+    )
+    return nothing
+end
+
+function TensorKit.add_transform_kernel!(
+        data_dst::A, data_src::A, p, transformer::TensorKit.AbelianTreeTransformer{T, N},
+        α, β, backend, allocator, scheduler
+    ) where {T, N, A <: AnyGPUArray}
+    # GPU-side object to hold the treetransformer information
+    device = device_transformer(data_dst, transformer, linearize(p))::DeviceAbelianTreeTransformer
+    _launch_abelian!(data_dst, data_src, device, α, β, Val(N))
+    return nothing
+end
+
+function TensorKit.add_transform_kernel!(
+        data_dst::A, data_src::A, p, transformer::TensorKit.GenericTreeTransformer{T, N},
+        α, β, backend, allocator, scheduler
+    ) where {T, N, A <: AnyGPUArray}
+    # GPU-side object to hold the treetransformer information
+    device = device_transformer(data_dst, transformer, linearize(p))::DeviceGenericTreeTransformer
+    # one-tree blocks are a scaled permutation, which the Abelian kernel already handles; the
+    # two kernels touch disjoint subblocks so the launch order does not matter
+    _launch_abelian!(data_dst, data_src, device.degenerate, α, β, Val(N))
+    _launch_generic!(data_dst, data_src, device, α, β, Val(N))
+    return nothing
+end
 
 end