twist! scales every fusion-tree subblock separately (slow on CPU, one kernel per subblock on GPU) #595

twist! loops over fusiontrees(t) and calls scale! on every subblock. Its cost therefore grows with the number of subblocks rather than with the data size: ≈ 1–1.7 µs per subblock on the CPU, and on a CuTensorMap a separate kernel launch (≈ 4–5 µs) per subblock. Fermionic tensors with a few U(1) charges have tens of thousands of subblocks. Since tensorcontract! twists one of its operands for fermionic sectors (twistA && twist!(Anew, …) in src/tensors/tensoroperations.jl), this is paid in every such contraction and in its pullback.

Measured on an H100 NVL host, TensorKit 0.17.2, CUDA.jl 6.4.2, fℤ₂ × U(1) tensors Vχ ⊗ VD ⊗ VD' ← Vχ ⊗ VD ⊗ VD' twisted on four indices, against one elementwise product of t.data with the vector of per-entry twist factors (fastest of 10 calls after a warm-up, times in ms):

The factors depend only on the space, the twisted indices and inv; computing them costs one twist!. Cached (MemoizationKit would take care of this), a repeated twist of the same space, the common case in iterative algorithms, is a single product: 230–1300× faster on the CPU and 4000–20000× on the GPU.

In a fermionic Hubbard PEPS (fℤ₂ × U(1) × U(1), D = 6, χ = 72) twist! was 67 % of the host time of a GPU CTMRG iteration. With cached factors in a local copy of TensorKit, the Zygote gradient of the energy at a fixed environment went from 42 to 12 s on one CPU thread.

Two possible fixes:

• Standalone twist!: cache the per-entry factors per (space, inds, inv), for instance with MemoizationKit.
• Contractions: fold the twist into the permutation that precedes it. In tensorcontract! a twisted operand is always copied first (copyA |= twistA), and that copy is a permutation whose kernels already multiply every destination subblock by a coefficient (α * blk.coeff in the batched GPU kernel, the recoupling matrix on the CPU). The twist is diagonal in the fusion-tree basis, with a phase that depends only on the uncoupled sectors of the twisted legs of the destination tree, so a twisting variant of treebraider/treetransposer (the twisted indices in its cache key) could scale those coefficients instead. That removes the separate pass over the operand and needs no per-entry factors; it also covers non-abelian sectors (scaling the rows of the recoupling matrix). Not implemented or measured yet.