From 4c97f9b89e1ddf5caa1abd36487ee049eb4d3bcf Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <kslimes@gmail.com>
Date: Tue, 15 Sep 2026 17:05:14 +0200
Subject: [PATCH 01/16] KA kernel for add_transform for Abelian symmetries

---
 Project.toml                 |   3 +
 ext/TensorKitGPUArraysExt.jl | 201 ++++++++++++++++++++++++++++++++++-
 2 files changed, 202 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 83093ea5b..d0471fdab 100644
--- a/Project.toml
+++ b/Project.toml
@@ -72,3 +72,6 @@ TimerOutputs = "1"
 TupleTools = "1.5"
 VectorInterface = "0.6"
 julia = "1.10"
+
+[sources]
+TensorOperations = {url = "https://github.com/QuantumKitHub/TensorOperations.jl", rev = "main"}
diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 158890aa8..61bfb3db8 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/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, DefaultAllocator
 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,201 @@ 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
+# -----------------------------
+# We 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}
+
+"""
+    AbelianTransformerBlock{T, N}
+
+`isbits` descriptor for a subblock that is a *single* scaled permutation:
+an entry of an `AbelianTreeTransformer`.
+"""
+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
+
+# device-side simple struct that GPU kernels can use
+struct DeviceAbelianTreeTransformer{VB <: AbstractVector{<:AbelianTransformerBlock}, VO <: AbstractVector{Int}}
+    blocks::VB
+    work_offsets::VO
+    nwork::Int
+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
+
+"""
+    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
+
+# 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
+            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
+            # be careful about the lifetime of these, since they live as long as their
+            # "parent" on the CPU, so they can persist beyond the call
+            GPUArrays.@uncached Adapt.adapt(
+                StorageAdaptor(proto), _device_transformer(transformer, p)
+            )
+        end
+    end
+end
+
+_device_transformer(t::TensorKit.AbelianTreeTransformer, p) = DeviceAbelianTreeTransformer(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
+
+# 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 TensorKit.add_transform_kernel!(
+        data_dst::A, data_src::A, p, transformer::TensorKit.AbelianTreeTransformer{T, N},
+        α, β, backend, allocator, ntasks::Int 
+    ) 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
 
 end

From 309f47b63d18600c0b3eee2d5f187551225cdde0 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.s.hyatt@gmail.com>
Date: Tue, 15 Sep 2026 18:10:24 +0200
Subject: [PATCH 02/16] Formatter

---
 ext/TensorKitGPUArraysExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 61bfb3db8..4d8627d8f 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -311,7 +311,7 @@ end
 
 function TensorKit.add_transform_kernel!(
         data_dst::A, data_src::A, p, transformer::TensorKit.AbelianTreeTransformer{T, N},
-        α, β, backend, allocator, ntasks::Int 
+        α, β, backend, allocator, ntasks::Int
     ) where {T, N, A <: AnyGPUArray}
     # GPU-side object to hold the treetransformer information
     device = device_transformer(data_dst, transformer, linearize(p))::DeviceAbelianTreeTransformer

From 200ff03d8612d89544fac513ad39df5e8300d7b0 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.s.hyatt@gmail.com>
Date: Tue, 15 Sep 2026 22:53:50 +0200
Subject: [PATCH 03/16] New TO version

---
 Project.toml | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index d0471fdab..fda238c20 100644
--- a/Project.toml
+++ b/Project.toml
@@ -67,11 +67,8 @@ Random = "1"
 ScopedValues = "1.3.0"
 Strided = "2.6.1"
 TensorKitSectors = "0.3.7"
-TensorOperations = "5.5.2, 5.6"
+TensorOperations = "5.8.1"
 TimerOutputs = "1"
 TupleTools = "1.5"
 VectorInterface = "0.6"
 julia = "1.10"
-
-[sources]
-TensorOperations = {url = "https://github.com/QuantumKitHub/TensorOperations.jl", rev = "main"}

From 72518ea6b0065a1cd0379fa18da04f36f0704056 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <kslimes@gmail.com>
Date: Thu, 17 Sep 2026 11:24:05 +0200
Subject: [PATCH 04/16] Touchups

---
 ext/TensorKitGPUArraysExt.jl | 34 ++++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 4d8627d8f..c6cd2c253 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -149,11 +149,11 @@ an entry of an `AbelianTreeTransformer`.
 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
+    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
@@ -164,17 +164,17 @@ struct DeviceAbelianTreeTransformer{VB <: AbstractVector{<:AbelianTransformerBlo
 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))
+_dense_strides(size::NTuple{N, Int}) where {N} = ntuple(n -> prod(size[1:(n - 1)]; init = 1), Val(N))
+_permute_strides(strides::NTuple{N, Int}, p) where {N} = ntuple(n -> strides[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
+        coeff::T, (size_dst, strides_dst, offsets_dst), (_, strides_src, offsets_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
+    return AbelianTransformerBlock{T, length(size_dst)}(
+        coeff, size_dst, _dense_strides(size_dst), strides_dst, offsets_dst,
+        _permute_strides(strides_src, p), offsets_src
     )
 end
 
@@ -254,7 +254,9 @@ _device_transformer(t::TensorKit.AbelianTreeTransformer, p) = DeviceAbelianTreeT
 # kernels are not reachable by coverage
 
 # largest `i` with `offsets[i] <= w`. This corresponds to the
-# block which  this kernel thread will work on.
+# 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
@@ -272,14 +274,14 @@ end
 # 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))
+@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 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
+@inline function _offset(coords::NTuple{N, Int}, strides::NTuple{N, Int}, offset) where {N}
+    return offset + sum(ntuple(n -> coords[n] * strides[n], Val(N))) + 1
 end
 
 # One thread per destination element in `data_dst`.

From b3070d6c98d5e5512326ff758896923da5db3853 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.s.hyatt@gmail.com>
Date: Thu, 17 Sep 2026 12:49:41 +0200
Subject: [PATCH 05/16] Forgot to change some variable names

---
 ext/TensorKitGPUArraysExt.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index c6cd2c253..ad0281ab4 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -292,9 +292,9 @@ end
     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)
+        coords = _coordinates(w - (@inbounds work_offsets[b]), blk.sz, blk.dense_strides)
+        i_dst = _offset(coords, blk.strides_dst, blk.offsets_dst)
+        i_src = _offset(coords, blk.permuted_strides_src, blk.offsets_src)
         @inbounds data_dst[i_dst] = α * blk.coeff * data_src[i_src] + β * data_dst[i_dst]
     end
 end

From 58f4f84086cb2ed58c8c255ec8a9c0344e6b44e1 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Mon, 28 Sep 2026 08:19:54 -0400
Subject: [PATCH 06/16] Updates on top of latest release

---
 ext/TensorKitGPUArraysExt.jl | 76 ++++++++++++++++++++----------------
 1 file changed, 43 insertions(+), 33 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index ad0281ab4..4042459b7 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -11,6 +11,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
 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}}}}
@@ -124,14 +125,15 @@ end
 
 # Device-side tree transformers
 # -----------------------------
-# We 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.
+# 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.
+# 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
@@ -141,12 +143,12 @@ end
 const TreeStructure{N} = Tuple{NTuple{N, Int}, Int}
 
 """
-    AbelianTransformerBlock{T, N}
+    UniqueTransformerBlock{T, N}
 
 `isbits` descriptor for a subblock that is a *single* scaled permutation:
-an entry of an `AbelianTreeTransformer`.
+an entry of a `UniqueTreeTransformer`.
 """
-struct AbelianTransformerBlock{T, N}
+struct UniqueTransformerBlock{T, N}
     coeff::T
     sz::NTuple{N, Int}
     dense_strides::NTuple{N, Int}
@@ -157,7 +159,7 @@ struct AbelianTransformerBlock{T, N}
 end
 
 # device-side simple struct that GPU kernels can use
-struct DeviceAbelianTreeTransformer{VB <: AbstractVector{<:AbelianTransformerBlock}, VO <: AbstractVector{Int}}
+struct DeviceUniqueTreeTransformer{VB <: AbstractVector{<:UniqueTransformerBlock}, VO <: AbstractVector{Int}}
     blocks::VB
     work_offsets::VO
     nwork::Int
@@ -169,10 +171,10 @@ _permute_strides(strides::NTuple{N, Int}, p) where {N} = ntuple(n -> strides[p[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(
+function _unique_block(
         coeff::T, (size_dst, strides_dst, offsets_dst), (_, strides_src, offsets_src), p
     ) where {T}
-    return AbelianTransformerBlock{T, length(size_dst)}(
+    return UniqueTransformerBlock{T, length(size_dst)}(
         coeff, size_dst, _dense_strides(size_dst), strides_dst, offsets_dst,
         _permute_strides(strides_src, p), offsets_src
     )
@@ -185,12 +187,14 @@ function _work_offsets(work)
     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]
+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 DeviceAbelianTreeTransformer(blocks, work_offsets, nwork)
+    return DeviceUniqueTreeTransformer(blocks, work_offsets, nwork)
 end
 
 """
@@ -209,8 +213,8 @@ function Adapt.adapt_storage(a::StorageAdaptor, x::AbstractArray)
     return copyto!(dst, x)
 end
 
-function Adapt.adapt_structure(to, t::DeviceAbelianTreeTransformer)
-    return DeviceAbelianTreeTransformer(
+function Adapt.adapt_structure(to, t::DeviceUniqueTreeTransformer)
+    return DeviceUniqueTreeTransformer(
         Adapt.adapt(to, t.blocks), Adapt.adapt(to, t.work_offsets), t.nwork
     )
 end
@@ -221,6 +225,8 @@ end
 # - `transformer.data`
 # - the storage type
 # - `p`, which is baked into the permuted source strides.
+# `conjsrc` is not part of the key: it is applied in the kernel, and the transformer data
+# for a conjugated source is already distinct (and cached separately) on the CPU side.
 # 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}}}()
@@ -248,7 +254,7 @@ function device_transformer(proto::AbstractArray, transformer, p)
     end
 end
 
-_device_transformer(t::TensorKit.AbelianTreeTransformer, p) = DeviceAbelianTreeTransformer(t, p)
+_device_transformer(t::UniqueTreeTransformer, p) = DeviceUniqueTreeTransformer(t, p)
 
 # COV_EXCL_START
 # kernels are not reachable by coverage
@@ -284,9 +290,10 @@ end
     return offset + sum(ntuple(n -> coords[n] * strides[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}
+# 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!(
+        data_dst, data_src, op, blocks, work_offsets, α, β, nwork, ::Val{N}
     ) where {N}
     w = @index(Global, Linear) - 1
     if w < nwork
@@ -295,29 +302,32 @@ end
         coords = _coordinates(w - (@inbounds work_offsets[b]), blk.sz, blk.dense_strides)
         i_dst = _offset(coords, blk.strides_dst, blk.offsets_dst)
         i_src = _offset(coords, blk.permuted_strides_src, blk.offsets_src)
-        @inbounds data_dst[i_dst] = α * blk.coeff * data_src[i_src] + β * data_dst[i_dst]
+        @inbounds data_dst[i_dst] = α * blk.coeff * op(data_src[i_src]) + β * data_dst[i_dst]
     end
 end
 
 # COV_EXCL_STOP
 
-function _launch_abelian!(data_dst, data_src, transformer, α, β, ::Val{N}) where {N}
+function _launch_unique!(data_dst, data_src, op, 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,
+    unique_batched_permute!(get_backend(data_dst))(
+        data_dst, data_src, op, transformer.blocks, transformer.work_offsets, α, β, nwork,
         Val(N); ndrange = nwork
     )
     return nothing
 end
 
+const GPUStridedSubblocks = StridedSubblocks{<:AnyGPUArray}
+
 function TensorKit.add_transform_kernel!(
-        data_dst::A, data_src::A, p, transformer::TensorKit.AbelianTreeTransformer{T, N},
-        α, β, backend, allocator, ntasks::Int
-    ) where {T, N, A <: AnyGPUArray}
+        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(data_dst, transformer, linearize(p))::DeviceAbelianTreeTransformer
-    _launch_abelian!(data_dst, data_src, device, α, β, Val(N))
+    device = device_transformer(dst.data, transformer, linearize(p))::DeviceUniqueTreeTransformer
+    op = conjsrc ? conj : identity
+    _launch_unique!(dst.data, src.data, op, device, α, β, Val(N))
     return nothing
 end
 

From 6d07a850f5592f7db1c1c0f7961212a0b81e8652 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <kshyatt@users.noreply.github.com>
Date: Mon, 28 Sep 2026 16:01:20 +0200
Subject: [PATCH 07/16] Update ext/TensorKitGPUArraysExt.jl

Co-authored-by: Jutho <Jutho@users.noreply.github.com>
---
 ext/TensorKitGPUArraysExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 4042459b7..dd4d3a3a5 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -166,7 +166,7 @@ struct DeviceUniqueTreeTransformer{VB <: AbstractVector{<:UniqueTransformerBlock
 end
 
 # strides of a dense array of shape `sz`
-_dense_strides(size::NTuple{N, Int}) where {N} = ntuple(n -> prod(size[1:(n - 1)]; init = 1), Val(N))
+_dense_strides(size::Dims) = (1, Base.front(cumprod(size))...)
 _permute_strides(strides::NTuple{N, Int}, p) where {N} = ntuple(n -> strides[p[n]], Val(N))
 
 # `permute(Vsrc, p) == Vdst` is enforced when the transformer is built, so the permuted

From 5b4fae0b14bef7bcaab8f4ac88630cb541aa34b7 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <kshyatt@users.noreply.github.com>
Date: Mon, 28 Sep 2026 18:11:26 +0200
Subject: [PATCH 08/16] Restore Project.toml

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index fda238c20..83093ea5b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -67,7 +67,7 @@ Random = "1"
 ScopedValues = "1.3.0"
 Strided = "2.6.1"
 TensorKitSectors = "0.3.7"
-TensorOperations = "5.8.1"
+TensorOperations = "5.5.2, 5.6"
 TimerOutputs = "1"
 TupleTools = "1.5"
 VectorInterface = "0.6"

From 328586bb980b292ceeb13e079e97a72ba15e2efe Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Mon, 28 Sep 2026 13:18:04 -0400
Subject: [PATCH 09/16] Use LRU cache

---
 ext/TensorKitGPUArraysExt.jl | 45 +++++++++++++++++-------------------
 1 file changed, 21 insertions(+), 24 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index dd4d3a3a5..a5dcaf7b5 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -4,6 +4,7 @@ using GPUArrays
 using GPUArrays: @allowscalar
 using GPUArrays.KernelAbstractions: @kernel, @index, get_backend
 using Adapt
+using TensorKit.LRUCache: LRU
 using Strided: StridedViews
 using MatrixAlgebraKit, Adapt
 using TensorKit
@@ -219,38 +220,34 @@ function Adapt.adapt_structure(to, t::DeviceUniqueTreeTransformer)
     )
 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`
+# 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.
-# `conjsrc` is not part of the key: it is applied in the kernel, and the transformer data
-# for a conjugated source is already distinct (and cached separately) on the CPU side.
-# 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()
+const DEVICE_TRANSFORMER_CACHE = LRU{Any, Any}(; maxsize = TensorKit.DEFAULT_GLOBALCACHE_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.data
-    return Base.@lock DEVICE_TRANSFORMER_LOCK begin
-        entry = get(DEVICE_TRANSFORMER_CACHE, objectid(key), nothing)
-        if isnothing(entry) || entry[1].value !== key
-            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
-            # be careful about the lifetime of these, since they live as long as their
-            # "parent" on the CPU, so they can persist beyond the call
-            GPUArrays.@uncached Adapt.adapt(
-                StorageAdaptor(proto), _device_transformer(transformer, p)
-            )
-        end
+    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)
+        )
     end
 end
 

From c2b8c6e176ce460316a9f65a83ca063d41d01f4d Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Mon, 28 Sep 2026 13:31:28 -0400
Subject: [PATCH 10/16] use Tupletools

---
 ext/TensorKitGPUArraysExt.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index a5dcaf7b5..8109fdbab 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -5,6 +5,7 @@ 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
 using TensorKit
@@ -168,7 +169,6 @@ end
 
 # strides of a dense array of shape `sz`
 _dense_strides(size::Dims) = (1, Base.front(cumprod(size))...)
-_permute_strides(strides::NTuple{N, Int}, p) where {N} = ntuple(n -> strides[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.
@@ -177,7 +177,7 @@ function _unique_block(
     ) where {T}
     return UniqueTransformerBlock{T, length(size_dst)}(
         coeff, size_dst, _dense_strides(size_dst), strides_dst, offsets_dst,
-        _permute_strides(strides_src, p), offsets_src
+        TupleTools.getindices(strides_src, p), offsets_src
     )
 end
 

From 1bf27100a64ddeafa328832f5912d5dc1605412a Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <kshyatt@users.noreply.github.com>
Date: Mon, 28 Sep 2026 23:06:52 +0200
Subject: [PATCH 11/16] Update ext/TensorKitGPUArraysExt.jl

Co-authored-by: Jutho <Jutho@users.noreply.github.com>
---
 ext/TensorKitGPUArraysExt.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 8109fdbab..eb9de6a1e 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -202,8 +202,8 @@ 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.
+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

From 5b233216c1090aed08efdc7fcb4badca993f4b9b Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Tue, 29 Sep 2026 02:32:20 -0400
Subject: [PATCH 12/16] Use copy rather than copyto

---
 ext/TensorKitGPUArraysExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index eb9de6a1e..78b61ea60 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -211,7 +211,7 @@ 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)
+    return copy!(dst, x)
 end
 
 function Adapt.adapt_structure(to, t::DeviceUniqueTreeTransformer)

From 40e1865c4fc5d93f7f9d4a1e7e9e766985b3c4ee Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Tue, 29 Sep 2026 02:33:40 -0400
Subject: [PATCH 13/16] _offset to _linear_index

---
 ext/TensorKitGPUArraysExt.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 78b61ea60..08793b02f 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -281,9 +281,9 @@ end
     return ntuple(n -> (w ÷ dense_strides[n]) % size[n], Val(N))
 end
 
-# finds the overall offset in the output and input arrays corresponding to the **sublock**
+# finds the overall linear index in the output and input arrays corresponding to the **sublock**
 # coordinates currently being worked on
-@inline function _offset(coords::NTuple{N, Int}, strides::NTuple{N, Int}, offset) where {N}
+@inline function _linear_index(coords::NTuple{N, Int}, strides::NTuple{N, Int}, offset) where {N}
     return offset + sum(ntuple(n -> coords[n] * strides[n], Val(N))) + 1
 end
 
@@ -297,8 +297,8 @@ end
         b = _searchblock(work_offsets, w)
         blk = @inbounds blocks[b]
         coords = _coordinates(w - (@inbounds work_offsets[b]), blk.sz, blk.dense_strides)
-        i_dst = _offset(coords, blk.strides_dst, blk.offsets_dst)
-        i_src = _offset(coords, blk.permuted_strides_src, blk.offsets_src)
+        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

From b5c3ab851f6ae6c3d0f7229dd99b33b2dc85271e Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Tue, 29 Sep 2026 02:35:17 -0400
Subject: [PATCH 14/16] Remove confusing comment

---
 ext/TensorKitGPUArraysExt.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index 08793b02f..e2afd4726 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -274,9 +274,6 @@ _device_transformer(t::UniqueTreeTransformer, p) = DeviceUniqueTreeTransformer(t
 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, size::NTuple{N, Int}, dense_strides::NTuple{N, Int}) where {N}
     return ntuple(n -> (w ÷ dense_strides[n]) % size[n], Val(N))
 end

From e11186f7127399f6fa8c5979a025b4316b6b1152 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Tue, 29 Sep 2026 02:36:10 -0400
Subject: [PATCH 15/16] Make the name of kernel more clear

---
 ext/TensorKitGPUArraysExt.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index e2afd4726..c5f5e2445 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -286,7 +286,7 @@ 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 function unique_batched_permute_kernel!(
         data_dst, data_src, op, blocks, work_offsets, α, β, nwork, ::Val{N}
     ) where {N}
     w = @index(Global, Linear) - 1
@@ -305,7 +305,7 @@ end
 function _launch_unique!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N}
     nwork = transformer.nwork
     nwork == 0 && return nothing
-    unique_batched_permute!(get_backend(data_dst))(
+    unique_batched_permute_kernel!(get_backend(data_dst))(
         data_dst, data_src, op, transformer.blocks, transformer.work_offsets, α, β, nwork,
         Val(N); ndrange = nwork
     )

From fde8b48a4b49a0363e4f35b26c5d75b30e34f186 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Tue, 29 Sep 2026 02:37:23 -0400
Subject: [PATCH 16/16] Remove type hint

---
 ext/TensorKitGPUArraysExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl
index c5f5e2445..3176336b1 100644
--- a/ext/TensorKitGPUArraysExt.jl
+++ b/ext/TensorKitGPUArraysExt.jl
@@ -319,7 +319,7 @@ function TensorKit.add_transform_kernel!(
         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))::DeviceUniqueTreeTransformer
+    device = device_transformer(dst.data, transformer, linearize(p))
     op = conjsrc ? conj : identity
     _launch_unique!(dst.data, src.data, op, device, α, β, Val(N))
     return nothing