Compile time plan

Compile time: plan

Issue #322, 0.8. Status: step 3 done, on branch core/compile-time (not yet pushed or in a pull request); the sweep's cold compile time is 13.8 s, from 15.5.

A kernel's first call lowers it from Python to Tack IR, runs Tack's passes, generates backend code and compiles it. The disk cache (tack.cache) makes later processes skip all of that, but a cold first call pays it in full: on a first run, in CI, after a Tack upgrade, and in MPI jobs, whose default mode only reads the cache. This plan makes cold compiles cheaper, for large kernels and for the many small ones.

The workload

The workload is AhmedML's analysis sweep, benchmarks/vtk_compare/ahmed_sweep.py on volume_1_appended.vtu: 21.5M cells, hexahedra, wedges and polyhedra in one mixed topology. The sweep runs 24 filters:

• gradients, flow quantities, projections, quality, integration;
• surface, slice, contour, clip, threshold, connected components, edges, tetrahedralization;
• locator, probe, streamlines;
• surface normals and feature edges.

It touches most of tack.dataset. All numbers below are on an M1 Max on the CPU with TACK_CACHE=off (four runs, 15.2–15.7 s).

Where the time goes

• Compiling takes about 15.5 s over 244 kernels, 337 variants and 309 lowerings.
• The sweep's first calls total 20.6 s cold, against 9.0 s with a warm disk cache.
• With a warm cache the slow path still costs 1.2 s: loading stored variants 0.86 s and lowered IR 0.35 s.

stage	s	share
Python lowering: AST to IR 2.2, template rewrite 1.4	3.6	23%
verify_ir: 5 times per variant, once per lowering	2.6	17%
Tack IR passes and checks: optimize_ir 1.2, clone_ir 0.3, _localize_outer_scalars 0.3, resolve_ir 0.3, annotate_types 0.3, shape_signature 0.2, the rest 0.5	3.2	20%
LLVM IR built and printed in Python (llvmlite.ir)	2.1	13%
LLVM itself: parse, O3, codegen and JIT	4.2	27%

About three quarters of the time is Tack's own Python, and only a quarter is LLVM.

Two populations of kernels

• Six large kernels take 7.1 s (45%). For them, LLVM is about 1 s of the 7.1; the rest is lowering, verification and passes.

  kernel	s	IR nodes
  _advect (streamlines)	2.67	29,081
  _find_from_points (locator search)	1.29	16,036
  _corner_gradients (point gradients)	0.99	7,586
  _values_from_points (probe)	0.79	2,982
  _test_points (locator)	0.72	3,977
  _kernel_points (locator)	0.63	8,298

• The tail is 208 kernels under 50 ms, 3.1 s in all, with a median of 13 ms. About half of the tail's time is LLVM's fixed cost per module (O3 0.86 s, codegen and JIT 1.02 s); the other half is Tack's fixed cost per variant.

Some filters compile many variants:

• slice: 45 variants, 10 of them of _lopez_trace;
• clip: 56 variants;
• extract surface: 28 variants.

Specialization diagnostics (#325) should say which are needed.

Since the issue was written

The issue's profile was of _advect alone: 47,849 nodes and 3.75 s, 29% of it in a recursive-generator IR walk.

• The walk has since become an explicit stack.
• Template methods now take only the parameters they read, and are resolved on demand (#298).
• _advect is now 29,081 nodes and 2.67 s.

So the issue's fourth item is done, and its shares need measuring again.

The plan

Each step is measured over the sweep (total, large kernels, tail) and checked by hashing every generated source before and after: the generated code must not change unless a step means it to.

1. A profiling harness in benchmarks/. It times every compile stage exclusively and attributes each to step, call, kernel and variant. It runs on the sweep and on a synthetic set for machines without AhmedML.
2. Profile the large kernels and the tail (cProfile under the harness) before changing anything, and re-rank the steps below by what the profile shows.
3. Verify less outside debug mode. Today verify_ir checks the lowered template and then the variant after resolve, inference, localization, optimization and typing. Keep the lowered and final checks; run the rest in debug mode and in the test suite. That could save up to about 2 s of the 2.6. Done: verify_stage runs the stages between in debug mode or with TACK_VERIFY_ALL (which the test suite sets). typed repeats every earlier stage's checks but packing's, so a broken pass is still stopped before code generation. Verification went from 2.6 s to 1.1 s, compiling from 15.5 s to 13.8 s, and _advect from 2.7 s to 2.3 s.
4. Drop branches decided by compile-time constants, during lowering. In the old profile they were 29–31% of _advect's nodes, and every later pass and the code generator walk them. Measure again on today's _advect and _find_from_points.
5. One-pass copy propagation: one substitution map applied in one walk. optimize_ir is 1.2 s, 0.8 s of it in the six large kernels.
6. Lowering itself (3.6 s): what the profile finds in ast_transform, the lowering mixins and template_rewrite.
7. LLVM IR without llvmlite.ir's object model (2.1 s): emit the text directly, or keep the model and avoid printing and parsing it again.
8. LLVM's fixed cost per module, for the tail: an optimization pipeline sized to the kernel, and fewer engine finalizations. Measure what each does to run time before taking it.
9. The GPUs: the same harness on CUDA, HIP and Level Zero, where NVRTC, hipRTC and ocloc replace LLVM (NVRTC's time on very large functions varies from run to run).

Targets

• _advect: a cold first call of about 1.5 s, from 2.67 s.
• The sweep: cold compile time halved, to under 8 s on the CPU, with no change to the generated code except where a step means it (step 4).
• The tail: a median small kernel under 8 ms, from 13.