storagetype-adapted treetransformers

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

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
 

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,19 +129,19 @@ 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)
     )
     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(