Precompilation#487
Conversation
lkdvos
force-pushed
the
ld-blascompile
branch
2 times, most recently
from
ce58308 to
f6b5b49
Compare
Codecov Report❌ Patch coverage is
... and 11 files with indirect coverage changes 🚀 New features to boost your workflow:
|
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) <noreply@anthropic.com>
| ```julia | ||
| using TensorKit, PrecompileTools | ||
| @compile_workload begin | ||
| TensorKit.precompile_contract(Vect[MyAnyon](s₀ => 1, s₁ => 1)) |
There was a problem hiding this comment.
Since we're defaulting to unit spaces for the sectors from TensorKit, maybe it's better to say that that suffices in this docstring?
Also, I prefer MySector or CustomSector to MyAnyon as a name.
On another note, shouldn't you call TensorKit.Precompilation.precompile_x, or otherwise put using .Precompilation underneath its include in TensorKit.jl?
| ```julia | ||
| using TensorKit, Preferences, PrecompileTools | ||
| # disable the workload entirely (fastest precompilation, no TTFX benefit) | ||
| PrecompileTools.set_preferences!(TensorKit, "precompile_workloads" => false; force=true) |
There was a problem hiding this comment.
I think TensorKit here needs to be replaced by PrecompileTools, at least from what I understand here. The other option is "precompile_workload" in singular form, which I think is more what you want to achieve here.
| 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]) |
There was a problem hiding this comment.
| return oneunit(Vect[I]) | |
| return unitspace(Vect[I]) |
| if p[1] === codomainind(tsrc) && p[2] === domainind(tsrc) | ||
| add!(tdst, tsrc, α, β) | ||
| else | ||
| p2 = (linearize(p), ()) |
There was a problem hiding this comment.
Sorry for my naive question here (I'm not too familiar with this part of the code), but I guess passing the same type of tuple, one containing the indices and the other always being empty, is beneficial for precompilation, and somehow doesn't matter within TensorOperations thatindices are grouped up like this now? So why would you not make the same changes to transpose! and braid in the cases that they don't pass add_transform! first?
2 participants
This PR adds a precompilation suite to TensorKit.
It takes about 300 seconds on my machine, and is enabled by default.
Given the general migration of the ecosystem towards more heavy precompilation workflows, I think this is warranted, and I'm hoping this improves our CI times as well as the general user experiences.
I've also structured it in such a way that you could conceivably reuse this locally for projects with specific symmetry needs, to ensure faster startup times.
As in general TensorKit is not updated all that often, I hope this is not something that has to be retriggered too much, therefore I felt the ~5mins of precompilation time to be reasonable.