Skip to content

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) → i64
  • IRConstant(3.14) → f32 unless it is part of a weak literal expression that meets an f64 operand, an f64 cast, or an f64 store, atomic or local; then every float literal in that expression is retyped to f64 (_meet, _adopt, _retype), following NumPy's NEP 50. Nodes carry _literal and _weak flags for this. Expressions without f64 operands annotate exactly as f32 literals would
  • IRArrayLoad → element type of the array
  • IRBinOp → promoted type, except integer / → f32, and shifts or integer ** → the left/base type independent of the count/exponent type
  • IRCast(value, ScalarType) → the target ScalarType
  • IRCall("sqrt", ...) → f32 (or f64 if any arg is f64)
  • IRCall("abs", [int_arg]) → preserves integer type
  • IRCall("min"/"max", ...) → promoted type of arguments
  • IRCall("pow", [int_base, int_exponent]) → the base type; otherwise the promoted floating precision
  • IRCompare, IRBoolOp → i32
  • IRName referencing 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).