Specialization-time metaprogramming — design notes

Motivation¶

Four related use cases all need the same underlying mechanism: generating specialized code from a compile-time-known shape or arity, instead of hand-duplicating code per arity — a problem that recurs elsewhere too (multi-dimensional __getitem__, array-API-style variadic functions like stack/concatenate, and per-field codegen such as constructors or GC traversal for fixed-size structs).

benchmark — a @blue.generic microbenchmarking helper. fct and repeat are bound as blue generic parameters, so benchmark[kernel, 100] fully specializes: fct becomes a statically known callee, repeat an unrollable compile-time constant.

from time import time
from operator import OpSpec
from __spy__ import UNROLL_RANGE

@blue.generic
def benchmark(fct, repeat):
    @blue.metafunc
    def meta(*m_args):
        def impl(*args: m_args.types):
            t_start = time()
            for _ in UNROLL_RANGE(repeat):
                fct(*args)
            return (time() - t_start) / repeat
        return OpSpec(impl)
    return meta

partial — a functools.partial-style utility, restricted to blue bound arguments. A general, runtime-value partial needs a struct/closure object plus an __call__-style operator overload to recombine bound and call-time values, which is a separate, larger feature and out of scope here.

@blue.generic
def partial(fct, *args):
    @blue.metafunc
    def meta(*m_args):
        def impl(*rest: m_args.types):
            return fct(*args, *rest)
        return OpSpec(impl)
    return meta

vectorize — turn a scalar kernel into a SIMD loop over N pointers/parameters, with explicit instantiation (fct[SIMD[dtype, W]](...)) rather than implicit dispatch:

from simd import SIMD, ptr_simd_load, ptr_simd_store, simd_width_of
from __spy__ import UNROLL_RANGE

@blue.generic
def load_simd_args(dtype, simd_width, types):
    @blue.metafunc
    def meta(*m_args):
        N = len(types)
        new_types = (*(
            SIMD[dtype, simd_width] if types[i] == ptr[dtype] else types[i]
            for i in UNROLL_RANGE(N)
        ),)
        RetT = tuple[*new_types]

        def impl(*args: types, i: int) -> RetT:
            return (*(
                ptr_simd_load[dtype, simd_width](args[k], i)
                if types[k] == ptr[dtype] else args[k]
                for k in UNROLL_RANGE(N)
            ),)
        return OpSpec(impl)
    return meta

# load_scalar_args: same shape, using args[k][i] instead of ptr_simd_load(...)

@blue.generic
def vectorize(fct, dtype, simd_width=None):
    @blue.metafunc
    def meta(*m_args):
        W = simd_width if simd_width is not None else simd_width_of[dtype]
        types = m_args.types

        def impl(*args: types, pout: ptr[dtype], n: int) -> None:
            i = 0
            while i + W <= n:
                fct_args = load_simd_args[dtype, W, types](*args, i)
                result = fct[SIMD[dtype, W]](*fct_args)
                ptr_simd_store(pout, i, result)
                i += W
            while i < n:
                fct_args = load_scalar_args[dtype, types](*args, i)
                pout[i] = fct[dtype](*fct_args)
                i += 1
        return OpSpec(impl)
    return meta

def kernel0[T](x: T, y: T, param0: i32) -> T:
    return x + param0 * y

def main() -> None:
    # alloc ptr_x, ptr_y, ptr_out
    ...
    vectorize[kernel0, f64, 4](ptr_x, ptr_y, 42, ptr_out, length)

convert — a per-field type-converting tuple transform. Small, but exercises every piece of the mechanism cleanly; used as the running example below.

@blue.metafunc
def convert(m_t):
    T_t = m_t.static_type
    items_T = T_t.types
    N = len(items_T)

    @blue
    def target_type(i):
        return i32 if items_T[i] == f64 else f64

    new_types = (*(target_type(i) for i in UNROLL_RANGE(N)),)
    new_tuple_type = tuple[*new_types]
    converters = new_types

    def impl(t: T_t) -> new_tuple_type:
        return (*(converters[i](t[i]) for i in UNROLL_RANGE(N)),)

    return OpSpec(impl)

Two kinds of tuple¶

A distinction that shaped the mechanism, confirmed by testing directly against main: interp tuples (W_InterpTuple) back blue varargs (*m_args, *items_T, m_args.types) and are blue-only — they cannot appear in compiled (red) code at all. Static tuples (spy/vm/modules/_tuple.py) are struct-backed, real red values: a tuple-typed parameter compiles to an ordinary by-value struct, indexing (t[0]) lowers to a plain generated field getter with no boxing, but __getitem__ requires a blue index. The only proven constructor is the literal display, (x, y).

Consequently: anything derived from m_args.types that needs to be a red parameter or return type must first become a static tuple[T0,...,Tn], and any per-element transform running inside a compiled body must reduce to already-working static-tuple construction/indexing with blue indices — never a runtime loop over a dynamically-sized collection.

The mechanism: UNROLL_RANGE as a splice generator¶

“Unroll a loop,” “spread a tuple into a call,” “build a tuple,” and “build a tuple type” collapse into one mechanism once UNROLL_RANGE is generalized slightly beyond its for-loop role: it must also behave as an ordinary blue callable (so blue-interpreted code can call it like range), making this syntax legal in the same three positions Starred needs (tuple elements, call args, subscript args):

*(elt_expr for idx in UNROLL_RANGE(n_expr))

Two implementations, chosen purely by whether the surrounding code is blue-interpreted or compiled:

• Blue-interpreted (inside meta, or any @blue body): ordinary tree-walking — UNROLL_RANGE(n) behaves like range(n), results collected into a real blue tuple. No compilation, same category as _tuple.py’s existing build(idx) recursion.

• Compiled (inside impl): n_expr must be blue-known by the time this runs; for k in range(n), clone elt_expr with idx → Constant(k) and splice the n results into the parent list. Sibling of the body-duplication pass already prototyped for UNROLL_RANGE in for-loops on experimental/spyapi-aws-lambda.

A useful downstream effect: after substitution, blue_value[k] (closure-captured, now literal index) constant-folds away, the same way _tuple.py already folds getattr(self, "_item" + str(idx)). So converters in convert never exists as a runtime object — the compiled code is just i32(t[0]), f64(t[1]), ....

Bare splats are the identity-bodied special case: *args ≡ *(args[i] for i in UNROLL_RANGE(len(args))). Spreading a blue tuple (partial’s fct(*args, *rest)) needs no compiled-context splice at all — handled once, generically, where a Call’s argument list is built (callop.py::w_CALL): unpack a blue Starred operand’s items directly into the list. Spreading a red tuple (benchmark’s fct(*args), vectorize’s fct[...](*fct_args)) goes through the compiled-context splice. Same recognition, same “which color backs this tuple” branch either way — no dispatch on how the tuple arose.

Building a tuple type (tuple[*new_types]) needs nothing beyond Subscript accepting Starred: once new_types is a resolved blue tuple, splatting it into tuple[...] is the blue-splat branch above, on a Subscript node instead of Call — Subscript-of-a-@blue.generic already desugars to a call.

*args: <type-tuple-expr> desugars purely syntactically to a single args: tuple[*types] parameter — no body-rewrite pass, since static tuples are already free-indexing, real parameters (confirmed by testing). Every args[i] inside the body is just ordinary static-tuple field access.

Implementation sketch¶

Prerequisite, upstream of astcompile: parser.py::from_py_expr_Call currently builds call arguments with [self.from_py_expr(py_arg) for py_arg in py_node.args] — no ast.Starred handling anywhere in parser.py. Same for tuple displays and (once it exists) subscript args. Starred needs to become a real spy.ast node first.

In astcompile, compile_expr_Tuple / compile_expr_Call / (future subscript handler) scan items/args for two shapes:

• plain Starred(expr): resolve names in expr, wrap in a preserved ast.Starred marker. Blue-vs-red is undecidable here — depends on expr’s color, unresolved until doppler.

• Starred(GeneratorExp(elt, [comprehension(target, iter=Call(UNROLL_RANGE, [n_expr]))])): recognized via the same is_unroll_range_call check UnrollFor uses (below), lowered to ast.SpliceUnroll(loc, target, n_expr, elt), left in place in the list.

Consumption happens in vm/astframe.py::eval_expr_Call / eval_expr_Tuple (and the future subscript equivalent) — today a flat comprehension with no splicing. Both need a branch expanding ast.Starred/ast.SpliceUnroll entries:

• Plain ASTFrame: blue Starred(expr) unpacks expr’s items directly; a SpliceUnroll is tree-walked, same as UnrollFor’s interp path.

• DopplerFrame: red Starred(expr) (static type tuple[T0,...,Tn]) rewrites to expr[0], ..., expr[n-1]; SpliceUnroll clones elt per literal k — same clone-and-splice logic as UnrollFor’s doppler path, worth sharing the substitution utility between the two.

m_args.types¶

@blue.metafunc and *m_args already work today (spy/tests/compiler/test_metafunc.py shows the singular m_x.static_type in active use). m_args.types is the small, additive plural form: map static_type over m_args, producing an interp tuple of types — blue-only, feeding the splice mechanism, never itself a red annotation.

UNROLL_RANGE in for-loops¶

Retained unchanged as one of UNROLL_RANGE’s roles. A prototype exists on experimental/spyapi-aws-lambda (unroll_range.py + a diff in astframe.py/doppler.py), but can’t be ported as-is: main has diverged substantially (~1200 commits), and for-loop desugaring moved to a new, early, purely-syntactic pass (astcompile.py::compile_stmt_For) that unconditionally lowers every for into a while before any blue value can be evaluated.

Implementation sketch¶

ASTCompiler has no access to self.vm — pure SymTable-driven name resolution and tree-shape rewriting, before doppler/redshift ever evaluates a blue value. So its role is recognition and marker insertion only, never evaluation. This isn’t merely a convenient place for the recognition, it’s the only one: ast.For is stage-restricted to "parsed", and compile_stmt_For erases it into a bare While with no trace of provenance. If UNROLL_RANGE isn’t recognized inside compile_stmt_For itself, that information is gone before doppler ever runs, and unrecoverable.

For the recognition, compile_expr_Name already turns a genuine imported-symbol reference into ast.NameImportRef(sym), with Symbol.impref an ImportRef(modname, attr). Resolving names first and checking the result rejects a shadowing local for free:

def compile_stmt_For(self, stmt: ast.For) -> list[ast.Stmt]:
    if isinstance(stmt.iter, ast.Call) and len(stmt.iter.args) == 1:
        compiled_func = self.compile_expr(stmt.iter.func)
        if (isinstance(compiled_func, ast.NameImportRef)
                and compiled_func.sym.impref == ImportRef("__spy__", "UNROLL_RANGE")):
            n_expr = self.compile_expr(stmt.iter.args[0])
            new_body = self.compile_body(stmt.body)
            return [ast.UnrollFor(stmt.loc, stmt.target, n_expr, new_body)]
    # ...existing For -> While desugaring, unchanged...

ast.UnrollFor is a new Stmt, valid only at >= astcompiled, carrying (target, n_expr, body) unexpanded. Consumption is one level down, in exec_stmt_* (shared base, overridden by DopplerFrame):

• Plain ASTFrame: UNROLL_RANGE(n) behaves like range(n) — a real loop, no duplication.

• DopplerFrame: evaluate n_expr (now blue-known), clone body per literal k in range(n), splice in place of the loop — the duplication logic already prototyped, re-targeted to this entry point.

This sketch is self-contained: it touches neither eval_expr_Call/eval_expr_Tuple nor anything Starred-shaped, and can be built independently of, in either order relative to, the mechanism above. The one shared piece of logic is the recognition check itself — worth factoring into a standalone is_unroll_range_call helper even while only compile_stmt_For calls it, so SpliceUnroll reuses it later rather than re-deriving it.

Parser scope note¶

Because a blue splat and a red splat can appear in the same call (fct(*args, *rest)), Starred support should accept any number of entries, in any position, mixed with plain args, across all three positions from the start — not something narrower that would need generalizing the moment a second consumer shows up.

Summary tables¶

Split by track, in implementation order, matching “Suggested sequencing” below.

Track A — UnrollFor¶

Track B — Starred / SpliceUnroll¶

Suggested sequencing¶

The two sketches above share no AST node and no consumption method — only the is_unroll_range_call idiom. Genuinely independent tracks, buildable in parallel, in either order, by different people; sharing substitution/cloning code between A3 and B5 is worth doing opportunistically if they land close together, not a hard dependency.

Track A — UnrollFor¶

A1. Recognize and preserve. compile_stmt_For emits ast.UnrollFor instead of desugaring to While. No consumer yet.

A2. Interp consumption. ASTFrame.exec_stmt_UnrollFor runs like an ordinary range-based loop — usable inside any blue-interpreted code:

@blue.metafunc
def meta(*m_args):
    total = 0
    for i in UNROLL_RANGE(3):
        total = total + i   # ordinary blue execution, no duplication needed
    ...

A3. Doppler consumption. DopplerFrame.exec_stmt_UnrollFor clones and splices the body into compiled code — independent of anything Starred-shaped:

def stencil(x: ptr[f64], WEIGHTS: tuple[f64, f64, f64]) -> f64:
    acc: f64 = 0.0
    for i in UNROLL_RANGE(3):
        acc = acc + x[i] * WEIGHTS[i]   # literal i after substitution
    return acc                          # already-working static-tuple indexing

Track A is complete after A3 — benchmark’s repeat-loop only needs Track B for fct(*args).

Track B — Starred / SpliceUnroll¶

B1. Parser. ast.Starred in Tuple.elts / Call.args / (eventually) Subscript args, per the parser scope note. Nothing executable yet.

B2. astcompile recognition. Plain Starred(expr) and the UNROLL_RANGE-genexpr shape recognized and preserved as markers. Still nothing executable.

B3. m_args.types. Small, independent accessor — listed here because the examples below need it.

B4. Blue-interpreted consumption. Unlocks everything that stays inside meta:

@blue.metafunc
def convert(m_t):
    T_t = m_t.static_type
    items_T = T_t.types
    N = len(items_T)

    @blue
    def target_type(i):
        return i32 if items_T[i] == f64 else f64

    new_types = (*(target_type(i) for i in UNROLL_RANGE(N)),)  # blue SpliceUnroll
    new_tuple_type = tuple[*new_types]                          # blue Starred, Subscript
    ...

— and partial’s fct(*args, *rest), when both operands are blue tuples reached by closure.

B5. Doppler consumption + *args: <type-tuple-expr> sugar. Makes code inside a compiled body (impl, not meta) work:

def impl(t: T_t) -> new_tuple_type:
    return (*(converters[i](t[i]) for i in UNROLL_RANGE(N)),)   # red SpliceUnroll

def impl(*args: m_args.types):        # sugar -> args: tuple[*m_args.types]
    for _ in UNROLL_RANGE(repeat):    # Track A
        fct(*args)                    # red Starred splat, this step

After both tracks land¶

benchmark, blue-only partial, convert, and vectorize should all fall out with little additional work — vectorize also depends on the SIMD generic work (SIMD[T,N], simd_width_of) landing independently on simd-pr4.