diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 5d1d68c09..2a853124a 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -28,6 +28,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Extend VectorInterface compat to include v0.7 - Require TensorKitSectors v0.3.10 +- `@cached` can be used from other modules, such as package extensions, which then own the cache of the methods they add; these caches are shown with their module in `global_cache_info` after registration in `__init__` (internal) +- The `TreeTransformer`s used in index manipulations are now constructed and cached per storagetype of the destination tensor, which is also a dispatch point for storage-specific transformers with their own cache (internal) ### Deprecated @@ -39,6 +41,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Performance +- Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. + ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 ### Added diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 7843460c3..948f9fa75 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -4,7 +4,6 @@ using GPUArrays using GPUArrays: @allowscalar using GPUArrays.KernelAbstractions: @kernel, @index, get_backend using Adapt -using TensorKit.LRUCache: LRU using TensorKit.TupleTools using Strided: StridedViews using MatrixAlgebraKit, Adapt @@ -13,7 +12,7 @@ using TensorKit.TensorOperations: linearize, DefaultAllocator using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check -using TensorKit: StridedSubblocks, UniqueTreeTransformer, GenericTreeTransformer +using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple 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}}}} @@ -272,70 +271,66 @@ function DeviceGenericTreeTransformer( ) end -""" - StorageAdaptor(proto) - -`Adapt` adaptor moving arrays onto the same device and array type as `proto`, preserving -their element type. For `proto::CuVector{Float64}` and `array::Vector{Int}`, -the call `adapt(typeof(proto), array)` would force-convert the element type `Int` -to `Float64`, while `adapt(StoreAdaptor(proto), array)` does not. -""" -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 copy!(dst, x) +# 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 -function Adapt.adapt_structure(to, t::DeviceUniqueTreeTransformer) - return DeviceUniqueTreeTransformer( - Adapt.adapt(to, t.blocks), Adapt.adapt(to, t.work_offsets), t.nwork - ) +@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 -function Adapt.adapt_structure(to, t::DeviceGenericTreeTransformer) - return DeviceGenericTreeTransformer( - Adapt.adapt(to, t.unique_blocks), 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) +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, + ) ) -end -# Copying a transformer to GPU is more expensive than running it, so we cache the device -# copy in a global LRU cache, registered in `TensorKit.GLOBAL_CACHES` so that -# `empty_globalcaches!` also frees the device memory. The key is: -# - transformer which is an immutable struct. It's hashed and compared by the identity of -# its fields, which avoids walking every recoupling matrix on every lookup. -# Holding it in the key also keeps it alive, so its `objectid` cannot be -# reused by a different transformer while the entry is cached. -# - the storage type -# - `p`, which is baked into the permuted source strides. -# TODO: should this live in the main package? -const DEVICE_TRANSFORMER_CACHE = LRU{Any, Any}(; maxsize = TensorKit.DEFAULT_GLOBALCACHE_SIZE[]) + # device-caches should probably have smaller sizes: + resize!(GLOBAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) -function __init__() - push!(TensorKit.GLOBAL_CACHES, :DEVICE_TRANSFORMER_CACHE => DEVICE_TRANSFORMER_CACHE) return nothing end -# 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, typeof(proto), p) - return get!(DEVICE_TRANSFORMER_CACHE, key) do - # be careful about the lifetime of these, since they live in a global cache and - # thus persist beyond the call - GPUArrays.@uncached Adapt.adapt( - StorageAdaptor(proto), _device_transformer(transformer, p) - ) +# Allocations must outlive any GPUArrays allocation-cache scope around the operation. +function device_transformer(A, transformer::UniqueTreeTransformer, p) + packed = DeviceUniqueTreeTransformer(transformer, p) + 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 -_device_transformer(t::UniqueTreeTransformer, p) = DeviceUniqueTreeTransformer(t, p) -_device_transformer(t::GenericTreeTransformer, p) = DeviceGenericTreeTransformer(t, p) +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 # COV_EXCL_START # kernels are not reachable by coverage @@ -373,7 +368,9 @@ end @kernel function unique_batched_permute_kernel!( data_dst, data_src, op, blocks, work_offsets, α, β, nwork, ::Val{N} ) where {N} - w = @index(Global, Linear) - 1 + # w = @index(Global, Linear) - 1 does not work due to JLArrays bug + w = @index(Global, Linear) + w -= 1 if w < nwork b = _searchblock(work_offsets, w) blk = @inbounds blocks[b] @@ -398,7 +395,9 @@ end data_dst, data_src, op, blocks, work_offsets, coeffs, structs_dst, structs_src, α, β, nwork, ::Val{N} ) where {N} - w = @index(Global, Linear) - 1 + # w = @index(Global, Linear) - 1 does not work due to JLArrays bug + w = @index(Global, Linear) + w -= 1 if w < nwork # bookkeeping to figure out where to read from and write to b = _searchblock(work_offsets, w) @@ -448,30 +447,32 @@ function _launch_generic!(data_dst, data_src, op, transformer, α, β, ::Val{N}) return nothing end -const GPUStridedSubblocks = StridedSubblocks{<:AnyGPUArray} +# 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} +) = (tdst.data, tsrc.data) function TensorKit.add_transform_kernel!( - dst::GPUStridedSubblocks, src::GPUStridedSubblocks, p, conjsrc::Bool, - transformer::UniqueTreeTransformer{T, N}, α, β, backend, allocator, ntasks::Int - ) where {T, N} - # GPU-side object to hold the treetransformer information - device = device_transformer(dst.data, transformer, linearize(p)) + dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool, + device::DeviceUniqueTreeTransformer, α, β, backend, allocator, ntasks::Int + ) + N = length(linearize(p)) op = conjsrc ? conj : identity - _launch_unique!(dst.data, src.data, op, device, α, β, Val(N)) + _launch_unique!(dst, src, op, device, α, β, Val(N)) return nothing end function TensorKit.add_transform_kernel!( - dst::GPUStridedSubblocks, src::GPUStridedSubblocks, p, conjsrc::Bool, - transformer::GenericTreeTransformer{T, N}, α, β, backend, allocator, ntasks::Int - ) where {T, N} - # GPU-side object to hold the treetransformer information - device = device_transformer(dst.data, transformer, linearize(p)) + 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.data, src.data, op, device.unique_blocks, α, β, Val(N)) - _launch_generic!(dst.data, src.data, op, device, α, β, Val(N)) + _launch_unique!(dst, src, op, device.unique_blocks, α, β, Val(N)) + _launch_generic!(dst, src, op, device, α, β, Val(N)) return nothing end diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index f3bec7844..bdc07ca08 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -54,9 +54,37 @@ function _cached_category(fname::Symbol) "symmetry" : "bookkeeping" end +# timer sections are controlled by the `timeit_debug_enabled` switch of the module using `@cached` +_timeit_expr(label, ex) = + Expr(:macrocall, GlobalRef(TimerOutputs, Symbol("@timeit_debug")), LineNumberNode(@__LINE__, @__FILE__), GlobalRef(@__MODULE__, :GLOBAL_TIMER), label, ex) + macro cached(ex) - Meta.isexpr(ex, :function) || - error("cached macro can only be used on function definitions") + return _cached(__module__, ex) +end + +function _cached_error(msg, ex) + ex isa Expr && (ex = Base.remove_linenums!(copy(ex))) + Meta.isexpr(ex, :kw) && (ex = Expr(:(=), ex.args...)) # prints as `x = 1` + throw(ArgumentError("@cached: $msg, got `$ex`")) +end + +# the arguments form the cache key, so they must be positional; unnamed ones get a name +function _cached_canonicalize_arg(arg) + Meta.isexpr(arg, :parameters) && _cached_error("keyword arguments are not supported", arg) + Meta.isexpr(arg, :kw) && _cached_error("default values are not supported", arg) + Meta.isexpr(arg, :...) && _cached_error("varargs are not supported", arg) + arg isa Symbol && return arg + Meta.isexpr(arg, :(::)) || _cached_error("unsupported argument", arg) + length(arg.args) == 1 && return Expr(:(::), gensym(:arg), arg.args[1]) + return arg +end + +function _cached(mod::Module, ex) + if !Meta.isexpr(ex, :function) + Meta.isexpr(ex, :(=)) && Meta.isexpr(ex.args[1], (:call, :where, :(::))) && + _cached_error("short-form function definitions are not supported, use `function ... end`", ex.args[1]) + _cached_error("expected a function definition", ex) + end fcall = ex.args[1] if Meta.isexpr(fcall, :where) hasparams = true @@ -72,24 +100,22 @@ macro cached(ex) else typed = false end - Meta.isexpr(fcall, :call) || - error("cached macro can only be used on function definitions") + Meta.isexpr(fcall, :call) || _cached_error("expected a function signature", fcall) fname = fcall.args[1] + # qualified names such as `TensorKit.treebraider` add methods to a function of another module + basename = Meta.isexpr(fname, :.) ? fname.args[end].value : fname + basename isa Symbol || + _cached_error("expected a function name or a qualified name such as `Module.f`", fname) + fargs = map(_cached_canonicalize_arg, fcall.args[2:end]) + fcall = Expr(:call, fname, fargs...) + fargnames = map(arg -> arg isa Symbol ? arg : arg.args[1], fargs) # timer labels for the cache lookup and the miss-path construction - lookuplabel = string("bookkeeping: cache ", fname) - misslabel = string(_cached_category(fname), ": compute ", fname) - fargs = fcall.args[2:end] - fargnames = map(fargs) do arg - if Meta.isexpr(arg, :(::)) - return arg.args[1] - else - return arg - end - end + lookuplabel = string("bookkeeping: cache ", basename) + misslabel = string(_cached_category(basename), ": compute ", basename) _fbody = ex.args[2] # actual implenetation, with underscore name - _fname = Symbol(:_, fname) + _fname = Symbol(:_, basename) _fcall = Expr(:call, _fname, fargs...) if hasparams _fcall = Expr(:where, _fcall, params...) @@ -103,7 +129,7 @@ macro cached(ex) end cachestylevar = gensym(:cachestyle) cachestyleex = Expr( - :(=), cachestylevar, Expr(:call, :CacheStyle, fname, fargnames...) + :(=), cachestylevar, Expr(:call, CacheStyle, fname, fargnames...) ) newfbody = Expr( :block, cachestyleex, Expr(:call, fname, fargnames..., cachestylevar) @@ -111,11 +137,11 @@ macro cached(ex) newfex = Expr(:function, newfcall, newfbody) # nocache implementation - fnocachecall = Expr(:call, fname, fargs..., :(::NoCache)) + fnocachecall = Expr(:call, fname, fargs..., :(::$NoCache)) if hasparams fnocachecall = Expr(:where, fnocachecall, params...) end - fnocachebody = :(@timeit_debug GLOBAL_TIMER $misslabel $(Expr(:call, _fname, fargnames...))) + fnocachebody = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) if typed T = gensym(:T) fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T)) @@ -124,16 +150,17 @@ macro cached(ex) # tasklocal cache implementation Dvar = gensym(:D) - flocalcachecall = Expr(:call, fname, fargs..., :(::TaskLocalCache{$Dvar})) + flocalcachecall = Expr(:call, fname, fargs..., :(::$TaskLocalCache{$Dvar})) if hasparams flocalcachecall = Expr(:where, flocalcachecall, params..., Dvar) else flocalcachecall = Expr(:where, flocalcachecall, Dvar) end - localcachename = Symbol(:_tasklocal_, fname, :_cache) + globalcachename = Symbol(:GLOBAL_, uppercase(string(basename)), :_CACHE) + localcachename = Symbol(:_tasklocal_, globalcachename) cachevar = gensym(:cache) getlocalcacheex = :( - $cachevar::$Dvar = get!(task_local_storage(), $localcachename) do + $cachevar::$Dvar = get!(task_local_storage(), $(QuoteNode(localcachename))) do return $Dvar() end ) @@ -143,11 +170,8 @@ macro cached(ex) else key = Expr(:tuple, fargnames...) end - getvalex = :( - @timeit_debug GLOBAL_TIMER $lookuplabel get!($cachevar, $key) do - return @timeit_debug GLOBAL_TIMER $misslabel $_fname($(fargnames...)) - end - ) + missex = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) + getvalex = _timeit_expr(lookuplabel, :(get!(() -> $missex, $cachevar, $key))) if typed T = gensym(:T) flocalcachebody = Expr( @@ -168,11 +192,10 @@ macro cached(ex) flocalcacheex = Expr(:function, flocalcachecall, flocalcachebody) # # global cache implementation - fglobalcachecall = Expr(:call, fname, fargs..., :(::GlobalLRUCache)) + fglobalcachecall = Expr(:call, fname, fargs..., :(::$GlobalLRUCache)) if hasparams fglobalcachecall = Expr(:where, fglobalcachecall, params...) end - globalcachename = Symbol(:GLOBAL_, uppercase(string(fname)), :_CACHE) getglobalcachex = Expr(:(=), cachevar, globalcachename) if typed T = gensym(:T) @@ -194,11 +217,14 @@ macro cached(ex) fglobalcacheex = Expr(:function, fglobalcachecall, fglobalcachebody) fglobalcachedef = Expr( :const, - Expr(:(=), globalcachename, :(LRU{Any, Any}(; maxsize = DEFAULT_GLOBALCACHE_SIZE[]))) - ) - fglobalcacheregister = Expr( - :call, :push!, :GLOBAL_CACHES, :($(QuoteNode(globalcachename)) => $globalcachename) + Expr(:(=), globalcachename, :($LRU{Any, Any}(; maxsize = $DEFAULT_GLOBALCACHE_SIZE[]))) ) + # Other modules register their caches explicitly from `__init__`, since + # mutations of this registry during their precompilation do not survive loading. + fglobalcacheregister = mod === (@__MODULE__) ? Expr( + :call, :push!, GlobalRef(@__MODULE__, :GLOBAL_CACHES), + :($(QuoteNode(globalcachename)) => $globalcachename) + ) : nothing # # total expression return esc( diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 106f020a4..61a4a6a12 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -705,21 +705,44 @@ function _add_transform_block!( ) end - # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear + # 2. Recoupling: buffer_dst = buffer_src * U^T (each output tree is a linear # combination of input trees weighted by the recoupling coefficients). - @timeit_debug GLOBAL_TIMER "dense: recouple mul!" begin - U′ = _adapt_recoupling(storagetype(dst), U) - mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) - end + @timeit_debug GLOBAL_TIMER "dense: recouple mul!" _recouple!(buffer, buffer_dst, buffer_src, U) # 3. Insert: scatter column j of buffer_dst into the destination, applying the - # actual index permutation p in the same tensoradd! call. + # actual index permutation p and the scaling α in the same tensoradd! call. @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (j, idst) in enumerate(inds_dst) TO.tensoradd!( dst[idst], sreshape(view(buffer_dst, :, j), sz_src), - p, false, One(), β, backend, allocator + p, false, α, β, backend, allocator ) end end return nothing end + +# computes `buffer_dst = buffer_src * transpose(U)`, where both are column-major views into `buffer` +_recouple!(buffer, buffer_dst, buffer_src, U) = + (mul!(buffer_dst, buffer_src, transpose(StridedView(U))); nothing) +# real coefficients acting on complex data: recouple the real and imaginary parts in a single real gemm +function _recouple!( + buffer::DenseVector{Complex{R}}, buffer_dst::StridedView, buffer_src::StridedView, U::AbstractMatrix{R} + ) where {R <: LinearAlgebra.BlasReal} + rbuffer = reinterpret(R, buffer) + mul!(_realview(rbuffer, buffer_dst), _realview(rbuffer, buffer_src), transpose(U)) + return nothing +end +# on the CPU, views of reinterpreted arrays are not `StridedMatrix`, so `mul!` would not use BLAS +function _recouple!( + buffer::CPUStorage{Complex{R}}, buffer_dst::StridedView, buffer_src::StridedView, U::Matrix{R} + ) where {R <: LinearAlgebra.BlasReal} + rbuffer = reinterpret(R, buffer) + LinearAlgebra.BLAS.gemm!('N', 'T', one(R), _realview(rbuffer, buffer_src), U, zero(R), _realview(rbuffer, buffer_dst)) + return nothing +end +# real view of the complex column-major view `b` into the reinterpreted buffer `rbuffer` +function _realview(rbuffer::AbstractVector, b::StridedView) + rows, cols = size(b) + offset = 2 * b.offset + return reshape(view(rbuffer, (offset + 1):(offset + 2 * rows * cols)), 2 * rows, cols) +end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 231171da8..990e50d2d 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -12,6 +12,27 @@ of `tsrc` itself; when `conjsrc` is `true` the fusion trees that are transformed """ abstract type TreeTransformer end +# `StridedSubblocks` report their storage as the `StridedView` parent type, which is `Memory` there +@static if isdefined(Core, :Memory) + const CPUStorage{T} = Union{Array{T}, Memory{T}} +else + const CPUStorage{T} = Array{T} +end + +""" + recoupling_scalartype(A::Type{<:AbstractVector}, Tₛ::Type{<:Number}) -> Type{<:Number} + +Scalar type used to store the recoupling coefficients with sector scalar type `Tₛ` in the +transformers for destination tensors with storagetype `A`. For storage with BLAS scalars, this is +the precision of the storage, where real coefficients are kept real also for complex storage, such +that they can be applied to the real and imaginary parts at once. +""" +function recoupling_scalartype(::Type{A}, ::Type{Tₛ}) where {A, Tₛ <: Number} + T = scalartype(A) + T <: BlasFloat || return Tₛ + return Tₛ <: Real ? real(T) : complex(T) +end + # (coefficient, destination position, source position) const UniqueTransformerData{T} = Tuple{T, Int, Int} @@ -29,7 +50,7 @@ struct UniqueTreeTransformer{T, N} <: TreeTransformer structure_src::Vector{StridedStructure{N}} end -# (recoupling matrix, destination positions, source positions): U[j, i] maps source i onto destination j +# (recoupling matrix, destination positions, source positions) const GenericTransformerData{T} = Tuple{Matrix{T}, Vector{Int}, Vector{Int}} """ @@ -37,9 +58,9 @@ const GenericTransformerData{T} = Tuple{Matrix{T}, Vector{Int}, Vector{Int}} Tree transformation for sectors with multiple fusion channels, where the subblocks of a [`FusionTreeBlock`](@ref) map onto the subblocks of the transformed block through a recoupling -matrix, stored as `(U, inds_dst, inds_src)`. The subblock structures of the destination and -source spaces are kept alongside, such that the [`StridedSubblocks`](@ref) of both tensors can be -created without further lookups. +matrix, stored as `(U, inds_dst, inds_src)`, where `U[j, i]` maps source `i` onto destination `j`. +The subblock structures of the destination and source spaces are kept alongside, such that the +[`StridedSubblocks`](@ref) of both tensors can be created without further lookups. """ struct GenericTreeTransformer{T, N} <: TreeTransformer data::Vector{GenericTransformerData{T}} @@ -47,19 +68,19 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer structure_src::Vector{StridedStructure{N}} end -function UniqueTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) +function UniqueTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A} t₀ = Base.time() spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) - T = sectorscalartype(sectortype(Vdst)) + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) data = Vector{UniqueTransformerData{T}}(undef, length(src_trees)) @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (isrc, (f₁, f₂)) in enumerate(src_trees) f_dst, coeff = transform(conjsrc ? (f₂, f₁) : (f₁, f₂)) _, (_, idst) = gettoken(dst_trees, f_dst) - data[isrc] = (coeff, idst, isrc) + data[isrc] = (convert(T, coeff), idst, isrc) end structure_dst = degeneracystructure(Vdst).subblockstructure @@ -72,7 +93,7 @@ function UniqueTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) return transformer end -function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) +function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A} t₀ = Base.time() spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) # the fusion blocks that are transformed are those of the adjoint space for a conjugated source @@ -80,7 +101,7 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) structure_dst = degeneracystructure(Vdst).subblockstructure structure_src = degeneracystructure(Vsrc).subblockstructure - T = sectorscalartype(sectortype(Vsrc)) + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vsrc))) fblocks = @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" fusionblocks(Vsrc′) nblocks = length(fblocks) @@ -103,7 +124,7 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) return idst end end - data[i] = (U, inds_dst, inds_src) + data[i] = (convert(Matrix{T}, U), inds_dst, inds_src) # cost model: L input blocks each going to L output blocks of a given length weights[i] = length(U) * prod(structure_dst[first(inds_dst)][1]) @@ -116,7 +137,7 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc::Bool) # sort by (approximate) weight to facilitate multi-threading strategies @timeit_debug GLOBAL_TIMER "bookkeeping: sort" Base.permute!(data, sortperm(weights; rev = true)) - transformer = GenericTreeTransformer(data, structure_dst, structure_src) + transformer = GenericTreeTransformer{T, numind(Vdst)}(data, structure_dst, structure_src) Δt = Base.time() - t₀ @debug( @@ -145,58 +166,48 @@ function buffersize(transformer::GenericTreeTransformer) end end -function treetransformertype(Vdst, Vsrc) +function treetransformertype(::Type{A}, Vdst, Vsrc) where {A} I = sectortype(Vdst) - T = sectorscalartype(I) + T = recoupling_scalartype(A, sectorscalartype(I)) N = numind(Vdst) return FusionStyle(I) == UniqueFusion() ? UniqueTreeTransformer{T, N} : GenericTreeTransformer{T, N} end function TreeTransformer( - transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S}, conjsrc::Bool - ) where {S} + ::Type{A}, transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S}, conjsrc::Bool + ) where {A, S} I = sectortype(Vdst) return FusionStyle(I) == UniqueFusion() ? - UniqueTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) : - GenericTreeTransformer(transform, p, Vdst, Vsrc, conjsrc) + UniqueTreeTransformer(A, transform, p, Vdst, Vsrc, conjsrc) : + GenericTreeTransformer(A, transform, p, Vdst, Vsrc, conjsrc) end # braid is special because it has levels function treebraider( tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple ) - return treebraider(space(tdst), space(tsrc), p, conjsrc, levels) + return treebraider(storagetype(tdst), space(tdst), space(tsrc), p, conjsrc, levels) end @cached function treebraider( - Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple - )::treetransformertype(Vdst, Vsrc) + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple + )::treetransformertype(A, Vdst, Vsrc) where {TA} Vsrc′, p′ = conjsrc ? (Vsrc', adjointtensorindices(Vsrc, p)) : (Vsrc, p) # levels are attached to the legs, so they follow the same relabeling as the permutation levels′ = conjsrc ? TupleTools.getindices(levels, adjointtensorindices(Vsrc′, allind(Vsrc′))) : levels levels″ = (TupleTools.getindices(levels′, codomainind(Vsrc′)), TupleTools.getindices(levels′, domainind(Vsrc′))) fusiontreebraider(f) = braid(f, p′, levels″) - return TreeTransformer(fusiontreebraider, p, Vdst, Vsrc, conjsrc) + return TreeTransformer(A, fusiontreebraider, p, Vdst, Vsrc, conjsrc) end function treetransposer(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool) - return treetransposer(space(tdst), space(tsrc), p, conjsrc) + return treetransposer(storagetype(tdst), space(tdst), space(tsrc), p, conjsrc) end @cached function treetransposer( - Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool - )::treetransformertype(Vdst, Vsrc) + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool + )::treetransformertype(A, Vdst, Vsrc) where {TA} p′ = conjsrc ? adjointtensorindices(Vsrc, p) : p fusiontreetransform(f) = transpose(f, p′) - return TreeTransformer(fusiontreetransform, p, Vdst, Vsrc, conjsrc) + return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end # default cachestyle is GlobalLRUCache - -# For CPU arrays the recoupling matrix can be used as is, also when the scalar types -# do not match, since Strided handles mixed-eltype mul! without the copy that -# Adapt.adapt would make (which additionally dispatches dynamically). Other storage -# types (e.g. GPU arrays) do require the conversion. -# TODO: transformers with dedicated storagetypes -# `StridedSubblocks` report their storage as the `StridedView` parent type, which is `Memory` there -const CPUStorage = @static isdefined(Core, :Memory) ? Union{Array, Memory} : Array -_adapt_recoupling(::Type{<:CPUStorage}, U::Matrix) = StridedView(U) -_adapt_recoupling(::Type{A}, U::Matrix) where {A} = Adapt.adapt(A, StridedView(U))