Skip to content

AST Transform

The AST transformer (ast_transform.py) converts a Python function's AST into Tack IR. It is the heart of the frontend.

KernelTransformer is one class assembled from mixins, one module each in lang/lowering/, grouped by job. Each sets up the state it owns in its own _init_... method and reaches the others through self:

Module Mixin What it lowers
ast_transform.py KernelTransformer the visit_* walk (assignment, calls, subscripts, names, operators) and the expression machinery the others use (_visit_expression, _visit_ordered, _capture_value, _source_error)
lowering/loops.py Loops the parallel loop, ndindex, sequential loops, while
lowering/arrays.py ArrayAccess indexing and linearization, vector and structure-of-arrays elements
lowering/vectors.py VectorValues vector and matrix values, components, methods, scalarized assignment
lowering/inlining.py Inlining @tack.func calls, return restructuring, renaming
lowering/local_arrays.py LocalArrays tack.shared, tack.local_array
lowering/checks.py LoweringChecks checks on emitted IR: effects outside the parallel loop, values left in scalar positions, unbound names
lowering/bounds.py BoundsChecks debug mode's IRBoundsCheck (_checked)
lowering/values.py — VectorValue and TupleValue

Entry Point

# In kernel.py, when get_ir() is called:
from tack.lang.ast_transform import transform_kernel
ir_module = transform_kernel(ast_tree, vector_arrays=..., texture_arrays=...,
                             python_func=kernel.func, template_funcs=...)

transform_kernel creates a KernelTransformer (an ast.NodeVisitor) and visits the module. The result is an IRModule containing one IRFunction.

call_bindings.py supplies static device-call resolution from the defining Python callable and transformation-local template method map. AST-only callers can supply an explicit bindings dictionary. Bare kernel intrinsics remain available without a Python import; device functions require a bound Func object. No global function-name registry is consulted.

Original kernel source is validated by source_validation.py before transformation. Original @tack.func bodies are also validated before return restructuring or variable renaming, so unreachable unsupported syntax cannot disappear. Unsupported constructs raise UnsupportedSyntaxError, a NotImplementedError subclass, with the function name, construct, and line/column in captured, dedented source. The transformer's default visitor also rejects unknown nodes. Docstrings and pass are explicit no-ops; decorators and annotations are host metadata.

KernelTransformer State

The transformer tracks several pieces of state:

  • _loop_depth — 0 at the top level. The for i in range(...) at depth 0 becomes IRParallelFor; nested loops become IRSequentialFor.
  • _in_parallel_loop, _parallel_loop_seen and _kernel_body — enforce the one-loop rule and the outside-the-loop rule. _visit_parallel_for accepts one for statement directly in the kernel body; a second one, or one inside a top-level if or while, raises UnsupportedSyntaxError, as does a kernel with none. While _in_parallel_loop is false, _visit_body passes each lowered statement to _check_outside_parallel_loop, which rejects an array store, atomic, barrier, block reduction or print (_effect_name) at the statement's source position. It runs on the lowered IR so effects from inlined device functions and template methods are seen; nested bodies are visited first, so the reported position is the innermost one.
  • _inline_counter — monotonic counter for generating unique variable names during @tack.func inlining.
  • _vector_vars — tracks which local variables are vectors or matrices (name → component count). _matrix_vars holds (rows, columns) for the ones that are matrices. _bind_vector sets both and _variable_value reads them, so use those rather than the dicts.
  • _vector_arrays — tracks which parameters are vector arrays (for scalarization), with _matrix_arrays beside it. The runtime passes a matrix array's entry as a (rows, columns) tuple; the constructor splits it. _array_shape(name) gives an element's shape either way.
  • _texture_arrays — tracks which parameters are textures (for tex.sample()).
  • _texture_origin — maps renamed texture names back to original parameter names (needed when textures pass through @tack.func inlining).
  • _shared_vars — names of tack.shared() and tack.local_array() variables (affects indexing — these are arrays, not scalars).
  • _pre_stmts — statements generated by expression-level inlining that must execute at the expression's original position. Reentrant body traversal preserves the surrounding expression's pending statements.

Expression Evaluation

_visit_expression isolates an expression's value and prerequisite statements. _visit_ordered visits operands left to right and materializes earlier loads before a later operand's inlined effects. Boolean and conditional expressions put those prerequisite statements inside IRIf guards, rather than moving them before the guard. Inlined while conditions are lowered into an always-entered loop with condition statements and a guarded break, so continue reevaluates the condition.

Augmented stores reuse a single evaluated address and retain the old load before right-hand-side effects. Sequential range bounds and steps are captured before the loop. CPU codegen also short-circuits pure IRBoolOp and IRIfExp expressions with branches and phi joins. Comparisons and logical expressions yield normalized i32 values 0/1.

See the kernel language contract for the supported subset and numerical decisions still open.

Loop Handling

Parallel For

At loop depth 0, for i in range(n) becomes:

IRParallelFor(var="i", start=IRConstant(0), end=IRName("n"), body=[...])

With range(start, end, step), Tack transforms the step into a contiguous-index parallel loop with a decomposition assignment:

# for i in range(0, 100, 3) becomes:
# for __step_idx__ in range(0, ceil(100/3)):
#     i = 0 + __step_idx__ * 3

Sequential For

At loop depth > 0, for loops become IRSequentialFor. The step attribute is preserved (used by codegens for the C for-loop increment). Bounds and steps are captured once before entry. Loops inside while are also sequential. Literal zero and negative steps are rejected; dynamic steps must be positive. A break from the outer parallel loop is rejected, while nested sequential break and both forms of continue are supported.

ndindex

A top-level for i, j in tack.ndindex(w, h) (or three dimensions) becomes one IRParallelFor that keeps its dimensions: dims names one index variable per dimension, slowest first, and extents holds their sizes as host-evaluated expressions, like the loop's end (their product). The body starts with i = __nd_0__ and j = __nd_1__ (plus a range's start, for a (start, end) pair), and nothing divides: each backend binds the dimension variables from a launch of the same shape (see lang/parallel_dims.py).

A kernel that uses workgroup primitives, a loop of four or more dimensions, and an ndindex nested inside the parallel loop keep the flat form, a loop over the product whose body recovers the indices:

j = __idx__ % h
i = __idx__ // h

Function Call Dispatch

visit_Call is the most complex method. It handles:

  1. Math builtins (sqrt, sin, etc.) → IRCall
  2. tack.atomic_add/min/max → IRAtomicOp
  3. tack.shared(dtype, size) → IRSharedAlloc
  4. tack.local_array(dtype, size) → IRLocalAlloc
  5. tack.barrier() → IRBarrier
  6. tack.thread_id() → IRThreadId
  7. print(...) → IRPrint
  8. int(), float() → IRCast
  9. tex.sample(u, v, w) → IRTextureSample
  10. Vector and matrix methods (.normalized(), .dot(), .cross(), .sum(), .transpose(), .inverse(), ...) → scalarized ops, in _emit_vector_method and _emit_matrix_method
  11. tack.vec2/vec3/vec4(...) → a VectorValue (_visit_vector_construct); tack.mat2/mat3/mat4(...), tack.mat(row, ...), tack.identity(n) → a VectorValue with a two-dimensional shape (_visit_matrix_construct, _visit_matrix_identity)
  12. @tack.func calls → inlined via _inline_func_call

A name that the function does not bind is resolved through CallBindings before anything else in visit_Name and visit_Attribute: one bound to a tack.constant lowers to an IRConstant (under an IRCast when the constant is typed), and math.pi/e/tau likewise.

Assignment Handling

visit_Assign checks for several special forms before falling through to a plain IRAssign:

  1. Shared/local allocation: smem = tack.shared(...), arr = tack.local_array(...)
  2. Vector construction: v = tack.vec3(a, b, c) (_try_parse_vector_construct)
  3. Multi-dimensional array store: array[i, j] = val → linearized index
  4. Tuple unpacking: a, b = func() → multiple IRAssign nodes. A target that is a subscript or a component gets its value through a temporary and is stored at its own position among the targets (_store_unpacked), after the whole right side has been evaluated.
  5. Components: v[c] = x, v.x = x, vf[i][c] = x, m[i, j] = x (_component_target turns the named forms into subscripts; _component_selector and _matrix_selector resolve the index)

Vector Scalarization

Vector operations are expanded at the AST level. A vector variable v with 3 components becomes 3 scalar variables v__0, v__1, v__2. Binary ops on vectors produce component-wise IR nodes:

# v = a + b  (where a, b are 3-vectors) becomes:
v__0 = a__0 + b__0
v__1 = a__1 + b__1
v__2 = a__2 + b__2

Array loads/stores are similarly scalarized: v = array[i] loads array[i*3], array[i*3+1], array[i*3+2].

During lowering a vector or matrix is a VectorValue: a list of component expressions with a shape, (n,) or (rows, columns). A source tuple (the indices of f[i, j], several results of a device function) is a TupleValue. Both are lists so that code mapping over components stays plain; the classes exist to tell them apart from each other and from a list of statements. Three rules follow from the representation:

  • Map, do not pass through. An operation given a VectorValue has to return one. _elementwise does that for anything componentwise, checks that the operands' shapes agree, and keeps the shape; use _like when building a result by hand.
  • Neither class may reach the IR. _check_single_values scans each lowered statement once and raises UnsupportedSyntaxError at the statement's position for one left in a scalar position. Without it the verifier reports "expected expr node" with a tree path.
  • Assignment is per component, but the right side is one value. _settle_components evaluates into a temporary any component that could observe an earlier component's write: a read of an earlier target component for a variable, any array load for an array element, since arrays may share storage.

Operations that use each component several times (@, determinant, inverse, outer_product) capture their operands into temporaries first, or the expression tree grows with every product in a chain.