Precompilation

Project.toml

@@ -13,6 +13,8 @@ LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
@@ -57,6 +59,8 @@ LinearAlgebra = "1"
 MatrixAlgebraKit = "0.6.9"
 Mooncake = "0.5.27"
 OhMyThreads = "0.8.0"
+PrecompileTools = "1"
+Preferences = "1"
 Printf = "1"
 Random = "1"
 ScopedValues = "1.3.0"

docs/make.jl

@@ -33,6 +33,7 @@ pages = [
             "man/indexmanipulations.md",
             "man/factorizations.md",
             "man/contractions.md",
+            "man/precompilation.md",
         ],
     ],
     "Library" => [

docs/src/man/precompilation.md

@@ -0,0 +1,44 @@
+# Precompilation
+
+The first tensor operation in a Julia session — a contraction, an index manipulation, or a factorization — incurs a significant compilation latency (time-to-first-execution, TTFX).
+To mitigate this, TensorKit ships a precompilation workload that is **enabled by default**:
+it runs a small set of representative index manipulations, contractions and factorizations for the `Trivial`, `Z2Irrep`, `SU2Irrep` and `FermionParity` symmetries over `Float64` and `ComplexF64`.
+
+Because the numerically-heavy kernels are mostly *sector-agnostic* (they depend only on the element type, the arity, and `sectorscalartype`),
+this workload also speeds up the first operation on symmetries that are *not* in the workload — including user-defined ones.
+
+The workload can be tuned or disabled through [Preferences](https://github.com/JuliaPackaging/Preferences.jl):
+
+```julia
+using TensorKit, Preferences, PrecompileTools
+# disable the workload entirely (fastest precompilation, no TTFX benefit)
+PrecompileTools.set_preferences!(TensorKit, "precompile_workloads" => false; force=true)
+# disable individual suites (each defaults to `true`)
+set_preferences!(TensorKit, "precompile_contract" => false; force=true)
+set_preferences!(TensorKit, "precompile_indexmanipulations" => false; force=true)
+set_preferences!(TensorKit, "precompile_factorizations" => false; force=true)
+# restrict the element types
+set_preferences!(TensorKit, "precompile_eltypes" => ["Float64"]; force=true)
+# restrict / change the symmetries that are precompiled
+set_preferences!(TensorKit, "precompile_sectors" => ["Trivial", "Z2Irrep"]; force=true)
+# highest tensor arity (leg count) to precompile, i.e. arities `1:n`; e.g. rank-6 for PEPS/PEPO
+set_preferences!(TensorKit, "precompile_ndims" => 6; force=true)
+```
+
+Changing a preference triggers recompilation of TensorKit the next time it is loaded.
+
+Downstream packages or startup files that define or heavily use their own symmetry can precompile TensorKit's operations for it by
+calling the `precompile_*` helpers inside their own `PrecompileTools.@compile_workload`:
+
+```julia
+using TensorKit, PrecompileTools
+@compile_workload begin
+    TensorKit.precompile_contract(Vect[MyAnyon](s₀ => 1, s₁ => 1))
+end
+```
+
+```@docs; canonical=false
+TensorKit.Precompilation.precompile_contract
+TensorKit.Precompilation.precompile_indexmanipulations
+TensorKit.Precompilation.precompile_factorizations
+```

src/TensorKit.jl

@@ -288,4 +288,6 @@ include("pullbacks/tensoroperations.jl")
 
 include("pullbacks/indexmanipulations.jl")
 
+include("precompile.jl")
+
 end

src/fusiontrees/duality_manipulations.jl

@@ -474,9 +474,10 @@ U = Utmp * U
 return dst, U
 ```
 =#
-@generated function repartition(src::Union{FusionTreePair, FusionTreeBlock}, ::Val{N}) where {N}
-    return _repartition_body(numout(src) - N)
-end
+# NOTE: `_repartition_body` must be defined *before* the `@generated repartition` that calls
+# it, so that it is visible to the generator's world age. Otherwise a precompile workload that
+# first triggers `repartition` fails with `UndefVarError: _repartition_body` (at runtime the
+# first call happens at a later world age, which masks the ordering issue).
 function _repartition_body(N)
     if N == 0
         ex = quote
@@ -503,6 +504,9 @@ function _repartition_body(N)
     end
     return ex
 end
+@generated function repartition(src::Union{FusionTreePair, FusionTreeBlock}, ::Val{N}) where {N}
+    return _repartition_body(numout(src) - N)
+end
 
 """
     transpose((f₁, f₂)::FusionTreePair{I}, p::Index2Tuple{N₁, N₂}) where {I, N₁, N₂}

src/precompile.jl

@@ -0,0 +1,71 @@
+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
+# -----------
+function _validate_precompile_eltypes(eltypes)
+    eltypes isa Vector{String} ||
+        throw(ArgumentError("`precompile_eltypes` should be a vector of strings, got $(typeof(eltypes)) instead"))
+    return map(eltypes) do Tstr
+        T = eval(Meta.parse(Tstr))
+        (T isa DataType && T <: Number) ||
+            error("Invalid precompile_eltypes entry: `$Tstr`")
+        return T
+    end
+end
+
+const PRECOMPILE_ELTYPES = _validate_precompile_eltypes(
+    @load_preference("precompile_eltypes", ["Float64", "ComplexF64"])
+)
+
+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
+
+function _precompile_space(name::AbstractString)
+    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", 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
+
+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
+# --------
+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
+    end
+end
+
+end # module Precompilation

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

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

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

src/tensors/indexmanipulations.jl

@@ -560,15 +560,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