IR Passes¶
Tack runs several IR passes between AST transformation and codegen. Each pass inspects the IR tree; transformation passes mutate a private copy.
Pass Order¶
1. ir_resolve — Replace compiled-in dimensions, set texture shapes, resolve array-like dtypes
2. type_inference — Annotate params with types from actual arguments
3. check_dispatch_types — Validate array dtypes against backend capabilities
4. scalar localization — Give scalar params and outer locals the loop assigns per-iteration local storage
5. workgroup participation — Prove the supported collective control-flow domain (GPU only)
6. ir_optimize — Conservative copy propagation
7. ir_pack_scalars — Group scalar params into array buffers (GPU only)
8. ir_type_annotate — Annotate scalar expressions and assignment storage types
9. ir_bounds.guard_accesses — Skip accesses whose checked index is negative (debug mode only, at the end of 8)
resolve_variant() runs the common passes only for a new compiled variant,
one neither in the backend's cache nor in the disk cache; a variant read
from disk carries its post-pass IR and runs none of them.
GPU packing works on a separate copy so launch-range resolution retains the
original parameter names. CPU skips packing. tack.inspect() also performs
the backend's texture decision, check_dispatch_types, scalar
localization and verification; its source preparation annotates once
before packing and again afterwards on GPU.
Copying and Template Ownership (ir_traversal.py)¶
Cached templates remain pristine. Variant preparation, assigned-scalar
localization, GPU packing and inspection use clone_ir() for their working
copies. This specializes Python deep copying for registered plain IR nodes
and list/dict containers. It copies every node attribute, including pass
annotations and metadata, rather than only structural children. One memo
preserves aliases and cycles across the entire graph, while mutable state is
independent of the original. ScalarType objects retain their identity;
other metadata uses Python's deepcopy protocol. Structural cycles remain
invalid and are rejected by the verifier.
Cache hits do not clone IR. All pass-boundary checks below remain in place. Adding a registered node with slots or a custom copying protocol requires revisiting the cloning fast path; the registered nodes currently use plain attribute dictionaries. Clone regression tests cover every registered kind, metadata ownership and cached dispatch.
uv run --no-sync python benchmarks/ir_clone.py --output /tmp/ir-clone.json
compares generic copying and specialized cloning on the pristine and typed
CPU IR of dataset filters: it runs a contour and a clip of VTK's wavelet once,
clears the variant cache, captures their kernels on a second run and
alternates copy methods on each graph. It reports median and individual
samples plus unique-node and occurrence counts. Setup, compilation and
filtering are outside those timings; this
measures copying cost, not total cold compilation or GPU performance.
Verification (ir_verify.py)¶
The production pipeline calls verify_ir() at these boundaries:
| Stage | Checks added to structural and loop checks |
|---|---|
lowered |
Valid statement/expression roles, known node kinds/operators, numeric constants, unique parameters, existing bindings, exactly one normalized top-level parallel loop, no store, atomic, block reduction, barrier or print outside it, and consistent break/continue targets |
resolved |
No dimensions or attributes in generated expressions; allocation dtypes and texture extents resolved |
inferred |
Scalar parameter types and array/scalar categories present; array accesses name array parameters or allocated arrays, including inlined pointer copies |
localized |
Recheck the invariants after assigned scalar parameters and outer locals become per-iteration locals |
optimized |
Recheck the invariants after copy propagation |
packed |
Recheck the invariants and ensure no scalar parameters remain on GPU |
typed |
Every generated scalar expression and scalar assignment has a type; logical results have i32 dtype |
The parallel loop end is evaluated by the host on each dispatch. Dimension queries in that expression remain legal and do not force specialization on a flat array's length. Array references are pointers and do not require a scalar expression dtype.
IRVerificationError reports the kernel, stage, node kind, and structural
path. Failed variants never reach compilation or enter the compiled cache.
Template verification runs before caching lowered IR; the remaining checks
run on variant misses. Warm dispatch performs no verification walks, and
neither does an entry read from the disk cache, which was verified when it
was built. Direct
low-level pass/codegen calls used to build IR fragments must invoke the
verifier themselves when they need the full kernel contract.
These checks establish tree and annotation invariants. Binding checks test whether a name exists anywhere in the function, not whether every path initializes it before a read. They do not prove bounds safety, alias safety of future optimizations, absence of races, barrier uniformity, or numerical equivalence. Differential tests and backend hardware validation remain necessary.
IR Resolve (ir_resolve.py)¶
Replaces IRDimSize(array_name, dim) with IRConstant(shape[dim]) using
the actual array shapes passed at call time. Also sets
IRTextureSample.shape from the Texture3D object. In debug mode it turns
each IRArrayExtent into a load from the hidden __tack_extents__ array
(three i64 extents per parameter, written by dispatch), so a debug kernel
does not specialize on the lengths it checks.
This pass runs early because subsequent passes (optimization, codegen) need concrete sizes.
Type Inference (type_inference.py, 81 lines)¶
Annotates each IRParam with:
- type_annotation — a ScalarType (f32, i32, i64, etc.)
- _is_array — True if the argument is an Array or Texture3D
- _is_texture — True if the argument is a Texture3D
Type rules:
- Array → dtype of the array
- Texture3D → dtype of the underlying array (always f32; tack.texture3d() rejects others)
- Python float or np.floating → f32 by default; auto-promotes to f64 if any array argument uses f64
- Python int or np.integer → i32 if it fits, else i64, else u64
Dispatch-Time Type Checking (type_inference.py)¶
check_dispatch_types() validates that all array argument dtypes are
supported by the target backend. This runs after type inference when
building a new compiled variant.
Each backend defines its supported dtypes (e.g., Metal excludes f64).
Unsupported dtypes produce a clear TypeError naming the kernel, parameter,
dtype, and backend.
Workgroup Participation (workgroup_participation.py)¶
GPU variant construction checks collective control flow after scalar
localization, before optimization and scalar packing. A structured uniformity
analysis follows assignments and branch joins, with monotone loop fixed
points for carried values and lane-dependent exits. It rejects unproven
participation before code generation and returns the full-256-lane launch
requirement, cached on KernelVariant. Cached dispatch checks logical counts
without repeating this pass. Public inspection follows the same boundary;
direct GPU generation checks mutable IR afresh.
Ordinary memory/texture loads are varying. Scalar arguments and immutable
scalar packs are uniform; pack_scalars marks generated parameters with
_is_scalar_pack so packed codegen retains that fact. Shape metadata and
collective results are uniform too. Conditional/short-circuit reductions and
while-condition reductions are restricted where codegen cannot preserve
evaluation. See the language contract and test_workgroup_participation.py.
This does not replace structural verification or prove race freedom.
IR Optimize (ir_optimize.py)¶
Copy propagation replaces subsequent uses of a single-assignment copy with its source when the source is never rebound in the containing block. Assignment counts include nested control flow, loop-variable bindings, and local/shared array bindings. Assignments remain in place, reads before the assignment are unchanged, and rewritten sequential loops retain their step.
This is useful after @tack.func inlining, which creates assignments such
as __func_x_0__ = x. Array aliases can then use the original parameter:
a = x followed by a[i] becomes x[i].
Tack's custom loop-invariant code motion and common subexpression elimination are disabled. An invariant address does not imply an invariant loaded value, and hoisting an assignment can change a zero-trip loop's behavior. Loads must also account for stores through overlapping arrays. These passes need memory and control-flow safety analyses before they can return. LLVM and vendor compilers continue to optimize generated code, without unconditional disjoint-storage promises on array parameters.
IR Type Annotate (ir_type_annotate.py)¶
Walks the IR and annotates every expression node with a dtype attribute
(a ScalarType). This also sets _resolved_type on each IRAssign.
After this pass, codegen backends read node.dtype directly instead of
reimplementing type inference heuristics. This eliminated ~40 lines of
duplicated _infer_c_type / _infer_expr_type logic per codegen backend.
Key rules:
IRConstant(42)→i32,IRConstant(2**31)→i64IRConstant(3.14)→f32unless it is part of a weak literal expression that meets anf64operand, anf64cast, or anf64store, atomic or local; then every float literal in that expression is retyped tof64(_meet,_adopt,_retype), following NumPy's NEP 50. Nodes carry_literaland_weakflags for this. Expressions withoutf64operands annotate exactly as f32 literals wouldIRArrayLoad→ element type of the arrayIRBinOp→ promoted type, except integer/→f32, and shifts or integer**→ the left/base type independent of the count/exponent typeIRCast(value, ScalarType)→ the target ScalarTypeIRCall("sqrt", ...)→f32(orf64if any arg is f64)IRCall("abs", [int_arg])→ preserves integer typeIRCall("min"/"max", ...)→ promoted type of argumentsIRCall("pow", [int_base, int_exponent])→ the base type; otherwise the promoted floating precisionIRCompare,IRBoolOp→i32IRNamereferencing an array param →None(array pointers aren't scalars)
Integer promotion preserves both full operand ranges; signed/u64 pairs require an explicit cast where promotion applies. True division and power use their distinct rules instead. Negative literal integer exponents are rejected here. See the language contract.
The pass joins all assignments to each local to one storage type, then
annotates expressions using the settled environment (var_name → ScalarType).
An IRBoundsCheck is i64.
Bounds Guards (ir_bounds.py, debug mode)¶
annotate_types ends with guard_accesses, which does nothing unless the
function was lowered with bounds checks (bounds_sites). Each user array
access whose index is an IRBoundsCheck is skipped when that index
is negative: a load becomes index >= 0 ? load : 0 with a zero of the
load's own type (hence after typing), and a store or atomic runs under
if index >= 0. The index is evaluated twice, which is harmless because
indices have no effects but the check's record of the first violation.
Accesses to the hidden arrays and to GPU scalar packs are left alone. It
runs once per function (_bounds_guarded), so annotating again before and
after packing does not guard twice.
Scalar Packing (ir_pack_scalars.py)¶
GPU backends only. Groups scalar parameters by type into packed array buffers to reduce buffer binding count (critical for Metal's 31-binding limit).
Algorithm:
1. Collect all params where _is_array == False, grouped by type
2. Create new IRParam for each group: __pack_f32__, __pack_i32__, etc.
3. Rewrite IRName("scalar_param") → IRArrayLoad(IRName("__pack_f32__"), IRConstant(idx))
4. Remove original scalar params, append pack params
The pass runs on a clone_ir copy of the IR (to preserve the cached original)
and stores pack_info metadata for the dispatch layer. Pack array buffers
are allocated once and cached alongside the compiled kernel — subsequent
calls just update the scalar values via from_numpy.
clone_ir preserves ScalarType identity, matching its __deepcopy__
protocol (type map lookups rely on object identity).