add full precompilation suite

lkdvos · lkdvos · commit 054aa332dd64 · 2026-07-14T14:57:08.000-04:00
diff --git a/src/TensorKit.jl b/src/TensorKit.jl
@@ -104,9 +104,6 @@ export notrunc, truncrank, trunctol, truncfilter, truncspace, truncerror
 # cache management
 export empty_globalcaches!
 
-# precompilation
-export precompile_contract
-
 # Imports
 #---------
 using TupleTools
diff --git a/src/precompile.jl b/src/precompile.jl
@@ -1,5 +1,13 @@
-using PrecompileTools: PrecompileTools, @compile_workload
-using Preferences: @load_preference, load_preference
+module Precompilation
+
+export precompile_indexmanipulations, precompile_contract, precompile_factorizations
+
+using ..TensorKit
+using ..TensorKit: TO
+using VectorInterface: One, Zero
+using TensorOperations: @tensor
+using PrecompileTools: @setup_workload, @compile_workload
+using Preferences: @load_preference
 
 # Preferences
 # -----------
@@ -18,150 +26,46 @@ const PRECOMPILE_ELTYPES = _validate_precompile_eltypes(
     @load_preference("precompile_eltypes", ["Float64", "ComplexF64"])
 )
 
-# Tensor arities (numbers of legs) to precompile the contraction/permutation machinery for.
-# The heavy transform/braid machinery is specialized per arity, so this controls how many
-# leg-counts get the full precompilation benefit (higher arities cost more precompile time).
-const PRECOMPILE_NDIMS = let n = @load_preference("precompile_ndims", [2, 3, 4])
-    (n isa Vector{Int} && all(≥(1), n)) ||
-        throw(ArgumentError("`precompile_ndims` should be a vector of positive `Int`s, got `$n`"))
+const PRECOMPILE_NDIMS = let n = @load_preference("precompile_ndims", 4)
+    (n isa Int && n ≥ 1) ||
+        throw(ArgumentError("`precompile_ndims` should be a positive `Int`, got `$n`"))
     n
 end
 
-# Representative space for each supported sector-name in the `precompile_sectors` preference.
-# One representative per (fusion style, coefficient type) suffices — the heavy kernels are
-# sector-agnostic (see `precompile_contract`). Other sectors can be precompiled by calling
-# `precompile_contract` directly.
 function _precompile_space(name::AbstractString)
-    name == "Trivial" && return ℂ^2
-    name == "Z2Irrep" && return Vect[Z2Irrep](0 => 1, 1 => 1)
-    name == "U1Irrep" && return Vect[U1Irrep](-1 => 1, 0 => 1, 1 => 1)
-    name == "SU2Irrep" && return Vect[SU2Irrep](0 => 1, 1 // 2 => 1)
-    name == "FermionParity" && return Vect[FermionParity](0 => 1, 1 => 1)
-    return error("unknown precompile_sectors entry `$name`; precompile it directly with `precompile_contract`")
+    I = Base.eval(TensorKit, Meta.parse(name))
+    (I isa DataType && I <: Sector) ||
+        error("invalid `precompile_sectors` entry `$name`; expected a `Sector` type")
+    return oneunit(Vect[I])
 end
 
 const PRECOMPILE_SECTORS = let s = @load_preference(
-        "precompile_sectors", ["Trivial", "Z2Irrep", "SU2Irrep", "FermionParity"]
+        "precompile_sectors", map(x -> repr(x; context = :module => TensorKit), [Trivial, Z2Irrep, SU2Irrep, FermionParity])
     )
     s isa Vector{String} ||
         throw(ArgumentError("`precompile_sectors` should be a vector of strings, got $(typeof(s))"))
     s
 end
 
-# Whether to run the precompile workload. Respects the ecosystem-wide
-# `PrecompileTools` switch and a TensorKit-local `precompile_workload` preference
-# (default `true`). Mirrors the pattern used by TensorOperations.
-function _precompile_workload_enabled(mod::Module = @__MODULE__)
-    return try
-        if load_preference(PrecompileTools, "precompile_workloads", true)
-            return load_preference(mod, "precompile_workload", true)
-        else
-            return false
-        end
-    catch
-        false
-    end
-end
+const PRECOMPILE_CONTRACT = @load_preference("precompile_contract", true)
+const PRECOMPILE_INDEXMANIPULATIONS = @load_preference("precompile_indexmanipulations", true)
+const PRECOMPILE_FACTORIZATIONS = @load_preference("precompile_factorizations", true)
 
 # Workload
 # --------
-"""
-    precompile_contract(V::IndexSpace; eltypes=[Float64, ComplexF64], ndims=[2, 3, 4])
-
-Run a small, representative set of tensor operations (contraction, trace, permutation) on
-tensors built from the space `V`, for each element type in `eltypes` and each tensor arity
-(number of legs) in `ndims`. This forces compilation of the contraction/permutation machinery
-for the sector type of `V`.
-
-The machinery is specialized per arity, so `ndims` controls which leg-counts get the full
-benefit; arities not covered still benefit partially from the sector- and arity-agnostic parts
-(most importantly the per-block BLAS `mul!`).
-
-The expensive, numerically-heavy kernels (`add_transform!`, the tree transformers, and the
-per-block BLAS `mul!`) are *sector-agnostic* — they depend only on the element type, the
-arity, and `sectorscalartype(sectortype(V))`. Consequently, calling this once for a
-representative symmetry precompiles code that is reused by every other symmetry with the same
-fusion style and coefficient type.
-
-Downstream packages that define their own symmetry can precompile TensorKit's contraction path
-for it by calling this inside their own `PrecompileTools.@compile_workload`, e.g.
-
-```julia
-@compile_workload begin
-    TensorKit.precompile_contract(Vect[MyAnyon](s₀ => 1, s₁ => 1))
-end
-```
-
-TensorKit's own workload (enabled by default) calls this for `Trivial`, `Z2Irrep`, `SU2Irrep`
-and `FermionParity`. It can be tuned or disabled via preferences:
-
-```julia
-using TensorKit, Preferences
-set_preferences!(TensorKit, "precompile_eltypes" => ["Float64", "ComplexF64"]; force=true)
-set_preferences!(TensorKit, "precompile_workload" => false; force=true)  # disable entirely
-```
-"""
-function precompile_contract(V::IndexSpace; eltypes = PRECOMPILE_ELTYPES, ndims = PRECOMPILE_NDIMS)
-    backend = TO.DefaultBackend()
-    allocator = TO.DefaultAllocator()
-    for T in eltypes
-        α, β = rand(T), rand(T)
-
-        # contraction + permutation for each requested arity (leg count)
-        for N in ndims
-            _precompile_contract_arity(V, T, Val(N), α, β, backend, allocator)
+include("precompile/indexmanipulations.jl")
+include("precompile/contract.jl")
+include("precompile/factorizations.jl")
+
+@setup_workload begin
+    for name in PRECOMPILE_SECTORS
+        V = _precompile_space(name)
+        @compile_workload begin
+            PRECOMPILE_INDEXMANIPULATIONS && precompile_indexmanipulations(V; eltypes = PRECOMPILE_ELTYPES)
+            PRECOMPILE_CONTRACT && precompile_contract(V; eltypes = PRECOMPILE_ELTYPES)
+            PRECOMPILE_FACTORIZATIONS && precompile_factorizations(V; eltypes = PRECOMPILE_ELTYPES)
         end
-
-        # the conjugated-operand branch (`conjA=true`) is a distinct runtime path (adjoint
-        # handling) that `conj(A) * B` networks hit; `@tensor` handles the space bookkeeping
-        A2 = randn(T, V ← V)
-        B2 = randn(T, V ← V)
-        @tensor Cc[a; c] := conj(A2[b; a]) * B2[b; c]
-
-        # partial trace (the two traced legs are mutually dual)
-        At = randn(T, V ⊗ V' ← V)
-        TO.tensortrace!(
-            TO.tensoralloc_add(T, At, ((3,), ()), false, Val(false)),
-            At, ((3,), ()), ((1,), (2,)), false, α, β, backend, allocator
-        )
     end
-    return nothing
-end
-
-# `V^{⊗M}` (the unit space for `M == 0`), with `M` a compile-time constant.
-_repeat_space(V, ::Val{0}) = one(V)
-_repeat_space(V, ::Val{M}) where {M} = foldl(⊗, ntuple(_ -> V, Val(M)))
-
-# Precompile the contraction and permutation machinery for tensors with `N` legs. `N` is a
-# `Val` so the index tuples are compile-time concrete (the machinery specializes per arity).
-function _precompile_contract_arity(V, ::Type{T}, ::Val{N}, α, β, backend, allocator) where {T, N}
-    W = _repeat_space(V, Val(N - 1))   # N-1 legs
-    # contraction of two arity-N tensors over their N-1 shared legs -> arity-2 result
-    A = randn(T, V ← W)
-    B = randn(T, W ← V)
-    pA = ((1,), ntuple(i -> i + 1, Val(N - 1)))
-    pB = (ntuple(identity, Val(N - 1)), (N,))
-    _precompile_contract!(A, pA, B, pB, ((1,), (2,)), α, β, backend, allocator)
-
-    # a non-trivial permutation exercises the repartition/braid/transform machinery at arity N
-    # (the contraction above takes the no-copy view path for its natural partition)
-    permute(A, (ntuple(i -> N - i + 1, Val(N)), ()))
-    return nothing
-end
-
-# both scalar-type paths: generic `(α,β)` and the identity `(One(), Zero())` fast path
-function _precompile_contract!(A, pA, B, pB, pAB, α, β, backend, allocator)
-    T = promote_type(scalartype(A), scalartype(B))
-    C = TO.tensoralloc_contract(T, A, pA, false, B, pB, false, pAB, Val(false))
-    TO.tensorcontract!(C, A, pA, false, B, pB, false, pAB, α, β, backend, allocator)
-    TO.tensorcontract!(C, A, pA, false, B, pB, false, pAB, One(), Zero(), backend, allocator)
-    return C
 end
 
-if _precompile_workload_enabled()
-    @compile_workload begin
-        for name in PRECOMPILE_SECTORS
-            precompile_contract(_precompile_space(name); eltypes = PRECOMPILE_ELTYPES)
-        end
-    end
-end
+end # module Precompilation
diff --git a/src/precompile/contract.jl b/src/precompile/contract.jl
@@ -0,0 +1,57 @@
+"""
+    precompile_contract(V::IndexSpace; eltypes=[Float64, ComplexF64], ndims=4)
+
+Run a small, representative set of contraction, trace, and permutation operations on tensors built
+from the space `V`, for each element type in `eltypes` and each tensor arity (number of legs) in `1:ndims`.
+
+Note that it can be beneficial to put a more comprehensive set of relevant symmetries in a startup file,
+for example by adding:
+
+```julia
+@compile_workload begin
+    TensorKit.precompile_contract(Vect[MyAnyon](s₀ => 1, s₁ => 1))
+end
+```
+
+See also [`precompile_indexmanipulations`](@ref TensorKit.precompile_indexmanipulations), [`precompile_factorizations`](@ref TensorKit.precompile_factorizations).
+"""
+function precompile_contract(V::IndexSpace; eltypes = PRECOMPILE_ELTYPES, ndims = PRECOMPILE_NDIMS)
+    backend = TO.DefaultBackend()
+    allocator = TO.DefaultAllocator()
+    for T in eltypes
+        α, β = rand(T), rand(T)
+
+        # contraction + permutation for each requested arity (leg count); `Val(N)`/`ntuple`
+        # keep the index tuples concrete so the machinery specializes per arity
+        for N in 1:ndims
+            W = V^(N - 1)   # N-1 legs
+            # contraction of two arity-N tensors over their N-1 shared legs -> arity-2 result
+            A = randn(T, V ← W)
+            B = randn(T, W ← V)
+            pA = ((1,), ntuple(i -> i + 1, Val(N - 1)))
+            pB = (ntuple(identity, Val(N - 1)), (N,))
+            C = TO.tensoralloc_contract(T, A, pA, false, B, pB, false, ((1,), (2,)), Val(false))
+            # both scalar-type paths: generic `(α,β)` and the identity `(One(), Zero())` fast path
+            TO.tensorcontract!(C, A, pA, false, B, pB, false, ((1,), (2,)), α, β, backend, allocator)
+            TO.tensorcontract!(C, A, pA, false, B, pB, false, ((1,), (2,)), One(), Zero(), backend, allocator)
+
+            # a non-trivial permutation exercises the repartition/braid/transform machinery at
+            # arity N (the contraction above takes the no-copy view path for its natural partition)
+            permute(A, (ntuple(i -> N - i + 1, Val(N)), ()))
+        end
+
+        # the conjugated-operand branch (`conjA=true`) is a distinct runtime path (adjoint
+        # handling) that `conj(A) * B` networks hit; `@tensor` handles the space bookkeeping
+        A2 = randn(T, V ← V)
+        B2 = randn(T, V ← V)
+        @tensor Cc[a; c] := conj(A2[b; a]) * B2[b; c]
+
+        # partial trace (the two traced legs are mutually dual)
+        At = randn(T, V ⊗ V' ← V)
+        TO.tensortrace!(
+            TO.tensoralloc_add(T, At, ((3,), ()), false, Val(false)),
+            At, ((3,), ()), ((1,), (2,)), false, α, β, backend, allocator
+        )
+    end
+    return nothing
+end
diff --git a/src/precompile/factorizations.jl b/src/precompile/factorizations.jl
@@ -0,0 +1,61 @@
+"""
+    precompile_factorizations(V::IndexSpace; eltypes=[Float64, ComplexF64])
+
+Run a small, representative set of tensor-factorization operations (singular value, QR, LQ,
+eigenvalue, orthogonal/null-space and polar decompositions) on tensors built from the space `V`,
+for each element type in `eltypes`.
+
+Note that it can be beneficial to put a more comprehensive set of relevant symmetries in a startup file,
+for example by adding:
+
+```julia
+@compile_workload begin
+    TensorKit.precompile_factorizations(Vect[MyAnyon](s₀ => 1, s₁ => 1))
+end
+```
+
+See also [`precompile_contract`](@ref TensorKit.precompile_contract), [`precompile_indexmanipulations`](@ref TensorKit.precompile_indexmanipulations).
+"""
+function precompile_factorizations(V::IndexSpace; eltypes = PRECOMPILE_ELTYPES)
+    for T in eltypes
+        # Three representative shapes: a square tensor for the bulk of the decompositions, a
+        # hermitian one for the `eigh` path, and rectangular ones so the null-space kernels
+        # operate on non-empty blocks.
+        W = V^2   # V ⊗ V
+        t = randn(T, W ← W)            # square
+        tr = randn(T, W ← V)           # tall (codomain larger) -> non-empty left null space
+        tw = randn(T, V ← W)           # wide (domain larger)   -> non-empty right null space
+        th = (t + t') / 2              # hermitian (square, Euclidean inner product)
+
+        # Singular value decomposition
+        svd_full(t)
+        svd_compact(t)
+        svd_vals(t)
+        svd_trunc(t; trunc = truncrank(1))
+
+        # QR / LQ decompositions (null-space variants on the appropriately shaped tensors)
+        qr_full(t)
+        qr_compact(t)
+        qr_null(tr)
+        lq_full(t)
+        lq_compact(t)
+        lq_null(tw)
+
+        # Eigenvalue decompositions (hermitian variants require a hermitian input)
+        eig_full(t)
+        eig_vals(t)
+        eigh_full(th)
+        eigh_vals(th)
+
+        # Orthogonal / null-space helpers
+        left_orth(t)
+        right_orth(t)
+        left_null(tr)
+        right_null(tw)
+
+        # Polar decompositions
+        left_polar(t)
+        right_polar(t)
+    end
+    return nothing
+end
diff --git a/src/precompile/indexmanipulations.jl b/src/precompile/indexmanipulations.jl
@@ -0,0 +1,47 @@
+"""
+    precompile_indexmanipulations(V::IndexSpace; eltypes=[Float64, ComplexF64], ndims=4)
+
+Run a small, representative set of index-manipulation operations on tensors built from the space `V`,
+for each element type in `eltypes` and each tensor arity (number of legs) in `1:ndims`.
+
+Note that it can be beneficial to put a more comprehensive set of relevant symmetries in a startup file,
+for example by adding:
+
+```julia
+@compile_workload begin
+    TensorKit.precompile_indexmanipulations(Vect[MyAnyon](s₀ => 1, s₁ => 1))
+end
+```
+
+See also [`precompile_contract`](@ref TensorKit.precompile_contract), [`precompile_factorizations`](@ref TensorKit.precompile_factorizations).
+"""
+function precompile_indexmanipulations(V::IndexSpace; eltypes = PRECOMPILE_ELTYPES, ndims = PRECOMPILE_NDIMS)
+    for T in eltypes
+        # `Val(N)`/`ntuple` keep the index tuples concrete so the machinery specializes per arity
+        for N in 1:ndims
+            N₁ = cld(N, 2)   # codomain legs
+            # a tensor split non-trivially into codomain/domain (N₁ | N-N₁)
+            t = randn(T, V^N₁ ← V^(N - N₁))
+
+            # a non-trivial reordering of all N legs (reverse of `1:N`), split back into (N₁ | N-N₁)
+            perm = ntuple(i -> N - i + 1, Val(N))
+            p1 = ntuple(i -> perm[i], Val(N₁))
+            p2 = ntuple(i -> perm[i + N₁], Val(N - N₁))
+
+            # `permute` and `braid` funnel through `add_transform!` with different transformers
+            permute(t, (p1, p2))
+            braid(t, (p1, p2), ntuple(identity, Val(N)))   # `levels` is a tuple over the source indices
+
+            # canonical transpose `(reverse(domain), reverse(codomain))` is always a valid cyclic transpose
+            tp1 = ntuple(i -> N - i + 1, Val(N - N₁))
+            tp2 = ntuple(i -> N₁ - i + 1, Val(N₁))
+            transpose(t, (tp1, tp2))
+
+            # repartition to a different codomain size, forcing the repartition/transpose path
+            repartition(t, N₁ < N ? N₁ + 1 : N₁ - 1)
+
+            twist(t, 1)
+        end
+    end
+    return nothing
+end
diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl
@@ -553,15 +553,16 @@ Base.@deprecate(
     if p[1] === codomainind(tsrc) && p[2] === domainind(tsrc)
         add!(tdst, tsrc, α, β)
     else
+        p2 = (linearize(p), ())
         if has_array_view(tdst) && has_array_view(tsrc)
-            TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend, allocator)
+            TO.tensoradd!(tdst[], tsrc[], p2, false, α, β, backend, allocator)
         else
             ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1
             scheduler = ntasks == 1 ? SerialScheduler() : DynamicScheduler(; ntasks, split = :roundrobin)
             if tdst isa TensorMap && tsrc isa TensorMap # unpack data fields to avoid specializing
-                add_transform_kernel!(tdst.data, tsrc.data, p, transformer, α, β, backend, allocator, scheduler)
+                add_transform_kernel!(tdst.data, tsrc.data, p2, transformer, α, β, backend, allocator, scheduler)
             else
-                add_transform_kernel!(tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler)
+                add_transform_kernel!(tdst, tsrc, p2, transformer, α, β, backend, allocator, scheduler)
             end
         end
     end
