Skip to content

Template System

The template system enables zero-cost polymorphism: the same kernel source compiles to different GPU code depending on the concrete types passed.

Components

Three pieces work together:

  1. @tack.device_class (device_class.py) — marks a class as passable to kernels. Collects @tack.func methods into cls._tack_func_methods.

  2. @tack.func (func.py) — captures a function's AST and original Python callable for inlining and static binding resolution. Class methods are stored on the class by @tack.device_class.

  3. Template rewrite (template_rewrite.py) — AST pre-pass that resolves template parameters before IR transformation.

Template Rewrite Flow

When backend.execute() detects a @tack.device_class argument:

# In cpu.py / metal.py / etc.:
template_args = _detect_template_args(kernel, args)
# → {0: ("grid", grid_obj)}  (param index → (name, object))

effective_args = _expand_template_args(args, template_args)
# → removes template objects, inserts their array attributes and runtime scalars

ir_module = kernel.get_ir(..., template_args=template_args)
# → calls template_rewrite.rewrite_templates() before AST transform

What rewrite_templates Does

Given a kernel:

@tack.kernel
def avg(cs, data, out, n):
    for c in range(n):
        total = 0.0
        for v in range(cs.points_per_cell):
            total = total + data[cs.get_point_id(c, v)]
        out[c] = total / float(cs.points_per_cell)

Called with CellSetStructured3D(nx=50, ny=50, nz=50):

  1. Remove template param from signature: cs is removed
  2. Add array params: cs.connectivity (if it exists) becomes a new parameter __tmpl_cs_connectivity__
  3. Add runtime scalar params: cs.width (instance scalar) becomes parameter __tmpl_cs_width__
  4. Replace class constant access: cs.points_per_cell → ast.Constant(8) (baked in)
  5. Replace instance scalar access: cs.width → ast.Name("__tmpl_cs_width__") (runtime parameter)
  6. Replace method calls: cs.get_point_id(c, v) → inline the method body with self references resolved

The rewrite returns the AST and a local map of resolved methods. Method names include the template parameter name, so repeated object arguments retain distinct synthetic array/scalar layouts. Generated calls carry an AST marker to distinguish them from user functions with the same spelling. The map belongs to this transformation; no registration or global cleanup is required. Resolved methods retain the original Python callable for globals/closure lookup and do not mutate the class's captured source.

Attribute Classification

classify_template_attrs(obj) splits an object's attributes into three categories:

scalars = {"points_per_cell": 8}              # class variables → compile-time constants
arrays = {"connectivity": <Array>}             # instance Arrays → kernel parameters
runtime_scalars = {"width": 512, "height": 512}  # instance scalars → runtime parameters
  • Class-level scalars (defined on the class, e.g., points_per_cell = 8) become AST constants, baked into generated code, and are part of the cache key. Changing them triggers recompilation.
  • Instance scalars (set in __init__, e.g., self.width = width) become synthetic kernel scalar parameters. They are passed at dispatch time and do not appear in the cache key — changing them reuses the compiled kernel.
  • Instance arrays become synthetic kernel buffer parameters (as before).

This distinction avoids unnecessary JIT recompilation when only runtime values (like image dimensions) change between calls.

Nested Templates

A template may hold templates as instance attributes. template_tree(obj) lists (path, object) for the object and everything nested in it, depth first, children by name, and refuses a cycle. classify_template_attrs flattens the tree: a nested object's constants, arrays and runtime scalars are named by their path joined with __ (base__step), so the runtime's argument expansion and vector-array detection, which read the flattened dicts, need nothing new, and the synthetic parameters follow one order. rewrite_templates builds a _Node per object (its own constants, the synthetic names of its own arrays and scalars, its resolved methods and function attributes, its children); every resolved method of the tree takes the whole tree's synthetic parameters, and self.a.b.method() and self.a.b.attr resolve by walking the children (_resolve_call, _resolve_attribute). A chain stops at the first attribute that is not a template, so self.data.shape[0] still reads an array's shape.

A dispatch reads a template several times -- its key, its expanded arguments, its vector arrays -- so resolve_variant opens a template_scope, in which each object's tree, classification, function attributes, structure and cache-key part are found once (memoized). The objects are the call's arguments, alive throughout, so their ids cannot be reused within the scope.

Only a call the dispatch fast path cannot answer enters the scope. A warm call reads each template into its signature instead (_template_signature in runtime/kernel_utils.py): the class token, the class constants' current values, and per attribute what the key and the expansion read, following an attribute layout learned once per class (_template_layouts, relearned when an instance's attribute names or types differ). It then expands the arguments through the plan stored with the signature.

@tack.func Inlining

When the AST transformer encounters a @tack.func call, it runs _inline_func_call():

Calls are first resolved by call_bindings.py using the defining callable's globals and closure cells, or the marked template-method map. Function identity determines the body. Arguments are lowered under the caller's bindings; body traversal temporarily switches to the callee's bindings. Validation uses the same resolver, including intrinsic shadowing checks. Only static module attributes are followed; object properties and runtime callables are unsupported. The active function identities detect recursion.

  1. Generate unique suffix: __{func_name}_{param}_{counter}__
  2. Build rename map: callee param names → unique names
  3. Deep-copy callee AST and rename all variables
  4. Propagate metadata: vector variables, texture arrays
  5. Emit parameter assignments: __func_x_0__ = caller_arg (skipped for texture params — they reference the original kernel parameter)
  6. Visit renamed body to produce IR statements
  7. Return result variable as IRName

The inlined statements are collected in _pre_stmts and hoisted before the statement containing the call.

Multi-Return Inlining

For functions returning tuples:

@tack.func
def minmax(a, b):
    if a < b:
        return a, b
    return b, a

The inliner detects n_returns > 1 and creates multiple result variables: __minmax_ret_0_0__, __minmax_ret_1_0__. Return statements become assignments to these variables.

Texture Propagation

When a texture array is passed through @tack.func, the inliner must:

  1. Propagate _texture_arrays: renamed param → same shape metadata
  2. Track origin: _texture_origin[renamed] = original_param_name
  3. Skip assignment: texture params can't be assigned to pointer variables
  4. Use original name: IRTextureSample.array_name uses the original kernel param name (via _texture_origin), not the renamed inline name

Local Array Propagation

When a tack.local_array is passed to a @tack.func, arrays can't be assigned in C (int arr2[8] = arr1; is invalid). The inliner handles this by aliasing the caller's array name directly:

  1. Before the renamer runs, check if the caller arg is in _shared_vars
  2. If so, override the rename map: rename_map[param_name] = caller_arg_name
  3. The renamer substitutes the callee's param name with the caller's array name throughout the inlined body
  4. Skip the parameter assignment (like textures)

This means the inlined body's pts[v] = ... compiles to caller_array[v] = ... with no intermediate variable. The same mechanism works for tack.shared() arrays.

Cache Keys

Each unique combination of template types produces a different compiled kernel. The cache key includes:

  • Kernel object identity (held weakly by the backend cache)
  • Type signature and array/scalar/texture category of all params
  • Actual template class identity, typed compile-time constants, array metadata, and runtime scalar attribute names
  • The paths and classes of the templates it holds (template_structure), whose tokens also retire the variants when a nested class is collected
  • Vector array widths
  • Texture shapes
  • Dimension sizes baked into code by shape resolution

Runtime scalar values do not specialize, but changing the set of runtime attributes changes the parameter layout and requires another variant.

This means avg(structured_cs, ...) and avg(explicit_cs, ...) produce two separately cached compilations with different generated code.

The disk cache cannot use class identity across processes, so it keys a template by content instead: the names along its class's MRO, the sources of its @tack.func methods and function attributes, and its constants. Each name lowering resolves is also recorded with a fingerprint of what it was bound to, so an edited method or class constant misses.