Compile time plan

Compile time: plan

Issue #322, 0.8. Status: in review, #384 (steps 3, 4, 5 and part of 6, CUBIN in the disk cache, and the CPU's chunk wrappers compiled with their kernels). The AhmedML sweep's cold compile on an M1 Max is 10.1 s on the CPU, from 15.5, and 9.2 s on Metal, from 12.7; on an RTX 6000 Ada host 36.5 s on CUDA, from 48.0, and 26.6 s on its CPU, from 37.1.

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.

Walks of the tree

cProfile counts each resumption of a generator as a call, so it overstated the IR walks. Timed directly, a walk costs 0.56 µs a node. The ~21–25 whole-tree walks a variant are about 2.5 s of the sweep. The duplicates that needed no design change are gone (step 2 in part); sharing one walk among passes would mean restructuring them.

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. Done:

   • Measured first: 26% of the sweep's lowered IR (52% of _advect, 62% of _find_from_points) sat in arms that integer and Boolean constants decide against: class constants compared, mask bits tested, and and chains with a false constant.
   • The fold: lowering now lowers only the arm such a condition selects, as C++'s if constexpr. The other arm is validated but not lowered, so it inlines nothing, requires nothing of the backend, and binds no names. Floating conditions are left to run.
   • Results: _advect went from 29K to 13.8K nodes and from 1.9 s to 1.3 s, and the sweep from 11.8 s to 10.8 s. The outputs are the same and the warm times within noise.

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. Done: each statement's own expressions take the copies known where it stands, and its bodies inherit them. Substitutions copy only the nodes above a replacement (substitute_ir), not every node of a changed statement. optimize_ir went to 0.58 s and the sweep to 11.8 s.

6. Lowering itself (3.6 s): what the profile finds in ast_transform, the lowering mixins and template_rewrite. Done in part:

   • each inlined function's bindings, validation, return count and assigned names are now found once a lowering, not at every call;
   • each template method's names are found once, not on every lowering;
   • the launch bounds no longer walk the body when they read only parameters.

   Together these went from 1.17 s to 0.27 s. Left: AST to IR, copy_ast, _prune_synthetic and ast.fix_missing_locations.

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. Measured: for a trivial kernel, O3 takes 1.5 ms and MCJIT's code generation 4.7 ms; the setup around them takes under 0.1 ms. Code generation costs the same at levels 1 to 3, and only level 0 (fast instruction selection and register allocation) is cheap. So the tail's LLVM time is real work, not overhead.

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).

Metal

Measured on an M1 Max, with Tack's disk cache off. Apple's Metal compiler keeps its own shader cache across processes. To measure a cold compile on Apple's side too, each MSL source was prefixed with a comment unique to the run.

• Totals:
  • with these steps, 9.5–10.3 s;
  • before steps 4–6, 12.7 s;
  • when Apple's cache holds the sources (an earlier process with the same Tack), 6.7–7.4 s.
• Of the 9.5 s:
  • Apple's compile takes 2.8 s (0.09 s from its cache);
  • generating MSL in Python takes 1.2 s;
  • lowering takes 1.15 s, the template rewrite 0.94 s, verification 0.62 s and optimize_ir 0.42 s.
• What follows:
  • Tack's own Python is most of Metal's cold compile, and nearly all of it once Apple's cache has the sources.
  • The steps that lighten lowering and the IR passes help Metal as much as the CPU.
  • Tack's disk cache stores MSL on Metal, so even with it warm, Metal compiles the source, which Apple's cache usually serves.

MSL generation

Profiled over the 301 kernels the Metal sweep compiles (their IR replayed offline, since generating MSL needs no GPU), MSL generation took 0.86 s, of which writing the text was 0.1 s. The rest:

• renaming the bindings, 0.35 s;
• a check for stores in sequential loops that was quadratic in loop nesting;
• the generator's own checks of the IR it is given.

The renaming no longer collects every name first or copies nodes through copy.copy; the store check is one walk; and the participation check runs only for kernels with collectives. MSL generation now takes 0.54 s, with the MSL of all 301 kernels byte-identical. The renaming is shared by CUDA, HIP and Level Zero.

CUDA

Measured on an RTX 6000 Ada with an EPYC 7453, NVRTC 13.4, the branch at 0716016a. The driver's JIT cache was off (CUDA_CACHE_DISABLE=1), which also turns off NVRTC's own cache (see below).

sweep, cold	base	branch
CPU	37.1 s	26.6 s (-28%)
CUDA	48.0 s	36.5 s (-24%)
of CUDA's: Tack's Python	20.9 s	10.3 s
of CUDA's: NVRTC (CUDA C to PTX)	18.2 s	17.3 s
of CUDA's: the driver (PTX to SASS at load)	7.8 s	7.8 s

The branch halved Tack's share on CUDA. NVRTC and the driver are now 69% of the cold compile. Every output is the same as the base's.

What else this showed:

• NVRTC caches too. NVRTC 13.4 caches its compiles across processes in the driver's cache directory, keyed on the preprocessed source. With the cache warm, NVRTC took 6.8–8.1 s instead of 17.3, about 23 ms a kernel.
• NVRTC's fixed cost is about 35 ms a program, even for a one-line kernel, so about 11 s of the 17.3 is that cost paid 313 times. Compiling 49 kernels as one program took 36% less, but needs the kernels known before the first call: a design question.
• Storing CUBIN instead of PTX (NVRTC compiles for the device's architecture where it knows it) costs the same cold. But a stored CUBIN loads in milliseconds, where the driver spends about 2 s on 49 kernels of PTX. With Tack's disk cache warm and the driver's cold (CI, a fresh node, a container), the sweep goes from 10.6 s to 2.1 s. In #384.
• --Ofast-compile cuts the cold sweep by 21%, but makes some kernels up to 3.9× slower when warm. Not pursued, not even as an option.

With the disk cache warm

The same sweep with Tack's disk cache written by an earlier process:

• The slow path is 1.0 s, against 10.1 s cold. It reads lowered IR (0.36 s), loads stored variants (0.34 s) and builds variant keys (0.26 s).

• First calls still cost 5.7 s more than warm calls, of which that compiling is 1.0 s. The other 4.9 s is one-time derivations cached on the topology: faces for extract_surface (1.2 s), edges for extract_edges (1.0 s), and others for tetrahedralize, cell geometry, point gradients, cell quality and the locator's leaves (0.2–0.6 s each).

• So with a warm cache, compiling is no longer what makes first calls slow. The derivations are their own item.

• Avoiding compiles altogether has two routes:

  • a cache warmed before first use, for a site or CI; the CPU's target names the host CPU, so a shipped cache would serve few machines;
  • compiling ahead in the background, which can overlap only the native stages (LLVM, Apple's compiler, NVRTC) and must guess the next kernels.

  Both are possible later; neither is planned here.

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.