Permuting an `AdjointTensorMap` seems much slower than it should be

[leburgel]

It seems the in-place permute! of an AdjointTensorMap with a plain TensorMap parent is much slower than it should be compared to the time and allocations it takes to permute! a plain TensorMap. This issue is specific to the in-place method, since an out-of-place permute seems to just unwrap the parent and follow the same path as a TensorMap permute. It also seems that the slowdown is much worse for TensorMaps with non-Abelian sectortypes, becoming increasingly severe with increasing complexity of the symmetry.

I first encountered this on TensorKit.jl v0.16.5, but it can be reproduced on the current main (v0.17.1). Running the following script,

Reproducer script

using TensorKit, BenchmarkTools
using TensorKit.VectorInterface: One, Zero

p = ((2, 3), (4, 1))

function compare(name, V)
    t   = randn(ComplexF64, V ⊗ V ← V ⊗ V)
    adj = t'         # AdjointTensorMap
    mat = copy(t')   # identical space and contents, but a genuine TensorMap

    Vdst    = TensorKit.permute(space(mat), p)
    dst_adj = similar(mat, Vdst)
    dst_mat = similar(mat, Vdst)

    # NB: permute(::AdjointTensorMap, p) unwraps the adjoint and is therefore NOT affected.
    # The slow path is the in-place one, which is what `tensoradd!`/`blas_contract!` use.
    f_adj() = permute!(dst_adj, adj, p, One(), Zero())
    f_mat() = permute!(dst_mat, mat, p, One(), Zero())
    f_adj(); f_mat()
    @assert convert(Array, dst_adj) ≈ convert(Array, dst_mat)

    b_adj = @benchmark $f_adj() samples = 500 seconds = 10
    b_mat = @benchmark $f_mat() samples = 500 seconds = 10
    println(rpad(name, 14),
            " adjoint ", round(minimum(b_adj).time / 1e6; digits = 4), " ms (", b_adj.allocs, " allocs)",
            "  tensormap ", round(minimum(b_mat).time / 1e6; digits = 4), " ms (", b_mat.allocs, " allocs)",
            "  ratio ", round(minimum(b_adj).time / minimum(b_mat).time; digits = 2), "x")
end

compare("fZ2",    Vect[FermionParity](0 => 8, 1 => 8))
compare("fZ2xU1", Vect[FermionParity ⊠ U1Irrep]((0,0) => 4, (1,1) => 4, (0,2) => 3, (1,-1) => 3))
compare("SU2",    Vect[SU2Irrep](0 => 4, 1//2 => 4, 1 => 3, 3//2 => 2))

# needs SUNRepresentations
using SUNRepresentations
compare("fZ2xU1xSU3", Vect[FermionParity ⊠ U1Irrep ⊠ SU3Irrep](
    (0, 0, SU3Irrep(0,0,0)) => 2, (1, 1, SU3Irrep(1,0,0)) => 1,
    (0, 2, SU3Irrep(1,1,0)) => 1, (1, 3, SU3Irrep(1,1,1)) => 1))

on a single thread using Julia 1.12.5 and TensorKit v0.17.1 (main) gives timings:

sector type	adjoint	TensorMap	ratio	allocs (adjoint / TensorMap)
fℤ₂	0.0717 ms	0.0539 ms	1.33×	122 / 58
fℤ₂ ⊠ U(1)	0.1232 ms	0.0319 ms	3.86×	794 / 310
SU(2)	2.4282 ms	0.3382 ms	7.18×	27 094 / 2 840
fℤ₂ ⊠ U(1) ⊠ SU(3)	3.8169 ms	0.0465 ms	82.08×	43 529 / 645

Diagnosis

(Updated from the original issue description after some more digging)

There seem to be two independent issues that cause this.

The first is the fact that an in-place permutation on an AdjointTensorMap falls back on a generic AbstractTensorMap path which doesn't use the same cached TreeTransformer as the regular plain TensorMap permute path. The generic path seems to only hit a cached method at a later stage (the GLOBAL_FSBRAID_CACHE), really performing repeated calls for every fusion tree pair on every call, rather than caching the whole transformer once per (Vdst, Vsrc, p, levels) specification (as contained in the GLOBAL_TREEBRAIDER_CACHE used on the fast TensorMap path).

The second is an issue with the keys used to access this GLOBAL_FSBRAID_CACHE, causing all queries to miss and rendering the cache effectively useless. This turned out to be due to the fact that FusionTreeBlock doesn't have a hash implementation. This in contrast to FusionTree which does have a hash, and is used to create keys for the GLOBAL_TREEBRAIDER_CACHE. This is the reason the cost is so much worse for non-Abelian sector types, because there the same symmetry data is being recomputed on every call and nothing is reused.

Potential solutions

The second issue is very easy to solve, just adding a Base.hash implementation for FusionTreeBlock resolves the GLOBAL_FSBRAID_CACHE misses and removes the additional overhead for non-Abelian sector types. This is done in #518.

To resolve the first issue, we should avoid AdjointTensorMap wrappers around plain TensorMaps that go through the fully generic path.

This could be done by just not instantiating an AdjointTensorMap in tensoradd! when conjA = true, and carrying the conj flag all the way into the eventual kernel. Some version of this is suggested in #520.

Another option is to just write a dedicated fast cached treebraider for an AdjointTensorMap wrapper of a plain TensorMap. Some version of this is suggested in #519.

I'm not quite sure which of these, if any, would be the best option. Any opinions would be very welcome, @Jutho and @lkdvos.

[lkdvos]

I'll try and investigate, but I'm really confused why this is causing cache misses, in principle this should be sharing the same cache as the other ones so unless the cache just isn't large enough this should hit the same amount?

[leburgel]

I looked a bit more, and I think the cache misses might be an issue with hashes. It seems the cache that is used for permutations of adjoint tensor maps just always has exactly zero hits: GLOBAL_FSBRAID_CACHE never gets a single hit. As in, doing the same permutation twice just seems to grow the cache. So cache size isn't the issue, increasing it just means it keeps growing in an unbounded way until the computation ends.

This seems to be related to the fact that FusionTreeBlock doesn't have a hash implementation. This in contrast to FusionTree which does have a hash, and is used to create keys for the GLOBAL_TREEBRAIDER_CACHE, which is (I think) the cache that is used for permuting regular tensor maps.

I filed the issue to point out that permuting an adjoint tensor map doesn't use the same tree braider cache as permutes of regular tensors, which I thought was weird. But the related performance issue seems purely due to non-matching keys for the cache used in the adjoint permutes. It just looked like this path never used a cache, because everything was always being recomputed. So it's actually slightly worse, it recomputes every time but also keeps everything in memory in a cache that only keeps filling up.

But, this is something I can actually try and fix, so I'll give that a go.

[leburgel]

Adding the hash for FusionTreeBlock definitely helps. It resolves the cache misses completely, but it still seems like the different path for adjoint tensor permutes inherently gives more allocations. With the hash fix, the reproduction script gives

sector type	adjoint	TensorMap	ratio	allocs (adjoint / TensorMap)
fℤ₂	0.0892 ms	0.0660 ms	1.35×	129 / 65
fℤ₂ ⊠ U(1)	0.1306 ms	0.0361 ms	3.61×	669 / 317
SU(2)	1.1784 ms	0.3621 ms	3.25×	7 864 / 2 847
fℤ₂ ⊠ U(1) ⊠ SU(3)	0.6892 ms	0.0526 ms	13.11×	8 011 / 652

So better, but still a problem it seems.

I can start by adding the hash fix since this is already a strict improvement, but it seems it would be better still if we could route AdjointTensorMap through the cached treebraider path.