TileFoundry Spec — Runtime¶
This spec owns the runtime contract outside the IR compile pipeline. It has two surfaces:
- the Python-side
RuntimeModule/ launcher ABI used bybuild(...)and examples/tests - the C++ runtime surface included by generated CUDA source
The C++ runtime is built on a vendored cutlass/include/{cute,cutlass}
snapshot.
1. Python Runtime Surface¶
The sections are one per file, and a file with more than one class states each
under its own heading. Each signature block is read out of that file by
scripts/runtime_spec_surface.py and written into its <!-- generated -->
region; a file's __all__ is what it states, and a docstring is the half the
spec says in its own words instead. A number nothing carries -- 1.5 and 1.7
-- is a section that became a heading under the file that owns it.
1.1 runtime/module.py¶
An ir Module (the semantic definition — @func bodies the evaluator runs) and
a RuntimeModule (the runtime instance — kernel bodies) are twins: same
name, same child tree, same entry. @runtime_module / @runtime_func
(tilefoundry.runtime.decorator, below) build one twin mechanically from the
other, validated one-to-one at decoration time; the correspondence is
additionally held by comparing the two numerically (§1.9), against bounds the
comparison's caller states — a RuntimeModule never runs the HIR evaluator.
RuntimeModule¶
# src/tilefoundry/runtime/module.py
class RuntimeModule:
name: str
entry: str | None
children: dict[str, RuntimeModule | tuple[RuntimeModule, ...]]
def __init__(
self,
name: str,
entry: str | None = None,
children: dict[str, RuntimeModule | tuple[RuntimeModule, ...]] | None = None,
) -> None: ...
@property
def module(self): ...
def forward(self, *args): ...
def __call__(self, *args): ...
def load(self, resource: RuntimeResource, *, placement: "Placement | None" = None) -> None: ...
- constraints:
- the base class is authored like a
torch.nn.Module: subclass it, build the child tree in__init__(children registered viachildren=), write the composition inforward. Function bodies areRuntimeFunctionattributes called fromforward.@runtime_module(below) generates this subclass mechanically from a semanticModuleand is the normal authoring path; a direct subclass remains available for special cases (e.g.CompiledModule, below). load(resource, *, placement=None)(base class): recurses into each child withresource.subtree(key)for a single slot, orresource.subtree(key).subtree(str(i))for each indexed instance, handing every one the same placement; the base class itself resolves nothing. One process is one program throughout the tree it loaded, so a child MUST NOT be given a placement of its own.- a twin with a sharded weight and no placement MUST refuse to read it.
Nothing else in the system says which of the mesh's positions this process
holds, and returning the whole tensor would silently give every program
every other program's data (§1.10).
Weight prefixes follow module paths, matching ir attribute addressing.
Lifecycle: construct (structure) →
load(remember source) → call. ARuntimeModuledoes not runprepare; it loads straight from the directory the semantic side prepared (§1.3). - correspondence contract: the runtime twin's
forwardand the semantic evaluator'sevaluate(LoadedModule, ...)must agree on the same inputs —measure.checkcomparing them against stated bounds (§1.9) is that contract. modulenames the authoredModulea twin was generated from, so a caller holding an implementation can reach what it is judged against. ARuntimeModulethat stands for no single authored Module — a compiled entry, a hand-written subclass — MUST reportNonerather than something chosen for it, and a caller that needs one MUST refuse instead of substituting.moduleis therefore reserved on a twin. An authoredModuleMAY name a function, child or methodmodule, and a generated twin binds each of those as an attribute, which would shadow the accessor.@runtime_moduleMUST reject such a Module when it is decorated, rather than generate a twin whose accessor answers something else.- the base class itself holds no weights or resource, and never runs the HIR evaluator.
- two origins: compiled —
tilefoundry.build/compile/jit→LinkedModule→ the loader binds aCompiledModule(aRuntimeModule, not aRuntimeFunction) (§1.4); handwritten — a@runtime_moduleclass (below), loading from a prepared checkpoint directory via aRuntimeResource(§1.8).
CompiledModule¶
# src/tilefoundry/runtime/module.py
class CompiledModule(RuntimeModule):
def __init__(self, type: CallableSignature, fn: Callable) -> None: ...
def load(self, resource: RuntimeResource, *, placement: "Placement | None" = None) -> None: ...
def forward(self, *args): ...
CompiledModule.forward(*args) uses the out-param ABI (type.output_count
trailing params are outputs):
- Auto-alloc:
len(args) == type.input_count— allocates output tensors from the first input's device/dtype, calls the entry, returns result(s). - Pre-alloc:
len(args) == len(type.params)— uses provided output tensors, returns same output(s). All outputs must be provided; partial →TypeError. - Return: single output → bare tensor; multiple outputs →
tuple.
Auto-alloc is torch-only; non-torch inputs raise TypeError. Output metadata
(dtype, shape) comes from CallableSignature.output_params — each
TensorSignature.type carries them (set by codegen from lowered IR, NOT
guessed at runtime). This convention is specific to CompiledModule; other
RuntimeModule.forward implementations are not bound by it.
1.2 runtime/function.py¶
RuntimeFunction is the base class for a node's function body; a body
subclasses it and overrides __call__. A handwritten torch / triton / CUDA
implementation takes whatever it needs (converted weights, caches) at
construction and returns its value(s) directly. RuntimeFunction.type is the
CallableSignature codegen produced for that function
(codegen §4.4).
# src/tilefoundry/runtime/function.py
class RuntimeFunction:
def __init__(self, type: CallableSignature) -> None: ...
def __call__(self, *args): ...
- constraints:
- the base
__call__raises; every usable body is a subclass. Agents may write any subclass whose__call__runs (torch / triton / CUDA / …). - the signature is codegen's product and the runtime's input: it is defined by codegen §4.4 and MUST NOT be redefined here.
1.3 runtime/decorator.py¶
# src/tilefoundry/runtime/decorator.py
def runtime_func(fn: Callable) -> Callable: ...
def runtime_module(sem: Module) -> Callable[[type], type]: ...
A weight's converter is registered per weight, not per module:
@<compute_fn>.converter("<weight_name>") decorates a throwaway def and
registers it on the base function's converters
(parser §3). Its parameters are
the raw-checkpoint names, annotated like any @func parameter; it returns
exactly the one declared ConstTensor's shape / dtype. A weight needing no
transform has no converter. Two converters registered for the same weight
name is an error.
A converter is an HIR body like any other and MAY reach another node's function.
load and prepare are methods on the IR Module owned by
core-ir §1. Execution of a loaded reading is entered
through evaluator.evaluate; this section defines that runtime-facing behavior
and the distinct loaded view:
class LoadedModule:
"""Bind one reading of an IR Module to its resource."""
module: Module
resource: RuntimeResource
children: dict[str, LoadedModule | tuple[LoadedModule, ...]]
target: Target | None
topologies: tuple[Topology, ...]
path: str
LoadedModule is a frozen resource binding: it stores no tensors. Attribute
access resolves a function to another callable LoadedModule whose entry
names that function, a single child to its reading, an indexed child to its
ordered tuple of readings, and an orchestration method
with types.MethodType. Calling a loaded module enters evaluate once and
passes only activations; constants come from its resource at first use.
- constraints:
- A reading covers its own Module's declared weights and holds one reading
per instance in
Module.children(core-ir §1). Calls follow core-ir §2.1, selecting the corresponding loaded child for an instance target. Each loop trip selects its evaluated index independently. A runtime index outside[0, N)MUST be refused with the slot key, index, and N. - Loading and preparation MUST carry the effective context under core-ir §1, storing it on each LoadedModule. Selecting a named function from a reading MUST re-enter with that stored context. Definitions and bodies remain shared; independent readings of the same definition MUST NOT borrow each other's resources.
LoadedModule.pathrecords only the diagnostic instance path: each single slot adds its key, each indexed slot addskey.i, joined with dots. A missing weight diagnostic MUST retainmissing declared weight '<name>'and include the complete path-plus-weight key, e.g.layers.3.w. This MUST NOT extend the RuntimeResource protocol.Module.prepare(semantic side, offline, once): walk every instance, children before their parent, using the same key and index segments as load. Each instance MUST run its definition's converters independently against its own raw subtree. For each declared weight, fetch its converter's parameters by their own (raw) names — a one-to-many alias is assembled here viatorch.stack(prepare's only reshaping) — run the converter through the evaluator, then strictly validate the result's shape and dtype against the weight's declaredConstTensortype. A weight with no converter is validated the same way against its raw (or stacked) value, unchanged. Output: one safetensors shard plusmodel.safetensors.index.json, keyed by clean, dot-joined module paths (e.g.layers.0.attention.w_kv). Plain directory — no content-hash cache / manifest. The runtime twin never prepares; both twinsloadstraight from the directory this writes.Module.load(resource): bind this node toresourceand recurse into each slot under its key and optional index segments, returning a frozenLoadedModuletree. It MUST NOT write bindings onto theModule, which stays pure IR: oneModulemay be read any number of times — two checkpoints, two devices — and each reading is independent of the others. A definition reached from two slots therefore yields oneLoadedModuleper instance rather than one binding the last instance wins. This is the semantic-side counterpart of theRuntimeModuletwin's ownload(§1.1).- The caller supplies
evaluate(..., device=...). Every activation and every weight used by the reading MUST already be on that device; a mismatch is rejected at the activation binding or weight's first-use point, naming the offending position or weight. Neither side moves tensors implicitly. - state is the caller's: a tensor that must survive across steps (e.g. a KV
cache) is an ordinary
Tensorparam passed in and returned, and a step MUST NOT mutate one it was given. Sharding such a tensor is therefore the same mechanism as for any other — its ownTensorType.layout— rather than a second description for an opaque state object. - The runtime twin's
forwardruns a registeredforwardorchestration method (Module.methods) when present, else its entry@func. A multi-node composition is chained by the caller, oneforwardper node. On the runtime twin, a bare step MUST be refused when neither aforwardmethod nor an entry is present, naming the functions and methods to call instead of reportingentryas wrong. On the semantic side,evaluate(loaded)andloaded(...)runentrywhen one is declared; neither treats aforwardmethod as an implicit default. A caller selects another HIR function or child throughloaded.<name>and calls it, while a named orchestration method remains ordinary host Python. - Python entry into an HIR
Functionthrough the runtime twin'sforwardis the runtime event governed by hir §1.1. A registered plain Python orchestration method stays on the host; every function it calls is a separate Python-to-HIR entry. An HIR-to-HIRCallremains inside the invocation and MUST NOT be surfaced as another runtime launch merely because its callee belongs to another Module. - a causal-LM root MAY define
init_caches,prepare_inputs_for_generation, andappend_cacheorchestration methods.prepare_inputs_for_generation(input_ids, step, caches, *, device)receives a one-dimensionaltorch.Tensorof token IDs. The model selects the token atstep, reshapes and places it, creates all other activation inputs in its ownforwardorder, and returns that positional tuple. The caller owns the cache and expands only the token-ID tensor; it MUST NOT reconstruct a model's positional, rotary, scaling, or state inputs. It passes the active token-ID prefix as a view; a method MUST NOT mutate that view or retain it across steps. These methods bind on the runtime twin in the same way asforward.
@runtime_module / @runtime_func¶
@runtime_module(sem) is a class decorator taking the semantic Module
instance; it returns a RuntimeModule subclass whose instances are sem's
runtime twin — same function names, same child tree, same entry.
@runtime_func tags a plain method as a kernel body: same call signature as
the semantic @func of the same name, weight params included.
# example
@runtime_module(attention_sem)
class Attention:
@runtime_func # weight params included, in the
def mla_kv_update(self, hidden, gamma_kv, w_kv, # semantic @func's own order
cos_pos, sin_pos, kv_cache0, cur_pos, s):
... # a real kernel body, e.g. a torch / triton / CUDA implementation
moe = SomeMoeRuntimeClass # a @runtime_module class, not an instance
- constraints:
- decoration-time validation is strictly one-to-one: the
@runtime_funcname set (a tagged method, or aRuntimeFunctioninstance class attribute — a heavy kernel that owns its own compilation state, standing in for a@runtime_func) MUST equalsem's function name set, and the child-attribute name set MUST equalsem.children's key set; missing or extra either MUST be rejected. Each slot MUST match both its single/indexed form and its number of elements, and a mismatch diagnostic MUST identify the key and the semantic and runtime counts. - Twin child declarations follow the semantic slots
(core-ir §1.3).
The generated constructor MUST build one runtime instance per semantic
element, passing
ir=<element definition>and inheritedtopologies=. Child classes MUST accept both constructor keywords. A child's explicit topology tuple replaces inheritance; an omitted declaration MUST inherit from its parent twin. Sharded input and weight projection MUST use this stored effective tuple rather than reading context from the shared definition. - weights are filled by name at call time from what
loadbound, so a kernel method's caller passes only activations. - orchestration methods (
forward/init_caches/ …) are reused from the semanticModule.methodsverbatim and are never rewritten on the runtime side: inside them,self.<fn>/self.<child>resolve to the runtime twin's own kernels / children. Absent an ownforward, the generated class runssem.methods["forward"]if present, else calls the entry function by name. A loaded semantic Module resolves the same named functions, children, and methods (§1.3).
1.4 runtime/loader.py¶
# src/tilefoundry/runtime/loader.py
def load_linked_module(linked: "LinkedModule") -> CompiledModule: ...
Module (IR) → codegen: per-target LinkableModule… → LinkedModule (.so + metadata)
LinkedModule → load → CompiledModule (fully-loaded, public, callable RuntimeModule)
LinkedModule is a codegen product
(codegen §4.3); the loader that turns it into a
CompiledModule is owned here. The loader and LinkedModule are not public
API; only CompiledModule (a RuntimeModule) is. load_linked_module returns
CompiledModule(type=replace(entry.signature, name=entry.name), fn=entry_callable),
where entry is the one LinkedFunction; name / entry are its registration
name. Its forward implements the out-param calling
convention directly (§1.1) — there is no separate function-body object it
delegates to. The compiled path has no resource / weights / states
(weights are ordinary entry arguments), so its load is the inherited no-op.
Launcher ABI¶
tilefoundry.build(mod) internally runs codegen and links the artifact (see
codegen), then loads it and binds the entry; these are
implementation details. Users interact only with RuntimeModule.__call__ /
RuntimeFunction.__call__.
Load contract:
- codegen produces the
LinkedModuleartifact (codegen §4.3) - loading uses
tvm_ffi.load_module(...) - immediately after loading, entry binding uses
tvm_ffi.get_global_func(entry.name);RuntimeModule.entryis this dotted registration name - the loader keeps every library it loads until process exit, because the global registration table does not keep the shared library loaded
- host registration allows override: another load replaces the global entry,
while an existing
CompiledModuleretains its own callable - the compiled loader requires exactly one
LinkedFunction - callable arguments are DLPack-compatible tensors;
torch.Tensoris one supported caller-side provider but is not the semantic contract itself
Generated host wrappers register entry functions with TVM FFI:
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::Function::SetGlobal(export_name,
tvm::ffi::Function::FromTyped(entry_function, export_name), /*override=*/true);
}
The exported function accepts flattened input/output tensor arguments. HIR
functions may be written as Function(params) -> tensor, but by the runtime
boundary the TIR/codegen surface is explicit input/output parameters.
Launch geometry (grid / block extents) is embedded into the generated host entry
or supplied by an authored launch(...); it is never carried as metadata past
codegen.
1.6 compile.py¶
# src/tilefoundry/compile.py
@dataclass(frozen=True)
class CompilerOptions:
target: Target
def canonical_text(self) -> str: ...
def normalize_to_module(fn_or_mod: HirFunction | Module) -> Module: ...
def build(mod: Module, /, *, target: Target | None = None) -> "RuntimeModule": ...
def compile(mod: Module, /, *, target: Target | None = None) -> "RuntimeModule": ...
def jit(
fn_or_mod, /, *, target: Target | None = None, options: CompilerOptions | None = None, **kwargs
) -> "RuntimeModule": ...
- constraints:
- accepts only TileFoundry IR (
Function/Module); raw Python functions raiseTypeError; the full input contract is stated below.
tilefoundry.jit(fn_or_mod, *, target=None, options=None) is the JIT
entry point. It accepts a hir.Function or Module, normalizes to a
Module, compiles with cache, and returns a callable RuntimeModule.
Input contract:
- Only TileFoundry IR objects (Function / Module) accepted.
- Raw Python functions raise TypeError — use @func first.
- A stated target MUST be a constructed Target instance. Strings
MUST be refused and MUST NOT select a backend or construct a Target.
- A Module-owned Target takes precedence and is retained as the exact instance;
a conflicting explicit Target or CompilerOptions.target MUST fail.
- Topology is declared by the Module; a single-function
@func(topologies=...) declares it through the implicit Module that
decorator yields (parser §2.1).
- Mesh layout is expressed in the DSL with lexical with Mesh(...) as mesh scopes.
- jit() has no cta_mesh / thread_mesh parameters.
Pipeline: jit() reuses passes §6's build()
pipeline (compile()). It auto-wraps a bare Function input into a
single-function Module that declares no execution context.
Cache: in-process dict cache keyed by
sha256(canonical_module_text + target_text + canonical_options_text).
canonical_module_text includes functions, declared topology hierarchies, and
with Mesh scopes. Independently compiled Modules declare their execution
context (core-ir §1). No
Python object identity and no
dedicated cta_mesh / thread_mesh key fields participate in the key.
jit.cache_clear() evicts; jit.cache_info() returns {"size": N}.
1.8 runtime/resource.py¶
Checkpoint aliasing is a base capability of every resource, not a wrapper
class: both implementations below take an alias={canonical: raw} table,
resolved by the same lookup order.
# src/tilefoundry/runtime/resource.py
@dataclasses.dataclass(frozen=True)
class Absolute:
name: str
@dataclasses.dataclass(frozen=True)
class Preprocessed:
name: "str | Absolute"
read: Convert
class RuntimeResource(Protocol):
def load(self, name: str) -> torch.Tensor: ...
def load_group(self, name: str) -> "tuple[torch.Tensor, ...] | None": ...
def subtree(self, seg: str) -> "RuntimeResource": ...
class DictResource:
def __init__(
self, data: Mapping[str, torch.Tensor], prefix: str = "", alias: AliasMap | None = None
) -> None: ...
def load(self, name: str) -> torch.Tensor: ...
def load_group(self, name: str) -> "tuple[torch.Tensor, ...] | None": ...
def subtree(self, seg: str) -> "DictResource": ...
class SafetensorsResource:
def __init__(
self, ckpt_dir: str, prefix: str = "", device: str = "cuda", alias: AliasMap | None = None
) -> None: ...
def load(self, name: str) -> torch.Tensor: ...
def load_group(self, name: str) -> "tuple[torch.Tensor, ...] | None": ...
def subtree(self, seg: str) -> "SafetensorsResource": ...
class DrawnResource:
def __init__(
self,
module: Module | ModuleList,
generator: torch.Generator,
device: str,
drawn: dict[str, torch.Tensor] | None = None,
prefix: str = "",
) -> None: ...
def load(self, name: str) -> torch.Tensor: ...
def load_group(self, name: str) -> None: ...
def subtree(self, seg: str) -> "DrawnResource": ...
def draw_tensor(declared: TensorType, generator: torch.Generator, device: str) -> torch.Tensor: ...
- constraints:
Preprocessedis a frozen dataclass carrying one raw name and its one-tensor read transform.load(name)returns the tensor for name; raisesKeyErrorif absent, and MUST raiseTypeError(namingload_group) if name resolves to a tuple-valued (one-to-many) alias.load_group(name)returns the tuple of raw tensors for a one-to-many alias entry (e.g. per-expert weight shards, in declared order), orNonewhen name has no tuple-valued alias — the ordinary, one-to-one case.subtree(seg)returns a view scoped under one more path segment (segis itself alias-resolved), so a childRuntimeModuleaddresses its own weights by their bare (unprefixed) name —RuntimeModule.__init__never sees a dotted name.- the resource resolves names and reads tensors. It MUST NOT stack, and it
reshapes only where a
Preprocessedalias entry states that the checkpoint stores that one tensor differently from how the Module declares it — a transpose, a slice of a fused tensor, or a dropped axis. Assembling a one-to-many group into one tensor staysprepare's job (§1.3), viatorch.stack, and so does any value that is a function of more than one raw tensor.PreprocessedMUST resolve to one name: a tuple-valued name is rejected at construction, naming the weight converter as the way to express it. Precision is not preprocessing: a read MUST return the checkpoint's own stored element type and be validated against the declaration (§1.3).
An alias entry renames one path segment or leaf within the current
scope, joined onto the caller's already-accumulated prefix: lookup order is
a path-qualified key (f"{prefix}{name}") first, then a bare name entry,
then identity (f"{prefix}{name}" unchanged). A bare entry therefore serves
every instance at that level uniformly (e.g. {"gamma_kv": "kv_norm.weight"}
fires under every layer); a per-instance name (one real decoder layer, one
per-expert shard group) needs one literal entry per instance instead. A
Preprocessed value is a one-to-one leaf: load applies its read callable
to the raw tensor. A tuple value is the one-to-many group load_group reads;
subtree's own segment resolution rejects a tuple-valued or Preprocessed hit
(a subtree segment MUST resolve to one relative path).
Aliasing therefore only ever reaches downward: a name resolved inside a
scope carries that scope's prefix, so a node cannot address a tensor its
parent owns — and a checkpoint may well store one there, such as a
layer-level norm weight a child consumes. Absolute(name) is the escape: an
alias whose value is Absolute MUST resolve to name as the whole raw key,
with no prefix joined onto it. It stays a leaf-only form — load_group reads
it as the one-to-one case and returns None, and subtree MUST reject it in
the same shape as a tuple-valued hit, because a subtree segment must resolve
to one relative name.
Three implementations:
class DictResource:
def __init__(
self, data: Mapping[str, torch.Tensor], prefix: str = "",
alias: "Mapping[str, AliasValue] | None" = None,
) -> None: ...
class SafetensorsResource:
def __init__(
self, ckpt_dir: str, prefix: str = "", device: str = "cuda",
alias: "Mapping[str, AliasValue] | None" = None,
) -> None: ...
class DrawnResource:
def __init__(
self, module: Module, generator, device: str,
drawn: "dict[str, torch.Tensor] | None" = None, prefix: str = "",
) -> None: ...
def draw_tensor(declared: TensorType, generator, device: str) -> torch.Tensor: ...
- constraints:
DictResource— in-memory / test double over a flat, dot-prefixed{"layer0.w": tensor, ...}mapping;subtreeonly extends the prefix eachload/load_groupname is joined onto, carryingaliasdown to every child view.SafetensorsResource— reads a safetensors checkpoint directory;load/load_groupopen at most one shard handle per shard file (mmap'd viasafetensors.safe_open, shared acrosssubtreeviews) and read only the requested tensor(s), placed on device. Two directory shapes MUST be accepted: N shard files with amodel.safetensors.index.jsonwhoseweight_mapnames the shard holding each key, and a single unshardedmodel.safetensorswith no index, whose own key list is that map — a published checkpoint is only sharded once it outgrows the writer's limit, so requiring an index would refuse the small ones. A directory with neither MUST be reported as such rather than as a missing index. A raw key is cached weakly: repeated reads return the same tensor while a caller holds it, and the entry may disappear after the last reference is collected; scoped views share the index, shard handles, and tensor cache.- device is resolved on every read, and both the shard handles and the
tensor cache are keyed by that resolved device alongside the name. A bare
"cuda"means whichever card the process has selected, which a single-process multi-card host changes between reads; a handle keyed by file alone would stay pinned to the card that first opened it, and a tensor keyed by name alone would hand a reader on one card what another card holds. DrawnResourcedraws a declared weight lazily withdraw_tensor, keeps the result strongly, and shares its generator and drawing ledger acrosssubtreeviews without reseeding. The strong ledger is a contract: drawing the same name again would produce another random tensor, so the two sides of a comparison would no longer receive the same value. It has no alias groups, soload_groupalways returnsNone.- every read tensor keeps the element type the checkpoint stores. A
declaration requiring a different precision uses a weight converter
(§1.3), and
loadvalidates the converted or raw result against that declaration.
1.9 runtime/measure.py¶
# src/tilefoundry/runtime/measure.py
@dataclass(frozen=True)
class Predicate:
name: ClassVar[str] = ""
bounds: ClassVar[tuple[str, ...]] = ()
needs_reference: ClassVar[bool] = True
discrete: ClassVar[bool] = True
guidance: ClassVar[str] = ""
def measure(
self, candidate: torch.Tensor, reference: torch.Tensor | None
) -> "PredicateResult": ...
@dataclass(frozen=True)
class PredicateResult:
predicate: Predicate
values: Mapping[str, float]
passed: bool
note: str | None = None
@dataclass(frozen=True)
class AllClose(Predicate):
atol: float
rtol: float
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class RelL2(Predicate):
max: float
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class Cosine(Predicate):
min: float
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class Equal(Predicate):
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class Ulp(Predicate):
max: float
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class MaxAbs(Predicate):
max: float
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class MaxRel(Predicate):
max: float
def measure(self, candidate, reference): ...
@dataclass(frozen=True)
class NanInf(Predicate):
def measure(self, candidate, reference=None): ...
PREDICATES: Mapping[str, type[Predicate]] = {
predicate.name: predicate
for predicate in (AllClose, RelL2, Cosine, Equal, Ulp, MaxAbs, MaxRel, NanInf)
}
@dataclass(frozen=True)
class OutputCheck:
path: str
shape: tuple[int, ...]
dtype: str
ref_norm: float | None
results: tuple[PredicateResult, ...]
@property
def passed(self) -> bool: ...
@dataclass(frozen=True)
class Report:
outputs: tuple[OutputCheck, ...]
@property
def passed(self) -> bool: ...
def flatten_outputs(x, path: str = "output") -> list[tuple[str, torch.Tensor]]: ...
def check(
candidate: Callable,
reference: Callable | None,
inputs: tuple,
*,
expect: Mapping[str, Sequence[Predicate]],
) -> Report: ...
- constraints:
PREDICATESMUST contain the built-inallclose,rel_l2,cosine,equal,ulp,max_abs,max_rel, andnan_infpredicates by name.checkrunscandidate(*inputs), andreference(*inputs)when there is a reference, and measures each output against the predicates expect states for it. Neither reference nor expect has a default.- an input MAY be a bare tensor or an arbitrarily nested tuple of tensors. Every leaf MUST be a tensor.
- a result MAY be a bare tensor or an arbitrarily nested tuple of tensors
(e.g.
forward's(logits, past_key_values)).checkflattens both results and MUST reject a candidate whose flattened structure, shape or dtype differs from the reference's. - every produced tensor MUST have exactly one non-empty list of predicates, and a path expect names that was not produced MUST be rejected. A result that flattens to nothing MUST be an error rather than a pass.
reference=Noneadmits only predicates whoseneeds_referenceis false; any other MUST be rejected.ref_normis then absent, having nothing to measure.- a predicate whose
discreteis false MUST be rejected on an integer or boolean output, naming exact comparison instead. - where a measure has no meaning at the values it was given — a relative
distance against a zero reference, a direction between two zero vectors — the
result MUST state what was measured instead through its
note, rather than return a number whose scale is an artefact of a clamp. OutputCheck.passedandReport.passedare read-only properties derived asallof their parts; neither is accepted as a constructor field, so a verdict cannot disagree with the measurements printed beside it.- it is not specific to
RuntimeModule: candidate / reference may be aRuntimeModulebound method, a raw torch callable, or an evaluator closure — anything callable on inputs.
1.10 runtime/tensor.py¶
Module.prepare writes one canonical tensor per weight and the HIR keeps the
global TensorType, so a checkpoint read hands back every program's data at
once. ShardTensor pairs that owning tensor with the type that shards it, and
to_local is the single place the pair is narrowed to one program.
# src/tilefoundry/runtime/tensor.py
@dataclass(frozen=True)
class ShardTensor:
tensor: torch.Tensor
type: TensorType
def to_local(
self, topologies: tuple[Topology, ...], placement: "Placement | None"
) -> torch.Tensor: ...
- constraints:
ShardTensorcarries noPlacement, coordinate, rank or device map. The program is an argument toto_local, so one owning tensor serves every program a host dispatches, andto_localis the only way to a slice. Named asDTensor.to_local, which answers the same question.- a
typewith noShardLayoutis every program's whole tensor, and so is an axis the mesh only broadcasts over or reduces across. - which part is one program's MUST come from
local_layout_and_offset(shard §7), the same algebra the device side runs under that name, rather than from a second walk over the attrs. Host and device read the one layout and the one offset, so a slice is not derived twice and cannot disagree with itself. - a level whose id is unfixed MUST be left undivided (shard §5.1).
- a
RuntimeModule's public boundary takes and returnsShardTensor, while a body -- a Python reference or a launched kernel alike -- only ever receives the localtorch.Tensor. A returned tensor is already one program's, so the type it is paired with states that shape and carries no layout: narrowing it again would take a shard of a shard. - a weight narrowed on first use is copied out of the read it came from, so the whole canonical tensor is released once every shard it contributes is held. A weight the mesh does not divide is kept as the read itself.
- what remembers a read weight is keyed by name, program ids and the current device together. One weight read on two cards is two tensors, and a key of the name alone hands a reader on one card what another card holds.
2. C++ Runtime Surface¶
The sections follow include/tilefoundry/runtime/, one per header. Each
declaration block is read out of that header by
scripts/runtime_spec_surface.py and written into its <!-- generated -->
region, so a signature is stated once, where it is compiled. What follows a
region is written by hand: the terms the declarations are spelled in, and the
decisions two reasonable implementations would differ on.
2.1 cpu/¶
2.1.1 cpu/runtime.h¶
Terms. The CPU target's launch boundary is raw pointers and shape scalars, so this header declares no surface. Host-side tensor compute is future work.
- constraints:
- The CPU target crosses the launch boundary as raw pointers and shape scalars, so nothing here is a type a caller names.
2.2 cuda/¶
2.2.1 cuda/runtime.cuh¶
// include/tilefoundry/runtime/cuda/runtime.cuh
inline constexpr int kWarpSize;
template <class...> inline constexpr bool dependent_false_v;
enum class TopologyScope {
gpu,
cta,
thread,
scope_count,
};
template <TopologyScope T>
CUTE_HOST_DEVICE constexpr auto program_dim() noexcept;
template <TopologyScope T>
CUTE_HOST_DEVICE constexpr auto program_shape() noexcept;
CUTE_HOST_DEVICE constexpr TopologyScope program_topology() noexcept;
CUTE_HOST_DEVICE constexpr auto program_topologies() noexcept;
struct ProgramMetaData {
int program_id[int(TopologyScope::scope_count)];
};
CUTE_HOST_DEVICE ProgramMetaData &program_meta();
CUTE_HOST_DEVICE void program_meta(ProgramMetaData const &meta);
template <TopologyScope T> CUTE_HOST_DEVICE size_t program_id() noexcept;
CUTE_HOST_DEVICE auto program_ids() noexcept;
Terms. A level is one parallel-resource level of the launch: gpu, cta or thread. An instance is one id of a level. A coordinate is one id per level, which is what program_ids() hands back.
- constraints:
gpuis the level the host places. A card carries no register naming which of the mesh's cards it is, soprogram_id<TopologyScope::gpu>()reads what the launch was told instead, and that value MUST equal the id the host'sPlacementgave thegpulevel for the same invocation (§1.10).- the value reaches the device as a by-value
ProgramMetaDatakernel parameter, and the generated kernel MUST hand it toprogram_metabefore any divergence. One thread writes and the barrier orders that write against every read. program_meta's slot is__shared__, so one name is one copy per block rather than one per device: two kernels running at once on one card each read what their own launch was told.scope_countstates how many levels there are, and everything indexed by a level MUST derive its length from it:ProgramMetaData.program_id, the tupleprogram_ids()returns, and the coord a mesh is asked about. A coord shorter than that is missing a level rather than naming a smaller mesh, and MUST be refused where it is used rather than left to read the neighbouring level's id.program_topology()is the coarsest level this program names andprogram_topologies()is the whole run from there toscope_count, coarsest first. Both derived answers are CuTe's own questions about that tuple -- how many levels iscute::rankof it, where levelSsits iscute::find(program_topologies(), cute::C<S>{})-- so neither has a name of its own, and a level this program does not name is refused the way CuTe refuses it:cute::findhands back the rank and thegetthat follows reports"Index out of range".-
that parameter MUST NOT appear in the wrapper's user-visible signature. A
Placementitself never reaches a kernel (§1.10); the id derived from it is an internal parameter the caller supplies. -
constraints:
- The enumeration is fixed to the
ctaandthreadprogram levels plus thescope_countsentinel. A warp-sized grouping is an axis of athreadmesh's layout, not a level of its own. - Positions are row-major: the last axis is the fastest. Extents
(8, 32)give(id / 32, id % 32), so a thread mesh names its warps first and its lanes last. Every layout in this surface is read that way, and CuTe's own algebra is not (§2.3.1). - One
.cuis one launch.program_dim<T>()is what a translation unit states;program_shape<T>()is derived from it as levelTand every level under it, which is the shape a mesh naming several levels is indexed against. The same mesh compiled into another translation unit belongs to another launch.
2.3 cuda/layout/¶
2.3.1 layout/cute_ext.cuh¶
// include/tilefoundry/runtime/cuda/layout/cute_ext.cuh
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr auto reverse(cute::Layout<Shape, Stride> const &l);
Terms. CuTe order is CuTe's own reading of a layout, mode zero fastest. A mesh is written row-major, last axis fastest, so the two are reverses of each other and CuTe's algebra reads a mesh backwards until it is turned.
- constraints:
- CuTe reads mode zero as the fastest and a mesh is written the other way round, so every CuTe function meaning "next to" reads a mesh backwards:
coalesceleaves(2,32):(32,1), 64 consecutive threads, unfolded, andlogical_dividecuts a domain that runs the other way. The runtime reverses a mesh's axes once and uses plain CuTe after that.
2.3.2 layout/mesh.cuh¶
// include/tilefoundry/runtime/cuda/layout/mesh.cuh
template <class TLayout, TopologyScope... Topos> struct Mesh {
using layout_type = TLayout;
static constexpr int level_count;
static constexpr TopologyScope scope;
TLayout layout;
};
template <TopologyScope S, class L, TopologyScope... Topos>
CUTE_HOST_DEVICE constexpr auto get(Mesh<L, Topos...> const &mesh);
template <class L, TopologyScope... Topos>
CUTE_HOST_DEVICE constexpr int offset(Mesh<L, Topos...> const &mesh);
template <class L, TopologyScope... Topos, class Coord>
CUTE_HOST_DEVICE constexpr bool contains(Mesh<L, Topos...> const &mesh,
Coord const &coord);
template <class L, TopologyScope... Topos, class Coord>
CUTE_HOST_DEVICE constexpr auto mesh_coords(Mesh<L, Topos...> const &mesh,
Coord const &coord);
template <class L, TopologyScope... Topos>
CUTE_HOST_DEVICE constexpr bool is_warped(Mesh<L, Topos...> const &mesh);
template <class L, TopologyScope... Topos>
CUTE_HOST_DEVICE constexpr auto as_warped(Mesh<L, Topos...> const &mesh);
template <TopologyScope Scope, class Extents>
CUTE_HOST_DEVICE constexpr auto make_mesh(Extents const &extents);
Terms. A mesh names one or more levels and carries a layout whose values are those levels' ids. A slice is a mesh whose layout has a constant term: it covers part of a level, starting at that term. Warped describes a thread mesh whose ids run a whole number of warps; its lane axis is one warp and its warp axis steps between them.
- constraints:
- A mesh is the levels it names and the layout whose values are their ids -- the same two facts the IR
Meshcarries beside its axis names (shard §5). Its instance count and shape arecute::sizeandcute::shapeof that layout, which read through a slice on their own. - A mesh naming several levels states its axes grouped one nest per level, each already in that level's own numbering, and
get<level>picks the nest. A mesh that names several without that grouping is refused: no rule says which axes are whose. It cannot also be sliced -- the slice and the level boundary would both be deciding which positions these are. mesh_coordsanswers where this program sits in the mesh's own coordinates: oneidx2crdper level, of that level's id against that level's nest, grouped the way the mesh is. The answer is congruent withcute::shapeof the mesh's layout whether it names one level or several, so a layout over that shape indexes with it as it stands. The ids it reads areprogram_ids(), one per level this program names, so a coord of another length is refused by name.containsandmesh_coordsare not the same question, and on a mesh narrower than its level not the same answer.idx2crdis(id / stride) % extent, so thread 64 of a 128-thread block is not in a 32-instance mesh and does act as its instance 0: the upper warps repeat what the lowest warp does. A slice does not repeat, so a coordinate from outside one reachingmesh_coordsis a codegen fault.containsasks every level the mesh names, because a mesh grouped one nest per level is one instance per level together.- A thread mesh has warps only when its ids run a whole number of them.
(32,..)stepping by one is one warp and the axis beyond it carries the warps;(128,..)holds four thatas_warpedsplits out. Lanes shorter than a warp, repeated across warps, MUST be refused by name -- they reach neitherbar.sync, which counts whole warps, nor one__syncwarp. as_warpedis idempotent: a mesh whose fastest axis is already one warp comes back unchanged.- Only one shape of
cute::ComposedLayoutis a mesh:cute::identityover a static offset. A swizzle in the first slot, or a dynamic offset, is a mesh whose first instance is not a compile-time number, and every reader wants it as one. - The offset MUST NOT be dropped. A mesh layout maps a coordinate to an instance within the mesh; only the offset turns that into an instance of the launch, and reading a coordinate off a raw id instead hands every instance of a slice the box its neighbour owns. A plain
cute::Layoutis the whole level at offset zero.
2.3.3 layout/shard_layout.cuh¶
// include/tilefoundry/runtime/cuda/layout/shard_layout.cuh
template <class TLayout, class TAttrs, class TMesh> struct ShardLayout {
using layout = TLayout;
using attrs = TAttrs;
using mesh = TMesh;
TLayout layout_value;
TMesh mesh_value;
};
template <int Axis> struct S {
static constexpr int axis;
};
struct B {};
template <class Reduction> struct P {
using reduction = Reduction;
};
struct Dynamic {};
template <size_t Ax, class L, class A, class M>
CUTE_HOST_DEVICE constexpr auto stride(ShardLayout<L, A, M> const &sl);
template <class SL> CUTE_HOST_DEVICE constexpr void check_shard_layout();
template <class Shape, class TMesh, class Attrs,
__CUTE_REQUIRES(!cute::is_layout<Shape>::value)>
CUTE_HOST_DEVICE constexpr auto
make_shard_layout(Shape const &shape, TMesh const &mesh, Attrs const &);
template <class TLayout, class TMesh, class Attrs,
__CUTE_REQUIRES(cute::is_layout<TLayout>::value)>
CUTE_HOST_DEVICE constexpr auto
make_shard_layout(TLayout const &layout, TMesh const &mesh, Attrs const &);
Terms. A shard layout is a tensor layout, one attr per mesh axis, and the mesh those attrs are indexed against. An attr says what its mesh axis does with the tensor: S<k> splits tensor axis k across it, B broadcasts, and P<R> leaves an unreduced partial.
- constraints:
TAttrsis acute::tupleof empty tags, so the aggregate holds no attribute value: the attributes are in the type and the layout is what carries data.- One attr per mesh axis. A shard layout that says nothing about an axis leaves the reader to guess whether that axis splits the tensor.
2.4 cuda/tensor_view/¶
2.4.1 tensor_view/shard_tensor.cuh¶
// include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh
template <class TEngine, class TGlobalLayout, class TShardLayout>
struct ShardTensor {
using engine_type = TEngine;
using global_layout_type = TGlobalLayout;
using shard_layout_type = TShardLayout;
TEngine engine;
TShardLayout shard_layout;
/// As ``cute::Tensor`` has it: the engine's own pointer, before any
/// instance's slice is projected out of it.
CUTE_HOST_DEVICE auto data();
CUTE_HOST_DEVICE auto data() const;
};
template <class T, class GL, class SL>
CUTE_HOST_DEVICE auto make_shard_tensor(T const &tensor, GL, SL shard_layout);
template <class T>
concept ShardTensorLike =
detail::is_shard_tensor<cute::remove_cvref_t<T>>::value;
template <class SL>
inline constexpr bool shard_layout_is_composed_v;
template <class T> CUTE_HOST_DEVICE decltype(auto) local_tensor(T &&t);
template <class T>
using local_view_t =
cute::remove_cvref_t<decltype(local_tensor(std::declval<T const &>()))>;
template <class T> CUTE_HOST_DEVICE constexpr int shard_mesh_instances();
Terms. A shard tensor is an engine, the global layout it came from, and a shard layout. Projecting it is resolving it to the slice this instance owns: local_tensor, whose type-level half is local_view_t, while ShardTensorLike is the one test for "is this operand sharded". All three are tilefoundry:: names, not detail ones.
- constraints:
engineis a CuTe tensor or view, never a raw pointer: residency lives on the engine type, anddata()drops it. Its layout MUST be the declared local shard layout read throughresolve_layout, using the shared register numbering for rmem; codegen MUST NOT choose a separate default engine order. Storage spans the layout codomain, including any pitch. A rank-0 address is encoded in CuTe as one singleton mode with stride zero.- A shard layout whose layout is a
cute::ComposedLayout-- which is how aSwizzlereaches the device -- is projected with its function kept:shard_layout_is_composed_vpicks that path, andlocal_tensorreturns a tensor over the engine's own pointer whose layout is the same composition, this instance's origin folded into its offset. Advancing the pointer instead would run the function on a slice-relative index and hand two instances one address, so the offset MUST NOT move to the engine. - A composed layout takes that path even when no mesh axis splits it. The composed path reconstructs the declared function and the participant offset over the engine pointer, including when every mesh attribute broadcasts.
local_tensor,local_view_tandShardTensorLikeare public: a step every op must take, and the word an op writes its own constraints in, are not implementation details.is_shard_tensor-- howShardTensorLikeis decided, which nothing outside needs -- stays indetail.- A tensor whose layout is a
ShardLayouthas distributed semantics -- it is the whole tensor, and each instance owns the slice its shard layout gives it. - One attr per mesh axis, against the mesh's axes flattened: a mesh naming several levels states them grouped one nest per level, and its rank is then how many levels it names rather than how many axes the attrs answer for.
- Several mesh axes MAY cut one tensor axis -- two axes of one level as a grid, or two levels one taking a block of what the other left. Each cut divides what the previous one left, so the axis's local extent is the whole divided by all of them, and one step of a mesh axis clears everything the axes inside it hold.
- Where an instance's slice begins is that mesh's own shape carrying, on each axis, how far along the tensor one step of it moves the origin: the attrs say which tensor axis and how much of it a step clears, the tensor layout says what that costs in elements. It is congruent with the mesh's shape, grouping and all, so
mesh_coordsindexes it as it stands and reading it iscrd2idx. It is notoffset(mesh), which is the first id a sliced mesh covers.
2.5 cuda/primitive/¶
2.5.1 primitive/unary.h¶
// include/tilefoundry/runtime/cuda/primitive/unary.h
struct rsqrt_op {
template <class T> __device__ T operator()(T x) const;
};
struct neg_op {
template <class T> __device__ T operator()(T x) const;
};
struct relu_op {
template <class T> __device__ T operator()(T x) const;
};
struct square_op {
template <class T> __device__ T operator()(T x) const;
};
struct sigmoid_op {
template <class T> __device__ T operator()(T x) const;
};
struct silu_op {
template <class T> __device__ T operator()(T x) const;
};
struct softplus_op {
template <class T> __device__ T operator()(T x) const;
};
struct exp_op {
template <class T> __device__ T operator()(T x) const;
};
struct log_op {
template <class T> __device__ T operator()(T x) const;
};
struct identity_op {
template <class T> __device__ T operator()(T x) const;
};
struct clamp_op {
float min_val, max_val;
template <class T> __device__ T operator()(T x) const;
};
Terms. A primitive is a callable tag an op instantiates. It takes scalars, not tensors, which is why it is not an op.
- constraints:
- A primitive takes scalars, not tensors, so it is not an op: it has no shard layout to read a behaviour off.
ops::closes on the entries that do (§2.6).
2.5.2 primitive/binary.h¶
// include/tilefoundry/runtime/cuda/primitive/binary.h
struct mul_op {
template <class T> __device__ T operator()(T a, T b) const;
};
struct add_op {
template <class T> __device__ T operator()(T a, T b) const;
};
struct sub_op {
template <class T> __device__ T operator()(T a, T b) const;
};
struct max_op {
__device__ float operator()(float a, float b) const;
template <class T> __device__ T operator()(T a, T b) const;
};
struct min_op {
__device__ float operator()(float a, float b) const;
template <class T> __device__ T operator()(T a, T b) const;
};
struct div_op {
template <class T> __device__ T operator()(T a, T b) const;
};
- constraints: none beyond the signature.
2.6 cuda/ops/¶
Codegen targets one public namespace function per runtime op/family:
flowchart LR
Codegen["generated target call"] --> Entry["ops::<op>(...) public entry"]
Entry --> Dispatch["internal trait / dispatch function"]
Dispatch --> ImplA["impl class / helper A"]
Dispatch --> ImplB["impl class / helper B"]
Entry --> SimpleImpl["single impl helper"]
What is an op. An op's behaviour is derivable from the layout system's
description of its arguments: its operands are ShardTensors or a Mesh, and
its tier is read off their layouts. An op is one indivisible dependency chain
or one per-element map; a composition of two ops is not a third one. What that
admits is a list and not a count — a number is a fact about today's list, and
the one time it was written down as the rule it argued a dependency chain into
elementwise to keep the count. ops:: holds these and nothing else:
| entry | operands the geometry comes from |
|---|---|
elementwise |
the destination's local domain, and a stride-0 mode wherever a source broadcasts |
copy / copy_async |
both operands' shard layouts: shape, strides, share and move width |
reduce |
the axes the destination broadcasts that the source splits |
dot |
the axes the operands' meshes contract |
mma |
one atom's per-lane operand fragments |
sync |
the mesh's scope, base and count |
copy_async_bulk |
both shard layouts, asserted: one contiguous run each, whole tiles, matching element types |
ldmatrix |
the source tile's declared shared-memory layout and the destination atom fragment |
Anything that takes a raw pointer, an int or a type and answers a question
about it is not an op but a utility, and belongs outside ops:: — the
warp-scoped primitives of §2.7 and the mbarrier
instructions a caller writes around copy_async_bulk are both that. In particular the
runtime publishes no predicate and no constant reporting which tier an op
selected or how wide a move it chose: a caller cannot use one to decide how to
build its operands, since the answer is a function of the operand types it would
already have had to build, and putting the tier behind one entry was the point.
Runtime-owned dispatch. Where an op has more than one implementation tier
(selected by scope or by operand layout), the runtime exposes exactly one
public entry — never one op per tier. The active tier is derived at compile
time from the operand ShardLayouts, together with any codegen-static geometry
passed as template parameters, through a template trait, and is selected inside
the entry (if constexpr). Codegen emits one uniform call per op and never
selects a tier, computes a per-tier parameter, or carries the selection on the
TIR op. ops::reduce derives its reduction level from the operand shard
layouts and ops::sync derives its participant predicate from the barrier
geometry; both are instances of this principle. A
target runtime implementation MAY select an internal optimized load/store path
(such as a wider vector copy) behind this single entry without changing the
public entry or its observable result. The codegen side is
codegen §3.
One pointwise op. Every elementwise operation — fill, cast, clamp, relu,
add, multiply, and every broadcast shape of those — is the single entry
ops::elementwise. Arity is the length of its source pack, the operation is
the callable it takes, and a broadcast operand is a stride-0 mode on that
operand's layout. None of those
three is a name: there is no ops::fill, ops::cast, ops::unary,
ops::binary, or ops::binary_bcast_*, and adding one would be re-encoding in
a symbol what an argument already states.
Annotation convention. ops::* public entries, their internal impl
functors, and op tags MUST be annotated __device__ (their bodies are
device-only). CUTE_HOST_DEVICE MUST be reserved for tensor-view / layout
helpers genuinely capable of host compilation (e.g. local_tensor,
make_shard_tensor).
2.6.1 ops/detail.cuh¶
Terms. ops::detail is the set of tensor-view implementation names an op may reach for: id_axes, is_partial_attr_v, is_shard_tensor, is_split_attr_v, positions_of, shard_attrs_match_mesh, shard_layout_is_full_broadcast. They are defined in tilefoundry::detail; this header states which of them cross into ops. Projecting an operand is local_tensor: a ShardTensor resolved to this instance's slice, and an operand the mesh never spread returned whole. It is not one of them -- it, local_view_t and ShardTensorLike are tilefoundry:: names an op calls unqualified by detail (§2.4.1).
- constraints:
- An op reaches for a tensor-view implementation name through
ops::detailand nowhere else, so what crosses that boundary is stated in one header. The gate holds only names an op may want and none must have; a step no op may skip is public instead, because a thing every op goes through is not an implementation detail. - An op projects every operand through
tilefoundry::local_tensor, sharded or not. An operand with no shard layout is not a case to refuse: it is one every instance holds entire, and the op's arithmetic over it is the same. So the choice is made once, where the operand is projected, rather than at each op.
2.6.2 ops/copy.cuh¶
// include/tilefoundry/runtime/cuda/ops/copy.cuh
template <class TPointer, class TLayout>
CUTE_HOST_DEVICE auto tensor_view(TPointer pointer, TLayout layout);
template <class TSrc, class TDst>
__device__ void copy(TSrc const &src, TDst &dst);
template <class TSrc, class TDst>
__device__ void copy_async(TSrc const &src, TDst &dst);
template <class ExecutionMesh, int Bytes, class TSrc, class TDst, class TIndex>
__device__ void copy_async(TSrc const &src, TDst &dst, TIndex const &index);
template <class ExecutionMesh, int Bytes, class TSrc, class TDst, class TIndex,
class TFill>
__device__ void copy_async(TSrc const &src, TDst &dst, TIndex const &index,
TFill fill);
For an RMEM ShardLayout, the shared ir.types.register_strides reader
used by codegen and Dot access relations MUST assign compact
register-engine steps to non-Split, non-unit modes in ascending order of their
physical strides read through resolve_layout; omitted strides therefore use
compact row-major. Split and unit modes have engine step zero. This
numbering is independent of mode order in the logical-domain correspondence.
tensor_view. Rebuilds the CuTe tensor engine that a TIR TensorView
describes from its typed pointer and emitted layout. Residency stays on the
pointer engine; a later make_shard_tensor adds the logical global and shard
layouts without recovering the source tensor.
copy.
One entry, taking tensors, with the transfer shape, the strides, the element
types, the share this instance owns and the move width all read off the operand
ShardLayouts.
- constraints:
- The move width is a property of the layouts, not of the call. It is the
largest power of two both operands run contiguously over —
max_common_vectorof the two projected layouts, which §2.4.1 derives statically from the shard layout and the mesh — capped so neither side's move exceeds 16 bytes. - A dtype change does not force the width to one element: the wide load still holds and the conversion happens on the way out. The cap follows the wider of the two element types.
- Register fragments take the element path. Addressing
&frag(i)under a loop the compiler cannot unroll spills the fragment to local memory, which costs more than the wider move saves. - Where the shard's offset lands is the one thing a layout cannot state, so the implementation tests the two base pointers against the width's alignment and falls back to the element path when it does not hold. That is one comparison, not a scan.
- The call site names no vector width and no thread index. A copy split across
a block's threads is a thread-scoped mesh in the operands' shard layouts;
the entry then copies exactly the slice
local_tensorhands each thread.
copy_async.
- constraints:
- The call is non-blocking; generated code orders later reads through
cp_async_commitandcp_async_wait. - The name is the contract: a pair of projected views
cp.asynccannot move a 4-byte-or-wider word of does not compile, and takesops::copyinstead. There is no synchronous tier hiding under this entry for the caller to wait on. - The two projected slices hold the same number of elements. Where an instance owns part of a wider destination, that is the destination's shard layout to state; the source's offset is not borrowed for it.
- Runtime implementation details such as vector width and tail handling live in code comments, not this spec entry.
2.6.3 ops/dot.cuh¶
// include/tilefoundry/runtime/cuda/ops/dot.cuh
template <class Lhs, class Rhs, class Dst,
class Ws = reduce_impl::no_workspace_t>
__device__ inline void dot(Lhs const &lhs, Rhs const &rhs, Dst &dst,
Ws &&ws = {});
dot.
dst = sum(lhs * rhs). A matrix-vector product's inner loop is one statement,
not an elementwise followed by a reduce: materialising the product first
would cost a register per element of the row, which is what makes that
pair the wrong spelling here.
- constraints:
- Two tiers, one entry. With no workspace the contraction lives inside a warp and one butterfly finishes it; with a workspace it spans the block, each warp posting a partial into a slot. Either way every participant leaves holding the total, so a caller never broadcasts it back.
- The load width is the operands' shard layouts', the same question
the
copyentry answers: the run both sides share, capped by what both are aligned for and by 16 bytes. A row split so that a lane owns four contiguous elements loads eight bytes at a time; one that leaves it eight loads sixteen. The width appears neither at the call site nor on the op surface — the runtime derives it inside the entry and publishes nothing that reports it, for the reason §2.6 gives for publishing no tier predicate. - The fold uses a fixed number of independent partial accumulators, and the tree that combines them is this op's summation order. One running sum would make a row a chain of dependent multiply-adds, and the loads could not run ahead of it: what a warp then waits on is the row's latency, not its bytes.
- Accumulation is f32 regardless of the operand type.
2.6.4 ops/elementwise.cuh¶
// include/tilefoundry/runtime/cuda/ops/elementwise.cuh
template <class Fn, class TOut, class... TIn>
__device__ void elementwise(TOut &dst, Fn fn, TIn const &...src);
elementwise.
dst(i) = fn(src(i)...) over the destination's local domain. The one pointwise
op: arity is the pack's length, the operation is fn, and a broadcast operand
is a stride-0 mode on that operand's layout.
- constraints:
dstandfnlead, because the sources are a variadic pack. Zero sources is a fill, one is a map, two a combine; the same instantiation serves all three.fnis anything callable. The primitives the runtime ships (ops/primitive/) are values, not ops, and are not in this section's list;reduceanddottake the arity-2 ones too. There is no separate tag path and lambda path.- The extent is
size()of the projected destination. No count parameter: a hand-passed one is a second statement of a fact the layout already holds, free to disagree with it. - No shape or domain check. A mismatched operand is a codegen bug, and shapes are computed on the compile side.
- The only conversion is the one
dst(i) = ...implies. The runtime performs nostatic_castof its own — it does not choose a compute precision. A caller that wants one states it infn(float(x) * scale), which mirrors TIR, where a cast is aCastnode. - Broadcast is a layout fact. A per-row scale is the operand read through
(M, K):(1, 0), a per-column weight through(M, K):(0, 1), a scalar through(N):(0).composestates this against the operand's own layout, so a strided or sharded operand is read where it actually lies.
2.6.5 ops/mma.cuh¶
// include/tilefoundry/runtime/cuda/ops/mma.cuh
template <class TA, class TB, class TC>
__device__ void mma(TA const &a, TB const &b, TC &c);
mma.
c += a @ b for one atom. The operands are the calling lane's already-gathered
fragments, and the entry issues one instruction.
- constraints:
a,b, andccontain exactly the atom's(8, 4, 4)values per lane. Other shapes are compile-time errors.- The fragment map is stated by each operand's
ShardLayout; the runtime publishes no fragment-coordinate function and no accumulator constructor. - Today's atom is
SM80_16x8x16_F32BF16BF16F32_TN: bf16 operands, f32 accumulate. Another instruction is another atom under the same entry, not another entry.
2.6.6 ops/reduce.cuh¶
// include/tilefoundry/runtime/cuda/ops/reduce.cuh
struct mean_op {};
struct absmax_op {};
template <class Op, class Axes, class Src, class Dst,
class Ws = reduce_impl::no_workspace_t>
__device__ inline void reduce(Src const &src, Dst &dst, Ws &&ws = {});
reduce.
- constraints:
reduceis the only public runtime reduce entry; tier names and helper functions are internal.- Sharded operands derive the active tier and warp grouping from
(src, dst)shard layouts inside the runtime.wsis a resource the layouts demand and never the tier selector: a reduce mesh that stays inside one warp takes the intra-warp tier and needs no workspace, and one that crosses warps does not compile without one. - Both operands name one mesh and give one attr per mesh axis; the reduced
set is
src[i]split or partial againstdst[i]broadcast. - The combine domain is
float—init,elem,combineandfinalizeall are — so an element typefloatcannot hold exactly (a 4- or 8-byte integer, a packed key) is refused at the entry rather than rounded. - Plain operands derive extents from the operand rank and size inside the runtime.
- A reduction whose reduced axis crosses CTA boundaries is not supported.
2.6.7 ops/sync.cuh¶
// include/tilefoundry/runtime/cuda/ops/sync.cuh
template <int Id> struct BarrierId {
static constexpr int value;
};
template <int Id> inline constexpr BarrierId<Id> bar_id;
struct no_resource_t {};
template <class T> struct barrier_id_traits {
static constexpr bool value;
};
template <class TMesh, TopologyScope... Topos, class TRes = no_resource_t>
__device__ inline void sync(Mesh<TMesh, Topos...> const &mesh,
TRes resource = {});
Terms. A tier is the barrier a mesh reaches: grid, block, warp, or a named barrier. A mesh that reaches none is refused by the name of what it lacks.
- constraints:
- A mesh that reaches no barrier is refused by the name of what it lacks, not approximated. Silence would read as "not implemented yet"; the refusal says which of the two it is.
sync.
- constraints:
- Which barrier runs is the mesh's own answer and never a call-site template
argument: the topology scope says the level, and the mesh layout says both
how many instances (
size(layout)) and which instance the run starts at (its offset — a sliced mesh is acute::ComposedLayoutmapping a coordinate tooffset + outer(c), (§2.3.2)). - The classification MUST agree case for case with
classifyinir/tir/sync.py: an un-slicedctamesh takes the grid barrier; a mesh covering the whole block takes the warp's convergence when the block is one warp wide and the block barrier otherwise; a mesh fitting inside one warp takes the warp's convergence; a warp-aligned run covering part of the block takes a named barrier. - Two meshes name no barrier and are rejected at the call site: a sliced
ctamesh (the CTAs outside it never arrive) and a cross-warp run whose base or count is not a multiple of 32 (part of a warp would be inside the barrier and part outside). Both are deadlocks, not slow paths. - The runtime states which tier needs which resource and allocates none. A caller that supplies none where one is needed MUST fail to compile with a message naming the resource; a caller that supplies the wrong one MUST fail the same way. A grid counter is a fact about the launch and a free named barrier a fact about the whole kernel; neither is knowable from one mesh.
- A named barrier id MUST be a compile-time value in
1..15. Id0is the one__syncthreadsarrives at and MUST be rejected:bar.synccounts arrivals per id, so a subset posted to it releases a whole-block barrier early. An id MUST NOT be derived from the mesh — two meshes sharing a base and differing in count would collide. - Codegen emits only
sync; it does not call lower-level barrier helpers. - Grid barriers require every CTA of the launch to be co-resident and to execute the barrier.
- A grid barrier's counter pair is zero-initialized before first use and is owned by the generated module. Whether one exists is a fact about the launch, so the caller passes it; with none, a cooperative launch's grid group is used instead.
2.6.8 ops/copy_async_bulk.cuh¶
// include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh
template <class Src, class Dst>
__device__ inline void copy_async_bulk(Src const &src, Dst &dst,
uint64_t *bar);
copy_async_bulk.
Stage a tile into shared memory and signal an mbarrier when it is readable. It
is not a tier of ops::copy_async: there every thread issues its own load and
a commit closes the group, so the thread that issues is the thread that waits;
here completion lands on a barrier, which is what lets a consumer wait for a
tile it did not fetch.
- constraints:
- The operands are tensors, not addresses: the entry takes the transfer shape,
the strides and the element type from the operand
ShardLayouts. A caller passing a pointer and a byte count would be choosing the tier itself, which is what §2.6 puts behind one entry. - One instruction, and a run-time hand-off.
cp.async.bulkmoves the run: one elected thread issues it, nothing blocks, and it is still in flight at the next statement. A byte count off the 16-byte grain hands off to the element path at run time, every thread in the block taking part. Which one runs is not a compile-time choice: the operands must satisfy the preconditions below, and a layout that fails them is refused, not routed elsewhere. It was a static tier once, selected on whether either operand coalesces to one run — but a layout that fails that test is exactly a layout the element path addresses wrongly, so the failing branch computed a wrong answer instead of rejecting the call. - Both paths leave
barcompleting when the data is readable, and both are safe to call from every thread in the block. Consumers wait on the phase and never learn which ran. - constraints on the operands, each a
static_assert: every mesh axis leaves both tiles whole (neither is split across instances); the destination is aShardTensor(the element path needs its mesh for the instance count and the barrier); both projections are one contiguous run; the element types match, since bytes move unconverted; and the two slices are the same size, because the count is read off the source and written at the destination's origin. baris a resource parameter, the wayops::reduce'swsis: a 64-bit shared-memory word, not a sharded tensor, so it carries noShardLayoutand there is nothing for a tier to read off it. Arming that word, testing its phase and releasing it are singlembarrier.*instructions which read nothing off a layout either, so they are not ops and the runtime publishes no entries for them; a caller writes them where it builds the ring. The phase parity alternates0, 1, 0, ...across successive completions, which is what lets a fixed ring of barriers serve a stream of any length: stagetof a ring ofnwaits on parity(t / n) & 1.- The bulk tier's arrival declares the byte count on the instruction that issues the copy, so the declared and delivered counts are one expression and cannot drift; the caller does not arrive separately.
- The transfer must be a whole number of 16-byte grains for the instruction to have defined behaviour. The extent is what the shard leaves behind, not a property of the layout type, so an off-grain extent is a run-time hand-off to the element path inside the same entry — same barrier, same result.
2.6.9 ops/ldmatrix.cuh¶
// include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh
template <class Src, class Dst>
__device__ inline void ldmatrix(Src const &src, Dst &dst);
ldmatrix.
One warp loads a dense shared-memory (16, 16) bf16 tile into the per-lane A
fragment consumed by ops::mma's current SM80 atom.
- constraints:
- All 32 lanes issue
ldmatrix.sync.aligned.m8n8.x4.shared.b16together. Each lane contributes the 16-byte-aligned address selected by the source's declared shard layout; a full-broadcast shard still uses that declaration, not the backing allocation's incidental layout. - The four returned registers are stored in the destination fragment's
layout order so that
ops::mmaobserves the PTX register order unchanged. - The TIR declaration fixes the source shape/dtype and destination fragment; the runtime entry does not choose an atom or a fragment layout.
2.7 cuda/utility/¶
tilefoundry::shuffle_xor, tilefoundry::shuffle_elect and
tilefoundry::warp_reduce (runtime/cuda/utility/warp.cuh) are the runtime's
warp-scoped primitives: part of its callable surface, and no part of its op
surface.
What is a warp primitive. An op takes tensors and reads its geometry off their shard layouts (§2.6). A warp primitive takes a value and a lane count — in one case not even that, only the calling thread's own id — so there is no layout for a tier selection to read and nothing an operand would have stated; each is one instruction or a short run of them, with one implementation and no tier at all. The runtime's warp-scoped surface holds these and nothing else:
| entry | what it takes instead of a layout |
|---|---|
shuffle_xor |
one lane's value, the lane-id bits to exchange across, the participating lanes |
shuffle_elect |
nothing; the answer is a fact about the calling thread's id |
warp_reduce |
one lane's value, a combine functor, and the lanes per fold |
Being outside ops:: is not licence to grow. A candidate taking a
ShardTensor or a Mesh is an op, and belongs in §2.6's closed
list or nowhere; a candidate offering a template parameter the hardware then
clamps to one value is not an entry at all, because a parameter every value but
one is clamped away from states nothing.
Why both surfaces exist. ops::reduce is the op for folding a
tensor over a mesh, and it is built on the butterfly of
warp_reduce — utility/warp.cuh is included before
ops/reduce.cuh so the intra-warp tier folds through warp_reduce rather than
spelling the butterfly a second time. A caller folding something it just
computed, rather than something it stored, has no tensor to hand an op and
reaches for the primitive.
Not ops. These take lanes and a width, not tensors and a mesh, so no layout describes them and no tier is read off anything.
2.7.1 utility/warp.cuh¶
// include/tilefoundry/runtime/cuda/utility/warp.cuh
template <class T>
inline constexpr bool is_shuffle_native_v;
template <class T> struct ShuffleXor {
__device__ T operator()(T value, int lane_mask,
unsigned member_mask) const;
};
struct Elect {
__device__ bool operator()() const;
};
template <class T>
__device__ inline T shuffle_xor(T value, int lane_mask,
unsigned member_mask = 0xFFFFFFFFu);
inline bool shuffle_elect();
template <class Combine, int Width = 32, class T>
__device__ inline T warp_reduce(T value);
shuffle_xor.
- constraints:
- The exchange is warp-scoped: a lane's result is defined only where the lane
it pairs with is in
member_mask. - A
Tthe shuffle intrinsic does not take directly is exchanged word-wise, and one whose size is not a whole number of 4-byte words MUST fail to compile. Which types the intrinsic takes is a code-level fact and lives in code comments, not this entry.
shuffle_elect.
- constraints:
- It elects one thread of the CTA, not one lane of each warp. An mbarrier armed for a single arrival is correct only under that reading.
- It takes no width and no participating-thread count. The elected thread is always in the first warp, so a leading run wider than a warp elects what 32 threads would, and electing among fewer is a different instruction rather than a parameter of this one.
- Which thread is asked is the linearised thread id
(§2.2.1) and not
threadIdx.x, so ids0..31are the first warp whatever shape the block has. A thread outside that run answersfalsewithout reaching the warp-scoped election, whose result is defined only for the lanes in its member mask.
warp_reduce.
- constraints:
- A butterfly over
Widthlanes —log2(Width)shuffles — which leaves the run's result in every lane of the run, not only in its first. WidthMUST be a power of two in2..32; a violation is astatic_assert, not a run-time check.- A
Widthunder 32 leaves the runs independent, which is what a per-row reduction over a tile laid out several rows to a warp needs. The alternatives are a shared round trip or leaving every lane but one row's idle. Combineis invoked asCombine{}(a, b)— the spellingops::elementwiseandops::reduce(§2.6) use, so a merge is defined once for all three and the op tags are validCombinevalues as they stand.