TileFoundry Spec — hir (@func pure SSA dataflow IR)

Defines HIR, the pure SSA-as-DAG dataflow IR: its Expr constructs — the Function container, the structured-SSA exceptions LoopRegion and MeshRegion, and the HIR Op subdirectories (math / tensor / nn / shape / sharding) — together with their HIR-specific typing rules. Mesh scope is authored in the parser (parser) and its Mesh / Topology are defined by shard (shard §5); HIR links to those owners where a construct carries the result.

flowchart TB
    Expr["<b>Expr</b><br/>(core-ir)"]
    Op["<b>Op</b><br/>(core-ir)"]
    Function["<b>Function</b> (Expr)"]
    LoopRegion["<b>LoopRegion</b> (Expr)"]
    MeshRegion["<b>MeshRegion</b> (Expr)"]
    HirOpBase["<b>hir.Op</b> subclasses<br/>math / tensor / nn / shape / sharding"]

    Expr --> Function
    Expr --> LoopRegion
    Expr --> MeshRegion
    Op --> HirOpBase

1. HIR Expr constructs

HIR values are Expr nodes (core-ir §2): a Function container, the loop-phi-shaped LoopRegion, the execution-region MeshRegion, and value Op calls. HIR is pure SSA-as-DAG — there are no Region / Block abstractions and no Stmt sequence; LoopRegion carries loop-phi-shaped SSA and MeshRegion carries the structured execution region.

1.1 Function

class Function(Expr):
    """HIR's function container; its value type is the function signature.

    Attributes:
        name: attribute; the function name; call sites resolve through Module's symbol table.
        params: attribute; each Var carries a type annotation.
        body: attribute; a single Expr — typically a Call DAG; None for a dispatch prototype.
        return_type: attribute; TensorType for single output, TupleType for multi.
        specializations: attribute; Dispatch patterns carried by a variant.
        variants: attribute; Shape-specialized implementations carried by a prototype.
        converters: attribute; Weight names paired with offline converter functions.
    """

    name: str
    params: tuple[Var, ...]
    body: Expr | None
    return_type: Type
    specializations: tuple[Pattern, ...] = field(default_factory=tuple)
    variants: tuple["Function", ...] = field(default_factory=tuple)
    converters: tuple[tuple[str, "Function"], ...] = field(default_factory=tuple)

- constraints: - an Expr subclass whose value type is the function signature; always returns by value (explicit output params are TIR-only). Typing and shape-dispatch rules are stated below. - mutable during the compiler's authorised typing, metadata, and specialization updates; fields that are not updated retain structural equality and hashing semantics. - value-range metadata is refreshed by a non-type-writing whole-function inference walk. Parser construction MAY attach an already-proved range to a leaf whose value is fixed by lexical geometry; MeshCoord uses its concrete mesh axis extent for this purpose. - a Function MUST NOT declare or override execution context. The Module that owns it declares the Target and the ordered Topology hierarchy its body runs against (core-ir §1).

Kernel invocation. Entering an HIR Function from Python begins one kernel invocation. A Call from one HIR Function to another is a device call inside the current kernel invocation and never begins another one, regardless of which Module owns the callee. A Module is a static container and an invocation is a dynamic event: Python entering one root twice is two invocations, and a Python loop entering it N times is N. Module ownership, topology equality, call-graph depth, source nesting, and how often a call site repeats take no part in this rule: a repeated site and a site a loop varies repeat device work inside the invocation they are already in.

Caller	Callee	Meaning
Python	HIR Function, directly or through a Module entry	Begin one kernel invocation
HIR Function	HIR Function	Same-kernel device call
Python	Plain Python Module method	Ordinary host orchestration
HIR Function	Plain Python method	Invalid

A plain Python Module method stays on the host. Each HIR Function it enters therefore begins its own invocation; the method itself is not interpreted, traced, or dummy-run as HIR.

Function.body is a single Expr (usually a Call DAG, possibly nested inside a LoopRegion). HIR has no Stmt sequence; name reuse lives in the parser's lexical environment, not the IR. The one exception is a dispatch prototype — a specialized function's base, whose body is None (written pass in the DSL); it declares the signature and dispatch envelope only, and its variants carry the implementations (see Shape dispatch and specializations below).

Function always returns by value; explicit output parameters are TIR-only (see tir). Crossing that boundary materialises the HIR return value into a TIR explicit output buffer parameter.

The return type MAY carry a Partial(reduction) in a TensorType, or in any tensor field of a nested TupleType. Function construction, type inference, and call elaboration MUST allow that state. A Function boundary MUST preserve the ShardLayout mesh and per-axis reduction; it MUST NOT complete the value or reject it merely because it is Partial.

A with Mesh(("cta",), layout=...) as cta: inside the body names a level of the execution domain the owning Module declares. The name MUST resolve to one of that Module's effective Topology levels, and the scope creates a parser-lexical mesh binding; ShardLayout.mesh MUST point at an active binding on the lexical path. A Mesh MAY map fewer levels than the domain declares, but it MUST NOT create a level or change one's extent.

Return type. Function.return_type is the HIR result Type: a TensorType for one result or a TupleType for multiple results. It is part of the function signature and is the result component projected into Function.type.

Value type. Function.type is the IR-level CallableType (types §7) projected from params + return_type. The projection is fixed at construction and stays consistent across construction sites.

Call typing — visitor-scoped inference. A same-module Call keeps its authored Function template as target; an instance call keeps its ModuleList slot or TupleGetItem(ModuleList, index) target (core-ir §2.1). Its result type is inferred by seeding a new visitor memo with the actual argument types bound to the callee's formal parameters, then walking the callee body in a child context that retains the call site's mesh. A callee without a mesh declaration runs in that scope; a MeshRegion in its body composes with it using the usual nested-region rules. This is type inference only: it does not rebuild a Function, mutate the target, or create a per-call instance. The callee body is read under the callee's effective Target and none of its nodes is written; only a walk that owns the caller's body stores the result on the caller's Call. Caller-supplied layout (sharding) flowing into a layout-unconstrained parameter propagates through the body, including through a Tuple or LoopRegion return.

Within one inference traversal, repeated calls to the same Function object with equal argument types and equal call-site mesh scopes MUST reuse the previously inferred result type. The cache key uses callee identity, the argument-type tuple, call-site mesh, and current Module with its effective target and topologies; the cached value is only a Type, never a derived Function or a replacement Call.target.

Argument types bind to supplied parameters in order, following the target and instance-resolution rules in core-ir §2.1. For an omitted ConstTensor parameter, its declared type seeds the callee visitor memo; its value comes from the instance's resource (runtime §1.3).

Caller and callee MUST resolve one effective Target and one effective topology hierarchy: a same-kernel call is one execution context, whichever Module owns each end. Inheritance is the canonical spelling, and a callee declaring the caller's hierarchy explicitly is accepted only when the resolved tuples are equal. A mismatch is invalid; it is not a nested launch.

Argument ↔ parameter binding is:

• Arity MUST match — exactly one argument per supplied parameter.
• A declared field that is not None MUST equal the corresponding argument field. A declared None leaves only that field undecided and binds it from the argument. This rule applies recursively:
• TensorType.shape and dtype always match exactly. storage matches exactly unless the parameter states UMAT, whose residency is undecided. layout=None leaves layout undecided; a stated layout recurses below.
• For Layout, stated shape and strides fields match independently. For ShardLayout, stated mesh, attrs, and nested layout fields match independently, and the nested layout follows this same rule.
• Any type not covered by that recursive structure requires exact equality.
• After a successful match, inference binds the parameter to the argument's full type. Fields the declaration constrained have already been proved equal; fields it left undecided retain the argument's concrete values.
• DimVar shapes keep envelope matching — inference does not monomorphize a dynamic shape into a concrete one (that is Shape dispatch and specializations below, unaffected by call typing).

The per-mesh-axis Partial state is part of the actual argument type. When a layout-unconstrained parameter binds to a sharded argument, inference MUST carry each Partial(reduction) at its original mesh-axis index through the body and into the concrete return type, including tuple fields. Only an explicit Reshard or allreduce may complete that state.

When the body cannot express a propagated sharding (e.g. a reshape whose layout factorization straddles a new axis), typeinfer fails at that op, not at the boundary. A dispatch-prototype callee (variants != (), body is None) is not walked: the call's result is the declared return_type and the None body is never inspected (variant selection is Shape dispatch and specializations below).

Call typing does not record provenance or bound dimensions because no call-site function instance is created. Explicit specialization remains the separate operation that produces a derived Function; such a derived function records its origin and chosen dimensions as specified by the Function specialization API below.

The placement sugar in a signature emits Layout(strides=None) (parser.md §2.1). Under the recursive rule, this leaves only strides undecided: the layout's stated shape, mesh, and shard attributes remain constraints, while the argument's concrete strides bind into the callee. A verbose Layout(strides=tuple) states a concrete stride contract and MUST match exactly. Signature binding does not materialize either form into a different layout.

SSA shape. HIR is pure SSA-as-DAG — sharing of intermediate results is expressed by Python object identity:

• Single use: nest the Calls. Call(Binary(kind=MUL), (Call(Binary(kind=ADD), (a, b)), c)) does not name the inner Binary result.
• Multiple uses: the parser binds c = tf.add(a, b) in its lexical env so subsequent mul(c, c) / sub(c, d) share the same Call node. The IR has no binding nodes; DAG edges express "same value".

There are no Region / Block abstractions in HIR. The single structured exception that carries loop-phi-shaped SSA is LoopRegion (§1.2). Everything else is a pure Call DAG.

Function typing rules. Function is not an Op and is not registered in the Op typeinfer registry. TypeInferVisitor handles it directly as a Call.target; structural signature rules are enforced by the HIR verifier:

• Function.body is a single Expr; Stmts MUST NOT appear.
• Function.params entries MUST be Vars.
• Within a Function signature, every occurrence of a same-name DimVar across params and return_type MUST agree on its (lo, hi) bounds; a disagreement is a verify error. A RangePattern specialization MUST anchor to a DimVar reachable from an input parameter and lie within that DimVar's envelope (see Shape dispatch and specializations below).

Shape dispatch and specializations. Function is the sole HIR function Expr. Shape-dispatch is carried on a single base Function through its variants field; there is no separate specialized-function type. The field is the IR-side carrier for the parser surface (parser.md).

The specializations, variants, and converters tuples are canonical Function fields. converters records (weight_name, converter) pairs in registration order and is sealed recursively with the base and its variants.

Structure. A Function is exactly one of three shapes:

• normal — specializations == (), variants == (), body is an Expr. An ordinary function.
• dispatch prototype (base) — specializations == (), variants != (), body is None. Declares the signature and dispatch envelope only; the implementations live in its variants.
• variant — specializations != (), variants == (), body is an Expr. A shape-specialized implementation registered on a base.

Nesting is exactly one level: a variant MUST NOT itself carry variants. In a sealed (verified) Module the invariant is body is None ⟺ variants != () — a function with no body and no variants is uncallable and invalid, and a real body combined with variants is invalid. During authoring the base is transiently body is None, variants == () between @func def f: pass and the first @f.specialize(...); this unsealed state is allowed only until the base enters a Module (see Authoring freeze below).

• variants is a canonical IR field — it participates in structural equality, hashing, and canonical printing.
• Parser-time Function roles use the following naming and handle contract:

role	purpose	lookup handle	_ allowed	binding uniqueness	parser ledger classification
kernel / prototype	Module execution unit	fn.name	no	unique within the Module class body	_Entry.bound[binding] = KERNEL; id(fn) is not in _Entry.owned
variant	Implementation for a dim range	f"{base}${dim}${lo}_{hi}"	no	unique within the base	_Entry.bound[binding] = VARIANT; id(fn) is in _Entry.owned
converter	Offline weight conversion	f"{base}.converter[{weight}]"	yes	reusable	_Entry.bound[binding] = CONVERTER; id(fn) is in _Entry.owned

The parser enforces this table during decoration. A parser-authored variant MUST carry a non-underscore display label, while a Function built directly in IR MAY omit the non-canonical label. The label MUST NOT participate in structural equality, hashing, or the canonical signature, and nothing MAY select an implementation by it. A printer emitting an unlabeled IR variant MUST choose a valid source binding so the emitted definition remains importable. - Every variant of a base MUST share the base's name, params, and return_type: a variant specializes the body, not the signature. A variant runs in the same execution domain as its base because both are owned by the same Module. - .specialize(pattern, *, mesh=None) accepts an explicit function mesh with the same semantics and Mesh type check as @func(mesh=...). It does not inherit the prototype's mesh; variant and base parameter and return types MUST still match. Authored layout forms and their acceptance rules are owned by shard §6 and parser §2. - A variant carries exactly one RangePattern in specializations. The canonical signature is ";".join(f"{p.dim_var}${p.lo}_{p.hi}" for p in specializations) (v0 allows only RangePattern). Two variants of one base MUST have distinct canonical signatures.

Envelope coverage. A dispatched function's parameter TensorType.shape carries a DimVar(name, lo, hi) whose (lo, hi) is the dispatch envelope; RangePattern references that DimVar by name. The variants' closed ranges MUST partition the envelope — pairwise disjoint and jointly complete (their union is exactly the DimVar's closed [lo, hi] envelope). Adjacent closed ranges are written [.., c - 1] then [c, ..]. Every in-envelope shape therefore selects exactly one variant.

Prototype body. A base's body is None: the prototype is never typeinferred, lowered, or evaluated as a body. Only its variants carry executable bodies. There is no base body to fall back to.

Dispatch resolution. A Call whose target is a dispatch prototype (variants != ()) is a dispatch call: the variant whose RangePattern matches is selected and is the call's result. Evaluation selects from the call's concrete argument shapes; specialization selects from the caller's stated dimension bindings. Both use the same variant table. A shape outside the envelope matches no variant and is an error; there is no base body to fall back to (the prototype body is None). A Call whose target has variants == () is a direct call to that body.

Authoring freeze. Variants accumulate during authoring, before the base Function enters a Module (core-ir §1). A sealed base rejects further variants. Because variants participates in hashing, a base MUST NOT be hashed while still accumulating variants. A top-level Module.functions entry MUST NOT be a variant: a top-level Function with specializations != () is a verifier error.

1.2 LoopRegion

class LoopRegion(Expr):
    """Loop-phi-shaped structured SSA folding a tile-style loop into one Expr value.

    Attributes:
        induction_var: attribute; loop induction Var, ranging over range(start, extent, step).
        params: attribute; entry bindings, carry slots followed by invariant captures.
        args: attribute; matching values evaluated in the enclosing scope.
        body: attribute; the loop body Expr.
        yield_values: attribute; loop-phi carry chain (equal lengths).
        extent: attribute; iteration-domain stop (half-open).
        step: attribute; induction-var stride.
        start: attribute; iteration-domain start (default 0).
    """

    induction_var: Var
    params: tuple[Var, ...]
    args: tuple[Expr, ...]
    body: Expr
    yield_values: tuple[Expr, ...]
    extent: ShapeDim
    step: ShapeDim
    start: ShapeDim = 0

- constraints: - the only HIR exception to pure Call DAG: loop-phi-shaped structured SSA that folds a tile-style loop into one Expr value; type is TensorType (single carry) or TupleType (multi-carry). - mutable during the compiler's authorised typing and metadata updates.

Iteration domain. Both DSL loop surfaces — for i in tf.tile(...) and for i in range(...) — lower to this one node; they share the domain (start, extent, step) and differ only in the loop-variable binding (tile binds a parser-side Python slice, while range binds a scalar; see parser §2.1). range is not unrolled. induction_var ranges over range(start, extent, step): start and extent are the half-open [start, extent) Python-range endpoints (so extent is the stop value, not a count). start defaults to 0 for tf.tile(stop, step) and range(stop); the tf.tile(start, stop, step) and range(start, stop[, step]) surfaces set it. Each of start / extent / step is a ShapeDim (types §4).

For tf.tile(stop, step) or tf.tile(start, stop, step), the parser-side window at one iteration is [induction_var, induction_var + step). The induction value is already a coordinate in range(start, stop, step), not an ordinal to multiply by step.

• When start / extent / step are static int, the trip count is recoverable from the node alone, without the parser-side window binding (parser §2.1).
• Every DimVar referenced by a ShapeDim start / extent / step MUST be bound by the enclosing Function's parameter shapes. Resolution substitutes each such DimVar with the corresponding argument-shape size and folds the dim Expr to a value n. The resolved start and extent MUST be non-negative integers and the resolved step MUST be a positive integer; otherwise resolution MUST raise. An unbound DimVar MUST raise.
• A ShapeDim start / extent / step is resolved by the evaluator at call time against concrete argument shapes; its trip count is not statically recoverable from the node alone.

Carry-out semantics. The parser populates the carry chain when an HIR grid-loop body contains an ast.Assign whose single Name target binds an outer-scope name:

• the carried name becomes a phi Var in params,
• the pre-loop binding of that name becomes the matching entry in args (the carry's value on the first iteration),
• inside the loop body the same name resolves to that phi Var,
• after the loop, the post-region binding refers to the LoopRegion itself (single carry) or a tuple_get_item of it (multi-carry, when len(yield_values) > 1).

args are value Exprs traversed and rewritten by the visitor / mutator; params and induction_var are binding sites. len(params) == len(args) and k = len(yield_values) <= len(params). The first k slots carry loop state: first-iteration values come from args[:k], and later values come from the matching yield. Remaining slots are invariant captures evaluated once outside the loop. A no-carry loop has k == 0 and may still have captured parameters.

A LoopRegion MUST be isolated like a MeshRegion: its body and yields MUST read outside values through entry parameters, never by embedding an args value. The parser captures each external expression binding across every region boundary; an external Call therefore belongs to its enclosing structural scope. Each structured region exposes its invariant bindings through captures(). MeshRegion.captures() pairs every parameter with its argument; LoopRegion.captures() skips the carry slots and pairs only the remaining parameters and arguments. Analysis stores those pairs on its iteration-scope tree, so capture resolution does not need to inspect the region kind.

Type inference derives every argument outside the loop and binds each parameter to its argument's type; a region parameter has no type of its own, so a type stored on it by an earlier inference does not constrain the entry value. It seeds the inner visitor with the induction binding and all parameter bindings, then derives the body and yields. Every yield MUST satisfy types_compatible(entry_type, yield_type) for the parameter it carries into, and k > len(params) is an error. The carry result is the entry type of the first k parameters: a more specific yield does not narrow it, and a loop that runs no iteration has that type. There is no fixed-point iteration.

LoopRegion.type is TensorType (single carry) or TupleType (multi-carry); the value is the Expr itself, not a Call. Parser-side rules: see parser §3.

1.2.1 MeshRegion

class MeshRegion(Expr):
    """Represent the execution domain of one structured HIR region."""

    mesh: Mesh
    params: tuple[Var, ...]
    args: tuple[Expr, ...]
    body: Expr

• constraints:
• mesh is the mesh opened by this region. Visitors compose it with the enclosing execution mesh at the body edge; the args edge is evaluated outside the region.
• params and args have equal length. The body may reference each captured value only through its corresponding parameter; an argument is never read directly from the body. This is the same binding boundary as Function. Each parameter takes its argument's current type, as for LoopRegion (§1.2).
• The region result is reachable through the values that escape its lexical body. A single escaping value is the region's body; multiple escaping values are carried by a Tuple and read through TupleGetItem projections.
• MeshRegion is type-transparent: MeshRegion.type == MeshRegion.body.type. Its execution domain never changes the result's ShardLayout.
• A body with no value is permitted only when at least one value escapes through a later binding; a body with no value and no escaping value is a parse error.

Minimal example — loop-carried accumulator:

# example
acc = tf.zeros((M,), "f32", storage="rmem")
for i in tf.tile(K, BLOCK):
    acc = acc + load_tile(x, i)
# After the loop, `acc` resolves to the LoopRegion value.

becomes (sketched):

# example
LoopRegion(
    induction_var = i,
    params        = (acc_phi, x_param),
    args          = (Call(Zeros(...), ()), x),   # pre-loop values
    body          = Call(Binary(kind=ADD), (acc_phi, load_tile(x_param, i))),
    yield_values  = (Call(Binary(kind=ADD), ...),),
    extent        = K,
    step          = BLOCK,
)

1.3 Op

HIR Ops are organised under tilefoundry.ir.hir.<namespace>/; the subdirectory is file organisation, not a separate IR layer. A custom Op records its full contract (fields, typing / verifier rules, worked examples) in its catalog entry below; a consensus Op needs only one sentence or a grouped external reference, per SPEC-RULES. The op name is the pointer — code carries no back-link to this catalog. ParamDef plumbing stays in code; the mechanism is owned by core-ir §2.3.

HIR-specific typing hooks. Each op's constraints are enforced by its registered @register_typeinfer(<OpClass>) body via ctx.error(...) (visitor-registry §4):

• Local(x): x.type.layout MUST be ShardLayout. The result shape contracts per the Split axes; dtype is preserved; layout becomes the corresponding local layout.
• Reshard(x, layout, storage): layout and storage are attributes (compile-time constants); the output preserves x.type.shape (logical). Architecture invariant: after HIR typeinfer runs, every ShardLayout reachable from a value's type has concrete layout.strides (never None) — the un-materialized (strides=None) parser sugar MUST be materialized by the owning typeinfer. The per-op (layout, storage) resolution table is in the Reshard op entry below.
• Any HIR Op MUST be value-form (core-ir §2.3); emitting an effect-form Call into HIR is a verify error.
• Each result layout MUST describe that result. A view derives its layout from its source when one is stated; an op producing a distinct value derives a layout over its own result shape. An op MUST NOT copy an input layout across a shape change.

Generic, analysis-wide typing behavior is owned by semantic-analysis: relation-driven type validity (semantic-analysis §1.1), output storage of multi-input ops, and operand layout / mesh ownership (semantic-analysis §3.3). HIR ops call these services; each op's registered typeinfer owns the layout / mesh compatibility and result layout it requires, and Reshard is the explicit op that changes a value's layout / mesh.

ir/hir/math/

Pointwise arithmetic and comparison, torch semantics with TileFoundry type-promotion. User-callable names (add / cmp_eq / logical_and / …) are surface aliases (core-ir §2.3) over the kinded Ops; there are no per-name IR classes. torch element-wise ops.

One spelling is preferred, so that two authors reading the same IR write it the same way: an arithmetic or comparison operand pair SHOULD be written with the Python operator (a + b, a * b, a < b), and a sub-tensor SHOULD be written as a subscript (x[:, :, j:j + 1], x[:, :, 3]). The named forms tf.add(a, b) and slice(x, begin=…, end=…, strides=…) remain the underlying surface — they are what the operator and subscript resolve to, and they stay available where a name must be computed — but they are not the form to reach for first. Both spellings build the same IR, so the choice carries no semantic weight; leaving it open is what lets one model read one way and its neighbour another.

Binary

class Binary(Op):
    """Kind-tagged pointwise binary operation; produces a Tensor.

    Attributes:
        lhs: input; input tensor.
        rhs: input; input tensor.
        kind: attribute; binary arithmetic, comparison, or boolean tag.
    """

    lhs: Tensor
    rhs: Tensor
    kind: BinaryKind

- constraints: - Both operands MUST use symmetric right-aligned broadcasting: missing axes and size-1 axes read the broadcast operand at zero, including rank-0 scalars. The result shape MUST be the broadcast of both operand shapes. - Values follow torch pointwise semantics; dtypes do not promote. Both operands MUST already carry the same dtype, and typeinfer MUST reject a mismatch. A Python literal is an ordinary operand of the dtype it is written with — f32 for a float, i64 for an integer — and the authoring surface MUST NOT adapt it to its peer. The rejection MUST name Cast as the remedy, so the dtype a value carries is the one the author wrote. - The elementwise min / max kinds are also surfaced as minimum / maximum. - Equal plain layouts, or one plain layout paired with layout=None, pass through only when that layout describes the broadcast result. Otherwise two non-sharded operands produce layout=None; broadcasting differently shaped views does not make either operand's layout describe the result. This fallback MUST NOT accept an incompatible ShardLayout pair. - If that fallback result is a shaped value in rmem or smem, Binary MUST give the newly produced value its own C-order Layout; a concrete local result cannot reach authored analysis with an unresolved layout. - A ShardLayout operand carrying Partial(reduction) propagates to the output only when kind provably commutes with reduction (op(reduction(x)) == reduction(op(x))); typeinfer rejects otherwise, naming the offending operand and the fix (an explicit Reshard to Broadcast). Decisions are made independently for each mesh axis. Partial states on different axes are not interchangeable; ADD rejects two Partial inputs when their states occupy different mesh axes. - ADD with both operands Partial: commutes (passes) only when both carry the same reduction="sum" on that mesh axis (max/min reject — max(x)+max(y) is not max(x+y)). - ADD with one Partial operand and the other plain/Broadcast: commutes (passes) for reduction in {"max", "min"} (adding a replicated constant is order-preserving) and rejects for "sum" (sum(x)+b != sum(x+b)). - MUL with one Partial operand and the other plain/Broadcast: commutes (passes) only for reduction="sum" (scaling by a replicated constant distributes over sum); rejects for "max"/"min" (the constant's sign is not statically provable, and a negative scale flips max to min). - Every other kind / operand-shape combination involving a Partial operand (including MUL with both operands Partial) rejects: not proven to commute with any reduction.

Unary

class Unary(Op):
    """Kind-tagged pointwise unary operation; produces a Tensor.

    Attributes:
        x: input; input tensor.
        kind: attribute; unary tag including neg, abs, logical_not, rsqrt,
            exp, log, ceil, round, exp2, and log2.
    """

    x: Tensor
    kind: UnaryKind

- constraints: - Behavior follows torch pointwise semantics with TileFoundry type promotion. - exp is the natural exponential e ** x; log is the natural logarithm; exp2 / log2 are the base-2 counterparts. ceil rounds toward positive infinity; round rounds to the nearest integer with ties to even (banker's rounding, matching torch's own round semantics). - A ShardLayout operand carrying Partial(reduction) propagates to the output only when kind provably commutes with reduction; typeinfer rejects otherwise, naming the offending operand and the fix (an explicit Reshard to Broadcast). exp / log / relu / ceil / round / exp2 / log2 are monotone non-decreasing, so they commute with max / min but not sum. neg is linear, so it commutes with sum but not max / min (negation reverses order). abs / square / rsqrt / logical_not are not proven to commute with any reduction and reject a Partial operand unconditionally.

Clamp

class Clamp(Op):
    """Clamp every element to a closed interval; produces a Tensor.

    Attributes:
        x: input; Source tensor.
        min_val: attribute; Lower bound.
        max_val: attribute; Upper bound.
    """

    x: Tensor
    min_val: float
    max_val: float

• constraints:
• The result MUST preserve x's shape, dtype, layout, and storage.
• Clamp is monotone non-decreasing, so it MAY preserve a Partial(max) or Partial(min) state and MUST reject Partial(sum).

Softplus

class Softplus(Op):
    """Apply pointwise softplus; produces a Tensor.

    Attributes:
        x: input; Source tensor.
    """

    x: Tensor

• constraints:
• The result MUST preserve x's type.
• Softplus is monotone non-decreasing, so it MAY preserve a Partial(max) or Partial(min) state and MUST reject Partial(sum).

ir/hir/tensor/

Tensor structural operations; consensus ops (Transpose / Slice / Concat / Stack / ShapeOf / Rank) follow torch / numpy (torch tensor manipulation ops).

Transpose, statically positioned Slice, and Reshape derive their result layout during type inference. Their result TensorType.layout MUST NOT be None; a Layout whose strides are None states that the arrangement is known while its steps are not. A missing input layout is the compact arrangement that HIR readers assign to that value. This requirement applies to these structural ops, not to every tensor-producing HIR op.

Pattern matching and lowering MUST read this result type rather than derive a second presentation of the result.

• Transpose produces a new value, not a storage alias. It MUST permute the tensor shape and use a fresh compact C-order result layout. A ShardLayout MUST remap its split positions through the registered relation with fresh strides. Lowering MUST copy the source's permuted-stride view into the new result storage.
• Bitcast(x, layout) reinterprets the same bytes with a new shape and layout, preserving the input dtype and storage. layout is a LayoutBase; the result shape for a plain layout is the product of the nested sizes in each top-level layout mode. For an rmem ShardLayout, the result preserves x.shape and changes only its layout. Its layout modes and Split positions MUST be valid for that logical shape under the existing layout-position-to-tensor-axis rules. Both layouts MUST be plain Layout / ComposedLayout without ShardLayout components, or both MUST be rmem ShardLayout on the same mesh containing only Split attributes. Every Layout component MUST have stated strides. For rmem shards, an address is the owning mesh coordinate paired with the offset within that unit; the two address sets MUST be equal, so every element stays in its original unit and register. Broadcast and Partial shards MUST be rejected. Rmem shards MUST have an underlying Layout with stated, non-nested strides; composed underlying layouts MUST be rejected. Ownership MUST use local_layout_and_offset on the layout's factored domain. Both shapes MUST be literal and have the same element count. Both address functions MUST be injective, and every result address MUST belong to the source address set, using the complete colex address functions including composed swizzles. Type inference MUST reject violations and identify the failing end and rule. Broadcast or overlapping layouts MUST NOT be bitcast. The result-to-source coordinate relation MUST pair equal addresses, forming a bijection. Bitcast MUST alias the source buffer, incur zero traffic, lower to a tensor view, and have no instruction candidate. Evaluation MUST reorder values by these same address functions. For a row-major (N, K) tensor, a bitcast with Layout((K, N), (1, K)) expresses the same bytes as (K, N) for Q @ K^T.
• Slice is normalized as Slice(x, starts, sizes=..., strides=...). starts is a tuple of rank-0 integer operands; sizes and strides are ShapeDim attributes stored in the same IR normal form as every other dim. Its result shape is exactly the normalized sizes and MUST NOT contain an induction Var. A start MAY be dim arithmetic over an induction Var — a window moved off that loop's window by a compile-time offset. That start is an address computed where it is read, not a value some op produces, so a walk over compute ops MUST leave it alone.
• A plain-layout Slice with static starts MUST produce a ComposedLayout: its offset is the source offset plus the starts multiplied by the source strides, and its outer layout carries the sliced shape and retained strides (multiplied by any slice step).
• A source axis written as a group of modes has no single stride: its start MUST be read through the group as a mixed-radix coordinate, and the window MUST be the group's least-stepping modes, as many as its size takes, in the order the group wrote them. A start that is not a multiple of the size, a size landing on no mode boundary, and a step MUST each be refused.
• ShardLayout slices and instruction projection MUST use one window derivation. On a logical axis containing a Split mode, the window MUST be one complete suffix block in the logical row-major mode order. A mode (a*b):(d) MAY be factored as (a,b):(b*d,d). Every Split mode MUST remain in the suffix, with its stride and mesh attribute preserved, and the slice step MUST be 1. The start MUST be a multiple of the block size: constants are checked directly, and induction expressions MUST be proved aligned with isl over the enclosing loop domain carried by type inference. Prefix coordinates contribute a per-thread displacement; they MUST NOT change which thread owns an element.
• On an axis without Split modes, the existing representable affine-window rules remain: a primitive mode takes the window size and stepped stride. This includes an unsplit 6 -> 4 QKV group window, non-unit steps, and runtime starts. Static nonzero offsets wrap the retained shard layout in a ComposedLayout; runtime starts without a known displacement preserve the bare distribution. A window crossing a Split mode or lacking its required alignment proof MUST be refused.
• ComposedLayout(inner=None, outer=ShardLayout(...)) is a sharded view. Relation-driven consumers, local projection, execution-domain discovery, and Partial checks MUST read distribution from its outer layout. Its offset is input addressing and MUST NOT be copied to a consumer's newly produced value; another view MAY compose and preserve that offset.
• Narrowing a logical axis targeted by any Split MUST fail type inference: the window need not align with that mesh division. A narrowed axis represented by more than one factored layout position MUST also fail rather than guess.
• Runtime starts MUST remain ordinary Call operands. They move the view's base address, not its strides: a plain input therefore still produces a result Layout, while the runtime offset is not carried by that layout. A safe sharded slice follows the preservation rule above. The result type describes a full window; whether a loop iteration can contain that window is an analysis-domain question, not a type-inference question.

Concat

Concat([inputs...], axis=a) materializes a rank-preserving tensor by joining each input segment along a. All inputs MUST have one common rank and dtype, and every non-concatenated dimension MUST match. Negative axis values resolve against that rank. The output's concatenated extent is the sum of the input extents; every input access map is defined only on its segment and subtracts the preceding segments' extent from that axis, while the output map is the identity.

The authored inputs value MUST be one explicit list, tuple, or supported static list comprehension. Its Tensor elements flatten into Call.args in source order; direct positional tensors and implicit iterable expansion are not part of this surface.

Type inference derives fresh output ownership from those access maps. A Split on a non-concatenated axis MAY propagate when shared ownership propagation proves a zero-offset projection. A Split on the concatenated axis MUST be rejected, including when it appears only on the zero-offset first input, and the diagnostic MUST require an explicit Reshard before Concat. Dynamic ownership, incompatible meshes, nonuniform Partial states, and an unrepresentable derived layout MUST likewise fail with the Reshard remedy. An output with no real sharding receives a fresh C-order Layout.

Rank and ShapeOf

• Rank produces a rank-0 i64; ShapeOf produces a rank-1 i64 vector with one entry per input axis.
• Both results are host shape metadata with layout=EMPTY_LAYOUT and storage=umat, not device-resident tensors.
• Evaluation MUST read the concrete runtime tensor rank and extents. A symbolic input TensorType.shape is the result bound, not the value returned at runtime.
• A compile-time integer subscript, tf.shape_of(x)[k], reads one dimension and produces the canonical rank-0 i64 umat scalar for that dimension.
• Neither reads an element of its input, and neither moves bytes.

Arange

Arange(type, start=0, step=1) produces exactly type.shape[0] integer values, starting at start and separated by step. type MUST be a rank-one TensorType with i32 or i64 dtype. start MUST be a static or symbolic ShapeDim; step MUST be a positive static integer. The complete result type, including layout and storage, is the supplied type; type inference MUST return it rather than constructing another type. Coordinates are synthesized without standalone traffic; their consumers own any materialization. The op has type, evaluation, access-relation, and cost semantics but no HIR-to-TIR lowering or codegen contract.

Where

Where(condition, input, other) applies right-aligned broadcasting across all three operands and selects elementwise from the two data branches. condition MUST be bool; input and other MUST have the same dtype. The data branches, not the condition, anchor result storage and distribution. A genuinely sharded condition MUST be compatible with the distribution derived from the data branches; every Partial operand is rejected. If the resulting shaped value is in rmem or smem and no branch layout describes the broadcast result, it receives a fresh C-order Layout. Cost counts one boolean selection per result element and complete reads of all three inputs plus one result write. The op has no HIR-to-TIR lowering or codegen contract.

ArgMax

class ArgMax(Op):
    """Produce indices of maximum values along an axis.

    Attributes:
        x: input; Source tensor.
        axis: attribute; Reduction axis.
    """

    x: Tensor
    axis: int = -1

• constraints:
• x MUST have rank at least one and axis MUST resolve within that rank.
• The result MUST remove axis, use dtype i64, preserve storage, and derive a layout over the reduced result shape. A Split on a surviving logical axis propagates through the registered relation.
• The reduction axis MUST NOT be Split-sharded. A winning index cannot be recovered from independent per-device winners, so typeinfer requires an explicit Reshard instead of silently completing that split.
• x MUST NOT carry a Partial state because a winning index cannot be recovered from an unpaired per-device partial reduction.

FullLike

class FullLike(Op):
    """Produce a tensor like an input filled with a scalar constant.

    Attributes:
        x: input; Type and shape template.
        value: attribute; Fill value.
    """

    x: Tensor
    value: float

• constraints:
• The result MUST have exactly x's type and every element MUST equal value converted to that dtype.
• x is read for its Type only: no element of x is read, and every element of the result is written.

Quant

class Quant(Op):
    """Produce per-token-group quantized values and scales.

    Attributes:
        x: input; Source tensor.
        scheme: attribute; Quantization scheme.
        group: attribute; Last-axis group size.
        target_dtype: attribute; Quantized element dtype.
    """

    x: Tensor
    scheme: str = "per_token_group"
    group: int = 128
    target_dtype: DType = DType.fp8e4m3

• constraints:
• x MUST have rank at least one and MUST NOT carry a Partial state.
• scheme is exactly "per_token_group"; packed block formats are a different operation boundary. group MUST be a positive, non-boolean integer, and target_dtype is exactly fp8e4m3.
• For a static last extent, group MUST divide that extent. For a symbolic extent, the scale extent is the symbolic floor division by group, and evaluation rejects an indivisible runtime extent before reshaping.
• The result MUST be (x_q, x_scale): x_q preserves x.shape with fp8e4m3; x_scale has dtype f32 and replaces the last extent by x.shape[-1] // group. Both fields preserve storage and receive freshly derived layouts over their own shapes. Outer-axis splits propagate. A last-axis split propagates only when its factorization proves that every owner holds complete groups and the scale split is representable; otherwise typeinfer requires an explicit Reshard. A fully Broadcast ShardLayout pins no mesh and produces unsharded result layouts.
• Evaluation computes each group's f32 absolute maximum and scale absmax / 448, divides, clamps to [-448, 448], then casts to fp8e4m3. An all-zero group uses scale one and produces no NaN.

RepeatInterleave

class RepeatInterleave(Op):
    """Repeat elements along one axis; produces a Tensor.

    Attributes:
        x: input; Source tensor.
        repeats: attribute; Repetitions per source element.
        axis: attribute; Axis to expand.
    """

    x: Tensor
    repeats: int
    axis: int

• constraints:
• axis MUST resolve within the rank and its result extent MUST be the input extent multiplied by repeats; all other extents and storage are preserved.
• The result layout is unsharded. A genuinely sharded input MUST be refused; an unsharded or fully broadcast input is accepted.

Split

class Split(Op):
    """Split a tensor into equal parts; produces a Tuple.

    Attributes:
        x: input; Source tensor.
        axis: attribute; Split axis.
        num_splits: attribute; Number of outputs.
    """

    x: Tensor
    axis: int
    num_splits: int

• constraints:
• axis MUST resolve in [-rank, rank), and num_splits MUST be positive.
• A static selected extent MUST be divisible by num_splits; every output field has that extent divided by num_splits and otherwise preserves the input type.
• A symbolic selected extent is retained in each output until a tighter symbolic quotient is available.
• Evaluation materializes num_splits equal tensors in axis order. Each element carries exactly its corresponding inferred TupleType field, including that field's storage and shard layout; the tuple aggregate has no independent layout or storage.

Stack

class Stack(Op):
    """Stack equal-shaped tensors along a new axis; produces a Tensor.

    Attributes:
        inputs: input; variadic tensors to stack.
        axis: attribute; inserted result axis.
    """

    inputs: Tuple[Tensor]
    axis: int

• constraints:
• The authored inputs value MUST be one explicit list, tuple, or supported static list comprehension. Its Tensor elements flatten into Call.args.
• At least one input is required; every input MUST have the same shape and dtype. axis MUST resolve in [-rank-1, rank].
• The operation materializes one distinct result. The inserted axis is local and unsharded, and the result layout is freshly derived over the result shape rather than copied from any input.
• Stack states a relation in which input i accesses the result slice whose inserted-axis coordinate is i; every old logical axis projects to the corresponding result axis on the other side of the insertion. Shared shard propagation derives any representable ownership from that relation.
• Compatible Split ownership carries to the shifted old logical axis. A fully replicated input contributes no sharding. Incompatible input meshes or per-axis states, a non-uniform Partial across result slices, and an unrepresentable result layout MUST fail naming the conflicting input and requiring an explicit Reshard; typeinfer MUST NOT select input zero's layout or silently discard real ownership.

TupleGetItem

class TupleGetItem(Op):
    """Extract one field from a tuple-typed expression by scalar index.

    Attributes:
        tuple_value: input; Tuple-typed expression.
        index: input; Scalar field index.
    """

    tuple_value: Expr
    index: Expr

• constraints:
• tuple_value.type MUST be TupleType. A constant index MUST be in range, and the result type is exactly the selected field type.
• A dynamic index requires a non-empty tuple whose field types are all equal; the result has that common field type.
• Access through a constant index reads the selected field's leaf span. Dynamic access conservatively reads every leaf because the selected field is known only at runtime.

Reshape

class Reshape(Op):
    """Reshape ``x`` to ``new_shape``; produces a Tensor.

    Attributes:
        x: input; source tensor.
        new_shape: attribute; target logical shape.
    """

    x: Tensor
    new_shape: tuple

- constraints: - The result shape is new_shape; size(new_shape) MUST equal size(x.shape). - Evaluation MUST resolve dimensions from the invocation bindings before reshaping. At most one unbound dimension may be inferred by element count. - A plain C-order input reshapes to a C-order Layout over new_shape. An input with no assigned layout, or a non-contiguous plain input whose regroup cannot be expressed, has a None result layout. - Removing only unit axes from a plain layout MUST preserve the strides of the remaining axes, including when the input is a strided tensor window. - A bare, fully-Broadcast ShardLayout input (every attr Broadcast, no genuine sharding) carries that ShardLayout through Reshape when the input layout positions can express new_shape by the view rules below. When they cannot (including while either layout shape is symbolic), the result layout is None and no error is raised because no genuine ownership is discarded. This rule does not extend to a ComposedLayout whose outer layout is a ShardLayout; that input follows the generic composed-layout rule and may produce a None layout. - A non-genuinely-sharded UMAT input reshaped to () remains the early scalar case and produces a plain UMAT scalar. - A genuine ShardLayout input (at least one non-Broadcast attr) carries through Reshape when the reshape is expressible as a view over the input's layout positions (layout.layout.shape, shard §7.1.1): - every layout position lies entirely within one new axis — non-size-1 new axes are the product of a contiguous run of whole layout positions, in either merge direction; size-1 axes insert/drop freely and hold no sharding; a Split layout-axis reference remaps to its new layout position; Partial / Broadcast carry through unchanged (mesh-axis states, no layout axis); OR - a Split-bound layout position divides across a new-axis boundary at a point its bound mesh extent evenly divides: the outer (earlier) sub-factor itself further factors into (mesh_ext, Split-bound, local extent 1) and (sub-factor / mesh_ext, plain), and the inner residual becomes a plain (non-Split) layout position — every Split-bound layout dim keeps local extent 1 (shard §7.1.1). - Arbitrary rank-N regroup — a Split-bound position whose device-owned block spans a boundary deeper than one divide, or two or more Split-bound positions interacting across the same regroup — is not yet supported and MUST fail closed. - A reshape of a genuine ShardLayout not expressible by the above MUST fail closed rather than fabricate or discard a layout. An unexpressible bare, fully-Broadcast ShardLayout instead follows the None-layout rule above.

Cast

class Cast(Op):
    """Convert the element dtype; produces a Tensor.

    Attributes:
        x: input; source tensor.
        dtype: attribute; target element dtype.
    """

    x: Tensor
    dtype: DType

- constraints: - Identity in shape / storage / layout; only the element dtype changes to dtype. A ShardLayout input keeps its layout (the relation is the identity). - Cast is the conversion boundary for the low-precision dtypes (fp8e4m3 / f8e8m0 / f4e2m1, see types §3), accepted as either the input or the target dtype. - The evaluator supports a dtype in {f32, f16, bf16, fp8e4m3, f8e8m0, i32, i64, bool}; evaluating a Cast to a dtype outside this set (e.g. f4e2m1) raises an unsupported-dtype error.

IndexAdd / IndexCopy / IndexSelect

class IndexAdd(Op):
    """Return dst with src slices accumulated at index along dim."""
    dst: Tensor
    index: Tensor
    src: Tensor
    dim: int = 0

class IndexCopy(Op):
    """Return dst with src slices copied to index along dim."""
    dst: Tensor
    index: Tensor
    src: Tensor
    dim: int = 0

class IndexSelect(Op):
    """Select whole slices from x at index along dim."""
    x: Tensor
    index: Tensor
    dim: int = 0
    fill_value: float | None = None

These are pure value forms of torch's whole-slice indexing family (torch.index_select, Tensor.index_add_, Tensor.index_copy_). They do not mutate an input. torch.gather is the separate elementwise-indexing operation and is not an HIR op.

• constraints:
• Every dim accepts torch-style negative indexing and MUST be in range for the data tensor's rank.
• IndexSelect.index MUST be rank 1 with dtype i32 or i64. Its result has x's rank, dtype, and storage; shape[dim] becomes index.shape[0] and all other extents are unchanged.
• A plain Layout input to IndexSelect MUST produce a row-major layout over the result shape, independent of the source strides.
• When IndexSelect.fill_value is not None, indices below zero or at least x.shape[dim] MUST produce slices filled with that value without reading the source. With None, selection retains torch's bounds behavior.
• IndexSelect produces a natural contiguous internal Layout for a ShardLayout input. Broadcast and Partial states carry through; a Split on dim becomes Partial(sum), and a Split on another dim keeps its target. Multiple Splits including dim, or a composed shard layout, MUST fail closed.
• IndexSelect is a pure value operation, not a copy. To stage dim-0 selections into shared memory, schedule its inputs as tf.schedule((x, tf.empty(Tensor[..., "smem"]), index), op=T.copy_async(fill=f)). The indexed instruction has no TMA candidate. Unscheduled IndexSelect MUST fail lowering with this scheduling form in the diagnostic; copying its already computed result with CopyAsync is not supported.
• IndexAdd and IndexCopy require rank-1 index, equal dst/src dtype and rank, equal non-dim extents, and index.shape[0] == src.shape[dim]. Their result type is exactly dst's.
• IndexAdd.index accepts i32 or i64. Repeated indices accumulate. Its value rule uses torch's default alpha=1; scaling src is explicit HIR.
• IndexCopy.index accepts i64 only. Repeated indices have undefined output because torch's last copy is nondeterministic; consumers MUST NOT rely on an order.
• IndexAdd and IndexCopy reject a Partial on dst, index, or src. Complete Split and Broadcast layouts remain logical HIR types.
• IndexAdd and IndexCopy have type inference, cost, and evaluator semantics only. No HIR-to-TIR lowering is registered.
• IndexAdd charges one add and reads src, index, and the addressed dst slices before writing those slices. IndexCopy charges no arithmetic and does not read the overwritten dst slices. Neither cost scales with the full dst extent.

Empty

class Empty(Op):
    """Declare storage without initializing its contents."""
    type: TensorType
    buffers: int = 1

• tf.empty(type, buffers=1) produces exactly the complete TensorType supplied by the author. Its contents are undefined; it declares the write-only operand of a schedule.
• buffers MUST be a positive integer. Staging ownership and rotation follow schedule §5.
• Its access relation states the same result boundary as Zeros, with no inputs. Its cost has no traffic and no arithmetic. Lowering declares storage without filling it.

Zeros

class Zeros(Op):
    """Allocate a zero-initialised tensor of one complete type.

    Attributes:
        type: attribute; complete output TensorType.
    """

    type: TensorType

- constraints: - Type inference MUST return type without reconstructing any part of it. - The result is zero-initialised and its logical shape is exactly type.shape. - Evaluation MUST resolve every symbolic entry of type.shape from the function's dimension bindings before allocating the concrete tensor.

Ops whose result type is determined entirely by attributes, including Zeros and Arange, MUST accept one complete TensorType; shape, dtype, layout, and storage MUST NOT be split across independent attributes. Transforming ops whose result type is partly inherited from an input, including Reshard, Cast, and Transpose, MUST instead carry exactly the portions they change and inherit the rest from that input.

Reduce

class Reduce(Op):
    """Reduce ``x`` over the selected axes; produces a Tensor.

    Attributes:
        x: input; input tensor.
        axes: attribute; reduced logical axes.
        keepdim: attribute; whether reduced axes remain as size-1 axes.
        kind: attribute; mean, sum, abs_max, or max.
    """

    x: Tensor
    axes: tuple
    keepdim: bool = True
    kind: ReduceKind = ReduceKind.MEAN

- constraints: - The logical result shape follows numpy reduction rules. - A reduction over an axis of no elements MUST return that kind's identity rather than fail. For max that is the least value its result dtype can hold, which is not always -inf: an integer dtype cannot hold -inf, bool's least value is False, and a finite-only float has no infinity. abs_max is 0, its results being magnitudes. sum is 0. mean has no identity — there is nothing to divide by — so it MUST NOT invent one. - Storage is preserved. - Plain input layout passes through unchanged. - For ShardLayout input, every split layout position that belongs to a reduced tensor axis collapses to broadcast with size-1 stride-0 output. - Non-default-stride sharded input must carry explicit producer strides, or typeinfer rejects it. - Lowering emits TIR Reduce; runtime dispatch is derived from operands, not from an HIR dispatch field. - An x mesh axis carrying Partial(reduction) (a pending cross-device reduction, orthogonal to the reduced tensor axes) propagates only when kind commutes with reduction: SUM / MEAN (both linear over the reduced axes) commute with reduction="sum" only; MAX commutes with reduction="max" only (the same associative operator applied over the combined tensor-axis and mesh-axis index set); ABS_MAX (a nonlinear abs composed with max) does not commute with any reduction. Typeinfer rejects a non-commuting combination, naming the offending reduction and the fix (an explicit Reshard to Broadcast).

InsertSlice

class InsertSlice(Op):
    """Write ``update`` into a window of ``dst``; produces a Tensor.

    Attributes:
        dst: input; target tensor (value form returns a tensor anchored on this
            buffer at lowering time).
        update: input; tensor written into the window.
        offsets: input; per-axis window starts — a rank-0 integer scalar for a
            rank-1 dst, or a tuple of rank-0 integer scalars (literal or
            runtime), one per axis, for rank N.
    """

    dst: Tensor
    update: Tensor
    offsets: Scalar

- constraints: - update has the same rank and dtype as dst; the window on each axis is [offset_axis, offset_axis + update.shape[axis]). - A rank-1 dst accepts a bare rank-0 scalar offset; a rank-N dst requires an offset tuple whose length equals the rank. - A literal (compile-time) offset that places a negative or out-of-bounds window on an axis fails typeinfer, naming the axis; a runtime offset is checked at eval/runtime. - The value form writes update into a slice view of dst's existing buffer (a loop-carried dst reuses one buffer with no replacement allocation). - When dst's ShardLayout carries a Partial(reduction) mesh axis, update MUST carry the identical mesh and the identical per-mesh-axis ShardAttr state (update's own cute layout may still differ, since its tensor shape is the smaller write window) for the write to type — writing a differently-sharded (or unsharded) update into a still-partial dst position under one output type is unrepresentable; typeinfer rejects otherwise. - When dst is complete, an update carrying a Partial MUST be rejected; the write result cannot preserve that secondary value state. An explicit Reshard(update, Broadcast) completes it. - ComputeCostMetadata charges no traffic for the untouched dst, one window-sized update read, the scalar or structural tuple offsets read, and one window-sized result write.

CacheUpdate

class CacheUpdate(Op):
    """Write a window of ``new`` into ``cache``; produces a same-shape cache.

    Attributes:
        cache: input; cache that receives the write.
        cur_pos: input; i32 scalar where the write begins.
        s: input; i32 scalar number of positions to write.
        new: input; source positions, taken as ``new[:, :s]``.
    """

    cache: Tensor
    cur_pos: Tensor
    s: Tensor
    new: Tensor

- constraints: - cache and new MUST be rank-4 [B, len, kv_heads, head_dim] tensors with the same dtype and equal B, kv_heads, and head_dim; typeinfer rejects a mismatch. When both lengths are static, new.len MUST NOT exceed cache.len. - cur_pos and s MUST be i32 scalar tensors. A scalar is rank-0 or has only literal size-1 dimensions; typeinfer rejects another dtype or shape. - The write interval is runtime data, never a shape dimension. The result has cache's same static shape; no context-length DimVar grows with a write. - cur_pos >= 0, 1 <= s <= new.len, and cur_pos + s <= cache.len MUST be checked at eval/runtime, not typeinfer, because their operands are runtime values. - This is a pure value-form op. Lowering MAY realize the output in place on cache's buffer. - A cache carrying Partial(reduction) on a mesh axis requires new to carry the identical mesh and per-mesh-axis state; a complete cache rejects a new carrying Partial. Typeinfer rejects either mismatch. - The registered boundary relation states the write as a window: the result and the two scalar controls are exact identities, while the cache and new boundaries name the operands that place the window. How much each moves does not depend on where it lands -- s rows written, the rest of the cache kept. An s that is not written down MUST carry the range this Op's own contract gives it, at least one row and no more than the fewer of new.len and cache.len; a written-down s outside that range MUST be refused. The cost evaluator supplies the same traffic it always did.

TopK

class TopK(Op):
    """Select the top ``k`` elements on ``axis``; produces ``(values, indices)``.

    Attributes:
        x: input; source tensor.
        k: attribute; elements kept on the selected axis.
        axis: attribute; selected axis.
        largest: attribute; greatest vs smallest selection.
        sorted: attribute; ordered selection.
    """

    x: Tensor
    k: ShapeDim
    axis: int = -1
    largest: bool = True
    sorted: bool = True

- constraints: - The result is a (values, indices) tuple; both shrink the selected axis to length k; values keep x's dtype and indices are i64. - k is a ShapeDim (types §4): a static int, or a dynamic k derived from a context-length DimVar (e.g. dim_min(512, CTX_LEN // 4)) — a first-class value propagated as the selected axis's symbolic length in the output shape, not a pad+mask workaround. k MUST satisfy the ShapeDim contract (int / DimVar / dim-arithmetic Expr); any other value fails typeinfer. - k MUST be non-negative (checked whenever k is static) and MUST NOT exceed the selected-axis length (checked whenever the axis length is static and k is either static or a symbolic value whose statically-derivable upper bound — DimVar.hi, composed through DimMin/DimMax/DimAdd/DimMul/DimFloorDiv/DimMod — is known). A symbolic k against a symbolic axis length, or an upper bound that does not statically compose (e.g. through DimSub, or a DimFloorDiv/ DimMod with a symbolic divisor), is not checked at typeinfer — it fails open, same as the pre-existing static-only check this widens. - A symbolic k's DimVar(s) MUST be resolvable from x's own (input) shape at evaluation time: narrower than LoopRegion's ShapeDim fields (§1.2), which resolve against the enclosing Function's full parameter shapes — a k expression whose DimVar appears only in some other argument, never in x, is not resolvable at TopK's evaluation site. - The selected axis MUST NOT be Split-sharded by a ShardLayout; a split selected axis fails typeinfer. - A ShardLayout output preserves the non-selected sharding and any replication; only the selected axis's layout extent becomes k, so the layout keeps size parity with the result shape. - sorted returns the selected elements ordered by largest; otherwise the same selected set is returned in an unspecified order. - x MUST NOT carry a Partial(reduction) mesh axis: indices identifies which position wins, which cannot be recovered from a per-device partial value without a paired value+device-identity reduction that a plain Partial attr cannot express; typeinfer rejects any Partial input regardless of reduction or k.

ir/hir/nn/

Neural-network value Ops following torch semantics (torch.nn.functional).

MatMul / Conv2D / ReLU / Sigmoid / Tanh / SoftMax / LayerNorm

Consensus torch.nn.functional ops. - constraints: - A ShardLayout operand carrying Partial(reduction) propagates to the output only when the op provably commutes with reduction; typeinfer rejects otherwise, naming the offending operand and the fix (an explicit Reshard to Broadcast). - ReLU / Sigmoid / Tanh are monotone non-decreasing elementwise, so they commute with max / min but not sum. - MatMul is linear in one value input when the other value-carrying input is Broadcast / replicated. On each mesh axis, one Partial(sum) is therefore allowed; a double-Partial input or a non-sum reduction is rejected. - MatMul.a_layout is "MK" (the default) or "KM"; MatMul.b_layout is "KN" (the default) or "NK". These literals state the physical order of each operand's final two axes. The access relation maps them to logical (M, K) and (K, N) before deriving the output, contraction ownership, and Partial(sum) state. Its cost is 2 * numel(local_output) * local_K, where local_K is reconstructed from that same logical-axis mapping. - MatMul requires one common operand dtype. MatMul.out_dtype names the result dtype; it defaults to None, which means the operand dtype, and a dtype name is resolved at construction, so an unknown name fails there. out_dtype states where the result is rounded, not which dtype hardware accumulates in. The flops cost stays keyed by the operand dtype; the result's traffic uses the result dtype. Evaluation with a result dtype equal to the operand dtype and an operand that is not an 8-bit float is torch.matmul on the operands; otherwise it is torch.matmul on the operands widened to f32, then cast once to the result dtype. - Under a CUDA Target, a MatMul with an operand in smem or rmem takes its result storage from the MMA instructions the Target supports for it: each declaration's A, B and C are read in every configuration its patterns admit, and a configuration whose A and B admit the operand dtypes and storages and whose C admits the result dtype contributes C's storage. One storage is the result storage; several are umat, for the author to resolve; none MUST be refused, naming the Target, both operands and the result dtype. Operands in gmem or umat keep the target-neutral rule, and so does every other Target. - Where no Target is in reach (ctx.resolve_target() is None), a MatMul with an operand in smem or rmem has umat result storage: which instruction computes it, and so where it lands, is decided once a Target is known. A shared child is typed under the visitor's current ValuationContext. - Conv2D requires rank-4 NCHW input and OIHW weight, a rank-1 bias, and one common operand dtype. stride and dilation are positive length-2 tuples, padding is a non-negative length-2 tuple, and groups is positive. Input and output channels MUST be divisible by groups, the weight's input channel extent MUST equal input_channels / groups, and the bias extent MUST equal the output channels. - Conv2D states its grouped input / weight / bias / output relation over output positions and the input-channel and kernel contraction dimensions. Shared shard propagation derives a fresh result layout: input batch and compatible weight/bias output-channel splits survive. A contraction split of the input-channel axis MUST be aligned on input and weight at the same mesh axis; a single-sided channel split is not a representable local convolution. A contraction split produces Partial(sum) only when the bias carries the same per-mesh-axis Partial(sum) state, so completing the partial result adds the bias exactly once. A translated/halo access, incompatible contraction, mesh, or per-axis state that the relation cannot represent MUST fail naming the operand and requiring an explicit Reshard; no input layout is copied or real ownership discarded. - Conv2D costs 2 * numel(output) * weight.shape[1] * kh * kw flops after topology projection, where the weight extents are reconstructed on logical I/KH/KW axes from any factorized local layout. Traffic has exactly four slots: complete reads of input, weight, and bias followed by one complete result write. - SoftMax normalizes across an axis (a non-monotonic combination of every value on that axis), so no reduction provably commutes; typeinfer rejects every Partial input. - LayerNorm(axis=a) normalizes over the complete trailing shape x.shape[a:], after resolving a in [-rank, rank). Weight and bias shapes MUST each equal that suffix exactly rather than merely broadcast to it. x MUST use f32, f16, or bf16. Weight and bias dtypes MUST agree and either match x, or be f32 when x is f16/bf16; the result dtype is x's. - LayerNorm rejects every Partial operand and every Split on a logical normalized-suffix axis of x, weight, or bias. A Split on an x prefix axis remains on the same-shape result.

Gelu

class Gelu(Op):
    """Gaussian Error Linear Unit.

    Attributes:
        x: input; tensor the activation applies to elementwise.
        approximate: attribute; ``"tanh"`` selects the tanh-based
            approximation (HF ``gelu_pytorch_tanh`` / Gemma-2 MLP activation).
    """

    x: Tensor
    approximate: str = "tanh"

- constraints: - Elementwise: the output type, shape, and layout are x's. - x * Phi(x) dips below zero before rising back through it near zero, so GELU is not monotone and commutes with no reduction — unlike the ReLU / Sigmoid / Tanh group above, which commutes with max / min. typeinfer rejects any Partial operand with a Reshard remedy.

Silu

class Silu(Op):
    """Sigmoid Linear Unit — ``x * sigmoid(x)`` as one op.

    Attributes:
        x: input; tensor the activation applies to elementwise.
    """

    x: Tensor

- constraints: - Elementwise: the output type, shape, and layout are x's. - Fused rather than decomposed into Sigmoid + Binary(MUL): the fused form does not round the intermediate sigmoid(x) to x's dtype, so at reduced precision the two differ by up to ~1 ULP per element. - x * sigmoid(x) has a minimum near x = -1.278, so SiLU is not monotone and commutes with no reduction; typeinfer rejects any Partial operand with a Reshard remedy, as Gelu does.

RMSNorm

class RMSNorm(Op):
    """Normalize ``x`` by the root-mean-square of its last axis, scaled by ``weight``.

    Attributes:
        x: input; tensor normalized over its last axis.
        weight: input; rank-1 scale, same length as ``x``'s last axis.
        eps: attribute; added to the mean square before the root.
    """

    x: Tensor
    weight: Tensor
    eps: float = 1e-6

- constraints: - weight MUST be rank-1 with the same length as x's last axis; every other x axis, including a dynamic (DimVar / dim-arithmetic) entry, flows through unchanged. - The normalization reduces the whole last axis at once (every output element depends on that axis's full mean of squares), so the reduced axis MUST stay inside a single op instance: it is never an iteration-domain axis of its own, only an existential range that both the read and the write cover in full. - x / weight normalize across an axis (a non-monotonic combination of every value on that axis), so no mesh-axis reduction provably commutes; typeinfer rejects any Partial operand.

RoPE

class RoPE(Op):
    """Rotate query and key tensors using position-indexed cos/sin caches."""

    q: Tensor
    k: Tensor
    cos_cache: Tensor
    sin_cache: Tensor
    pos_ids: Tensor

- constraints: - q and k MUST have the same rank, at least 1, and the same even head_dim (last axis). Their other extents MAY differ, so grouped-query attention gives q more heads than k. - The rotation is the rotate-half form: the last axis splits in two halves and the pair (x[i], x[i + d/2]) rotates together. This is the unqualified HF convention (apply_rotary_pos_emb / rotate_half); the interleaved form (rotate_every_two, GPT-J / CodeGen) is a different Op, not an attribute of this one. - The result is (q_rope, k_rope) and each branch preserves the layout of its corresponding q or k input. - On each mesh axis, a branch MAY preserve one Partial(sum) on its corresponding query or key input only when cos_cache, sin_cache, and pos_ids are Broadcast / replicated on that axis. - A non-sum Partial, multiple value-carrying Partials, or a Partial on a secondary cache/index input MUST be rejected with a Reshard remedy.

ir/hir/sharding/

ShardLayout and Mesh are type-system constructs, not Expr inputs (shard §5).

MeshCoord

class MeshCoord(Op):
    """Return the current participant's coordinate along one mesh axis."""

    mesh: Mesh
    axis: Expr

• constraints:
• The result is an i64 scalar with layout=None and no placement.
• Type inference MUST reject a mesh that is not bound by the current MeshRegion.
• Evaluation MUST raise EvalError, because the evaluator models values without selecting one physical participant.

Reshard

class Reshard(Op):
    """Convert ``x`` to a target layout / storage; produces a Tensor.

    Attributes:
        x: input; input tensor.
        layout: attribute; optional target LayoutBase (Layout / ComposedLayout / ShardLayout).
        storage: attribute; optional target storage kind.
    """

    x: Tensor
    layout: LayoutBase = None
    storage: StorageKind = None

- constraints: - Omitting layout preserves x.layout; omitting storage preserves x.storage. - The output preserves the input logical TensorType.shape. - Supplied layout is a LayoutBase (Layout / ComposedLayout / ShardLayout). - Memory analysis and lowering MUST treat a Reshard as a view iff storage is unchanged, the shape is literal, and both structural address boxes agree. A box records each zero-stride Split's mesh axis, logical tensor axis, position within that axis, and extent. Its physical address layout is the coalesced layout after removing these independent modes; swizzle and offset are preserved. No element or participant enumeration is required. The result MUST reuse the source buffer, lower to a view, and require no instruction candidate. Broadcast and nonzero-stride Split ownership do not change a shared address, including a device ownership change with the same physical box. This alias decision does not estimate redistribution traffic. Zero-stride Split changes require a copy.

| Storage | Broadcast ↔ Split | Repartition | Transpose |
|---|---|---|---|
| gmem | view | view | copy |
| rmem with independent Split modes | copy | copy | copy |

• Every other Reshard MUST retain its independent buffer and instruction selection requirements.
• Destination storage is concrete, not unmaterialized.
• The single op covers zero-copy view, cross-storage copy, cross-CTA redistribute, and mixed cases; typeinfer and the recursive-local Cost Evaluator classify the call.

Stride resolution. Storage direction follows the physical addressability hierarchy rmem < smem < gmem (per-thread / per-CTA / per-program). Typeinfer dispatches on (layout, storage):

• layout=None, storage unchanged → x.type (no-op).
• layout=None, storage changed → error; a storage change MUST carry an explicit layout=.
• a layout with unspecified strides, storage unchanged → dest strides match the form already on x.layout: a Split-axes-zero source ⇒ per-instance form; otherwise ⇒ shared-engine C-order over the canonical global shape. When x.layout is None (plain kernel-param), fall back to shared-engine C-order.
• a layout with unspecified strides, low → high level → dest strides = C-order over layout.shape (shared-engine form).
• a layout with unspecified strides, high → low level → dest strides[k]=0 for every Split axis k; non-Split axes follow C-order over shard_layout_local_shape(layout) with size-1 → 0 (per-instance form).
• layout=Layout(strides=tuple) (verbose) → dest strides are taken verbatim; typeinfer MUST NOT rewrite them (e.g. SM80 MMA fragment layouts).

Cost classification. A Reshard is a zero-copy view with zero traffic only when no element moves under the alias rule above, including the layout=None no-op. Every other Reshard reports one full source read plus one full destination write, including a same-storage change of addresses or storage instances. Cross-unit traffic is recorded alongside storage traffic; an aliased Reshard has neither storage traffic nor cross-unit traffic.

Cross-CTA fence. The grid fence for a cross-CTA reshard is owned by the reshard lowering, not by a separately authored sync. When a reshard reads a gmem shard produced under a different CTA ownership (an ownership change across a cta mesh), the lowering MUST emit a grid barrier before the reshard so every CTA's prior shard writes are visible. The reshard lowering owns only the fence; cross-CTA data redistribution (all-to-all / gather across CTAs) is not part of this op.

Local

class Local(Op):
    """Take the current device's local view of a sharded tensor; produces a Tensor.

    Attributes:
        x: input; input tensor with ShardLayout.
    """

    x: Tensor

- constraints: - The result shape contracts each Split axis by that mesh axis's extent. - dtype and storage are preserved. - The result uses the corresponding local Layout, with its domain matching the contracted result shape. A noncanonical input MUST be rejected at the Local call with its source location. - Static split sizes divide by mesh extent; symbolic sizes pass through. - Local only names the current topology position's existing view; it materializes no value and moves no bytes.

2. Function specialization API

class SpecializationError(ValueError):
    """Report that a function cannot be specialized as requested."""


PROVENANCE = "_specialized_from"
BOUND_DIMS = "_specialized_dims"


def origin_of(function: object) -> Function | None:
    """Return the source function a rebuilt function came from.

    Args:
        function: Candidate derived function.

    Returns:
        Its recorded origin, or None.
    """
    ...


def bound_dims_of(function: object) -> tuple[tuple[str, int], ...] | None:
    """Return the sorted dimensions a rebuild chose, if it chose any.

    Args:
        function: Candidate derived function.

    Returns:
        Recorded bindings, or None.
    """
    ...


def variant_for(fn: Function, dims: Mapping[str, int]) -> Function:
    """Select the unique implementation covering concrete dimensions.

    Args:
        fn: Function or dispatch prototype.
        dims: Concrete extents by dimension name.

    Returns:
        The selected implementation.
    """
    ...


class InstantiateDimensionCloner(ExprCloner):
    """Clone a reachable program with dimension bindings, without type inference."""

    def __init__(self, dims: Mapping[str, int] | None): ...
    def module(self, module: Module) -> Module: ...
    def function(self, function: Function) -> Function: ...


def residual_dims(fn: Function, frame: ValuationContext | None = None) -> tuple[str, ...]:
    """Return every dimension still stated as a range.

    Args:
        fn: Function to inspect recursively.
        frame: Current module and effective execution context, required for instance targets.

    Returns:
        Residual dimension names.
    """
    ...


def dim_vars_reached(fn: Function, frame: ValuationContext | None = None) -> dict[str, object]:
    """Return residual dimension declarations by name.

    Args:
        fn: Function to inspect recursively.
        frame: Current module and effective execution context, required for instance targets.

    Returns:
        Residual dimension declarations.
    """
    ...

• constraints:
• variant_for MUST return an ordinary function unchanged and otherwise MUST select exactly one matching variant. Missing or ambiguous coverage and an unstated dimension used by a pattern MUST raise SpecializationError.
• InstantiateDimensionCloner MUST select variants by variant_for and substitute dimensions in types, Op attributes, loop bounds, and meshes. It MUST replace Function targets with clones deduplicated by identity; one shared definition MUST be cloned once. It MUST preserve ModuleList and TupleGetItem instance targets, except for rebuilding their computed index operand. It MUST NOT perform type inference while cloning.
• Each cloned function MUST retain its original name and record the chosen implementation and sorted bindings, recoverable through origin_of and bound_dims_of. Provenance and bound-dimension records MUST NOT participate in structural equality or hashing. A recorded origin MUST NOT be re-pointed after cloning.
• Dimension validation and check errors belong to analysis §2.
• residual_dims and dim_vars_reached MUST inspect the whole function graph, including signatures, bodies, Op attributes, loop bounds, variants, and called functions, including return types. Runtime values without a DimVar are rejected later by their consumer.
• Instance dimension scans MUST receive a ValuationContext and recurse through element entries using the current frame's children (core-ir §2.1). Function-only graphs retain the frame-free query API.