Skip to content

TileFoundry Spec — Codegen

Codegen turns verified, lowered tir.PrimFunctions into a loadable artifact. It owns the whole producer side of the build: emitting per-target source, assembling each target's translation unit, and linking those units into one host-callable shared library. Loading that artifact and exposing it as a RuntimeModule is owned by runtime.

flowchart LR
    TIR["verified <b>tir.PrimFunction</b>s"]
    Emit["per-target <b>Emitter</b>"]
    LM["<b>LinkableModule</b><br/>(per target)"]
    Link["<b>link</b>"]
    Linked["<b>LinkedModule</b><br/>artifact + metadata"]
    RM["<b>RuntimeModule</b><br/>(see runtime)"]

    TIR --> Emit --> LM --> Link --> Linked
    Linked -. runtime load .-> RM

1. Pipeline

  • Input is verified TIR. HIR Ops MUST NOT reach codegen.
  • A module's functions are grouped by their target (cuda / cpu). Each group is emitted by its target's emitter into one LinkableModule.
  • The link step compiles every LinkableModule with its own toolchain and links them into one LinkedModule — a host-callable shared library plus the host-visible metadata the loader needs.
  • Codegen does not run passes, does not load or launch device code, and does not own the user-facing entry points (compile / build / jit).
  • Host / device boundary. A host LinkableModule MUST NOT reference CUDA or CuTe symbols or types. A CUDA LinkableModule owns the kernels and their C-ABI launch shims. The host module invokes device code only through that C-ABI shim.

The target-specific emitter behavior — how a cpu vs cuda function emits, the dispatch and shape-scalar ABI, program-shape / dynamic-CTA accessors, and the ShardLayout runtime mapping — is owned by target.

2. Emitter

An emitter walks a verified tir.PrimFunction and produces source plus the metadata the link step needs.

2.1 Emitter registry

def get_emitter(target: str) -> Emitter: ...    # resolve the emitter registered for a target
# Emitter: Callable[[tuple[PrimFunction, ...]], LinkableModule]
  • constraints:
  • each target registers one emitter, resolved by target name.

An emitter MUST consume only TIR and MUST return a LinkableModule for its target. The emitter file layout mirrors the IR file layout (codegen/<target>/tir/... parallels ir/tir/...); the mirror rule is owned by code-organization.

2.2 Per-Op handler registry

def handler(call: Call, ctx: CodegenContext) -> None: ...   # Call (wrapped Op inside Evaluate) + CodegenContext
  • constraints:
  • a handler is registered on a target's emitter with the per-target register_codegen_* decorator (visitor-registry §6); dispatch (matching Evaluate and selecting the handler) is owned by visitor-registry.

Dispatch is owned by visitor-registry §6. A handler receives the Call (the wrapped Op inside Evaluate) plus a CodegenContext, and MUST emit through ctx.emit(...); raw print / direct file writes are prohibited.

2.3 CodegenContext

class CodegenContext:
    def emit(self, line: str) -> None: ...               # append a target source line
    def expr(self, node: Expr) -> str: ...               # render an Expr as a target expression string
    def dtype_to_cpp(self, dtype_name: str) -> str: ...  # backend dtype mapping
    def name_for(self, var) -> str: ...                  # allocate / return the target-side identifier for an SSA var
  • constraints:
  • the per-walk state object passed to handlers; the single source of truth for target-side type strings, so handlers do not read the IR for them directly.

A handler MUST NOT reach into the IR for type strings on its own; the context is the single source of truth. Other helpers MAY be added per target.

2.4 Effect Op dispatch

Effect Ops (Copy, Fill, Mma, tir.nn.*, ...) appear in Stmt position as Evaluate(op, args) rather than as Stmt subclasses. The walker matches Evaluate and dispatches on type(callable) through the handler registry. Handlers stay small; the runtime function they call carries the semantic load.

3. Runtime-owned op dispatch

Where more than one runtime template implements an op, codegen emits one uniform runtime op call, passing the operand ShardLayouts (and any codegen-static participant geometry) as compile-time template parameters. The runtime template dispatches on those layouts at compile time; codegen does not select a tier, compute a per-tier parameter, or carry the selection on the TIR op. This is the codegen side of the runtime-owned dispatch principle, whose contract lives in runtime.md §3. The target-side emission that produces these calls is owned by target.

4. Codegen products

4.1 LinkableFunction

One lowered function's pre-link source.

class LinkableFunction:
    name: str      # the function / kernel symbol
    source: str    # that function's emitted text
  • constraints:
  • MUST be the function / kernel symbol.
  • MUST be that function's emitted text.

4.2 LinkableModule

One target's pre-link translation unit.

class LinkableModule:
    target: str                              # the function target name (cuda / cpu)
    language: str                            # the source language (cu / cpp)
    source: str                              # the assembled translation-unit text
    functions: tuple[LinkableFunction, ...]  # the module's constituent LinkableFunctions
  • constraints:
  • MUST be the function target name (cuda / cpu).
  • MUST be the source language: cu for a CUDA translation unit, cpp for a host translation unit.
  • MUST be the assembled translation-unit text the link step compiles.
  • MUST list the module's constituent LinkableFunctions, in emission order.

A LinkableModule is a build artifact, not a runtime object and not a user-callable.

4.3 LinkedModule

The link output: a loadable artifact plus the host-visible metadata the loader needs.

class LinkedModule:
    library_path: Path                  # the produced shared library
    source: str                         # the assembled host + device source
    entry: EntryABI                     # the host-visible ABI of the module entry
  • constraints:
  • MUST point at the produced shared library.
  • MUST carry the assembled host + device source — the diagnostic source the runtime exposes as RuntimeModule.source (runtime).
  • MUST be the host-visible ABI of the module entry.

The entry EntryABI is a host-visible ABI metadata type owned by runtime; codegen references it on LinkedModule and MUST NOT redefine it.

The link step consumes the per-target LinkableModules, compiles each with its own toolchain, and links them into one LinkedModule. LinkedModule is consumed by the runtime loader (runtime); the concrete compiler commands are an implementation detail and not part of the contract.