Tiled shared-memory kernel for GPU index manipulations

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

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