From 024d7308ac1d93c217ab2cafc5feff2cf171bb2c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 11 Jul 2026 22:08:21 -0400 Subject: [PATCH 1/3] Precompile the tensor-contraction path via a PrecompileTools workload Add `src/precompile.jl`: a default-on `@compile_workload` that runs representative contractions, traces and permutations for `Trivial`, `Z2Irrep`, `SU2Irrep` and `FermionParity` over `Float64`/`ComplexF64` and arities [2,3,4]. The heavy kernels (tree transformers, `add_transform!`, per-block BLAS `mul!`) are sector-agnostic, so this also speeds up the first contraction of other symmetries -- including user-defined ones. The exported `precompile_contract(V; eltypes, ndims)` helper is the reusable unit; downstream packages can call it for their own symmetry inside their own `@compile_workload`. The workload is tunable/disable-able through Preferences (`precompile_workload`, `precompile_eltypes`, `precompile_sectors`, `precompile_ndims`), mirroring TensorOperations. Also move `_repartition_body` above the `@generated repartition` that calls it at generation time: the workload is the first code to trigger `repartition` at precompile time, which exposed a latent world-age ordering bug (`UndefVarError` that runtime masks via a later world age). Adds PrecompileTools and Preferences dependencies. Co-Authored-By: Claude Opus 4.8 (1M context) --- Project.toml | 4 + docs/src/man/contractions.md | 42 ++++++ src/TensorKit.jl | 5 + src/fusiontrees/duality_manipulations.jl | 10 +- src/precompile.jl | 167 +++++++++++++++++++++++ 5 files changed, 225 insertions(+), 3 deletions(-) create mode 100644 src/precompile.jl diff --git a/Project.toml b/Project.toml index 4344242d5..7db1ed492 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/man/contractions.md b/docs/src/man/contractions.md index dcaf13bf5..d5eddaacd 100644 --- a/docs/src/man/contractions.md +++ b/docs/src/man/contractions.md @@ -317,3 +317,45 @@ Since in this case `@tensor` and `@planar` yield different results, the `@planso [^Mortier]: Mortier, Q., Devos, L., Burgelman, L., et al. (2025). Fermionic Tensor Network Methods. SciPost Physics 18, no. 1. [10.21468/SciPostPhys.18.1.012](https://doi.org/10.21468/SciPostPhys.18.1.012). + +## Precompilation + +The first tensor contraction in a Julia session 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 contractions, traces and +permutations for the `Trivial`, `Z2Irrep`, `SU2Irrep` and `FermionParity` symmetries over +`Float64` and `ComplexF64`. + +Because the numerically-heavy kernels are *sector-agnostic* (they depend only on the element +type, the arity, and `sectorscalartype`), this workload also speeds up the first contraction of +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 +# disable the workload entirely (fastest precompilation, no TTFX benefit) +set_preferences!(TensorKit, "precompile_workload" => 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) +# change the tensor arities (leg counts) that are precompiled, e.g. up to rank-6 for PEPS/PEPO +set_preferences!(TensorKit, "precompile_ndims" => [2, 3, 4, 5, 6]; force=true) +``` + +Changing a preference triggers recompilation of TensorKit the next time it is loaded. + +Downstream packages that define their own symmetry can precompile TensorKit's contraction path +for it by calling [`precompile_contract`](@ref) 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 +precompile_contract +``` diff --git a/src/TensorKit.jl b/src/TensorKit.jl index f20859dc6..9d8fbba18 100644 --- a/src/TensorKit.jl +++ b/src/TensorKit.jl @@ -104,6 +104,9 @@ export notrunc, truncrank, trunctol, truncfilter, truncspace, truncerror # cache management export empty_globalcaches! +# precompilation +export precompile_contract + # Imports #--------- using TupleTools @@ -288,4 +291,6 @@ include("pullbacks/tensoroperations.jl") include("pullbacks/indexmanipulations.jl") +include("precompile.jl") + end diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index 4187868ad..6faff78a0 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/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₂} diff --git a/src/precompile.jl b/src/precompile.jl new file mode 100644 index 000000000..251bdbe55 --- /dev/null +++ b/src/precompile.jl @@ -0,0 +1,167 @@ +using PrecompileTools: PrecompileTools, @compile_workload +using Preferences: @load_preference, 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"]) +) + +# 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`")) + 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`") +end + +const PRECOMPILE_SECTORS = let s = @load_preference( + "precompile_sectors", ["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 + +# 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) + 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 From ab3fc2ca0394f77b7725de747f17407df8895cb0 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 14 Jul 2026 14:30:36 -0400 Subject: [PATCH 2/3] add full precompilation suite --- src/TensorKit.jl | 3 - src/precompile.jl | 162 ++++++--------------------- src/precompile/contract.jl | 57 ++++++++++ src/precompile/factorizations.jl | 61 ++++++++++ src/precompile/indexmanipulations.jl | 47 ++++++++ src/tensors/indexmanipulations.jl | 7 +- 6 files changed, 202 insertions(+), 135 deletions(-) create mode 100644 src/precompile/contract.jl create mode 100644 src/precompile/factorizations.jl create mode 100644 src/precompile/indexmanipulations.jl diff --git a/src/TensorKit.jl b/src/TensorKit.jl index 9d8fbba18..e5e3a4dc5 100644 --- 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 index 251bdbe55..7b81ac701 100644 --- 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 new file mode 100644 index 000000000..d6a206ceb --- /dev/null +++ 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 new file mode 100644 index 000000000..f4abaa000 --- /dev/null +++ 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 new file mode 100644 index 000000000..d08d73c60 --- /dev/null +++ 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 index 7155bb4be..6c0252acb 100644 --- a/src/tensors/indexmanipulations.jl +++ b/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 From ccf6b65fd09142889b38177b8946af2e93c0da09 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 14 Jul 2026 14:32:02 -0400 Subject: [PATCH 3/3] update precompilation docs --- docs/make.jl | 1 + docs/src/man/contractions.md | 42 -------------------------------- docs/src/man/precompilation.md | 44 ++++++++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 42 deletions(-) create mode 100644 docs/src/man/precompilation.md diff --git a/docs/make.jl b/docs/make.jl index d72b851a6..08e5bd5d5 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -33,6 +33,7 @@ pages = [ "man/indexmanipulations.md", "man/factorizations.md", "man/contractions.md", + "man/precompilation.md", ], ], "Library" => [ diff --git a/docs/src/man/contractions.md b/docs/src/man/contractions.md index d5eddaacd..dcaf13bf5 100644 --- a/docs/src/man/contractions.md +++ b/docs/src/man/contractions.md @@ -317,45 +317,3 @@ Since in this case `@tensor` and `@planar` yield different results, the `@planso [^Mortier]: Mortier, Q., Devos, L., Burgelman, L., et al. (2025). Fermionic Tensor Network Methods. SciPost Physics 18, no. 1. [10.21468/SciPostPhys.18.1.012](https://doi.org/10.21468/SciPostPhys.18.1.012). - -## Precompilation - -The first tensor contraction in a Julia session 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 contractions, traces and -permutations for the `Trivial`, `Z2Irrep`, `SU2Irrep` and `FermionParity` symmetries over -`Float64` and `ComplexF64`. - -Because the numerically-heavy kernels are *sector-agnostic* (they depend only on the element -type, the arity, and `sectorscalartype`), this workload also speeds up the first contraction of -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 -# disable the workload entirely (fastest precompilation, no TTFX benefit) -set_preferences!(TensorKit, "precompile_workload" => 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) -# change the tensor arities (leg counts) that are precompiled, e.g. up to rank-6 for PEPS/PEPO -set_preferences!(TensorKit, "precompile_ndims" => [2, 3, 4, 5, 6]; force=true) -``` - -Changing a preference triggers recompilation of TensorKit the next time it is loaded. - -Downstream packages that define their own symmetry can precompile TensorKit's contraction path -for it by calling [`precompile_contract`](@ref) 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 -precompile_contract -``` diff --git a/docs/src/man/precompilation.md b/docs/src/man/precompilation.md new file mode 100644 index 000000000..344bc37b9 --- /dev/null +++ b/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 +```