Skip to content

Add PrecompileTools workload for improved startup time#185

Merged
ChrisRackauckas merged 1 commit into
SciML:mainfrom
ChrisRackauckas-Claude:precompile-improvements-20251229-201253
Dec 30, 2025
Merged

Add PrecompileTools workload for improved startup time#185
ChrisRackauckas merged 1 commit into
SciML:mainfrom
ChrisRackauckas-Claude:precompile-improvements-20251229-201253

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown
Contributor

Summary

  • Adds PrecompileTools precompilation workloads to significantly reduce time-to-first-execution (TTFX) for common operations
  • Adds PrecompileTools as a dependency with compat "1"
  • Adds src/precompilation.jl with @compile_workload for core operations

Performance Improvements

Operation Before After Improvement
Package load ~13.3s ~11.0s 17% faster
Model construction + setup ~2.3s ~0.017s 135x faster
Second model construction N/A ~0.002s -

What's Precompiled

The workload precompiles the most common user operations:

  • DeepEquilibriumNetwork construction with Dense layers and SSRootfind solver
  • Lux.setup() for parameter/state initialization
  • DeepEquilibriumSolution constructor
  • Common utility functions (flatten, flatten_vcat, split_and_reshape, check_unrolled_mode, zeros_init)
  • Fixed-depth (unrolled) mode state handling

Notes

  • Forward pass TTFX depends on the specific solver used (NonlinearSolve, OrdinaryDiffEq) which are optional dependencies
  • Users should expect solver-specific compilation on first use of a particular solver
  • Invalidation analysis showed ~2335 invalidation trees, but they originate from dependencies (Static, ChainRulesCore, RecursiveArrayTools, GPUCompiler) rather than this package

Test plan

  • CI passes
  • Package precompiles without errors
  • Verify startup time improvement on local machine

cc @ChrisRackauckas

🤖 Generated with Claude Code

Adds PrecompileTools precompilation workloads to significantly reduce
time-to-first-execution (TTFX) for common operations.

## Changes
- Added PrecompileTools as a dependency
- Added src/precompilation.jl with @compile_workload for:
  - DEQ model construction with Dense layers and SSRootfind solver
  - Lux.setup() for parameter/state initialization
  - DeepEquilibriumSolution constructor
  - Common utility functions (flatten, flatten_vcat, split_and_reshape)
  - check_unrolled_mode for both regular and fixed-depth modes

## Performance Improvements
Before:
- Package load: ~13.3s
- Model construction + setup: ~2.3s

After:
- Package load: ~11.0s (17% improvement)
- Model construction + setup: ~0.017s (135x improvement!)

The forward pass TTFX depends on the specific solver used (NonlinearSolve,
OrdinaryDiffEq) which are optional dependencies - users should expect
solver-specific compilation on first use.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@ChrisRackauckas ChrisRackauckas merged commit 03da7a9 into SciML:main Dec 30, 2025
2 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants