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/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/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 +``` diff --git a/src/TensorKit.jl b/src/TensorKit.jl index f20859dc6..e5e3a4dc5 100644 --- a/src/TensorKit.jl +++ b/src/TensorKit.jl @@ -288,4 +288,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..7b81ac701 --- /dev/null +++ b/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 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