Regression on `main`: `permute`/`braid` ~1.7–1.9× slower than v0.16.5 for non-Abelian sectors

[leburgel]

On main (v0.17.1), permuting a TensorMap with a non-Abelian sector type seems to be ~1.7–1.9× slower than on v0.16.5, with ~40–55 % more allocations. Abelian sector types are unaffected.

Measurements

Benchmarking a simple permute and braid call back to back using both version, single threaded on Julia v1.12.5 gives:

permute(t, p) with t::TensorMap, same tensors:

sector type	v0.16.5	main	ratio	allocs 0.16.5 → main
fℤ₂	0.0840 ms	0.0806 ms	0.96×	65 → 63
fℤ₂ ⊠ U(1)	0.0360 ms	0.0367 ms	1.02×	317 → 315
SU(2)	0.1904 ms	0.3584 ms	1.88×	1 841 → 2 845
fℤ₂ ⊠ U(1) ⊠ SU(3)	0.0281 ms	0.0471 ms	1.67×	470 → 649

braid(t, p, levels) with t::TensorMap, same tensors:

sector type	v0.16.5	main	ratio
fℤ₂	0.0815 ms	0.0858 ms	1.05×
fℤ₂ ⊠ U(1)	0.0361 ms	0.0365 ms	1.01×
SU(2)	0.1981 ms	0.3581 ms	1.81×
fℤ₂ ⊠ U(1) ⊠ SU(3)	0.0281 ms	0.0473 ms	1.68×

Potential cause

The regression is in the generic (non-abelian) transform kernel, and appears to come from the rewrite of add_kernel_nonthreaded!/_add_transform_multi! into add_transform_kernel!: the per-fusion-block buffers, which v0.16.5 allocated once per call and reused, are now allocated and freed per fusion block under a ReentrantLock.

Comparing the generic non-abelian kernel between versions:

v0.16.5 — add_kernel_nonthreaded!(::FusionStyle, …, ::GenericTreeTransformer, …)
allocates buffers once, outside the loop, and reuses them for every fusion block:

buffers = allocate_buffers(tdst, tsrc, transformer)   # once per call
for subtransformer in transformer.data
    ...
    _add_transform_multi!(tdst, tsrc, p, subtransformer, buffers, α, β, backend...)
end

and the pack step is a direct strided copy:

copy!(bufblock_src, subblock_src)

main — add_transform_kernel!(…, ::GenericTreeTransformer, …) allocates and frees a
buffer inside the per-block loop, under a lock, and uses a general tensoradd! for the
pack step:

OhMyThreads.@tasks for subtransformer in transformer.data
    ...
    buffer = @lock buffer_lock TO.tensoralloc(typeof(data_dst), blocksize * (rows + cols), Val(true), allocator)
    ...
    TO.tensoradd!(sreshape(view(buffer_src, :, i), sz_src), …, ptriv, false, One(), Zero(), backend, allocator)
    ...
    U′ = Adapt.adapt(storagetype(tdst), StridedView(U))
    ...
    @lock buffer_lock TO.tensorfree!(buffer, allocator)
end

Candidate contributors, in rough order of suspicion:

1. Per-block tensoralloc/tensorfree! instead of one preallocation per call, which also
   matches the allocation counts (SU(2): +1004 allocs, SU(3): +179).
2. ReentrantLock acquired twice per multi-tree block, paid even when the scheduler is
   serial.
3. Adapt.adapt(storagetype(tdst), StridedView(U)) per block.
4. The pack step going through a general TO.tensoradd! with a trivial permutation rather
   than copy!.

The abelian kernel (AbelianTreeTransformer) is a plain tforeach with no buffers and no
lock, which is consistent with abelian sectors showing no regression.

Reproduction

Run the following in two environments — one pinned to TensorKit v0.16.5, one tracking
main — and compare:

using TensorKit, BenchmarkTools, SUNRepresentations

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

function bench(name, V)
    t = randn(ComplexF64, V ⊗ V ← V ⊗ V)
    levels = ntuple(identity, 4)
    bA = @benchmark TensorKit.permute($t, $p)        samples = 500 seconds = 10
    bB = @benchmark braid($t, $p, $levels)           samples = 500 seconds = 10
    println(rpad(name, 14),
            " permute ", round(minimum(bA).time / 1e6; digits = 5), " ms (", bA.allocs, " allocs)",
            "  braid ",  round(minimum(bB).time / 1e6; digits = 5), " ms (", bB.allocs, " allocs)")
end

bench("fZ2",    Vect[FermionParity](0 => 8, 1 => 8))
bench("fZ2xU1", Vect[FermionParity ⊠ U1Irrep]((0,0) => 4, (1,1) => 4, (0,2) => 3, (1,-1) => 3))
bench("SU2",    Vect[SU2Irrep](0 => 4, 1//2 => 4, 1 => 3, 3//2 => 2))
bench("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))

[lkdvos]

For context, most of these changes were related to me trying to generalize some of our calls and making everything work for GPU out of the box, and with arbitrary allocator frameworks. Realistically I can definitely tone down the generality on that, which gets rid of the locks etc. I'll have a look and propose something next week. Do we have specific benchmark cases we could add to catch this and other regressions easier in the future?