Skip to content

IR Design

The Tack IR (ir.py) is a tree of Python objects with plain attributes. ir_traversal.py declares structural children explicitly and provides shared walks and rewrites. ir_verify.py checks invariants at compilation boundaries. Control-flow-sensitive passes retain their own ordered walks.

Structure

IRModule
  └── functions: list[IRFunction]
        ├── name: str
        ├── params: list[IRParam]   (name, type_annotation, _is_array, etc.)
        └── body: list[IRNode]      (statements)

Node Types

Loops and Control Flow

Node Description
IRParallelFor Outermost loop → GPU thread dispatch. Has var, start, end, body.
IRSequentialFor Inner loop → C for-loop. Has var, start, end, step, body.
IRWhile While loop. Has condition, body.
IRIf Conditional. Has condition, then_body, else_body.
IRBreak Break out of loop.
IRContinue Continue to next iteration.
IRReturn Legacy return node; rejected in kernel IR. Device returns are eliminated during inlining.

Expressions

Node Description
IRBinOp Binary operation: left op right (op is +, -, *, /, //, %, **).
IRUnaryOp Unary: -x, not x.
IRCompare Comparison: left op right (op is <, <=, >, >=, ==, !=).
IRBoolOp Boolean: a and b, a or b.
IRIfExp Ternary: a if cond else b.
IRCall Math builtin call: sqrt(x), sin(x), etc.
IRCast Type cast: int(x), float(x), tack.f64(x), etc. dtype is a ScalarType.
IRConstant Literal: 42, 3.14. dtype is a ScalarType, set by type annotation pass.
IRName Variable reference.
IRAttribute Attribute access: obj.attr (used for array.shape).

Array Operations

Node Description
IRArrayLoad Read: array[index]. array is an IRName, index is any expression.
IRArrayStore Write: array[index] = value.
IRAtomicOp Atomic: atomic_add(array, index, value).
IRDimSize Dimension query: array.shape[dim] (resolved to constant before codegen).
IRTextureSample Texture read: tex.sample(u, v, w). Stores array_name, coords, shape.

Debug Mode

Lowering emits these only for a kernel lowered with debug=True (tack.set_debug(True) or TACK_DEBUG=1); see lang/ir_bounds.py.

Node Description
IRBoundsCheck An access's index: flat as i64 when each of indices lies in [0, extent) of the extent beside it, otherwise -1, after recording the first violation (site, an index into the function's bounds_sites) in the hidden errors array.
IRArrayExtent An array parameter's extent at run time: LEADING (dimension 0), ELEMENTS (a flat index's bound) or STORAGE (scalar components). Never compiled in: ir_resolve turns it into a load from the hidden extents array, so checking does not specialize on lengths.

GPU Primitives

Node Description
IRSharedAlloc Allocate threadgroup memory: tack.shared(dtype, size) or tack.shared_like(array, size). Has optional array_name for shared_like (resolved in ir_resolve).
IRLocalAlloc Allocate per-thread memory: tack.local_array(dtype, size) or tack.local_array_like(array, size). Has optional array_name for local_array_like (resolved in ir_resolve).
IRBarrier Threadgroup sync: tack.barrier().
IRThreadId Local thread index: tack.thread_id().

Other

Node Description
IRAssign Assignment: target = value. After type annotation, has _resolved_type.
IRPrint Debug print (emits printf on CPU/CUDA/HIP, no-op on Metal).

Annotations

IR nodes are annotated by various passes:

  • IRParam.type_annotation — ScalarType (f32, i32, etc.), set by type inference.
  • IRParam._is_array — True for array/texture params, False for scalars.
  • IRParam._is_texture — True for Texture3D params.
  • IRParam._texture_shape — (W, H, D) tuple for texture params.
  • IRAssign._resolved_type — ScalarType or None, set by type annotation pass.
  • Expression dtype — ScalarType on all expression nodes (IRBinOp, IRArrayLoad, IRCall, IRCast, etc.), set by the type annotation pass. Codegens read this instead of reimplementing type inference.
  • IRTextureSample.shape — (W, H, D), set by IR resolve.
  • IRFunction.bounds_sites — debug mode only: one ir_bounds.BoundsSite per checked access (kernel or device function, file, line, source text), set by lowering and read by debug_bounds to word the IndexError. The function also gains two hidden array parameters, __tack_errors__ and __tack_extents__, after the kernel's own.

Invariants

  • The IR has exactly one IRParallelFor at the top level (the thread dispatch loop).
  • Sequential loops appear only inside the parallel loop's body.
  • IRDimSize nodes are resolved to IRConstant before codegen, except in the host-evaluated parallel loop end, which remains dynamic. IRArrayExtent nodes are resolved to loads of the hidden extents array.
  • After type inference, every IRParam has type_annotation and _is_array.
  • After scalar packing, no scalar params remain (GPU backends only).

verify_ir(function, stage) enforces structural, loop-target, binding, resolution, parameter, and annotation checks. A failed check raises IRVerificationError with the kernel, stage, node kind, and tree path. See IR Passes for the stages and their limits.

Shared Traversal

CHILD_FIELDS records each node's child attributes and their roles: expressions, optional expressions, statement lists, parameter lists, and function lists. walk_ir(root) yields nodes in preorder; transform_ir(root, rewrite) rewrites children before applying rewrite(node). Rewrites mutate the tree in place and preserve metadata on unchanged nodes. A replacement's children are not visited again. Repeated references to a node are visited at each occurrence; cycles and unregistered node kinds raise errors. copy_nodes=True copies each occurrence before rewriting it, preserving metadata while allowing statement-specific substitutions. Copy propagation uses this mode so a later rewrite cannot alter an earlier shared expression.

Only structural children are followed. Types, texture extents, cached shape dependencies, and debug/source metadata do not become accidental subtrees. Resolution, scalar packing, copy substitution, and shape-dependency queries use this schema. Analyses that depend on statement order or control-flow context must still implement that logic explicitly.