diff --git a/docs/spec/architecture.md b/docs/spec/architecture.md index e5a531a3..a2bd80c1 100644 --- a/docs/spec/architecture.md +++ b/docs/spec/architecture.md @@ -116,7 +116,7 @@ the layout / shard hierarchy diagram lives in [shard](./shard.md). The parser turns Python DSL source into a `core_ir.Module`. There are two layers: a **module layer** (`parse_module`, the sole top entry) that assembles the compilation unit, and a **function layer** -(`parse_func` / `parse_prim_func`) that turns each `ast.FunctionDef` +(`parse_function(fn, context)`) that turns each authored Python `FunctionType` into an `hir.Function` or `tir.PrimFunction`. The DSL surface (authoring namespace, OpSchema registry, dispatch tokens, AST subset, sugar disambiguation) and the lexical-env rules for diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 10b136af..2b2720fe 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -179,7 +179,7 @@ each get their own file. Codegen consumes TIR only. **Rule 5 — `/__init__.py` re-export rules:** real Op submodules are re-exported; aliases are imported only for registration side effects; user imports -go through [parser §2](./parser.md#2-dsl-namespace-surface). +go through [parser §2](./parser.md#2-syntax-and-rules). **Rule 6 — one pass = one file.** A pass class lives in `passes/transforms/.py`; internal visitors / mutators stay in that file. @@ -227,7 +227,7 @@ from tilefoundry.dsl import tf, T, Tensor sugar; it is owned by `tilefoundry.dsl` (defined under `tilefoundry.dsl._tensor`, re-exported as `tilefoundry.dsl.Tensor`). It is **not** the IR tensor type — the IR type carrier is - `tilefoundry.ir.types.TensorType`. See [parser §1.4](./parser.md#14-tensor-and-consttensor-annotations) for + `tilefoundry.ir.types.TensorType`. See [parser §2.1](./parser.md#21-syntax) for the annotation grammar. - `DType` is **not** re-exported. dtype values use string form in DSL source (`Tensor[(8,), "bf16"]`, `zeros((1, 64), "bf16", ...)`); diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index bb915c0a..11e9c95c 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -91,7 +91,7 @@ invocation. Owning a `Function`, attaching a child `Module`, or declaring a Python-to-HIR entry and HIR-to-HIR device calls are governed by [hir §1.1](./hir.md#11-function). -- `parse_module` (see [parser §1](./parser.md#1-dsl-syntax)) returns a `Module`. +- `parse_module` (see [parser §2](./parser.md#2-syntax-and-rules)) returns a `Module`. - A bare `@func` / `@prim_func` becomes an implicit single-function `Module` whose `entry` is set to that function. A function that declares execution context of its own is therefore already a `Module`. @@ -131,7 +131,7 @@ chain and is not copied onto each Module or Function. renamed by index) and one prototype serve any number of independent builds. - `methods` collects plain Python functions (orchestration methods, e.g. `forward` / `init_caches`; full collection rule in - [parser §2.7](./parser.md#27-module-authoring-surface)). A function name, + [parser §3](./parser.md#3-implementation-overview)). A function name, a child module name, and a method name MUST be disjoint at one `Module`'s own level — all three resolve through the same attribute surface ([§1.1](#11-function-access) below), so a name used by more than one would be ambiguous. @@ -305,7 +305,7 @@ class SourceSpanMetadata(IRMetadata): - `SourceSpanMetadata` records the parser source range before type inference. - `ExecutionDomainMetadata` records the `with Mesh(...)` scopes a `Call` was written inside, outermost first - ([parser §1.6](./parser.md#16-with-mesh-as-m)). `at(level)` returns the + ([parser §2.1](./parser.md#21-syntax)). `at(level)` returns the innermost scope naming *level*, or `None` when none does. It states where the work ran, which is not what the result's layout states -- a value laid out across threads may have been produced by work one CTA did -- so the two diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 07592990..aca2e59c 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -200,7 +200,7 @@ across sessions. `Tensor[..., (sugar)]` annotation on a parameter or return appears at the kernel boundary, where the underlying engine is a shared buffer handed across the FFI surface. When the surface sugar emits -`Layout(strides=None)` ([parser.md §1.5](./parser.md#15-layout-sugar)), +`Layout(strides=None)` ([parser.md §2.1](./parser.md#21-syntax)), function-signature binding MUST materialize it to **shared-engine C-order over the canonical global shape** before the resulting `TensorType` enters the body. Verbose `Layout(strides=tuple)` @@ -355,7 +355,7 @@ class GridRegionExpr(Expr): `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 §1.7](./parser.md#17-for-i-in-tile--for-i-in-range-hir-only)). `range` is not unrolled. `induction_var` ranges +[parser §2.1](./parser.md#21-syntax)). `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` (`tile(...)` and `range(stop)`); the @@ -369,7 +369,7 @@ already a coordinate in `range(0, extent, step)`, not an ordinal to multiply by - When `start` / `extent` / `step` are static `int`, the trip count is recoverable from the node alone, without the parser-side window binding - ([parser §1.7](./parser.md#17-for-i-in-tile--for-i-in-range-hir-only)). + ([parser §2.1](./parser.md#21-syntax)). - 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 @@ -404,7 +404,7 @@ parser scope. `GridRegionExpr.type` is `TensorType` (single carry) or `TupleType` (multi-carry); the value is the Expr itself, not a `Call`. Parser-side rules: see -[parser §5.1](./parser.md#51-gridregionexpr-carry-out-lifting). +[parser §3](./parser.md#3-implementation-overview). **Minimal example** — loop-carried accumulator: @@ -508,7 +508,7 @@ class Binary(Op): - 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 float scalar is given the other operand's float dtype by the authoring - surface, before it is an operand at all ([parser §1.9](./parser.md#19-compile-time-values)); a Python + surface, before it is an operand at all ([parser §2.1](./parser.md#21-syntax)); a Python integer is not. - The elementwise `min` / `max` kinds are also surfaced as `minimum` / `maximum`. - Equal plain layouts, or one plain layout paired with `layout=None`, pass diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index ff0b5916..e1a7d25c 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -135,7 +135,7 @@ resolves to — an inherited Target or topology hierarchy prints nothing, so the declaration-versus-inheritance split survives the round trip. A declared Target prints as the `@module(target=...)` argument and a declared hierarchy as the `@module(topologies=...)` argument -([parser §2.7](./parser.md#27-module-authoring-surface)). +([parser §3](./parser.md#3-implementation-overview)). Every dimension referenced only by a declared topology expression MUST still be emitted in the dimension prelude. A topology `ShapeDim` MUST use the same DSL expression text as tensor and Mesh geometry, including public constructors @@ -144,7 +144,7 @@ such as `ceildiv`, so importing restores the same expression tree. - constraints: - The decorator MUST print in its called form, `@module()` included. A bare decorator has not run while the class body is evaluated, so a body naming a - child call could not resolve it ([parser §1.1](./parser.md#11-decorators)). + child call could not resolve it ([parser §2.1](./parser.md#21-syntax)). - A nested Module MUST print before the owner's Functions, because a body calling one names the attribute it is bound to and a class body binds in the order it is written. @@ -154,7 +154,7 @@ such as `ceildiv`, so importing restores the same expression tree. from the attached entry's identity and the recorded origin ([hir §1.1](./hir.md#11-function)), not from the name they share and not from the parser's consumed authoring record - ([parser §4.2](./parser.md#42-closure-then-registry-callee-resolution)): + ([parser §3](./parser.md#3-implementation-overview)): anything may be called the same, and two attributes may hold copies of one Module. - Such a call MUST print exactly the arguments `Call.args` carries @@ -177,7 +177,7 @@ output supports two modes derived from the same pretty-print core: - `canonical` — round-trippable text used by `as_script()`, pass dumps, and viewer detail `code` blocks: the `Tensor[...]` form of - [parser §1.5](./parser.md#15-layout-sugar) (storage as the string + [parser §2.1](./parser.md#21-syntax) (storage as the string slot, `gmem` omitted). - `compact` — abbreviated, **display-only / non-round-trip** text for summaries / labels: `dtype[shape] {value-state?} @storage`. It inlines @@ -245,7 +245,7 @@ def small_sequence(x: Tensor[(S,), "f32"]) -> Tensor[(S,), "f32"]: The pattern prints in its constructor form (`DimVarRangePat("S", 1, 4)`; other `Pattern` subclasses fall back to `repr(pattern)`). The emitted binding -mirrors the authoring surface ([parser.md §1.1](./parser.md#11-decorators)); +mirrors the authoring surface ([parser.md §2.1](./parser.md#21-syntax)); when an IR variant has no display label, the printer synthesizes a valid binding from its canonical specialization signature. Because a dispatch prototype has a `DimVar` parameter, its rendering is a @@ -267,7 +267,7 @@ MUST agree over - DType annotations and op attributes preserve the selected descriptor singleton through their canonical names - Partial layouts preserve mesh names through the canonical - [parser §1.5](./parser.md#15-layout-sugar) value-state form, and preserve `Partial.reduction` plus the + [parser §2.1](./parser.md#21-syntax) value-state form, and preserve `Partial.reduction` plus the attrs-position mesh axis in the underlying IR **Display-only** — the rendering of a function with a `DimVar` parameter, and diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 5be88d5d..1b044d70 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -1,1203 +1,286 @@ -# TileFoundry Spec — Parser +# TileFoundry Spec - Parser -The parser turns Python source decorated with TileFoundry's IR decorators -into a `core_ir.Module`. This document covers the user-facing DSL -syntax ([§1](#1-dsl-syntax)), the DSL namespace surface ([§2](#2-dsl-namespace-surface)), the parser architecture -([§3](#3-parser-architecture)), the shared machinery both IRs use ([§4](#4-shared-parsing-machinery)), and the per-IR parser -bodies ([§5](#5-hir-parser) / [§6](#6-tir-parser)). Validation and rejection rules are in [§7](#7-validation-and-rejection). +The Parser accepts authored Python functions and produces HIR or TIR through one typed API. -## 1. DSL syntax +## 1. Public API -This section is the language reference for TileFoundry DSL source. -Each subsection introduces one construct using a grammar production -followed by a short description and an example. Productions use -the convention `lhs ::= rhs`; literal terminals are quoted; `*` is -zero-or-more, `?` is optional. - -### 1.1 Decorators - -``` -decorator-form ::= '@tilefoundry.module' - | '@tilefoundry.func' - | '@tilefoundry.prim_func' -``` - -`@tilefoundry.module(entry="")` decorates a class and evaluates to a -`core_ir.Module`: the decorated name binds to the Module itself -(not the class). It collects the class body's `@func` / `@prim_func` results, -child `Module`s, and plain Python functions, in definition order, into the -result's `functions` / `modules` / `methods` (full contract in -[§2.7](#27-module-authoring-surface)); `entry` is an optional argument naming -which collected function is the default step. A member MAY call a sibling -**defined above it** — -the call lowers to a `Call` targeting that sibling function; forward -references (a callee declared below the caller) are unresolved (see -[§3.3](#33-description)). Functions are reached by name on the result (see -[core-ir §1.1](./core-ir.md#11-function-access)). - -`@tilefoundry.func` and `@tilefoundry.prim_func` decorate functions and -evaluate to the parsed IR directly: `@func` to a `hir.Function` -([hir.md §1.1](./hir.md#11-function)), `@prim_func` to a -`tir.PrimFunction`. The decorated name binds to that IR node, not to -the original Python function. Passing a standalone function to `compile` / -`jit` lifts it into an implicit single-function `Module` whose `entry` is -that function. `@func` parses with dispatch token `"hir"`, `@prim_func` -with `"tir"`. - -A standalone `@func` MAY declare its own execution context — -`@func(target=..., topologies=(...))`. Declaring either makes that function -its own execution domain, so the decorated name binds to the implicit -single-function `Module` carrying the declaration rather than to the -`hir.Function`. A plain `@func` inside a `@module` class body never does -this: the class declares the domain, and the member stays a `hir.Function` -so it remains callable by its siblings and specializable through -`.specialize`. -A class-body `@func(topologies=(...))` instead binds a single-function child -`Module`: it declares its own hierarchy and inherits the owning Module's Target. -The printer MAY render the child as a nested `@module` class rather than -reproduce the member-function form. - -```python -# example -@tilefoundry.module(entry="f") -class M: - @tilefoundry.func - def g(...): ... - - @tilefoundry.func - def f(...): - return g(...) # sibling g is declared above → resolves to g's Function -``` - -**Specialization decorators.** A function specializes its body per input -shape through `Function.specialize`. The base function is defined with -`@tilefoundry.func`; each variant is added by decorating a `def` -with `@base.specialize(pattern)`: - -```python -# example -S = DimVar("S", 1, 9) # envelope [1, 9) = 1..8 - -@tilefoundry.func -def f(x: Tensor[(S,), "f32"]) -> Tensor[(S,), "f32"]: - pass # prototype base - -@f.specialize(DimVarRangePat("S", 1, 5)) -def small_s(x: Tensor[(S,), "f32"]) -> Tensor[(S,), "f32"]: - return small_impl(x) # variant [1, 5) = 1..4 - -@f.specialize(DimVarRangePat("S", 5, 9)) -def _(x: Tensor[(S,), "f32"]) -> Tensor[(S,), "f32"]: - return large_impl(x) # variant [5, 9), unlabelled -``` - -- `@tilefoundry.func` evaluates to the base `hir.Function`. `func()` has no - `specializations=` parameter; specialization is reachable only through - `.specialize`. -- The prototype base body is `pass`: it declares the signature and dispatch - envelope only and parses to `Function.body is None`. The implementations - live in the variants ([hir.md §1.1](./hir.md#11-function)). A - `pass` body is legal only for a function that receives variants; a `pass` - body with no variants, or a real body combined with variants, is rejected. -- `base.specialize(pattern)` returns a decorator. It parses the decorated - `def` into a variant `hir.Function` (same `name` as the base, - `specializations=(pattern,)`), registers it on `base.variants`, and - returns the variant. -- The decorated identifier is the variant's **display label** and nothing more. - Variant naming, binding uniqueness, and lookup handles follow the contract in - [hir §1.1](./hir.md#11-function). The label MUST NOT become the variant's - `name`, and MUST NOT take part in equality, hashing, dispatch, or TIR - identity: the variants of one base share that base's name, and which one runs - is decided by the pattern alone. -- `pattern` MUST be a single `DimVarRangePat` (see - [core-ir §3](./core-ir.md#3-pattern)); other `Pattern` subclasses are rejected for - v0. The referenced `DimVar` and its `(lo, hi)` envelope live on the base - parameter's shape. Each variant's range MUST fall within that envelope, - and the full variant set MUST partition it — disjoint and complete (see - [hir.md §1.1](./hir.md#11-function)). Two variants with the - same canonical signature are rejected. -- A `DimVar` shape entry MAY be written **inline** (`DimVar("S", lo, hi)` - AST node in the shape tuple) or as a **named alias** (`S = DimVar("S", - lo, hi)` then `Tensor[(S,), ...]`). Both resolve to the same `DimVar` - instance (type-level cache keyed by `(name, lo, hi)`). A second `DimVar` - with the same `name` but conflicting `(lo, hi)` is a hard parse-time - error. -- Variants accumulate on the base only during authoring. Once the base - enters a `Module` (see [core-ir §1](./core-ir.md#1-module)) it is sealed and - `.specialize` raises. A variant lives only inside `base.variants`; it is - never a separate `Module.functions` entry. - -### 1.2 DSL namespace - -``` -import-form ::= 'from tilefoundry.dsl.tf import *' - | 'from tilefoundry.dsl.T import *' - | 'from tilefoundry.dsl import' ('tf' | 'T') - -namespace-callee ::= ('tf' | 'T') '.' op-name -``` - -`tilefoundry.dsl.tf` (HIR) and `tilefoundry.dsl.T` (TIR) are the only -DSL-facing entries to the Op catalogue. The mechanism that backs -this surface is described in [§2](#2-dsl-namespace-surface). - -### 1.3 Op call - -``` -op-call ::= callee '(' arg-list ')' -callee ::= op-name ; bare-name path - | namespace-callee ; namespace-attribute path -op-name ::= identifier - | identifier '_' ; trailing-underscore = effect form (TIR only) -``` - -A bare-name callee MUST resolve to an `_op_schema`-bearing surface -value (an `Op` class for real-Op schemas, or an alias builder -function carrying `_op_schema` for surface-alias schemas) through -the function's closure (typically established by an `import-form`) or, -failing that, through `dispatch.resolve_callable` -([§4.2](#42-closure-then-registry-callee-resolution)/[§4.3](#43-opschema-and-overload-resolution)) — the -path a trailing-underscore `op-name` always takes, since nothing binds -a literal `foo_` name in the closure. The namespace-callee form -resolves on the namespace package directly ([§2](#2-dsl-namespace-surface)). When `op-name` is -registered with multiple schemas ("overloads"), the parser uses -first-match dispatch ([§4.3](#43-opschema-and-overload-resolution)). The trailing-underscore selector is -gated to the TIR token; using it in HIR is a verify error. - -### 1.4 `Tensor[...]` and `ConstTensor[...]` annotations - -`Tensor` and `ConstTensor` are parser-owned **DSL authoring type sugar**, -imported from `tilefoundry.dsl` (`from tilefoundry.dsl import Tensor`). It -is distinct from the IR-level `tilefoundry.ir.types.TensorType` (the -runtime carrier on `Expr.type`). The parser resolves a `Tensor[...]` literal -into a `TensorType` both in parameter/return annotations and in expression -positions that bind a `TensorType` operation attribute. - -``` -tensor-annot ::= ('Tensor' | 'ConstTensor') '[' shape ',' dtype (',' layout)? (',' storage)? ']' -shape ::= '(' (dim (',' dim)*)? ')' ; '()' is rank-0 -dim ::= integer-literal | dim-Expr ; dim-Expr per types §4 -dtype ::= '"f32"' | '"f16"' | '"bf16"' | … ; see types §3 -layout ::= layout-sugar ; see §1.5 - | 'ShardLayout(' … ')' ; verbose, see shard §7 -storage ::= '"host"' | '"gmem"' | '"smem"' | '"rmem"' | '"tmem"' | '"umat"' - | 'host' | 'gmem' | 'smem' | 'rmem' | 'tmem' | 'umat' -``` - -`Tensor[...]` and `ConstTensor[...]` resolve to the same ordinary `TensorType`; -the latter sets `Var.is_const=True` on a function parameter. `is_const` marks -external residency semantics only and does not embed a payload. `Tensor[...]` is -the carrier of optional layout sugar; it does not -own the sugar (which lives at [§1.5](#15-layout-sugar)). A rank-0 (scalar) tensor is -written `Tensor[(), "bf16"]`; the form `Tensor["bf16"]` (without -shape) is rejected. +`@module` executes its Python class body and finalizes the collected Function declarations, +child Modules, and ordinary methods. Module authoring is two-phase: class execution records +Function, specialization, and converter declarations; finalization attaches all child Modules +and puts them in parser scope before parsing Functions in source order. `@func` produces an HIR +Function; `@prim_func` produces a TIR PrimFunction. +`specialize` and `converter` register variants and weight converters on an existing HIR Function. ```python -Tensor[(4096, 2048), "bf16"] -Tensor[(4096, 2048), "bf16", (4096 @ gpu.cta, 2048)] -Tensor[(4096, 2048), "bf16", (4096 @ gpu.cta, 2048), "smem"] -Tensor[(4096 @ gpu.cta, 2048), "bf16", "smem"] # placement in shape -Tensor[(), "bf16"] # scalar -``` - -- constraints: - - The dtype slot MUST use a canonical quoted `DType.name` from - [types §3](./types.md#3-dtype). - - The parser MUST normalize that string to the corresponding process-lifetime - descriptor before constructing `TensorType`. - - An unknown dtype string MUST be rejected; it MUST NOT fall back to another - descriptor. - - Bare storage names MUST resolve to the six constants exported by - `tilefoundry.dsl.storage`; they are equivalent to the quoted spellings. - -### 1.5 Layout sugar - -``` -layout-sugar ::= axis-tuple ; implicit strides, no value-state - | '(' axis-tuple ',' stride-tuple ')' ; explicit strides, no value-state - | '(' axis-tuple ',' value-state ')' ; implicit strides + value-state - | '(' axis-tuple ',' stride-tuple ',' value-state ')' ; explicit strides + value-state -axis-tuple ::= '(' axis-spec (',' axis-spec)* ','? ')' -axis-spec ::= axis-extent ; a layout dim, not split (axis placement only) - | shape-extent '@' mesh-axis ; Split(axis_index) on mesh-axis - | shape-extent '@' '(' mesh-axis (',' mesh-axis)* ')' ; sequential decomposition -axis-extent ::= dim-expr -shape-extent ::= dim-expr -dim-expr ::= dim-atom | dim-expr dim-op dim-atom -dim-atom ::= integer-literal | dim-ref | dim-call | '(' dim-expr ')' -dim-ref ::= identifier ; closure-resolved ShapeDim -dim-call ::= identifier '(' dim-expr (',' dim-expr)* ')' -dim-op ::= '+' | '-' | '*' | '//' | '%' ; bool rejected -stride-tuple ::= '(' integer-literal (',' integer-literal)* ','? ')' -value-state ::= '{' partial-spec (',' partial-spec)* ','? '}' ; a set; only the last outer item -partial-spec ::= mesh-axis '@' 'P(' '"' reduction '"' ')' ; Partial(reduction) on mesh-axis -mesh-axis ::= identifier '.' identifier ; e.g. gpu.cta -reduction ::= 'sum' | 'max' | 'min' | … -``` - -A layout or mesh extent is a `ShapeDim`: a static integer or a -closure-resolved `DimVar` / dim expression. A bare axis is `Broadcast` -(carries no mesh binding). A symbolic split is canonicalizable when the split -extent and mesh-axis extent are the same expression, which gives local extent -one without symbolic division. Other symbolic split combinations MUST be -rejected with a diagnostic that asks the author to bind dimensions first. -The restricted static evaluator MUST keep the five `dim-op` forms as canonical -dimension arithmetic when either operand is symbolic, including when a -`dim-call` such as `ceildiv(...)` produced that operand. - -Implicit C-order strides retain symbolic products. HIR type inference may also -retain an unconsumed symbolic local extent in a per-instance stride. Lowering -and code generation keep the stricter boundary: all local extents MUST be -concrete after dimension binding before storage is materialized. - -The `axis-tuple` carries only **axis placement** (`Split` inlined as -`size @ mesh.axis`; a bare `size` is a non-split layout dim). The optional -`{...}` **value-state** set carries the mesh-axis `Partial` states -(`mesh.axis @ P("reduction")`). It is a Python `set` literal recognized at -the AST level (its element order carries no meaning) and MUST be the **last -item of the outer tuple**; it is never mixed into the `axis-tuple`. A mesh -axis named in no `Split` and no `Partial` is `Broadcast` (the default) — -Broadcast is never written. There is no `_ @ B(...)` / `_ @ P(...)` form. - -Outer-tuple discrimination: a bare `axis-tuple` is implicit-strides with no -value-state; an outer length-2 tuple whose second item is a `stride-tuple` -is explicit strides; length-2 whose second item is a `value-state` set is -implicit-strides + value-state; length-3 `(axis-tuple, stride-tuple, -value-state)` is explicit strides + value-state. - -`dim @ (m.a, m.b)` expands to one Split axis per mesh axis -(each with extent = mesh extent), followed by a bare remainder axis -of size `dim / ∏(mesh_extents)`. The remainder axis is always -appended last; the mesh-axis order in the tuple determines the tensor -axis order. - -**Several meshes in one layout.** A value may be distributed at more than one -level at once -- a CTA owns a tile and a lane owns part of that tile -- and -saying so names axes of more than one Mesh. Those meshes MUST be the enclosing -`with Mesh(...)` scopes, and they compose outermost first -([shard §5](./shard.md#5-mesh)) into the one Mesh the resulting `ShardLayout` -carries, so this and one Mesh naming both levels produce the same IR. Which is -inside which MUST come from the scopes rather than from the layout: a layout -naming a mesh that is not a scope it is written inside is refused, because there -is nothing that says how the two nest. - -**Canonicalization (single-mesh-axis form)**. Surface sugar -`N @ m.a` where `N > mesh_extent(a)` MUST be expanded at parse time -into the factorised pair `(mesh_extent(a) @ m.a, N // mesh_extent(a))` -before the `ShardLayout` is constructed. The first element becomes -a `Split` axis with `local_shape = 1`; the second becomes a bare -residual axis (non-`Split`). `N // mesh_extent(a)` MUST divide `N` -exactly; otherwise the sugar is rejected. The factorisation is -opaque to the user: input `N @ m.a` and input -`(mesh_extent(a) @ m.a, N // mesh_extent(a))` produce the same IR. - -#### Stride materialization (parser surface) - -The first sugar form -(`'(' axis-spec ... ')'`) emits `Layout(shape=canonical, -strides=None)` — the layout strides are deferred to `Reshard` -typeinfer, which fills them in based on the storage-level direction -(see [hir.md §1.3](./hir.md#13-op)). The -verbose form (`'(' axis-tuple ',' stride-tuple ')'`) emits a -concrete `strides` tuple; typeinfer respects it verbatim. The -parser does NOT inspect `storage` or do any physical-materialization -logic — that responsibility lives entirely in `Reshard` typeinfer. - -Spec: [shard.md §7.1.1](./shard.md#711-layoutshape), -[hir.md §1.3](./hir.md#13-op). - -Layout sugar is accepted **anywhere the expected surface value is -a `ShardLayout`**. Dispatch is annotation-driven (see [§4.4](#44-annotation-driven-sugar-dispatch)): the -parser consults the position's expected `ParamDef.annotation` (or -the `Tensor[...]` layout slot) to decide whether to invoke the -sugar parser. Omitted mesh axes default to `Broadcast`. Sugar -forms that would lose mesh / layout information fall through to -the verbose `ShardLayout(...)` constructor. - -```python -# Tensor[...] annotation slot -Tensor[(4096, 2048), "bf16", (4096 @ gpu.cta, 2048)] - -# value-state set: a Partial on a mesh axis (implicit strides) -Tensor[(4, 64), "f32", ((4 @ trd.l, 64), {trd.t @ P("sum")}), "smem"] - -# Op attribute slot whose ParamDef.annotation is ShardLayout -reshard(x, layout=((2048 @ gpu.cta, 64), {gpu.warp @ P("sum")})) -``` - -### 1.6 `with Mesh(...) as m` - -``` -with-mesh ::= 'with' 'Mesh' '(' mesh-args ')' 'as' identifier ':' suite -``` - -The `with Mesh(...) as m` grammar is shared by both dialects; the -binding name `m` is visible only inside `suite`. The two dialects differ -in what it lowers to: - -- **HIR** treats it as an **active mesh context** — a parser-lexical - alias for the constructed `Mesh`, so layout sugar ([§1.5](#15-layout-sugar)) may bind axes - with `… @ m.axis` and tensors authored under it reuse the one `Mesh`. In a layout - position, `m.axis` names the static mesh axis used by `dim @ m.axis`. In an HIR - Expr position, the same name is the current rank-0 coordinate along that axis; the - parser synthesizes the invariant index vector and its local view. Such a coordinate - may index an unplaced tensor in a runtime slice; indexing an already placed tensor - is rejected because its data-dependent mesh ownership is unresolved. - It is not a tensor-binding scope and emits **no IR node**, but every `Call` - built inside `suite` records the enclosing stack of these scopes, outermost - first, as `ExecutionDomainMetadata` - ([core-ir §2](./core-ir.md#2-expr)): where an occurrence was written is what - says which participants run it, and nothing else on the Call says it. Ordinary - values assigned inside `suite` follow normal function-body visibility - (not confined); `return` inside `suite` returns from the enclosing - `@func` (no mesh-region result), and a `@func` MUST NOT be defined - inside `suite`. A tensor's mesh/layout lives on its - `TensorType.layout`, not on the block it is written in; `reshard` is - the explicit boundary, and op typeinfer - ([hir §1.3](./hir.md#13-op)) decides whether values - combine. -- **TIR** lowers it to an explicit `MeshScope` Stmt ([§6](#6-tir-parser)) carrying the - `Mesh` and the binding `Var`. - -### 1.7 `for i in tile(...)` / `for i in range(...)` (HIR-only) - -``` -for-loop ::= 'for' identifier 'in' 'tile' '(' tile-args ')' ':' suite - | 'for' identifier 'in' 'range' '(' range-args ')' ':' suite -tile-args ::= extent-Expr ',' step-Expr -range-args ::= stop-Expr - | start-Expr ',' stop-Expr - | start-Expr ',' stop-Expr ',' step-Expr -``` - -Loop arguments are positional-only at the IR authoring surface; keyword -arguments are rejected. - -`tile(...)` and `range(...)` share **one** loop domain `(start, extent, -step)` and lower to the **same** `GridRegionExpr` ([hir §1.2](./hir.md#12-gridregionexpr)) — -`range` is not a separate construct and is **not** statically unrolled. The -only difference is the loop-variable binding: - -- `range(...)` binds `i` to a **scalar** induction var (`i: i64`); use it as - `x[i]` or write the window manually (`i : i + step`). Args follow Python - `range`: `range(stop)` (start `0`, step `1`), `range(start, stop)` (step - `1`), `range(start, stop, step)`. -- `tile(extent, step)` binds `i` to the standard Python - `slice(iv, iv + step, 1)` so `x[:, i]` lifts to a `Slice` over - the current window. The grid domain already advances `iv` by `step`; the - binding MUST NOT multiply it again. In any other Expr position, including an - `insert_slice` offset tuple, `i` resolves to the scalar `slice.start`. The - Python `slice` is parser-only and does not reach IR; the grid-domain `start` - is `0`. A single-argument `tile(extent)` is rejected; use `range(extent)` for - scalar iteration. - -`start-Expr` / `extent-Expr` (the **stop** endpoint, not a length — the -domain is half-open `[start, extent)`) / `step-Expr` MAY be any `ShapeDim` -([types §4](./types.md#4-dim--symbolic-shape-dimensions)), including a dim expression such as `C // N`; the -value is carried verbatim into `GridRegionExpr.start` / `.extent` / `.step` -and resolved at evaluate time ([hir §1.2](./hir.md#12-gridregionexpr)). - -A tensor subscript `x[slice0, …]` inside a loop body lifts to a -`hir.tensor.Slice` Op call. - -A subscript axis MAY **move** a tile window by a compile-time integer: `x[:, i + C]` -reads `[lo + C, lo + C + step)`, the window `i` alone reads translated by `C`. The -base is compile-time either way, so the moved start is a dim expression over `iv` -and the axis keeps the static extent `i` alone gives it. Offsets accumulate, so -`i + A + B` and `A + B + i` name one move by one sum, and each term is a -compile-time integer on its own. `C - i` MUST be refused: it reverses the window -rather than moving it. The loop domain, the window length, the offset and the axis -extent are all compile-time, so a move whose first window would start before the -axis, or whose last window would end past a static extent, MUST be refused where it -is written — unlike an unmoved window's own tail, which the axis extent catches at -evaluate time. In an Expr position `i + C` is ordinary scalar arithmetic over -`slice.start` and carries none of this. - -An `ast.Assign` whose single Name target is -bound in *outer* scope is a loop-carried rebinding (see [§5](#5-hir-parser) for the carry-out -lift). A **nested** `for ... in tile/range(...)` is allowed and lifts to a -nested `GridRegionExpr`; the carry scan recurses into nested loops, so an -outer-scope name rebound only inside a nested loop is still carried across the -outer loop (and the nested loop carries it too). - -### 1.8 Hard schedule constraints - -``` -constraint-annotation ::= 'where' '(' constraint-field (',' constraint-field)* ')' -constraint-field ::= 'layout' '=' layout-constraint - | 'mesh' '=' mesh-expression - | 'storage' '=' storage-expression -layout-constraint ::= layout-axis-tuple - | '(' layout-axis-tuple ',' binding-set ')' -binding-set ::= '{' binding (',' binding)* '}' -binding ::= topology '@' 'B()' - | topology '@' 'P(' string-literal ')' -``` - -`where(...)` is keyword-only and non-empty. A layout axis is `_`, an integer -or symbolic extent, or `extent @ topology`. The split form binds an existing -`Split` attribute to the physical layout position. `_` is a private -constraint wildcard and is never stored as an entry in `Layout.shape`. `B()` -and `P(...)` reuse the existing `Broadcast` and `Partial` -`ShardAttr` values; a topology may be bound at most once in one layout -constraint. `mesh` resolves to a `Mesh`, and `storage` resolves through the -current storage-kind registry. CTA capability checks do not occur in this -parser surface. - -### 1.9 Compile-time values - -A **compile-time value** is a Python number the parser can reach without building -any IR: a numeric literal, a name captured from the enclosing scope, an attribute -of a captured object, and arithmetic over those (`+ - * / // % **`, unary `-`). - -``` -compile-time-expr ::= number-literal | identifier | compile-time-expr '.' identifier - | compile-time-expr binary-arith-op compile-time-expr - | '-' compile-time-expr -``` - -- A compile-time value MAY appear anywhere a number is required: an attribute - argument, a shape or extent, a subscript index, or an op input, where it - becomes a rank-0 unmaterialized `Constant`. -- A body-local assignment whose right-hand side is a compile-time value binds - **the value**, not an `Expr`; the name is then usable in every position above. - A tuple target binds one number per name when every element is a compile-time - number. -- An op input that is a **Python float** carries no precision of its own: the - parser MUST give it the float dtype of the operands it is used with, and MUST - reject the call when those name more than one float dtype. A Python **integer** - keeps its own dtype, so combining one with a float tensor MUST still be - rejected ([hir §1.3](./hir.md#13-op)). -- Evaluation MUST NOT call anything reached from a speculative position: a value - that is not statically reachable is parsed as IR instead. -- Dimension arithmetic has one canonical spelling after it enters IR, including - shape annotations and op attributes. Slice endpoints that contain the same - runtime scalar value MAY cancel to a static window size; an unrelated runtime - endpoint remains invalid under the ordinary `ShapeDim` rule. - -A **compile-time list** holds `Expr` elements and never reaches the IR: - -``` -compile-time-list ::= '[' expr (',' expr)* ']' - | '[' expr 'for' identifier 'in' compile-time-sequence ']' -``` - -- A comprehension MUST declare exactly one `for` clause, no `if` guard, and a - plain-name target. Its sequence MUST be either the builtin `range(...)` over - compile-time integers or a compile-time tuple / list; any other call MUST be - rejected rather than evaluated. -- Subscripting the bound name with a compile-time integer selects one element. - Python's negative indexing applies. -- The comprehension form is the fixed-length unrolled spelling; it does not - affect `for` ([§1.7](#17-for-i-in-tile--for-i-in-range-hir-only)), which always builds a `GridRegionExpr`. - -A tensor subscript resolves per axis: an `ast.Slice` keeps the axis, a -compile-time integer **drops** it (as in torch), and a negative integer counts -back from the axis extent, which MUST therefore be static. - -For a tensor slice, a run-time rank-0 integer is permitted as an endpoint only -when the resulting window size is a compile-time dimension. The canonical -spelling is `start:start + K`; the parser MUST reject an unrelated stop -endpoint because `Slice.sizes` is a static attribute. The slice stride remains -compile-time. An endpoint MAY also be a compile-time dimension, including the -axis's own symbolic extent; its window size is symbolic accordingly. A tile -window keeps its own length and MAY be **moved** by a compile-time offset instead -([§1.7](#17-for-i-in-tile--for-i-in-range-hir-only)). - -## 2. DSL namespace surface - -### 2.1 Model - -The two namespaces are real Python modules; resolution is module -`__getattr__` over the OpSchema registry. There is no -`DslNamespace` class. - -```python -# tilefoundry/ir/core/op_registry.py -def _register_schema(schema: OpSchema, *, prepend: bool = False) -> None: ... -def get_schemas(dialect: str, name: str) -> list[OpSchema]: ... -def iter_schema_names(dialect: str) -> Iterable[str]: ... - -# tilefoundry/ir/core/op_schema.py — a frozen dataclass -class OpSchema: - """Describe one registered callable schema. - - Attributes: - name: attribute; Canonical callable name. - dialect: attribute; Surface dialect, `tf` or `T`. - category: attribute; Organizational group. - signature: attribute; Parameters in declaration order. - builder: attribute; Callable that constructs the IR operation. - op_class: attribute; Registered Op class, or None for an alias schema. - """ - - name: str - dialect: str - category: str - signature: tuple[ParamDef, ...] - builder: Callable[..., Any] - op_class: type | None = None - -# tilefoundry/dsl/tf/__init__.py (TIR symmetric in tilefoundry/dsl/T) -def __getattr__(name: str) -> type[Op] | Callable: ... -def __dir__() -> list[str]: ... -``` - -### 2.2 Class diagram +def parse_function( + fn: FunctionType, context: FuncParserContext +) -> hir.Function | tir.PrimFunction: ... +``` + +`FuncParserContext` carries the dialect, Function role, closure, topology scope, target, and +optional base/key for one parse. `FunctionRole` is `ROOT`, `VARIANT`, or `CONVERTER`. +`ParseError` is the single authored-source diagnostic type and includes source location and +recursive parse situation. These are the only public parser symbols. + +## 2. Syntax and Rules + +### 2.1 Syntax + + +```ebnf +; root: function +; literal: Python ast.Constant syntax, e.g. 1, "bf16", or None +; name: Python variable name; primary: name/attribute base for calls and subscripts +; expression: Python syntax composed from literals, names, primaries, and operators +; runtime-expression: expression lowered to a TileFoundry IR Expr +mesh-axis ::= identifier + | identifier '.' identifier +dim-expr ::= integer-literal + | identifier + | primary '.' identifier + | dim-expr ('+' | '-' | '*' | '//' | '%') dim-expr + | (identifier | primary '.' identifier) '(' (dim-expr (',' dim-expr)*)? + ')' +placed-layout ::= '(' ((expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | mesh-axis) | + dim-expr) (',' (expression '@' ('(' mesh-axis (',' mesh-axis)* ')' | + mesh-axis) | dim-expr))*)? ')' +shape ::= '(' (dim-expr (',' dim-expr)*)? ')' + | identifier + | primary '.' identifier +tensor-shape-layout ::= placed-layout + | shape +dtype ::= string-literal + | primary +literal ::= None + | Ellipsis + | boolean-literal + | integer-literal + | float-literal + | complex-literal + | string-literal + | bytes-literal +primary ::= identifier + | primary '.' identifier +sequence ::= '(' (expression (',' expression)*)? ')' + | '[' (expression (',' expression)*)? ']' + | '{' (expression (',' expression)*)? '}' +dict ::= '{' (expression ':' expression (',' expression ':' expression)*)? '}' +binary-operation ::= expression ('+' | '-' | '*' | '/' | '//' | '%' | '**') expression +unary-operation ::= ('+' | '-' | 'not') expression +slice ::= (expression)? ':' (expression)? (':' expression)? +subscript ::= expression '[' expression ']' +expression ::= literal + | primary + | sequence + | dict + | binary-operation + | unary-operation + | call + | slice + | subscript +call ::= expression '(' ((expression | keyword-name '=' expression) (',' + (expression | keyword-name '=' expression))*)? ')' +explicit-layout ::= '(' (tensor-shape-layout | shape) ',' shape ')' +plain-layout ::= '(' (dim-expr (',' dim-expr)*)? ')' +layout ::= None + | primary + | call + | explicit-layout + | placed-layout + | plain-layout +storage ::= string-literal + | primary +tensor-optional-slot ::= layout + | storage +tensor ::= tensor-head '[' '(' (tensor-shape-layout ',' dtype | tensor-shape-layout + ',' dtype ',' tensor-optional-slot | tensor-shape-layout ',' dtype ',' + tensor-optional-slot ',' tensor-optional-slot) ')' ']' +scalar-type ::= primary +type-annotation ::= tensor + | scalar-type +signature ::= (name ':' type-annotation (',' name ':' type-annotation)*)? +return-type ::= type-annotation +loop-carry-statement ::= expression '=' expression + | 'for' name 'in' expression ':' loop-carry + | statement +loop-carry ::= (loop-carry-statement (newline loop-carry-statement)*)? +loop-header ::= 'for' identifier 'in' ('tile' | 'range') '(' expression (',' expression)* + ')' ':' loop-carry +loop-body ::= (statement (newline statement)*)? +for ::= 'for' name 'in' expression ':' loop-body +mesh-context ::= ('Mesh' | primary '.' identifier) '(' (expression | ('layout' | 'names') + '=' expression) (',' (expression | ('layout' | 'names') '=' + expression))* ')' + | expression +with ::= 'with' mesh-context ('as' identifier)? ':' block +op-call ::= primary '(' ((expression | keyword-name '=' expression) (',' (expression | + keyword-name '=' expression))*)? ')' +launch ::= callee '(' ')' +slice-endpoint-binary ::= index-endpoint dim-op index-endpoint +mesh-coordinate ::= identifier '.' identifier +index-endpoint ::= literal + | primary + | slice-endpoint-binary + | mesh-coordinate + | runtime-expression + | expression +index-slice ::= (index-endpoint)? ':' (index-endpoint)? (':' index-endpoint)? +subscript-index ::= '(' ((index-slice | index-endpoint) (',' (index-slice | + index-endpoint))*)? ')' + | index-slice + | index-endpoint +subscript-expression ::= runtime-expression '[' subscript-index ']' +binary-expression ::= runtime-expression binary-op runtime-expression + | runtime-expression comparison-op runtime-expression + | runtime-expression boolean-op runtime-expression +unary-expression ::= unary-op runtime-expression +name ::= identifier +constant ::= boolean-literal + | integer-literal + | float-literal +tuple-expression ::= '(' (runtime-expression (',' runtime-expression)*)? ')' +runtime-expression ::= op-call + | launch + | subscript-expression + | binary-expression + | unary-expression + | mesh-coordinate + | name + | constant + | tuple-expression + | tensor + | primary '.' identifier +tuple-assignment ::= '(' identifier (',' identifier)* ')' '=' runtime-expression +where-annotation ::= 'where' '(' ')' +statement ::= for + | with + | tuple-assignment + | identifier '=' (runtime-expression | expression) + | identifier ':' (where-annotation | type-annotation) ('=' + (runtime-expression | expression))? + | 'return' (runtime-expression)? + | runtime-expression + | 'pass' +block ::= (statement (newline statement)*)? +function ::= 'def' name '(' signature ')' ('->' return-type)? ':' block +``` + + +### 2.2 Rules + + +| Owner | Situation | Rule | Statement | Source | +| --- | --- | --- | --- | --- | +| binary_expression | expression | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | expression | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | expression | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | slice_endpoint | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | slice_endpoint | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | slice_endpoint | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| binary_expression | subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| dim_expr | dim_expr | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | layout_extent | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | layout_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | tensor_dim_expr | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | tensor_optional_slot | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dim_expr | tensor_shape | ShapeDimRule | A shape dimension must be an integer, DimVar, or expression. | src/tilefoundry/parser/ast_pattern.py | +| dtype | tensor_dtype | CanonicalDTypeRule | A dtype must resolve to a canonical DType. | src/tilefoundry/parser/ast_pattern.py | +| explicit_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| explicit_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| function | function | FunctionDialectRule | A function kind and constructed value must agree with the active dialect. | src/tilefoundry/parser/pattern_nodes.py | +| function | function | FunctionRegistrationRule | A validated function must be registered exactly once in its owning scope. | src/tilefoundry/parser/pattern_nodes.py | +| function | function | FunctionReturnRule | A HIR function body's inferred type must match its return type. | src/tilefoundry/parser/pattern_nodes.py | +| function | function | FunctionRoleValidationRule | A root, variant, or converter must satisfy its role before registration. | src/tilefoundry/parser/pattern_nodes.py | +| function | function | FunctionSignatureRule | A function must construct an ordered parameter tuple. | src/tilefoundry/parser/pattern_nodes.py | +| layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| op_call | expression | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | expression | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | expression | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | slice_endpoint | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | slice_endpoint | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | slice_endpoint | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| op_call | subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| placed_layout | layout_shape | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | layout_shape | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | tensor_shape | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| placed_layout | tensor_shape | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| plain_layout | tensor_optional_slot | LayoutPositionRule | A layout must be legal for its parser position. | src/tilefoundry/parser/ast_pattern.py | +| plain_layout | tensor_optional_slot | LayoutShapeRule | A layout must have a valid non-boolean shape. | src/tilefoundry/parser/ast_pattern.py | +| shape | layout_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | +| shape | layout_strides | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | +| shape | tensor_shape | ShapeTupleRule | A shape must construct a tuple of dimensions. | src/tilefoundry/parser/ast_pattern.py | +| storage | tensor_optional_slot | StorageValueRule | Storage must resolve to a StorageKind. | src/tilefoundry/parser/ast_pattern.py | +| tensor | annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | +| tensor | annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +| tensor | expression | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | +| tensor | expression | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +| tensor | slice_endpoint | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | +| tensor | slice_endpoint | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +| tensor | subscript_index | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | +| tensor | subscript_index | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +| tensor | type_annotation | TensorLayoutStorageRule | A tensor type must contain compatible layout and storage values. | src/tilefoundry/parser/ast_pattern.py | +| tensor | type_annotation | TensorPositionRule | A tensor type's storage must be legal for its dialect and position. | src/tilefoundry/parser/ast_pattern.py | +| unary_expression | expression | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | expression | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | expression | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | slice_endpoint | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | slice_endpoint | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | slice_endpoint | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | subscript_index | CallBindingRule | A call must bind its arguments into a Call tuple. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | subscript_index | CallExpectedTypeRule | A call's inferred type must satisfy the expected expression type. | src/tilefoundry/parser/pattern_nodes.py | +| unary_expression | subscript_index | CallTypeInferenceRule | A call's result type must be inferred from its binding. | src/tilefoundry/parser/pattern_nodes.py | +| module | module_function | ModuleFunctionValidationRule | A module function must satisfy its root, variant, or converter role before mutation. | src/tilefoundry/parser/ast_pattern.py | +| module | module_function | ModuleFunctionRegistrationRule | A validated module function must be recorded in declaration order. | src/tilefoundry/parser/ast_pattern.py | +| module | module_finalization | ModuleFinalizationRule | A module declaration must contain valid unique members and a resolvable entry. | src/tilefoundry/parser/ast_pattern.py | + + + +## 3. Implementation Overview + +| Component | Responsibility | +| --- | --- | +| Parser API and Context | Receives authored Functions and carries dialect, role, scope, and recursion inputs. | +| Executable Pattern Graph | Composes concrete AST elements into the Function root pattern. | +| Match and Construction | Matches recursively into `AstMatch`, then constructs owner values on return. | +| Ordered Rules | Validates and normalizes each owner value after construction. | +| Module Build | Lets Python execute the class body, collects declarations, resolves child Modules first, then parses Functions in source order and finalizes the Module. | +| Pattern Visitor | Traverses the same graph to render this section's generated grammar and constraints. | ```mermaid classDiagram - class op_registry { - get_schemas(dialect, name) list~OpSchema~ - iter_schema_names(dialect) - _register_schema(schema) - } - class OpSchema { name; dialect; category; signature; builder; op_class } - class tilefoundry_dsl_tf { __getattr__(name); __dir__() } - class tilefoundry_dsl_T { __getattr__(name); __dir__() } - class Op { <> } - - op_registry "1" --> "*" OpSchema : indexes - tilefoundry_dsl_tf ..> op_registry : queries dialect="tf" - tilefoundry_dsl_T ..> op_registry : queries dialect="T" - OpSchema --> Op : op_class (None for alias schemas) -``` - -### 2.3 Resolution algorithm - -The CPython attribute-access path runs `module.__getattr__(name)` -for any name not found on the module's own namespace. Both -namespaces implement it identically: - -```python -def __getattr__(name: str) -> type[Op] | Callable: - """Resolve a dialect name to its Op class or alias builder.""" - ... - -def __dir__() -> list[str]: - """Return the dialect's registered schema names.""" - ... -``` - -`__getattr__` looks up `op_registry.get_schemas(_DIALECT, name)` and raises -`AttributeError` on a miss; for a single real-Op schema it returns the `Op` -class, for a surface-alias schema (`op_class is None`) the alias builder fn, and -for more than one schema an overload resolver. - -Both forms (Op class for real-Op schemas, alias builder fn for -alias schemas) carry an `_op_schema` attribute, so the parser's -bare-name resolver looks the schema up with a single -`getattr(val, "_op_schema", None)` regardless of which form was -bound. - -`from tilefoundry.dsl.tf import *` invokes `__dir__` and then -`__getattr__` for every returned name; `from tilefoundry.dsl import tf` -binds the module object itself, leaving each later `tf.` -attribute access to `__getattr__`. - -### 2.4 `.pyi` stub regeneration - -The dynamic `__getattr__` surface is invisible to static analysers -and editors. To restore IDE completion / type inference, -`tilefoundry.dsl._stub_gen` emits per-namespace `.pyi` stubs derived -from the OpSchema registry: - -``` -tilefoundry/dsl/tf/__init__.pyi # generated, gitignored -tilefoundry/dsl/T/__init__.pyi # generated, gitignored -``` - -The CLI is `python -m tilefoundry.dsl regen`. The generator walks -every `OpSchema` registered for the dialect and emits one -`def (: [, ...]) -> Expr: ...` signature per -schema. Multi-schema overloads emit `@typing.overload` stubs in -registration order, followed by a final non-overload signature -that matches the runtime resolver. - -Conventions: - -- `kind="input"` ParamDefs render as `Expr` regardless of their - declared `annotation` (operands are always Exprs at the DSL - surface). -- `kind="attribute"` ParamDefs render their `annotation` verbatim - (`int` / `str` / `ShardLayout` / …). Referenced types - are auto-imported in the generated header so the `.pyi` is - self-contained. -- A `DType` attribute renders as `Literal[] | DType`, with the - `Literal` members derived from the closed descriptor set in - [types §3](./types.md#3-dtype). The string form is the canonical DSL authoring - path and the parser normalizes it to the corresponding descriptor at the call - boundary. A descriptor value remains accepted as the IR-canonical attribute - value in direct Python expressions. -- Any other string-valued enum attribute, such as `ReduceKind`, renders as - `Literal[] | `. Its `Literal` members derive from - the enum, and the parser normalizes a string to the corresponding enum member - at the call boundary. - -Stubs are not part of the runtime resolution path; the parser still -goes through [§2.3](#23-resolution-algorithm). They exist solely so editors can show typed -completions for `tf.(...)`. - -### 2.5 Invariants - -- **Dialect isolation**. `tilefoundry.dsl.tf` MUST surface only - schemas with `dialect="tf"`; `tilefoundry.dsl.T` MUST surface only - schemas with `dialect="T"`. [§4.6](#46-per-dialect-strict-resolution)'s strict per-dialect resolution - depends on this. -- **Late-registration visibility**. An Op registered after the - namespace module is first imported is visible on the next - `__getattr__` call; the namespace MUST NOT cache resolutions in - a way that would hide it. -- **Implementation independence**. The DSL surface MUST NOT depend - on the `tilefoundry.ir..` directory layout. DSL - source addresses Ops only through `(dialect, name)`. -- **Single-schema identity**. For a single-schema name `n`, - `getattr(tilefoundry.dsl.tf, n) is get_schemas("tf", n)[0].op_class`. - No wrapper class is interposed. - -### 2.6 Platform sub-namespaces - -`tilefoundry.dsl.T` exposes platform-specific instruction and atom -surfaces under a fixed set of **platform sub-namespaces** (e.g. -`T.cuda`). `dsl.T.__getattr__` resolves a platform name **before** the -OpSchema registry lookup ([§2.3](#23-resolution-algorithm)): a name in the platform set returns the -platform namespace object; every other name falls through to the -registry. - -- A platform sub-namespace is not an `Op` and carries no `_op_schema`. - It surfaces platform-specific descriptors (instruction specs, atoms) - only, never catalogue Ops. [§2.5](#25-invariants)'s dialect-isolation invariant is - unaffected — the platform set is disjoint from registered Op names. -- The platform set is fixed; a name outside it MUST resolve as an - ordinary `dialect="T"` Op name, preserving late-registration - visibility ([§2.5](#25-invariants)). -- `T.cuda.mma` is the CUDA MMA surface: `T.cuda.mma.` is an - `MmaOpSpec` and `T.cuda.mma.atom(op=...)` an `MmaAtom` - ([tir §2.3](./tir.md#23-tir-ops)). - The folder name (`cuda`) matches `codegen/cuda/` and the runtime tree. - -In a `@prim_func` body a chain rooted at a platform sub-namespace is a -**compile-time static binding**, not a runtime value: `op = T.cuda.mma.` -and `atom = T.cuda.mma.atom(op=op)` bind Python descriptor objects in -the parser environment and emit no `LetStmt`. A subsequent `atom.A/B/C` -attribute access resolves statically against the bound descriptor. - -The `.pyi` stub generator ([§2.4](#24-pyi-stub-regeneration)) emits the platform sub-namespace -surface so editors complete `T.cuda.mma.` and `.atom(...)`. - -### 2.7 `@module` authoring surface - -`@module(entry="", target=..., topologies=(Topology(...), ...))` collects -a class body into a `Module` -([core-ir §1](./core-ir.md#1-module)). The decorated name binds to the -resulting `Module`. `target` MUST be a constructed Target instance; -a string MUST be refused and MUST NOT be resolved or constructed. The value -declares the hardware the domain runs on. `topologies` declares the ordered -`Topology` hierarchy (see below). - -A file of `@module` classes is an ordinary importable Python module: every name -its class bodies read — the shape configuration above all — MUST resolve within -that file, and the file MUST NOT require execution with a namespace injected into -its globals. The Modules it defines are module-level values, so `import` reaches -them, a linter sees them, and a CLI selector addresses them by name. A file that -has to be executed with a namespace injected is reachable by none of those. - -A `@module` class body is evaluated once where it stands, so a body at file scope -states one shape. A model asked about more than one structural configuration — -shapes differing in a submodule count or a per-layer tuple, not only in a tensor -axis — MAY instead place its class bodies in a function of that same file that -takes the configuration as a parameter, and publish that function; each call -states the same source at the shape its caller names. This is not injection: the -file is still an ordinary import and the configurations are values its own -package publishes. `@func` MUST resolve the names in its signature and body -against the locals of every scope enclosing it, innermost first, so a parameter -of that function is as visible to a nested class body as a module-level import -is. - -- Every non-dunder class member MUST be one of three kinds: an `@func` / - `@prim_func` result (an `hir.Function` / `tir.PrimFunction`); a child - `Module` — or a tuple/list of them, how a factory attaches N identical - instances under one attribute (each already named by the factory, e.g. - `renamed(f"layer{i}")`); or a plain Python function (an orchestration - method, e.g. `forward` / `init_caches`). Any other member — a stray - attribute, an undecorated method that is neither a DSL function nor a - plain orchestration function, … — MUST be rejected. A specialization - variant (a `@base.specialize` def) and a per-weight converter (a - `@base.converter(name)` def, - [runtime §1.1.2](./runtime.md#112-weight-converter-and-prepare--forward)) - are not standalone members — both live on their base function and are - skipped when collecting. -- A nested `class` statement is a legal member when it is itself decorated - with `@module(...)`: by the time the outer class body finishes running, - the inner decorator has already replaced the name with a `Module` - instance, so it is collected as an ordinary child `Module`. An - *undecorated* nested class is rejected (it is none of the three kinds). -- `@func` / `@prim_func` results are collected in **definition order** into - `Module.functions`; the class body MUST contain at least one. A child - `Module` is collected into `Module.modules`, renamed to the attribute it - is attached under — torch / HuggingFace checkpoint-naming semantics: - assigning a child to `self.self_attn` names it `self_attn` in the tree, - independent of the child's own class name - ([core-ir §1](./core-ir.md#1-module)). A plain Python function is - collected into `Module.methods` by its own name. A duplicate function - name, or duplicate child module name, across the class body MUST be - rejected. -- A class body MUST declare at least one function, child `Module`, or plain - method; only an empty body MUST be rejected. A methods-only Module is therefore - valid — it composes the children it is given. -- `entry` is optional. Supplied, it MUST name exactly one collected function and - an unknown name MUST be rejected. Omitted, the Module has no default step. -- A method's name is free, but `forward` is the one a bare `(...)` - delegates to; any other name is reached only by naming it. A class-body - `__call__` MUST be rejected: Python resolves a dunder on the type, so one - attached to the built Module instance would never run. -- An HIR member MAY call a sibling HIR function **defined above it**. The call - resolves to a `Call` targeting that Function and remains inside the current - kernel invocation ([hir §1.1](./hir.md#11-function)); a forward reference to a - sibling defined below stays unresolved and MUST fail. -- An HIR member MAY equally call a child `Module` bound **above it** in the same - class body; that form, what it resolves to, and what it refuses are - [§4.2](#42-closure-then-registry-callee-resolution). -- The `topologies=(Topology(...), ...)` decorator argument declares the - domain's complete ordered hierarchy. Omitting it inherits the owning class's - hierarchy; `()` declares an explicitly topology-free domain. The value MUST - be a tuple of `Topology`; a value that is not MUST be rejected rather than - read as an empty hierarchy. -- The printer emits this surface: shared meshes at module level (before the - class) so the class body stays function-and-nested-Module-only, then the - `@module(...)` decorator, then the functions and nested Modules. What that - decorator and that order print exactly is - [inspection §2.2](./inspection.md#22-module-printer). - -#### Design rationale - -`entry` is a function-name forward reference rather than a function object -because a class decorator's arguments are evaluated before the class body runs, -so the entry function does not yet exist when `@module(entry=...)` is called. - -The `topologies` decorator argument is evaluated before the class body runs. -The resulting declaration is retained while that body is evaluated, so a -function body MAY name a level of its domain (`with Mesh(("cta",), ...)`) when -`@func` parses it. Nested `@module` bodies resolve the declaration belonging to -their own class body first; omission walks outward to the owning class, while -an explicit `()` stops inheritance. `target` needs no such early lookup because -nothing consumes it until after the `Module` exists. - -## 3. Parser architecture - -### 3.1 Model - -The implementation lives under `tilefoundry/parser/`. The current -function-level entry points are independent calls; there is no -`ModuleContext` / `FunctionDecl` data class. - -```python -# tilefoundry/parser/hir_parser.py -def parse_func(fn, *, topologies=(), specializations=(), extra_closure=None) -> hir.Function: ... -class _HirBodyVisitor(BaseExprVisitor): ... - -# tilefoundry/parser/tir_parser.py -def parse_prim_func(fn, *, target=None, extra_closure=None) -> tir.PrimFunction: ... -class _TirBodyVisitor(BaseExprVisitor): ... - -# tilefoundry/parser/base.py -def extract_ast(fn) -> ast.FunctionDef: ... -class BaseExprVisitor: ... - -# tilefoundry/parser/symtab.py -class LexicalEnv: - def push_frame(self) -> None: ... - def pop_frame (self) -> dict[str, Any]: ... - def define (self, name, value) -> None: ... - def lookup (self, name) -> object: ... - def innermost_mesh(self) -> Mesh | None: ... - -# tilefoundry/parser/dispatch.py -def resolve_op (name) -> type | None: ... -def resolve_stmt (name) -> type | None: ... -def resolve_schema(name, dialect: str = "tf") -> OpSchema | None: ... -def resolve_callable(name, token: Literal["hir", "tir"]) -> tuple[str, type]: ... + ParserAPI --> FuncParserContext + ParserAPI --> FunctionPattern + AstPattern <|.. Element + Element o-- AstPattern + Element o-- AstRule + AstPattern --> AstMatch + PatternVisitor ..> AstPattern + ParserAPI ..> ModuleBuild ``` -Each function-level parser collects a closure dict from the live -Python function (`_collect_closure(fn) -> dict[str, Any]`), reads -the AST via `extract_ast(fn)`, and walks the body with the -dialect's `BaseExprVisitor` subclass. - -### 3.2 Class diagram - ```mermaid -classDiagram - class parse_func - class parse_prim_func - class BaseExprVisitor { <> } - class _HirBodyVisitor - class _TirBodyVisitor - class LexicalEnv - class dispatch { resolve_op; resolve_stmt; resolve_schema; resolve_callable } - - parse_func ..> _HirBodyVisitor : drives - parse_prim_func ..> _TirBodyVisitor : drives - BaseExprVisitor <|-- _HirBodyVisitor - BaseExprVisitor <|-- _TirBodyVisitor - _HirBodyVisitor o-- LexicalEnv - _TirBodyVisitor o-- LexicalEnv - _HirBodyVisitor ..> dispatch : callee lookup - _TirBodyVisitor ..> dispatch : callee lookup -``` - -### 3.3 Description - -`parse_func` / `parse_prim_func` consume a live Python function (`fn`); -`topologies` supplies the parse-time namespace a body may name and is not -retained on the resulting `Function`. The `@tilefoundry.module` decorator -builds a `core_ir.Module` from the class's already-parsed `@func` / -`@prim_func` methods and nested Modules. A standalone `@func` returns an -`hir.Function`, or the implicit single-function `core_ir.Module` when its -decorator declares execution context. - -The closure dict supplies same-module callee lookup. Names defined -in the user's Python module (other `@func` / `@prim_func` -functions, mesh / topology objects, Op classes imported from -`tilefoundry.dsl.tf` / `T`) are visible through the closure. The -closure also includes the `@func` / `@prim_func` bindings present in -the *definition frame* when the decorator runs — for a -`@tilefoundry.module` class body, that is the sibling methods declared -above the one being parsed. Each such binding **is** the sibling's -`hir.Function` / `tir.PrimFunction` IR node (the decorator evaluates to -the IR directly, [§1.1](#11-decorators)), so a sibling callee resolves to that `Function` -and becomes the `Call` target directly. This is what makes -callee-before-caller sibling calls work; a forward reference is simply -absent from the closure and fails as an unresolved callee. The merge is -additive: it never shadows the function's own globals / freevars. - -`LexicalEnv` is a frame stack used by both body visitors for -parser-time bindings (Mesh axes, the Python `slice` from two-argument `tile`, SSA -aliasing). Frame push / pop matches the Python-source scope -(`with Mesh(...)`, HIR grid loops). - -`dispatch.resolve_callable(name, token)` performs strict -per-dialect Op resolution against `op_registry`; both body visitors' -`ast.Name` callee resolution (`BaseExprVisitor._resolve_call_target`, -shared by HIR and TIR) delegates to it once the closure path misses, and -the TIR top-level-statement dispatch (`_call_as_top_level_stmt`) delegates -to it for a bare-name effect Stmt / intrinsic before falling back to -`call_to_op_call`. It never resolves an arbitrary Python name — only a -name already cataloged as an Op / Stmt / intrinsic under the body's own -dialect — so an undefined name still raises *unknown Op name*. - -There is no parser-side intermediate IR; function bodies translate -directly into `core_ir` nodes plus dialect-specific subclasses. - -## 4. Shared parsing machinery - -### 4.1 Lexical environment - -Both parsers use the same lexical-env stack. `define(name, expr_node)` -binds a Python name to an `Expr` object. Subsequent uses of that name -reuse the same `Expr`, which is how HIR's SSA-as-DAG sharing falls -out for free. - -### 4.2 Closure-then-registry callee resolution - -Bare-name callees resolve through the lexical env + the function's -closure first — the common case, covering every name reached via a -star-import or an explicit `tf.` / `T.` binding. When that -misses, resolution falls through to `dispatch.resolve_callable(name, -token)` ([§4.6](#46-per-dialect-strict-resolution)): dialect-strict dispatch against the Op / Stmt / intrinsic -catalogue, not an arbitrary-name lookup, so a name that is neither bound -in the closure nor cataloged under the body's own dialect still raises -*unknown Op name*. - -A bare-name callee bound to a `Function` is that Function, and one bound to a -`Module` is the function that Module's `entry` names -([core-ir §1](./core-ir.md#1-module)). Those two are the only callee forms that -resolve to a function, and either yields an ordinary Function `Call` -([hir §1.1](./hir.md#11-function)). A `Module` callee belongs to the HIR body -parser alone; in TIR a name bound to one is not a callee and still raises -*unknown Op name*. - -- constraints: - - In an HIR body, a bare name bound to a `Module` MUST resolve to the function - that Module's `entry` names, and only while the calling function is being - authored inside an active `@module()` / `@module(...)` class body — the body - that attaches the callee and so gives the call a child to reach. Elsewhere - the callee MUST be refused naming that, whether or not the calling function - declares execution context of its own and whether or not something later - lifts it into a Module. Which declaration is active MUST be decided by the - declaring frame, not by one being open somewhere, so a declaration a failed - class body left open attaches nothing. Only a direct `name = ` - binding in that body attaches a callable name; a name it does not attach - MUST be refused, naming what it does bind. - - A Module declaring no entry, and one whose entry is not an HIR `Function`, - MUST each be refused at the call site naming that reason, rather than - falling through to *unknown Op name* — the callee resolved, and what it - could not answer is which function a call of it runs. - - The parser MUST carry the binding name it reached that Module through as - private authoring state on the resulting `Call`. A class body is parsed - before its children are attached and attaching renames — so copies — the - Module the body named ([core-ir §1](./core-ir.md#1-module)), and that state - is what `@module` collection needs to rebuild every such call, inside a - specialization variant included, against the child attached under that name. - Two attributes may hold copies of one Module, and then the binding is the - only thing that says which copy a call meant. It is not part of the IR - contract: collection consumes it and takes it off the call, the rebuilt - `Call.target` having stated the callee. - - The parser MUST check a written call's arity against the supplied-parameter - set HIR resolves for it ([hir §1.1](./hir.md#11-function)), and where that - set is narrower than the declared parameters its diagnostic MUST count in - activations. - - An attribute callee whose base name is bound to a `Module` MUST be refused, - naming that a Module is called through its bare binding and its entry. This - holds for the attribute that names the entry as well: which name was written - is not what decides where a call goes. - -The closure binding for a name from `tilefoundry.dsl.tf` / -`tilefoundry.dsl.T` is whatever its module `__getattr__` returns: - -- a real-Op class for single-schema names whose schema has an - `op_class` (e.g. `tf.matmul` → `MatMul`); -- the alias's builder function for surface-alias schemas (e.g. - `tf.add` → `_add_alias`); the function carries `_op_schema` so - the parser still recovers the schema by attribute lookup. - -Both forms expose `_op_schema`, so the parser's -`_resolve_call_target` returns an `OpSchema` uniformly. Namespace- -attribute callees (`tf.add` / `T.copy`) skip the closure binding and -go directly through `dispatch.resolve_schema(name, dialect)`, which -honours alias prepend order — an alias schema (if any) wins over a -legacy real-Op schema sharing the same name. - -### 4.3 OpSchema and overload resolution - -A registered Op has one or more `OpSchema` entries indexed by -`(dialect, name)`. Each schema lists the Op's `ParamDef` descriptors -(see [core-ir §2.3](./core-ir.md#23-op)). When the parser sees a callee: - -1. Look up the schema list via - `op_registry.get_schemas(dialect, name)`. -2. Filter by arity. `ParamDef.is_required` (i.e. `default is - MISSING`) sets `n_min`; `optional` does NOT lower `n_min`. -3. For the surviving candidates, walk each input ParamDef and run - `pattern.match(arg_type)`. `pattern is None` accepts any. -4. Return the first schema whose every input pattern matches. - Registration order is the tiebreaker; there is no "best match" - search. - -### 4.4 Annotation-driven sugar dispatch - -At each attribute slot of a call, the parser consults the matched schema's -`ParamDef.annotation` and invokes `parse_sugar(node, expected, ...)`: - -- `annotation=ShardLayout` → `expected=ShardLayout` -- `annotation=Layout` → `expected=Layout` -- `annotation=TensorType` → `expected=TensorType` - -```python -def parse_sugar( - node: ast.AST, - expected: type, - *, - closure: dict[str, Any] | None = None, - mesh_resolver: Callable[[str], Mesh] | None = None, -) -> Layout | ShardLayout | TensorType | None: - """Parse one type-directed layout or tensor-type sugar form. - - Args: - node: Static tuple syntax. - expected: Required result family. - closure: Optional static-name bindings. - mesh_resolver: Optional lexical mesh lookup. - - Returns: - The parsed value, or None when a TensorType head is not recognised. - """ - ... -``` - -- constraints: - - The `Layout` branch MUST resolve closure-bound integer extents and - MUST derive C-order strides when the tuple omits them. - -Sugar dispatch is annotation-driven, not name-driven: an attribute -called `shape` will not be parsed as layout sugar unless its -ParamDef declares a layout annotation. Ops without a registered -schema fall through to a small legacy heuristic -(`attr_name == "layout"` ⇒ `ShardLayout` sugar) until they migrate. - -Attribute string normalization is also annotation-driven: - -- `annotation=DType` resolves a canonical string by descriptor `name` and - rejects any other string with a `VerifyError`. -- A string-valued Enum annotation resolves by Enum value and rejects any other - string with a `VerifyError`. - -### 4.5 `Tensor[...]` type-literal surface - -`Tensor[shape, dtype, layout, storage]` is recognised by `parse_sugar`, the -single public entry point for `Layout`, `ShardLayout`, and `TensorType` sugar. -The caller supplies the expected result type plus contextual closure and mesh -resolution. Both `@func` and `@prim_func` annotations and body expression -positions use this entry point. The shape and dtype slots are required; layout -and storage are optional. Placement sugar MAY appear directly in the shape -slot, in which case the third slot is storage and no separate layout slot is -accepted. - -The dtype slot MUST resolve to the canonical descriptor named by its string. -Unknown names MUST be rejected and MUST NOT silently select `DType.f32`. - -### 4.6 Per-dialect strict resolution - -`parser.dispatch.resolve_callable(name, token)` does NOT fall back -across dialects. An HIR-only Op (e.g. `rope`) raises *unknown TIR -callable* in a TIR body, and a TIR-only Op (e.g. `copy`) raises -*unknown HIR callable* in an HIR body. The trailing-underscore -selector is gated to the TIR token only. - -### 4.7 Restricted static evaluation - -```python -def eval_static( - node: ast.AST, - *, - closure: dict[str, Any], - lookup: Callable[[str], Any] | None = None, - allowed_nodes: tuple[type, ...] = ALL_NODES, - div: DivMode = "true", - attr_resolver: Callable[[Any, str], Any] | None = None, - on_closure_name: Callable[[Any, str], None] | None = None, -) -> Any: - """Evaluate the parser's restricted static-AST subset. - - Args: - node: AST node to evaluate. - closure: Python closure bindings. - lookup: Optional parser-lexical name resolver. - allowed_nodes: Admitted AST node classes. - div: Division policy. - attr_resolver: Optional attribute resolver. - on_closure_name: Optional closure-use callback. - - Returns: - The statically evaluated value. - """ - ... -``` - -- constraints: - - Lexical `lookup` MUST run before closure fallback. - - A node outside `allowed_nodes`, an unresolved name, or an unsupported - operator MUST raise `VerifyError`. - - `div="floor"` MUST affect `/`; `//` is always floor division. - -## 5. HIR parser - -The HIR parser walks an `@func` body. The body is a sequence of -Python statements that the parser folds into a single `Expr` tree. - -| Python | HIR action | -|---|---| -| `x = expr` | `define(x, expr_node)`; no IR node. Subsequent `x` reuses the same `Expr` (SSA-as-DAG). | -| `x + y` | `Call(Binary(kind=ADD), (x, y))`; the parser maps Python AST `BinOp` / `Compare` / `BoolOp` directly to a `Binary` instance with the matching `BinaryKind`. `UnaryOp` USub / Not maps similarly to `Unary(kind=NEG)` / `Unary(kind=NOT)`. AST `@` (matmul) routes to `MatMul` (a real Op, not kinded). | -| `foo(a, b)` | `Call(target_op, args)` where `target_op` is constructed by the resolved schema's `builder` ([§4.2](#42-closure-then-registry-callee-resolution) / [§4.3](#43-opschema-and-overload-resolution)). For surface aliases (e.g. `add` / `cmp_eq` / `neg`), the alias's builder returns the kinded target Op (`Binary(kind=...)` / `Unary(kind=...)`); for real Ops, the default builder is the Op class itself. | -| `for i in tile(...)` / `for i in range(...)` | `GridRegionExpr` (see [§1.7](#17-for-i-in-tile--for-i-in-range-hir-only) and below). | -| `with Mesh(...) as m` | Push `m` onto the parser-lexical stack; pop on exit. No IR node. Every `Call` built inside carries the stack as `ExecutionDomainMetadata`. | -| `return expr` | Sets `Function.body`. A `return` without a value is rejected. | -| `return (a, b)` / `return a, b` | A literal tuple return (both spellings are the same AST) folds to a core `Tuple` body ([core-ir §2.2](./core-ir.md#22-var--constant--tuple)); `Function.return_type` is the `TupleType` of the element types. Callers destructure via the existing tuple-unpack rule (`o, s = f(...)`). | -| `pass` | Accepted only as the **entire** body: sets `Function.body = None`, declaring a dispatch prototype whose implementations are registered via `.specialize` -([hir §1.1](./hir.md#11-function)). A `pass` mixed with any other statement is rejected. | - -An assignment whose RHS computes a new expression records the LHS as that -expression's binding. A bare-name RHS computes nothing: rebinding an existing -value under either a new or existing name (for example `acc = x` to initialize -a loop carry, or the final `m = m_new` carry update) MUST reuse the same `Expr` -object without replacing its binding metadata. A new name updates the parser's -symbol table; it does not add or rename an IR node. - -`for` / `if` / `while` over arbitrary ranges, conditionals, and other -Stmt forms are TIR-only. They are rejected by the HIR parser. - -A `pass` body yields `Function.body is None`, declaring a **dispatch -prototype** that awaits variants. Immediately after `@func def f: pass` -and before any `@f.specialize(...)`, the base is transiently -`body is None, variants == ()` — a valid *unsealed authoring* state. The -sealed (verified) invariant is `body is None` ⟺ `variants != ()` -([hir.md §1.1](./hir.md#11-function)); the verifier rejects a -`body is None` function with no variants, a variant whose body is `pass` -(a variant MUST carry a real body), and a real body combined with -variants. The `@base.specialize(...)` parse rejects a `pass`-bodied -variant directly. - -### 5.1 GridRegionExpr carry-out lifting - -Inside an HIR grid-loop body, an `ast.Assign` whose single -`Name` target is already bound in *outer* scope is a loop-carried -rebinding. The parser: - -1. Allocates a fresh phi `Var` per carry name (same type, same name). -2. Records the phi in `GridRegionExpr.carried_args`. -3. Snapshots the final RHS bound to that name as a `yield_value`. -4. After the loop, rebinds the carry name in the outer frame to the - `GridRegionExpr` (single carry) or projects each carry value out - of its `TupleType` result (multi-carry). - -Only `=` assignments are accepted; `+=` is rejected. `return` and a -nested `with` inside a loop body are rejected. A nested `for ... in -tile/range(...)` IS allowed and lifts to a nested `GridRegionExpr`; the -carry scan recurses into it, so an outer-scope name rebound only inside the -nested loop is carried across both loops. - -### 5.2 Constraint attachment - -The HIR parser attaches one `ScheduleConstraintMetadata` record to one -existing tensor `Expr`. Attachment MUST update that node in place: it MUST NOT -rebuild the annotated value or any consumer in its SSA DAG. An annotation with -a value still follows the assignment rule above, so it may annotate a newly -computed RHS or bind another name to an existing value without copying it. -Tensor parameters, tensor-valued intermediate SSA -values, and bound tensor-valued `TupleGetItem` values are valid subjects. -Whole tuples, shape scalars, unit values, direct subscripts, and unresolved -names are rejected. Inline and standalone annotations for the same Expr are -duplicates, not merged declarations. Diagnostics identify the subject and -retain the authored source location. Constraint metadata does not alter the -tensor type or introduce an HIR node. - -## 6. TIR parser - -The TIR parser walks a `@prim_func` body. The body is a sequence of -imperative statements that fold into a `Sequential` of Stmts. - -| Python | TIR action | -|---|---| -| `x = expr` | `LetStmt(var=x, value=expr, body=)`. The remaining body of the function is nested as `body`. | -| `a = Tensor(...)` | `LetStmt(var=a, value=Call(tir.memory.AllocTensor, (), attrs=), body=)`. See [tir §2.3](./tir.md#23-tir-ops). | -| `foo(a, b)` (effect Op) | `Evaluate(target_op, args)` Stmt. | -| `foo(a, b)` (value Op) | `Call` Expr embedded in the right-hand side of a `LetStmt` or another Stmt's Expr field. | -| `for i in range(...)` | `For(induction_var=i, start, stop, step, body)`. | -| `if/elif/else` | `If(cond, then_body, else_body)`. | -| `while` | `While(cond, body)`. | -| `with Mesh(...) as m` | `MeshScope(mesh, binding=m, body)`. | -| `return` | `Return()` Stmt. A `return value` is rejected. | - -TIR has no SSA-as-DAG sharing rule; every binding is an explicit -`LetStmt`. `for i in tile(...)` is HIR-only and is rejected here. - -## 7. Validation and rejection - -- Any `ast` node not in [§4](#4-shared-parsing-machinery) / [§5](#5-hir-parser) is rejected — the unsupported - forms include `try` / `with` over non-Mesh contexts / `lambda` / - list / dict / set comprehensions / `yield` / `async`. -- Cross-dialect callees fall through to *unknown callable* ([§4.6](#46-per-dialect-strict-resolution)). -- A bare-name callee that the lexical env / closure does not - resolve to an `_op_schema`-bearing surface value (`Op` subclass or - alias builder function) is *unknown Op name*. -- `Tensor[...]` with the wrong number of slots, an unknown dtype, a - non-injective layout, or an `ast.Slice` shape element is - rejected. -- `for tile` is HIR-only; emitting it inside a `@prim_func` is - rejected. `with Mesh(...) as m` is accepted in both dialects — in a - `@prim_func` it lowers to a `MeshScope` Stmt ([§6](#6-tir-parser)), unlike the - no-IR-node HIR sugar ([§1.6](#16-with-mesh-as-m)). -- Layout sugar that would lose mesh information falls through to the - verbose form ([§1.4](#14-tensor-and-consttensor-annotations)); if neither is acceptable, the type is - rejected. +flowchart TD + API["parse_function(fn, context)"] --> AST["Extract FunctionDef AST"] + AST --> ROOT["FunctionPattern.match"] + ROOT --> TREE["AstMatch tree"] + TREE --> BACKWARD["construct children, then apply Rules"] + BACKWARD --> FUNCTION["HIR Function / TIR PrimFunction"] + FUNCTION --> MODULE{"Module authoring context?"} + MODULE -->|yes| FINALIZE["defer declaration"] + FINALIZE --> CHILDREN["attach child Modules and bind module scope"] + CHILDREN --> ORDERED["parse roots in source order; then variants/converters"] + ORDERED --> BUILT["construct final Module and verify"] + BUILT --> RETURN + MODULE -->|no| RETURN["return standalone result"] +``` + +Pattern combinators serve both runtime matching and Spec traversal. `AstMatch` separates syntax +matching from object construction, while each Rule reads the recursive context after its owner +value exists. Module class control flow remains Python execution; no Module AST grammar exists. diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index 040ad75c..50e6c828 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -166,7 +166,7 @@ class RuntimeFunction: A weight's converter is registered **per weight**, not per module: `@.converter("")` decorates a throwaway `def` and registers it on the base function's `converters` -([parser §2.7](./parser.md#27-module-authoring-surface)). Its parameters are +([parser §3](./parser.md#3-implementation-overview)). 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 @@ -382,7 +382,7 @@ entry point. It accepts a `hir.Function` or `Module`, normalizes to a 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 §1.1](./parser.md#11-decorators)). + decorator yields ([parser §2.1](./parser.md#21-syntax)). - Mesh layout is expressed in the DSL with lexical `with Mesh(...) as mesh` scopes. - `jit()` has no `cta_mesh` / `thread_mesh` parameters. diff --git a/docs/spec/shard.md b/docs/spec/shard.md index f51fc1e8..28527bea 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -475,7 +475,7 @@ Let `sl: ShardLayout`, `T: TensorType`, and `G = sl.layout.shape`. `N > mesh_extent(a)` is canonicalized at parse time into a factorised form (`(mesh_extent(a) @ m.a, N // mesh_extent(a))`); the factorised residual axis enters the IR as a non-`Split` layout dim. See - [parser §1.5](./parser.md#15-layout-sugar). + [parser §2.1](./parser.md#21-syntax). - `local_shape(sl)[k] = G[k] / sl.mesh.layout.shape[a] = 1` iff some mesh axis `a` has `sl.attrs[a] = Split(k)`. - `local_shape(sl)[k] = G[k]` otherwise. @@ -543,7 +543,7 @@ not `Reshard`. Logical tensor `(2, 1536)` reshards via surface sugar `(2 @ m.x, 12 @ m.y, 128 @ m.t)` with `mesh=(x=2, y=4, t=32)`. Parser -canonicalization ([§7.1.1](#711-layoutshape), [parser §1.5](./parser.md#15-layout-sugar)) expands `12 @ m.y` into +canonicalization ([§7.1.1](#711-layoutshape), [parser §2.1](./parser.md#21-syntax)) expands `12 @ m.y` into `(4 @ m.y, 3)` and `128 @ m.t` into `(32 @ m.t, 4)` and emits `Layout(shape=(2, 4, 3, 32, 4), strides=None)` — un-materialized because the user wrote sugar ([§7.1.2](#712-layoutstrides)). Reshard typeinfer diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 2f78a089..8e161c32 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -917,7 +917,7 @@ class MmaOpSpec: The realized atom for an `op` (the CuTe `MMA_Atom` analog), built by `T.cuda.mma.atom(op=...)` -([parser §2.6](./parser.md#26-platform-sub-namespaces)). +([parser §2](./parser.md#2-syntax-and-rules)). ```python class MmaAtom: @@ -981,7 +981,7 @@ with Mesh((Topology("thread", 32),), Layout(shape=(4, 8), strides=(1, 4))) as wa `atom.A/B/C` layout and fills it with its own `T.copy`. The accumulator is initialised with `Fill` and then read-modify-written. - `atom` is a compile-time attribute on the `Mma` Op - ([parser §2.6](./parser.md#26-platform-sub-namespaces)), not a runtime + ([parser §2](./parser.md#2-syntax-and-rules)), not a runtime operand. When absent, lowering takes the bare-`Mma` per-target path. ##### Verify diff --git a/docs/spec/visitor-registry.md b/docs/spec/visitor-registry.md index 49fe43f4..4ebf4440 100644 --- a/docs/spec/visitor-registry.md +++ b/docs/spec/visitor-registry.md @@ -176,6 +176,20 @@ class FunctionScope: function: Function +@dataclass +class CallFeed(Generic[T]): + """Values supplied to one callee, keyed by formal parameter identity.""" + + by_param: Mapping[int, T] + + def value_for(self, param: Param) -> T: ... + + +class CallFeedProvider(Protocol[T]): + def build_call_feed(self, callee: Function, supplied: tuple[T, ...]) -> CallFeed[T]: ... + def scope_for(self, callee: Function) -> FunctionScope | None: ... + + @dataclass class TypeInferContext: """Walk-local type cache, mesh scope, and elaboration cache. @@ -192,8 +206,12 @@ class TypeInferContext: cache: dict[int, Type] = field(default_factory=dict) mesh_scope: tuple = () elaboration_cache: dict[tuple, Any] = field(default_factory=dict) + call_feed_provider: CallFeedProvider[Type] | None = None + feed: CallFeed[Type] | None = None def type_of(self, expr: Expr) -> Type: ... + def build_call_feed(self, callee: Function, supplied: tuple[Type, ...]) -> CallFeed[Type]: ... + def scope_for(self, callee: Function) -> FunctionScope | None: ... def error(self, node: Expr | Stmt, msg: str) -> NoReturn: ... ``` @@ -208,10 +226,11 @@ nothing of that kind rather than guessing. - `scope` MUST be the only context state describing where a walk is reading, and the pair MUST be reachable from the package root together, since one is how the other is constructed. - - A walk-visible query answering a question about one construct — which Module - a particular kind of callee belongs to, how a particular call binds its - arguments ([hir §1.1](./hir.md#11-function)) — MUST NOT be added to the - context; such a question is resolved by whoever asks it, from `scope`. + - A `CallFeed` MUST contain only the formal parameter identity to value mapping; + it MUST NOT carry a Module, scope, reading, or parser metadata. + - `CallFeedProvider` is context-owned: Parser, Type Inference, and function-level + Evaluator provide their own value type and ownership rules. HIR consumes + `TypeInferContext.build_call_feed()` and does not own a resolver. - `type_of` is a walk-local cache only — it holds no dispatch rule of its own. A cache miss delegates to `TypeInferVisitor(self).visit(expr)` (below), whose `visit_Call` is what consults diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index c6f80010..bf26c296 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -15,9 +15,8 @@ Var, get_metadata, ) -from tilefoundry.ir.core.module import Module, owning_module, subtree +from tilefoundry.ir.core.module import Module, child_module_of, owning_module, subtree from tilefoundry.ir.core.pattern import DimVarRangePat -from tilefoundry.ir.hir._call_binding import binding_for from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.hir.specialize import ( @@ -589,19 +588,9 @@ def visit(owner: Module, prefix: str) -> None: def _call_reading(module: Module, caller: Function, call: Call) -> Module | None: - binding = binding_for( - call.target, - call, - TypeInferContext(scope=FunctionScope(module, caller)), - ) - if not binding.from_reading: - return None - owner = owning_module(module, call.target) + owner = child_module_of(module, caller, call.target) if owner is None: - raise AnalysisError( - f"Function call {call.target.name!r} has no unique owner in Module " - f"{module.name!r}" - ) + return None return owner diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index ac292b39..5c522f6d 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -268,12 +268,16 @@ def _emit(let: LetStmt, ctx: CodegenContext) -> None: dst_local = shard_layout_local_shape(dst_layout) split_axes = {a.axis for a in dst_layout.attrs if isinstance(a, Split)} non_split = [a for a in range(len(dst_local)) if a not in split_axes] - if len(logical_coords) != len(non_split): + if len(logical_coords) == len(dst_local): + coordinate_axes = range(len(dst_local)) + elif len(logical_coords) == len(non_split): + coordinate_axes = non_split + else: raise ValueError( f"tensor_view: {len(logical_coords)} offsets for " - f"{len(non_split)} local axes" + f"{len(non_split)} or {len(dst_local)} local axes" ) - entries = tuple(zip(non_split, logical_coords, let.var.type.shape)) + entries = tuple(zip(coordinate_axes, logical_coords, let.var.type.shape)) kept = tuple( entry for entry in entries if int(upper_bound(dst_local[entry[0]])) != 1 ) diff --git a/src/tilefoundry/dsl/T/__init__.py b/src/tilefoundry/dsl/T/__init__.py index a2e042ba..682b46fc 100644 --- a/src/tilefoundry/dsl/T/__init__.py +++ b/src/tilefoundry/dsl/T/__init__.py @@ -18,7 +18,7 @@ def _resolve_platform(name: str): Platform sub-namespaces (``T.cuda``, later other targets) are compile-time instruction descriptors rather than callable Ops, so they are resolved before the OpSchema registry - ([parser §2.6](docs/spec/parser.md#26-platform-sub-namespaces)). Importing + ([parser §2](docs/spec/parser.md#2-syntax-and-rules)). Importing ``_platforms`` pulls in the CUDA MMA IR modules, which should only happen on first actual ``T.`` access, not on ``import tilefoundry.dsl.T``. diff --git a/src/tilefoundry/dsl/__init__.py b/src/tilefoundry/dsl/__init__.py index 4e6706f4..b11f9da9 100644 --- a/src/tilefoundry/dsl/__init__.py +++ b/src/tilefoundry/dsl/__init__.py @@ -4,7 +4,7 @@ annotations and string dtype sugar complete the source surface. Generated stubs provide static completion. -See [parser §2](docs/spec/parser.md#2-dsl-namespace-surface). +See [parser §2](docs/spec/parser.md#2-syntax-and-rules). """ from __future__ import annotations diff --git a/src/tilefoundry/dsl/_namespace.py b/src/tilefoundry/dsl/_namespace.py index 911d857e..266e7551 100644 --- a/src/tilefoundry/dsl/_namespace.py +++ b/src/tilefoundry/dsl/_namespace.py @@ -5,7 +5,7 @@ Both dialect namespaces resolve names on demand against the OpSchema registry with the same algorithm; this module ships that algorithm once so ``dsl.tf`` / ``dsl.T`` shrink to a dialect string (and, for ``T``, a -platform-sub-namespace pre-resolver, [parser §2.6](docs/spec/parser.md#26-platform-sub-namespaces)). +platform-sub-namespace pre-resolver, [parser §2](docs/spec/parser.md#2-syntax-and-rules)). """ from __future__ import annotations @@ -13,7 +13,7 @@ from typing import Any, Callable, Iterable from tilefoundry.ir.core.op_registry import get_schemas, iter_schema_names -from tilefoundry.parser.overload import resolve +from tilefoundry.ir.core.overload import resolve PreResolver = Callable[[str], Any] @@ -28,7 +28,7 @@ def make_dialect_namespace( runtime resolver. ``__all__`` is computed on demand so later registrations remain visible. - See [parser §2.3](docs/spec/parser.md#23-resolution-algorithm). + See [parser §2](docs/spec/parser.md#2-syntax-and-rules). """ def __getattr__(name: str) -> Any: diff --git a/src/tilefoundry/dsl/_stub_gen.py b/src/tilefoundry/dsl/_stub_gen.py index 1d730b9f..31cbd050 100644 --- a/src/tilefoundry/dsl/_stub_gen.py +++ b/src/tilefoundry/dsl/_stub_gen.py @@ -3,8 +3,8 @@ Inputs use ``Expr`` while attributes retain declared types and required imports. Overloads preserve registration order and include a runtime fallback signature. Generated files are gitignored and rebuilt through the DSL CLI. See -[parser §2.4](docs/spec/parser.md#24-pyi-stub-regeneration) and -[parser §1.7](docs/spec/parser.md#17-for-i-in-tile--for-i-in-range-hir-only). +[parser §2](docs/spec/parser.md#2-syntax-and-rules) and +[parser §2.1](docs/spec/parser.md#21-syntax). """ from __future__ import annotations diff --git a/src/tilefoundry/dsl/_tensor.py b/src/tilefoundry/dsl/_tensor.py index 9bf8a69c..5e41df53 100644 --- a/src/tilefoundry/dsl/_tensor.py +++ b/src/tilefoundry/dsl/_tensor.py @@ -3,7 +3,7 @@ Both resolve to ``TensorType``; ``ConstTensor`` additionally marks the parsed parameter as an external constant. -See [parser §1.4](docs/spec/parser.md#14-tensor-and-consttensor-annotations). +See [parser §2.1](docs/spec/parser.md#21-syntax). """ from __future__ import annotations diff --git a/src/tilefoundry/evaluator/__init__.py b/src/tilefoundry/evaluator/__init__.py index 280ef3b7..4ed9c2bb 100644 --- a/src/tilefoundry/evaluator/__init__.py +++ b/src/tilefoundry/evaluator/__init__.py @@ -3,7 +3,7 @@ from typing import Any -from tilefoundry.evaluator.context import EvalContext +from tilefoundry.evaluator.context import EvalContext, FunctionEvalContext from tilefoundry.evaluator.registry import eval_registry, register_eval from tilefoundry.evaluator.value import ( EvalError, @@ -23,6 +23,7 @@ "TensorValue", "TupleValue", "EvalContext", + "FunctionEvalContext", "EvalError", "to_torch_dtype", "as_layout_view", diff --git a/src/tilefoundry/evaluator/context.py b/src/tilefoundry/evaluator/context.py index e078bf76..16752f01 100644 --- a/src/tilefoundry/evaluator/context.py +++ b/src/tilefoundry/evaluator/context.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import Any +from tilefoundry.visitor_registry.contexts import CallFeed + @dataclass(frozen=True) class EvalContext: @@ -18,3 +20,17 @@ class EvalContext: def __post_init__(self) -> None: if self.dim_bindings is None: object.__setattr__(self, "dim_bindings", {}) + + +@dataclass(frozen=True) +class FunctionEvalContext: + """Runtime state for one recursive function invocation.""" + + feed: CallFeed[Any] + loaded_module: Any | None = None + device: str = "cpu" + dim_bindings: dict[str, int] | None = None + + def __post_init__(self) -> None: + if self.dim_bindings is None: + object.__setattr__(self, "dim_bindings", {}) diff --git a/src/tilefoundry/evaluator/interpreter.py b/src/tilefoundry/evaluator/interpreter.py index 575f522b..ed0e03e5 100644 --- a/src/tilefoundry/evaluator/interpreter.py +++ b/src/tilefoundry/evaluator/interpreter.py @@ -8,7 +8,7 @@ import torch -from tilefoundry.evaluator.context import EvalContext +from tilefoundry.evaluator.context import EvalContext, FunctionEvalContext from tilefoundry.evaluator.dim import resolve_dim from tilefoundry.evaluator.registry import eval_registry from tilefoundry.evaluator.value import ( @@ -24,6 +24,7 @@ from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.visitor import ExprVisitor +from tilefoundry.visitor_registry.contexts import CallFeed def _default_device() -> str: @@ -56,21 +57,21 @@ def _bind_dim_vars(params, values) -> dict[str, int]: return binding -def child_reading(reading, callee: Function): - """The child reading a call to *callee* runs against, else ``None``. +def child_module_instance(loaded_module, callee: Function): + """The child module instance a call to *callee* runs against, else ``None``. ``None`` is a same-owner call, which binds every declared parameter. A collected call carries no binding record, so which reading supplies the constants is answered by which child owns the callee. """ - if reading is None or reading.module.owns(callee, derived=True): + if loaded_module is None or loaded_module.module.owns(callee, derived=True): return None matches = tuple( - child for child in reading.modules if child.module.owns(callee, derived=True) + child for child in loaded_module.modules if child.module.owns(callee, derived=True) ) if len(matches) > 1: raise EvalError( - f"evaluator: {reading.name!r} holds {len(matches)} readings owning " + f"evaluator: {loaded_module.name!r} holds {len(matches)} child modules owning " f"{callee.name!r}; one call reaches one child" ) return matches[0] if matches else None @@ -83,12 +84,14 @@ def __init__( self, env: dict[int, Value], device: str, dim_env: dict[str, int] | None = None, reading=None, + function_context: FunctionEvalContext | None = None, ) -> None: super().__init__() self.env = env self.device = device self.dim_env = dim_env or {} - self.reading = reading + self.function_context = function_context + self.reading = reading if function_context is None else function_context.loaded_module def visit_Var(self, var: Var) -> Value: try: return self.env[id(var)] @@ -126,7 +129,7 @@ def visit_Call(self, call: Call) -> Value: ) def _call_function(self, callee: Function, arg_exprs) -> Value: - child = child_reading(self.reading, callee) + child = child_module_instance(self.reading, callee) supplied = [p for p in callee.params if not (child is not None and p.is_const)] if len(arg_exprs) != len(supplied): kind = "activation(s)" if child is not None else "args" @@ -142,10 +145,20 @@ def _call_function(self, callee: Function, arg_exprs) -> Value: for param in callee.params ] target = _select_variant(callee, args) if callee.variants else callee - sub_env = {id(param): arg for param, arg in zip(target.params, args)} - sub_dim_env = _bind_dim_vars(target.params, args) + feed = CallFeed({id(param): arg for param, arg in zip(target.params, args)}) + function_context = FunctionEvalContext( + feed=feed, + loaded_module=child if child is not None else self.reading, + device=self.device, + dim_bindings=_bind_dim_vars(target.params, args), + ) + sub_env = dict(feed.by_param) return Evaluator( - sub_env, self.device, sub_dim_env, child if child is not None else self.reading + sub_env, + self.device, + function_context.dim_bindings, + function_context.loaded_module, + function_context=function_context, ).visit(target.body) def _resolve_loop_field(self, dim, what: str) -> int: @@ -217,17 +230,17 @@ def iter_env(i: int, carried) -> dict: return carried[0] if len(carried) == 1 else TupleValue(tuple(carried)) -def _child_constant(reading, callee: Function, param) -> TensorValue: +def _child_constant(loaded_module, callee: Function, param) -> TensorValue: """*param*'s constant, read from the child *callee* belongs to. Wrapped without a device argument: placement is settled before execution and this must not be where a weight quietly moves. """ try: - value = reading.constants[param.name] + value = loaded_module.constants[param.name] except KeyError: raise EvalError( - f"evaluator: {reading.name!r} has no binding for {param.name!r} of " + f"evaluator: {loaded_module.name!r} has no binding for {param.name!r} of " f"{callee.name!r}; a child call takes its ConstTensor parameters " f"from that child's own resources" ) from None diff --git a/src/tilefoundry/ir/core/context.py b/src/tilefoundry/ir/core/context.py index 0d2ee855..a058e4eb 100644 --- a/src/tilefoundry/ir/core/context.py +++ b/src/tilefoundry/ir/core/context.py @@ -2,6 +2,11 @@ from __future__ import annotations -from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext +from tilefoundry.visitor_registry.contexts import ( + CallFeed, + CallFeedProvider, + FunctionScope, + TypeInferContext, +) -__all__ = ["FunctionScope", "TypeInferContext"] +__all__ = ["CallFeed", "CallFeedProvider", "FunctionScope", "TypeInferContext"] diff --git a/src/tilefoundry/ir/core/module.py b/src/tilefoundry/ir/core/module.py index a8c15c31..a9030ba2 100644 --- a/src/tilefoundry/ir/core/module.py +++ b/src/tilefoundry/ir/core/module.py @@ -23,6 +23,7 @@ ModuleFunction = Union[HirFunction, PrimFunction] + _MISSING_PREPARED_WEIGHT = ( "[runtime §1.1.2](docs/spec/runtime.md#112-weight-converter-and-prepare--forward)" ) @@ -664,10 +665,7 @@ def _reached( is what lets a nested dispatch be resolved the way execution resolves it. """ from tilefoundry.evaluator.interpreter import ( # noqa: PLC0415 -- IR→evaluator - child_reading, - ) - from tilefoundry.ir.hir._call_binding import ( # noqa: PLC0415 -- cycle - bound_params, + child_module_instance, ) found: list[tuple[str, LoadedModule, object]] = [] @@ -681,8 +679,12 @@ def visit(path: str, reading: "LoadedModule", function, extents: dict) -> None: found.append((path, reading, function)) for call in _calls_in(function): declared = call.target - child = child_reading(reading, declared) - supplied = bound_params(declared, from_reading=child is not None) + child = child_module_instance(reading, declared) + supplied = tuple( + param + for param in declared.params + if child is None or not param.is_const + ) inner = _extended_dims( supplied, tuple(arg.type for arg in call.args), extents ) diff --git a/src/tilefoundry/ir/core/op_registry.py b/src/tilefoundry/ir/core/op_registry.py index 61359a03..10a0427d 100644 --- a/src/tilefoundry/ir/core/op_registry.py +++ b/src/tilefoundry/ir/core/op_registry.py @@ -16,7 +16,7 @@ def _register_schema(schema: "OpSchema", *, prepend: bool = False) -> None: Ops append; aliases prepend and therefore win first-match resolution. Concrete-class lookup skips aliases whose ``op_class`` is ``None``. - See [parser §4.3](docs/spec/parser.md#43-opschema-and-overload-resolution). + See [parser §3](docs/spec/parser.md#3-implementation-overview). """ key = (schema.dialect, schema.name) bucket = _schemas_by_dialect_name.setdefault(key, []) diff --git a/src/tilefoundry/ir/core/op_schema.py b/src/tilefoundry/ir/core/op_schema.py index d1f96c27..c95fb343 100644 --- a/src/tilefoundry/ir/core/op_schema.py +++ b/src/tilefoundry/ir/core/op_schema.py @@ -23,7 +23,7 @@ class OpSchema: node; aliases have no ``op_class`` and may build another operation type. ``category`` organizes documentation but is not part of the surface path. - See [parser §2.1](docs/spec/parser.md#21-model). + See [parser §2](docs/spec/parser.md#2-syntax-and-rules). """ name: str diff --git a/src/tilefoundry/parser/overload.py b/src/tilefoundry/ir/core/overload.py similarity index 100% rename from src/tilefoundry/parser/overload.py rename to src/tilefoundry/ir/core/overload.py diff --git a/src/tilefoundry/ir/core/param_def.py b/src/tilefoundry/ir/core/param_def.py index 4a49e351..a2c6796a 100644 --- a/src/tilefoundry/ir/core/param_def.py +++ b/src/tilefoundry/ir/core/param_def.py @@ -5,7 +5,7 @@ non-``MISSING`` default permits omission. ``__set_name__`` supplies the canonical parameter name. -See [parser §2.1](docs/spec/parser.md#21-model). +See [parser §2](docs/spec/parser.md#2-syntax-and-rules). """ from __future__ import annotations diff --git a/src/tilefoundry/ir/core/register.py b/src/tilefoundry/ir/core/register.py index ef53e163..6d94125e 100644 --- a/src/tilefoundry/ir/core/register.py +++ b/src/tilefoundry/ir/core/register.py @@ -4,7 +4,7 @@ them explicitly. Names default to the lowercase class name. Registration is the only route into the callable schema registry. -See [parser §2.1](docs/spec/parser.md#21-model). +See [parser §2](docs/spec/parser.md#2-syntax-and-rules). """ from __future__ import annotations diff --git a/src/tilefoundry/parser/static_eval.py b/src/tilefoundry/ir/core/static_eval.py similarity index 100% rename from src/tilefoundry/parser/static_eval.py rename to src/tilefoundry/ir/core/static_eval.py diff --git a/src/tilefoundry/ir/hir/_call_binding.py b/src/tilefoundry/ir/hir/_call_binding.py deleted file mode 100644 index 24294fb2..00000000 --- a/src/tilefoundry/ir/hir/_call_binding.py +++ /dev/null @@ -1,93 +0,0 @@ -"""How one call binds its arguments, resolved without asking a public question. - -A call into a child Module supplies activations alone and leaves the callee's -``ConstTensor`` parameters to that child's own reading. Which calls those are is -stated, never counted: while a class body is being authored the parser's own -record says so, and afterwards ownership within the walk's scope does. Both the -record and the ownership rule stay here, behind one result. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -from tilefoundry.visitor_registry.contexts import FunctionScope - - -def bound_params(callee, *, from_reading: bool) -> tuple: - """The parameters a call site supplies, in the order it supplies them.""" - if not from_reading: - return callee.params - return tuple(param for param in callee.params if not param.is_const) - - -@dataclass(frozen=True) -class CallBinding: - """What one call site supplies, and where its callee is read. - - ``params`` are the parameters its arguments bind to, in order. ``scope`` is - where the callee's body is read, which is what lets a call the callee makes - in turn be resolved the same way. ``from_reading`` says the parameters left - out come from a Module reading rather than from this call. - """ - - params: tuple - scope: FunctionScope | None - from_reading: bool - - -_authoring_reader: Any = None - - -def set_authoring_reader(reader) -> None: - """Install the reader for records written while a class body is authored. - - Called once by the parser that writes them. Without it nothing is authored, - so ownership is the only answer -- which is correct for a tree already built. - """ - global _authoring_reader - _authoring_reader = reader - - -def _authored_owner(call): - return None if _authoring_reader is None else _authoring_reader(call) - - -def _owned_child(ctx, callee): - """The child of the scope's function's owner that owns *callee*.""" - scope = getattr(ctx, "scope", None) - if scope is None or scope.module is None: - return None - from tilefoundry.ir.core.module import ( # noqa: PLC0415 — avoid import cycle - child_module_of, - ) - - return child_module_of(scope.module, scope.function, callee) - - -def binding_for(callee, call, ctx) -> CallBinding: - """How a call on *callee* binds its arguments in *ctx*. - - Fails closed: a call whose callee no single child of the caller's owner owns - binds every declared parameter, so a short argument list is refused rather - than reinterpreted. - """ - owner = None - if call is not None: - owner = _authored_owner(call) - if owner is None: - owner = _owned_child(ctx, callee) - if owner is not None: - return CallBinding( - bound_params(callee, from_reading=True), - FunctionScope(module=owner, function=callee), - True, - ) - scope = getattr(ctx, "scope", None) - tree = None if scope is None else scope.module - return CallBinding( - callee.params, - None if tree is None else FunctionScope(module=tree, function=callee), - False, - ) diff --git a/src/tilefoundry/ir/hir/function.py b/src/tilefoundry/ir/hir/function.py index 8e1019fa..8a205586 100644 --- a/src/tilefoundry/ir/hir/function.py +++ b/src/tilefoundry/ir/hir/function.py @@ -110,7 +110,6 @@ def seal(self) -> None: conv.seal() -from tilefoundry.ir.hir._call_binding import binding_for # noqa: E402 from tilefoundry.ir.visitor import ExprMutator # noqa: E402 @@ -179,48 +178,49 @@ def elaborate( arg_types: tuple[Type, ...], ctx: TypeInferContext | None = None, call: Call | None = None, + *, + feed=None, ) -> "Function": """Construct a concrete callee for one call site's argument types. Dispatch prototypes and already-equal bindings return unchanged. Other templates rebuild per distinct type tuple and reuse the construction session's elaboration cache. ``call`` anchors binding errors when present, - and is what says how the call binds: one whose callee's constants come from - a Module reading binds the non-constant parameters alone. + and is what says how the call binds: the context supplies a complete + ``CallFeed`` for the callee's formal parameters. See [hir §1.1](docs/spec/hir.md#11-function). """ if ctx is None: ctx = TypeInferContext() - binding = binding_for(callee, call, ctx) - expected = len(binding.params) - got = len(arg_types) - if got != expected: - kind = "activation(s)" if binding.from_reading else "parameter(s)" - ctx.error( - call if call is not None else callee, - f"hir Function call {callee.name!r}: arity mismatch — " - f"callee declares {expected} {kind}, call passed {got}", - ) - given = iter(enumerate(arg_types)) + feed = feed or ctx.build_call_feed(callee, arg_types) + supplied_params = ( + callee.params + if len(arg_types) == len(callee.params) + else tuple(param for param in callee.params if not param.is_const) + ) + supplied = iter(enumerate(arg_types)) bound_types = [] + supplied_ids = {id(param) for param in supplied_params} for param in callee.params: - if binding.from_reading and param.is_const: - bound_types.append(param.type) + if id(param) not in supplied_ids: + bound_types.append(feed.value_for(param)) continue - index, arg_ty = next(given) + index, arg_ty = next(supplied) bound_types.append(_bind_param_type(ctx, callee, index, param, arg_ty, call)) if callee.variants or callee.body is None: return callee if all(bt == p.type for bt, p in zip(bound_types, callee.params)): return callee - cache_key = (id(callee), arg_types, binding.from_reading) + cache_key = (id(callee), arg_types) cached = ctx.elaboration_cache.get(cache_key) if cached is not None: return cached - instance = _elaborate_from_bound_types(callee, bound_types, ctx, scope=binding.scope) + instance = _elaborate_from_bound_types( + callee, bound_types, ctx, scope=ctx.scope_for(callee) + ) ctx.elaboration_cache[cache_key] = instance return instance @@ -536,14 +536,13 @@ def _typeinfer_hir_function_call(call: Call, ctx) -> Type: """ callee: Function = call.target # type: ignore[assignment] arg_types = tuple(ctx.type_of(a) for a in call.args) - binding = binding_for(callee, call, ctx) - instance = elaborate(callee, arg_types, ctx, call=call) + feed = ctx.build_call_feed(callee, arg_types) + instance = elaborate(callee, arg_types, ctx, call=call, feed=feed) if instance.body is None: return instance.return_type - inner = binding.scope - body_ctx = TypeInferContext( - scope=None if inner is None else dataclasses.replace(inner, function=instance) - ) + body_ctx = ctx.child(callee, feed) + if body_ctx.scope is not None: + body_ctx.scope = dataclasses.replace(body_ctx.scope, function=instance) return body_ctx.type_of(instance.body) diff --git a/src/tilefoundry/ir/hir/verify.py b/src/tilefoundry/ir/hir/verify.py index 74cfb448..6752b462 100644 --- a/src/tilefoundry/ir/hir/verify.py +++ b/src/tilefoundry/ir/hir/verify.py @@ -1,6 +1,6 @@ from __future__ import annotations -from tilefoundry.ir.core import Expr, TypeInferContext, VerifyError +from tilefoundry.ir.core import Expr, FunctionScope, TypeInferContext, VerifyError from tilefoundry.ir.core.expr import Call, Var from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.tir.stmt import Stmt @@ -12,7 +12,7 @@ from .function import Function, canonical_specialization_signature -def verify_function(fn: Function) -> None: +def verify_function(fn: Function, *, module=None) -> None: """Verify HIR parameters, symbolic dimensions, body, and variants.""" for p in fn.params: if not isinstance(p, Var): @@ -24,7 +24,7 @@ def verify_function(fn: Function) -> None: f"hir Function {fn.name!r}: a function with variants must have " f"no body (a dispatch prototype's body is None / `pass`)" ) - _verify_variants(fn) + _verify_variants(fn, module=module) return if fn.body is None: return @@ -34,10 +34,11 @@ def verify_function(fn: Function) -> None: ) _reject_stmt_nodes(fn.body) - TypeInferContext().type_of(fn.body) + scope = FunctionScope(module, fn) if module is not None else None + TypeInferContext(scope=scope).type_of(fn.body) -def _verify_variants(base: Function) -> None: +def _verify_variants(base: Function, *, module=None) -> None: """Verify a dispatch prototype's variants and their envelope partition.""" base_param_types = tuple(p.type for p in base.params) sigs: dict[str, Function] = {} @@ -69,7 +70,7 @@ def _verify_variants(base: Function) -> None: f"hir Function {base.name!r}: duplicate variant canonical signature {sig!r}" ) sigs[sig] = v - verify_function(v) + verify_function(v, module=module) _verify_partition(base) diff --git a/src/tilefoundry/module.py b/src/tilefoundry/module.py index b2badff8..e3721f0d 100644 --- a/src/tilefoundry/module.py +++ b/src/tilefoundry/module.py @@ -56,93 +56,6 @@ def enclosing_declaration(frame: FrameType | None) -> _Entry | None: return None -def _retarget_module_calls(owner: str, functions, attached: dict) -> None: - """Rebuild each marked call against the child attached under its binding. - - Runs before ``Module`` construction seals the functions. The record is - repointed at the attached child before the rebuild reads it, because - attaching may have copied the Module the class body named. A binding the - class body does not attach is refused: there is no child to rebuild against, - and collecting it would leave the call pointing outside the tree being built. - """ - from tilefoundry.ir.core import ( # noqa: PLC0415 — avoid import cycle - Expr, - FunctionScope, - TypeInferContext, - get_metadata, - ) - from tilefoundry.ir.core.expr import ( # noqa: PLC0415 — avoid import cycle - Call, - ) - from tilefoundry.ir.hir.function import Function as HirFunction # noqa: PLC0415 - from tilefoundry.ir.hir.function import elaborate # noqa: PLC0415 - from tilefoundry.ir.visitor import ExprWalker # noqa: PLC0415 - from tilefoundry.parser.base import _ModuleCallee # noqa: PLC0415 - - unattached: list[str] = [] - - class _RetargetVisitor(ExprWalker[None]): - def visit(self, expr): - if expr is None or not isinstance(expr, Expr): - return None - return super().visit(expr) - - def visit_Call(self, expr: Call) -> None: - record = get_metadata(expr, _ModuleCallee) - if isinstance(expr.target, HirFunction) and record is None: - self.visit(expr.target) - elif isinstance(expr.target, HirFunction) and record.binding not in attached: - unattached.append(record.binding) - elif isinstance(expr.target, HirFunction): - child = attached[record.binding] - entry = child.entry_function() - object.__setattr__( - expr, - "metadata", - tuple( - _ModuleCallee(record.binding, child) - if isinstance(m, _ModuleCallee) - else m - for m in expr.metadata - ), - ) - object.__setattr__( - expr, - "target", - elaborate( - entry, - tuple(a.type for a in expr.args), - TypeInferContext(scope=FunctionScope(child, entry)), - call=expr, - ), - ) - object.__setattr__( - expr, - "metadata", - tuple(m for m in expr.metadata if not isinstance(m, _ModuleCallee)), - ) - self.visit_operands(expr) - - def visit_Function(self, fn) -> None: - self.visit_operands(fn) - for variant in fn.variants: - self.visit(variant) - for converter in fn.converters: - if isinstance(converter, tuple): - self.visit(converter[-1]) - - visitor = _RetargetVisitor() - for fn in functions: - visitor.visit(fn) - if unattached: - raise ValueError( - f"@module {owner!r}: call(s) to Module(s) {sorted(set(unattached))} that " - f"no class-body binding attaches; it binds {sorted(attached)}. A Module " - f"call is rebuilt against the child attached under the binding it names, " - f"so a name nothing binds has no child to call" - ) - - def _validate(topologies) -> tuple: from tilefoundry.ir.types.shard.mesh import Topology # noqa: PLC0415 @@ -161,7 +74,7 @@ def module( """Collect a class body into a ``Module``. Members may be DSL functions, child modules, or orchestration methods. See - [parser §2.7](docs/spec/parser.md#27-module-authoring-surface). + [parser §3](docs/spec/parser.md#3-implementation-overview). ``entry`` optionally names which collected function is the default step. @@ -179,6 +92,29 @@ def module( target_instance(target) resolved_target = target declared_topologies = None if topologies is UNDECLARED else _validate(topologies) + if cls is None: + from tilefoundry.parser.ast_pattern import create_module_context # noqa: PLC0415 + + owner_frame = sys._getframe(1) + context = create_module_context( + entry=entry, + target=resolved_target, + topologies=declared_topologies, + owner_frame=owner_frame, + source_filename=owner_frame.f_code.co_filename, + ) + + def _wrap_with_context(cls_inner): + try: + return context.finalize(cls_inner) + except Exception: + from tilefoundry.parser.ast_pattern import consume_module_context # noqa: PLC0415 + + consume_module_context(context) + raise + + return _wrap_with_context + mine = _Entry(declared_topologies, sys._getframe(1)) _DECLARING.append(mine) @@ -255,7 +191,6 @@ def _wrap(cls_inner): f"{mod_dupes} (a class-body alias of a nested @module is not " f"allowed; one name maps to one child module)" ) - _retarget_module_calls(cls_inner.__name__, functions, attached) if entry is not None and entry not in names: raise ValueError( f"@module {cls_inner.__name__!r}: entry {entry!r} names no " diff --git a/src/tilefoundry/parser/__init__.py b/src/tilefoundry/parser/__init__.py index dbcc9002..d441269b 100644 --- a/src/tilefoundry/parser/__init__.py +++ b/src/tilefoundry/parser/__init__.py @@ -1,6 +1,4 @@ -from __future__ import annotations +from .ast_pattern import FuncParserContext, FunctionRole, ParseError +from .parser_visitor import parse_function -from .hir_parser import parse_func -from .tir_parser import parse_prim_func - -__all__ = ["parse_func", "parse_prim_func"] +__all__ = ["parse_function", "FuncParserContext", "FunctionRole", "ParseError"] diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py new file mode 100644 index 00000000..610deda3 --- /dev/null +++ b/src/tilefoundry/parser/ast_pattern.py @@ -0,0 +1,1704 @@ +"""Executable AST patterns for the parser rewrite prototype. + +Patterns choose one local grammar branch. The resulting :class:`AstMatch` +constructs real TileFoundry values after its declared children have been +constructed, then applies the owning pattern's immutable rules in order. +""" + +# ruff: noqa: PLC0415, D202, E402, F403, F405 + +from __future__ import annotations + +import ast +import dataclasses +import operator +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from enum import StrEnum +from types import FrameType, SimpleNamespace +from typing import Any, ClassVar, Generic, Protocol, TypeVar +from typing import Literal as TypingLiteral + +from tilefoundry.ir.core import ( + BindingMetadata, + Call, + Constant, + ExecutionDomainMetadata, + Expr, + SourceSpanMetadata, + Var, + VerifyError, + replace_metadata, +) +from tilefoundry.ir.core.expr import Tuple as IrTuple +from tilefoundry.ir.core.kinds import BinaryKind, UnaryKind +from tilefoundry.ir.core.module import Module +from tilefoundry.ir.core.op_schema import OpSchema +from tilefoundry.ir.hir.function import Function, elaborate +from tilefoundry.ir.hir.grid_region import GridRegionExpr +from tilefoundry.ir.hir.math.binary import Binary +from tilefoundry.ir.hir.math.unary import Unary +from tilefoundry.ir.hir.sharding.local import Local +from tilefoundry.ir.hir.sharding.reshard import Reshard +from tilefoundry.ir.hir.specialize import DISPLAY_NAME +from tilefoundry.ir.hir.tensor.arange import Arange +from tilefoundry.ir.hir.tensor.reshape import Reshape +from tilefoundry.ir.hir.tensor.slice import Slice, slice_size +from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem +from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.tir.stmts import Evaluate, LetStmt, MeshScope, Return, Sequential +from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType +from tilefoundry.ir.types.dim import ( + DimAdd, + DimFloorDiv, + DimMod, + DimMul, + DimSub, + DimVar, + dim_expr, + simplify_dim, +) +from tilefoundry.ir.types.dim_isl import normalize_dim +from tilefoundry.ir.types.shard import ( + Broadcast, + Layout, + Mesh, + ShardLayout, + Split, + c_order_strides, + canonical_shard_layout, + composed, +) +from tilefoundry.ir.types.shard.layout import LayoutBase +from tilefoundry.ir.types.storage import StorageKind, resolve_storage +from tilefoundry.visitor_registry.contexts import CallFeed, FunctionScope, TypeInferContext +from tilefoundry.visitor_registry.visitors import TypeInferVisitor + +T = TypeVar("T") +_RETURN_TYPE = "" +_TYPE_INFER_CONTEXT = "" + + +@dataclass(frozen=True) +class PatternFailure: + """A recognized pattern whose nested validation failed.""" + + pattern_id: str + node: object + detail: str + causes: tuple["PatternFailure", ...] = () + + def render(self, indent: int = 0) -> str: + lines = [f"{' ' * indent}{self.pattern_id}: {self.detail}"] + for cause in self.causes: + lines.append(cause.render(indent + 1)) + return "\n".join(lines) + + +def _pattern_label(pattern: object) -> str: + label = getattr(pattern, "element_name", None) + if isinstance(label, str) and label: + return label + pattern_id = getattr(pattern, "pattern_id", None) + if isinstance(pattern_id, str) and pattern_id: + return pattern_id + node_type = getattr(pattern, "node_type", None) + if isinstance(node_type, type): + return node_type.__name__ + return type(pattern).__name__ + + +def _wrap_failure(pattern: object, node: object, failure: PatternFailure) -> PatternFailure: + return PatternFailure( + pattern_id=_pattern_label(pattern), + node=node, + detail="nested pattern failed", + causes=(failure,), + ) + + +def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContext") -> object: + """Attach source identity without copying lexical values such as parameter Vars.""" + if not isinstance(value, Expr): + return value + owner = context.function or context.module + source_filename = owner.source_filename if owner is not None else "" + line = getattr(node, "lineno", None) + column = getattr(node, "col_offset", None) + if isinstance(line, int) and isinstance(column, int): + value = replace_metadata( + value, + SourceSpanMetadata( + source_filename, + line, + column + 1, + getattr(node, "end_lineno", None), + getattr(node, "end_col_offset", None), + ), + ) + if context.binding_name: + value = replace_metadata(value, BindingMetadata(context.binding_name)) + return value + + +runtime = SimpleNamespace( + Call=Call, + Broadcast=Broadcast, + Binary=Binary, + BinaryKind=BinaryKind, + Constant=Constant, + DType=DType, + DimAdd=DimAdd, + DimFloorDiv=DimFloorDiv, + DimMod=DimMod, + DimMul=DimMul, + DimSub=DimSub, + DimVar=DimVar, + Evaluate=Evaluate, + Expr=Expr, + ExecutionDomainMetadata=ExecutionDomainMetadata, + Function=Function, + GridRegionExpr=GridRegionExpr, + Arange=Arange, + IrTuple=IrTuple, + Layout=Layout, + LayoutBase=LayoutBase, + Local=Local, + LetStmt=LetStmt, + MeshScope=MeshScope, + Mesh=Mesh, + Module=Module, + OpSchema=OpSchema, + PrimFunction=PrimFunction, + Reshard=Reshard, + Reshape=Reshape, + Return=Return, + Sequential=Sequential, + ShardLayout=ShardLayout, + Slice=Slice, + Split=Split, + StorageKind=StorageKind, + TensorType=TensorType, + TupleType=TupleType, + TypeInferContext=TypeInferContext, + FunctionScope=FunctionScope, + TypeInferVisitor=TypeInferVisitor, + TupleGetItem=TupleGetItem, + Unary=Unary, + UnaryKind=UnaryKind, + UnitType=UnitType, + Var=Var, + DISPLAY_NAME=DISPLAY_NAME, + c_order_strides=c_order_strides, + canonical_shard_layout=canonical_shard_layout, + composed=composed, + dim_expr=dim_expr, + elaborate=elaborate, + normalize_dim=normalize_dim, + slice_size=slice_size, + simplify_dim=simplify_dim, + resolve_storage=resolve_storage, + ) + + +class AstRule(Protocol[T]): + STATEMENT: ClassVar[str] + + def apply( + self, + value: T, + *, + match: AstMatch[T], + context: MatchContext, + ) -> T: ... + + +class AstPattern(Protocol[T]): + element_name: str | None + + def accept(self, visitor: PatternVisitor[Any]) -> Any: ... + + def match( + self, node: object, context: MatchContext + ) -> AstMatch[T] | PatternFailure | None: ... + + +class PatternVisitor(Protocol[T]): + def visit(self, pattern: AstPattern[Any]) -> T: ... + + +class CombinatorPattern(AstPattern[Any]): + """Shared base for executable AstPattern composition nodes.""" + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + element_name: str | None = None + + def accept(self, visitor: PatternVisitor[T]) -> T: + return visitor.visit(self) + + @staticmethod + def construct( + match: AstMatch[Any], children: Mapping[str, object], context: MatchContext + ) -> object: + if children: + return tuple(children.values()) + return match.captures.get("value") + + @staticmethod + def _merge( + owner: AstPattern[Any], + node: object, + matches: tuple[AstMatch[Any], ...], + *, + pattern_id: str, + branch_id: str, + ) -> AstMatch[Any]: + captures: dict[str, object] = {} + children: list[AstChild] = [] + construct_context = None + selected_branch = branch_id + selected_pattern_id = pattern_id + structural_ids = { + "sequence", + "field", + "repeat", + "optional", + "capture", + "child", + "predicate", + } + for matched in matches: + captures.update(matched.captures) + children.extend(matched.children) + construct_context = matched.construct_context or construct_context + if matched.branch_id not in structural_ids: + selected_branch = matched.branch_id + if matched.pattern_id not in structural_ids: + selected_pattern_id = matched.pattern_id + return AstMatch( + owner, + selected_pattern_id, + node, + captures, + selected_branch, + tuple(children), + construct_context, + ) + + +class ElementPattern(CombinatorPattern, Generic[T]): + """A named grammar production backed by one executable syntax graph.""" + + syntax: ClassVar[AstPattern[Any] | None] = None + + def match(self, node: object, context: MatchContext) -> AstMatch[T] | None: + syntax = type(self).syntax + if syntax is None: + raise TypeError(f"{type(self).__name__} has no executable syntax") + matched = syntax.match(node, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + if isinstance(matched.pattern, ElementPattern): + return matched + return AstMatch( + self, + matched.pattern_id, + node, + matched.captures, + matched.branch_id, + matched.children, + matched.construct_context, + ) + + +class LazyPattern(CombinatorPattern): + """Resolve one statically bound forward or recursive Pattern reference.""" + + def __init__(self, factory: Callable[[], AstPattern[Any]]): + self.factory = factory + self._resolved: AstPattern[Any] | None = None + + @property + def pattern(self) -> AstPattern[Any]: + if self._resolved is None: + self._resolved = self.factory() + return self._resolved + + def match( + self, node: object, context: MatchContext + ) -> AstMatch[Any] | PatternFailure | None: + matched = self.pattern.match(node, context) + return matched + + +class AstNodePattern(CombinatorPattern): + def __init__(self, node_type: type, *parts: AstPattern[Any]): + self.node_type = node_type + self.parts = tuple(parts) + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not isinstance(node, self.node_type): + return None + matches: list[AstMatch[Any]] = [] + for part in self.parts: + matched = part.match(node, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + matches.append(matched) + return self._merge( + self, + node, + tuple(matches), + pattern_id=self.node_type.__name__, + branch_id=self.node_type.__name__.lower(), + ) + + +class FieldPattern(CombinatorPattern): + def __init__(self, name: str, pattern: AstPattern[Any]): + self.name = name + self.pattern = pattern + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not hasattr(node, self.name): + return None + value = getattr(node, self.name) + matched = self.pattern.match(value, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + return AstMatch( + self, + matched.pattern_id, + node, + matched.captures, + matched.branch_id, + matched.children, + matched.construct_context, + ) + + +class LiteralPattern(CombinatorPattern): + def __init__( + self, + value: object = dataclasses.MISSING, + *, + value_type: type | tuple[type, ...] | None = None, + ): + self.value = value + self.value_type = value_type + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + raw = node.value if isinstance(node, ast.Constant) else node + if self.value is not dataclasses.MISSING and raw != self.value: + return None + if self.value_type is not None and not isinstance(raw, self.value_type): + return None + return AstMatch(self, "literal", node, {"value": raw}, "literal") + + +class ReferencePattern(CombinatorPattern): + def __init__( + self, *, resolve: bool = False, expected: type | tuple[type, ...] | None = None + ): + self.resolve = resolve + self.expected = expected + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not isinstance(node, (ast.Name, ast.Attribute)): + return None + captures: dict[str, object] = {} + if self.resolve: + try: + value = _resolve_reference(node, context) + except ParseError: + return None + if self.expected is not None and not isinstance(value, self.expected): + return None + captures["reference"] = value + return AstMatch(self, "reference", node, captures, "reference") + + +class SequencePattern(CombinatorPattern): + def __init__(self, *patterns: AstPattern[Any]): + self.patterns = tuple(patterns) + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not isinstance(node, (tuple, list)) or len(node) != len(self.patterns): + return None + matches: list[AstMatch[Any]] = [] + for value, pattern in zip(node, self.patterns): + matched = pattern.match(value, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + matches.append(matched) + return self._merge( + self, + node, + tuple(matches), + pattern_id="sequence", + branch_id="sequence", + ) + + +class ChoicePattern(CombinatorPattern): + def __init__(self, *patterns: AstPattern[Any]): + self.patterns = tuple(patterns) + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + failures: list[PatternFailure] = [] + for pattern in self.patterns: + matched = pattern.match(node, context) + if isinstance(matched, PatternFailure): + failures.append(matched) + continue + if matched is not None: + return matched + failures.append( + PatternFailure( + pattern_id=_pattern_label(pattern), + node=node, + detail="pattern did not match", + ) + ) + if failures: + return PatternFailure( + pattern_id="choice", + node=node, + detail="no choice matched", + causes=tuple(failures), + ) + return None + + +class ConditionPattern(CombinatorPattern): + """Run one sub-pattern only when an explicit context condition holds.""" + + def __init__( + self, + label: str, + test: Callable[[object, MatchContext], bool], + pattern: AstPattern[Any], + ): + self.label = label + self.test = test + self.pattern = pattern + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not self.test(node, context): + return None + return self.pattern.match(node, context) + + +class OptionalPattern(CombinatorPattern): + def __init__(self, pattern: AstPattern[Any]): + self.pattern = pattern + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if node is None: + return AstMatch(self, "optional", node, {}, "optional") + matched = self.pattern.match(node, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + return AstMatch( + self, + matched.pattern_id, + node, + matched.captures, + matched.branch_id, + matched.children, + matched.construct_context, + ) + + +class RepeatPattern(CombinatorPattern): + def __init__(self, pattern: AstPattern[Any], *, minimum: int = 0): + self.pattern = pattern + self.minimum = minimum + + @staticmethod + def _index_child(child: AstChild, index: int) -> AstChild: + return dataclasses.replace(child, name=child.name.format(index=index)) + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not isinstance(node, (tuple, list)) or len(node) < self.minimum: + return None + matches: list[AstMatch[Any]] = [] + for index, value in enumerate(node): + matched = self.pattern.match(value, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + matches.append( + dataclasses.replace( + matched, + children=tuple( + self._index_child(child, index) for child in matched.children + ), + ) + ) + return self._merge( + self, + node, + tuple(matches), + pattern_id="repeat", + branch_id="repeat", + ) + + +class PredicatePattern(CombinatorPattern): + def __init__(self, label: str, predicate: Callable[[object, MatchContext], bool]): + self.label = label + self.predicate = predicate + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + if not self.predicate(node, context): + return None + return AstMatch(self, "predicate", node, {}, "predicate") + + +class CapturePattern(CombinatorPattern): + def __init__(self, name: str, extractor: Callable[[object, MatchContext], object]): + self.name = name + self.extractor = extractor + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + try: + value = self.extractor(node, context) + except (AttributeError, KeyError, TypeError, ValueError): + return None + return AstMatch(self, "capture", node, {self.name: value}, "capture") + + +class ChildPattern(CombinatorPattern): + def __init__( + self, + name: str, + pattern: AstPattern[Any] | Callable[[], AstPattern[Any]], + situation: str, + role: str | None = None, + *, + expected_type: object | Callable[[object, MatchContext], object] | None = None, + values: Mapping[str, object] + | Callable[[object, MatchContext], Mapping[str, object]] + | None = None, + isolated_scope: bool = False, + transform: Callable[[object], object] | None = None, + ): + self.name = name + self._pattern = pattern + self.situation = situation + self.role = role + self.expected_type = expected_type + self.values = values + self.isolated_scope = isolated_scope + self.transform = transform + + @property + def pattern(self) -> AstPattern[Any]: + pattern = self._pattern + if callable(pattern): + pattern = pattern() + self._pattern = pattern + return pattern + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + value = self.transform(node) if self.transform is not None else node + if not isinstance(value, ast.AST): + return None + values = ( + self.values(value, context) if callable(self.values) else self.values or {} + ) + expected_type = ( + self.expected_type(value, context) + if callable(self.expected_type) + else self.expected_type + ) + child = AstChild( + self.name, + self.pattern, + value, + self.situation, + self.role, + expected_type=expected_type, + values=values, + isolated_scope=self.isolated_scope, + ) + return AstMatch(self, "child", node, {}, "child", (child,)) + + +class BranchPattern(CombinatorPattern): + def __init__( + self, branch_id: str, pattern: AstPattern[Any], *, pattern_id: str | None = None + ): + self.branch_id = branch_id + self.pattern = pattern + self.pattern_id = pattern_id or branch_id + + def match( + self, node: object, context: MatchContext + ) -> AstMatch[Any] | PatternFailure | None: + matched = self.pattern.match(node, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + return AstMatch( + self, + self.pattern_id, + node, + matched.captures, + self.branch_id, + matched.children, + matched.construct_context, + ) + + +class BindPattern(CombinatorPattern): + """Add semantic captures/children after an executable structural match.""" + + def __init__( + self, + pattern: AstPattern[Any], + binder: Callable[[object, MatchContext, AstMatch[Any]], AstMatch[Any] | None], + ): + self.pattern = pattern + self.binder = binder + + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | PatternFailure | None: + matched = self.pattern.match(node, context) + if isinstance(matched, PatternFailure): + return _wrap_failure(self, node, matched) + if matched is None: + return None + bound = self.binder(node, context, matched) + return bound + + +class LexicalScope: + """Parser-local lexical frames shared by sequential child construction.""" + + def __init__(self, frames: tuple[Mapping[str, object], ...] | None = None): + source = frames or ({},) + self._frames = [dict(frame) for frame in source] + + def define(self, name: str, value: object) -> None: + self._frames[-1][name] = value + + def lookup(self, name: str) -> object | None: + for frame in reversed(self._frames): + if name in frame: + return frame[name] + return None + + def fork(self) -> LexicalScope: + return LexicalScope(tuple(self._frames) + ({},)) + + def push_frame(self) -> None: + self._frames.append({}) + + def pop_frame(self) -> dict[str, object]: + if len(self._frames) == 1: + raise RuntimeError("cannot pop the root lexical frame") + return self._frames.pop() + + def items(self): + merged: dict[str, object] = {} + for frame in self._frames: + merged.update(frame) + return merged.items() + + +class FunctionRole(StrEnum): + ROOT = "root" + VARIANT = "variant" + CONVERTER = "converter" + + +@dataclass +class ParserState: + """Mutable state owned by one Function parse.""" + + mesh_stack: list[object] = field(default_factory=list) + mesh_coordinates: dict[tuple[int, int], object] = field(default_factory=dict) + + +@dataclass(frozen=True) +class LoopFrame: + kind: str + target: str + induction_var: object + start: object + extent: object + step: object + carry_names: tuple[str, ...] + phi_vars: tuple[object, ...] + init_args: tuple[object, ...] + + +@dataclass(frozen=True) +class FuncParserContext: + """Function-level authored inputs known before walking a FunctionDef.""" + + dialect: TypingLiteral["hir", "tir"] + role: FunctionRole = FunctionRole.ROOT + closure: Mapping[str, object] = field(default_factory=dict) + topologies: Mapping[str, object] = field(default_factory=dict) + source_filename: str = "" + module_scope: object | None = None + module: ModuleBuildContext | None = None + base: object | None = None + key: object | None = None + target: object | None = None + output_count: int = 1 + hardware_context: Mapping[str, object] = field(default_factory=dict) + binding_name: str | None = None + base_name: str | None = None + function_kind: str | None = field(default=None, repr=False) + state: ParserState = field(default_factory=ParserState) + + def __post_init__(self) -> None: + role = self.role + if self.function_kind is not None and role is FunctionRole.ROOT: + role = self.function_kind + if isinstance(role, str) and role not in {item.value for item in FunctionRole}: + role = ( + FunctionRole.ROOT + if role in {"func", "prim_func", "kernel"} + else FunctionRole(role) + ) + if role is not self.role: + object.__setattr__(self, "role", role) + if self.function_kind is None: + object.__setattr__( + self, + "function_kind", + "prim_func" + if role is FunctionRole.ROOT and self.dialect == "tir" + else "func" + if role is FunctionRole.ROOT + else role.value, + ) + + @property + def specializations(self) -> tuple[object, ...]: + return ( + () + if self.role is not FunctionRole.VARIANT or self.key is None + else (self.key,) + ) + + @property + def converter(self) -> object | None: + return self.key if self.role is FunctionRole.CONVERTER else None + + +@dataclass(frozen=True) +class ParserCallFeedProvider: + """Build call feeds from the authored module scope during parsing.""" + + module_scope: object | None + + def _child_for(self, callee: object): + if self.module_scope is None: + return None + for _name, child in self.module_scope.items(): + if getattr(child, "owns", lambda *_args, **_kwargs: False)( + callee, derived=True + ): + return child + return None + + def build_call_feed(self, callee: object, supplied: tuple[object, ...]) -> CallFeed: + child = self._child_for(callee) + params = tuple(p for p in callee.params if not (child is not None and p.is_const)) + if len(supplied) != len(params): + kind = "activation(s)" if child is not None else "parameter(s)" + raise VerifyError( + f"hir Function call {callee.name!r}: arity mismatch — " + f"callee declares {len(params)} {kind}, call passed {len(supplied)}" + ) + given = iter(supplied) + return CallFeed( + { + id(param): param.type if child is not None and param.is_const else next(given) + for param in callee.params + } + ) + + def scope_for(self, callee: object) -> FunctionScope | None: + child = self._child_for(callee) + return None if child is None else FunctionScope(child, callee) + + +@dataclass(frozen=True) +class ModuleFunctionValidationRule: + STATEMENT: ClassVar[str] = ( + "A module function must satisfy its root, variant, or converter role before mutation." + ) + + def apply( + self, + function: object, + *, + context: FuncParserContext, + module: ModuleBuildContext, + ) -> object: + module._validate_function(function, context) + return function + + +@dataclass(frozen=True) +class ModuleFunctionRegistrationRule: + STATEMENT: ClassVar[str] = ( + "A validated module function must be recorded in declaration order." + ) + + def apply( + self, + function: object, + *, + context: FuncParserContext, + module: ModuleBuildContext, + ) -> object: + module._commit_function(function, context) + return function + + +@dataclass(frozen=True) +class ModuleFinalizationRule: + STATEMENT: ClassVar[str] = ( + "A module declaration must contain valid unique members and a resolvable entry." + ) + + def apply(self, cls: type, *, module: ModuleBuildContext) -> object: + return module._finalize(cls) + + +@dataclass +class ModuleBuildContext: + """Authoring ledger for one Python ``@module`` class body. + + The context is deliberately independent from recursive ``MatchContext``. + It is created while the decorator expression is evaluated, looked up by + declaring frame during class-body execution, and consumed exactly once by + :meth:`finalize`. + """ + + owner_frame: FrameType + owner_name: str | None = None + closure: Mapping[str, object] = field(default_factory=dict) + source_filename: str = "" + module_scope: LexicalScope = field(default_factory=LexicalScope) + entry: str | None = None + target: object | None = None + topologies: tuple[object, ...] | None = None + roots: list[object] = field(default_factory=list) + bindings: dict[str, FunctionRole] = field(default_factory=dict) + binding_values: dict[str, object] = field(default_factory=dict) + owned: dict[int, FunctionRole] = field(default_factory=dict) + variant_keys: dict[int, set[object]] = field(default_factory=dict) + converter_keys: dict[int, set[str]] = field(default_factory=dict) + declarations: list[object] = field(default_factory=list) + _consumed: bool = False + FUNCTION_RULES: ClassVar[tuple[object, ...]] = ( + ModuleFunctionValidationRule(), + ModuleFunctionRegistrationRule(), + ) + FINALIZATION_RULES: ClassVar[tuple[object, ...]] = (ModuleFinalizationRule(),) + + def function_context( + self, + *, + dialect: TypingLiteral["hir", "tir"], + role: FunctionRole, + binding_name: str, + closure: Mapping[str, object] | None = None, + base: object | None = None, + key: object | None = None, + ) -> FuncParserContext: + topology_scope = { + getattr(topology, "name", str(index)): topology + for index, topology in enumerate(self.topologies or ()) + } + return FuncParserContext( + dialect=dialect, + role=role, + closure=closure if closure is not None else self.closure, + topologies=topology_scope, + source_filename=self.source_filename, + module_scope=self.module_scope, + module=self, + base=base, + key=key, + target=self.target if dialect == "tir" else None, + binding_name=binding_name, + ) + + @staticmethod + def _binding_error(role: FunctionRole, binding: str, owner: str | None) -> str: + return f"@module {owner or ''!r}: duplicate {role.value} binding {binding!r}" + + def validate_function(self, function: object, context: FuncParserContext) -> None: + self.FUNCTION_RULES[0].apply(function, context=context, module=self) + + def _validate_function(self, function: object, context: FuncParserContext) -> None: + role = context.role + binding = context.binding_name or getattr(function, "name", "") + if role is FunctionRole.ROOT: + if binding == "_": + raise ValueError( + f"@module {self.owner_name or ''!r}: a root binding may not be named '_'" + ) + if binding in self.bindings: + raise ValueError(self._binding_error(role, binding, self.owner_name)) + if any( + getattr(root, "name", None) == getattr(function, "name", None) + for root in self.roots + ): + raise ValueError( + self._binding_error( + role, getattr(function, "name", binding), self.owner_name + ) + ) + expected = ( + runtime.PrimFunction if context.dialect == "tir" else runtime.Function + ) + if not isinstance(function, expected): + raise TypeError( + f"root {binding!r} constructed {type(function).__name__}, expected {expected.__name__}" + ) + return + base = context.base + if base is None or not isinstance(base, runtime.Function): + raise ValueError(f"{role.value} {binding!r}: base is not a HIR Function") + if getattr(base, "_sealed", False): + raise RuntimeError( + f"base {base.name!r}: cannot register {role.value} after seal" + ) + if binding == "_" and role is FunctionRole.VARIANT: + raise ValueError( + f"base {base.name!r}: a variant binding may not be named '_'" + ) + if binding in self.bindings and role is not FunctionRole.CONVERTER: + raise ValueError(self._binding_error(role, binding, self.owner_name)) + if role is FunctionRole.VARIANT: + if getattr(function, "body", None) is None: + raise ValueError(f"base {base.name!r}: a variant must have a real body") + keys = self.variant_keys.setdefault(id(base), set()) + key = context.key + if key in keys: + raise ValueError( + f"base {base.name!r}: duplicate specialization key {key!r}" + ) + return + if getattr(function, "body", None) is None: + raise ValueError(f"base {base.name!r}: a converter must have a real body") + key = context.key + if not isinstance(key, str): + raise TypeError(f"base {base.name!r}: converter weight key must be str") + keys = self.converter_keys.setdefault(id(base), set()) + if key in keys: + raise ValueError(f"base {base.name!r}: duplicate converter weight {key!r}") + + def commit_function(self, function: object, context: FuncParserContext) -> None: + self.FUNCTION_RULES[1].apply(function, context=context, module=self) + + def _commit_function(self, function: object, context: FuncParserContext) -> None: + role = context.role + binding = context.binding_name or getattr(function, "name", "") + if role is FunctionRole.ROOT: + self.roots.append(function) + self.bindings[binding] = role + self.binding_values[binding] = function + return + base = context.base + assert base is not None + if role is FunctionRole.VARIANT: + base.add_variant(function) + self.variant_keys.setdefault(id(base), set()).add(context.key) + else: + base.add_converter(context.key, function) + self.converter_keys.setdefault(id(base), set()).add(context.key) + self.bindings[binding] = role + self.binding_values[binding] = function + self.owned[id(function)] = role + + def finalize(self, cls: type) -> object: + value: object = cls + for rule in self.FINALIZATION_RULES: + value = rule.apply(value, module=self) + return value + + def _finalize(self, cls: type) -> object: + if self._consumed: + raise RuntimeError( + f"@module {self.owner_name or cls.__name__!r}: declaration context already consumed" + ) + self._consumed = True + consume_module_context(self) + functions: list[object] = [] + modules: list[object] = [] + module_bindings: dict[str, object] = {} + methods: dict[str, object] = {} + for name, value in vars(cls).items(): + if name == "__call__": + raise TypeError( + f"@module {cls.__name__!r}: a class-body __call__ has no effect " + "-- name the method `forward` instead" + ) + if name.startswith("__") and name.endswith("__"): + continue + if isinstance(value, runtime.Module): + child = value if value.name == name else value.renamed(name) + modules.append(child) + module_bindings[name] = child + elif ( + isinstance(value, (tuple, list)) + and value + and all(isinstance(item, runtime.Module) for item in value) + ): + modules.extend(value) + for child in value: + module_bindings[child.name] = child + elif isinstance(value, (runtime.Function, runtime.PrimFunction)): + if id(value) in self.owned: + continue + functions.append(value) + elif getattr(value, "_tilefoundry_deferred", False): + continue + elif callable(value): + methods[name] = value + else: + raise TypeError( + f"@module {cls.__name__!r}: member {name!r} is a " + f"{type(value).__name__}, not an @func / @prim_func result, a " + "Module (or tuple/list of Modules), or a plain function; a " + "@module class body may contain only these three member kinds" + ) + for binding, role in self.bindings.items(): + value = vars(cls).get(binding) + expected = self.binding_values.get(binding) + if role is not FunctionRole.CONVERTER and value is not expected: + raise ValueError( + f"@module {cls.__name__!r}: registered {role.value} {binding!r} " + "was overwritten or aliased" + ) + for name, child in module_bindings.items(): + self.module_scope.define(name, child) + for declaration in self.declarations: + parsed = declaration.parse() + if declaration.role is FunctionRole.ROOT: + functions.append(parsed) + + names = [getattr(fn, "name", None) for fn in functions] + duplicates = sorted({name for name in names if names.count(name) > 1}) + if duplicates: + raise ValueError( + f"@module {cls.__name__!r}: duplicate function name(s) {duplicates} " + "(a class-body alias of a DSL function is not allowed)" + ) + if not functions and not modules and not methods: + raise TypeError( + f"@module {cls.__name__!r}: empty class body; declare a Function, " + "child Module, or orchestration method" + ) + if self.entry is not None and self.entry not in names: + raise ValueError( + f"@module {cls.__name__!r}: entry {self.entry!r} names no collected function (have {names})" + ) + result = runtime.Module( + name=cls.__name__, + functions=tuple(functions), + entry=self.entry, + modules=tuple(modules), + target=self.target, + topologies=self.topologies, + methods=methods, + ) + from tilefoundry.ir.hir.verify import verify_function # noqa: PLC0415 + from tilefoundry.ir.tir.verify import verify_prim_function # noqa: PLC0415 + + prim_functions = tuple( + function for function in functions if isinstance(function, runtime.PrimFunction) + ) + for function in functions: + if isinstance(function, runtime.Function): + verify_function(function, module=result) + elif isinstance(function, runtime.PrimFunction): + verify_prim_function(function, module_fns=prim_functions) + return result + + +_MODULE_CONTEXTS: list[ModuleBuildContext] = [] + + +def register_module_context(context: ModuleBuildContext) -> None: + _MODULE_CONTEXTS.append(context) + + +def create_module_context( + *, + entry: str | None = None, + target: object | None = None, + topologies: tuple[object, ...] | None = None, + closure: Mapping[str, object] | None = None, + owner_frame: FrameType | None = None, + owner_name: str | None = None, + source_filename: str = "", +) -> ModuleBuildContext: + """Create and register a context from a module decorator expression.""" + + frame = owner_frame or __import__("sys")._getframe(1) + context = ModuleBuildContext( + owner_frame=frame, + owner_name=owner_name, + closure=closure or {}, + source_filename=source_filename, + entry=entry, + target=target, + topologies=topologies, + ) + register_module_context(context) + return context + + +def module_context_for_frame(frame: FrameType | None) -> ModuleBuildContext | None: + while frame is not None: + if "__qualname__" in frame.f_locals: + for context in reversed(_MODULE_CONTEXTS): + if context.owner_frame is frame.f_back: + context.owner_name = frame.f_locals["__qualname__"].rsplit(".", 1)[-1] + return context + return None + if frame.f_code.co_name == "": + return None + frame = frame.f_back + return None + + +def consume_module_context(context: ModuleBuildContext) -> None: + try: + _MODULE_CONTEXTS.remove(context) + except ValueError: + pass + + +@dataclass(frozen=True) +class MatchContext: + """Inherited context for one recursive grammar position.""" + + function: FuncParserContext | None + module: ModuleBuildContext | None + situation: str + role: str | None = None + binding_name: str | None = None + expected_type: object | None = None + lexical_scope: LexicalScope = field(default_factory=LexicalScope) + parent: MatchContext | None = None + values: Mapping[str, object] = field(default_factory=dict) + + @classmethod + def from_function(cls, function: FuncParserContext) -> MatchContext: + scope = LexicalScope() + provider = ( + ParserCallFeedProvider(function.module_scope) + if function.module_scope is not None + else None + ) + scope.define( + _TYPE_INFER_CONTEXT, + runtime.TypeInferContext(call_feed_provider=provider), + ) + return cls( + function=function, + module=None, + situation="function", + role="function", + lexical_scope=scope, + values=function.hardware_context, + ) + + def child( + self, + *, + situation: str, + role: str | None = None, + binding_name: str | None = None, + expected_type: object | None = None, + values: Mapping[str, object] | None = None, + isolated_scope: bool = False, + function: FuncParserContext | None = None, + module: ModuleBuildContext | None = None, + ) -> MatchContext: + merged = dict(self.values) + if values: + merged.update(values) + switching_function = function is not None and function is not self.function + if switching_function: + scope = LexicalScope() + provider = ( + ParserCallFeedProvider(function.module_scope) + if function is not None and function.module_scope is not None + else None + ) + scope.define( + _TYPE_INFER_CONTEXT, + runtime.TypeInferContext(call_feed_provider=provider), + ) + else: + scope = self.lexical_scope.fork() if isolated_scope else self.lexical_scope + if expected_type is None and role == "return_value": + expected_type = scope.lookup(_RETURN_TYPE) + return MatchContext( + function=function or self.function, + module=module or self.module, + situation=situation, + role=role, + binding_name=binding_name, + expected_type=expected_type, + lexical_scope=scope, + parent=self, + values=merged, + ) + + def resolve_lexical(self, name: str) -> object: + value = self.lexical_scope.lookup(name) + if value is None: + raise ParseError.from_node( + ast.Name(id=name, ctx=ast.Load()), + self, + f"undefined lexical name {name!r}", + ) + return value + + def resolve_static( + self, node: ast.AST, expected: type[T] | tuple[type[Any], ...] + ) -> T: + if not isinstance(node, (ast.Name, ast.Attribute)): + raise ParseError.from_node( + node, self, "static references use Name or Attribute" + ) + value = _resolve_reference(node, self) + if not isinstance(value, expected): + raise ParseError.from_node( + node, + self, + f"static reference resolved to {type(value).__name__}, " + f"expected {_expected_name(expected)}", + ) + return value + + +@dataclass(frozen=True) +class AstChild: + name: str + pattern: AstPattern[Any] + node: ast.AST + situation: str + role: str | None = None + expected_type: object | None = None + values: Mapping[str, object] = field(default_factory=dict) + isolated_scope: bool = False + function_context: FuncParserContext | None = None + module_context: ModuleBuildContext | None = None + + +@dataclass(frozen=True) +class AstMatch(Generic[T]): + pattern: AstPattern[T] + pattern_id: str + node: ast.AST + captures: Mapping[str, object] + branch_id: str + children: tuple[AstChild, ...] = () + construct_context: MatchContext | None = None + + def construct(self, children: Mapping[str, object], context: MatchContext) -> T: + pattern_constructor = getattr(self.pattern, "construct", None) + if not callable(pattern_constructor): + raise TypeError( + f"Pattern {type(self.pattern).__name__} has no construct method" + ) + value = pattern_constructor(self, children, context) + for rule in self.pattern.RULES: + value = rule.apply(value, match=self, context=context) + return value + + +class ParseError(VerifyError): + """A grammar or context rule failure anchored to an authored AST node.""" + + def __init__( + self, + *, + node: ast.AST, + context: MatchContext, + detail: str | None = None, + ): + line = getattr(node, "lineno", None) + column = getattr(node, "col_offset", None) + location = "" + owner = context.function or context.module + source_filename = owner.source_filename if owner is not None else "" + if isinstance(line, int): + location = f" at {source_filename}:{line}" + if isinstance(column, int): + location += f":{column + 1}" + message = detail or f"no AST pattern matched situation {context.situation!r}" + if context.role: + message += f" (role {context.role!r})" + super().__init__(message + location) + self.node = node + self.context = context + self.detail = detail + + @classmethod + def from_node( + cls, node: ast.AST, context: MatchContext, detail: str | None = None + ) -> ParseError: + return cls(node=node, context=context, detail=detail) + + +def parse_node(pattern: AstPattern[T], node: ast.AST, context: MatchContext) -> T: + """Select, recursively construct, and apply rules for one local pattern.""" + + matched = pattern.match(node, context) + if isinstance(matched, PatternFailure): + raise ParseError.from_node(node, context, matched.render()) + if matched is None: + raise ParseError.from_node(node, context) + active_context = matched.construct_context or context + children: dict[str, object] = {} + for child in matched.children: + if child.name in children: + raise RuntimeError(f"duplicate AstChild name {child.name!r}") + binding_name = None + if child.name == "value": + binding_name = matched.captures.get("name") + if binding_name is None: + names = matched.captures.get("names") + if isinstance(names, tuple): + binding_name = ", ".join(names) + child_context = active_context.child( + situation=child.situation, + role=child.role, + binding_name=binding_name, + expected_type=child.expected_type, + values=child.values, + isolated_scope=child.isolated_scope, + function=child.function_context, + module=child.module_context, + ) + children[child.name] = parse_node(child.pattern, child.node, child_context) + return matched.construct(children, active_context) + + +def _expected_name(expected: type | tuple[type, ...]) -> str: + if isinstance(expected, tuple): + return " | ".join(item.__name__ for item in expected) + return expected.__name__ + + +def _decorator_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + +def _resolve_reference(node: ast.AST, context: MatchContext) -> object: + if isinstance(node, ast.Name): + lexical = context.lexical_scope.lookup(node.id) + if lexical is not None: + return lexical + function = context.function + module = context.module + module_scope = ( + function.module_scope + if function is not None + else module.module_scope + if module is not None + else None + ) + if isinstance(module_scope, Mapping) and node.id in module_scope: + return module_scope[node.id] + lookup = getattr(module_scope, "lookup", None) + if callable(lookup): + try: + value = lookup(node.id) + except (KeyError, ValueError): + pass + else: + if value is not None: + return value + closure = ( + function.closure + if function is not None + else module.closure + if module is not None + else {} + ) + if node.id in closure: + return closure[node.id] + raise ParseError.from_node(node, context, f"undefined static name {node.id!r}") + if isinstance(node, ast.Attribute): + owner = _resolve_reference(node.value, context) + try: + return getattr(owner, node.attr) + except AttributeError as error: + raise ParseError.from_node( + node, context, f"{type(owner).__name__} has no attribute {node.attr!r}" + ) from error + raise ParseError.from_node(node, context, "expected static Name or Attribute") + + +_BINARY_OPERATORS: Mapping[type[ast.operator], Callable[[Any, Any], Any]] = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, + ast.Mod: operator.mod, + ast.Pow: operator.pow, +} +_UNARY_OPERATORS: Mapping[type[ast.unaryop], Callable[[Any], Any]] = { + ast.UAdd: operator.pos, + ast.USub: operator.neg, + ast.Not: operator.not_, +} + + +@dataclass(frozen=True) +class CanonicalDTypeRule: + STATEMENT: ClassVar[str] = "A dtype must resolve to a canonical DType." + + def apply(self, value, *, match, context): + if not isinstance(value, runtime.DType): + raise ParseError.from_node( + match.node, context, "dtype did not construct DType" + ) + canonical = runtime.DType._members().get(value.name) + if canonical is not value: + raise ParseError.from_node( + match.node, context, f"non-canonical dtype {value.name!r}" + ) + return value + + +@dataclass(frozen=True) +class LayoutShapeRule: + STATEMENT: ClassVar[str] = "A layout must have a valid non-boolean shape." + + def apply(self, value, *, match, context): + if value is not None and not isinstance(value, runtime.LayoutBase): + raise ParseError.from_node( + match.node, context, "layout did not construct LayoutBase" + ) + if value is not None and not isinstance(value.shape, tuple): + raise ParseError.from_node( + match.node, context, "layout shape is not a tuple" + ) + return value + + +@dataclass(frozen=True) +class LayoutPositionRule: + STATEMENT: ClassVar[str] = "A layout must be legal for its parser position." + + def apply(self, value, *, match, context): + if context.role in {"storage", "dtype", "shape"}: + raise ParseError.from_node( + match.node, context, "layout used in a non-layout role" + ) + return value + + +@dataclass(frozen=True) +class StorageValueRule: + STATEMENT: ClassVar[str] = "Storage must resolve to a StorageKind." + + def apply(self, value, *, match, context): + if not isinstance(value, runtime.StorageKind): + raise ParseError.from_node( + match.node, context, "storage did not construct StorageKind" + ) + return value + + +@dataclass(frozen=True) +class TensorLayoutStorageRule: + STATEMENT: ClassVar[str] = ( + "A tensor type must contain compatible layout and storage values." + ) + + def apply(self, value, *, match, context): + if not isinstance(value, runtime.TensorType): + raise ParseError.from_node( + match.node, context, "tensor did not construct TensorType" + ) + if not isinstance(value.storage, runtime.StorageKind): + raise ParseError.from_node( + match.node, context, "tensor storage is not StorageKind" + ) + if value.layout is not None and not isinstance( + value.layout, runtime.LayoutBase + ): + raise ParseError.from_node( + match.node, context, "tensor layout is not LayoutBase" + ) + return value + + +@dataclass(frozen=True) +class TensorPositionRule: + STATEMENT: ClassVar[str] = ( + "A tensor type's storage must be legal for its dialect and position." + ) + + def apply(self, value, *, match, context): + if context.role == "storage": + raise ParseError.from_node( + match.node, context, "TensorType used in a storage role" + ) + allowed = context.values.get("allowed_storage") + if allowed is not None and value.storage not in allowed: + rendered = tuple(str(item) for item in allowed) + raise ParseError.from_node( + match.node, + context, + f"storage {value.storage} is not allowed by hardware context {rendered}", + ) + return value + + +@dataclass(frozen=True) +class ShapeDimRule: + STATEMENT: ClassVar[str] = ( + "A shape dimension must be an integer, DimVar, or expression." + ) + + def apply(self, value, *, match, context): + if isinstance(value, bool) or not isinstance( + value, (int, runtime.DimVar, runtime.Expr) + ): + raise ParseError.from_node( + match.node, + context, + f"shape dimension must be int, DimVar, or Expr, got {type(value).__name__}", + ) + return value + + +@dataclass(frozen=True) +class ShapeTupleRule: + STATEMENT: ClassVar[str] = "A shape must construct a tuple of dimensions." + + def apply(self, value, *, match, context): + if not isinstance(value, tuple): + raise ParseError.from_node(match.node, context, "shape is not a tuple") + return value + + +__all__ = [ + "AstChild", + "AstMatch", + "AstPattern", + "AstRule", + "BlockPattern", + "CallPattern", + "ConstantPattern", + "DTypePattern", + "FuncParserContext", + "FunctionPattern", + "FunctionRole", + "LayoutPattern", + "LexicalScope", + "MatchContext", + "ModuleBuildContext", + "NamePattern", + "ParseError", + "PatternFailure", + "RenderVisitor", + "ScalarTypePattern", + "ShapePattern", + "SignaturePattern", + "StatementPattern", + "StoragePattern", + "TensorOptionalSlotPattern", + "TensorPattern", + "TypeAnnotationPattern", + "consume_module_context", + "create_module_context", + "module_context_for_frame", + "parse_node", + "register_module_context", + "render_grammar", +] + + +def _constant(value): + if isinstance(value, bool): + dtype = runtime.DType.bool + elif isinstance(value, int): + dtype = runtime.DType.i64 + elif isinstance(value, float): + dtype = runtime.DType.f32 + else: + raise TypeError(type(value).__name__) + return runtime.Constant( + type=runtime.TensorType.scalar(dtype, storage=runtime.StorageKind.UMAT), + value=value, + ) + + +def _infer_call(operation, args, context): + placeholder_type = getattr(operation, "type", None) + if placeholder_type is None and args: + placeholder_type = args[0].type + if placeholder_type is None: + placeholder_type = runtime.TensorType.scalar(runtime.DType.f32) + metadata = () + if context.function is not None and context.function.state.mesh_stack: + metadata = ( + runtime.ExecutionDomainMetadata(tuple(context.function.state.mesh_stack)), + ) + placeholder = runtime.Call( + type=placeholder_type, target=operation, args=tuple(args), metadata=metadata + ) + infer_context = context.lexical_scope.lookup(_TYPE_INFER_CONTEXT) + if not isinstance(infer_context, runtime.TypeInferContext): + infer_context = runtime.TypeInferContext() + inferred = runtime.TypeInferVisitor(infer_context).visit(placeholder) + return dataclasses.replace(placeholder, type=inferred) + + +def _slice_size(begin, end, stride, context, node): + try: + value = runtime.normalize_dim( + runtime.slice_size( + runtime.dim_expr(begin), + runtime.dim_expr(end), + runtime.dim_expr(stride), + ) + ) + except (TypeError, ValueError) as error: + raise ParseError.from_node(node, context, str(error)) from error + if isinstance(value, runtime.Constant): + return value.value + return value + + +from .grammar_render import RenderVisitor, render_grammar +from .pattern_nodes import * diff --git a/src/tilefoundry/parser/base.py b/src/tilefoundry/parser/base.py deleted file mode 100644 index a41145d1..00000000 --- a/src/tilefoundry/parser/base.py +++ /dev/null @@ -1,1472 +0,0 @@ -from __future__ import annotations - -import ast -import dataclasses -import enum -import inspect -import logging -import textwrap -from typing import Any, Callable - -from tilefoundry.ir.core import ( - BindingMetadata, - Call, - Constant, - ExecutionDomainMetadata, - Expr, - IRMetadata, - SourceSpanMetadata, - Tuple, - TypeInferContext, - Var, - VerifyError, - get_metadata, - replace_metadata, -) -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.core.op_schema import OpSchema -from tilefoundry.ir.hir._call_binding import bound_params, set_authoring_reader -from tilefoundry.ir.hir.function import Function as HirFunction -from tilefoundry.ir.hir.function import elaborate -from tilefoundry.ir.hir.math.binary import Binary -from tilefoundry.ir.hir.math.unary import Unary -from tilefoundry.ir.hir.tensor.reshape import Reshape -from tilefoundry.ir.hir.tensor.slice import Slice -from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem -from tilefoundry.ir.types import DType, TensorType, TupleType -from tilefoundry.ir.types.dim import DimAdd, dim_expr, is_dim_expr, simplify_dim -from tilefoundry.ir.types.dim_isl import normalize_dim -from tilefoundry.ir.types.dtype import FloatDType -from tilefoundry.ir.types.shape_helpers import i64_const -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, shard_layout_of -from tilefoundry.ir.types.storage import StorageKind, resolve_storage -from tilefoundry.ir.types.substitute import canonicalize_dims -from tilefoundry.visitor_registry.visitors import TypeInferVisitor - -from .dispatch import ( - Token, - _binary_kind_for_ast_op, - _unary_kind_for_ast_op, - resolve_callable, - resolve_op, - resolve_schema, - resolve_stmt, -) -from .static_eval import eval_static -from .sugar import ( - LayoutSugarError, - _is_tuple_sugar, - parse_sugar, -) -from .symtab import LexicalEnv - -logger = logging.getLogger(__name__) - - -@dataclasses.dataclass(frozen=True) -class _ModuleCallee(IRMetadata): - """The class-body binding a call reached its callee's Module through. - - Authoring state private to the parser: a class body is parsed before its - children are attached, and attaching one copies it, so the binding name is - what says which attached child the call meant. ``@module`` collection - rebuilds against that child and takes this record off. - """ - - binding: str - owner: Module - - -def _authored_child_call(call): - """The child Module *call* was written through, else ``None``. - - The authoring phase's answer to which calls carry activations only. It holds - until ``@module`` collection consumes the record, and asks nothing of the - walk: the record is on the call site itself. - """ - record = get_metadata(call, _ModuleCallee) - return None if record is None else record.owner - - -set_authoring_reader(_authored_child_call) - - -_IR_OBJECT_TYPES = { - "Topology": None, - "Mesh": None, - "ShardLayout": None, - "Layout": None, -} - - -def _warn_if_ir_object(val: Any, name: str) -> None: - """Warn when a preconstructed IR object is resolved from closure. - - Canonical DSL source should use AST constructor syntax instead of capturing - prebuilt Python objects in the closure. - """ - type_name = type(val).__name__ - if type_name in _IR_OBJECT_TYPES: - logger.warning( - "Closure-captured IR object %r of type %s — " - "this is not canonical. Prefer declaring in DSL source or " - "using topology-name string resolution.", - name, type_name, - ) - - -def extract_ast(fn) -> ast.FunctionDef: - source_lines, start_line = inspect.getsourcelines(fn) - src = textwrap.dedent("".join(source_lines)) - mod = ast.parse(src) - ast.increment_lineno(mod, start_line - 1) - - - for node in ast.walk(mod): - if isinstance(node, ast.FunctionDef): - return node - raise VerifyError("cannot locate FunctionDef in source") - - -def _collect_closure(fn, extra: dict[str, Any] | None = None) -> dict[str, Any]: - """Collect a live Python function's name-resolution namespace. - - Shared by ``parse_func`` (HIR) and ``parse_prim_func`` (TIR): ``extra`` - (sibling ``@func`` / ``@prim_func`` bindings from a ``@module`` class - body's definition frame) sits below the function's own globals / - freevars so it cannot shadow them. - """ - closure: dict[str, Any] = {} - if extra: - closure.update(extra) - if fn.__globals__ is not None: - closure.update(fn.__globals__) - if fn.__closure__ is not None: - for name, cell in zip(fn.__code__.co_freevars, fn.__closure__): - try: - closure[name] = cell.cell_contents - except ValueError: - pass - return closure - - -def _annotation_head_name(node: ast.AST) -> str | None: - """Annotation head name. - - Return the subscript base identifier (``Tensor`` / ``ConstTensor``), - resolving through an attribute path such as ``dsl.ConstTensor``. - """ - if isinstance(node, ast.Name): - return node.id - if isinstance(node, ast.Attribute): - return node.attr - return None - - -def _is_const_tensor_annotation(node: ast.AST) -> bool: - """``ConstTensor[...]`` marks a parameter ``is_const=True``. - - ``ConstTensor[...]`` marks a parameter ``is_const=True``; ``Tensor[...]`` - and every other annotation form leave it ``False``. - """ - return ( - isinstance(node, ast.Subscript) - and _annotation_head_name(node.value) == "ConstTensor" - ) - - -def _resolve_tensor_type( - node: ast.AST, - closure: dict[str, Any], - *, - mesh_resolver=None, - default_mesh: Mesh | None = None, - mesh_order: "tuple[Mesh, ...]" = (), - name_resolver=None, -) -> TensorType: - """Resolve tensor annotations identically for HIR and TIR functions. - - Parse compact layout sugar directly from AST or evaluate the verbose shard - layout form in *closure*. See - [parser §1.4](docs/spec/parser.md#14-tensor-and-consttensor-annotations) and - [parser §1.5](docs/spec/parser.md#15-layout-sugar). - """ - bindings = dict(closure) - if name_resolver is not None: - for candidate in ast.walk(node): - if isinstance(candidate, ast.Name): - value = name_resolver(candidate.id) - if value is not None: - bindings[candidate.id] = value - result = parse_sugar( - node, - TensorType, - closure=bindings, - mesh_resolver=mesh_resolver, - default_mesh=default_mesh, - mesh_order=mesh_order, - ) - if result is not None: - return canonicalize_dims(result) - try: - code = compile(ast.Expression(body=node), "", "eval") - val = eval(code, closure) # noqa: S307 — controlled internal eval - except Exception as exc: - raise VerifyError(f"failed to resolve type annotation: {exc}") - if isinstance(val, TensorType): - return canonicalize_dims(val) - raise VerifyError(f"annotation did not resolve to TensorType, got {type(val).__name__}") - - -def _build_params( - node: ast.FunctionDef, - closure: dict[str, Any], - resolve_annotation: Callable[[ast.AST, dict[str, Any]], TensorType], - *, - decorator_name: str, -) -> tuple[Var, ...]: - """Build ``Var`` parameters from a function's AST arg annotations. - - Shared by ``parse_func`` (HIR) and ``parse_prim_func`` (TIR); both pass - :func:`_resolve_tensor_type` as *resolve_annotation* so a ``Tensor[...]`` - layout-sugar annotation works identically on ``@func`` and ``@prim_func`` - params. - """ - out: list[Var] = [] - for a in node.args.args: - if a.annotation is None: - raise VerifyError(f"{decorator_name} param {a.arg!r} must be annotated") - ann_type = resolve_annotation(a.annotation, closure) - is_const = _is_const_tensor_annotation(a.annotation) - out.append(Var(type=ann_type, name=a.arg, is_const=is_const)) - return tuple(out) - - -def _i64(value: int) -> Constant: - return i64_const(value) - - -def _constant_from_py(value: Any) -> Constant: - - - if isinstance(value, bool): - return Constant(type=TensorType.scalar(DType.bool, storage=StorageKind.UMAT), value=value) - if isinstance(value, int): - return Constant(type=TensorType.scalar(DType.i64, storage=StorageKind.UMAT), value=value) - if isinstance(value, float): - return Constant(type=TensorType.scalar(DType.f32, storage=StorageKind.UMAT), value=value) - raise VerifyError(f"unsupported literal type {type(value).__name__}") - - - - -_STATIC_ARITH_NODES: tuple[type, ...] = ( - ast.Constant, ast.Name, ast.Attribute, ast.UnaryOp, ast.BinOp, -) - - -def _is_python_float_scalar(expr: Expr) -> bool: - """Whether *expr* is a Python float scalar, which carries no precision.""" - ty = expr.type - return ( - isinstance(expr, Constant) - and isinstance(ty, TensorType) - and ty.shape == () - and ty.storage is StorageKind.UMAT - and isinstance(ty.dtype, FloatDType) - ) - - -def _with_python_float_dtypes(args: tuple[Expr, ...]) -> tuple[Expr, ...]: - """Give each Python float scalar the float dtype its fellow operands carry. - - Applies to floats only; a Python integer keeps its own dtype. - """ - floats = {i for i, arg in enumerate(args) if _is_python_float_scalar(arg)} - if not floats: - return args - others = { - arg.type.dtype - for i, arg in enumerate(args) - if i not in floats - and isinstance(arg.type, TensorType) - and isinstance(arg.type.dtype, FloatDType) - } - if len(others) != 1: - return args - dtype = next(iter(others)) - out = list(args) - for i in floats: - if out[i].type.dtype != dtype: - out[i] = dataclasses.replace( - out[i], type=dataclasses.replace(out[i].type, dtype=dtype) - ) - return tuple(out) - - -class BaseExprVisitor: - """Shared visitor for Expr-returning AST nodes. Emits core_ir Expr.""" - - token: Token - - resolves_module_callees = False - - def __init__( - self, env: LexicalEnv, closure: dict[str, Any], *, in_module_body: bool = False - ): - self.env = env - self.closure = closure - self.in_module_body = in_module_body - - - - self._ctx = TypeInferContext() - - - - self._explicit_binding_call_ids: set[int] = set() - - - - - self._call_dsl_names: dict[int, str] = {} - self._scalar_index_ids: set[int] = set() - - - - self._tile_windows: dict[int, tuple[Any, Any]] = {} - self._active_source_node: ast.AST | None = None - self._active_binding_hint: str | None = None - self._mesh_scopes: tuple[Any, ...] = () - self.source_filename = "" - - def _tuple_expr_expr(self, node: ast.Tuple): - """Build a ``Tuple`` from an AST tuple literal.""" - elements = tuple(self.expr(e) for e in node.elts) - field_types = tuple(e.type for e in elements) - return Tuple(type=TupleType(fields=field_types), elements=elements) - - def _resolve_body_mesh(self, name: str): - """Resolve a mesh by variable name from the lexical env only. - - Body sugar (``reshard(layout=(... @ mesh.axis, ...))``) must use - meshes from lexical ``with Mesh(...) as name`` scopes. Closure / - global mesh IR objects are NOT accepted for body sugar. - """ - val = self.env.lookup(name) - if isinstance(val, Mesh): - return val - return None - - def _current_default_mesh(self): - """Return the innermost Mesh from the lexical scope, or None. - - Used as the *default_mesh* for all-Broadcast ShardLayout sugar. - """ - return self.env.innermost_mesh() - - def _contains_mesh_coordinate(self, node: ast.AST) -> bool: - """Whether *node* contains a dialect-specific mesh coordinate.""" - return False - - - - def expr(self, node: ast.AST) -> Expr: - method = getattr(self, f"visit_{type(node).__name__}", None) - if method is None: - raise VerifyError(f"unsupported AST node in expression: {type(node).__name__}") - previous = self._active_source_node - self._active_source_node = node - try: - return method(node) - finally: - self._active_source_node = previous - - def expr_with_binding(self, node: ast.AST, name: str) -> Expr: - """Parse one RHS while making its authored LHS available to errors.""" - previous = self._active_binding_hint - self._active_binding_hint = name - try: - return self.expr(node) - finally: - self._active_binding_hint = previous - - @staticmethod - def _attach_metadata(expr: Expr, value: IRMetadata) -> None: - """Attach parser-authored metadata without rebuilding the SSA node.""" - kept = tuple(item for item in expr.metadata if type(item) is not type(value)) - object.__setattr__(expr, "metadata", (*kept, value)) - - def _source_span(self) -> SourceSpanMetadata | None: - node = self._active_source_node - if node is None or not hasattr(node, "lineno"): - return None - return SourceSpanMetadata( - file=self.source_filename, - line=node.lineno, - - column=node.col_offset + 1, - end_line=getattr(node, "end_lineno", None), - end_column=( - getattr(node, "end_col_offset", None) + 1 - if getattr(node, "end_col_offset", None) is not None - else None - ), - ) - - def _source_metadata(self) -> tuple: - metadata = [] - span = self._source_span() - if span is not None: - metadata.append(span) - if self._active_binding_hint is not None: - metadata.append(BindingMetadata(self._active_binding_hint)) - if self._mesh_scopes: - metadata.append(ExecutionDomainMetadata(self._mesh_scopes)) - return tuple(metadata) - - @staticmethod - def _with_binding(expr: Expr, name: str) -> Expr: - return replace_metadata(expr, BindingMetadata(name)) - - - - def visit_Constant(self, node: ast.Constant) -> Expr: - return self._constant_expr(node.value) - - def _constant_expr(self, value: Any) -> Expr: - constant = _constant_from_py(value) - span = self._source_span() - return replace_metadata(constant, span) if span is not None else constant - - def _static_number(self, node: ast.AST): - """The number *node* already is, or ``None`` to leave it to the IR path.""" - try: - value = eval_static( - node, - closure=self.closure, - lookup=self.env.lookup, - allowed_nodes=_STATIC_ARITH_NODES, - attr_resolver=self._resolve_static_attribute, - ) - except VerifyError: - return None - return value if isinstance(value, (int, float)) else None - - def _static_iterable(self, node: ast.AST): - """The compile-time sequence a comprehension walks. - - The compile-time sequence a comprehension walks: builtin ``range`` over - compile-time integers, or a tuple / list of compile-time values. - - No other call is evaluated here — resolving one would run it. - """ - if isinstance(node, ast.Call): - shadowed = self.env.lookup("range") is not None or "range" in self.closure - if ( - not isinstance(node.func, ast.Name) - or node.func.id != "range" - or shadowed - or node.keywords - ): - raise VerifyError( - f"compile-time list comprehension iterates `range(...)` or a " - f"compile-time sequence, not {ast.unparse(node)!r}" - ) - bounds = [self._static_number(arg) for arg in node.args] - if not bounds or any( - isinstance(bound, bool) or not isinstance(bound, int) for bound in bounds - ): - raise VerifyError("`range(...)` here takes compile-time integers") - return range(*bounds) - value = eval_static( - node, - closure=self.closure, - lookup=self.env.lookup, - allowed_nodes=(*_STATIC_ARITH_NODES, ast.Tuple, ast.List), - attr_resolver=self._resolve_static_attribute, - ) - if not isinstance(value, (tuple, list, range)): - raise VerifyError( - f"compile-time list comprehension iterates a compile-time sequence, " - f"got {type(value).__name__}" - ) - return value - - def _static_expr_list(self, node: ast.AST) -> "list[Expr] | None": - """A compile-time list of IR expressions, or ``None`` when *node* is not one. - - The list stays Python; only its elements are Exprs. - """ - if isinstance(node, ast.List): - return [self.expr(el) for el in node.elts] - if not isinstance(node, ast.ListComp): - return None - if len(node.generators) != 1: - raise VerifyError("compile-time list comprehension takes one `for` clause") - generator = node.generators[0] - if generator.ifs or generator.is_async: - raise VerifyError( - "compile-time list comprehension takes no `if` guard and is not async" - ) - if not isinstance(generator.target, ast.Name): - raise VerifyError("compile-time list comprehension binds one plain name") - values = list(self._static_iterable(generator.iter)) - self.env.push_frame() - try: - items = [] - for value in values: - self.env.define(generator.target.id, value) - items.append(self.expr(node.elt)) - finally: - self.env.pop_frame() - return items - - - - def visit_Name(self, node: ast.Name) -> Expr: - val = self.env.lookup(node.id) - from_closure = False - if val is None: - val = self.closure.get(node.id) - from_closure = True - if val is None: - raise VerifyError(f"undefined name {node.id!r}") - if isinstance(val, Expr): - return val - if isinstance(val, slice): - return val.start - if isinstance(val, (int, float, bool)): - return _constant_from_py(val) - - - if from_closure and type(val).__name__ in _IR_OBJECT_TYPES: - _warn_if_ir_object(val, node.id) - raise VerifyError(f"name {node.id!r} resolved to non-Expr Python value {type(val).__name__}") - - - - def visit_Attribute(self, node: ast.Attribute) -> Expr: - - value = self._static_number(node) - if value is not None: - return self._constant_expr(value) - - raise VerifyError(f"attribute access {ast.unparse(node)!r} not valid as Expr") - - - - def visit_Subscript(self, node: ast.Subscript) -> Expr: - """Resolve ``expr[idx]`` to a ``TupleGetItem`` or ``Slice`` Call. - - - a compile-time list + integer index → the element it holds. - - ``TupleType`` value + int constant index → ``TupleGetItem``. - - ``TensorType`` + slices / tile-window Names → ``Slice``. - - Compile-time integers and scalar range induction Names also reshape - away their selected axis, matching torch indexing. - """ - if _annotation_head_name(node.value) in ("Tensor", "ConstTensor"): - return _resolve_tensor_type( - node, - self.closure, - mesh_resolver=self._resolve_body_mesh, - default_mesh=self._current_default_mesh(), - mesh_order=self._mesh_scopes, - name_resolver=self.env.lookup, - ) - if isinstance(node.value, ast.Name): - bound = self.env.lookup(node.value.id) - if isinstance(bound, list): - return self._list_element(node.value.id, bound, node.slice) - value = self.expr(node.value) - if isinstance(value.type, TupleType): - slc = node.slice - if not (isinstance(slc, ast.Constant) and isinstance(slc.value, int) - and not isinstance(slc.value, bool)): - raise VerifyError( - "subscript on TupleType requires an integer constant index" - ) - return self._build_call(TupleGetItem(index=slc.value), (value,)) - if isinstance(value.type, TensorType): - return self._lift_tensor_subscript(value, node.slice) - raise VerifyError( - f"subscript only supported on TupleType / TensorType (got " - f"{type(value.type).__name__})" - ) - - def _lift_tensor_subscript(self, value, slc: ast.AST): - """Lift ``x[slice0, slice1, ...]`` to a ``Slice`` Op call. - - Each subscript element is one of: - - ``ast.Slice`` — full or partial ``start:stop[:step]``; - - an ``ast.Name`` resolving to a Python ``slice`` parser-side - binding (``for ok in tile(extent, step)``). - - Other forms (constants, computed Expr indices, ellipsis, lists) - are deferred to indexed read/write ops and raise here. - """ - if isinstance(slc, ast.Tuple): - elts = list(slc.elts) - else: - elts = [slc] - - x_ty = value.type - if not isinstance(x_ty, TensorType): # pragma: no cover — guarded above - raise VerifyError("tensor subscript: value must be TensorType") - if len(elts) != len(x_ty.shape): - raise VerifyError( - f"tensor subscript rank {len(elts)} != tensor rank " - f"{len(x_ty.shape)}" - ) - if shard_layout_of(x_ty.layout) is not None and any( - self._contains_mesh_coordinate(el) for el in elts - ): - raise VerifyError( - "tensor subscript uses a mesh coordinate to index an already " - "placed tensor; data-dependent mesh ownership is unresolved" - ) - - starts: list[Expr] = [] - sizes: list[Any] = [] - strides: list[Any] = [] - collapsed: list[int] = [] - for axis, (el, dim) in enumerate(zip(elts, x_ty.shape)): - index = self._integer_index(el, dim) - if index is not None: - starts.append(i64_const(index)) - sizes.append(1) - strides.append(1) - collapsed.append(axis) - continue - scalar_index = self._scalar_index(el) - if scalar_index is not None: - starts.append(scalar_index) - sizes.append(1) - strides.append(1) - collapsed.append(axis) - continue - b, e, s = self._slicer_for_dim(el, dim, axis) - b_expr = dim_expr(b) - e_expr = dim_expr(e) - s_expr = dim_expr(s) - starts.append(b_expr) - from tilefoundry.ir.hir.tensor.slice import slice_size # noqa: PLC0415 - - size = normalize_dim(slice_size(b_expr, e_expr, s_expr)) - if not is_dim_expr(size): - raise VerifyError( - f"tensor subscript axis {axis}: a run-time start needs the " - "stop endpoint written as `start + K` for a compile-time K, " - "because Slice takes a static size" - ) - sizes.append(int(size.value) if isinstance(size, Constant) else size) - strides.append(s) - - starts_expr = Tuple( - type=TupleType(fields=tuple(start.type for start in starts)), - elements=tuple(starts), - ) - sliced = self._build_call( - Slice(sizes=tuple(sizes), strides=tuple(strides)), - (value, starts_expr), - ) - if not collapsed: - return sliced - kept = tuple( - dim for axis, dim in enumerate(sliced.type.shape) if axis not in collapsed - ) - return self._build_call(Reshape(new_shape=kept), (sliced,)) - - def _integer_index(self, el: ast.AST, dim: Any) -> "int | None": - """The compile-time integer this element is, counted from the front. - - ``ast.Slice`` and a tile-window ``slice`` name keep their axis and are left - to ``_slicer_for_dim``. - """ - if isinstance(el, ast.Slice): - return None - if isinstance(el, ast.Name) and isinstance(self.env.lookup(el.id), slice): - return None - value = self._static_number(el) - if isinstance(value, bool) or not isinstance(value, int): - return None - if value < 0 or isinstance(dim, int): - if not isinstance(dim, int): - raise VerifyError( - f"tensor subscript index {value}: counting back from the end needs " - f"a static extent, and this axis is {dim}" - ) - normalized = value + dim if value < 0 else value - if not 0 <= normalized < dim: - raise VerifyError( - f"tensor subscript index {value} is out of range for extent {dim}" - ) - return normalized - return value - - def _scalar_index(self, el: ast.AST) -> "Expr | None": - """A registered scalar induction index from a ``range`` loop.""" - if not isinstance(el, ast.Name): - return None - value = self.env.lookup(el.id) - type_ = value.type if isinstance(value, Expr) else None - if ( - id(value) in self._scalar_index_ids - and isinstance(type_, TensorType) - and type_.shape == () - and type_.dtype is DType.i64 - ): - return value - return None - - def _list_element(self, name: str, items: list, slc: ast.AST) -> Expr: - """One element of a compile-time list, taken by compile-time index.""" - index = self._static_number(slc) - if isinstance(index, bool) or not isinstance(index, int): - raise VerifyError( - f"{name!r} is a compile-time list, so its index must be a " - f"compile-time integer" - ) - if not -len(items) <= index < len(items): - raise VerifyError( - f"{name!r} holds {len(items)} entries; index {index} is out of range" - ) - return items[index] - - def _slicer_for_dim(self, el: ast.AST, dim: Any, axis: int): - """Resolve one subscript element to ``(begin, end, stride)``. - - ``dim`` is the input tensor's static shape value at this axis - (used as the default upper bound for ``:``). - """ - if isinstance(el, ast.Slice): - - if el.lower is None: - begin = 0 - else: - begin = self._slicer_endpoint(el.lower) - if el.upper is None: - end = dim - else: - end = self._slicer_endpoint(el.upper) - if el.step is None: - stride = 1 - else: - stride = self._eval_static(el.step) - if not is_dim_expr(stride): - raise VerifyError( - f"tensor subscript axis {axis}: slice stride must be a " - "compile-time dimension" - ) - if all( - isinstance(value, int) and not isinstance(value, bool) - for value in (dim, begin, end, stride) - ) and stride > 0: - begin, end, stride = slice(begin, end, stride).indices(dim) - return begin, end, stride - if isinstance(el, ast.Name): - val = self.env.lookup(el.id) - if isinstance(val, slice): - return val.start, val.stop, val.step - moved = self._moved_tile_window(el, dim, axis) - if moved is not None: - return moved.start, moved.stop, moved.step - raise VerifyError( - f"tensor subscript axis {axis}: unsupported indexer " - f"{ast.dump(el)} (expected `:`, `a:b`, a tile-window slice, or a " - f"tile window moved by a compile-time integer)" - ) - - def _slicer_endpoint(self, node: ast.AST): - """Resolve a slice endpoint, admitting runtime scalar dim arithmetic.""" - try: - return self._eval_static(node, allow_runtime_scalar=True) - except VerifyError: - return self.expr(node) - - def _tile_window(self, node: ast.AST) -> "slice | None": - """The tile window *node* names, else ``None``.""" - if not isinstance(node, ast.Name): - return None - value = self.env.lookup(node.id) - if isinstance(value, slice) and id(value.start) in self._tile_windows: - return value - return None - - def _window_move(self, el: ast.AST) -> "tuple[slice, int] | None": - """The tile window *el* reads and the offset that moves it. - - ``None`` when *el* names no window. Offsets accumulate, so - ``i + QN + KN`` and ``QN + KN + i`` name one move by one sum, and each - term is a compile-time integer on its own. - """ - window = self._tile_window(el) - if window is not None: - return window, 0 - if not isinstance(el, ast.BinOp) or not isinstance(el.op, (ast.Add, ast.Sub)): - return None - sign = -1 if isinstance(el.op, ast.Sub) else 1 - moved = self._window_move(el.left) - offset_node = el.right - if moved is None: - moved = self._window_move(el.right) - if moved is None: - return None - if sign == -1: - raise VerifyError( - f"{ast.unparse(el)!r}: an offset moves a window, so the window " - f"is what the offset is added to -- subtracting it from " - f"{ast.unparse(el.left)!r} reverses the window rather than " - f"moving it" - ) - offset_node = el.left - offset = self._static_number(offset_node) - if isinstance(offset, bool) or not isinstance(offset, int): - raise VerifyError( - f"{ast.unparse(el)!r}: a tile window moves by a compile-time " - f"integer, and {ast.unparse(offset_node)!r} is not one" - ) - window, carried = moved - return window, carried + sign * offset - - def _moved_tile_window(self, el: ast.AST, dim: Any, axis: int) -> "slice | None": - """The window ``i + C`` reads, or ``None`` when *el* names no window. - - A tile window is a length bound to a moving base, so a compile-time - offset moves the base and leaves the length alone: ``i + C`` reads - ``[lo + C, lo + C + step)``. The base was already computed at compile - time, so this axis keeps the static extent ``i`` alone gives it. - """ - move = self._window_move(el) - if move is None: - return None - window, offset = move - if offset == 0: - return window - extent, length = self._tile_windows[id(window.start)] - self._check_moved_window(el, axis, offset, extent, length, dim) - base = simplify_dim(DimAdd, (window.start, offset)) - return slice(base, simplify_dim(DimAdd, (base, length)), window.step) - - @staticmethod - def _check_moved_window( - el: ast.AST, axis: int, offset: int, extent: Any, length: Any, dim: Any - ) -> None: - """Refuse a move that reads off the axis. - - The loop domain, the window length, the offset and the axis extent are - all compile-time, so the last window a moved read touches is too. A - symbolic extent leaves the bound to evaluate time, which is where an - unmoved window's own tail is caught. - """ - if offset < 0: - raise VerifyError( - f"tensor subscript axis {axis}: {ast.unparse(el)!r} moves the " - f"window back by {-offset}, and this loop's first window starts " - f"at 0, so it would begin before the axis" - ) - if not all( - isinstance(value, int) and not isinstance(value, bool) - for value in (extent, length, dim) - ): - return - if extent <= 0 or length <= 0: - return - last = (extent - 1) // length * length - if last + offset + length > dim: - raise VerifyError( - f"tensor subscript axis {axis}: {ast.unparse(el)!r} reads " - f"[{last + offset}, {last + offset + length}) on its last " - f"iteration, and the axis is {dim} long" - ) - - - - def visit_BinOp(self, node: ast.BinOp) -> Expr: - - value = self._static_number(node) - if value is not None: - return self._constant_expr(value) - opname = type(node.op).__name__ - - if opname == "MatMult": - matmul_cls = resolve_op("matmul") - if matmul_cls is None: - raise VerifyError("matmul op not registered") - left = self.expr(node.left) - right = self.expr(node.right) - return self._build_call(matmul_cls(), (left, right)) - kind = _binary_kind_for_ast_op(opname) - if kind is None: - raise VerifyError(f"unsupported binary op {opname}") - left = self.expr(node.left) - right = self.expr(node.right) - return self._build_call(self._make_binary(kind), (left, right)) - - - - def visit_Compare(self, node: ast.Compare) -> Expr: - if len(node.ops) != 1: - raise VerifyError("chained comparison not supported in V1") - opname = type(node.ops[0]).__name__ - kind = _binary_kind_for_ast_op(opname) - if kind is None: - raise VerifyError(f"unsupported comparison {opname}") - left = self.expr(node.left) - right = self.expr(node.comparators[0]) - return self._build_call(self._make_binary(kind), (left, right)) - - def visit_BoolOp(self, node: ast.BoolOp) -> Expr: - opname = type(node.op).__name__ - kind = _binary_kind_for_ast_op(opname) - if kind is None: - raise VerifyError(f"unsupported bool op {opname}") - if len(node.values) != 2: - raise VerifyError("bool op requires exactly 2 operands in V1") - left = self.expr(node.values[0]) - right = self.expr(node.values[1]) - return self._build_call(self._make_binary(kind), (left, right)) - - @staticmethod - def _make_binary(kind): - return Binary(kind=kind) - - @staticmethod - def _make_unary(kind): - return Unary(kind=kind) - - - - def visit_UnaryOp(self, node: ast.UnaryOp) -> Expr: - opname = type(node.op).__name__ - kind = _unary_kind_for_ast_op(opname) - if kind is None: - raise VerifyError(f"unsupported unary op {opname}") - operand = self.expr(node.operand) - return self._build_call(self._make_unary(kind), (operand,)) - - def visit_Call(self, node: ast.Call) -> Expr: - return self.call_to_op_call(node) - - - - def _resolve_call_target(self, func: ast.AST): - """Resolve a bare or namespaced callee to an operation schema. - - Lexical and closure bindings precede dialect-strict registry lookup. - Return schemas so aliases and operation classes share one builder path; - match namespace modules by identity. See - [parser §4.2](docs/spec/parser.md#42-closure-then-registry-callee-resolution) - and [parser §4.6](docs/spec/parser.md#46-per-dialect-strict-resolution). - """ - if isinstance(func, ast.Name): - val = self.env.lookup(func.id) - if val is None: - val = self.closure.get(func.id) - schema = self._schema_from_value(val) - if schema is not None: - return schema - try: - _kind, cls = resolve_callable(func.id, self.token) - except VerifyError: - return None - return getattr(cls, "_op_schema", None) - if ( - isinstance(func, ast.Attribute) - and isinstance(func.value, ast.Name) - ): - ns = self.env.lookup(func.value.id) - if ns is None: - ns = self.closure.get(func.value.id) - if ns is None: - return None - - - # noqa cycle: tilefoundry.dsl pulls tilefoundry.parser.overload, which - - import tilefoundry.dsl as _dsl # noqa: PLC0415 - if ns is _dsl.tf: - return resolve_schema(func.attr, "tf") - if ns is _dsl.T: - return resolve_schema(func.attr, "T") - return None - - def _resolve_function_target(self, func: ast.AST): - """Resolve function target. - - Return ``(hir.Function, binding, owner)`` behind a callee AST, or - ``(None, None, None)`` when the callee is not an ``@func``-decorated function. - ``@func`` evaluates to the ``hir.Function`` directly, so a sibling - callee binding *is* that Function (see :func:`tilefoundry.script.func`). - A bare name bound to a ``Module`` is that Module's entry function and - reports the name as *binding* and that Module as *owner*; the attribute - spelling is refused. - """ - val: Any = None - if isinstance(func, ast.Name): - val = self.env.lookup(func.id) - if val is None: - val = self.closure.get(func.id) - if isinstance(val, Module) and self.resolves_module_callees: - if not self.in_module_body: - raise VerifyError( - f"{func.id!r}: a Module is called only from a function " - f"authored in a @module class body, which is what attaches " - f"{val.name!r} as a child and gives the call a child to " - f"reach; this function declares none" - ) - return self._module_entry_target(val, func.id), func.id, val - elif isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): - owner = self.env.lookup(func.value.id) - if owner is None: - owner = self.closure.get(func.value.id) - if isinstance(owner, Module) and self.resolves_module_callees: - declares = ( - f"entry {owner.entry!r}" if owner.entry is not None else "no entry" - ) - raise VerifyError( - f"{ast.unparse(func)!r}: a Module is called through its bare " - f"binding, which calls its entry function; Module {owner.name!r} " - f"declares {declares}, so write {func.value.id}(...) rather than " - f"reaching in for {func.attr!r}" - ) - if isinstance(val, HirFunction): - return val, None, None - return None, None, None - - @staticmethod - def _module_entry_target(owner: Module, name: str) -> HirFunction: - """The hir Function *owner* is entered through. - - A Module that cannot answer is refused here, naming what it could not - answer, rather than falling through to *unknown Op name*. - """ - try: - entry = owner.entry_function() - except ValueError as exc: - raise VerifyError( - f"{name!r}: calling a Module calls its entry function -- {exc}" - ) from exc - if not isinstance(entry, HirFunction): - raise VerifyError( - f"{name!r}: calling a Module calls its entry function, and Module " - f"{owner.name!r} enters through {entry.name!r}, a " - f"{type(entry).__name__} rather than an hir Function" - ) - return entry - - def _build_function_call( - self, callee: Any, node: ast.Call, name: str, - module_binding: str | None = None, owner: Module | None = None, - ) -> Expr: - """Build a nested HIR function call with an elaborated target. - - Enforce arity before binding argument types so the call targets its - per-site function instance. Accept only positional IR arguments plus the - explicit ``loc=`` binding label. See - [hir §1.1](docs/spec/hir.md#11-function). - """ - explicit_loc: str | None = None - explicit_loc_given = False - extra_kwargs: list[str] = [] - for k in node.keywords: - if k.arg == "loc": - explicit_loc = self._eval_static(k.value) - explicit_loc_given = True - continue - extra_kwargs.append(k.arg) - if extra_kwargs: - raise VerifyError( - f"{name!r}: nested @func call does not accept keyword args " - f"{extra_kwargs!r} (positional-only at the IR level)" - ) - module_call = module_binding is not None - expected = len(bound_params(callee, from_reading=module_call)) - got = len(node.args) - if got != expected and module_call: - raise VerifyError( - f"{name!r}: Module {name!r} takes {expected} activation(s) — its " - f"{len(callee.params) - expected} ConstTensor parameter(s) come from " - f"that Module's own bindings — but got {got}" - ) - if got != expected: - raise VerifyError( - f"{name!r}: nested @func call arity mismatch — callee " - f"declares {expected} parameter(s), call passed {got}" - ) - input_args = tuple(self.expr(a) for a in node.args) - records = () if owner is None else (_ModuleCallee(module_binding, owner),) - call_for_errors = Call( - type=callee.return_type, target=callee, args=input_args, - metadata=(*self._source_metadata(), *records), - ) - if explicit_loc_given: - call_for_errors = replace_metadata( - call_for_errors, BindingMetadata(explicit_loc) - ) - instance = elaborate( - callee, tuple(a.type for a in input_args), self._ctx, - call=call_for_errors, - ) - call = self._build_call(instance, input_args, records=records) - if explicit_loc_given: - call = replace_metadata(call, BindingMetadata(explicit_loc)) - self._explicit_binding_call_ids.add(id(call)) - - self._call_dsl_names[id(call)] = name - return call - - @staticmethod - def _schema_from_value(val): - """Extract an ``OpSchema`` from a bound DSL surface value. - - Accepts: - - an ``OpSchema`` instance directly; - - any object carrying an ``_op_schema`` attribute (Op class - set by ``@register_op``; alias builder fn set by - ``@register_alias``). - Returns ``None`` for anything else. - """ - if isinstance(val, OpSchema): - return val - schema = getattr(val, "_op_schema", None) - if isinstance(schema, OpSchema): - return schema - return None - - def call_to_op_call(self, node: ast.Call) -> Expr: - """Resolve ``foo(...)`` to a ``Call`` on an hir Op. - - Dispatches a callee that is an ``hir.Function`` (nested ``@func`` call) - to :meth:`_build_function_call`; otherwise resolves an ``OpSchema`` via - :meth:`_resolve_call_target` and binds positional / keyword args to the - schema's input / attribute ParamDefs. Raises when the callee is a tir - Stmt op (the caller handles Stmt position) or is unresolved. - """ - if isinstance(node.func, ast.Name): - name = node.func.id - elif isinstance(node.func, ast.Attribute): - name = ast.unparse(node.func) - else: - raise VerifyError("only Name / Attribute callees supported in V1") - - - - - - callee_func, module_binding, owner = self._resolve_function_target(node.func) - if callee_func is not None: - return self._build_function_call( - callee_func, node, name, module_binding, owner - ) - - schema = self._resolve_call_target(node.func) - if schema is None: - - if isinstance(node.func, ast.Name) and resolve_stmt(name) is not None: - raise VerifyError( - f"{name!r} is an effect Stmt op; cannot appear in Expr position " - f"(wrap in Assign or use as top-level Stmt)" - ) - raise VerifyError(f"unknown Op name {name!r}") - - - param_infos = schema.signature - input_params = [p for p in param_infos if p.kind == "input"] - attr_params = [p for p in param_infos if p.kind == "attribute"] - - - - - - - - - is_variadic = bool(getattr(getattr(schema, "op_class", None), "is_variadic", False)) - - - pos_args = list(node.args) - input_args = [] - attr_kwargs: dict[str, Any] = {} - - if is_variadic: - if len(input_params) != 1: - raise VerifyError( - f"{name!r}: variadic op schema must declare exactly one " - f"input ParamDef, got {len(input_params)}" - ) - for arg in pos_args: - input_args.append(self.expr(arg)) - else: - for i, arg in enumerate(pos_args): - if i < len(input_params): - if ( - isinstance(arg, ast.Tuple) - and ( - ( - schema.name == "insert_slice" - and input_params[i].name == "offsets" - ) - or ( - schema.name == "slice" - and input_params[i].name == "starts" - ) - ) - ): - - - - - input_args.append(self._tuple_expr_expr(arg)) - else: - input_args.append(self.expr(arg)) - else: - attr_idx = i - len(input_params) - if attr_idx >= len(attr_params): - raise VerifyError( - f"{name!r}: too many positional arguments " - f"(expected at most {len(input_params) + len(attr_params)}, got {len(pos_args)})" - ) - attr_name = attr_params[attr_idx].name - if attr_name in attr_kwargs: - raise VerifyError( - f"{name!r}: duplicate binding for attribute {attr_name!r}" - ) - attr_kwargs[attr_name] = self._eval_static_or_sugar( - attr_name, arg, schema=schema - ) - - - explicit_loc: str | None = None - explicit_loc_given = False - - - for k in node.keywords: - if k.arg == "loc": - explicit_loc = self._eval_static(k.value) - explicit_loc_given = True - continue - if k.arg in attr_kwargs: - raise VerifyError( - f"{name!r}: duplicate binding for attribute {k.arg!r} " - f"(both positional and keyword)" - ) - attr_kwargs[k.arg] = self._eval_static_or_sugar(k.arg, k.value, schema=schema) - - - - if "storage" in attr_kwargs: - attr_kwargs["storage"] = resolve_storage(attr_kwargs["storage"]) - - op_inst = self._build_op_instance(schema, attr_kwargs) - call = self._build_call(op_inst, tuple(input_args)) - if explicit_loc_given: - call = replace_metadata(call, BindingMetadata(explicit_loc)) - self._explicit_binding_call_ids.add(id(call)) - - - - - self._call_dsl_names[id(call)] = schema.name - return call - - def _build_op_instance(self, schema, attr_kwargs): - """Construct an Op instance from a resolved schema and attr kwargs. - - There is a single path — every schema (real Op or surface - alias) carries a ``builder`` callable. Real-Op schemas default - to ``cls`` itself; alias schemas have a custom builder that - constructs the kinded target Op. - """ - return schema.builder(**attr_kwargs) - - - - def _maybe_autofill_binding(self, expr: Expr, name: str) -> Expr: - """Set a Call binding label to *name* unless ``loc=`` was explicit. - - Returns *expr* unchanged when it is not a Call or already has an - explicit binding. - """ - if not isinstance(expr, Call): - return expr - if id(expr) in self._explicit_binding_call_ids: - return expr - return self._with_binding(expr, name) - - def _maybe_autofill_binding_default(self, expr: Expr) -> Expr: - """Maybe autofill binding default. - - Set the Call binding label to the DSL callable name (default) when the - user did not supply ``loc=`` explicitly. Used for tuple-unpack - parents where there is no single LHS variable name to inherit. - """ - if not isinstance(expr, Call): - return expr - if id(expr) in self._explicit_binding_call_ids: - return expr - dsl_name = self._call_dsl_names.get(id(expr)) - if dsl_name is None: - return expr - binding = get_metadata(expr, BindingMetadata) - if binding is not None and binding.name == dsl_name: - return expr - return self._with_binding(expr, dsl_name) - - def _build_call( - self, op_inst, args: tuple[Expr, ...], *, records: tuple[IRMetadata, ...] = () - ) -> Call: - """Build a Call with type eagerly populated via the typeinfer registry. - - *records* are carried on the node the typeinfer walk sees, because a - record stating how the call binds its arguments has to be there before - the type is derived from them. - """ - args = _with_python_float_dtypes(args) - - - placeholder = Call( - type=TensorType.scalar(DType.f32), target=op_inst, args=args, - metadata=(*self._source_metadata(), *records), - ) - computed = TypeInferVisitor(self._ctx).visit(placeholder) - return dataclasses.replace(placeholder, type=computed) - - def _eval_static_or_sugar( - self, - attr_name: str, - node: ast.AST, - *, - schema=None, - op_cls: type | None = None, - ): - """Evaluate an attribute with annotation-driven layout-sugar parsing. - - Prefer alias-aware schema annotations and retain operation classes for - legacy callers. Without an annotation, a ``layout`` attribute still - attempts shard-layout sugar before static evaluation. See - [parser §4.4](docs/spec/parser.md#44-annotation-driven-sugar-dispatch). - """ - annotation = self._lookup_param_annotation( - schema=schema, op_cls=op_cls, attr_name=attr_name - ) - if annotation is TensorType and isinstance(node, ast.Subscript): - return self.expr(node) - if annotation is not None and _is_tuple_sugar(node): - sugar = self._sugar_parser_for_annotation(annotation) - if sugar is not None: - try: - return sugar(node) - except LayoutSugarError: - - - raise - except ValueError: - pass - elif annotation is None and attr_name == "layout" and _is_tuple_sugar(node): - - try: - return parse_sugar( - node, - ShardLayout, - mesh_resolver=self._resolve_body_mesh, - default_mesh=self._current_default_mesh(), - mesh_order=self._mesh_scopes, - closure=self.closure, - ) - except LayoutSugarError: - raise - except ValueError: - pass - value = self._eval_static(node) - - if ( - isinstance(value, str) - and isinstance(annotation, type) - and issubclass(annotation, enum.Enum) - ): - try: - return annotation(value) - except ValueError: - valid = ", ".join(repr(e.value) for e in annotation) - raise VerifyError( - f"{annotation.__name__}: unknown value {value!r}; " - f"valid values are {valid}" - ) from None - if isinstance(value, str) and annotation is DType: - try: - return DType.from_name(value) - except ValueError as exc: - raise VerifyError(str(exc)) from None - return value - - def _lookup_param_annotation( - self, - *, - schema=None, - op_cls: type | None = None, - attr_name: str, - ) -> type | None: - """Return the ``ParamDef.annotation`` for *attr_name*. - - Prefers the explicit ``schema`` argument (alias-aware); falls - back to ``op_cls._op_schema.signature`` for legacy callers. - Returns ``None`` when no schema/ParamDef matches. - """ - if schema is None and op_cls is not None: - schema = getattr(op_cls, "_op_schema", None) - if schema is None: - return None - for pd in schema.signature: - if pd.name == attr_name: - return pd.annotation - return None - - def _sugar_parser_for_annotation(self, annotation: type): - """Return the sugar parser for a Layout-family annotation, else None.""" - if annotation is ShardLayout: - return lambda n: parse_sugar( - n, - ShardLayout, - mesh_resolver=self._resolve_body_mesh, - default_mesh=self._current_default_mesh(), - mesh_order=self._mesh_scopes, - closure=self.closure, - ) - if annotation is Layout: - return lambda n: parse_sugar(n, Layout, closure=self.closure) - return None - - def _resolve_static_attribute(self, owner, attr: str): - """Resolve a static ``owner.attr`` access during ``_eval_static``. - - Default: plain ``getattr``. Dialect visitors override this to add - context-sensitive resolution (e.g. the TIR parser checks that an MMA - fragment ``atom.A`` is used inside a compatible enclosing mesh scope). - """ - return getattr(owner, attr) - - def _eval_static(self, node: ast.AST, *, allow_runtime_scalar: bool = False): - """Eval static. - - Evaluate an AST node statically for attribute kwargs (axis=1, - new_shape=(M,K), layout=ShardLayout(...), etc.). - - Thin policy wrapper over :func:`eval_static` (parser/static_eval.py): - the full node set, ``Name`` resolution through the lexical env before - the closure, closure-captured-IR-object warnings, and true division - for ``ast.Div``. - """ - return eval_static( - node, - closure=self.closure, - lookup=self.env.lookup, - attr_resolver=self._resolve_static_attribute, - on_closure_name=_warn_if_ir_object, - allow_runtime_scalar=allow_runtime_scalar, - ) - - -__all__ = ["BaseExprVisitor", "extract_ast", "Token"] diff --git a/src/tilefoundry/parser/dispatch.py b/src/tilefoundry/parser/dispatch.py deleted file mode 100644 index c6baf4c6..00000000 --- a/src/tilefoundry/parser/dispatch.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Resolve DSL names and Python operators through dialect registries. - -HIR resolves only HIR operations; TIR resolves only TIR statements and user -intrinsics, with no cross-dialect fallback. Binary and unary AST operators map -directly to kinded IR. Registered schemas provide flat-name operation and -statement lookup. -See [parser §4.6](docs/spec/parser.md#46-per-dialect-strict-resolution). -""" -from __future__ import annotations - -from typing import Literal - -from tilefoundry.ir.core import VerifyError -from tilefoundry.ir.core.kinds import BinaryKind, UnaryKind -from tilefoundry.ir.core.op_registry import ( - _first_schema, - get_op_by_name, - get_stmt_by_name, -) -from tilefoundry.ir.core.op_schema import OpSchema -from tilefoundry.ir.tir.intrinsic import _intrinsic_dispatch - -Token = Literal["hir", "tir"] - - - - - - - -def _binary_kind_for_ast_op(ast_op_name: str): - _MAP = { - "Add": BinaryKind.ADD, "Sub": BinaryKind.SUB, - "Mult": BinaryKind.MUL, "Div": BinaryKind.DIV, - "FloorDiv": BinaryKind.FLOOR_DIV, "Mod": BinaryKind.MOD, - "Eq": BinaryKind.EQ, "NotEq": BinaryKind.NE, - "Lt": BinaryKind.LT, "LtE": BinaryKind.LE, - "Gt": BinaryKind.GT, "GtE": BinaryKind.GE, - "And": BinaryKind.AND, "Or": BinaryKind.OR, - } - return _MAP.get(ast_op_name) - - -def _unary_kind_for_ast_op(ast_op_name: str): - _MAP = {"USub": UnaryKind.NEG, "Not": UnaryKind.NOT} - return _MAP.get(ast_op_name) - - -def resolve_op(name: str) -> type | None: - """Resolve a DSL bare-call name to an HIR Op subclass, or ``None``. - - Skips alias schemas (``op_class=None``); this returns the concrete - legacy class. Use :func:`resolve_schema` to honour aliases. - """ - return get_op_by_name(name) - - -def resolve_schema(name: str, dialect: str = "tf") -> OpSchema | None: - """Resolve a DSL bare-call name to its first ``OpSchema`` (alias-aware). - - A surface name may map to a surface-alias schema - (``schema.op_class is None``) prepended over a legacy real-Op - schema. Parser dispatch uses this resolver so the alias wins - first-match — its ``builder`` constructs the kinded target Op - (e.g. ``Binary(kind=ADD)``) instead of the legacy class. - """ - return _first_schema(dialect, name) - - -def resolve_stmt(name: str) -> type | None: - """Resolve a DSL bare-call name to a TIR Stmt subclass, or ``None``. - - Falls through to user-registered intrinsics (``@intrinsic`` decorator) - so user-defined effect Stmts continue to participate in TIR dispatch - without going through the canonical opt-in registry. - """ - cls = get_stmt_by_name(name) - if cls is not None: - return cls - return _intrinsic_dispatch.get(name) - - -def resolve_callable(name: str, token: Token) -> tuple[str, type]: - """Dispatch *name* within one strict DSL dialect. - - Return an operation or statement kind and class, or raise ``VerifyError``. - A trailing underscore explicitly selects effect form in TIR. HIR names never - fall back to TIR, and TIR names never fall back to HIR. - See [parser §1.3](docs/spec/parser.md#13-op-call). - """ - if token == "tir": - - - - if name.endswith("_") and not name.startswith("_"): - base = name[:-1] - stmt = resolve_stmt(base) - if stmt is not None: - return ("stmt", stmt) - stmt = resolve_stmt(name) - if stmt is not None: - return ("stmt", stmt) - raise VerifyError( - f"unknown TIR callable {name!r} in @tilefoundry.prim_func body " - f"(bare HIR Op fallback removed; use tf.. " - f"namespace if this is meant to be an HIR Op)" - ) - op = resolve_op(name) - if op is not None: - return ("op", op) - raise VerifyError(f"unknown HIR callable {name!r} in @tilefoundry.func body") - - - - - - - - - - -__all__ = [ - "resolve_op", - "resolve_schema", - "resolve_stmt", - "resolve_callable", - "_binary_kind_for_ast_op", - "_unary_kind_for_ast_op", -] diff --git a/src/tilefoundry/parser/grammar_render.py b/src/tilefoundry/parser/grammar_render.py new file mode 100644 index 00000000..53b499d9 --- /dev/null +++ b/src/tilefoundry/parser/grammar_render.py @@ -0,0 +1,503 @@ +"""Render the executable Pattern graph as deterministic surface EBNF.""" + +# ruff: noqa: PLC0415 + +from __future__ import annotations + +import ast +import dataclasses +import textwrap +from typing import Any + +from .ast_pattern import ( + AstNodePattern, + BindPattern, + BranchPattern, + CapturePattern, + ChildPattern, + ChoicePattern, + ConditionPattern, + ElementPattern, + FieldPattern, + LazyPattern, + LiteralPattern, + OptionalPattern, + PredicatePattern, + ReferencePattern, + RepeatPattern, + SequencePattern, +) + + +def _grammar_name(value: str) -> str: + return value.replace("_", "-").replace(" ", "-").lower() + + +@dataclasses.dataclass(frozen=True) +class _Expr: + text: str + alternatives: tuple[str, ...] = () + + def flat(self) -> str: + if self.alternatives: + return "(" + " | ".join(self.alternatives) + ")" + return self.text + + +def _text(value: str) -> _Expr: + return _Expr(value) + + +def _concat(*parts: _Expr | None) -> _Expr: + return _text(" ".join(part.flat() for part in parts if part and part.flat())) + + +def _choice(*parts: _Expr) -> _Expr: + alternatives: list[str] = [] + for part in parts: + values = part.alternatives or (part.text,) + for value in values: + if value and value not in alternatives: + alternatives.append(value) + if len(alternatives) == 1: + return _text(alternatives[0]) + return _Expr("", tuple(alternatives)) + + +def _terminal(value: str) -> _Expr: + return _text(repr(value)) + + +def _optional(value: _Expr) -> _Expr: + return _text(f"({value.flat()})?") + + +def _delimited(left: str, value: _Expr, right: str) -> _Expr: + return _concat(_terminal(left), value, _terminal(right)) + + +def _separated(items: tuple[_Expr, ...], separator: _Expr) -> _Expr: + if not items: + return _text("") + output = items[0] + for item in items[1:]: + output = _concat(output, separator, item) + return output + + +def _repeated(item: _Expr, minimum: int, separator: _Expr) -> _Expr: + required = _concat(item, _text(f"({separator.flat()} {item.flat()})*")) + return required if minimum else _optional(required) + + +class RenderVisitor: + """Project executable AST patterns into readable Python authoring EBNF.""" + + def __init__(self, *, line_width: int = 100): + self.line_width = line_width + self._seen: set[str] = set() + self._productions: list[tuple[str, _Expr]] = [] + + def _element(self, pattern: ElementPattern[Any]) -> _Expr: + name = pattern.element_name + if not name: + raise TypeError(f"{type(pattern).__name__} has no explicit element_name") + grammar_name = _grammar_name(name) + if name in self._seen: + return _text(grammar_name) + self._seen.add(name) + if pattern.syntax is None: + raise TypeError(f"{type(pattern).__name__} has no executable syntax") + rhs = self.visit(pattern.syntax) + self._productions.append((grammar_name, rhs)) + return _text(grammar_name) + + def _list_pattern(self, pattern: object) -> _Expr: + comma = _terminal(",") + if isinstance(pattern, RepeatPattern): + return _repeated(self.visit(pattern.pattern), pattern.minimum, comma) + if isinstance(pattern, SequencePattern): + return _separated( + tuple(self.visit(item) for item in pattern.patterns), comma + ) + if isinstance(pattern, ChoicePattern): + return _choice(*(self._list_pattern(item) for item in pattern.patterns)) + return self.visit(pattern) + + @staticmethod + def _fields(pattern: AstNodePattern) -> dict[str, object]: + return { + part.name: part.pattern + for part in pattern.parts + if isinstance(part, FieldPattern) + } + + def _field(self, fields: dict[str, object], name: str, fallback: str) -> _Expr: + pattern = fields.get(name) + return _text(fallback) if pattern is None else self.visit(pattern) + + def _optional_field( + self, fields: dict[str, object], name: str, fallback: str + ) -> _Expr: + pattern = fields.get(name) + if isinstance(pattern, OptionalPattern): + pattern = pattern.pattern + return _text(fallback) if pattern is None else self.visit(pattern) + + def _ast_node(self, pattern: AstNodePattern) -> _Expr: + node_type = pattern.node_type + for part in pattern.parts: + self.visit(part) + fields = self._fields(pattern) + + if node_type is ast.Constant: + predicates = [ + part for part in pattern.parts if isinstance(part, PredicatePattern) + ] + if predicates: + return self.visit(predicates[-1]) + value = fields.get("value") + if isinstance(value, CapturePattern) or value is None: + return _choice( + _text("None"), + _text("Ellipsis"), + _text("boolean-literal"), + _text("integer-literal"), + _text("float-literal"), + _text("complex-literal"), + _text("string-literal"), + _text("bytes-literal"), + ) + return self.visit(value) + if node_type is ast.Name: + identifier = fields.get("id") + if isinstance(identifier, (LiteralPattern, ChoicePattern)): + return self.visit(identifier) + return _text("identifier") + if node_type is ast.Attribute: + value = fields.get("value") + if isinstance(value, AstNodePattern) and value.node_type is ast.Name: + return _concat( + _text("identifier"), + _terminal("."), + _text("identifier"), + ) + return _concat( + _text("primary"), + _terminal("."), + _text("identifier"), + ) + if node_type in {ast.Tuple, ast.List, ast.Set}: + delimiters = { + ast.Tuple: ("(", ")"), + ast.List: ("[", "]"), + ast.Set: ("{", "}"), + } + left, right = delimiters[node_type] + values_pattern = fields.get("elts") + if values_pattern is None: + structural = [ + part + for part in pattern.parts + if not isinstance( + part, (CapturePattern, FieldPattern, PredicatePattern) + ) + ] + if structural: + return self.visit(structural[-1]) + values = _text("") + else: + values = self._list_pattern(values_pattern) + return _delimited(left, values, right) + if node_type is ast.Dict: + key_pattern = fields.get("keys", LiteralPattern()) + value_pattern = fields.get("values", LiteralPattern()) + key = ( + self.visit(key_pattern.pattern) + if isinstance(key_pattern, RepeatPattern) + else self.visit(key_pattern) + ) + value = ( + self.visit(value_pattern.pattern) + if isinstance(value_pattern, RepeatPattern) + else self.visit(value_pattern) + ) + entry = _concat(key, _terminal(":"), value) + return _delimited("{", _repeated(entry, 0, _terminal(",")), "}") + if node_type is ast.BinOp: + return _concat( + self._field(fields, "left", "expression"), + self._field(fields, "op", "binary-op"), + self._field(fields, "right", "expression"), + ) + if node_type is ast.UnaryOp: + return _concat( + self._field(fields, "op", "unary-op"), + self._field(fields, "operand", "expression"), + ) + if node_type is ast.Compare: + return _concat( + self._field(fields, "left", "expression"), + self._list_pattern(fields.get("ops", SequencePattern())), + self._list_pattern(fields.get("comparators", SequencePattern())), + ) + if node_type is ast.BoolOp: + values = fields.get("values", SequencePattern()) + if isinstance(values, SequencePattern) and len(values.patterns) == 2: + return _concat( + self.visit(values.patterns[0]), + self._field(fields, "op", "boolean-op"), + self.visit(values.patterns[1]), + ) + return self._list_pattern(values) + if node_type is ast.Call: + argument_patterns: list[_Expr] = [] + minimum = 0 + for name in ("args", "keywords"): + values = fields.get(name) + if isinstance(values, RepeatPattern): + argument_patterns.append(self.visit(values.pattern)) + minimum = max(minimum, values.minimum) + arguments = ( + _repeated(_choice(*argument_patterns), minimum, _terminal(",")) + if argument_patterns + else _text("") + ) + return _concat( + self._field(fields, "func", "callee"), + _delimited("(", arguments, ")"), + ) + if node_type is ast.keyword: + return _concat( + self._field(fields, "arg", "name"), + _terminal("="), + self._field(fields, "value", "expression"), + ) + if node_type is ast.Slice: + lower = self._field(fields, "lower", "") + upper = self._field(fields, "upper", "") + step = self._optional_field(fields, "step", "") + return _concat( + lower, + _terminal(":"), + upper, + _optional(_concat(_terminal(":"), step)), + ) + if node_type is ast.Subscript: + return _concat( + self._field(fields, "value", "expression"), + _delimited("[", self._field(fields, "slice", "expression"), "]"), + ) + if node_type is ast.arguments: + return self._list_pattern(fields.get("args", SequencePattern())) + if node_type is ast.arg: + return _concat( + _text("name"), + _terminal(":"), + self._field(fields, "annotation", "type-annotation"), + ) + if node_type is ast.Assign: + return _concat( + self._list_pattern(fields.get("targets", SequencePattern())), + _terminal("="), + self._field(fields, "value", "expression"), + ) + if node_type is ast.AnnAssign: + value = self._optional_field(fields, "value", "expression") + return _concat( + self._field(fields, "target", "name"), + _terminal(":"), + self._field(fields, "annotation", "type-annotation"), + _optional(_concat(_terminal("="), value)), + ) + if node_type is ast.Return: + return _concat(_terminal("return"), self._field(fields, "value", "")) + if node_type is ast.Expr: + return self._field(fields, "value", "expression") + if node_type is ast.Pass: + return _terminal("pass") + if node_type is ast.For: + return _concat( + _terminal("for"), + self._field(fields, "target", "name"), + _terminal("in"), + self._field(fields, "iter", "expression"), + _terminal(":"), + self._field(fields, "body", "block"), + ) + if node_type is ast.withitem: + alias = _concat( + _terminal("as"), + self._field(fields, "optional_vars", "name"), + ) + return _concat( + self._field(fields, "context_expr", "expression"), + _optional(alias), + ) + if node_type is ast.With: + return _concat( + _terminal("with"), + self._list_pattern(fields.get("items", SequencePattern())), + _terminal(":"), + self._field(fields, "body", "block"), + ) + if node_type is ast.Module: + body = fields.get("body", SequencePattern()) + if isinstance(body, RepeatPattern): + return _repeated( + self.visit(body.pattern), body.minimum, _text("newline") + ) + return self._list_pattern(body) + if node_type is ast.FunctionDef: + returns = self._optional_field(fields, "returns", "return-type") + return _concat( + _terminal("def"), + _text("name"), + _delimited("(", self._field(fields, "args", "signature"), ")"), + _optional(_concat(_terminal("->"), returns)), + _terminal(":"), + self._field(fields, "body", "block"), + ) + if node_type is ast.MatMult: + return _terminal("@") + operator = { + ast.Add: "+", + ast.Sub: "-", + ast.Mult: "*", + ast.Div: "/", + ast.FloorDiv: "//", + ast.Mod: "%", + ast.Pow: "**", + ast.UAdd: "+", + ast.USub: "-", + ast.Not: "not", + }.get(node_type) + if operator is not None: + return _terminal(operator) + if node_type is ast.Load: + return _text("") + if node_type in {ast.expr, ast.stmt}: + structural = [ + part + for part in pattern.parts + if not isinstance(part, (CapturePattern, PredicatePattern)) + ] + if structural: + return self.visit(structural[-1]) + predicates = [ + part for part in pattern.parts if isinstance(part, PredicatePattern) + ] + if predicates: + return self.visit(predicates[-1]) + return _text("expression" if node_type is ast.expr else "statement") + + parts = tuple(self.visit(part) for part in pattern.parts) + return _concat(_text(_grammar_name(node_type.__name__)), *parts) + + def visit(self, pattern: Any) -> _Expr: + if isinstance(pattern, ElementPattern): + return self._element(pattern) + if isinstance(pattern, AstNodePattern): + return self._ast_node(pattern) + if isinstance(pattern, FieldPattern): + return self.visit(pattern.pattern) + if isinstance(pattern, LiteralPattern): + if pattern.value_type is not None: + types = ( + pattern.value_type + if isinstance(pattern.value_type, tuple) + else (pattern.value_type,) + ) + names = { + bool: "boolean-literal", + int: "integer-literal", + float: "float-literal", + str: "string-literal", + } + return _choice( + *( + _text(names.get(item, f"{item.__name__}-literal")) + for item in types + ) + ) + if pattern.value is dataclasses.MISSING: + return _text("literal") + return _text(repr(pattern.value)) + if isinstance(pattern, ReferencePattern): + return _text("primary") + if isinstance(pattern, SequencePattern): + return _separated( + tuple(self.visit(item) for item in pattern.patterns), _terminal(",") + ) + if isinstance(pattern, ChoicePattern): + return _choice(*(self.visit(item) for item in pattern.patterns)) + if isinstance(pattern, ConditionPattern): + return self.visit(pattern.pattern) + if isinstance(pattern, OptionalPattern): + return _optional(self.visit(pattern.pattern)) + if isinstance(pattern, RepeatPattern): + return _repeated( + self.visit(pattern.pattern), pattern.minimum, _terminal(",") + ) + if isinstance(pattern, (ChildPattern, BranchPattern, BindPattern, LazyPattern)): + return self.visit(pattern.pattern) + if isinstance(pattern, PredicatePattern): + return _text(_grammar_name(pattern.label)) + if isinstance(pattern, CapturePattern): + return _text(_grammar_name(pattern.name)) + raise TypeError(f"unsupported executable pattern {type(pattern).__name__}") + + def _format_production(self, name: str, rhs: _Expr, width: int) -> list[str]: + prefix = f"{name:<{width}} ::= " + continuation = " " * (width + 5) + if rhs.alternatives: + lines: list[str] = [] + for index, alternative in enumerate(rhs.alternatives): + marker = "" if index == 0 else "| " + lines.extend( + textwrap.wrap( + marker + alternative, + width=self.line_width, + initial_indent=prefix if index == 0 else continuation, + subsequent_indent=continuation + " ", + break_long_words=False, + break_on_hyphens=False, + ) + ) + return lines + return textwrap.wrap( + rhs.text, + width=self.line_width, + initial_indent=prefix, + subsequent_indent=continuation, + break_long_words=False, + break_on_hyphens=False, + ) or [prefix.rstrip()] + + def render(self, root: ElementPattern[Any]) -> str: + root_name = root.element_name + if not root_name: + raise TypeError(f"{type(root).__name__} has no explicit element_name") + self._element(root) + width = max(len(name) for name, _ in self._productions) + lines = [ + f"; root: {_grammar_name(root_name)}", + '; literal: Python ast.Constant syntax, e.g. 1, "bf16", or None', + "; name: Python variable name; primary: name/attribute base for calls and subscripts", + "; expression: Python syntax composed from literals, names, primaries, and operators", + "; runtime-expression: expression lowered to a TileFoundry IR Expr", + ] + for name, rhs in self._productions: + lines.extend(self._format_production(name, rhs, width)) + return "\n".join(lines) + + +def render_grammar(root: ElementPattern[Any] | None = None) -> str: + if root is None: + from .pattern_nodes import FunctionPattern + + root = FunctionPattern() + return RenderVisitor().render(root) + + +__all__ = ["RenderVisitor", "render_grammar"] diff --git a/src/tilefoundry/parser/hir_parser.py b/src/tilefoundry/parser/hir_parser.py deleted file mode 100644 index 15d9edb2..00000000 --- a/src/tilefoundry/parser/hir_parser.py +++ /dev/null @@ -1,1086 +0,0 @@ -from __future__ import annotations - -import ast -import dataclasses -from typing import Any - -from tilefoundry.ir.constraints import ( - ConstraintProvenance, - LayoutConstraint, - MeshConstraint, - ScheduleConstraint, - ScheduleConstraintMetadata, - SourceLocation, - StorageConstraint, -) -from tilefoundry.ir.constraints.layout import _LAYOUT_WILDCARD -from tilefoundry.ir.core import BindingMetadata, Expr, Var, VerifyError, get_metadata -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.grid_region import GridRegionExpr -from tilefoundry.ir.hir.sharding.local import Local -from tilefoundry.ir.hir.sharding.reshard import Reshard -from tilefoundry.ir.hir.tensor.arange import Arange -from tilefoundry.ir.hir.tensor.reshape import Reshape -from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem -from tilefoundry.ir.types import DType, TensorType, TupleType -from tilefoundry.ir.types.dim import ( - DimAdd, - DimFloorDiv, - DimMod, - DimMul, - DimSub, - DimVar, - is_dim_expr, - simplify_dim, -) -from tilefoundry.ir.types.dim_isl import normalize_dim -from tilefoundry.ir.types.shard import ( - Broadcast, - Layout, - Mesh, - Partial, - ShardLayout, - Split, - Topology, -) -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.utils.spec_ref import spec_ref_render - -from .base import ( - BaseExprVisitor, - _build_params, - _collect_closure, - _resolve_tensor_type, - extract_ast, -) -from .sugar import _is_tuple_sugar, _resolve_mesh_axis, parse_sugar -from .symtab import LexicalEnv - -_HIR_FUNCTION = "[hir §1.1](docs/spec/hir.md#11-function)" -_PARSER_MESH = "[parser §1.6](docs/spec/parser.md#16-with-mesh-as-m)" - - - -_AST_DIM_OPS = { - ast.Add: DimAdd, - ast.Sub: DimSub, - ast.Mult: DimMul, - ast.FloorDiv: DimFloorDiv, - ast.Mod: DimMod, -} - - -def parse_func(fn, *, topologies=(), specializations=(), extra_closure=None) -> Function: - """@tilefoundry.func parser entry. Parse fn's source into hir.Function. - - ``topologies`` is the declared topology namespace this body may name in - ``Mesh(("...",), ...)``; it is a parse-time namespace, not state the - resulting Function keeps. ``extra_closure`` adds names to the resolution - namespace below ``fn``'s own globals/freevars; it lets an ``@func`` defined - in a ``@module`` class body resolve sibling ``@func`` methods (which are - ``hir.Function`` values) as nested-call targets. - """ - return _parse_func( - fn, topologies=topologies, specializations=specializations, - extra_closure=extra_closure, - ) - - -def _parse_func( - fn, *, topologies=(), specializations=(), extra_closure=None, in_module_body=False -) -> Function: - """`parse_func` plus whether a ``@module`` class body is being authored. - - Only the decorators can answer that, and only they may state it: it decides - whether a name bound to a ``Module`` is a callee at all. - """ - node = extract_ast(fn) - closure = _collect_closure(fn, extra_closure) - return _parse_func_node( - node, closure, topologies=topologies, - specializations=specializations, - source_filename=getattr(getattr(fn, "__code__", None), "co_filename", ""), - in_module_body=in_module_body, - ) - - - -_NOT_STATIC = object() - -def _parse_func_node( - node: ast.FunctionDef, - closure: dict[str, Any], - *, - topologies=(), - specializations=(), - source_filename: str = "", - in_module_body: bool = False, -) -> Function: - env = LexicalEnv() - params = _build_params( - node, closure, _resolve_tensor_type, decorator_name="@tilefoundry.func" - ) - for p in params: - env.define(p.name, p) - - topo_ns: dict[str, "Topology"] = {} - for t in topologies: - if t.name in topo_ns: - raise VerifyError(f"duplicate topology name {t.name!r}") - topo_ns[t.name] = t - visitor = _HirBodyVisitor( - env, closure, topo_ns=topo_ns, source_filename=source_filename, - in_module_body=in_module_body, - ) - if _is_pass_body(node.body): - - - body_expr = None - else: - body_expr = visitor.visit_body(node.body) - return_type = _resolve_return_type(node, closure, body_expr) - return Function.build( - name=node.name, - params=params, - body=body_expr, - return_type=return_type, - specializations=tuple(specializations), - ) - - -def _constraint_value(node: ast.AST): - """Read a stage-neutral layout extent or topology reference.""" - if isinstance(node, ast.Constant): - return node.value - if isinstance(node, ast.Name): - return node.id - if isinstance(node, ast.Attribute): - return ast.unparse(node) - raise VerifyError( - f"where layout extent must be a literal or symbolic name, got " - f"{type(node).__name__}" - ) - - -def _parse_partial_value(node: ast.AST) -> Partial: - if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): - raise VerifyError('Partial binding must use P("reduction")') - if node.func.id != "P" or len(node.args) != 1 or node.keywords: - raise VerifyError('Partial binding must use P("reduction")') - value = _constraint_value(node.args[0]) - if not isinstance(value, str) or not value: - raise VerifyError("partial reduction must be a non-empty string") - return Partial(value) - - -def _parse_binding_set(node: ast.AST) -> list[tuple[str, Broadcast | Partial]]: - if not isinstance(node, ast.Set): - raise VerifyError("layout bindings must be a set") - out: list[tuple[str, Broadcast | Partial]] = [] - for item in node.elts: - if not isinstance(item, ast.BinOp) or not isinstance(item.op, ast.MatMult): - raise VerifyError('layout bindings must use `topology @ B()` or `P()`') - topology = _constraint_value(item.left) - if not isinstance(topology, str) or not topology: - raise VerifyError("layout binding topology must be symbolic") - if ( - isinstance(item.right, ast.Call) - and isinstance(item.right.func, ast.Name) - and item.right.func.id == "B" - and not item.right.args - and not item.right.keywords - ): - attr: Broadcast | Partial = Broadcast() - else: - attr = _parse_partial_value(item.right) - out.append((topology, attr)) - return out - - -def _parse_layout_constraint( - node: ast.AST, - resolve_extent, -) -> LayoutConstraint: - if not isinstance(node, ast.Tuple): - raise VerifyError("layout constraint must be a tuple") - dims_node = node - extras: tuple[ast.AST, ...] = () - if node.elts and isinstance(node.elts[0], ast.Tuple): - dims_node = node.elts[0] - extras = tuple(node.elts[1:]) - if len(extras) > 1: - raise VerifyError("layout constraint accepts one binding set") - if not dims_node.elts: - raise VerifyError("layout constraint cannot be empty") - shape: list[object] = [] - bindings: list[tuple[str, Split | Broadcast | Partial]] = [] - for index, item in enumerate(dims_node.elts): - if isinstance(item, ast.Name) and item.id == "_": - shape.append(_LAYOUT_WILDCARD) - continue - if isinstance(item, ast.Name) and item.id == "D": - raise VerifyError("layout broadcast must use a `{topology @ B()}` binding") - if isinstance(item, ast.BinOp) and isinstance(item.op, ast.MatMult): - extent = resolve_extent(item.left) - topology = _constraint_value(item.right) - if not isinstance(topology, str) or not topology: - raise VerifyError("layout topology binding must be symbolic") - shape.append(extent) - bindings.append((topology, Split(index))) - continue - shape.append(resolve_extent(item)) - if extras: - bindings.extend(_parse_binding_set(extras[0])) - if len({topology for topology, _ in bindings}) != len(bindings): - raise VerifyError("layout constraint cannot bind one topology more than once") - return LayoutConstraint(layout=Layout(shape=tuple(shape)), bindings=tuple(bindings)) - - -def _dim_operand_str(value: object) -> str: - """Name a loop-domain operand the way its author wrote it. - - A parsed operand carries its whole ``TensorType`` in its repr, which buries - the one thing the reader needs — which operand it was, and what kind of - thing it turned out to be. - """ - name = getattr(value, "name", None) - if isinstance(name, str) and name: - return f"{name} ({type(value).__name__})" - return f"{type(value).__name__}" - - -def _is_pass_body(stmts: list[ast.stmt]) -> bool: - """A dispatch-prototype body is exactly ``pass``. - - A dispatch-prototype body is exactly ``pass``. A ``pass`` mixed with any - other statement is rejected (it is not a partial body form). - """ - if not any(isinstance(s, ast.Pass) for s in stmts): - return False - if len(stmts) != 1: - raise VerifyError( - "@tilefoundry.func: `pass` must be the entire body — it declares a " - "dispatch prototype (signature only); mixing it with other " - "statements is not allowed" - ) - return True - - -def _resolve_return_type(node: ast.FunctionDef, closure, body_expr) -> TensorType: - if node.returns is not None: - return _resolve_tensor_type(node.returns, closure) - if body_expr is None: - raise VerifyError( - "@tilefoundry.func: a `pass` prototype must annotate its return type" - ) - - t = getattr(body_expr, "type", None) - if t is None: - raise VerifyError("@tilefoundry.func: cannot determine return_type") - return t - -class _HirBodyVisitor(BaseExprVisitor): - token = "hir" - resolves_module_callees = True - - def __init__( - self, env, closure, *, topo_ns=None, source_filename="", - in_module_body=False, - ): - super().__init__(env, closure, in_module_body=in_module_body) - self.topo_ns: dict[str, "Topology"] = topo_ns or {} - self.source_filename = source_filename - self.pending_constraints: dict[int, ScheduleConstraintMetadata] = {} - self._mesh_coordinate_cache: dict[tuple[int, int], Expr] = {} - - def _mesh_axis_node(self, node: ast.AST): - """Resolve ``mesh.axis`` when it names a lexical Mesh object.""" - if not isinstance(node, ast.Attribute) or not isinstance(node.value, ast.Name): - return None - mesh = self.env.lookup(node.value.id) - if not isinstance(mesh, Mesh): - return None - return mesh, _resolve_mesh_axis(mesh, node.attr) - - def _contains_mesh_coordinate(self, node: ast.AST) -> bool: - return any( - self._mesh_axis_node(candidate) is not None for candidate in ast.walk(node) - ) - - def _mesh_coordinate(self, node: ast.Attribute) -> Expr | None: - """Build or reuse the current rank-0 coordinate for a mesh axis. - - Arange is an invariant source; the explicit C-order stride keeps the - generated layout within the injective layout contract. - """ - resolved = self._mesh_axis_node(node) - if resolved is None: - return None - mesh, axis = resolved - extent = mesh.layout.shape[axis] - if not isinstance(extent, int) or isinstance(extent, bool): - raise VerifyError( - f"mesh coordinate {ast.unparse(node)!r} requires a concrete axis extent" - ) - cache_key = (id(mesh), axis) - cached = self._mesh_coordinate_cache.get(cache_key) - if cached is not None: - return cached - - vector = self._build_call( - Arange( - type=TensorType( - shape=(extent,), - dtype=DType.i64, - layout=None, - storage=StorageKind.GMEM, - ) - ), - (), - ) - attrs = tuple( - Split(axis=0) if mesh_axis == axis else Broadcast() - for mesh_axis in range(len(mesh.layout.shape)) - ) - layout = ShardLayout( - layout=Layout(shape=(extent,), strides=(1,)), - attrs=attrs, - mesh=mesh, - ) - placed = self._build_call( - Reshard(layout=layout, storage=StorageKind.RMEM), (vector,) - ) - local = self._build_call(Local(), (placed,)) - coordinate = self._build_call(Reshape(new_shape=()), (local,)) - self._mesh_coordinate_cache[cache_key] = coordinate - return coordinate - - def _resolve_static_attribute(self, owner, attr: str): - if isinstance(owner, Mesh): - _resolve_mesh_axis(owner, attr) - raise VerifyError( - f"mesh coordinate {attr!r} is a run-time Expr, not a static value" - ) - return super()._resolve_static_attribute(owner, attr) - - def visit_Attribute(self, node: ast.Attribute) -> Expr: - coordinate = self._mesh_coordinate(node) - if coordinate is not None: - return coordinate - return super().visit_Attribute(node) - - def _slicer_endpoint(self, node: ast.AST): - try: - return self._eval_static(node, allow_runtime_scalar=True) - except VerifyError: - return self.expr(node) - - - - - - def visit_body(self, stmts: list[ast.stmt]) -> Expr: - """Fold an HIR function body into its tail expression DAG. - - Nested function definitions are rejected syntactically across the whole - body. See [hir §1](docs/spec/hir.md#1-hir-expr-constructs) and - [parser §5](docs/spec/parser.md#5-hir-parser). - """ - for stmt in stmts: - for sub in ast.walk(stmt): - if isinstance(sub, (ast.FunctionDef, ast.AsyncFunctionDef)): - raise VerifyError( - "hir: nested function definition not allowed in an " - "@tilefoundry.func body (helper functions are " - "module/function-level definitions)" - ) - return self._visit_chain(stmts, 0) - - def _visit_chain( - self, stmts: list[ast.stmt], idx: int, require_return: bool = True - ) -> Expr | None: - """Fold a statement chain into a single tail ``Expr``. - - ``require_return=True`` (the function body) requires a terminal - ``return`` and raises when the chain runs out. A ``with Mesh(...)`` - suite is folded with ``require_return=False``: a setup-only suite that - carries no ``return`` yields ``None`` so the caller can continue folding - the post-``with`` tail in the outer frame. - """ - if idx >= len(stmts): - if require_return: - raise VerifyError("@tilefoundry.func body must end with `return`") - return None - node = stmts[idx] - if isinstance(node, ast.Return): - if node.value is None: - raise VerifyError( - f"{spec_ref_render(_HIR_FUNCTION)}: @tilefoundry.func return must carry a value" - ) - if isinstance(node.value, ast.Tuple): - return self._tuple_expr_expr(node.value) - return self.expr(node.value) - if isinstance(node, ast.Assign): - if len(node.targets) != 1: - raise VerifyError("hir: only single-target assignments supported in V1") - target = node.targets[0] - if isinstance(target, ast.Name): - tgt = target.id - bound = self._static_body_value(node.value) - if bound is not _NOT_STATIC: - - - self.env.define(tgt, bound) - return self._visit_chain(stmts, idx + 1, require_return) - rhs = self._assignment_rhs(node.value, tgt) - - self.env.define(tgt, rhs) - return self._visit_chain(stmts, idx + 1, require_return) - if isinstance(target, ast.Tuple): - - - - - - if not self._static_tuple_assign(target, node.value): - self._visit_tuple_assign(target, node.value) - return self._visit_chain(stmts, idx + 1, require_return) - raise VerifyError("hir: only single-target Name or Tuple assignments supported in V1") - if isinstance(node, ast.AnnAssign): - return self._visit_annotated_assignment(node, stmts, idx, require_return) - if isinstance(node, ast.With): - return self._visit_with(node, stmts, idx, require_return) - if isinstance(node, ast.Expr): - if isinstance(node.value, (ast.Yield, ast.YieldFrom)): - raise VerifyError( - "hir: `yield` is not an HIR statement — a @func body builds one value, " - "so return it instead of yielding it" - ) - raise VerifyError("hir: bare expression statement not allowed; use assign or return") - if isinstance(node, ast.For): - return self._visit_loop_for(node, stmts, idx, require_return) - raise VerifyError(f"hir: unsupported statement {type(node).__name__}") - - def _visit_annotated_assignment( - self, - node: ast.AnnAssign, - stmts: list[ast.stmt], - idx: int, - require_return: bool, - ) -> Expr | None: - if not isinstance(node.target, ast.Name): - raise VerifyError( - "where annotation target must be a bound plain Name; " - "subscripts are not annotation lvalues" - ) - if node.value is None: - target = self.env.lookup(node.target.id) - if not isinstance(target, Expr): - raise VerifyError( - f"where annotation target {node.target.id!r} is unresolved " - f"at {self.source_filename}:{node.lineno}:{node.col_offset}" - ) - else: - target = self._assignment_rhs(node.value, node.target.id) - self.env.define(node.target.id, target) - self._record_annotated_assignment(node, target) - return self._visit_chain(stmts, idx + 1, require_return) - - def _record_annotated_assignment(self, node: ast.AnnAssign, target: Expr) -> None: - metadata = self._parse_where_annotation(node.annotation, node) - if not isinstance(target.type, TensorType): - binding = get_metadata(target, BindingMetadata) - label = binding.name if binding is not None else self.source_filename - raise VerifyError( - f"where annotation requires a tensor-valued Expr at " - f"{label}:{node.lineno}:{node.col_offset + 1}" - ) - previous_metadata = get_metadata(target, ScheduleConstraintMetadata) - if previous_metadata is not None: - previous = previous_metadata.source_loc.describe() - current = metadata.source_loc.describe() - binding = get_metadata(target, BindingMetadata) - label = binding.name if binding is not None else "" - raise VerifyError( - f"duplicate where annotation for Expr {label!r} " - f"at {current}; first annotation at {previous}" - ) - self._attach_metadata(target, metadata) - - def _assignment_rhs(self, node: ast.AST, target_name: str) -> Expr: - """Parse an assignment RHS without turning a name alias into a node.""" - if isinstance(node, ast.Name): - value = self.env.lookup(node.id) - if isinstance(value, Expr): - return value - return self._maybe_autofill_binding( - self.expr_with_binding(node, target_name), target_name - ) - - def _parse_where_annotation( - self, annotation: ast.AST, node: ast.AnnAssign - ) -> ScheduleConstraintMetadata: - if not isinstance(annotation, ast.Call) or not isinstance( - annotation.func, ast.Name - ) or annotation.func.id != "where": - raise VerifyError( - "annotations must use `where(...)`; positional or other forms " - "are not supported" - ) - if annotation.args: - raise VerifyError("where(...) accepts keyword arguments only") - if not annotation.keywords: - raise VerifyError("where(...) cannot be empty") - source_loc = SourceLocation( - filename=self.source_filename, - line=node.lineno, - column=node.col_offset, - end_line=getattr(node, "end_lineno", None), - end_column=getattr(node, "end_col_offset", None), - ) - constraints: list[ScheduleConstraint] = [] - fields: set[str] = set() - for keyword in annotation.keywords: - if keyword.arg is None: - raise VerifyError("where(...) does not accept **kwargs") - if keyword.arg in fields: - raise VerifyError( - f"where(...) repeats keyword {keyword.arg!r} at " - f"{source_loc.describe()}" - ) - fields.add(keyword.arg) - if keyword.arg == "layout": - layout = _parse_layout_constraint( - keyword.value, self._resolve_layout_extent - ) - constraints.append( - dataclasses.replace( - layout, - source_loc=source_loc, - provenance=ConstraintProvenance.AUTHOR, - ) - ) - elif keyword.arg == "mesh": - try: - mesh = self._eval_static(keyword.value) - except (TypeError, ValueError, VerifyError) as exc: - raise VerifyError( - f"where mesh constraint could not be resolved at " - f"{source_loc.describe()}: {exc}" - ) from exc - constraints.append( - MeshConstraint( - mesh=mesh, - source_loc=source_loc, - provenance=ConstraintProvenance.AUTHOR, - ) - ) - elif keyword.arg == "storage": - try: - storage = self._eval_static(keyword.value) - except (TypeError, ValueError, VerifyError) as exc: - raise VerifyError( - f"where storage constraint could not be resolved at " - f"{source_loc.describe()}: {exc}" - ) from exc - constraints.append( - StorageConstraint( - storage=storage, - source_loc=source_loc, - provenance=ConstraintProvenance.AUTHOR, - ) - ) - else: - raise VerifyError( - f"where(...) has unknown field {keyword.arg!r}; use " - "layout=..., mesh=..., or storage=..." - ) - return ScheduleConstraintMetadata( - constraints=tuple(constraints), source_loc=source_loc - ) - - def _resolve_layout_extent(self, node: ast.AST) -> int | DimVar: - """Resolve one ``where`` shape entry to a concrete ``int``/``DimVar``. - - Resolve one ``where(layout=...)`` shape entry to a concrete - ``int``/``DimVar`` (the ``_`` wildcard is handled by the caller before - this is reached). A literal resolves directly; a symbolic name - resolves through the lexical env / authoring closure, matching - ``_eval_static``'s ``Name`` resolution. - """ - if isinstance(node, ast.Constant): - value = node.value - if isinstance(value, bool) or not isinstance(value, int): - raise VerifyError( - "layout dimensions must use `_`, an integer, or a " - "symbolic extent with `@ topology`" - ) - return value - if isinstance(node, ast.Name): - try: - resolved = self._eval_static(node) - except (TypeError, ValueError, VerifyError) as exc: - raise VerifyError( - f"where layout extent {node.id!r} could not be resolved: {exc}" - ) from exc - if isinstance(resolved, bool) or not isinstance(resolved, (int, DimVar)): - raise VerifyError( - f"where layout extent {node.id!r} must resolve to an " - f"int or DimVar, got {type(resolved).__name__}" - ) - return resolved - raise VerifyError( - "layout dimensions must use `_`, an integer, or a symbolic " - "extent with `@ topology`" - ) - - def _resolve_loop_bound(self, node: ast.AST): - """Resolve a ``tile`` / ``range`` bound to an ``int``, ``DimVar``, or dim ``Expr``. - - Resolve a ``tile`` / ``range`` bound (extent / step / start) to an - ``int``, ``DimVar``, or dim ``Expr``. - - Unlike ``_eval_static`` (which only folds numeric constants), a - ``BinOp`` whose operands reach a ``DimVar`` builds a dim expression via - ``simplify_dim`` (e.g. ``C // NUM_SPLITS`` → ``DimFloorDiv(C, N)``). The - IR / evaluator already resolve dim-expression loop bounds; this lets the - DSL surface write them. - """ - if isinstance(node, ast.BinOp): - op = _AST_DIM_OPS.get(type(node.op)) - if op is None: - raise VerifyError( - f"loop bound: unsupported operator {type(node.op).__name__} " - f"(use + - * // %)" - ) - left = self._resolve_loop_bound(node.left) - right = self._resolve_loop_bound(node.right) - - - if isinstance(left, int) and not isinstance(left, bool) and \ - isinstance(right, int) and not isinstance(right, bool): - return self._eval_static(node) - return simplify_dim(op, (left, right)) - return self._eval_static(node) - - def _visit_loop_for(self, node: ast.For, stmts, idx, require_return: bool = True): - """Lower tile or range loops to a grid region and continue the chain. - - Both forms share start, extent, and step. Tile loops bind a range slice - for indexed use; range loops bind a scalar induction variable. Neither - form is unrolled. - See [parser §1.7](docs/spec/parser.md#17-for-i-in-tile--for-i-in-range-hir-only). - """ - grid = self._build_grid_for(node) - if idx + 1 < len(stmts): - return self._visit_chain(stmts, idx + 1, require_return) - - - - - return grid if require_return else None - - def _build_grid_for(self, node: ast.For) -> Expr: - """Build a grid region and rebind loop-carried names in this frame. - - Assigning an outer name creates a phi and yield; one carry maps to the - grid and multiple carries to tuple projections after the loop. Nested - loops compose recursively. The body accepts assignments and nested loops - but no return; the caller processes sibling statements. - See [parser §5.1](docs/spec/parser.md#51-gridregionexpr-carry-out-lifting). - """ - if not isinstance(node.iter, ast.Call) or not isinstance(node.iter.func, ast.Name): - raise VerifyError("hir For: iter must be a `tile(...)` or `range(...)` call") - loop_kind = node.iter.func.id - if loop_kind not in ("tile", "range"): - raise VerifyError( - f"hir For: iter must be `tile(...)` or `range(...)`, got " - f"{loop_kind!r}" - ) - if node.iter.keywords: - raise VerifyError( - f"{loop_kind}() does not accept keyword args " - "(positional-only at the IR level)" - ) - if not isinstance(node.target, ast.Name): - raise VerifyError("hir For: target must be a Name") - iv = Var(type=TensorType.scalar(DType.i64), name=node.target.id) - - loop_args = node.iter.args - iv_binding: Expr | slice - if loop_kind == "range": - - - if len(loop_args) == 1: - start, extent, step = 0, self._resolve_loop_bound(loop_args[0]), 1 - elif len(loop_args) == 2: - start = self._resolve_loop_bound(loop_args[0]) - extent = self._resolve_loop_bound(loop_args[1]) - step = 1 - elif len(loop_args) == 3: - start = self._resolve_loop_bound(loop_args[0]) - extent = self._resolve_loop_bound(loop_args[1]) - step = self._resolve_loop_bound(loop_args[2]) - else: - raise VerifyError( - f"range() takes 1-3 arguments (stop | start, stop[, step]), " - f"got {len(loop_args)}" - ) - iv_binding = iv - else: - start = 0 - if len(loop_args) == 2: - extent = self._resolve_loop_bound(loop_args[0]) - step = self._resolve_loop_bound(loop_args[1]) - step_expr = self._constant_expr(step) if isinstance(step, int) else step - iv_binding = slice( - iv, - simplify_dim(DimAdd, (iv, step_expr)), - 1, - ) - elif len(loop_args) == 1: - raise VerifyError( - "tile(extent) is not supported; use range(extent) for " - "scalar iteration" - ) - else: - raise VerifyError( - f"tile() takes 2 arguments (extent, step), got {len(loop_args)}" - ) - start, extent, step = ( - normalize_dim(value) for value in (start, extent, step) - ) - if not (is_dim_expr(start) and is_dim_expr(extent) and is_dim_expr(step)): - offending = ", ".join( - f"{label}={_dim_operand_str(value)}" - for label, value in (("start", start), ("extent", extent), ("step", step)) - if not is_dim_expr(value) - ) - raise VerifyError( - f"{loop_kind}(): start / extent / step must be a dim expression " - f"(int / DimVar / dim-op Expr), and {offending} is not one" - ) - - - - - - carry_names: list[str] = [] - carry_seen: set[str] = set() - - def _scan_carries(body_stmts: list[ast.stmt]) -> None: - for stmt in body_stmts: - if isinstance(stmt, ast.Assign) and len(stmt.targets) == 1 \ - and isinstance(stmt.targets[0], ast.Name): - name = stmt.targets[0].id - if name not in carry_seen: - carry_seen.add(name) - if isinstance(self.env.lookup(name), Expr): - carry_names.append(name) - elif isinstance(stmt, ast.For): - _scan_carries(stmt.body) - - _scan_carries(node.body) - - - - phi_vars: list[Var] = [] - init_exprs: list[Expr] = [] - for name in carry_names: - outer_expr = self.env.lookup(name) - phi_vars.append(Var(type=outer_expr.type, name=name)) - init_exprs.append(outer_expr) - - - self.env.push_frame() - if loop_kind == "range": - self._scalar_index_ids.add(id(iv)) - else: - - - - self._tile_windows[id(iv)] = (extent, step) - try: - self.env.define(node.target.id, iv_binding) - for cname, phi in zip(carry_names, phi_vars): - self.env.define(cname, phi) - body_expr = self._visit_grid_body(node.body) - - - yield_exprs: list[Expr] = [] - for cname in carry_names: - v = self.env.lookup(cname) - if not isinstance(v, Expr): - raise VerifyError( - f"tile-for: carry name {cname!r} did not resolve to " - f"an Expr at end of body (got {type(v).__name__})" - ) - yield_exprs.append(v) - finally: - self._scalar_index_ids.discard(id(iv)) - self._tile_windows.pop(id(iv), None) - self.env.pop_frame() - - if not carry_names: - return GridRegionExpr( - type=body_expr.type, - induction_var=iv, - carried_args=(), - init_args=(), - body=body_expr, - yield_values=(), - start=start, - extent=extent, - step=step, - ) - - - - if len(carry_names) == 1: - grid_type = phi_vars[0].type - else: - grid_type = TupleType(fields=tuple(p.type for p in phi_vars)) - grid = GridRegionExpr( - type=grid_type, - induction_var=iv, - carried_args=tuple(phi_vars), - init_args=tuple(init_exprs), - body=body_expr, - yield_values=tuple(yield_exprs), - start=start, - extent=extent, - step=step, - ) - - if len(carry_names) == 1: - self.env.define(carry_names[0], grid) - else: - - for i, cname in enumerate(carry_names): - proj = self._build_call(TupleGetItem(index=i), (grid,)) - self.env.define(cname, proj) - return grid - - def _visit_grid_body(self, body_stmts: list[ast.stmt]) -> Expr: - """Process tile-for body statements and return the final body Expr. - - Body must be a sequence of Assigns (single Name or Tuple targets); - ``return`` / bare expression statements / for / with are rejected. - Returns the last bound RHS Expr (or first Tuple-unpack RHS if the - last stmt is a Tuple unpack). - """ - last_expr: Expr | None = None - for stmt in body_stmts: - if isinstance(stmt, ast.Return): - raise VerifyError( - "hir tile-for body must not contain `return` " - "(use a final assignment to the carry variable instead)" - ) - if isinstance(stmt, ast.Expr): - raise VerifyError( - "hir tile-for body: bare expression statement not allowed" - ) - if isinstance(stmt, ast.With): - raise VerifyError( - "hir tile-for body: nested With not supported in v1" - ) - if isinstance(stmt, ast.For): - - - - - last_expr = self._build_grid_for(stmt) - continue - if isinstance(stmt, ast.AugAssign): - raise VerifyError( - "hir tile-for body: augmented assignment (+= etc.) " - "not supported in v1; rewrite as `x = add(x, ...)`" - ) - if isinstance(stmt, ast.Assign): - if len(stmt.targets) != 1: - raise VerifyError( - "hir tile-for body: only single-target assignments " - "supported in v1" - ) - target = stmt.targets[0] - if isinstance(target, ast.Name): - rhs = self._assignment_rhs(stmt.value, target.id) - self.env.define(target.id, rhs) - last_expr = rhs - continue - if isinstance(target, ast.Tuple): - - - - rhs = self._visit_tuple_assign(target, stmt.value) - last_expr = rhs - continue - raise VerifyError( - "hir tile-for body: assignment target must be Name or Tuple" - ) - raise VerifyError( - f"hir tile-for body: unsupported statement {type(stmt).__name__}" - ) - if last_expr is None: - raise VerifyError( - "hir tile-for body must contain at least one assignment" - ) - return last_expr - - def _static_body_value(self, node: ast.AST): - """A body-local name's compile-time value. - - A body-local name's compile-time value — a number or a list of Exprs — - or ``_NOT_STATIC`` when the right-hand side belongs to the IR. - """ - number = self._static_number(node) - if number is not None: - return number - items = self._static_expr_list(node) - return _NOT_STATIC if items is None else items - - def _static_tuple_assign(self, target: ast.Tuple, value: ast.AST) -> bool: - """Bind ``a, b = , ``, reporting whether it applied.""" - if not isinstance(value, ast.Tuple) or len(target.elts) != len(value.elts): - return False - if not all(isinstance(elt, ast.Name) for elt in target.elts): - return False - numbers = [self._static_number(el) for el in value.elts] - if any(number is None for number in numbers): - return False - for elt, number in zip(target.elts, numbers): - self.env.define(elt.id, number) - return True - - def _visit_tuple_assign(self, target: ast.Tuple, value: ast.AST) -> Expr: - """Tuple-unpack inside tile body (mirrors _visit_chain Tuple branch).""" - names: list[str] = [] - for elt in target.elts: - if not isinstance(elt, ast.Name): - raise VerifyError( - "hir: tuple unpack targets must all be plain names" - ) - names.append(elt.id) - rhs = self.expr_with_binding(value, ", ".join(names)) - if not isinstance(rhs.type, TupleType): - raise VerifyError( - f"hir: tuple unpack requires RHS of TupleType, " - f"got {type(rhs.type).__name__}" - ) - if len(names) != len(rhs.type.fields): - raise VerifyError( - f"hir: tuple unpack arity mismatch — RHS has " - f"{len(rhs.type.fields)} fields, LHS binds {len(names)} names" - ) - rhs = self._maybe_autofill_binding_default(rhs) - last_item: Expr = rhs - for i, nm in enumerate(names): - item = self._build_call(TupleGetItem(index=i), (rhs,)) - item = self._maybe_autofill_binding(item, nm) - self.env.define(nm, item) - last_item = item - return last_item - - def _visit_with(self, node: ast.With, stmts, idx, require_return: bool = True): - """Parse an active mesh context with suite-local mesh binding. - - Ordinary suite bindings escape to the function frame, while the mesh - alias does not. See [parser §1.6](docs/spec/parser.md#16-with-mesh-as-m). - """ - if len(node.items) != 1: - raise VerifyError("hir: only single-item `with` supported") - item = node.items[0] - if item.optional_vars is None or not isinstance(item.optional_vars, ast.Name): - raise VerifyError("hir: `with Mesh(...) as name` requires a single Name binding") - - - mesh = self._resolve_mesh_context(item.context_expr) - if not isinstance(mesh, Mesh): - raise VerifyError( - f"hir: `with` context must evaluate to a Mesh " - f"({spec_ref_render(_PARSER_MESH)}), got {type(mesh).__name__}" - ) - name = item.optional_vars.id - - - - - - - - - - - - self.env.push_frame() - scopes = self._mesh_scopes - self._mesh_scopes = (*scopes, mesh) - try: - self.env.define(name, mesh) - body_result = self._visit_chain(list(node.body), 0, require_return=False) - finally: - self._mesh_scopes = scopes - suite_frame = self.env.pop_frame() - for bound_name, bound_value in suite_frame.items(): - if bound_name != name: - self.env.define(bound_name, bound_value) - if body_result is not None: - return body_result - return self._visit_chain(stmts, idx + 1, require_return) - - def _resolve_mesh_context(self, node: ast.AST) -> Mesh: - """Resolve a ``Mesh(...)`` call with tuple topology-name sugar.""" - if not isinstance(node, ast.Call): - return self._eval_static(node) - - - - - - def _eval_mesh_arg(arg_node: ast.AST, *, is_layout_slot: bool = True): - if is_layout_slot and _is_tuple_sugar(arg_node): - return parse_sugar(arg_node, Layout, closure=self.closure) - return self._eval_static(arg_node) - - def _resolve_string_topology(name: str) -> object: - obj = self.topo_ns.get(name) - if obj is None: - raise VerifyError( - f"topology {name!r} not declared in function/module topologies " - f"(available: {list(self.topo_ns.keys())})" - ) - return obj - - if any(keyword.arg in {"topology", "topologies"} for keyword in node.keywords): - raise VerifyError( - 'hir: Mesh requires a tuple of declared topology names, ' - 'for example Mesh(("cta",), layout=(128,))' - ) - if ( - not node.args - or not isinstance(node.args[0], ast.Tuple) - or not node.args[0].elts - or not all( - isinstance(entry, ast.Constant) and isinstance(entry.value, str) - for entry in node.args[0].elts - ) - ): - raise VerifyError( - 'hir: Mesh requires a tuple of declared topology names, ' - 'for example Mesh(("cta",), layout=(128,))' - ) - - mesh_fn = self._eval_static(node.func) - topologies = tuple(_resolve_string_topology(entry.value) for entry in node.args[0].elts) - pos_args = [topologies] - for index, argument in enumerate(node.args[1:], start=1): - pos_args.append(_eval_mesh_arg(argument, is_layout_slot=(index == 1))) - pos_kw = { - keyword.arg: _eval_mesh_arg( - keyword.value, is_layout_slot=(keyword.arg == "layout") - ) - for keyword in node.keywords - } - return mesh_fn(*pos_args, **pos_kw) - -__all__ = ["parse_func"] diff --git a/src/tilefoundry/parser/parser_visitor.py b/src/tilefoundry/parser/parser_visitor.py new file mode 100644 index 00000000..6d85c4d7 --- /dev/null +++ b/src/tilefoundry/parser/parser_visitor.py @@ -0,0 +1,59 @@ +"""Function parser entry points for the AST pattern prototype.""" + +from __future__ import annotations + +import ast +import inspect +import textwrap +from types import FunctionType +from typing import Any + +from .ast_pattern import ( + FuncParserContext, + FunctionPattern, + MatchContext, + parse_node, +) + + +class FuncParserVisitor: + """Walk one authored function from its selected root pattern.""" + + def __init__(self, context: FuncParserContext): + self.context = context + self.root_pattern = FunctionPattern() + + def visit(self, node: ast.AST) -> Any: + return parse_node( + self.root_pattern, node, MatchContext.from_function(self.context) + ) + + def visit_function(self, node: ast.FunctionDef) -> Any: + return self.visit(node) + + +def _extract_function_def(fn: FunctionType) -> ast.FunctionDef: + if not isinstance(fn, FunctionType): + raise TypeError(f"parse_function expects a Python function, got {type(fn).__name__}") + try: + source_lines, start_line = inspect.getsourcelines(fn) + source = textwrap.dedent("".join(source_lines)) + except (OSError, TypeError) as error: + raise TypeError("parse_function requires authored source for the function") from error + module = ast.parse(source, filename=inspect.getsourcefile(fn) or "") + ast.increment_lineno(module, start_line - 1) + functions = [node for node in ast.walk(module) if isinstance(node, ast.FunctionDef)] + if len(functions) != 1: + raise TypeError("parse_function requires exactly one authored FunctionDef") + return functions[0] + + +def parse_function(fn: FunctionType, context: FuncParserContext) -> Any: + """Parse one authored Python function using its typed parser context.""" + return FuncParserVisitor(context).visit_function(_extract_function_def(fn)) + + +__all__ = [ + "FuncParserVisitor", + "parse_function", +] diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py new file mode 100644 index 00000000..a1512a8a --- /dev/null +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -0,0 +1,3922 @@ +"""Concrete executable AST Pattern nodes. + +Pattern matching and construction live together here; shared parser state, +rules, and runtime helpers remain in ast_pattern. +""" + +from __future__ import annotations + +import ast +import dataclasses +import enum +import operator +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, ClassVar + +from tilefoundry.ir.constraints import ( + ConstraintProvenance, + LayoutConstraint, + MeshConstraint, + ScheduleConstraintMetadata, + SourceLocation, + StorageConstraint, +) +from tilefoundry.ir.constraints.layout import _LAYOUT_WILDCARD +from tilefoundry.ir.core import ( + BindingMetadata, + ExecutionDomainMetadata, + get_metadata, + replace_metadata, +) +from tilefoundry.ir.tir.launch import launch_call +from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.shard import Broadcast, Layout, Partial, Split + +from .ast_pattern import ( + _BINARY_OPERATORS, + _RETURN_TYPE, + _TYPE_INFER_CONTEXT, + _UNARY_OPERATORS, + AstChild, + AstMatch, + AstNodePattern, + AstPattern, + AstRule, + BindPattern, + BranchPattern, + CanonicalDTypeRule, + CapturePattern, + ChildPattern, + ChoicePattern, + ConditionPattern, + ElementPattern, + FieldPattern, + FuncParserContext, + FunctionRole, + LayoutPositionRule, + LayoutShapeRule, + LazyPattern, + LiteralPattern, + LoopFrame, + MatchContext, + OptionalPattern, + ParseError, + PatternFailure, + PredicatePattern, + ReferencePattern, + RepeatPattern, + SequencePattern, + ShapeDimRule, + ShapeTupleRule, + StorageValueRule, + TensorLayoutStorageRule, + TensorPositionRule, + _constant, + _infer_call, + _resolve_reference, + _slice_size, + attach_authored_metadata, + runtime, +) + + +class DimExprPattern(ElementPattern): + element_name = "dim_expr" + syntax = LazyPattern( + lambda: AstNodePattern( + ast.expr, + ChoicePattern( + BranchPattern( + "dim_literal", + AstNodePattern( + ast.Constant, + PredicatePattern( + "integer-literal", + lambda node, context: ( + isinstance(node.value, int) and not isinstance(node.value, bool) + ), + ), + CapturePattern("value", lambda node, context: node.value), + ), + pattern_id="dim.literal", + ), + BranchPattern( + "dim_name", + AstNodePattern( + ast.Name, + FieldPattern("id", CapturePattern("name", lambda value, context: value)), + ), + pattern_id="dim.name", + ), + BranchPattern( + "dim_reference", + AstNodePattern(ast.Attribute), + pattern_id="dim.reference", + ), + BranchPattern( + "dim_binary", + AstNodePattern( + ast.BinOp, + FieldPattern( + "op", + ChoicePattern( + AstNodePattern(ast.Add), + AstNodePattern(ast.Sub), + AstNodePattern(ast.Mult), + AstNodePattern(ast.FloorDiv), + AstNodePattern(ast.Mod), + ), + ), + FieldPattern( + "left", + ChildPattern("left", lambda: DimExprPattern(), "dim_expr", "left"), + ), + FieldPattern( + "right", + ChildPattern("right", lambda: DimExprPattern(), "dim_expr", "right"), + ), + CapturePattern( + "operator", + lambda node, context: _BINARY_OPERATORS[type(node.op)], + ), + ), + pattern_id="dim.binary", + ), + BranchPattern( + "dim_call", + AstNodePattern( + ast.Call, + FieldPattern( + "func", + ChoicePattern( + AstNodePattern( + ast.Name, + FieldPattern( + "id", + CapturePattern("callee", lambda value, context: value), + ), + ), + AstNodePattern(ast.Attribute), + ), + ), + FieldPattern( + "args", + RepeatPattern( + ChildPattern( + "arg_{index}", + lambda: DimExprPattern(), + "dim_expr", + "argument", + ) + ), + ), + FieldPattern("keywords", SequencePattern()), + ), + pattern_id="dim.call", + ), + ), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "dim_literal": + return match.captures["value"] + if match.branch_id == "dim_name": + return _resolve_reference(match.node, context) + if match.branch_id == "dim_reference": + return _resolve_reference(match.node, context) + if match.branch_id == "dim_binary": + try: + return match.captures["operator"](children["left"], children["right"]) + except (TypeError, ValueError, ZeroDivisionError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + if match.branch_id == "dim_call": + callee = _resolve_reference(match.node.func, context) + args = tuple(value for name, value in children.items() if name.startswith("arg_")) + if not callable(callee): + raise ParseError.from_node( + match.node, context, "dimension call target is not callable" + ) + try: + return callee(*args) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = (ShapeDimRule(),) + + +class ShapePattern(ElementPattern): + element_name = "shape" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "tuple_children", + AstNodePattern( + ast.Tuple, + FieldPattern("ctx", AstNodePattern(ast.Load)), + PredicatePattern( + "shape-tuple", + lambda node, context: ( + not any(isinstance(item, ast.Tuple) for item in node.elts) + ), + ), + FieldPattern( + "elts", + RepeatPattern( + ChildPattern( + "dim_{index}", + lambda: DimExprPattern(), + "tensor_dim_expr", + "dim_expr", + ) + ), + ), + ), + pattern_id="tensor.shape", + ), + BranchPattern( + "shape_reference", + ChoicePattern(AstNodePattern(ast.Name), AstNodePattern(ast.Attribute)), + pattern_id="tensor.shape.reference", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "shape_reference": + return _resolve_reference(match.node, context) + return tuple(children.values()) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = (ShapeTupleRule(),) + + +class DTypePattern(ElementPattern): + element_name = "dtype" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "dtype_literal", + AstNodePattern( + ast.Constant, + FieldPattern("value", LiteralPattern(value_type=str)), + ), + pattern_id="tensor.dtype.literal", + ), + BranchPattern( + "dtype_reference", + ReferencePattern(), + pattern_id="tensor.dtype.reference", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "dtype_literal": + try: + return runtime.DType.from_name(match.node.value) + except ValueError as error: + raise ParseError.from_node(match.node, context, str(error)) from error + elif match.branch_id == "dtype_reference": + if isinstance(match.node, ast.Name) and match.node.id in runtime.DType._members(): + return runtime.DType.from_name(match.node.id) + value = _resolve_reference(match.node, context) + if isinstance(value, str): + try: + return runtime.DType.from_name(value) + except ValueError as error: + raise ParseError.from_node(match.node, context, str(error)) from error + return value + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = (CanonicalDTypeRule(),) + + +class ExplicitLayoutPattern(ElementPattern): + element_name = "explicit_layout" + syntax = LazyPattern( + lambda: BranchPattern( + "explicit_layout", + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + SequencePattern( + AstNodePattern( + ast.Tuple, + ChoicePattern( + ConditionPattern( + "active Mesh", + lambda node, context: ( + context.function is not None + and bool(context.function.state.mesh_stack) + ), + ChildPattern( + "shape", + lambda: TensorShapeLayoutPattern(), + "layout_shape", + "layout_shape", + ), + ), + ChildPattern( + "shape", + lambda: ShapePattern(), + "layout_shape", + "layout_shape", + ), + ), + ), + AstNodePattern( + ast.Tuple, + ChildPattern( + "strides", + lambda: ShapePattern(), + "layout_strides", + "layout_strides", + ), + ), + ), + ), + ), + pattern_id="tensor.layout.explicit", + ) + ) + + @staticmethod + def construct(match, children, context): + shape_or_layout = children["shape"] + strides = children["strides"] + if isinstance(shape_or_layout, runtime.ShardLayout): + shape = shape_or_layout.layout.shape + else: + shape = shape_or_layout + if len(shape) != len(strides): + raise ParseError.from_node(match.node, context, "layout shape/stride rank mismatch") + layout = runtime.Layout(shape=shape, strides=strides) + if isinstance(shape_or_layout, runtime.ShardLayout): + return runtime.ShardLayout( + layout=layout, + attrs=shape_or_layout.attrs, + mesh=shape_or_layout.mesh, + ) + return layout + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + LayoutShapeRule(), + LayoutPositionRule(), + ) + + +class PlainLayoutPattern(ElementPattern): + element_name = "plain_layout" + syntax = LazyPattern( + lambda: BranchPattern( + "plain_layout", + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + AstNodePattern( + ast.expr, + PredicatePattern( + "layout-extent", + lambda node, context: not isinstance(node, (ast.Tuple, ast.Slice)), + ), + ChildPattern( + "extent_{index}", + lambda: DimExprPattern(), + "layout_extent", + "layout_extent", + ), + ) + ), + ), + ), + pattern_id="tensor.layout.literal", + ) + ) + + @staticmethod + def construct(match, children, context): + shape = tuple(children.values()) + layout = runtime.Layout( + shape=shape, + strides=runtime.c_order_strides(shape, mul=operator.mul), + ) + if ( + context.situation != "mesh_layout" + and context.function is not None + and context.function.state.mesh_stack + ): + mesh = context.function.state.mesh_stack[-1] + return runtime.ShardLayout( + layout=layout, + attrs=tuple(runtime.Broadcast() for _ in mesh.layout.shape), + mesh=mesh, + ) + return layout + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + LayoutShapeRule(), + LayoutPositionRule(), + ) + + +class MeshAxisPattern(ElementPattern): + element_name = "mesh_axis" + syntax = LazyPattern( + lambda: BranchPattern( + "mesh_axis", + ChoicePattern( + AstNodePattern(ast.Name), + AstNodePattern( + ast.Attribute, + FieldPattern("value", AstNodePattern(ast.Name)), + ), + ), + pattern_id="tensor.layout.mesh_axis", + ) + ) + + @staticmethod + def construct(match, children, context): + node = match.node + if isinstance(node, ast.Name): + binding = node.id + axis_name = None + else: + binding = node.value.id + axis_name = node.attr + mesh = context.lexical_scope.lookup(binding) + if mesh is None: + try: + mesh = _resolve_reference(node.value, context) + except ParseError: + mesh = None + if not isinstance(mesh, runtime.Mesh): + raise ParseError.from_node(node, context, f"{binding!r} is not an active Mesh") + if axis_name is None: + if len(mesh.layout.shape) != 1: + raise ParseError.from_node( + node, context, "bare Mesh placement requires a one-axis mesh" + ) + return mesh, 0 + try: + axis = mesh.names.index(axis_name) + except ValueError as error: + raise ParseError.from_node( + node, context, f"Mesh {binding!r} has no axis {axis_name!r}" + ) from error + return mesh, axis + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class PlacedLayoutPattern(ElementPattern): + element_name = "placed_layout" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + ChoicePattern( + AstNodePattern( + ast.BinOp, + FieldPattern("op", AstNodePattern(ast.MatMult)), + FieldPattern("left", AstNodePattern(ast.expr)), + FieldPattern( + "right", + ChoicePattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + MeshAxisPattern(), + minimum=1, + ), + ), + ), + MeshAxisPattern(), + ), + ), + ), + DimExprPattern(), + ) + ), + ), + ), + PlacedLayoutPattern._bind, + ) + ) + + @staticmethod + def _placement_parts(node: ast.AST) -> tuple[ast.AST, tuple[ast.AST, ...]] | None: + """Flatten ``extent @ axis @ axis`` into one extent and its axes.""" + if not isinstance(node, ast.BinOp) or not isinstance(node.op, ast.MatMult): + return None + left = PlacedLayoutPattern._placement_parts(node.left) + if left is None: + extent = node.left + axes: tuple[ast.AST, ...] = () + else: + extent, axes = left + right_axes = tuple(node.right.elts) if isinstance(node.right, ast.Tuple) else (node.right,) + return extent, (*axes, *right_axes) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.Tuple) + children: list[AstChild] = [] + bindings: list[tuple[str, int]] = [] + found_placement = False + for tensor_axis, item in enumerate(node.elts): + placement = PlacedLayoutPattern._placement_parts(item) + extent_node = item + axis_nodes: tuple[ast.AST, ...] = () + if placement is not None: + found_placement = True + extent_node, axis_nodes = placement + extent_context = context.child(situation="layout_extent", role="layout_extent") + if DimExprPattern().match(extent_node, extent_context) is None: + return None + children.append( + AstChild( + f"extent_{tensor_axis}", + DimExprPattern(), + extent_node, + "layout_extent", + "layout_extent", + ) + ) + for mesh_axis, axis_node in enumerate(axis_nodes): + axis_context = context.child(situation="mesh_axis", role="mesh_axis") + if MeshAxisPattern().match(axis_node, axis_context) is None: + return None + child_name = f"binding_{tensor_axis}_{mesh_axis}" + bindings.append((child_name, tensor_axis)) + children.append( + AstChild( + child_name, + MeshAxisPattern(), + axis_node, + "mesh_axis", + "mesh_axis", + ) + ) + if not found_placement: + return None + return dataclasses.replace( + matched, + pattern_id="tensor.layout.placed", + branch_id="placed_layout", + captures={ + **matched.captures, + "rank": len(node.elts), + "bindings": tuple(bindings), + }, + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + rank = match.captures["rank"] + shape = tuple(children[f"extent_{axis}"] for axis in range(rank)) + bindings = tuple( + (*children[child_name], tensor_axis) + for child_name, tensor_axis in match.captures["bindings"] + ) + referenced_ids = {id(mesh) for mesh, _, _ in bindings} + if context.function is None: + raise ParseError.from_node( + match.node, context, "placed layout requires function context" + ) + meshes = tuple( + mesh for mesh in context.function.state.mesh_stack if id(mesh) in referenced_ids + ) + if len(meshes) != len(referenced_ids): + meshes = tuple(dict.fromkeys(mesh for mesh, _, _ in bindings)) + if len(meshes) != len(referenced_ids): + raise ParseError.from_node( + match.node, context, "placement references an inactive Mesh" + ) + mesh = meshes[0] if len(meshes) == 1 else runtime.composed(meshes) + source_offsets: dict[int, int] = {} + offset = 0 + for source in meshes: + source_offsets[id(source)] = offset + offset += len(source.layout.shape) + attrs: list[object] = [runtime.Broadcast() for _ in mesh.layout.shape] + for source, source_axis, tensor_axis in bindings: + target_axis = source_offsets[id(source)] + source_axis + if not isinstance(attrs[target_axis], runtime.Broadcast): + raise ParseError.from_node(match.node, context, "mesh axis is bound more than once") + attrs[target_axis] = runtime.Split(tensor_axis) + try: + canonical = runtime.canonical_shard_layout(shape, mesh, tuple(attrs)) + return runtime.ShardLayout( + layout=runtime.Layout(shape=canonical.layout.shape, strides=None), + attrs=canonical.attrs, + mesh=canonical.mesh, + ) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + LayoutShapeRule(), + LayoutPositionRule(), + ) + + +class LayoutPattern(ElementPattern): + element_name = "layout" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "none", + AstNodePattern( + ast.Constant, + FieldPattern("value", LiteralPattern(None)), + ), + pattern_id="tensor.layout.none", + ), + BranchPattern( + "layout_reference", + ReferencePattern(), + pattern_id="tensor.layout.reference", + ), + BranchPattern( + "identity", + ConditionPattern( + "layout call", + lambda node, context: isinstance(node, ast.Call), + ChildPattern( + "value", + lambda: StaticCallPattern(), + "static_layout", + "layout", + ), + ), + pattern_id="tensor.layout.call", + ), + ExplicitLayoutPattern(), + PlacedLayoutPattern(), + PlainLayoutPattern(), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "none": + return None + elif match.branch_id == "layout_reference": + return context.resolve_static(match.node, runtime.LayoutBase) + elif match.branch_id == "identity": + return children["value"] + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + LayoutShapeRule(), + LayoutPositionRule(), + ) + + +class StoragePattern(ElementPattern): + element_name = "storage" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "storage", + AstNodePattern( + ast.Constant, + FieldPattern("value", LiteralPattern(value_type=str)), + ), + pattern_id="tensor.storage.literal", + ), + BranchPattern( + "storage", + ReferencePattern(), + pattern_id="tensor.storage.reference", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + if isinstance(match.node, ast.Name) and match.node.id in { + str(k) for k in runtime.StorageKind + }: + raw = match.node.id + elif isinstance(match.node, ast.Constant): + raw = match.node.value + else: + raw = _resolve_reference(match.node, context) + try: + return runtime.resolve_storage(raw) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = (StorageValueRule(),) + + +class TensorOptionalSlotPattern(ElementPattern): + element_name = "tensor_optional_slot" + syntax = LazyPattern( + lambda: ChoicePattern( + ConditionPattern( + "role == layout", + lambda node, context: context.role == "layout", + LayoutPattern(), + ), + ConditionPattern( + "role == storage", + lambda node, context: context.role == "storage", + StoragePattern(), + ), + ConditionPattern( + "role == layout_or_storage and value is layout", + lambda node, context: ( + context.role == "layout_or_storage" + and TensorOptionalSlotPattern._slot_kind(node, context) == "layout" + ), + LayoutPattern(), + ), + ConditionPattern( + "role == layout_or_storage and value is storage", + lambda node, context: ( + context.role == "layout_or_storage" + and TensorOptionalSlotPattern._slot_kind(node, context) == "storage" + ), + StoragePattern(), + ), + ) + ) + + @staticmethod + def _slot_kind(node: object, context: MatchContext) -> str | None: + if isinstance(node, ast.Constant): + if node.value is None: + return "layout" + if isinstance(node.value, str): + return "storage" + if isinstance(node, (ast.Tuple, ast.Call)): + return "layout" + if isinstance(node, (ast.Name, ast.Attribute)): + try: + value = _resolve_reference(node, context) + except ParseError: + if isinstance(node, ast.Name) and node.id in { + str(item) for item in runtime.StorageKind + }: + return "storage" + raise + if isinstance(value, runtime.StorageKind): + return "storage" + if isinstance(value, runtime.LayoutBase): + return "layout" + return None + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class TensorShapeLayoutPattern(ElementPattern): + element_name = "tensor_shape_layout" + syntax = LazyPattern( + lambda: ChoicePattern( + PlacedLayoutPattern(), + ShapePattern(), + ) + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class TensorPattern(ElementPattern): + element_name = "tensor" + syntax = LazyPattern( + lambda: BranchPattern( + "tensor", + AstNodePattern( + ast.Subscript, + FieldPattern( + "value", + AstNodePattern( + ast.expr, + PredicatePattern( + "tensor-head", + lambda node, context: ( + TensorPattern._head(node) in {"Tensor", "ConstTensor"} + ), + ), + CapturePattern( + "head", + lambda node, context: TensorPattern._head(node), + ), + ), + ), + FieldPattern( + "slice", + AstNodePattern( + ast.Tuple, + CapturePattern("slot_count", lambda node, context: len(node.elts)), + FieldPattern( + "elts", + ChoicePattern( + SequencePattern( + ChildPattern( + "shape_or_layout", + lambda: TensorShapeLayoutPattern(), + "tensor_shape", + "tensor_shape_or_layout", + ), + ChildPattern( + "dtype", + lambda: DTypePattern(), + "tensor_dtype", + "dtype", + ), + ), + SequencePattern( + ChildPattern( + "shape_or_layout", + lambda: TensorShapeLayoutPattern(), + "tensor_shape", + "tensor_shape_or_layout", + ), + ChildPattern( + "dtype", + lambda: DTypePattern(), + "tensor_dtype", + "dtype", + ), + ChildPattern( + "optional_0", + lambda: TensorOptionalSlotPattern(), + "tensor_optional_slot", + "layout_or_storage", + ), + ), + SequencePattern( + ChildPattern( + "shape_or_layout", + lambda: TensorShapeLayoutPattern(), + "tensor_shape", + "tensor_shape_or_layout", + ), + ChildPattern( + "dtype", + lambda: DTypePattern(), + "tensor_dtype", + "dtype", + ), + ChildPattern( + "optional_0", + lambda: TensorOptionalSlotPattern(), + "tensor_optional_slot", + "layout", + ), + ChildPattern( + "optional_1", + lambda: TensorOptionalSlotPattern(), + "tensor_optional_slot", + "storage", + ), + ), + ), + ), + ), + ), + ), + pattern_id="tensor.annotation", + ) + ) + + @staticmethod + def construct(match, children, context): + shape_or_layout = children["shape_or_layout"] + if isinstance(shape_or_layout, runtime.LayoutBase): + shape = shape_or_layout.shape + layout = shape_or_layout + else: + shape = shape_or_layout + layout = None + storage = runtime.StorageKind.GMEM + third = children.get("optional_0") + fourth = children.get("optional_1") + if isinstance(third, runtime.StorageKind): + storage = third + elif third is None or isinstance(third, runtime.LayoutBase): + layout = third + else: + raise ParseError.from_node( + match.node, context, "third Tensor slot is not layout/storage" + ) + if fourth is not None: + storage = fourth + return runtime.TensorType( + shape=shape, dtype=children["dtype"], layout=layout, storage=storage + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + TensorLayoutStorageRule(), + TensorPositionRule(), + ) + + @staticmethod + def _head(node: ast.expr) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + +class ScalarTypePattern(ElementPattern): + element_name = "scalar_type" + syntax = LazyPattern( + lambda: BranchPattern( + "type_reference", + ReferencePattern(), + pattern_id="type.reference", + ) + ) + + @staticmethod + def construct(match, children, context): + value = _resolve_reference(match.node, context) + if not isinstance(value, (runtime.TensorType, runtime.TupleType, runtime.UnitType)): + raise ParseError.from_node(match.node, context, "annotation did not resolve to IR Type") + return value + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class TypeAnnotationPattern(ElementPattern): + element_name = "type_annotation" + syntax = LazyPattern( + lambda: ChoicePattern( + TensorPattern(), + ScalarTypePattern(), + ) + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +def _constraint_value(node: ast.AST): + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return ast.unparse(node) + raise ValueError( + f"where layout extent must be a literal or symbolic name, got {type(node).__name__}" + ) + + +def _parse_partial_constraint(node: ast.AST): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): + raise ValueError('Partial binding must use P("reduction")') + if node.func.id != "P" or len(node.args) != 1 or node.keywords: + raise ValueError('Partial binding must use P("reduction")') + value = _constraint_value(node.args[0]) + if not isinstance(value, str) or not value: + raise ValueError("partial reduction must be a non-empty string") + return Partial(value) + + +def _parse_constraint_bindings(node: ast.AST): + if not isinstance(node, ast.Set): + raise ValueError("layout bindings must be a set") + bindings = [] + for item in node.elts: + if not isinstance(item, ast.BinOp) or not isinstance(item.op, ast.MatMult): + raise ValueError("layout bindings must use `topology @ B()` or `P()`") + topology = _constraint_value(item.left) + if not isinstance(topology, str) or not topology: + raise ValueError("layout binding topology must be symbolic") + if ( + isinstance(item.right, ast.Call) + and isinstance(item.right.func, ast.Name) + and item.right.func.id == "B" + and not item.right.args + and not item.right.keywords + ): + attribute = Broadcast() + else: + attribute = _parse_partial_constraint(item.right) + bindings.append((topology, attribute)) + return bindings + + +def _where_static(node: ast.AST, context: MatchContext): + if isinstance(node, ast.Constant): + return node.value + return _resolve_reference(node, context) + + +def _parse_layout_constraint(node: ast.AST, context: MatchContext): + if not isinstance(node, ast.Tuple): + raise ValueError("layout constraint must be a tuple") + dims_node = node + extras: tuple[ast.AST, ...] = () + if node.elts and isinstance(node.elts[0], ast.Tuple): + dims_node = node.elts[0] + extras = tuple(node.elts[1:]) + if len(extras) > 1: + raise ValueError("layout constraint accepts one binding set") + if not dims_node.elts: + raise ValueError("layout constraint cannot be empty") + + def resolve_extent(item: ast.AST): + value = _where_static(item, context) + if isinstance(value, bool) or not isinstance(value, (int, DimVar)): + raise ValueError("layout dimensions must use `_`, an integer, or a symbolic extent") + return value + + shape = [] + bindings = [] + for index, item in enumerate(dims_node.elts): + if isinstance(item, ast.Name) and item.id == "_": + shape.append(_LAYOUT_WILDCARD) + continue + if isinstance(item, ast.Name) and item.id == "D": + raise ValueError("layout broadcast must use a `{topology @ B()}` binding") + if isinstance(item, ast.BinOp) and isinstance(item.op, ast.MatMult): + extent = resolve_extent(item.left) + topology = _constraint_value(item.right) + if not isinstance(topology, str) or not topology: + raise ValueError("layout topology binding must be symbolic") + shape.append(extent) + bindings.append((topology, Split(index))) + continue + shape.append(resolve_extent(item)) + if extras: + bindings.extend(_parse_constraint_bindings(extras[0])) + if len({topology for topology, _ in bindings}) != len(bindings): + raise ValueError("layout constraint cannot bind one topology more than once") + return LayoutConstraint(layout=Layout(shape=tuple(shape)), bindings=tuple(bindings)) + + +class WhereAnnotationPattern(ElementPattern): + element_name = "where_annotation" + syntax = BranchPattern( + "where_annotation", + AstNodePattern( + ast.Call, + FieldPattern( + "func", + AstNodePattern( + ast.Name, + FieldPattern("id", LiteralPattern("where")), + ), + ), + ), + pattern_id="annotation.where", + ) + + @staticmethod + def construct(match, children, context): + node = match.node + source = context.function.source_filename if context.function else "" + location = SourceLocation( + filename=source, + line=getattr(node, "lineno", 0), + column=getattr(node, "col_offset", 0), + end_line=getattr(node, "end_lineno", None), + end_column=getattr(node, "end_col_offset", None), + ) + if node.args: + raise ParseError.from_node(node, context, "where(...) accepts keyword arguments only") + if not node.keywords: + raise ParseError.from_node(node, context, "where(...) cannot be empty") + constraints = [] + fields = set() + try: + for keyword in node.keywords: + if keyword.arg is None: + raise ValueError("where(...) does not accept **kwargs") + if keyword.arg in fields: + raise ValueError(f"where(...) repeats keyword {keyword.arg!r}") + fields.add(keyword.arg) + if keyword.arg == "layout": + constraints.append( + dataclasses.replace( + _parse_layout_constraint(keyword.value, context), + source_loc=location, + provenance=ConstraintProvenance.AUTHOR, + ) + ) + elif keyword.arg == "mesh": + constraints.append( + MeshConstraint( + mesh=_where_static(keyword.value, context), + source_loc=location, + provenance=ConstraintProvenance.AUTHOR, + ) + ) + elif keyword.arg == "storage": + constraints.append( + StorageConstraint( + storage=_where_static(keyword.value, context), + source_loc=location, + provenance=ConstraintProvenance.AUTHOR, + ) + ) + else: + raise ValueError( + f"where(...) has unknown field {keyword.arg!r}; use " + "layout=..., mesh=..., or storage=..." + ) + except (TypeError, ValueError) as error: + raise ParseError.from_node(node, context, str(error)) from error + return ScheduleConstraintMetadata(constraints=tuple(constraints), source_loc=location) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class ReturnTypePattern(ElementPattern): + element_name = "return_type" + syntax = LazyPattern( + lambda: BranchPattern( + "return_type", + ChildPattern( + "type", + lambda: TypeAnnotationPattern(), + "type_annotation", + "return", + values={"position": "hir_output"}, + ), + pattern_id="function.return_type", + ) + ) + + @staticmethod + def construct(match, children, context): + value = children["type"] + context.lexical_scope.define(_RETURN_TYPE, value) + return value + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class SignaturePattern(ElementPattern): + element_name = "signature" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.arguments, + FieldPattern("posonlyargs", SequencePattern()), + FieldPattern( + "args", + RepeatPattern( + AstNodePattern( + ast.arg, + FieldPattern( + "annotation", + AstNodePattern( + ast.expr, + ChildPattern( + "parameter_{index}", + TypeAnnotationPattern(), + "type_annotation", + "parameter", + ), + ), + ), + ) + ), + ), + FieldPattern("vararg", LiteralPattern(None)), + FieldPattern("kwonlyargs", SequencePattern()), + FieldPattern("kw_defaults", SequencePattern()), + FieldPattern("kwarg", LiteralPattern(None)), + FieldPattern("defaults", SequencePattern()), + ), + SignaturePattern._bind, + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.arguments) + if context.function is None: + return None + names = tuple(argument.arg for argument in node.args) + constness = tuple( + isinstance(argument.annotation, ast.Subscript) + and TensorPattern._head(argument.annotation.value) == "ConstTensor" + for argument in node.args + ) + children: list[AstChild] = [] + for index, argument in enumerate(node.args): + assert argument.annotation is not None + if context.function.dialect == "hir": + position = "hir_input" + else: + first_output = max(0, len(node.args) - context.function.output_count) + position = "tir_output" if index >= first_output else "tir_input" + children.append( + AstChild( + f"parameter_{index}", + TypeAnnotationPattern(), + argument.annotation, + "type_annotation", + "parameter", + values={"position": position}, + ) + ) + return dataclasses.replace( + matched, + pattern_id="function.signature", + branch_id="signature", + captures={ + **matched.captures, + "names": names, + "constness": constness, + }, + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + names = match.captures["names"] + constness = match.captures["constness"] + params = tuple( + runtime.Var( + type=children[f"parameter_{index}"], + name=name, + is_const=constness[index], + ) + for index, name in enumerate(names) + ) + for param in params: + context.lexical_scope.define(param.name, param) + return params + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticLiteralPattern(ElementPattern): + element_name = "literal" + syntax = LazyPattern( + lambda: BranchPattern( + "static_literal", + AstNodePattern( + ast.Constant, + FieldPattern( + "value", + CapturePattern("value", lambda value, context: value), + ), + ), + pattern_id="static.literal", + ) + ) + + @staticmethod + def construct(match, children, context): + return match.captures["value"] + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticReferencePattern(ElementPattern): + element_name = "primary" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "static_name", + AstNodePattern( + ast.Name, + FieldPattern( + "id", + CapturePattern("name", lambda value, context: value), + ), + ), + pattern_id="static.name", + ), + BranchPattern( + "static_attribute", + AstNodePattern( + ast.Attribute, + FieldPattern( + "attr", + CapturePattern("attribute", lambda value, context: value), + ), + FieldPattern( + "value", + ChildPattern( + "owner", + lambda: StaticValuePattern(), + "static_owner", + "static_owner", + ), + ), + ), + pattern_id="static.attribute", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "static_name": + return _resolve_reference(match.node, context) + elif match.branch_id == "static_attribute": + owner = children["owner"] + attribute = match.captures["attribute"] + try: + return getattr(owner, attribute) + except AttributeError as error: + raise ParseError.from_node( + match.node, + context, + f"{type(owner).__name__} has no attribute {attribute!r}", + ) from error + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticSequencePattern(ElementPattern): + element_name = "sequence" + syntax = LazyPattern( + lambda: ChoicePattern( + StaticSequencePattern._branch(ast.Tuple, tuple, "static.tuple"), + StaticSequencePattern._branch(ast.List, list, "static.list"), + StaticSequencePattern._branch(ast.Set, set, "static.set"), + ) + ) + + @staticmethod + def _branch(node_type: type, constructor: type, pattern_id: str) -> AstPattern[Any]: + return BranchPattern( + "static_sequence", + AstNodePattern( + node_type, + CapturePattern("constructor", lambda node, context: constructor), + FieldPattern( + "elts", + RepeatPattern( + ChildPattern( + "item_{index}", + StaticValuePattern(), + "static_item", + "static_item", + ) + ), + ), + ), + pattern_id=pattern_id, + ) + + @staticmethod + def construct(match, children, context): + return match.captures["constructor"](children.values()) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticDictPattern(ElementPattern): + element_name = "dict" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.Dict, + FieldPattern("keys", RepeatPattern(StaticValuePattern())), + FieldPattern("values", RepeatPattern(StaticValuePattern())), + ), + StaticDictPattern._bind, + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.Dict) + if any(key is None for key in node.keys): + return None + children: list[AstChild] = [] + for index, (key, value) in enumerate(zip(node.keys, node.values)): + assert key is not None + children.extend( + ( + AstChild( + f"key_{index}", + StaticValuePattern(), + key, + "static_key", + "static_key", + ), + AstChild( + f"value_{index}", + StaticValuePattern(), + value, + "static_value", + "static_value", + ), + ) + ) + return dataclasses.replace( + matched, + pattern_id="static.dict", + branch_id="static_dict", + captures={**matched.captures, "length": len(node.keys)}, + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + return { + children[f"key_{index}"]: children[f"value_{index}"] + for index in range(match.captures["length"]) + } + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticBinaryPattern(ElementPattern): + element_name = "binary_operation" + syntax = LazyPattern( + lambda: BranchPattern( + "static_binary", + AstNodePattern( + ast.BinOp, + FieldPattern( + "op", + ChoicePattern( + AstNodePattern(ast.Add), + AstNodePattern(ast.Sub), + AstNodePattern(ast.Mult), + AstNodePattern(ast.Div), + AstNodePattern(ast.FloorDiv), + AstNodePattern(ast.Mod), + AstNodePattern(ast.Pow), + ), + ), + CapturePattern("operator", lambda node, context: _BINARY_OPERATORS[type(node.op)]), + FieldPattern( + "left", + ChildPattern("left", StaticValuePattern(), "static_operand"), + ), + FieldPattern( + "right", + ChildPattern("right", StaticValuePattern(), "static_operand"), + ), + ), + pattern_id="static.binary", + ) + ) + + @staticmethod + def construct(match, children, context): + try: + return match.captures["operator"](children["left"], children["right"]) + except (TypeError, ValueError, ZeroDivisionError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticUnaryPattern(ElementPattern): + element_name = "unary_operation" + syntax = LazyPattern( + lambda: BranchPattern( + "static_unary", + AstNodePattern( + ast.UnaryOp, + FieldPattern( + "op", + ChoicePattern( + AstNodePattern(ast.UAdd), + AstNodePattern(ast.USub), + AstNodePattern(ast.Not), + ), + ), + CapturePattern("operator", lambda node, context: _UNARY_OPERATORS[type(node.op)]), + FieldPattern( + "operand", + ChildPattern("operand", StaticValuePattern(), "static_operand"), + ), + ), + pattern_id="static.unary", + ) + ) + + @staticmethod + def construct(match, children, context): + try: + return match.captures["operator"](children["operand"]) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticCallPattern(ElementPattern): + element_name = "call" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.Call, + FieldPattern("func", StaticValuePattern()), + FieldPattern("args", RepeatPattern(StaticValuePattern())), + FieldPattern( + "keywords", + RepeatPattern( + AstNodePattern( + ast.keyword, + FieldPattern( + "arg", + PredicatePattern( + "keyword-name", + lambda value, context: isinstance(value, str), + ), + ), + FieldPattern("value", StaticValuePattern()), + ) + ), + ), + ), + StaticCallPattern._bind, + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.Call) + keyword_names = [keyword.arg for keyword in node.keywords] + if len(set(keyword_names)) != len(keyword_names): + return None + children = [AstChild("callee", StaticValuePattern(), node.func, "static_callee")] + children.extend( + AstChild( + f"arg_{index}", + StaticValuePattern(), + argument, + "static_argument", + ) + for index, argument in enumerate(node.args) + ) + children.extend( + AstChild( + f"kw_{keyword.arg}", + StaticValuePattern(), + keyword.value, + "static_argument", + keyword.arg, + ) + for keyword in node.keywords + if keyword.arg is not None + ) + return dataclasses.replace( + matched, + pattern_id="static.call", + branch_id="static_call", + captures={ + **matched.captures, + "arg_count": len(node.args), + "keywords": tuple(keyword_names), + }, + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + callee = children["callee"] + if not callable(callee) and not isinstance(callee, type): + raise ParseError.from_node( + match.node, + context, + "static calls require a callable target", + ) + args = tuple(children[f"arg_{index}"] for index in range(match.captures["arg_count"])) + kwargs = {name: children[f"kw_{name}"] for name in match.captures["keywords"]} + try: + return callee(*args, **kwargs) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticSlicePattern(ElementPattern): + element_name = "slice" + syntax = LazyPattern( + lambda: BranchPattern( + "static_slice", + AstNodePattern( + ast.Slice, + FieldPattern( + "lower", + OptionalPattern( + ChildPattern("lower", StaticValuePattern(), "static_slice", "lower") + ), + ), + FieldPattern( + "upper", + OptionalPattern( + ChildPattern("upper", StaticValuePattern(), "static_slice", "upper") + ), + ), + FieldPattern( + "step", + OptionalPattern( + ChildPattern("step", StaticValuePattern(), "static_slice", "step") + ), + ), + ), + pattern_id="static.slice", + ) + ) + + @staticmethod + def construct(match, children, context): + return slice(children.get("lower"), children.get("upper"), children.get("step")) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticSubscriptPattern(ElementPattern): + element_name = "subscript" + syntax = LazyPattern( + lambda: BranchPattern( + "static_subscript", + AstNodePattern( + ast.Subscript, + FieldPattern( + "value", + ChildPattern("owner", StaticValuePattern(), "static_owner"), + ), + FieldPattern( + "slice", + ChildPattern("key", StaticValuePattern(), "static_key"), + ), + ), + pattern_id="static.subscript", + ) + ) + + @staticmethod + def construct(match, children, context): + try: + return children["owner"][children["key"]] + except (IndexError, KeyError, TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StaticValuePattern(ElementPattern): + element_name = "expression" + syntax = LazyPattern( + lambda: ChoicePattern( + StaticLiteralPattern(), + StaticReferencePattern(), + StaticSequencePattern(), + StaticDictPattern(), + StaticBinaryPattern(), + StaticUnaryPattern(), + StaticCallPattern(), + StaticSlicePattern(), + StaticSubscriptPattern(), + ) + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class NamePattern(ElementPattern): + element_name = "name" + syntax = LazyPattern( + lambda: BranchPattern( + "name", + AstNodePattern( + ast.Name, + FieldPattern("id", CapturePattern("name", lambda value, context: value)), + ), + pattern_id="expression.name", + ) + ) + + @staticmethod + def construct(match, children, context): + name = match.captures["name"] + value = context.lexical_scope.lookup(name) + if value is None: + value = context.function.closure.get(name) + if isinstance(value, slice): + value = value.start + if isinstance(value, runtime.Expr): + return value + if isinstance(value, (bool, int, float)): + return _constant(value) + raise ParseError.from_node(match.node, context, f"name {name!r} is not an Expr") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class ConstantPattern(ElementPattern): + element_name = "constant" + syntax = LazyPattern( + lambda: BranchPattern( + "constant", + AstNodePattern( + ast.Constant, + FieldPattern("value", LiteralPattern(value_type=(bool, int, float))), + ), + pattern_id="expression.constant", + ) + ) + + @staticmethod + def construct(match, children, context): + return _constant(match.captures["value"]) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class TupleExpressionPattern(ElementPattern): + element_name = "tuple_expression" + syntax = LazyPattern( + lambda: BranchPattern( + "expr_tuple", + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + ChildPattern( + "element_{index}", + ExpressionPattern(), + "expression", + "tuple_item", + ) + ), + ), + ), + pattern_id="expression.tuple", + ) + ) + + @staticmethod + def construct(match, children, context): + elements = tuple(children.values()) + return runtime.IrTuple( + type=runtime.TupleType(fields=tuple(item.type for item in elements)), + elements=elements, + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +@dataclass(frozen=True) +class CallBindingRule: + STATEMENT: ClassVar[str] = "A call must bind its arguments into a Call tuple." + + def apply(self, value, *, match, context): + if not isinstance(value, runtime.Call): + raise ParseError.from_node(match.node, context, "call did not construct Call") + if not isinstance(value.args, tuple): + raise ParseError.from_node(match.node, context, "call arguments are not a tuple") + if context.function is not None and context.function.state.mesh_stack: + value = replace_metadata( + value, + ExecutionDomainMetadata(tuple(context.function.state.mesh_stack)), + ) + return value + + +def _types_compatible(actual: object, expected: object) -> bool: + if actual == expected: + return True + if isinstance(actual, runtime.TensorType) and isinstance(expected, runtime.TensorType): + try: + actual_shape = tuple(runtime.normalize_dim(dim) for dim in actual.shape) + expected_shape = tuple(runtime.normalize_dim(dim) for dim in expected.shape) + except (TypeError, ValueError): + actual_shape = actual.shape + expected_shape = expected.shape + return actual_shape == expected_shape and actual.dtype == expected.dtype + if isinstance(actual, runtime.TupleType) and isinstance(expected, runtime.TupleType): + return len(actual.fields) == len(expected.fields) and all( + _types_compatible(left, right) for left, right in zip(actual.fields, expected.fields) + ) + return False + + +@dataclass(frozen=True) +class CallTypeInferenceRule: + STATEMENT: ClassVar[str] = "A call's result type must be inferred from its binding." + + def apply(self, value, *, match, context): + if not isinstance(value, runtime.Call): + return value + value = attach_authored_metadata(value, match.node, context) + infer_context = context.lexical_scope.lookup(_TYPE_INFER_CONTEXT) + if not isinstance(infer_context, runtime.TypeInferContext): + infer_context = runtime.TypeInferContext() + computed = runtime.TypeInferVisitor(infer_context).visit(value) + return dataclasses.replace(value, type=computed) + + +class CallExpectedTypeRule: + STATEMENT: ClassVar[str] = "A call's inferred type must satisfy the expected expression type." + + def apply(self, value, *, match, context): + expected = context.expected_type + actual = getattr(value, "type", None) + if expected is not None and not _types_compatible(actual, expected): + raise ParseError.from_node( + match.node, + context, + f"expression type {actual!r} does not match expected type {expected!r}", + ) + return value + + +class CallPattern(ElementPattern): + element_name = "op_call" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.Call, + FieldPattern("func", ReferencePattern()), + FieldPattern("args", RepeatPattern(AstNodePattern(ast.expr))), + FieldPattern( + "keywords", + RepeatPattern( + AstNodePattern( + ast.keyword, + FieldPattern( + "arg", + PredicatePattern( + "keyword-name", + lambda value, context: isinstance(value, str), + ), + ), + FieldPattern("value", AstNodePattern(ast.expr)), + ) + ), + ), + ), + CallPattern._bind, + ) + ) + + @staticmethod + def _pattern_for_param(param: object, node: ast.AST) -> AstPattern[Any]: + annotation = param.annotation + if annotation is runtime.TensorType and isinstance(node, ast.Subscript): + return TensorPattern() + if annotation is runtime.DType: + return DTypePattern() + if annotation is runtime.StorageKind or param.name == "storage": + return StoragePattern() + if annotation in (runtime.Layout, runtime.ShardLayout, runtime.LayoutBase): + return LayoutPattern() + return StaticValuePattern() + + @staticmethod + def _schema_children(node: ast.Call, schema: object) -> tuple[AstChild, ...] | None: + params = tuple(schema.signature) + inputs = [param for param in params if param.kind == "input"] + attrs = [param for param in params if param.kind == "attribute"] + variadic = bool(getattr(schema.op_class, "is_variadic", False)) + children: list[AstChild] = [] + bound_attrs: set[str] = set() + for index, argument in enumerate(node.args): + if variadic or index < len(inputs): + name = inputs[0].name if variadic else inputs[index].name + children.append( + AstChild( + f"input_{index}", + ExpressionPattern(), + argument, + "call_argument", + name, + ) + ) + continue + attr_index = index - len(inputs) + if attr_index >= len(attrs): + return None + param = attrs[attr_index] + bound_attrs.add(param.name) + children.append( + AstChild( + f"attr_{param.name}", + CallPattern._pattern_for_param(param, argument), + argument, + "call_attribute", + "allocation" if param.annotation is runtime.TensorType else param.name, + ) + ) + for keyword in node.keywords: + if keyword.arg is None or keyword.arg in bound_attrs: + return None + param = next((item for item in attrs if item.name == keyword.arg), None) + if param is None: + return None + bound_attrs.add(param.name) + children.append( + AstChild( + f"attr_{param.name}", + CallPattern._pattern_for_param(param, keyword.value), + keyword.value, + "call_attribute", + "allocation" if param.annotation is runtime.TensorType else param.name, + ) + ) + if not variadic and len(node.args) < len(inputs): + return None + return tuple(children) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.Call) + module_owner = None + try: + callee = _resolve_reference(node.func, context) + except ParseError: + return None + if isinstance(callee, runtime.Module): + module_owner = callee + if node.keywords: + return None + try: + callee = callee.entry_function() + except ValueError: + return None + if isinstance(callee, runtime.Function): + if node.keywords: + return None + return dataclasses.replace( + matched, + pattern_id="call.function", + branch_id="function_call", + captures={ + **matched.captures, + "callee": callee, + "module_owner": locals().get("module_owner"), + "module_binding": (node.func.id if isinstance(node.func, ast.Name) else None), + }, + children=tuple( + AstChild( + f"arg_{index}", + ExpressionPattern(), + argument, + "call_argument", + "argument", + ) + for index, argument in enumerate(node.args) + ), + ) + schema = ( + callee if isinstance(callee, runtime.OpSchema) else getattr(callee, "_op_schema", None) + ) + if not isinstance(schema, runtime.OpSchema): + return None + children = CallPattern._schema_children(node, schema) + if children is None: + return None + return dataclasses.replace( + matched, + pattern_id="call.operation", + branch_id="operation_call", + captures={**matched.captures, "schema": schema}, + children=children, + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "operation_call": + schema = match.captures["schema"] + inputs = tuple(value for name, value in children.items() if name.startswith("input_")) + attrs = { + name.removeprefix("attr_"): value + for name, value in children.items() + if name.startswith("attr_") + } + for name, value in tuple(attrs.items()): + parameter = next((item for item in schema.signature if item.name == name), None) + annotation = None if parameter is None else parameter.annotation + if ( + isinstance(annotation, type) + and issubclass(annotation, enum.Enum) + and isinstance(value, str) + ): + try: + attrs[name] = annotation(value) + except ValueError as error: + raise ParseError.from_node(match.node, context, str(error)) from error + try: + operation = schema.builder(**attrs) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + placeholder_type = context.expected_type + if placeholder_type is None: + placeholder_type = runtime.TensorType.scalar(runtime.DType.f32) + return runtime.Call(type=placeholder_type, target=operation, args=inputs) + elif match.branch_id == "function_call": + callee = match.captures["callee"] + args = tuple(children.values()) + module_owner = match.captures.get("module_owner") + if module_owner is not None and context.function.module is None: + raise ParseError.from_node( + match.node, + context, + "a Module call is only valid inside a @module class body", + ) + placeholder = runtime.Call( + type=callee.return_type, + target=callee, + args=args, + ) + infer_context = context.lexical_scope.lookup(_TYPE_INFER_CONTEXT) + if not isinstance(infer_context, runtime.TypeInferContext): + infer_context = runtime.TypeInferContext() + instance = runtime.elaborate( + callee, tuple(arg.type for arg in args), infer_context, placeholder + ) + return dataclasses.replace(placeholder, target=instance) + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + CallBindingRule(), + CallTypeInferenceRule(), + CallExpectedTypeRule(), + ) + + +_EXPR_BINARY_KINDS: Mapping[type[ast.AST], str] = { + ast.Add: "ADD", + ast.Sub: "SUB", + ast.Mult: "MUL", + ast.Div: "DIV", + ast.FloorDiv: "FLOOR_DIV", + ast.Mod: "MOD", + ast.Eq: "EQ", + ast.NotEq: "NE", + ast.Lt: "LT", + ast.LtE: "LE", + ast.Gt: "GT", + ast.GtE: "GE", + ast.And: "AND", + ast.Or: "OR", +} + + +class BinaryExpressionPattern(ElementPattern): + element_name = "binary_expression" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "binary_expression", + AstNodePattern( + ast.BinOp, + FieldPattern( + "op", + PredicatePattern( + "binary-op", + lambda op, context: type(op) in _EXPR_BINARY_KINDS, + ), + ), + CapturePattern( + "kind", + lambda node, context: _EXPR_BINARY_KINDS[type(node.op)], + ), + FieldPattern( + "left", + ChildPattern("left", ExpressionPattern(), "expression"), + ), + FieldPattern( + "right", + ChildPattern("right", ExpressionPattern(), "expression"), + ), + ), + pattern_id="expression.binary", + ), + BranchPattern( + "binary_expression", + AstNodePattern( + ast.Compare, + FieldPattern( + "ops", + SequencePattern( + PredicatePattern( + "comparison-op", + lambda op, context: type(op) in _EXPR_BINARY_KINDS, + ) + ), + ), + CapturePattern( + "kind", + lambda node, context: _EXPR_BINARY_KINDS[type(node.ops[0])], + ), + FieldPattern( + "left", + ChildPattern("left", ExpressionPattern(), "expression"), + ), + FieldPattern( + "comparators", + SequencePattern(ChildPattern("right", ExpressionPattern(), "expression")), + ), + ), + pattern_id="expression.binary", + ), + BranchPattern( + "binary_expression", + AstNodePattern( + ast.BoolOp, + FieldPattern( + "op", + PredicatePattern( + "boolean-op", + lambda op, context: type(op) in _EXPR_BINARY_KINDS, + ), + ), + CapturePattern( + "kind", + lambda node, context: _EXPR_BINARY_KINDS[type(node.op)], + ), + FieldPattern( + "values", + SequencePattern( + ChildPattern("left", ExpressionPattern(), "expression"), + ChildPattern("right", ExpressionPattern(), "expression"), + ), + ), + ), + pattern_id="expression.binary", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + return runtime.Call( + type=children["left"].type, + target=runtime.Binary(kind=runtime.BinaryKind[match.captures["kind"]]), + args=(children["left"], children["right"]), + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + CallBindingRule(), + CallTypeInferenceRule(), + CallExpectedTypeRule(), + ) + + +class UnaryExpressionPattern(ElementPattern): + element_name = "unary_expression" + syntax = LazyPattern( + lambda: BranchPattern( + "unary_expression", + AstNodePattern( + ast.UnaryOp, + FieldPattern( + "op", + PredicatePattern( + "unary-op", + lambda op, context: type(op) in {ast.USub, ast.Not}, + ), + ), + CapturePattern( + "kind", + lambda node, context: {ast.USub: "NEG", ast.Not: "NOT"}[type(node.op)], + ), + FieldPattern( + "operand", + ChildPattern("operand", ExpressionPattern(), "expression"), + ), + ), + pattern_id="expression.unary", + ) + ) + + @staticmethod + def construct(match, children, context): + operand = children["operand"] + return runtime.Call( + type=operand.type, + target=runtime.Unary(kind=runtime.UnaryKind[match.captures["kind"]]), + args=(operand,), + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + CallBindingRule(), + CallTypeInferenceRule(), + CallExpectedTypeRule(), + ) + + +class SliceEndpointBinaryPattern(ElementPattern): + element_name = "slice_endpoint_binary" + syntax = LazyPattern( + lambda: BranchPattern( + "dimension_binary", + AstNodePattern( + ast.BinOp, + FieldPattern( + "op", + PredicatePattern( + "dim-op", + lambda op, context: ( + type(op) in {ast.Add, ast.Sub, ast.Mult, ast.FloorDiv, ast.Mod} + ), + ), + ), + CapturePattern("operator", lambda node, context: _BINARY_OPERATORS[type(node.op)]), + CapturePattern( + "dimension_operator", + lambda node, context: { + ast.Add: runtime.DimAdd, + ast.Sub: runtime.DimSub, + ast.Mult: runtime.DimMul, + ast.FloorDiv: runtime.DimFloorDiv, + ast.Mod: runtime.DimMod, + }[type(node.op)], + ), + FieldPattern( + "left", + ChildPattern("left", IndexEndpointPattern(), "slice_endpoint"), + ), + FieldPattern( + "right", + ChildPattern("right", IndexEndpointPattern(), "slice_endpoint"), + ), + ), + pattern_id="expression.slice_endpoint.binary", + ) + ) + + @staticmethod + def construct(match, children, context): + left = children["left"] + right = children["right"] + if ( + isinstance(left, slice) + and type(match.node.op) in {ast.Add, ast.Sub} + and isinstance(right, (int, runtime.Expr)) + ): + offset = right + if type(match.node.op) is ast.Sub: + offset = runtime.simplify_dim(runtime.DimMul, (-1, offset)) + try: + start = runtime.simplify_dim(runtime.DimAdd, (left.start, offset)) + stop = runtime.simplify_dim(runtime.DimAdd, (left.stop, offset)) + except (TypeError, ValueError, ZeroDivisionError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + return slice(start, stop, left.step) + numeric = all( + isinstance(value, (int, float)) and not isinstance(value, bool) + for value in (left, right) + ) + try: + if numeric: + return match.captures["operator"](left, right) + return runtime.simplify_dim(match.captures["dimension_operator"], (left, right)) + except (TypeError, ValueError, ZeroDivisionError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class IndexEndpointPattern(ElementPattern): + element_name = "index_endpoint" + syntax = LazyPattern( + lambda: ChoicePattern( + ConditionPattern( + "constant", + lambda node, context: isinstance(node, ast.Constant), + StaticLiteralPattern(), + ), + ConditionPattern( + "name", + lambda node, context: isinstance(node, ast.Name), + StaticReferencePattern(), + ), + SliceEndpointBinaryPattern(), + MeshCoordinatePattern(), + ExpressionPattern(), + StaticValuePattern(), + ) + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class IndexSlicePattern(ElementPattern): + element_name = "index_slice" + syntax = LazyPattern( + lambda: BranchPattern( + "static_slice", + AstNodePattern( + ast.Slice, + *( + FieldPattern( + name, + OptionalPattern( + ChildPattern( + name, + IndexEndpointPattern(), + "slice_endpoint", + name, + ) + ), + ) + for name in ("lower", "upper", "step") + ), + ), + pattern_id="expression.index.slice", + ) + ) + + @staticmethod + def construct(match, children, context): + return slice(children.get("lower"), children.get("upper"), children.get("step")) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class SubscriptIndexPattern(ElementPattern): + element_name = "subscript_index" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "tuple_children", + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern( + ChoicePattern( + ConditionPattern( + "slice", + lambda node, context: isinstance(node, ast.Slice), + ChildPattern( + "index_{index}", + IndexSlicePattern(), + "subscript_index", + "subscript_index", + ), + ), + ChildPattern( + "index_{index}", + IndexEndpointPattern(), + "subscript_index", + "subscript_index", + ), + ) + ), + ), + ), + pattern_id="expression.index.tuple", + ), + BranchPattern( + "identity", + ChoicePattern( + ConditionPattern( + "slice", + lambda node, context: isinstance(node, ast.Slice), + ChildPattern( + "value", + IndexSlicePattern(), + "subscript_index", + "subscript_index", + ), + ), + ChildPattern( + "value", + IndexEndpointPattern(), + "subscript_index", + "subscript_index", + ), + ), + pattern_id="expression.index", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "identity": + return children["value"] + elif match.branch_id == "tuple_children": + return tuple(children.values()) + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class SubscriptExpressionPattern(ElementPattern): + element_name = "subscript_expression" + syntax = LazyPattern( + lambda: BranchPattern( + "subscript_expression", + AstNodePattern( + ast.Subscript, + FieldPattern("value", ChildPattern("value", ExpressionPattern(), "expression")), + FieldPattern( + "slice", + ChildPattern("index", SubscriptIndexPattern(), "subscript_index"), + ), + ), + pattern_id="expression.subscript", + ) + ) + + @staticmethod + def construct(match, children, context): + value = children["value"] + index = children["index"] + if isinstance(value.type, runtime.TupleType): + if isinstance(index, bool) or not isinstance(index, int): + raise ParseError.from_node( + match.node, context, "Tuple subscript requires an integer literal" + ) + normalized = index + len(value.type.fields) if index < 0 else index + return _infer_call(runtime.TupleGetItem(index=normalized), (value,), context) + if not isinstance(value.type, runtime.TensorType): + raise ParseError.from_node( + match.node, context, "subscript requires TensorType or TupleType" + ) + indices = index if isinstance(index, tuple) else (index,) + if len(indices) != len(value.type.shape): + raise ParseError.from_node( + match.node, + context, + f"tensor subscript rank {len(indices)} != tensor rank {len(value.type.shape)}", + ) + starts = [] + sizes = [] + strides = [] + collapsed = [] + for axis, (component, extent) in enumerate(zip(indices, value.type.shape)): + if isinstance(component, slice): + begin = 0 if component.start is None else component.start + end = extent if component.stop is None else component.stop + stride = 1 if component.step is None else component.step + starts.append(runtime.dim_expr(begin)) + sizes.append(_slice_size(begin, end, stride, context, match.node)) + strides.append(stride) + continue + if isinstance(component, bool): + raise ParseError.from_node(match.node, context, "bool is not a tensor index") + start = runtime.dim_expr(component) + starts.append(start) + sizes.append(1) + strides.append(1) + collapsed.append(axis) + starts_expr = runtime.IrTuple( + type=runtime.TupleType(fields=tuple(start.type for start in starts)), + elements=tuple(starts), + ) + sliced = _infer_call( + runtime.Slice(sizes=tuple(sizes), strides=tuple(strides)), + (value, starts_expr), + context, + ) + if not collapsed: + return sliced + new_shape = tuple( + extent for axis, extent in enumerate(sliced.type.shape) if axis not in collapsed + ) + return _infer_call(runtime.Reshape(new_shape=new_shape), (sliced,), context) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class MeshCoordinatePattern(ElementPattern): + element_name = "mesh_coordinate" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.Attribute, + FieldPattern("value", AstNodePattern(ast.Name)), + FieldPattern("attr", LiteralPattern(value_type=str)), + ), + MeshCoordinatePattern._bind, + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.Attribute) + assert isinstance(node.value, ast.Name) + if context.function is None: + return None + mesh = context.lexical_scope.lookup(node.value.id) + if not isinstance(mesh, runtime.Mesh): + return None + axis = next( + (index for index, name in enumerate(mesh.names) if name == node.attr), + None, + ) + if axis is None and node.attr in {"x", "y", "z"}: + candidate = ("x", "y", "z").index(node.attr) + if candidate < len(mesh.layout.shape): + axis = candidate + if axis is None: + return None + return dataclasses.replace( + matched, + pattern_id="expression.mesh_coordinate", + branch_id="mesh_coordinate", + captures={**matched.captures, "mesh": mesh, "axis": axis}, + ) + + @staticmethod + def construct(match, children, context): + if context.function is None: + raise ParseError.from_node(match.node, context, "mesh coordinate lacks context") + mesh = match.captures["mesh"] + axis = match.captures["axis"] + extent = mesh.layout.shape[axis] + if isinstance(extent, bool) or not isinstance(extent, int): + raise ParseError.from_node( + match.node, context, "mesh coordinate requires a concrete axis extent" + ) + cache_key = (id(mesh), axis) + cached = context.function.state.mesh_coordinates.get(cache_key) + if cached is not None: + return cached + vector_type = runtime.TensorType( + shape=(extent,), + dtype=runtime.DType.i64, + layout=None, + storage=runtime.StorageKind.GMEM, + ) + vector = _infer_call(runtime.Arange(type=vector_type), (), context) + attrs = tuple( + runtime.Split(axis=0) if mesh_axis == axis else runtime.Broadcast() + for mesh_axis in range(len(mesh.layout.shape)) + ) + layout = runtime.ShardLayout( + layout=runtime.Layout(shape=(extent,), strides=(1,)), + attrs=attrs, + mesh=mesh, + ) + placed = _infer_call( + runtime.Reshard(layout=layout, storage=runtime.StorageKind.RMEM), + (vector,), + context, + ) + local = _infer_call(runtime.Local(), (placed,), context) + coordinate = _infer_call(runtime.Reshape(new_shape=()), (local,), context) + context.function.state.mesh_coordinates[cache_key] = coordinate + return coordinate + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class ExpressionPattern(ElementPattern): + element_name = "runtime_expression" + syntax = LazyPattern( + lambda: ChoicePattern( + CallPattern(), + LaunchPattern(), + SubscriptExpressionPattern(), + BinaryExpressionPattern(), + UnaryExpressionPattern(), + MeshCoordinatePattern(), + NamePattern(), + ConstantPattern(), + TupleExpressionPattern(), + TensorPattern(), + BranchPattern( + "attribute_expr", + AstNodePattern(ast.Attribute), + pattern_id="expression.attribute", + ), + ) + ) + + @staticmethod + def construct(match, children, context): + value = _resolve_reference(match.node, context) + if isinstance(value, runtime.Expr): + return value + if isinstance(value, (bool, int, float)): + return _constant(value) + raise ParseError.from_node(match.node, context, "attribute did not resolve to Expr") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class MeshContextPattern(ElementPattern): + element_name = "mesh_context" + syntax = LazyPattern( + lambda: ChoicePattern( + BindPattern( + AstNodePattern( + ast.Call, + FieldPattern( + "func", + ChoicePattern( + AstNodePattern( + ast.Name, + FieldPattern("id", LiteralPattern("Mesh")), + ), + AstNodePattern( + ast.Attribute, + FieldPattern("attr", LiteralPattern("Mesh")), + ), + ), + ), + FieldPattern("args", RepeatPattern(AstNodePattern(ast.expr), minimum=1)), + FieldPattern( + "keywords", + RepeatPattern( + AstNodePattern( + ast.keyword, + FieldPattern( + "arg", + ChoicePattern( + LiteralPattern("layout"), + LiteralPattern("names"), + ), + ), + ) + ), + ), + ), + MeshContextPattern._bind, + ), + BranchPattern( + "mesh_reference", + AstNodePattern( + ast.expr, + PredicatePattern( + "reference", + lambda node, context: isinstance( + node, (ast.Name, ast.Attribute, ast.Subscript) + ), + ), + ChildPattern( + "value", + StaticValuePattern(), + "static_mesh", + "mesh", + ), + ), + pattern_id="mesh.reference", + ), + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.Call) + if not node.args or not isinstance(node.args[0], ast.Tuple): + return None + positional = list(node.args[1:]) + keywords = {keyword.arg: keyword.value for keyword in node.keywords} + layout_node = keywords.get("layout") + names_node = keywords.get("names") + if positional: + if layout_node is not None: + return None + layout_node = positional.pop(0) + if positional: + if names_node is not None: + return None + names_node = positional.pop(0) + if positional or layout_node is None: + return None + children = [ + AstChild( + "topology_names", + StaticValuePattern(), + node.args[0], + "mesh_topologies", + "topologies", + ), + AstChild("layout", LayoutPattern(), layout_node, "mesh_layout", "layout"), + ] + if names_node is not None: + children.append( + AstChild( + "names", + StaticValuePattern(), + names_node, + "mesh_names", + "names", + ) + ) + return dataclasses.replace( + matched, + pattern_id="mesh.context", + branch_id="mesh_context", + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "mesh_context": + if context.function is None: + raise ParseError.from_node(match.node, context, "Mesh requires function context") + topology_names = children["topology_names"] + if not isinstance(topology_names, tuple): + raise ParseError.from_node( + match.node, + context, + "Mesh topologies must be a tuple", + ) + if all(isinstance(name, str) for name in topology_names): + try: + topologies = tuple(context.function.topologies[name] for name in topology_names) + except KeyError as error: + raise ParseError.from_node( + match.node, + context, + f"topology {error.args[0]!r} not declared by @module", + ) from error + elif all(hasattr(topology, "name") for topology in topology_names): + topologies = topology_names + else: + raise ParseError.from_node( + match.node, + context, + "Mesh topologies must be names or Topology objects", + ) + names = children.get("names", ()) + try: + mesh = runtime.Mesh( + topologies=topologies, + layout=children["layout"], + names=names, + ) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + binding = context.values.get("mesh_binding") + if isinstance(binding, str): + context.lexical_scope.define(binding, mesh) + context.function.state.mesh_stack.append(mesh) + return mesh + elif match.branch_id == "mesh_reference": + if context.function is None: + raise ParseError.from_node(match.node, context, "Mesh requires function context") + mesh = children["value"] + if not isinstance(mesh, runtime.Mesh): + raise ParseError.from_node(match.node, context, "with context is not Mesh") + binding = context.values.get("mesh_binding") + if isinstance(binding, str): + context.lexical_scope.define(binding, mesh) + context.function.state.mesh_stack.append(mesh) + return mesh + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class WithPattern(ElementPattern): + element_name = "with" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.With, + FieldPattern( + "items", + SequencePattern( + AstNodePattern( + ast.withitem, + FieldPattern("optional_vars", AstNodePattern(ast.Name)), + FieldPattern( + "context_expr", + ChildPattern( + "mesh", + MeshContextPattern(), + "mesh_context", + "mesh", + ), + ), + ) + ), + ), + FieldPattern( + "body", + ChildPattern( + "body", + BlockPattern(), + "block", + "with_body", + transform=_module_from_body, + ), + ), + ), + WithPattern._bind, + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.With) + item = node.items[0] + assert isinstance(item.optional_vars, ast.Name) + binding = item.optional_vars.id + return dataclasses.replace( + matched, + pattern_id="statement.with_mesh", + branch_id="with_mesh", + captures={**matched.captures, "binding": binding}, + children=( + AstChild( + "mesh", + MeshContextPattern(), + item.context_expr, + "mesh_context", + "mesh", + values={"mesh_binding": binding}, + ), + AstChild( + "body", + BlockPattern(), + _module_from_body(node.body), + "block", + "with_body", + ), + ), + ) + + @staticmethod + def construct(match, children, context): + if context.function is None or not context.function.state.mesh_stack: + raise ParseError.from_node(match.node, context, "Mesh stack is unbalanced") + mesh = context.function.state.mesh_stack.pop() + if context.function.dialect == "hir": + return children["body"] + binding = runtime.Var( + type=runtime.TensorType.scalar(runtime.DType.i64, storage=runtime.StorageKind.RMEM), + name=match.captures["binding"], + ) + return runtime.MeshScope(mesh=mesh, binding=binding, body=children["body"]) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class LaunchPattern(ElementPattern): + """TIR host launch statement lowered from the authored ``launch`` call.""" + + element_name = "launch" + syntax = LazyPattern(lambda: BindPattern(AstNodePattern(ast.Call), LaunchPattern._bind)) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]): + assert isinstance(node, ast.Call) + if not isinstance(node.func, ast.Name) or node.func.id != "launch": + return None + if not node.args: + return None + keywords = {keyword.arg: keyword.value for keyword in node.keywords} + if any(name is None for name in keywords): + return None + if "grid" not in keywords or "block" not in keywords: + return None + if set(keywords) - {"grid", "block", "cluster", "dynamic_smem", "stream", "attrs"}: + return None + children = [ + AstChild("callee", StaticValuePattern(), node.args[0], "launch_callee"), + *( + AstChild(f"arg_{index}", ExpressionPattern(), argument, "launch_argument") + for index, argument in enumerate(node.args[1:]) + ), + AstChild("grid", StaticValuePattern(), keywords["grid"], "launch_extent"), + AstChild("block", StaticValuePattern(), keywords["block"], "launch_extent"), + ] + for name in ("cluster", "dynamic_smem", "stream", "attrs"): + if name in keywords: + children.append( + AstChild(name, StaticValuePattern(), keywords[name], "launch_option") + ) + return dataclasses.replace( + matched, + pattern_id="statement.launch", + branch_id="launch", + captures={"arg_count": len(node.args) - 1}, + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + callee = children["callee"] + if isinstance(callee, runtime.Module): + callee = callee.entry_function() + options = { + name: children[name] + for name in ("cluster", "dynamic_smem", "stream", "attrs") + if name in children + } + return launch_call( + callee, + tuple(children[f"arg_{index}"] for index in range(match.captures["arg_count"])), + children["grid"], + children["block"], + **options, + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class LoopCarryStatementPattern(ElementPattern): + element_name = "loop_carry_statement" + syntax = LazyPattern( + lambda: ChoicePattern( + BranchPattern( + "loop_carry_statement", + AstNodePattern( + ast.Assign, + FieldPattern( + "targets", + SequencePattern(AstNodePattern(ast.expr)), + ), + CapturePattern("names", LoopCarryStatementPattern._target_names), + ), + pattern_id="loop.carry_statement", + ), + BranchPattern( + "loop_carry_statement", + AstNodePattern( + ast.For, + CapturePattern("names", lambda node, context: ()), + FieldPattern( + "body", + ChildPattern( + "nested", + LoopCarryPattern(), + "loop_carry", + transform=_module_from_body, + ), + ), + ), + pattern_id="loop.carry_statement", + ), + BranchPattern( + "loop_carry_statement", + AstNodePattern( + ast.stmt, + CapturePattern("names", lambda node, context: ()), + ), + pattern_id="loop.carry_statement", + ), + ) + ) + + @staticmethod + def _target_names(node: object, context: MatchContext) -> tuple[str, ...]: + assert isinstance(node, ast.Assign) + target = node.targets[0] + targets = target.elts if isinstance(target, ast.Tuple) else (target,) + return tuple(item.id for item in targets if isinstance(item, ast.Name)) + + @staticmethod + def construct(match, children, context): + return (*match.captures["names"], *children.get("nested", ())) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class LoopCarryPattern(ElementPattern): + element_name = "loop_carry" + syntax = LazyPattern( + lambda: BranchPattern( + "loop_carry", + AstNodePattern( + ast.Module, + FieldPattern( + "body", + RepeatPattern( + ChildPattern( + "statement_{index}", + LoopCarryStatementPattern(), + "loop_carry_statement", + ) + ), + ), + ), + pattern_id="loop.carry", + ) + ) + + @staticmethod + def construct(match, children, context): + names: list[str] = [] + for statement_names in children.values(): + for name in statement_names: + if name not in names: + names.append(name) + return tuple(names) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class LoopHeaderPattern(ElementPattern): + element_name = "loop_header" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.For, + FieldPattern("target", AstNodePattern(ast.Name)), + FieldPattern( + "iter", + AstNodePattern( + ast.Call, + FieldPattern( + "func", + AstNodePattern(ast.Name), + ), + FieldPattern("keywords", RepeatPattern(AstNodePattern(ast.keyword))), + FieldPattern("args", RepeatPattern(AstNodePattern(ast.expr), minimum=1)), + ), + ), + FieldPattern( + "body", + ChildPattern( + "carry", + LoopCarryPattern(), + "loop_carry", + transform=_module_from_body, + ), + ), + ), + LoopHeaderPattern._bind, + ) + ) + + @staticmethod + def _bind( + node: object, context: MatchContext, matched: AstMatch[Any] + ) -> AstMatch[Any] | PatternFailure | None: + assert isinstance(node, ast.For) + assert isinstance(node.target, ast.Name) + assert isinstance(node.iter, ast.Call) + assert isinstance(node.iter.func, ast.Name) + kind = node.iter.func.id + count = len(node.iter.args) + if kind not in {"tile", "range"}: + return PatternFailure( + "loop_header", + node.iter.func, + "loop iterator must be tile(...) or range(...)", + ) + if node.iter.keywords: + return PatternFailure( + "loop_header", + node.iter, + "tile()/range() does not accept keyword args (positional-only at the IR level)", + ) + if (kind == "tile" and count != 2) or (kind == "range" and count not in {1, 2, 3}): + if kind == "tile" and count == 1: + detail = "tile(extent) is not supported; use range(extent)" + elif kind == "tile": + detail = f"tile() takes 2 arguments (extent, step), got {count}" + else: + detail = f"range() takes 1 to 3 arguments, got {count}" + return PatternFailure( + "loop_header", + node.iter, + detail, + ) + if kind == "tile": + fields = ("extent", "step") + defaults = {"start": 0} + elif count == 1: + fields = ("extent",) + defaults = {"start": 0, "step": 1} + elif count == 2: + fields = ("start", "extent") + defaults = {"step": 1} + else: + fields = ("start", "extent", "step") + defaults = {} + children = [ + AstChild( + "carry", + LoopCarryPattern(), + _module_from_body(node.body), + "loop_carry", + ) + ] + children.extend( + AstChild( + field_name, + StaticValuePattern(), + argument, + "loop_bound", + field_name, + ) + for field_name, argument in zip(fields, node.iter.args) + ) + return dataclasses.replace( + matched, + pattern_id=f"loop.header.{kind}", + branch_id="loop_header", + captures={ + **matched.captures, + "kind": kind, + "target": node.target.id, + "fields": fields, + "defaults": defaults, + }, + children=tuple(children), + ) + + @staticmethod + def construct(match, children, context): + if context.function is None or context.function.dialect != "hir": + raise ParseError.from_node(match.node, context, "loops require HIR context") + values = dict(match.captures["defaults"]) + values.update((name, value) for name, value in children.items() if name != "carry") + try: + start = runtime.normalize_dim(values["start"]) + extent = runtime.normalize_dim(values["extent"]) + step = runtime.normalize_dim(values["step"]) + except (TypeError, ValueError) as error: + raise ParseError.from_node(match.node, context, str(error)) from error + induction_var = runtime.Var( + type=runtime.TensorType.scalar(runtime.DType.i64), + name=match.captures["target"], + ) + carry_names = tuple( + name + for name in children["carry"] + if isinstance(context.lexical_scope.lookup(name), runtime.Expr) + ) + init_args = tuple(context.lexical_scope.lookup(name) for name in carry_names) + phi_vars = tuple( + runtime.Var(type=value.type, name=name) for name, value in zip(carry_names, init_args) + ) + context.lexical_scope.push_frame() + if match.captures["kind"] == "tile": + stop = runtime.simplify_dim(runtime.DimAdd, (induction_var, runtime.dim_expr(step))) + binding = slice(induction_var, stop, 1) + else: + binding = induction_var + context.lexical_scope.define(match.captures["target"], binding) + for name, phi in zip(carry_names, phi_vars): + context.lexical_scope.define(name, phi) + return LoopFrame( + kind=match.captures["kind"], + target=match.captures["target"], + induction_var=induction_var, + start=start, + extent=extent, + step=step, + carry_names=carry_names, + phi_vars=phi_vars, + init_args=init_args, + ) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class LoopBodyPattern(ElementPattern): + element_name = "loop_body" + syntax = LazyPattern( + lambda: BranchPattern( + "loop_body", + AstNodePattern( + ast.Module, + PredicatePattern( + "assignment-suite", + lambda node, context: ( + not any( + isinstance(statement, (ast.Return, ast.With, ast.Expr, ast.Pass)) + for statement in node.body + ) + ), + ), + FieldPattern( + "body", + RepeatPattern( + ChildPattern( + "statement_{index}", + StatementPattern(), + "loop_statement", + "loop_statement", + ) + ), + ), + ), + pattern_id="loop.body", + ) + ) + + @staticmethod + def construct(match, children, context): + if not children: + raise ParseError.from_node(match.node, context, "loop body cannot be empty") + value = tuple(children.values())[-1] + if not isinstance(value, runtime.Expr): + raise ParseError.from_node(match.node, context, "loop body must yield an Expr") + return value + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class ForPattern(ElementPattern): + element_name = "for" + syntax = LazyPattern( + lambda: BranchPattern( + "loop", + AstNodePattern( + ast.For, + ChildPattern("header", LoopHeaderPattern(), "loop_header"), + FieldPattern( + "body", + ChildPattern( + "body", + LoopBodyPattern(), + "loop_body", + transform=_module_from_body, + ), + ), + ), + pattern_id="statement.for", + ) + ) + + @staticmethod + def construct(match, children, context): + frame = children["header"] + body = children["body"] + yield_values = tuple(context.lexical_scope.lookup(name) for name in frame.carry_names) + context.lexical_scope.pop_frame() + if frame.carry_names: + result_type = ( + frame.phi_vars[0].type + if len(frame.phi_vars) == 1 + else runtime.TupleType(fields=tuple(phi.type for phi in frame.phi_vars)) + ) + else: + result_type = body.type + grid = runtime.GridRegionExpr( + type=result_type, + induction_var=frame.induction_var, + carried_args=frame.phi_vars, + init_args=frame.init_args, + body=body, + yield_values=yield_values, + start=frame.start, + extent=frame.extent, + step=frame.step, + ) + if len(frame.carry_names) == 1: + context.lexical_scope.define(frame.carry_names[0], grid) + else: + for index, name in enumerate(frame.carry_names): + projection = _infer_call(runtime.TupleGetItem(index=index), (grid,), context) + context.lexical_scope.define(name, projection) + return grid + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class TupleAssignmentPattern(ElementPattern): + element_name = "tuple_assignment" + syntax = LazyPattern( + lambda: BranchPattern( + "tuple_assignment", + AstNodePattern( + ast.Assign, + FieldPattern( + "targets", + SequencePattern( + AstNodePattern( + ast.Tuple, + FieldPattern( + "elts", + RepeatPattern(AstNodePattern(ast.Name), minimum=1), + ), + ) + ), + ), + CapturePattern( + "names", + lambda node, context: tuple(item.id for item in node.targets[0].elts), + ), + FieldPattern( + "value", + ChildPattern( + "value", + ExpressionPattern(), + "expression", + "assignment_value", + ), + ), + ), + pattern_id="statement.tuple_assign", + ) + ) + + @staticmethod + def construct(match, children, context): + value = children["value"] + names = match.captures["names"] + if not isinstance(value.type, runtime.TupleType): + raise ParseError.from_node( + match.node, context, "tuple assignment requires a TupleType value" + ) + if len(names) != len(value.type.fields): + raise ParseError.from_node( + match.node, context, "tuple assignment arity does not match TupleType" + ) + schema = getattr(type(value.target), "_op_schema", None) if isinstance(value, runtime.Call) else None + parent_name = getattr(schema, "name", None) or ", ".join(names) + value = replace_metadata(value, BindingMetadata(parent_name)) + for index, name in enumerate(names): + projection = _infer_call(runtime.TupleGetItem(index=index), (value,), context) + projection = replace_metadata(projection, BindingMetadata(name)) + context.lexical_scope.define(name, projection) + return value + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class StatementPattern(ElementPattern): + element_name = "statement" + syntax = LazyPattern( + lambda: ChoicePattern( + ForPattern(), + WithPattern(), + TupleAssignmentPattern(), + BranchPattern( + "assignment", + AstNodePattern( + ast.Assign, + FieldPattern( + "targets", + SequencePattern( + AstNodePattern( + ast.Name, + FieldPattern( + "id", + CapturePattern("name", lambda value, context: value), + ), + ) + ), + ), + FieldPattern( + "value", + ChildPattern( + "value", + ChoicePattern(ExpressionPattern(), StaticValuePattern()), + "expression", + "assignment_value", + ), + ), + ), + pattern_id="statement.assign", + ), + BranchPattern( + "assignment", + AstNodePattern( + ast.AnnAssign, + FieldPattern( + "target", + AstNodePattern( + ast.Name, + FieldPattern( + "id", + CapturePattern("name", lambda value, context: value), + ), + ), + ), + FieldPattern( + "annotation", + ChildPattern( + "annotation", + ChoicePattern(WhereAnnotationPattern(), TypeAnnotationPattern()), + "annotation", + "local", + ), + ), + FieldPattern( + "value", + OptionalPattern( + ChildPattern( + "value", + ChoicePattern(ExpressionPattern(), StaticValuePattern()), + "expression", + "assignment_value", + ) + ), + ), + ), + pattern_id="statement.annassign", + ), + BranchPattern( + "return", + AstNodePattern( + ast.Return, + FieldPattern( + "value", + OptionalPattern( + ChildPattern( + "value", + ExpressionPattern(), + "expression", + "return_value", + ) + ), + ), + ), + pattern_id="statement.return", + ), + BranchPattern( + "expr_statement", + AstNodePattern( + ast.Expr, + FieldPattern( + "value", + ChildPattern( + "value", + ExpressionPattern(), + "expression", + "statement_value", + ), + ), + ), + pattern_id="statement.expr", + ), + BranchPattern("none", AstNodePattern(ast.Pass), pattern_id="statement.pass"), + ) + ) + + @staticmethod + def construct(match, children, context): + if match.branch_id == "assignment": + value = children.get("value") + annotation = children.get("annotation") + if value is None: + raise ParseError.from_node(match.node, context, "assignment requires a value") + if context.function.dialect == "tir" and not isinstance(value, runtime.Expr): + context.lexical_scope.define(match.captures["name"], value) + return None + if isinstance(annotation, ScheduleConstraintMetadata): + if context.function.dialect != "hir" or not isinstance(value.type, TensorType): + raise ParseError.from_node( + match.node, + context, + "where annotation requires a tensor-valued HIR Expr", + ) + previous = get_metadata(value, ScheduleConstraintMetadata) + if previous is not None: + binding = get_metadata(value, BindingMetadata) + label = binding.name if binding is not None else "" + raise ParseError.from_node( + match.node, + context, + f"duplicate where annotation for Expr {label!r}", + ) + value = replace_metadata(value, BindingMetadata(match.captures["name"])) + object.__setattr__(value, "metadata", (*value.metadata, annotation)) + elif annotation is not None and value.type != annotation: + raise ParseError.from_node( + match.node, context, "annotated assignment type mismatch" + ) + name = match.captures["name"] + if context.function.dialect == "hir": + if isinstance(value, runtime.Call) and get_metadata(value, BindingMetadata) is None: + value = replace_metadata(value, BindingMetadata(name)) + context.lexical_scope.define(name, value) + return value + variable = runtime.Var(type=value.type, name=name) + context.lexical_scope.define(name, variable) + return runtime.LetStmt(variable, value, runtime.Sequential(body=())) + elif match.branch_id == "return": + if context.function.dialect == "tir": + if "value" in children: + raise ParseError.from_node(match.node, context, "prim_func return must be bare") + return runtime.Return() + if "value" not in children: + raise ParseError.from_node(match.node, context, "func return must carry a value") + return children["value"] + elif match.branch_id == "expr_statement": + value = children["value"] + if context.function.dialect == "hir": + raise ParseError.from_node( + match.node, context, "HIR does not allow expression statements" + ) + if isinstance(value, runtime.Evaluate): + return value + if not isinstance(value, runtime.Call) or not isinstance(value.type, runtime.UnitType): + raise ParseError.from_node( + match.node, context, "TIR expression statement must be unit Call" + ) + return runtime.Evaluate(callable=value.target, args=value.args) + elif match.branch_id == "none": + return None + raise RuntimeError(f"no constructor branch for {match.branch_id!r}") + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +class BlockPattern(ElementPattern): + element_name = "block" + syntax = LazyPattern( + lambda: BranchPattern( + "block", + AstNodePattern( + ast.Module, + CapturePattern( + "pass_only", + lambda node, context: ( + len(node.body) == 1 and isinstance(node.body[0], ast.Pass) + ), + ), + CapturePattern( + "terminal_children", + lambda node, context: tuple( + f"statement_{index}" + for index, statement in enumerate(node.body) + if isinstance(statement, (ast.Return, ast.With, ast.For)) + ), + ), + FieldPattern( + "body", + RepeatPattern( + ChildPattern( + "statement_{index}", + StatementPattern(), + "statement", + "statement", + ) + ), + ), + ), + pattern_id="function.block", + ) + ) + + @staticmethod + def construct(match, children, context): + values = list(children.values()) + if context.function.dialect == "hir": + if match.captures["pass_only"]: + return None + for child_name in reversed(match.captures["terminal_children"]): + value = children[child_name] + if value is not None: + return value + if context.role == "with_body": + return None + raise ParseError.from_node(match.node, context, "HIR body must end with return") + + + def fold(index: int) -> list[object]: + output: list[object] = [] + while index < len(values): + value = values[index] + if value is None: + index += 1 + continue + if isinstance(value, runtime.LetStmt): + nested = runtime.Sequential(body=tuple(fold(index + 1))) + output.append(dataclasses.replace(value, body=nested)) + return output + output.append(value) + index += 1 + return output + + return runtime.Sequential(body=tuple(fold(0))) + + RULES: ClassVar[tuple[AstRule[Any], ...]] = () + + +@dataclass(frozen=True) +class FunctionSignatureRule: + STATEMENT: ClassVar[str] = "A function must construct an ordered parameter tuple." + + def apply(self, value, *, match, context): + if not isinstance(getattr(value, "params", None), tuple): + raise ParseError.from_node(match.node, context, "function params were not constructed") + return value + + +@dataclass(frozen=True) +class FunctionReturnRule: + STATEMENT: ClassVar[str] = "A HIR function body's inferred type must match its return type." + + def apply(self, value, *, match, context): + if isinstance(value, runtime.Function) and value.body is not None: + infer_context = context.lexical_scope.lookup(_TYPE_INFER_CONTEXT) + if not isinstance(infer_context, runtime.TypeInferContext): + infer_context = runtime.TypeInferContext() + body_type = runtime.TypeInferVisitor(infer_context).visit(value.body) + if not _types_compatible(body_type, value.return_type): + raise ParseError.from_node( + match.node, + context, + f"function body type {body_type!r} does not match return {value.return_type!r}", + ) + return value + + +@dataclass(frozen=True) +class FunctionDialectRule: + STATEMENT: ClassVar[str] = ( + "A function kind and constructed value must agree with the active dialect." + ) + + def apply(self, value, *, match, context): + kind = context.function.function_kind + if context.function.dialect == "hir" and kind == "prim_func": + raise ParseError.from_node(match.node, context, "prim_func requires tir dialect") + if context.function.dialect == "tir" and kind != "prim_func": + raise ParseError.from_node(match.node, context, f"{kind} requires hir dialect") + expected = runtime.Function if context.function.dialect == "hir" else runtime.PrimFunction + if not isinstance(value, expected): + raise ParseError.from_node( + match.node, + context, + f"{context.function.dialect} context constructed {type(value).__name__}", + ) + return value + + +@dataclass(frozen=True) +class FunctionRoleValidationRule: + STATEMENT: ClassVar[str] = ( + "A root, variant, or converter must satisfy its role before registration." + ) + + @staticmethod + def _validate_standalone( + function: object, + function_context: FuncParserContext, + node: ast.AST, + match_context: MatchContext, + ) -> None: + if function_context.role is FunctionRole.ROOT: + return + base = function_context.base + if not isinstance(base, runtime.Function): + raise ParseError.from_node(node, match_context, "standalone role lacks a HIR base") + if getattr(base, "_sealed", False): + raise ParseError.from_node(node, match_context, f"base {base.name!r} is sealed") + if getattr(function, "body", None) is None: + raise ParseError.from_node( + node, + match_context, + f"{function_context.role.value} must have a real body", + ) + if function_context.role is FunctionRole.CONVERTER and not isinstance( + function_context.key, str + ): + raise ParseError.from_node(node, match_context, "converter key must be a weight name") + + def apply(self, value, *, match, context): + if context.function is None: + raise ParseError.from_node(match.node, context, "function lacks parser context") + if context.function.module is not None: + context.function.module.validate_function(value, context.function) + else: + self._validate_standalone(value, context.function, match.node, context) + return value + + +@dataclass(frozen=True) +class FunctionRegistrationRule: + STATEMENT: ClassVar[str] = ( + "A validated function must be registered exactly once in its owning scope." + ) + + @staticmethod + def _commit_standalone( + function: object, + function_context: FuncParserContext, + ) -> None: + if function_context.role is FunctionRole.VARIANT: + assert function_context.base is not None + function_context.base.add_variant(function) + elif function_context.role is FunctionRole.CONVERTER: + assert function_context.base is not None + function_context.base.add_converter(function_context.key, function) + + def apply(self, value, *, match, context): + if context.function is None: + raise ParseError.from_node(match.node, context, "function lacks parser context") + if context.function.module is not None: + context.function.module.commit_function(value, context.function) + else: + self._commit_standalone(value, context.function) + return value + + +class FunctionPattern(ElementPattern): + element_name = "function" + syntax = LazyPattern( + lambda: BindPattern( + AstNodePattern( + ast.FunctionDef, + FieldPattern("name", CapturePattern("name", lambda value, context: value)), + FieldPattern( + "args", + ChildPattern("signature", SignaturePattern(), "signature", "function"), + ), + FieldPattern( + "returns", + OptionalPattern( + ChildPattern( + "return", + ReturnTypePattern(), + "type_annotation", + "return", + ) + ), + ), + FieldPattern( + "body", + ChildPattern( + "body", + BlockPattern(), + "block", + "body", + transform=_module_from_body, + ), + ), + ), + FunctionPattern._bind, + ) + ) + + @staticmethod + def _bind(node: object, context: MatchContext, matched: AstMatch[Any]) -> AstMatch[Any] | None: + assert isinstance(node, ast.FunctionDef) + function_context = context.function + if function_context is None: + return None + active_context = context.child( + situation="function", + role=function_context.function_kind, + function=function_context, + ) + return dataclasses.replace( + matched, + pattern_id="function", + branch_id="function", + construct_context=active_context, + ) + + @staticmethod + def construct(match, children, context): + if context.function is None: + raise ParseError.from_node(match.node, context, "function lacks context") + params = children["signature"] + body = children["body"] + declared_return = children.get("return") + specializations = context.function.specializations + converter = context.function.converter + if context.function.dialect == "hir": + if declared_return is None: + if body is None: + raise ParseError.from_node( + match.node, + context, + "HIR pass prototype requires a return annotation", + ) + declared_return = body.type + elif ( + isinstance(declared_return, runtime.TensorType) + and isinstance(body, runtime.Expr) + and isinstance(body.type, runtime.TensorType) + and declared_return.shape == body.type.shape + and declared_return.dtype == body.type.dtype + and declared_return.storage == body.type.storage + and isinstance(declared_return.layout, runtime.ShardLayout) + and isinstance(body.type.layout, runtime.ShardLayout) + and isinstance(declared_return.layout.layout, runtime.Layout) + and isinstance(body.type.layout.layout, runtime.Layout) + and declared_return.layout.attrs == body.type.layout.attrs + and declared_return.layout.mesh == body.type.layout.mesh + and declared_return.layout.layout.strides is None + and body.type.layout.layout.strides is not None + ): + declared_return = runtime.TensorType( + shape=declared_return.shape, + dtype=declared_return.dtype, + layout=body.type.layout, + storage=declared_return.storage, + ) + function_name = ( + getattr(context.function.base, "name", None) + or context.function.base_name + or match.captures["name"] + ) + function = runtime.Function.build( + name=function_name, + params=params, + body=body, + return_type=declared_return, + specializations=specializations, + ) + if context.function.role is FunctionRole.VARIANT: + object.__setattr__(function, runtime.DISPLAY_NAME, match.captures["name"]) + object.__setattr__(function, "name", function_name) + elif context.function.role is FunctionRole.CONVERTER: + object.__setattr__(function, "name", f"{function_name}.converter[{converter}]") + binding = context.function.binding_name or match.captures["name"] + define = getattr(context.function.module_scope, "define", None) + if callable(define): + define(binding, function) + return function + if declared_return is not None: + raise ParseError.from_node(match.node, context, "prim_func cannot return a value type") + kwargs = {} + if context.function.target is not None: + kwargs["target"] = context.function.target + function = runtime.PrimFunction( + name=match.captures["name"], + params=params, + body=body, + output_count=context.function.output_count, + **kwargs, + ) + define = getattr(context.function.module_scope, "define", None) + if callable(define): + define(context.function.binding_name or match.captures["name"], function) + return function + + RULES: ClassVar[tuple[AstRule[Any], ...]] = ( + FunctionSignatureRule(), + FunctionReturnRule(), + FunctionDialectRule(), + FunctionRoleValidationRule(), + FunctionRegistrationRule(), + ) + + +def _module_from_body(body: object) -> ast.Module: + assert isinstance(body, list) + return ast.Module(body=body, type_ignores=[]) + + +__all__ = [ + "BinaryExpressionPattern", + "BlockPattern", + "CallBindingRule", + "CallExpectedTypeRule", + "CallPattern", + "CallTypeInferenceRule", + "ConstantPattern", + "DTypePattern", + "DimExprPattern", + "ExplicitLayoutPattern", + "ExpressionPattern", + "ForPattern", + "FunctionDialectRule", + "FunctionPattern", + "FunctionRegistrationRule", + "FunctionReturnRule", + "FunctionRoleValidationRule", + "FunctionSignatureRule", + "IndexEndpointPattern", + "IndexSlicePattern", + "LayoutPattern", + "LoopBodyPattern", + "LoopCarryPattern", + "LoopCarryStatementPattern", + "LoopHeaderPattern", + "MeshAxisPattern", + "MeshContextPattern", + "MeshCoordinatePattern", + "NamePattern", + "PlacedLayoutPattern", + "PlainLayoutPattern", + "ReturnTypePattern", + "ScalarTypePattern", + "ShapePattern", + "SignaturePattern", + "StatementPattern", + "StaticBinaryPattern", + "StaticCallPattern", + "StaticDictPattern", + "StaticLiteralPattern", + "StaticReferencePattern", + "StaticSequencePattern", + "StaticSlicePattern", + "StaticSubscriptPattern", + "StaticUnaryPattern", + "StaticValuePattern", + "StoragePattern", + "SubscriptExpressionPattern", + "SubscriptIndexPattern", + "TensorOptionalSlotPattern", + "TensorPattern", + "TensorShapeLayoutPattern", + "TupleAssignmentPattern", + "TupleExpressionPattern", + "TypeAnnotationPattern", + "UnaryExpressionPattern", + "WithPattern", +] diff --git a/src/tilefoundry/parser/spec.py b/src/tilefoundry/parser/spec.py new file mode 100644 index 00000000..7283558d --- /dev/null +++ b/src/tilefoundry/parser/spec.py @@ -0,0 +1,292 @@ +"""Generate the private parser grammar and constraint reference.""" + +from __future__ import annotations + +import argparse +import difflib +import inspect +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from .ast_pattern import ( + AstNodePattern, + BindPattern, + BranchPattern, + CapturePattern, + ChildPattern, + ChoicePattern, + ConditionPattern, + ElementPattern, + FieldPattern, + LazyPattern, + LiteralPattern, + ModuleBuildContext, + OptionalPattern, + PredicatePattern, + ReferencePattern, + RepeatPattern, + SequencePattern, +) +from .grammar_render import render_grammar +from .pattern_nodes import FunctionPattern + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[3] + + +@dataclass(frozen=True, order=True) +class RuleRow: + owner: str + situation: str + rule: str + statement: str + source: str + + +def _source(rule: object) -> str: + filename = inspect.getsourcefile(type(rule)) + if filename is None: + return "" + path = Path(filename).resolve() + try: + return path.relative_to(_REPOSITORY_ROOT).as_posix() + except ValueError: + return path.as_posix() + + +def _row(owner: str, situation: str, rule: object) -> RuleRow: + return RuleRow( + owner=owner, + situation=situation, + rule=type(rule).__name__, + statement=rule.STATEMENT, + source=_source(rule), + ) + + +class _RuleVisitor: + def __init__(self) -> None: + self._seen_elements: set[tuple[str, str]] = set() + self._rows: set[RuleRow] = set() + + def visit(self, pattern: object, situation: str) -> None: + if isinstance(pattern, ElementPattern): + name = pattern.element_name + if not name: + raise TypeError(f"{type(pattern).__name__} has no element_name") + key = (name, situation) + if key in self._seen_elements: + return + self._seen_elements.add(key) + self._rows.update(_row(name, situation, rule) for rule in pattern.RULES) + if pattern.syntax is None: + raise TypeError(f"{type(pattern).__name__} has no executable syntax") + self.visit(pattern.syntax, situation) + return + if isinstance(pattern, LazyPattern): + self.visit(pattern.pattern, situation) + return + if isinstance(pattern, ChildPattern): + self.visit(pattern.pattern, pattern.situation) + return + if isinstance(pattern, AstNodePattern): + for part in pattern.parts: + self.visit(part, situation) + return + if isinstance(pattern, (ChoicePattern, SequencePattern)): + for item in pattern.patterns: + self.visit(item, situation) + return + if isinstance( + pattern, + ( + BindPattern, + BranchPattern, + ConditionPattern, + FieldPattern, + OptionalPattern, + RepeatPattern, + ), + ): + self.visit(pattern.pattern, situation) + return + if isinstance( + pattern, + (CapturePattern, LiteralPattern, PredicatePattern, ReferencePattern), + ): + return + raise TypeError(f"unsupported executable pattern {type(pattern).__name__}") + + def rows(self) -> tuple[RuleRow, ...]: + return tuple(sorted(self._rows)) + + +def _collect_rule_rows(root: ElementPattern[Any]) -> tuple[RuleRow, ...]: + visitor = _RuleVisitor() + visitor.visit(root, "function") + return visitor.rows() + + +def _collect_module_rule_rows() -> tuple[RuleRow, ...]: + rows = [ + _row("module", "module_function", rule) + for rule in ModuleBuildContext.FUNCTION_RULES + ] + rows.extend( + _row("module", "module_finalization", rule) + for rule in ModuleBuildContext.FINALIZATION_RULES + ) + return tuple(rows) + + +def _escape_cell(value: str) -> str: + return value.replace("|", "\\|").replace("\n", " ") + + +def render_spec_content() -> str: + root = FunctionPattern() + rows = (*_collect_rule_rows(root), *_collect_module_rule_rows()) + lines = [ + "# Parser Grammar and Constraints", + "", + "```ebnf", + render_grammar(root), + "```", + "", + "| Owner | Situation | Rule | Statement | Source |", + "| --- | --- | --- | --- | --- |", + ] + lines.extend( + "| " + + " | ".join( + _escape_cell(value) + for value in ( + row.owner, + row.situation, + row.rule, + row.statement, + row.source, + ) + ) + + " |" + for row in rows + ) + return "\n".join(lines) + "\n" + + +def render_parser_document() -> str: + """Render the checked three-section Parser Spec document.""" + generated = render_spec_content().removeprefix("# Parser Grammar and Constraints\n\n") + return f'''# TileFoundry Spec - Parser + +The Parser accepts authored Python functions and produces HIR or TIR through one typed API. + +## 1. Public API + +`@module` executes its Python class body and finalizes the collected Functions, child Modules, +and ordinary methods. `@func` produces an HIR Function; `@prim_func` produces a TIR PrimFunction. +`specialize` and `converter` register variants and weight converters on an existing HIR Function. + +```python +def parse_function( + fn: FunctionType, context: FuncParserContext +) -> hir.Function | tir.PrimFunction: ... +``` + +`FuncParserContext` carries the dialect, Function role, closure, topology scope, target, and +optional base/key for one parse. `FunctionRole` is `ROOT`, `VARIANT`, or `CONVERTER`. +`ParseError` is the single authored-source diagnostic type and includes source location and +recursive parse situation. These are the only public parser symbols. + +## 2. Syntax and Rules + +### 2.1 Syntax + + +{generated.split('| Owner | Situation | Rule | Statement | Source |', 1)[0].rstrip()} + + +### 2.2 Rules + + +| Owner | Situation | Rule | Statement | Source | +| --- | --- | --- | --- | --- | +{generated.split('| Owner | Situation | Rule | Statement | Source |', 1)[1].split('| --- | --- | --- | --- | --- |', 1)[1].lstrip()} + + +## 3. Implementation Overview + +| Component | Responsibility | +| --- | --- | +| Parser API and Context | Receives authored Functions and carries dialect, role, scope, and recursion inputs. | +| Executable Pattern Graph | Composes concrete AST elements into the Function root pattern. | +| Match and Construction | Matches recursively into `AstMatch`, then constructs owner values on return. | +| Ordered Rules | Validates and normalizes each owner value after construction. | +| Module Build | Lets Python execute the class body, records Functions, and finalizes the Module. | +| Pattern Visitor | Traverses the same graph to render this section's generated grammar and constraints. | + +```mermaid +classDiagram + ParserAPI --> FuncParserContext + ParserAPI --> FunctionPattern + AstPattern <|.. Element + Element o-- AstPattern + Element o-- AstRule + AstPattern --> AstMatch + PatternVisitor ..> AstPattern + ParserAPI ..> ModuleBuild +``` + +```mermaid +flowchart TD + API["parse_function(fn, context)"] --> AST["Extract FunctionDef AST"] + AST --> ROOT["FunctionPattern.match"] + ROOT --> TREE["AstMatch tree"] + TREE --> BACKWARD["construct children, then apply Rules"] + BACKWARD --> FUNCTION["HIR Function / TIR PrimFunction"] + FUNCTION --> MODULE{{"Module authoring context?"}} + MODULE -->|yes| FINALIZE["registration / finalization"] + MODULE -->|no| RETURN["return standalone result"] +``` + +Pattern combinators serve both runtime matching and Spec traversal. `AstMatch` separates syntax +matching from object construction, while each Rule reads the recursive context after its owner +value exists. Module class control flow remains Python execution; no Module AST grammar exists. +''' + + +def _parse_args(argv: list[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + output = parser.add_mutually_exclusive_group() + output.add_argument("--write", type=Path, metavar="PATH") + output.add_argument("--check", type=Path, metavar="PATH") + return parser.parse_args(argv) + + +def _main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + generated = render_spec_content() + if args.write is not None: + args.write.write_text(render_parser_document()) + return 0 + if args.check is not None: + expected = render_parser_document() + actual = args.check.read_text() if args.check.exists() else "" + if actual == expected: + return 0 + sys.stderr.writelines( + difflib.unified_diff( + actual.splitlines(keepends=True), + expected.splitlines(keepends=True), + fromfile=str(args.check), + tofile="generated parser spec", + ) + ) + return 1 + sys.stdout.write(generated) + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/src/tilefoundry/parser/sugar.py b/src/tilefoundry/parser/sugar.py deleted file mode 100644 index f9ac9d87..00000000 --- a/src/tilefoundry/parser/sugar.py +++ /dev/null @@ -1,850 +0,0 @@ -"""Unified layout sugar parser. - -Core model: tuple sugar is a type-directed layout literal. -Shared bottom layer ``_parse_layout_literal()`` extracts shape + strides -from a tuple AST node. Target-specific entry points lower the literal -to ``Layout`` or ``ShardLayout``. - -Consumers call :func:`parse_sugar` with the expected result type; contextual -differences are closure lookup and mesh resolution. -""" - -from __future__ import annotations - -import ast -from typing import Any, Callable - -from tilefoundry.ir.core import VerifyError -from tilefoundry.ir.core.expr import Expr -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim -from tilefoundry.ir.types.shape_dim import ShapeDim -from tilefoundry.ir.types.shard import c_order_strides -from tilefoundry.ir.types.shard.layout import Layout -from tilefoundry.ir.types.shard.mesh import Mesh, composed -from tilefoundry.ir.types.shard.shard_layout import ( - Broadcast, - Partial, - ShardAttr, - ShardLayout, - Split, -) -from tilefoundry.ir.types.storage import StorageKind, resolve_storage -from tilefoundry.utils.spec_ref import spec_ref_render - -from .static_eval import eval_static - -_SHARD_ATTR = "[shard §6](docs/spec/shard.md#6-shardattr)" - - -class LayoutSugarError(VerifyError): - """Report a structurally recognized but malformed layout-sugar node. - - A layout-sugar node was recognized structurally but is malformed - (e.g. a dynamic ``DimVar`` / ``bool`` static extent). - - It subclasses ``VerifyError`` (itself a ``ValueError``) so both - still catch it, but callers that speculatively try sugar parsing (and fall - back to generic static evaluation on a plain ``ValueError``) MUST let this - propagate so the real diagnostic is not masked by a downstream error. - """ - - -def _is_constant(node: ast.AST) -> bool: - return isinstance(node, ast.Constant) - - -def _is_layout_slot_constant(node: ast.AST) -> bool: - """Whether a constant in the third annotation slot is a layout rather than storage. - - ``Tensor[shape, dtype, storage]`` puts a storage name where a layout would - go, and ``Tensor[shape, dtype, None, storage]`` leaves that slot empty - ([parser §1.4](docs/spec/parser.md#14-tensor-and-consttensor-annotations) - makes both optional and independent). Neither is layout sugar, so neither - may pull the annotation onto the sugar path. - """ - return _is_constant(node) and node.value is not None and not isinstance(node.value, str) - - -def _is_matmul(node: ast.AST) -> bool: - return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) - - -def _is_placeholder(node: ast.AST) -> bool: - return isinstance(node, ast.Name) and node.id == "_" - - -def _is_strided_layout_tuple(node: ast.AST) -> bool: - return ( - isinstance(node, ast.Tuple) - and len(node.elts) == 2 - and isinstance(node.elts[0], ast.Tuple) - and isinstance(node.elts[1], ast.Tuple) - ) - - -def _is_tuple_sugar(node: ast.AST) -> bool: - """Check whether an AST node is a tuple literal that could be layout sugar. - - Returns True for Tuple nodes (which may contain ``@`` operators). - Bare Constant (single int) is checked separately by the consumer - based on whether meshes are available. - """ - if isinstance(node, ast.Tuple): - return True - return False - - -def _has_sugar(node: ast.AST) -> bool: - """Check whether an AST node contains a ``@`` sugar operator.""" - found = False - - def visitor(n: ast.AST): - nonlocal found - if found: - return - if _is_matmul(n): - found = True - return - for _field, child in ast.iter_fields(n): - if isinstance(child, ast.AST): - visitor(child) - elif isinstance(child, list): - for item in child: - if isinstance(item, ast.AST): - visitor(item) - - visitor(node) - return found - - -_EVAL_AST_NODES_NO_CLOSURE = (ast.Constant, ast.Tuple, ast.UnaryOp) -_EVAL_AST_NODES_WITH_CLOSURE = ( - *_EVAL_AST_NODES_NO_CLOSURE, - ast.Name, - ast.Attribute, - ast.Call, - ast.BinOp, -) - - -def _eval_ast(node: ast.AST, closure: dict[str, Any] | None = None) -> Any: - """Evaluate layout literals, inline dimensions, and closure dimensions. - - Translate the shared static evaluator's ``VerifyError`` to ``ValueError``. - Names and attributes resolve only when a closure is supplied; inline - ``DimVar`` calls are recognized directly rather than as general callees. - """ - if ( - closure is not None - and isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "DimVar" - ): - pos = [_eval_ast(a, closure) for a in node.args] - kw = {k.arg: _eval_ast(k.value, closure) for k in node.keywords} - return DimVar(*pos, **kw) - allowed = _EVAL_AST_NODES_WITH_CLOSURE if closure is not None else _EVAL_AST_NODES_NO_CLOSURE - try: - return eval_static(node, closure=closure or {}, allowed_nodes=allowed) - except VerifyError as exc: - raise VerifyError(str(exc)) from exc - - -def _is_shape_dim(v: Any) -> bool: - """True for a valid layout axis extent: ``ShapeDim = int | DimVar | Expr``. - - ``bool`` is a subclass of ``int`` but is rejected (never a real extent). - """ - if isinstance(v, bool): - return False - return isinstance(v, (int, DimVar, Expr)) - - -def _name_of(node: ast.AST) -> str: - """Extract bare Name id from an AST node.""" - if isinstance(node, ast.Name): - return node.id - raise VerifyError(f"expected Name, got {ast.dump(node)}") - - -def _dim_mul(a: ShapeDim, b: ShapeDim) -> ShapeDim: - if isinstance(a, int) and isinstance(b, int): - return a * b - return simplify_dim(DimMul, (a, b)) - - -def _auto_strides(shape: tuple[ShapeDim, ...]) -> tuple[ShapeDim, ...]: - """C-order contiguous strides: ``(d0, d1, d2)`` → ``(d1*d2, d2, 1)``.""" - return c_order_strides(shape, mul=_dim_mul) - - -def _resolve_dtype_ast(node: ast.AST, closure: dict[str, Any]) -> DType | None: - """Resolve a dtype from an AST node (bare name, string, or DType.attr).""" - if isinstance(node, ast.Constant) and isinstance(node.value, str): - return DType._members().get(node.value) - if isinstance(node, ast.Name): - val = closure.get(node.id) - if isinstance(val, DType): - return val - if isinstance(val, str): - return DType._members().get(val) - return DType._members().get(node.id) - if isinstance(node, ast.Attribute): - try: - val = _eval_ast(node, closure) - if isinstance(val, DType): - return val - except ValueError: - pass - return None - - -def _resolve_mesh(name: str, mesh_by_name: dict[str, Mesh]) -> Mesh: - """Look up a Mesh by variable name.""" - mesh = mesh_by_name.get(name) - if mesh is None: - available = list(mesh_by_name.keys()) - raise VerifyError(f"undefined mesh {name!r}; available: {available}") - if not isinstance(mesh, Mesh): - raise VerifyError(f"{name!r} is not a Mesh, got {type(mesh).__name__}") - return mesh - - -def _resolve_mesh_axis(mesh: Mesh, axis_name: str) -> int: - """Resolve a mesh axis by name (preferred) or x/y/z position fallback. - - - If *mesh* has ``names``, resolve by matching name. - - If *axis_name* is ``"x"``, ``"y"``, or ``"z"``, resolve by position. - - Otherwise, raises ``ValueError``. - """ - for index, name in enumerate(mesh.names): - if name == axis_name: - return index - if axis_name in ("x", "y", "z"): - index = ("x", "y", "z").index(axis_name) - if index < len(mesh.layout.shape): - return index - available = list(mesh.names) if mesh.names else ["x", "y", "z"][: len(mesh.layout.shape)] - raise VerifyError(f"mesh has no axis named {axis_name!r}; available: {available}") - - -def _parse_layout_literal( - node: ast.AST, *, closure: dict[str, Any] | None = None -) -> tuple[tuple[ShapeDim, ...], tuple[ShapeDim, ...] | None]: - """Parse layout shape and optional strides without choosing a final type. - - Accept flat dimensions, explicit ``(dimensions, strides)``, or a scalar 1-D - form. For placement sugar, extract the left-hand dimension while a target - parser handles axis binding. A closure may resolve named static dimensions. - """ - if isinstance(node, ast.Tuple): - if ( - len(node.elts) == 2 - and isinstance(node.elts[0], ast.Tuple) - and isinstance(node.elts[1], ast.Tuple) - ): - dim_nodes = list(node.elts[0].elts) - strides = _eval_ast(node.elts[1], closure) - shape = tuple(_extract_dim(dn, closure=closure) for dn in dim_nodes) - else: - dim_nodes = list(node.elts) - strides = None - shape = tuple(_extract_dim(dn, closure=closure) for dn in dim_nodes) - elif _is_constant(node): - shape = (_extract_dim(node, closure=closure),) - strides = None - else: - try: - resolved = _eval_ast(node, closure) - except ValueError: - resolved = None - if not isinstance(resolved, tuple) or not all(_is_shape_dim(dim) for dim in resolved): - raise VerifyError(f"expected tuple layout literal, got {ast.dump(node)}") - shape = resolved - strides = None - - if strides is not None: - if not isinstance(strides, tuple): - raise VerifyError(f"strides must be a tuple, got {strides!r}") - if len(strides) != len(shape): - raise VerifyError(f"strides rank {len(strides)} != layout shape rank {len(shape)}") - - return shape, strides - - -def _extract_dim(node: ast.AST, *, closure: dict[str, Any] | None = None) -> ShapeDim: - """Extract a static or symbolic dimension from a layout dim node. - - Handles: plain ``Constant(32)``, closure/global ``Name`` references bound to - a ``ShapeDim``, and ``BinOp(, MatMult, ...)`` sugar forms where only the - left operand is the dimension. ``bool`` and non-dimension values are rejected - with a clear diagnostic rather than a raw AST / attribute error. - """ - dim_node = node.left if _is_matmul(node) else node - if _is_constant(dim_node): - val = dim_node.value - else: - try: - val = _eval_ast(dim_node, closure) - except ValueError: - raise VerifyError(f"expected shape dimension, got {ast.dump(node)}") from None - if not _is_shape_dim(val): - if isinstance(val, bool): - raise LayoutSugarError( - f"layout dim must be a shape dimension, and bool {val!r} is not one; bool is " - "an int subclass, so it is refused here rather than silently becoming 0 or 1" - ) - raise LayoutSugarError( - f"layout dim must be a shape dimension (int / DimVar / dim-op Expr), got " - f"{type(val).__name__} {val!r}" - ) - return val - - -def _parse_layout_sugar(node: ast.AST, *, closure: dict[str, Any] | None = None) -> Layout: - shape, strides = _parse_layout_literal(node, closure=closure) - if strides is None: - strides = _auto_strides(shape) - return Layout(shape=shape, strides=strides) - - -MeshResolver = Callable[[str], Mesh] - - -def _parse_shard_layout_sugar( - node: ast.AST, - mesh_resolver: MeshResolver, - *, - default_mesh: Mesh | None = None, - closure: dict[str, Any] | None = None, - mesh_order: "tuple[Mesh, ...]" = (), -) -> ShardLayout: - """Parse placement and value-state sugar into a shard layout. - - Resolve named meshes, using *default_mesh* only for all-broadcast layouts. - Bare dimensions broadcast, ``dim @ mesh.axis`` splits, and a final set maps - mesh axes to partial reductions. Unmentioned mesh axes broadcast. Missing - both explicit and default mesh information is an error. - See [parser §1.5](docs/spec/parser.md#15-layout-sugar) and - [shard §6](docs/spec/shard.md#6-shardattr). - """ - axis_node, strides, value_set_node = _split_layout_outer(node) - - dim_nodes = _get_dim_nodes(axis_node) - - canonicalize = strides is None - parsed: list[_LayoutItem] = [] - for dn in dim_nodes: - parsed.extend( - _parse_layout_item(dn, mesh_resolver, canonicalize=canonicalize, closure=closure) - ) - - value_states = ( - _parse_value_state(value_set_node, mesh_resolver) if value_set_node is not None else [] - ) - - named: list[Mesh] = [] - for _d, mesh, _mi, _k, _r in parsed: - if mesh is not None and not any(item is mesh for item in named): - named.append(mesh) - for mesh, _mi, _r in value_states: - if not any(item is mesh for item in named): - named.append(mesh) - - if not named: - if default_mesh is None: - raise VerifyError( - "all-Broadcast ShardLayout sugar requires a mesh from " - "context; use verbose ShardLayout(...) to disambiguate" - ) - named = [default_mesh] - ordered = _nesting_order(named, mesh_order) - - shape: list[int] = [] - axis_of: list[int | None] = [] - for dim, _mesh, _m_axis, _kind, _reduction in parsed: - if dim is not None: - shape.append(dim) - axis_of.append(len(shape) - 1) - else: - axis_of.append(None) - - attrs_list: list[ShardAttr] = [] - for mesh in ordered: - mesh_rank = len(mesh.layout.shape) - own: list[ShardAttr] = [Broadcast() for _ in range(mesh_rank)] - for index, (_dim, item_mesh, m_axis, kind, _reduction) in enumerate(parsed): - if kind != "split" or item_mesh is not mesh: - continue - layout_axis = axis_of[index] - if m_axis is None or m_axis >= mesh_rank: - raise VerifyError(f"layout dim {layout_axis}: invalid mesh axis {m_axis}") - if not isinstance(own[m_axis], Broadcast): - raise VerifyError( - f"mesh axis {m_axis} already bound; " - f"one layout dim per mesh axis ({spec_ref_render(_SHARD_ATTR)})" - ) - own[m_axis] = Split(layout_axis) - for item_mesh, m_axis, reduction in value_states: - if item_mesh is not mesh: - continue - if m_axis >= mesh_rank: - raise VerifyError(f"value-state: invalid mesh axis {m_axis}") - if not isinstance(own[m_axis], Broadcast): - raise VerifyError(f"mesh axis {m_axis} already bound") - own[m_axis] = Partial(reduction or "sum") - attrs_list.extend(own) - - try: - resolved_mesh = composed(tuple(ordered)) - except ValueError as error: - raise VerifyError(str(error)) from None - return ShardLayout( - layout=Layout(shape=tuple(shape), strides=strides), - attrs=tuple(attrs_list), - mesh=resolved_mesh, - ) - - -def _nesting_order(named: list[Mesh], mesh_order: "tuple[Mesh, ...]") -> list[Mesh]: - """The meshes one layout names, outermost scope first. - - A value can be distributed at more than one level at once -- a CTA owns a - tile and a lane owns part of that tile -- and saying so takes both meshes. - Which is inside which is not something a layout can be read for, so it is - taken from the scopes the layout was written in; two meshes that are not - nested have no such answer and are refused rather than ordered by guess. - """ - if len(named) == 1: - return named - position = {id(mesh): index for index, mesh in enumerate(mesh_order)} - missing = [mesh for mesh in named if id(mesh) not in position] - if missing: - raise VerifyError( - "a layout naming several meshes needs them nested in one another, so " - "which distributes which is stated rather than guessed; " - f"{len(missing)} of them is not a scope this layout is written inside" - ) - return sorted(named, key=lambda mesh: position[id(mesh)]) - - -def _split_layout_outer( - node: ast.AST, -) -> tuple[ast.AST, "tuple | None", "ast.Set | None"]: - """Split a layout-sugar node into (axis-tuple node, strides, value-state set). - - Outer-tuple grammar (parser layout sugar): - - ``(d0, d1, ...)`` → implicit strides, no value-state - - ``((dims), (strides))`` → explicit strides - - ``((dims), {value-state})`` → implicit strides + value-state - - ``((dims), (strides), {value-state})`` → explicit strides + value-state - - The value-state `set` literal (if present) MUST be the last outer item. - """ - if _is_constant(node) or _is_matmul(node): - return node, None, None - if not isinstance(node, ast.Tuple): - raise VerifyError(f"expected tuple layout, got {ast.dump(node)}") - - if node.elts and isinstance(node.elts[0], ast.Tuple): - axis_node = node.elts[0] - strides = None - value_set: ast.Set | None = None - for elt in node.elts[1:]: - if value_set is not None: - raise VerifyError("layout sugar: the value-state set must be the last outer item") - if isinstance(elt, ast.Set): - value_set = elt - elif isinstance(elt, ast.Tuple): - if strides is not None: - raise VerifyError("layout sugar: at most one stride tuple") - strides = _eval_ast(elt) - else: - raise VerifyError( - f"layout sugar outer item must be a stride tuple or value-state " - f"set, got {ast.dump(elt)}" - ) - return axis_node, strides, value_set - - return node, None, None - - -def _parse_value_state(node: "ast.Set", mesh_resolver: MeshResolver) -> list[tuple[Mesh, int, str]]: - """Parse value-state entries into mesh-axis reductions. - - Parse a ``{mesh.axis @ P("reduction"), ...}`` value-state set into a list - of ``(mesh, mesh_axis_index, reduction)``. Element order carries no meaning. - """ - if not isinstance(node, ast.Set): - raise VerifyError(f"value-state must be a set literal, got {ast.dump(node)}") - out: list[tuple[Mesh, int, str]] = [] - for elt in node.elts: - if not ( - _is_matmul(elt) - and isinstance(elt.right, ast.Call) - and isinstance(elt.right.func, ast.Name) - and elt.right.func.id == "P" - ): - raise VerifyError( - f'value-state entry must be `mesh.axis @ P("reduction")`, got {ast.dump(elt)}' - ) - if len(elt.right.args) != 1: - raise VerifyError( - "value-state P(...) requires exactly one reduction argument, " - 'e.g. `mesh.axis @ P("sum")`' - ) - mesh_name, axis_name = _parse_axis_ref(elt.left) - mesh = mesh_resolver(mesh_name) - if mesh is None: - raise VerifyError(f"undefined mesh {mesh_name!r}") - axis = _resolve_mesh_axis(mesh, axis_name) - reduction = _eval_ast(elt.right.args[0]) - out.append((mesh, axis, reduction)) - return out - - -def _get_dim_nodes(node: ast.AST) -> list[ast.AST]: - """Extract dimension sub-nodes from a layout sugar tuple. - - Accepts both Tuple and BinOp (for standalone sugar like - ``1536 @ (m.w, m.t)`` without a wrapping tuple). - """ - if isinstance(node, ast.Tuple): - if ( - len(node.elts) == 2 - and isinstance(node.elts[0], ast.Tuple) - and isinstance(node.elts[1], ast.Tuple) - ): - return list(node.elts[0].elts) - return list(node.elts) - if _is_constant(node) or _is_matmul(node): - return [node] - raise VerifyError(f"expected tuple layout, got {ast.dump(node)}") - - -_LayoutItem = tuple[ShapeDim | None, Mesh | None, int | None, str, str | None] - - -def _parse_layout_item( - node: ast.AST, - mesh_resolver: MeshResolver, - *, - canonicalize: bool = True, - closure: dict[str, Any] | None = None, -) -> list[_LayoutItem]: - """Parse a single layout-dim element into one or more layout items. - - Returns a list of (dim_size_or_none, mesh, mesh_axis_index, kind, reduction). - The axis-tuple carries only placement; value states (`Partial`) live in the - separate ``{...}`` set parsed by ``_parse_value_state``. - - Forms:: - dim → [(dim, None, None, "broadcast", None)] - dim @ mesh.axis → [(dim, mesh, axis_idx, "split", None)] - dim @ (mesh.axis, ...) → [split items…, bare remainder item] - """ - if _is_constant(node): - return [(_extract_dim(node, closure=closure), None, None, "broadcast", None)] - - if _is_matmul(node): - rhs = node.right - dim = None if _is_placeholder(node.left) else _extract_dim(node.left, closure=closure) - if dim is None: - raise VerifyError( - "layout placeholder `_` is not valid in the axis tuple; " - 'value states go in the `{mesh.axis @ P("reduction")}` set' - ) - if isinstance(rhs, ast.Tuple): - return _expand_multi_axis_sugar(dim, rhs.elts, mesh_resolver) - - if isinstance(rhs, ast.Attribute): - mesh_name, axis_name = _parse_axis_ref(rhs) - mesh = mesh_resolver(mesh_name) - if mesh is None: - raise VerifyError(f"undefined mesh {mesh_name!r}") - axis = _resolve_mesh_axis(mesh, axis_name) - if not canonicalize: - return [(dim, mesh, axis, "split", None)] - return _canonicalize_single_axis(dim, mesh, axis) - - if isinstance(rhs, ast.Name): - mesh = mesh_resolver(rhs.id) - if mesh is None: - raise VerifyError(f"undefined mesh {rhs.id!r}") - mesh_rank = len(mesh.layout.shape) - if mesh_rank != 1: - raise VerifyError( - f"``int @ {rhs.id}`` shorthand requires a single-axis mesh " - f"(found {mesh_rank} axes); write ``{rhs.id}.`` explicitly" - ) - if not canonicalize: - return [(dim, mesh, 0, "split", None)] - return _canonicalize_single_axis(dim, mesh, 0) - - if closure is not None: - try: - dim = _eval_ast(node, closure) - except ValueError: - dim = None - if _is_shape_dim(dim): - return [(dim, None, None, "broadcast", None)] - - raise VerifyError(f"unexpected layout dim AST: {ast.dump(node)}") - - -def _canonicalize_single_axis( - dim: ShapeDim, - mesh: Mesh, - axis: int, -) -> list[_LayoutItem]: - """Factor an oversized single-axis split by its mesh extent. - - Expand ``N @ m.a`` so the split-bound dimension has local size one. ``N`` - must divide evenly by the mesh extent or parsing raises ``ValueError``. - See [parser §1.5](docs/spec/parser.md#15-layout-sugar) and - [shard §7.1.1](docs/spec/shard.md#711-layoutshape). - """ - extent = mesh.layout.shape[axis] - if not isinstance(dim, int) or not isinstance(extent, int): - if dim == extent: - return [(dim, mesh, axis, "split", None)] - raise LayoutSugarError( - f"split layout dim {dim!r} and mesh extent {extent!r} do not have " - "a decidable divisibility relation; bind symbolic dimensions before " - "authoring this split" - ) - if dim % extent != 0: - raise VerifyError( - f"dim {dim} not divisible by mesh extent {extent} on axis " - f"{axis}; cannot canonicalize ``{dim} @ m.``" - ) - if dim == extent: - return [(dim, mesh, axis, "split", None)] - residual = dim // extent - return [ - (extent, mesh, axis, "split", None), - (residual, None, None, "broadcast", None), - ] - - -def _expand_multi_axis_sugar( - dim: ShapeDim, - axis_nodes: list[ast.AST], - mesh_resolver: MeshResolver, -) -> list[_LayoutItem]: - """Expand ``dim @ (mesh.axis, ...)`` into split + remainder items. - - Each mesh axis gets extent = mesh_extent (Split). The remainder - ``dim / ∏(mesh_extents)`` becomes a bare (Broadcast) value axis - appended at the end. - - Raises ``ValueError`` if *dim* is not divisible by the product of - all mesh extents. - """ - items: list[_LayoutItem] = [] - remaining = dim - - for i, ax_node in enumerate(axis_nodes): - mesh, axis = _resolve_axis_node(ax_node, mesh_resolver) - extent = mesh.layout.shape[axis] - if not isinstance(remaining, int) or not isinstance(extent, int): - if remaining == extent and i == len(axis_nodes) - 1: - items.append((extent, mesh, axis, "split", None)) - remaining = 1 - continue - raise LayoutSugarError( - f"split layout dim {dim!r} and mesh extent {extent!r} at axis " - f"position {i} do not have a decidable divisibility relation; " - "bind symbolic dimensions before authoring this split" - ) - if remaining % extent != 0: - raise VerifyError( - f"dim {dim} not divisible by mesh extent {extent} " - f"at axis position {i}; remaining={remaining}" - ) - per_axis = extent - remaining //= extent - items.append((per_axis, mesh, axis, "split", None)) - - if remaining > 0: - items.append((remaining, None, None, "broadcast", None)) - - return items - - -def _resolve_axis_node( - node: ast.AST, - mesh_resolver: MeshResolver, -) -> tuple[Mesh, int]: - """Resolve a mesh-axis reference node to a mesh and layout-axis index. - - Accepts ``mesh.axis`` attribute references and single-axis - ``mesh`` name shorthand. - """ - if isinstance(node, ast.Attribute): - mesh_name, axis_name = _parse_axis_ref(node) - mesh = mesh_resolver(mesh_name) - if mesh is None: - raise VerifyError(f"undefined mesh {mesh_name!r}") - axis = _resolve_mesh_axis(mesh, axis_name) - return (mesh, axis) - if isinstance(node, ast.Name): - mesh = mesh_resolver(node.id) - if mesh is None: - raise VerifyError(f"undefined mesh {node.id!r}") - mesh_rank = len(mesh.layout.shape) - if mesh_rank != 1: - raise VerifyError( - f"``int @ (..., {node.id}, ...)`` shorthand requires a " - f"single-axis mesh (found {mesh_rank} axes); " - f"write ``{node.id}.`` explicitly" - ) - return (mesh, 0) - raise VerifyError(f"expected mesh.axis, got {ast.dump(node)}") - - -def _parse_axis_ref(node: ast.AST) -> tuple[str, str]: - """Parse a mesh-qualified axis reference. - - ``gpu.cluster`` → ``("gpu", "cluster")`` - ``gpu.x`` → ``("gpu", "x")`` - """ - if isinstance(node, ast.Attribute): - mesh_name = _name_of(node.value) - axis_name = node.attr - return (mesh_name, axis_name) - raise VerifyError(f"expected mesh.axis (e.g. gpu.cluster), got {ast.dump(node)}") - - -def _parse_tensor_type_sugar( - node: ast.AST, - closure: dict[str, Any], - *, - mesh_resolver: MeshResolver | None = None, - default_mesh: Mesh | None = None, - mesh_order: "tuple[Mesh, ...]" = (), -) -> TensorType | None: - """Parse a ``Tensor[...]`` or ``ConstTensor[...]`` type literal.""" - if not isinstance(node, ast.Subscript): - return None - head = node.value.id if isinstance(node.value, ast.Name) else None - if isinstance(node.value, ast.Attribute): - head = node.value.attr - if head not in ("Tensor", "ConstTensor"): - return None - if not isinstance(node.slice, ast.Tuple): - raise VerifyError("Tensor[...] requires shape and dtype slots") - elts = node.slice.elts - if len(elts) not in (2, 3, 4): - raise VerifyError("Tensor[...] requires shape, dtype, and optional layout/storage") - - shape, _ = _parse_layout_literal(elts[0], closure=closure) - dtype_val = _resolve_dtype_ast(elts[1], closure) - if dtype_val is None: - raise VerifyError(f"unknown tensor dtype {ast.unparse(elts[1])!r}") - - meshes = {key: value for key, value in closure.items() if isinstance(value, Mesh)} - resolver = mesh_resolver or meshes.get - if default_mesh is None and len(meshes) == 1: - default_mesh = next(iter(meshes.values())) - - layout = None - storage = StorageKind.GMEM - embedded_layout = _has_sugar(elts[0]) - if embedded_layout: - layout = _parse_shard_layout_sugar( - elts[0], - resolver, - default_mesh=default_mesh, - closure=closure, - mesh_order=mesh_order, - ) - - if len(elts) >= 3: - third = elts[2] - if isinstance(third, ast.Constant) and third.value is None: - pass - else: - try: - third_value = _eval_ast(third, closure) - except ValueError: - third_value = None - try: - storage_value = resolve_storage(third_value) - except (TypeError, ValueError): - storage_value = None - if storage_value is not None: - storage = storage_value - else: - if embedded_layout: - raise VerifyError( - "Tensor[...] cannot specify placement in both shape and layout slots" - ) - if isinstance(third_value, (Layout, ShardLayout)): - layout = third_value - else: - layout = _parse_shard_layout_sugar( - third, resolver, default_mesh=default_mesh, closure=closure - ) - if len(elts) == 4: - if embedded_layout: - raise VerifyError( - "Tensor[...] with placement in its shape takes storage as the third slot" - ) - storage = resolve_storage(_eval_ast(elts[3], closure)) - - if ( - isinstance(layout, ShardLayout) - and isinstance(layout.layout, Layout) - and layout.layout.strides is None - ): - layout = ShardLayout( - layout=Layout( - shape=layout.layout.shape, - strides=_auto_strides(layout.layout.shape), - ), - attrs=layout.attrs, - mesh=layout.mesh, - ) - return TensorType(shape=shape, dtype=dtype_val, layout=layout, storage=storage) - - -def parse_sugar( - node: ast.AST, - expected: type, - *, - closure: dict[str, Any] | None = None, - mesh_resolver: MeshResolver | None = None, - default_mesh: Mesh | None = None, - mesh_order: "tuple[Mesh, ...]" = (), -) -> Layout | ShardLayout | TensorType | None: - """Parse one type-directed layout or tensor-type sugar form.""" - closure = closure or {} - if expected is Layout: - return _parse_layout_sugar(node, closure=closure) - if expected is ShardLayout: - if mesh_resolver is None: - meshes = {key: value for key, value in closure.items() if isinstance(value, Mesh)} - mesh_resolver = meshes.get - if default_mesh is None and len(meshes) == 1: - default_mesh = next(iter(meshes.values())) - return _parse_shard_layout_sugar( - node, - mesh_resolver, - default_mesh=default_mesh, - closure=closure, - mesh_order=mesh_order, - ) - if expected is TensorType: - return _parse_tensor_type_sugar( - node, - closure, - mesh_resolver=mesh_resolver, - default_mesh=default_mesh, - mesh_order=mesh_order, - ) - raise TypeError(f"unsupported sugar result type {expected!r}") - - -__all__ = ["LayoutSugarError", "parse_sugar"] diff --git a/src/tilefoundry/parser/symtab.py b/src/tilefoundry/parser/symtab.py deleted file mode 100644 index 2ee47aa1..00000000 --- a/src/tilefoundry/parser/symtab.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from tilefoundry.ir.types.shard.mesh import Mesh - - -class LexicalEnv: - """Stack of dict frames for parser-time name resolution. - - Used by hir parser for Mesh scope ([parser §1.6](docs/spec/parser.md#16-with-mesh-as-m), - parser-only lexical env) and - by tir parser for Var / Mesh tracking. - """ - - def __init__(self) -> None: - self._frames: list[dict[str, Any]] = [{}] - - def push_frame(self) -> None: - self._frames.append({}) - - def pop_frame(self) -> dict[str, Any]: - if len(self._frames) <= 1: - raise RuntimeError("cannot pop root frame") - return self._frames.pop() - - def define(self, name: str, value: Any) -> None: - self._frames[-1][name] = value - - def lookup(self, name: str) -> Any: - for frame in reversed(self._frames): - if name in frame: - return frame[name] - return None - - def innermost_mesh(self): - """Return the innermost Mesh in the lexical scope, or None.""" - for frame in reversed(self._frames): - for val in frame.values(): - if isinstance(val, Mesh): - return val - return None - - -__all__ = ["LexicalEnv"] diff --git a/src/tilefoundry/parser/tir_parser.py b/src/tilefoundry/parser/tir_parser.py deleted file mode 100644 index 82d35c0b..00000000 --- a/src/tilefoundry/parser/tir_parser.py +++ /dev/null @@ -1,566 +0,0 @@ -from __future__ import annotations - -import ast -import dataclasses -from dataclasses import dataclass -from typing import Any, Union - -from tilefoundry.ir.core import Call, Expr, Var, VerifyError -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function as HirFunction -from tilefoundry.ir.tir.launch import LaunchAttrs, launch_call -from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.tir.shape import shape_var_name -from tilefoundry.ir.tir.stmt import Stmt -from tilefoundry.ir.tir.stmts import ( - Evaluate, - For, - If, - LetStmt, - MeshScope, - Return, - Sequential, - While, -) -from tilefoundry.ir.tir.symbol_ref import symbol_call -from tilefoundry.ir.types import DType, TensorType, UnitType -from tilefoundry.ir.types.dim import ( - DimVar, - is_dim_expr, -) -from tilefoundry.ir.types.shard.mesh import Mesh -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.ir.visitor import StmtVisitor -from tilefoundry.target import CudaTarget, default_target -from tilefoundry.utils.spec_ref import spec_ref_render - -from .base import ( - BaseExprVisitor, - _build_params, - _collect_closure, - _i64, - _resolve_tensor_type, - extract_ast, -) -from .dispatch import resolve_callable -from .symtab import LexicalEnv - -_TIR_PRIM = "[tir §1.3](docs/spec/tir.md#13-primfunction)" - - -def _module_target(mod: Module): - """The Target a launch callee's Module declares, if it declares one.""" - try: - return mod.resolve_target() - except ValueError: - return None - - - -@dataclass(frozen=True) -class _Bind: - var: Var - value: Expr - - -_Item = Union[_Bind, Stmt] - - -def _fold_items(items: list[_Item]) -> Sequential: - """Fold ``_Bind`` markers into nested ``LetStmt``; plain Stmts stay.""" - def fold(i: int) -> list[Stmt]: - out: list[Stmt] = [] - while i < len(items): - item = items[i] - if isinstance(item, _Bind): - inner = fold(i + 1) - out.append( - LetStmt( - var=item.var, - value=item.value, - body=Sequential(body=tuple(inner)), - ) - ) - return out - out.append(item) - i += 1 - return out - - return Sequential(body=tuple(fold(0))) - - -def _is_device_target(target) -> bool: - """Whether *target* is a CUDA device Target, including subclasses.""" - return isinstance(target, CudaTarget) - - -def _dim_var_names_in_type(ty: Any) -> set[str]: - """DimVar names appearing in a ``TensorType``'s shape. - - DimVar names appearing in a ``TensorType``'s shape (and the shape of its - ``ShardLayout``'s inner layout). - """ - names: set[str] = set() - for d in getattr(ty, "shape", None) or (): - if isinstance(d, DimVar): - names.add(d.name) - inner = getattr(getattr(ty, "layout", None), "layout", None) - for d in getattr(inner, "shape", None) or (): - if isinstance(d, DimVar): - names.add(d.name) - return names - - -class _DimVarRefCollector(StmtVisitor): - """Represent DimVarRefCollector. - - Collect DimVar names the body actually references via the types of bound - results and op operands. A dynamic tensor dim surfaces in these - ``TensorType`` shapes; the codegen plumbs each as a ``_shape_`` - runtime scalar, so the kernel signature must declare it. - """ - - def __init__(self) -> None: - self.names: set[str] = set() - - def visit_LetStmt(self, s: LetStmt) -> None: - self.names |= _dim_var_names_in_type(getattr(s.var, "type", None)) - self.generic_visit(s) - - def visit_Evaluate(self, s: Evaluate) -> None: - for a in s.args: - self.names |= _dim_var_names_in_type(getattr(a, "type", None)) - self.generic_visit(s) - - -def _shape_scalar_params( - params: tuple[Var, ...], referenced: set[str] -) -> tuple[Var, ...]: - """Shape scalar params. - - Hidden ``_shape_`` i32 scalar params for each referenced - dynamic dim, mirroring the HIR→TIR lowering. Each DimVar maps to its first - occurrence in a tensor param's shape (the same rule codegen uses to source - the runtime extent). Idempotent: a param that already exists is skipped. - """ - existing = {p.name for p in params} - scalar_i32 = TensorType.scalar(dtype=DType.i32, storage=StorageKind.RMEM) - seen: set[str] = set() - extra: list[Var] = [] - for p in params: - ty = p.type - if not isinstance(ty, TensorType): - continue - for axis, dim in enumerate(ty.shape): - if not isinstance(dim, DimVar) or dim.name not in referenced: - continue - if dim.name in seen: - continue - seen.add(dim.name) - name = shape_var_name(p.name, axis) - if name not in existing: - extra.append(Var(type=scalar_i32, name=name)) - return tuple(extra) - - -def _scope_str(mesh) -> str: - """A thread scope as its topologies and the lane layout they are viewed through.""" - topologies = ", ".join(f"{t.name}({t.size})" for t in mesh.topologies) - layout = mesh.layout - return f"{topologies} viewed as shape {tuple(layout.shape)} strides {tuple(layout.strides)}" - - -def _scope_mismatch_str(current, required) -> str: - """Say which part of the required thread scope the enclosing one fails to provide. - - A fragment's scope can be wrong by topology, by lane count, or by the exact - thread-value decomposition, and reporting the last of those for all three - sends the reader looking at the layout when the topology is what differs. - """ - head = ( - f"enclosing mesh scope [{_scope_str(current)}] does not provide the atom's " - f"required thread scope [{_scope_str(required)}]" - ) - if current.topologies[0].name != required.topologies[0].name: - return ( - f"{head} — the scope is a {current.topologies[0].name} scope and the atom needs " - f"a {required.topologies[0].name} one" - ) - if current.topologies[0].size != required.topologies[0].size: - return ( - f"{head} — the scope has {current.topologies[0].size} lanes and the atom needs " - f"{required.topologies[0].size}" - ) - return ( - f"{head} — the lane count matches, so what differs is the thread-value " - "decomposition, which must match shape and strides exactly" - ) - - -def parse_prim_func(fn, *, target=None, extra_closure=None) -> PrimFunction: - node = extract_ast(fn) - closure = _collect_closure(fn, extra_closure) - env = LexicalEnv() - params = _build_params( - node, closure, _resolve_tensor_type, decorator_name="@tilefoundry.prim_func" - ) - for p in params: - env.define(p.name, p) - visitor = _TirBodyVisitor(env, closure) - body = _fold_items(visitor.visit_body(node.body)) - - - - - if _is_device_target(target): - collector = _DimVarRefCollector() - collector.visit(body) - params = (*params, *_shape_scalar_params(params, collector.names)) - kwargs = {} if target is None else {"target": target} - return PrimFunction(name=node.name, params=params, body=body, **kwargs) - - -def _is_none(node: ast.AST) -> bool: - """True for a literal ``None`` AST node.""" - return isinstance(node, ast.Constant) and node.value is None - - - -_LAUNCH_CONFIG_KEYS = frozenset( - {"grid", "block", "cluster", "dynamic_smem", "stream", "attrs"} -) - - -class _TirBodyVisitor(BaseExprVisitor): - token = "tir" - - def visit_body(self, stmts: list[ast.stmt]) -> list[_Item]: - out: list[_Item] = [] - for node in stmts: - s = self._visit_stmt(node) - if s is not None: - out.append(s) - return out - - def _visit_stmt(self, node: ast.stmt) -> _Item | None: - if isinstance(node, ast.Assign): - if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): - raise VerifyError("tir: only single-target Name assignments supported in V1") - tgt = node.targets[0].id - - - - - if self._is_platform_rooted(node.value): - self.env.define(tgt, self._eval_static(node.value)) - return None - rhs = self.expr(node.value) - var = Var(type=rhs.type, name=tgt) - self.env.define(tgt, var) - - - return _Bind(var=var, value=rhs) - - if isinstance(node, ast.Expr): - inner = node.value - if not isinstance(inner, ast.Call): - raise VerifyError("tir: bare expression statement must be a call") - return self._call_as_top_level_stmt(inner) - - if isinstance(node, ast.Return): - if node.value is not None: - raise VerifyError( - f"{spec_ref_render(_TIR_PRIM)}: @tilefoundry.prim_func return must be bare " - f"(no value)" - ) - return Return() - - if isinstance(node, ast.For): - return self._visit_for(node) - - if isinstance(node, ast.While): - cond = self.expr(node.test) - body = _fold_items(self.visit_body(node.body)) - return While(cond=cond, body=body) - - if isinstance(node, ast.If): - cond = self.expr(node.test) - then_body = _fold_items(self.visit_body(node.body)) - else_body = _fold_items(self.visit_body(node.orelse)) if node.orelse else Sequential(body=()) - return If(cond=cond, then_body=then_body, else_body=else_body) - - if isinstance(node, ast.With): - return self._visit_with(node) - - raise VerifyError(f"tir: unsupported statement {type(node).__name__}") - - def _call_as_top_level_stmt(self, node: ast.Call) -> Stmt: - """Resolve a TIR top-level call as a special form, statement, or effect. - - Name-only forms include launch, registered statements, and sibling - primitive functions. Other calls must produce unit type to become - ``Evaluate``. See [parser §1.3](docs/spec/parser.md#13-op-call) and - [parser §4.6](docs/spec/parser.md#46-per-dialect-strict-resolution). - """ - if isinstance(node.func, ast.Name): - name = node.func.id - if name == "launch": - return self._launch_as_stmt(node) - - - - - - - - try: - _kind, stmt_cls = resolve_callable(name, "tir") - except VerifyError: - stmt_cls = None - if stmt_cls is not None and getattr(stmt_cls, "_op_schema", None) is None: - kwargs = {k.arg: self._eval_static(k.value) for k in node.keywords} - pos = [self.expr(a) for a in node.args] - field_names = [f.name for f in dataclasses.fields(stmt_cls) if f.name != "loc"] - if len(pos) > len(field_names): - raise VerifyError(f"tir: {name!r} too many positional args") - bound = dict(zip(field_names, pos)) - bound.update(kwargs) - return stmt_cls(**bound) - - - callee_ir = self.closure.get(name) - if isinstance(callee_ir, PrimFunction): - if node.keywords: - raise VerifyError(f"tir: call to {name!r} does not support kwargs") - args = tuple(self.expr(a) for a in node.args) - return symbol_call(callee_ir, args) - - - - - expr = self.call_to_op_call(node) - if isinstance(expr, Call) and isinstance(expr.type, UnitType): - return Evaluate(callable=expr.target, args=expr.args) - disp = node.func.id if isinstance(node.func, ast.Name) else ast.unparse(node.func) - raise VerifyError( - f"tir: value op {disp!r} cannot be top-level Stmt; wrap with `=`" - ) - - def _is_platform_rooted(self, node: ast.AST) -> bool: - """True when ``node`` is a ``T....`` expression. - - True when ``node`` is a ``T....`` expression — an - attribute/call chain whose root Name resolves to the ``dsl.T`` module - and whose first attribute is a platform name (``cuda``, later other - targets). Such expressions are compile-time descriptors - (``MmaOpSpec`` / ``MmaAtom``), bound statically rather than lowered to - a ``LetStmt``. - """ - cur = node - first_attr_on_root: str | None = None - while isinstance(cur, (ast.Attribute, ast.Call)): - if isinstance(cur, ast.Call): - cur = cur.func - else: - first_attr_on_root = cur.attr - cur = cur.value - if not isinstance(cur, ast.Name): - return False - root = self.env.lookup(cur.id) - if root is None: - root = self.closure.get(cur.id) - import tilefoundry.dsl as _dsl # noqa: PLC0415 - if root is not _dsl.T: - return False - from tilefoundry.dsl.T._platforms import PLATFORM_NAMESPACES # noqa: PLC0415 - return first_attr_on_root in PLATFORM_NAMESPACES - - def _launch_as_stmt(self, node: ast.Call) -> Evaluate: - - - if not node.args: - raise VerifyError("tir: launch(...) needs a device function first argument") - callee_node = node.args[0] - if not isinstance(callee_node, ast.Name): - raise VerifyError("tir: launch(...) first argument must be a function name") - callee_ir = self.closure.get(callee_node.id) - - - - callee_target = None - if isinstance(callee_ir, Module): - callee_target = _module_target(callee_ir) - callee_ir = callee_ir.entry_function() - elif isinstance(callee_ir, PrimFunction): - callee_target = callee_ir.target - if not isinstance(callee_ir, (HirFunction, PrimFunction)): - raise VerifyError( - f"tir: launch(...) callee {callee_node.id!r} must be a @func or " - f"@prim_func device function" - ) - effective_target = ( - callee_target if callee_target is not None else default_target() - ) - if not isinstance(effective_target, CudaTarget): - raise VerifyError( - f"tir: launch(...) callee {callee_node.id!r} must target a CUDA device" - ) - - args = tuple(self.expr(a) for a in node.args[1:]) - kw: dict[str, ast.AST] = {} - for k in node.keywords: - if k.arg is None: - raise VerifyError("tir: launch(...) does not accept `**kwargs`") - if k.arg not in _LAUNCH_CONFIG_KEYS: - raise VerifyError( - f"tir: launch(...) got unexpected keyword {k.arg!r}; " - f"allowed: {sorted(_LAUNCH_CONFIG_KEYS)}" - ) - kw[k.arg] = k.value - if "grid" not in kw or "block" not in kw: - raise VerifyError("tir: launch(...) requires `grid=` and `block=`") - cluster = ( - self._launch_dim(kw["cluster"]) - if "cluster" in kw and not _is_none(kw["cluster"]) - else None - ) - dynamic_smem = self.expr(kw["dynamic_smem"]) if "dynamic_smem" in kw else 0 - stream = ( - self.expr(kw["stream"]) - if "stream" in kw and not _is_none(kw["stream"]) - else None - ) - if "attrs" in kw: - attrs = self._eval_static(kw["attrs"]) - if not isinstance(attrs, LaunchAttrs): - raise VerifyError( - f"tir: launch(...) `attrs=` must be a LaunchAttrs, got " - f"{type(attrs).__name__}" - ) - else: - attrs = LaunchAttrs() - return launch_call( - callee_ir, - args, - self._launch_dim(kw["grid"]), - self._launch_dim(kw["block"]), - cluster=cluster, - dynamic_smem=dynamic_smem, - stream=stream, - attrs=attrs, - ) - - def _launch_dim(self, value: ast.AST) -> tuple[Expr, Expr, Expr]: - """Normalize a launch grid / block spec to a 3-tuple of extents. - - Grid / block extents are compile-time shape arithmetic, not runtime - expressions, so each element is evaluated statically (a literal, a - ``DimVar``, or a dim-arithmetic ``Expr`` such as ``ceildiv(S, tile)``). - A scalar or 1-/2-tuple is right-padded with the constant ``1``. - """ - nodes = value.elts if isinstance(value, ast.Tuple) else [value] - elts = [self._eval_launch_extent(n) for n in nodes] - if len(elts) > 3: - raise VerifyError("tir: launch grid/block accepts at most 3 dimensions") - while len(elts) < 3: - elts.append(_i64(1)) - return tuple(elts) - - def _eval_launch_extent(self, node: ast.AST): - """Eval launch extent. - - Statically evaluate one grid / block extent to an ``int`` (wrapped as - a constant), a ``DimVar``, or a dim-arithmetic ``Expr``; reject anything - else loudly (extents are shape arithmetic, not arbitrary values). - """ - val = self._eval_static(node) - if not is_dim_expr(val): - raise VerifyError( - f"tir: launch grid/block extent must be an int, DimVar, or dim " - f"expression, got {type(val).__name__}" - ) - return _i64(val) if isinstance(val, int) else val - - def _visit_for(self, node: ast.For) -> For: - if not isinstance(node.target, ast.Name): - raise VerifyError("tir: For target must be a Name") - if not isinstance(node.iter, ast.Call) or not isinstance(node.iter.func, ast.Name): - raise VerifyError("tir: For iter must be a plain `range(...)` call") - if node.iter.func.id != "range": - raise VerifyError(f"tir: For iter must be `range(...)`, got {node.iter.func.id!r}") - args = node.iter.args - if len(args) == 1: - start, stop, step = _i64(0), self.expr(args[0]), _i64(1) - elif len(args) == 2: - start, stop, step = self.expr(args[0]), self.expr(args[1]), _i64(1) - elif len(args) == 3: - start, stop, step = self.expr(args[0]), self.expr(args[1]), self.expr(args[2]) - else: - raise VerifyError("tir: range() expects 1-3 args") - iv = Var( - type=TensorType.scalar(DType.i64, storage=StorageKind.RMEM), - name=node.target.id, - ) - self.env.push_frame() - try: - self.env.define(node.target.id, iv) - body = _fold_items(self.visit_body(node.body)) - finally: - self.env.pop_frame() - return For(induction_var=iv, start=start, stop=stop, step=step, body=body) - - def _visit_with(self, node: ast.With) -> MeshScope: - if len(node.items) != 1: - raise VerifyError("tir: only single-item `with` supported") - item = node.items[0] - mesh = self._eval_static(item.context_expr) - if not isinstance(mesh, Mesh): - raise VerifyError( - f"tir: `with` context must evaluate to a Mesh, got {type(mesh).__name__}" - ) - if item.optional_vars is None or not isinstance(item.optional_vars, ast.Name): - raise VerifyError("tir: `with Mesh(...) as name` requires a single Name binding") - binding_name = item.optional_vars.id - binding = Var( - type=TensorType.scalar(DType.i64, storage=StorageKind.RMEM), - name=binding_name, - ) - self.env.push_frame() - try: - self.env.define(binding_name, mesh) - body = _fold_items(self.visit_body(node.body)) - finally: - self.env.pop_frame() - return MeshScope(mesh=mesh, binding=binding, body=body) - - def _resolve_static_attribute(self, owner, attr: str): - """TIR static attribute resolution. - - An MMA fragment access ``atom.A/B/C`` returns the atom's layout - contract **as-is** (no rebind). But because a fragment is only valid in - a thread scope that can host the atom, we check the enclosing mesh - scope against ``atom.required_scope`` here, at the use point — rejecting - e.g. a ``cta`` or wrong-sized ``thread`` scope. The match is structural - (thread participation), independent of binding/axis names. - """ - from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom # noqa: PLC0415 - from tilefoundry.ir.types.shard.scope_match import ( # noqa: PLC0415 - mesh_scope_matches_required_scope, - ) - - val = getattr(owner, attr) - if isinstance(owner, MmaAtom) and attr in ("A", "B", "C"): - mesh = self._current_default_mesh() - if mesh is None: - raise VerifyError( - f"mma fragment `atom.{attr}` must be used inside a " - f"`with Mesh(...)` thread scope" - ) - if not mesh_scope_matches_required_scope(mesh, owner.required_scope): - raise VerifyError( - f"mma fragment `atom.{attr}`: {_scope_mismatch_str(mesh, owner.required_scope)}" - ) - return val - - -__all__ = ["parse_prim_func"] diff --git a/src/tilefoundry/schedule/partition/program.py b/src/tilefoundry/schedule/partition/program.py index fc0acd31..bacfd8c3 100644 --- a/src/tilefoundry/schedule/partition/program.py +++ b/src/tilefoundry/schedule/partition/program.py @@ -133,13 +133,20 @@ def visit_Call(self, expr: Call) -> tuple[int, ...]: self.owner._expr_values[key] = refs self.owner._record_requirement(refs, expr) return refs + output_types = tensor_leaves(expr.type) + if output_types and all( + type.storage is StorageKind.UMAT for _path, type in output_types + ): + self.owner._expr_values[key] = () + self.owner._record_requirement((), expr) + return () site_id = self.owner._next_site self.owner._next_site += 1 output_refs = tuple( self.owner._new_value( expr, type, path, self.function_path, producer_site_id=site_id ) - for path, type in tensor_leaves(expr.type) + for path, type in output_types ) self.owner.sites.append( OperationSite(site_id, expr, self.function_path, arg_refs, output_refs) diff --git a/src/tilefoundry/script.py b/src/tilefoundry/script.py index d9e49985..8ec958f4 100644 --- a/src/tilefoundry/script.py +++ b/src/tilefoundry/script.py @@ -1,6 +1,6 @@ """Define parser-backed ``@func`` and ``@prim_func`` decorators. -The surface follows [parser §1](docs/spec/parser.md#1-dsl-syntax). A decorator +The surface follows [parser §2](docs/spec/parser.md#2-syntax-and-rules). A decorator returns the parsed and verified IR node, not the original Python function. """ @@ -9,7 +9,8 @@ import sys from dataclasses import dataclass from enum import StrEnum -from typing import Any, Callable, ClassVar, Literal +from types import FunctionType +from typing import Any, Callable, ClassVar, Literal, Mapping from tilefoundry.ir.core.module import Module from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern @@ -18,9 +19,9 @@ from tilefoundry.ir.hir.verify import verify_function from tilefoundry.ir.tir.intrinsic import intrinsic as _intrinsic from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.module import UNDECLARED, _Entry, enclosing_declaration -from tilefoundry.parser import parse_prim_func -from tilefoundry.parser.hir_parser import _parse_func +from tilefoundry.module import UNDECLARED, _Entry +from tilefoundry.parser import FuncParserContext, FunctionRole, parse_function +from tilefoundry.parser.ast_pattern import LexicalScope, module_context_for_frame from tilefoundry.target.base import target_instance @@ -147,14 +148,8 @@ def _register( *, base: HirFunction | None = None, ) -> None: - """Validate and record a parser-time Function in its enclosing Module.""" - entry = _enclosing_declaration() - ParsedFuncRules.check(kind, ir, binding_name, key, entry=entry, base=base) - if entry is None: - return - entry.bound[binding_name] = kind - if kind is not ParsedFuncKind.KERNEL: - entry.owned.add(id(ir)) + """Compatibility no-op; FunctionPattern owns validation and registration.""" + return None def _validate_one_pattern(pattern: Any) -> Pattern: @@ -214,7 +209,7 @@ def _enclosing_declaration(): here = __file__ while frame is not None and frame.f_code.co_filename == here: frame = frame.f_back - return enclosing_declaration(frame) + return module_context_for_frame(frame) def _enclosing_topologies() -> tuple | None: @@ -223,6 +218,149 @@ def _enclosing_topologies() -> tuple | None: return entry.topologies if entry is not None else None +def _parse_authored( + fn_inner: FunctionType, + *, + dialect: Literal["hir", "tir"], + role: FunctionRole, + binding_name: str, + base: HirFunction | None = None, + key: object | None = None, + target: object | None = None, + topologies: tuple | None = None, + module_context=None, + closure: Mapping[str, Any] | None = None, +): + """Build one typed context and route every authored function through the API.""" + module_context = module_context if module_context is not None else _enclosing_declaration() + closure = dict(closure) if closure is not None else _definition_namespace() + if fn_inner.__globals__ is not None: + closure.update(fn_inner.__globals__) + if fn_inner.__closure__ is not None: + for name, cell in zip(fn_inner.__code__.co_freevars, fn_inner.__closure__): + try: + closure[name] = cell.cell_contents + except ValueError: + pass + use_owner_context = module_context is not None and ( + role is not FunctionRole.ROOT or (target is None and topologies is None) + ) + if use_owner_context: + context = module_context.function_context( + dialect=dialect, + role=role, + binding_name=binding_name, + closure=closure, + base=base, + key=key, + ) + else: + topology_scope = { + getattr(topology, "name", str(index)): topology + for index, topology in enumerate(topologies or ()) + } + context = FuncParserContext( + dialect=dialect, + role=role, + closure=closure, + topologies=topology_scope, + module_scope=LexicalScope(), + base=base, + key=key, + target=target if dialect == "tir" else None, + binding_name=binding_name, + ) + return parse_function(fn_inner, context) + + +def _capture_function_closure(fn_inner: FunctionType) -> dict[str, Any]: + closure = _definition_namespace() + closure.update(fn_inner.__globals__) + if fn_inner.__closure__ is not None: + for name, cell in zip(fn_inner.__code__.co_freevars, fn_inner.__closure__): + try: + closure[name] = cell.cell_contents + except ValueError: + pass + return closure + + +@dataclass +class _DeferredFunction: + module_context: Any + fn_inner: FunctionType + dialect: Literal["hir", "tir"] + role: FunctionRole + binding_name: str + closure: Mapping[str, Any] + base: object | None = None + key: object | None = None + parsed: object | None = None + _tilefoundry_deferred: bool = True + + def parse(self): + base = self.base.parsed if isinstance(self.base, _DeferredFunction) else self.base + if self.role is FunctionRole.CONVERTER: + _validate_converter_weight_name(base, self.key) + self.parsed = _parse_authored( + self.fn_inner, + dialect=self.dialect, + role=self.role, + binding_name=self.binding_name, + base=base, + key=self.key, + module_context=self.module_context, + closure=self.closure, + ) + if self.dialect == "hir": + if self.role is FunctionRole.VARIANT: + object.__setattr__(self.parsed, DISPLAY_NAME, self.binding_name) + object.__setattr__(self.parsed, "name", base.name) + elif self.role is FunctionRole.CONVERTER: + object.__setattr__( + self.parsed, + "name", + f"{base.name}.converter[{self.key}]", + ) + return self.parsed + + def specialize(self, pattern: Any): + pat = _validate_one_pattern(pattern) + + def _wrap_variant(fn_inner): + declaration = _DeferredFunction( + self.module_context, + fn_inner, + "hir", + FunctionRole.VARIANT, + fn_inner.__name__, + _capture_function_closure(fn_inner), + base=self, + key=pat, + ) + self.module_context.declarations.append(declaration) + return declaration + + return _wrap_variant + + def converter(self, weight_name: str): + def _wrap_converter(fn_inner): + declaration = _DeferredFunction( + self.module_context, + fn_inner, + "hir", + FunctionRole.CONVERTER, + fn_inner.__name__, + _capture_function_closure(fn_inner), + base=self, + key=weight_name, + ) + self.module_context.declarations.append(declaration) + return declaration + + return _wrap_converter + + def func(fn=None, *, topologies=UNDECLARED, target=None): """Decorator: parse an ``@func``-decorated function into HIR. @@ -238,19 +376,25 @@ def func(fn=None, *, topologies=UNDECLARED, target=None): declared_topologies = None if topologies is UNDECLARED else tuple(topologies) def _wrap(fn_inner): - ParsedFuncRules.NAMING[ParsedFuncKind.KERNEL].check( - fn_inner.__name__, - _enclosing_declaration(), - kind=ParsedFuncKind.KERNEL, - ) - extra_closure = _definition_namespace() - parse_topologies = declared_topologies - if parse_topologies is None: - parse_topologies = _enclosing_topologies() - ir = _parse_func( - fn_inner, topologies=parse_topologies or (), - extra_closure=extra_closure, - in_module_body=_enclosing_declaration() is not None, + module_context = _enclosing_declaration() + if module_context is not None and not declares_context: + declaration = _DeferredFunction( + module_context, + fn_inner, + "hir", + FunctionRole.ROOT, + fn_inner.__name__, + _capture_function_closure(fn_inner), + ) + module_context.declarations.append(declaration) + return declaration + ir = _parse_authored( + fn_inner, + dialect="hir", + role=FunctionRole.ROOT, + binding_name=fn_inner.__name__, + target=resolved_target, + topologies=declared_topologies, ) verify_function(ir) _register(ParsedFuncKind.KERNEL, ir, fn_inner.__name__, None) @@ -280,17 +424,14 @@ def _specialize(self: HirFunction, pattern: Any): pat = _validate_one_pattern(pattern) def _wrap_variant(fn_inner): - ParsedFuncRules.NAMING[ParsedFuncKind.VARIANT].check( - fn_inner.__name__, - _enclosing_declaration(), - kind=ParsedFuncKind.VARIANT, - base_name=self.name, - ) - extra_closure = _definition_namespace() - ir = _parse_func( - fn_inner, topologies=_enclosing_topologies() or (), - specializations=(pat,), extra_closure=extra_closure, - in_module_body=_enclosing_declaration() is not None, + ir = _parse_authored( + fn_inner, + dialect="hir", + role=FunctionRole.VARIANT, + binding_name=fn_inner.__name__, + base=self, + key=pat, + topologies=_enclosing_topologies(), ) if ir.body is None: raise TypeError( @@ -301,14 +442,7 @@ def _wrap_variant(fn_inner): object.__setattr__(ir, DISPLAY_NAME, fn_inner.__name__) object.__setattr__(ir, "name", self.name) verify_function(ir) - _register( - ParsedFuncKind.VARIANT, - ir, - fn_inner.__name__, - pat, - base=self, - ) - self.add_variant(ir) + _register(ParsedFuncKind.VARIANT, ir, fn_inner.__name__, pat, base=self) return ir return _wrap_variant @@ -327,10 +461,14 @@ def _converter(self: HirFunction, weight_name: str): _validate_converter_weight_name(self, weight_name) def _wrap_converter(fn_inner): - extra_closure = _definition_namespace() - ir = _parse_func( - fn_inner, extra_closure=extra_closure, - in_module_body=_enclosing_declaration() is not None, + ir = _parse_authored( + fn_inner, + dialect="hir", + role=FunctionRole.CONVERTER, + binding_name=fn_inner.__name__, + base=self, + key=weight_name, + topologies=_enclosing_topologies(), ) if ir.body is None: raise TypeError( @@ -340,14 +478,7 @@ def _wrap_converter(fn_inner): object.__setattr__(ir, "name", f"{self.name}.converter[{weight_name}]") verify_function(ir) - _register( - ParsedFuncKind.CONVERTER, - ir, - fn_inner.__name__, - weight_name, - base=self, - ) - self.add_converter(weight_name, ir) + _register(ParsedFuncKind.CONVERTER, ir, fn_inner.__name__, weight_name, base=self) return ir return _wrap_converter @@ -368,8 +499,25 @@ def prim_func(fn=None, *, target=None): resolved_target = target def _wrap(fn_inner): - extra_closure = _definition_namespace() - ir = parse_prim_func(fn_inner, target=resolved_target, extra_closure=extra_closure) + module_context = _enclosing_declaration() + if module_context is not None and resolved_target is None: + declaration = _DeferredFunction( + module_context, + fn_inner, + "tir", + FunctionRole.ROOT, + fn_inner.__name__, + _capture_function_closure(fn_inner), + ) + module_context.declarations.append(declaration) + return declaration + ir = _parse_authored( + fn_inner, + dialect="tir", + role=FunctionRole.ROOT, + binding_name=fn_inner.__name__, + target=resolved_target, + ) verify_prim_function(ir) return ir diff --git a/src/tilefoundry/visitor_registry/contexts.py b/src/tilefoundry/visitor_registry/contexts.py index 644d7ad9..05e53e31 100644 --- a/src/tilefoundry/visitor_registry/contexts.py +++ b/src/tilefoundry/visitor_registry/contexts.py @@ -12,7 +12,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any, NoReturn, Union +from typing import Any, Generic, NoReturn, Protocol, TypeVar, Union from tilefoundry.ir.core.errors import VerifyError from tilefoundry.ir.core.expr import Call, Expr @@ -27,6 +27,29 @@ from tilefoundry.ir.types.tensor_type import DType, TensorType, Type from tilefoundry.ir.types.utils import local_type_of +T = TypeVar("T") + + +@dataclass(frozen=True) +class CallFeed(Generic[T]): + """Values supplied to one callee, keyed by formal parameter identity.""" + + by_param: Mapping[int, T] + + def value_for(self, param: object) -> T: + try: + return self.by_param[id(param)] + except KeyError: + raise KeyError(f"no call-feed value for parameter {getattr(param, 'name', param)!r}") from None + + +class CallFeedProvider(Protocol[T]): + """Context-owned call binding and callee-scope construction.""" + + def build_call_feed(self, callee: object, supplied: tuple[T, ...]) -> CallFeed[T]: ... + + def scope_for(self, callee: object) -> FunctionScope | None: ... + def _constant_type(value: object) -> TensorType: if isinstance(value, bool): @@ -70,6 +93,54 @@ class TypeInferContext: cache: dict[int, Type] = field(default_factory=dict) mesh_scope: tuple = () elaboration_cache: dict[tuple, Any] = field(default_factory=dict) + call_feed_provider: CallFeedProvider[Type] | None = None + feed: CallFeed[Type] | None = None + + def build_call_feed(self, callee: object, supplied: tuple[Type, ...]) -> CallFeed[Type]: + """Build the type values visible to *callee* in this walk.""" + if self.call_feed_provider is not None: + return self.call_feed_provider.build_call_feed(callee, supplied) + + child = self._child_module(callee) + params = tuple(p for p in callee.params if not (child is not None and p.is_const)) + if len(supplied) != len(params): + kind = "activation(s)" if child is not None else "parameter(s)" + raise VerifyError( + f"hir Function call {callee.name!r}: arity mismatch — " + f"callee declares {len(params)} {kind}, call passed {len(supplied)}" + ) + values = iter(supplied) + return CallFeed( + { + id(param): param.type if child is not None and param.is_const else next(values) + for param in callee.params + } + ) + + def scope_for(self, callee: object) -> FunctionScope | None: + """Return the runtime scope in which *callee*'s body is read.""" + if self.call_feed_provider is not None: + return self.call_feed_provider.scope_for(callee) + if self.scope is None: + return None + child = self._child_module(callee) + return FunctionScope(child, callee) if child is not None else FunctionScope(self.scope.module, callee) + + def child(self, callee: object, feed: CallFeed[Type]) -> "TypeInferContext": + """Create the recursive context while retaining this provider.""" + return TypeInferContext( + scope=self.scope_for(callee), + mesh_scope=self.mesh_scope, + call_feed_provider=self.call_feed_provider, + feed=feed, + ) + + def _child_module(self, callee: object): + if self.scope is None or self.scope.module is None: + return None + from tilefoundry.ir.core.module import child_module_of # noqa: PLC0415 + + return child_module_of(self.scope.module, self.scope.function, callee) def type_of(self, expr: Expr) -> Type: key = id(expr) @@ -214,6 +285,8 @@ def __post_init__(self) -> None: __all__ = [ + "CallFeed", + "CallFeedProvider", "FunctionScope", "TypeInferContext", "VerifyContext", diff --git a/tests/dsl/test_module_decorator.py b/tests/dsl/test_module_decorator.py index 5e38de25..69205445 100644 --- a/tests/dsl/test_module_decorator.py +++ b/tests/dsl/test_module_decorator.py @@ -57,7 +57,7 @@ def test_attribute_access_ambiguous_name_and_real_fields(): A duplicated function name is ambiguous under attribute access (raises), real Module fields are never intercepted, and ``function_named`` returns all - matches — the core-ir [parser §2.1](docs/spec/parser.md#21-model) ambiguity rule. + matches — the core-ir [parser §2](docs/spec/parser.md#2-syntax-and-rules) ambiguity rule. """ base = _Demo.lookup("leaf") dup_a = dataclasses.replace(base, name="dup") diff --git a/tests/inspection/test_roundtrip.py b/tests/inspection/test_roundtrip.py index 7ef41823..92e75fe6 100644 --- a/tests/inspection/test_roundtrip.py +++ b/tests/inspection/test_roundtrip.py @@ -7,9 +7,7 @@ """ from tests._source import import_dsl -from tests.fixtures.placed.gqa_decode import GqaOnline from tilefoundry.inspection import as_script -from tilefoundry.ir.hir.specialize import specialize_concretely from tilefoundry.ir.types import DType _HEADER = ( @@ -243,25 +241,3 @@ def test_carry_updates_print_last_without_shadowing_the_old_value() -> None: ) assert repr(fn.body) == repr(rebuilt.body) assert as_script(rebuilt) == printed - - -def test_gqa_correction_reads_the_old_carry_and_unique_yield() -> None: - function = specialize_concretely(GqaOnline.entry_function(), {"ctx_len": 8}) - printed = as_script(function) - - assert printed.count(" m_new = max(m, score_2)") == 1 - assert " = sub(m, m_new)" in printed - lines = printed.splitlines() - start = lines.index(" for i in range(8):") - end = next( - index - for index in range(start + 1, len(lines)) - if lines[index].startswith(" ") - and not lines[index].startswith(" ") - ) - assert lines[end - 3 : end] == [ - " l = l_3", - " o = o_4", - " m = m_new", - ] - assert lines[end] == ' k_n = cast(k_new, dtype="f32")' diff --git a/tests/ir/core/test_overload.py b/tests/ir/core/test_overload.py index d5735a1d..183400c5 100644 --- a/tests/ir/core/test_overload.py +++ b/tests/ir/core/test_overload.py @@ -1,4 +1,4 @@ -"""``parser.overload`` — F3 first-match contract.""" +"""``ir.core.overload`` — F3 first-match contract.""" from __future__ import annotations @@ -8,9 +8,9 @@ import pytest from tilefoundry.ir.core.op_schema import OpSchema +from tilefoundry.ir.core.overload import OverloadError, filter_candidates, resolve from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.core.pattern import Scalar, Tensor, TensorPat -from tilefoundry.parser.overload import OverloadError, filter_candidates, resolve @dataclass(frozen=True) diff --git a/tests/ir/types/test_mesh.py b/tests/ir/types/test_mesh.py index 9498e076..6a031b1e 100644 --- a/tests/ir/types/test_mesh.py +++ b/tests/ir/types/test_mesh.py @@ -1,14 +1,11 @@ from __future__ import annotations -import ast - import pytest from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard import ( Layout, Mesh, - Partial, Split, Topology, make_mesh, @@ -20,7 +17,6 @@ states_consistent_positions, ) from tilefoundry.ir.types.shard.shard_layout import ShardLayout -from tilefoundry.parser.sugar import parse_sugar from tilefoundry.schedule.partition.problem import _placement_relation @@ -71,26 +67,6 @@ def test_mesh_slice_keeps_the_parent_topologies() -> None: assert sliced.layout.shape == (1, 32) -def test_named_mesh_axis_sugar_carries_a_layout_index() -> None: - mesh = make_mesh((8,), names=("cta",), topology="cta") - node = ast.parse("(8 @ cta.cta,)", mode="eval").body - - layout = parse_sugar( - node, - ShardLayout, - mesh_resolver=lambda name: mesh if name == "cta" else None, - ) - partial_node = ast.parse('((8,), {cta.cta @ P("sum")})', mode="eval").body - partial_layout = parse_sugar( - partial_node, - ShardLayout, - mesh_resolver=lambda name: mesh if name == "cta" else None, - ) - - assert layout.attrs == (Split(0),) - assert partial_layout.attrs == (Partial("sum"),) - - def test_mesh_value_equality_is_usable_by_partition() -> None: left = make_mesh((8,), topology="thread") right = make_mesh((8,), topology="thread") diff --git a/tests/models/gemma2_2b/model.py b/tests/models/gemma2_2b/model.py index 3d43b84b..eec3c4f5 100644 --- a/tests/models/gemma2_2b/model.py +++ b/tests/models/gemma2_2b/model.py @@ -225,14 +225,15 @@ def self_attention( # normalisation instead. Written out per group because the two are # differently shaped and a @func binds its parameter shapes exactly. q_e = tf.reshape(q_s, new_shape=(1, S, config.num_attention_heads, 1, config.head_dim)) + softcap = tf.cast(ATTN_SOFTCAP, dtype=_DT) z_ctx = ( - tf.reduce(q_e * k_ctx, axes=(-1,), keepdim=True, kind="sum") / ATTN_SOFTCAP + tf.reduce(q_e * k_ctx, axes=(-1,), keepdim=True, kind="sum") / softcap ) - score_ctx = tf.tanh(z_ctx) * ATTN_SOFTCAP + score_ctx = tf.tanh(z_ctx) * softcap z_new = ( - tf.reduce(q_s * k_new, axes=(-1,), keepdim=True, kind="sum") / ATTN_SOFTCAP + tf.reduce(q_s * k_new, axes=(-1,), keepdim=True, kind="sum") / softcap ) - score_new = tf.tanh(z_new) * ATTN_SOFTCAP + score_new = tf.tanh(z_new) * softcap # Log-sum-exp merge of the two groups against their joint max. peak = tf.max( @@ -325,7 +326,7 @@ def embed( row = tf.reshape( tf.index_select(w_embed, token_ids, dim=0), new_shape=(1, S, config.hidden_size) ) - return row * EMBED_SCALE + return row * tf.cast(EMBED_SCALE, dtype=_DT) @func def final_rms_norm( @@ -343,8 +344,9 @@ def lm_head( # Soft-capped as `Gemma2ForCausalLM.forward` caps it, at # `final_logit_softcapping` rather than attention's cap. logits = tf.matmul(tf.reshape(hidden, new_shape=(1, config.hidden_size)), w_head) - z = logits / LOGIT_SOFTCAP - return tf.tanh(z) * LOGIT_SOFTCAP + softcap = tf.cast(LOGIT_SOFTCAP, dtype=_DT) + z = logits / softcap + return tf.tanh(z) * softcap @lm_head.converter("w_head") def _( diff --git a/tests/models/minicpm3_4b/model.py b/tests/models/minicpm3_4b/model.py index 59ce9f2c..2ee4f290 100644 --- a/tests/models/minicpm3_4b/model.py +++ b/tests/models/minicpm3_4b/model.py @@ -364,7 +364,7 @@ def embed( row = tf.reshape( tf.index_select(w_embed, token_ids, dim=0), new_shape=(1, S, config.hidden_size) ) - return row * EMBED_SCALE + return row * tf.cast(EMBED_SCALE, dtype=_DT) @func def final_rms_norm( @@ -385,7 +385,9 @@ def lm_head( ) -> Tensor[(1, config.vocab_size), _DT]: # `MiniCPM3ForCausalLM.forward` divides the hidden state by # `logits_scaling` before the head, not after. - scaled = tf.reshape(hidden, new_shape=(1, config.hidden_size)) / LOGITS_SCALING + scaled = tf.reshape(hidden, new_shape=(1, config.hidden_size)) / tf.cast( + LOGITS_SCALING, dtype=_DT + ) return tf.matmul(scaled, w_head) @lm_head.converter("w_head") diff --git a/tests/parser/README.md b/tests/parser/README.md deleted file mode 100644 index 296effd5..00000000 --- a/tests/parser/README.md +++ /dev/null @@ -1,62 +0,0 @@ -# tests/parser - -Three shapes, and the rule for picking one. - -## A new parser feature goes into an existing program - -`programs.py` holds five programs, each parsed once at import. Find the one -whose `features` list names the area you are touching — `HirExpressions`, -`HirGrid`, `HirSharded`, `HirModule`, or the `tir_*` prim funcs — and add a -line to its body plus an entry to `features`. Then regenerate the goldens: - -``` -pytest tests/parser --update-golden -``` - -`golden/.py` is the program printed back as DSL source by -`tilefoundry.inspection.as_script`. Read the diff before keeping it: that diff -*is* the assertion, so an unexplained change there is the finding. The goldens -are excluded from ruff, because reformatting recorded printer output edits the -evidence rather than the code. - -There is no golden for the TIR programs — nothing prints a `PrimFunction` — so -`test_tir_programs.py` asserts on the parsed nodes directly. - -## A new refusal goes into the table - -`error_cases.py` holds every subject `tests/parser` refuses, as one row each: -the subject, the exception, and the message it must carry. `subject` is DSL -source to import, or a builder for the ones with no source to feed — a -definition that must fail while it is decorated, a hand-forged TIR node, an -operand check that never reaches the parser. `test_refused_programs.py` runs -the table and is the only entry point; nothing else here carries a bare -`pytest.raises`. - -## Only what neither can express gets its own test - -`test_programs.py` and `test_tir_programs.py` carry the rest, each saying in -its docstring why a golden cannot hold it. In practice that is node identity, -a target's registered `Op` class, a canonicalisation the printer renders back -as the sugar it was written as, and the things that were never one program -parsing: evaluating against torch, printer equality, re-elaboration, and -reading a generated `.pyi`. - -A new file is the last resort, not the first. Adding one is a claim that the -subject fits none of the three shapes above. - -## Measuring - -The gauge is a whole round of `tests/parser`, not per-test contexts — module-level -programs are parsed at import, and that coverage belongs to no test's context. - -``` -COVERAGE_FILE=/parser.coverage python -m pytest tests/parser -q \ - -p no:randomly --cov=tilefoundry --cov-branch --cov-report= -COVERAGE_FILE=/parser.coverage python -m coverage report \ - --include='*/tilefoundry/parser/*' -COVERAGE_FILE=/parser.coverage python -m coverage report -``` - -Read both totals. Plenty of what these tests pin lives outside the parser -package — `dsl/`, `evaluator/`, `ir/tir/`, `codegen/` — and the parser-only -number is blind to all of it. diff --git a/tests/parser/conftest.py b/tests/parser/conftest.py deleted file mode 100644 index 2f7c7494..00000000 --- a/tests/parser/conftest.py +++ /dev/null @@ -1,65 +0,0 @@ -"""The golden-file fixture for the parser programs. - -A golden is the program printed back as DSL source, so what a reviewer reads is -the program itself rather than a list of node assertions. ``--update-golden`` -rewrites the files from what the parser produced instead of asserting against -them; run it whenever the printer changes, and read the diff before keeping it. -""" - -from __future__ import annotations - -import difflib -from dataclasses import dataclass -from pathlib import Path - -import pytest - - -def pytest_addoption(parser: pytest.Parser) -> None: - """Register ``--update-golden``.""" - parser.addoption( - "--update-golden", - action="store_true", - help="rewrite tests/parser/golden/*.py from what the parser produced", - ) - - -@dataclass(frozen=True) -class GoldenFiles: - """The recorded output of each program, under one directory.""" - - root: Path - update: bool - - def check(self, name: str, actual: str) -> None: - """Compare *actual* against the recorded ``name``, or record it.""" - path = self.root / name - if self.update: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(actual) - return - assert path.exists(), f"no golden at {path}; rerun with --update-golden" - expected = path.read_text() - if actual == expected: - return - diff = "".join( - difflib.unified_diff( - expected.splitlines(keepends=True), - actual.splitlines(keepends=True), - fromfile=f"{name} (recorded)", - tofile=f"{name} (parsed now)", - ) - ) - raise AssertionError( - f"{name} no longer matches what the parser produces:\n\n{diff}\n" - "If the new output is right, rerun with --update-golden." - ) - - -@pytest.fixture -def golden(request: pytest.FixtureRequest) -> GoldenFiles: - """The golden directory beside this file, honouring ``--update-golden``.""" - return GoldenFiles( - root=Path(__file__).parent / "golden", - update=bool(request.config.getoption("--update-golden")), - ) diff --git a/tests/parser/error_cases.py b/tests/parser/error_cases.py deleted file mode 100644 index 3ea9b3e9..00000000 --- a/tests/parser/error_cases.py +++ /dev/null @@ -1,1403 +0,0 @@ -"""Every subject ``tests/parser`` refuses, in one table. - -A row is one subject and the diagnostic it must raise. ``subject`` is DSL source -to import, or a builder for the ones with no source to feed: a definition that -must fail while it is decorated, a hand-forged TIR node fed to the verifier, an -operand check that never reaches the parser. ``test_refused_programs.py`` is the -only entry point, so a new refusal is a new row rather than a new file. The -subjects the rows share live here too, and the surviving test files import them -from here instead of keeping a copy each. -""" - -from __future__ import annotations - -import ast -import textwrap -from collections.abc import Callable -from dataclasses import dataclass - -import pytest -import torch - -import tilefoundry.codegen.cuda # noqa: F401 — trigger emitter autodiscovery -from tests._source import import_dsl -from tests.fixtures.logical.hir_composition import Expert -from tests.fixtures.shapes.window_programs import tile_window_add -from tilefoundry import func, module, prim_func -from tilefoundry.codegen.cuda.context import CodegenContext -from tilefoundry.dsl import ConstTensor, DimVar, T, Tensor, tf -from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op names used by the @func bodies -from tilefoundry.evaluator import EvalError, evaluate -from tilefoundry.ir.core import Var, VerifyError -from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.tir.stmts import Evaluate, MeshScope, Return, Sequential -from tilefoundry.ir.tir.sync import Sync, classify -from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimAdd, simplify_dim -from tilefoundry.ir.types.shard import Layout, Mesh, P, ShardLayout, Topology -from tilefoundry.ir.types.shard.layout import ComposedLayout -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.module import _DECLARING -from tilefoundry.parser.sugar import parse_sugar -from tilefoundry.target import CudaTarget - - -@dataclass(frozen=True) -class ParseErrorCase: - """One refused subject: what is fed in, and the diagnostic it must raise.""" - - id: str - subject: str | Callable[[], object] - """DSL source to import, or a builder for the ones with no source to feed.""" - raises: type[BaseException] - match: str - """The diagnostic, tight enough that no other row in the table satisfies it. - - A pattern loose enough to accept a neighbour's message stops asserting - which thing went wrong, and a row can then drift off the check it was - written for without anything noticing. - """ - - -def run_parse_error_case(case: ParseErrorCase) -> None: - """Run one ``ParseErrorCase``: the subject is refused with that diagnostic.""" - with pytest.raises(case.raises, match=case.match): - if callable(case.subject): - case.subject() - else: - import_dsl(textwrap.dedent(case.subject).lstrip("\n")) - - -HIR_PRELUDE = """from tilefoundry import func -from tilefoundry.dsl.tf import * -from tilefoundry.dsl import Tensor -""" - -_ONE_TENSOR = 'x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]' - - -def hir_source(*body: str, signature: str = _ONE_TENSOR, prelude: str = HIR_PRELUDE) -> str: - """A one-``@func`` script. - - *signature* closes the parameter list and states the return annotation, - *body* lines carry their own nesting, *prelude* holds module-level bindings - the parser resolves through ``fn.__globals__``. - """ - lines = "\n".join(f" {line}" for line in body) - return f"{prelude}\n@func\ndef f({signature}:\n{lines}\n" - - -_WHERE_PRELUDE = ( - "from tilefoundry.ir.types.shard import Layout, Mesh, Topology\n\n" - 'cta_mesh = Mesh((Topology("cta", 8),), Layout((8,), (1,)))' -) - - -def where_source(body: str, preamble: str = _WHERE_PRELUDE, ret: str = '(8, 16), "bf16"') -> str: - """A one-``@func`` script carrying ``where(...)`` annotations on its body.""" - return f"""from __future__ import annotations -from tilefoundry import func -from tilefoundry.dsl import Tensor, tf - -{preamble} - -@func -def candidate(x: Tensor[(8, 16), "bf16"]) -> Tensor[{ret}]: -{body} -""" - - -@module(entry="entry") -class Callee: - """A child Module whose entry calls a sibling by its bare class-body binding.""" - - @func - def helper(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return tf.mul(x, x) - - @func - def entry(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return helper(x) # noqa: F821 — sibling binding in the class body - - -@module -class NoEntry: - """A Module with no callable entry, so calling it has nothing to reach.""" - - @func - def only(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return tf.add(x, x) - - -@module(entry="device") -class PrimEntry: - """A Module whose entry is a device ``@prim_func`` rather than an HIR Function.""" - - @prim_func(target=CudaTarget("nvidia.h200_sxm")) - def device(x: Tensor[(8,), "f32"]) -> None: # noqa: ARG001 - with Mesh((Topology("thread", 8),), Layout(shape=(8,), strides=(1,))) as m: - T.sync(m) - - -def _arg_type_mismatch() -> None: - """A nested call whose argument dtype does not match the callee's parameter.""" - - @func - def _inner_double(x: Tensor[(8, 64), "f32"]) -> Tensor[(8, 64), "f32"]: - return add(x, x) # noqa: F821 - - @func - def _bad_dtype(x: Tensor[(8, 64), "bf16"]) -> Tensor[(8, 64), "f32"]: - return _inner_double(x) # noqa: F841 - - -def _reach_a_child_entry_by_name() -> None: - @module(entry="reach") - class _ReachEntry: - leaf = Callee - - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf.entry(x) # noqa: F821 - - -def _reach_a_child_helper_by_name() -> None: - @module(entry="reach") - class _ReachHelper: - leaf = Callee - - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf.helper(x) # noqa: F821 - - -def _reach_a_module_member_by_class() -> None: - @func - def _reach_by_class(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return Callee.entry(x) - - -def _call_a_module_with_no_entry() -> None: - @module(entry="reach") - class _CallsEntryless: - leaf = NoEntry - - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf(x) # noqa: F821 - - -def _call_a_module_whose_entry_is_a_prim_func() -> None: - @module(entry="reach") - class _CallsPrim: - leaf = PrimEntry - - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf(x) # noqa: F821 - - -def _bare_decorator_leaves_the_child_unattached() -> None: - @module - class _BareDecorated: - leaf = Callee - - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf(x) # noqa: F821 - - -def _module_call_with_no_binding() -> None: - """A Module named in a body no class-body binding attaches. - - This row and the list-bound one below own the only arcs in ``module.py`` - that nothing else under ``tests/parser`` reaches. A third subject reaching - the same message from a ``@run.converter`` body owned none, so it is not a - row. - """ - - @module(entry="reach") - class _Unattached: - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return Callee(x) - - -def _module_call_bound_only_inside_a_list() -> None: - @module(entry="reach") - class _ListAttached: - kids = [Callee] - - @func - def reach(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return Callee(x) - - -def _declaration_left_open_by_a_failed_class_body() -> None: - """A class body that raises leaves its declaration open, resolving nothing. - - The two halves cannot be separate rows: the second only means anything while - the first has left a declaration on the stack, and the stack has to be - restored either way. So the setup asserts its own ``RuntimeError`` here and - the row pins what the leaked declaration must *not* do. The standalone - ``@func`` below is also the whole of a subject that used to be its own test - and owned no arcs of its own, so that one is not a row. - """ - open_declarations = len(_DECLARING) - try: - with pytest.raises(RuntimeError, match="boom"): - - @module(entry="never") - class _Boom: - raise RuntimeError("boom") - - @func - def _after(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return Callee(x) - finally: - del _DECLARING[open_declarations:] - - -def _direct_call_of_the_wrong_arity() -> None: - """A sibling ``@func`` called with fewer arguments than it declares. - - A standalone ``@func`` reaching the same site with the shorter ``arity - mismatch`` message owned no arcs of its own, and that message is a substring - of this one, so it is not a row. - """ - - @module(entry="root", target=CudaTarget("nvidia.h200_sxm")) - class _Direct: - @func - def leaf(x: Tensor[(4, 8), "f32"], w: ConstTensor[(8, 8), "f32"]) -> Tensor[(4, 8), "f32"]: - return tf.matmul(x, w) - - @func - def root(x: Tensor[(4, 8), "f32"]) -> Tensor[(4, 8), "f32"]: - return leaf(x) # noqa: F821 - - -def _child_call_of_the_wrong_width() -> None: - @module(entry="fused", target=CudaTarget("nvidia.h200_sxm")) - class _TooMany: - mlp = Expert - - @func - def fused(x: Tensor[(4, 8), "f32"]) -> Tensor[(4, 8), "f32"]: - return mlp(x, x) # noqa: F821 - - -CTX_LEN = DimVar("CTX_LEN", 1, 4097) - - -def _evaluate_a_non_divisible_tile_window() -> None: - evaluate(tile_window_add, torch.ones((10, 4)), torch.ones((4, 4)), device="cpu") - - -_M_MULTI = Mesh((Topology("thread", 6 * 32),), Layout((6, 32), (32, 1)), names=("w", "t")) -_M_STATE = Mesh( - (Topology("thread", 4 * 2 * 16),), Layout((4, 2, 16), (32, 16, 1)), names=("l", "g", "t") -) -_S_DYN = DimVar("seq_len", 1, 4) -_MESH_DIM_W = DimVar("W", 1, 8) - - -def _multi_axis_split_not_divisible() -> None: - """A dim must be divisible by the product of the mesh extents; ``100 @`` is rejected.""" - - @func - def _bad( - a: Tensor[(1, 100), "f32", (1, 100 @ (_M_MULTI.w, _M_MULTI.t)), "smem"], - ) -> Tensor[(1, 100), "f32"]: - return a - - -def _value_state_not_final() -> None: - """The ``{...}`` value-state set is valid only as the last outer item.""" - - @func - def _bad( - a: Tensor[ - (4, 64), - "f32", - ((4 @ _M_STATE.l, 64), {_M_STATE.t @ P("sum")}, (64, 1)), - "smem", - ], - ) -> Tensor[(4, 64), "f32"]: - return a - - -def _value_state_bare_p() -> None: - """``P(...)`` in the value-state set requires its reduction argument.""" - - @func - def _bad( - a: Tensor[(4, 64), "f32", ((4 @ _M_STATE.l, 64), {_M_STATE.t @ P()}), "smem"], - ) -> Tensor[(4, 64), "f32"]: - return a - - -def _mesh_coordinate_slices_a_placed_tensor() -> None: - @func(topologies=(Topology("cta", 8),)) - def _bad(x: Tensor[(8,), "i64"]) -> Tensor[(4,), "i64"]: - with Mesh(("cta",), layout=(8,), names=("w",)) as cta: - placed = reshard(x, (8 @ cta.w,), "rmem") # noqa: F405, F821 - return placed[cta.w : cta.w + 4] - - -def _unresolved_dynamic_split_axis() -> None: - """Different symbolic extents cannot establish divisibility before binding.""" - other = DimVar("other", 1, 4) - cta = Mesh((Topology("cta", other),), Layout((other,), (1,)), names=("cta",)) - node = ast.parse("(1, S @ cta, 32, 128)", mode="eval").body - parse_sugar( - node, - ShardLayout, - mesh_resolver=lambda n: cta if n == "cta" else None, - closure={"S": _S_DYN}, - ) - - -def mesh_dims_reshard_func(warps, lanes): - """A reshard whose mesh shape comes from ``layout=(warps, lanes)``. - - The dims may be integer literals or closure Names — a closure int must - resolve like the literal, and a dynamic ``DimVar`` in that static-extent - position must be rejected. - """ - - @func(topologies=(Topology("thread", 128),)) - def _f(x: Tensor[(1, 128), "bf16"]) -> Tensor[(1, 128), "bf16"]: - with Mesh(("thread",), layout=(warps, lanes), names=("w", "t")) as m: - xr = reshard(x, (1, 128 @ (m.w, m.t)), "rmem") # noqa: F405, F821 - return reshard(xr, (1, 128), "gmem") # noqa: F405, F821 - - return _f - - -def literal_reshard_func(): - """The all-literal twin of ``mesh_dims_reshard_func``. - - ``layout=(4, 32)`` mesh dims spelled out, so the closure-resolved builder - above has something to print equal to. - """ - - @func(topologies=(Topology("thread", 128),)) - def _f(x: Tensor[(1, 128), "bf16"]) -> Tensor[(1, 128), "bf16"]: - with Mesh(("thread",), layout=(4, 32), names=("w", "t")) as m: - xr = reshard(x, (1, 128 @ (m.w, m.t)), "rmem") # noqa: F405, F821 - return reshard(xr, (1, 128), "gmem") # noqa: F405, F821 - - return _f - - -def _symbolic_mesh_extent() -> None: - """A symbolic mesh extent is valid even when a later split is undecidable.""" - mesh_dims_reshard_func(_MESH_DIM_W, 32) - - -def _bool_split_extent_single_axis() -> None: - """A ``bool`` split extent in the single-axis form (``True @ m.w``) is rejected.""" - - @func(topologies=(Topology("thread", 128),)) - def _f(x: Tensor[(1, 128), "bf16"]) -> Tensor[(1, 128), "bf16"]: - with Mesh(("thread",), layout=(4, 32), names=("w", "t")) as m: - xr = reshard(x, (1, True @ m.w), "rmem") # noqa: F405, F821 - return reshard(xr, (1, 128), "gmem") # noqa: F405, F821 - - -def _float_split_extent_single_axis() -> None: - """A non-``bool`` split extent that is not a shape dimension names its own type. - - ``bool`` earns a separate diagnostic because it *is* an int; every other - wrong type gets the plain one, and this row is what keeps that plain one - from going unwritten. - """ - - @func(topologies=(Topology("thread", 128),)) - def _f(x: Tensor[(1, 128), "bf16"]) -> Tensor[(1, 128), "bf16"]: - with Mesh(("thread",), layout=(4, 32), names=("w", "t")) as m: - xr = reshard(x, (1, 1.5 @ m.w), "rmem") # noqa: F405, F821 - return reshard(xr, (1, 128), "gmem") # noqa: F405, F821 - - -def _duplicate_topology_name() -> None: - @func(topologies=(Topology("cta", 128), Topology("cta", 64))) - def _dup(a: Tensor[(1, 1536), "f32"]) -> Tensor[(1, 1536), "f32"]: - return a - - -def _mesh_on_an_undeclared_topology() -> None: - @func(topologies=(Topology("cta", 128),)) - def _unk(a: Tensor[(1, 1536), "f32"]) -> Tensor[(1, 1536), "f32"]: - with Mesh(("nonexistent",), layout=Layout(shape=(128,), strides=(1,))) as m: # noqa: F841 - return a - - -def _mesh_topology_source(mesh_source: str) -> str: - return ( - "from tilefoundry import func\n" - "from tilefoundry.dsl import Mesh, Tensor\n" - "from tilefoundry.ir.types.shard import Topology\n\n" - "@func(topologies=(Topology('cta', 128),))\n" - "def f(a: Tensor[(128,), 'f32']):\n" - f" with {mesh_source} as cta:\n" - " return a\n" - ) - - -_DTYPE_HEADER = """ -from tilefoundry import func -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import * -""" - -_BAD_CALL_DTYPE = ( - _DTYPE_HEADER - + """ -@func -def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "bf16"]: - return cast(x, dtype="float32") -""" -) - -_BAD_ANNOTATION_DTYPE = ( - _DTYPE_HEADER - + """ -@func -def f(x: Tensor[(8,), "float32"]) -> Tensor[(8,), "f32"]: - return cast(x, dtype="f32") -""" -) - -_BAD_REDUCE_KIND = ( - _DTYPE_HEADER - + """ -@func -def g(x: Tensor[(8,), "f32"]) -> Tensor[(1,), "f32"]: - return reduce(x, axes=(0,), keepdim=True, kind="plus") -""" -) - - -def alloc_frag_kernel(topology, mesh_layout, names=()): - """A kernel that allocs a fragment via ``atom.A`` inside the given scope.""" - - def kernel(a: Tensor[(16, 16), "bf16"]): # noqa: ARG001 - atom = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) - with Mesh((topology,), mesh_layout, names=names) as warp: # noqa: F841 - frag = T.alloc_tensor( # noqa: F841 - TensorType( - shape=(16, 16), dtype=DType.bf16, layout=atom.A, storage=StorageKind.RMEM - ) - ) - - return kernel - - -def _mma_scope(topology, layout) -> Callable[[], object]: - """A builder that parses ``alloc_frag_kernel`` under one candidate thread scope.""" - return lambda: prim_func(target=CudaTarget("nvidia.h200_sxm"))( - alloc_frag_kernel(topology, layout) - ) - - -def _atom_outside_a_mesh_scope() -> None: - def kernel(a: Tensor[(16, 16), "bf16"]): # noqa: ARG001 - atom = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) - frag = T.alloc_tensor( # noqa: F841 - TensorType(shape=(16, 16), dtype=DType.bf16, layout=atom.A, storage=StorageKind.RMEM) - ) - - prim_func(target=CudaTarget("nvidia.h200_sxm"))(kernel) - - -def thread_mesh() -> Mesh: - """A 128-thread block viewed as (4 warps, 32 lanes).""" - return Mesh((Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t")) - - -def cta_mesh() -> Mesh: - """A 128-CTA grid.""" - return Mesh(topologies=(Topology("cta", 128),), layout=Layout(shape=(128,), strides=(1,))) - - -def _binding(name: str = "m") -> Var: - return Var(type=TensorType.scalar(DType.i64, storage=StorageKind.RMEM), name=name) - - -def scoped_sync(mesh: Mesh, sync_mesh: Mesh) -> PrimFunction: - """A ``PrimFunction`` syncing on *sync_mesh* inside a mesh scope over *mesh*. - - The three rows built on this that share the ``enclosing`` diagnostic — no - enclosing mesh at all, a forged sub-box exceeding its parent, a forged - topology tuple — each own arcs in ``ir/tir/sync.py`` no other subject under - ``tests/parser`` reaches. They are three judgements in front of one raise, - so none of them stands in for another. - """ - return PrimFunction( - name="fn", - params=(), - body=Sequential( - body=( - MeshScope( - mesh=mesh, - binding=_binding(), - body=Sequential( - body=(Evaluate(callable=Sync(mesh=sync_mesh), args=()), Return()) - ), - ), - ) - ), - ) - - -def _sync_argument_is_not_a_mesh() -> None: - def kernel(a: Tensor[(128,), "f32"]): # noqa: ARG001 - with Mesh( - (Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t") - ) as m: # noqa: F841 - T.sync(a) - - prim_func(target=CudaTarget("nvidia.h200_sxm"))(kernel) - - -def _sync_with_no_enclosing_mesh() -> None: - verify_prim_function( - PrimFunction( - name="fn", - params=(), - body=Sequential(body=(Evaluate(callable=Sync(mesh=cta_mesh()), args=()), Return())), - ) - ) - - -def _sync_on_a_non_contiguous_slice() -> None: - """A lane subset across warps (``m[:, 1:3]``) is not a contiguous thread interval.""" - m = thread_mesh() - verify_prim_function(scoped_sync(m, m[:, 1:3])) - - -def _sync_on_a_mesh_the_scope_does_not_bind() -> None: - """An un-sliced mesh the enclosing scope simply does not bind. - - This is the middle of the three ways to miss an enclosing mesh: a scope is - in force, and the synced mesh is a whole mesh rather than a sub-box, so what - is wrong is the identity of the mesh and not the shape of a slice. - """ - enclosing = thread_mesh() - other = Mesh((Topology("thread", 64),), Layout(shape=(64,), strides=(1,))) - verify_prim_function(scoped_sync(enclosing, other)) - - -def _sync_on_a_forged_subbox_exceeding_its_parent() -> None: - """A (1, 64) sub-box of a (4, 32) parent is not constructible by ``Mesh.__getitem__``.""" - enclosing = thread_mesh() - forged = Mesh( - topologies=enclosing.topologies, - layout=ComposedLayout(inner=None, offset=0, outer=Layout((1, 64), (32, 1))), - names=enclosing.names, - ) - verify_prim_function(scoped_sync(enclosing, forged)) - - -def _sync_on_a_forged_topology_mismatch() -> None: - """A forged sync mesh sharing the primary topology but not the full topology tuple.""" - enclosing = Mesh( - topologies=(Topology("warp", 4), Topology("thread", 32)), - layout=Layout(shape=(4, 32), strides=(32, 1)), - ) - forged = Mesh( - topologies=(Topology("warp", 4),), - layout=ComposedLayout(inner=None, offset=0, outer=Layout((2, 32), (32, 1))), - names=enclosing.names, - ) - verify_prim_function(scoped_sync(enclosing, forged)) - - -def _sync_on_a_cross_warp_unaligned_slice() -> None: - """A contiguous but cross-warp-unaligned range (lanes 16..47).""" - m = Mesh((Topology("thread", 64),), Layout(shape=(64,), strides=(1,))) - verify_prim_function(scoped_sync(m, m[16:48])) - - -def _classify_a_partial_cta_slice() -> None: - """A cta slice is a subset of CTAs, and no barrier covers that.""" - classify(cta_mesh()[0:64]) - - -def _exhaust_the_named_barriers() -> None: - ctx = CodegenContext() - ctx.reset_barrier_ids() - for _ in range(15): - ctx.alloc_barrier_id() - ctx.alloc_barrier_id() - - -_HIR_BODY_STATEMENTS: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="yield-in-hir-body", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - yield x - """, - raises=VerifyError, - match=r"`yield` is not an HIR statement", - ), - ParseErrorCase( - id="lambda-in-hir-body", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - g = lambda y: y - return g(x) - """, - raises=VerifyError, - match="Lambda", - ), - ParseErrorCase( - id="augassign-in-grid-body", - subject=hir_source("o = relu(x)", "for i in range(8):", " o += x", "return o"), - raises=VerifyError, - match="augmented assignment", - ), - ParseErrorCase( - id="return-in-grid-body", - subject=hir_source("for i in range(8):", " return x", "return x"), - raises=VerifyError, - match="must not contain `return`", - ), -) - - -_OP_CALL_SURFACE: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="unknown-op-name", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return totally_undefined_op(x) - """, - raises=VerifyError, - match=r"unknown HIR callable|unknown Op name", - ), - ParseErrorCase( - id="tuple-input-for-a-plain-tensor-param", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return relu((x, x)) - """, - raises=VerifyError, - match=r"unsupported AST node in expression: Tuple", - ), - ParseErrorCase( - id="integer-literal-keeps-its-own-dtype", - subject=""" - from tilefoundry import func - from tilefoundry.dsl import Tensor - from tilefoundry.dsl.tf import * - - @func - def f(x: Tensor[(1, 8), 'f32']) -> Tensor[(1, 8), 'f32']: - return mul(x, 2) - """, - raises=VerifyError, - match=r"Binary MUL: dtype mismatch \(f32 vs i64\)", - ), - ParseErrorCase( - id="unknown-dtype-string-at-a-call", - subject=_BAD_CALL_DTYPE, - raises=VerifyError, - match=r"DType: unknown value 'float32'", - ), - ParseErrorCase( - id="unknown-dtype-string-in-an-annotation", - subject=_BAD_ANNOTATION_DTYPE, - raises=ValueError, - match=r"DType: unknown value 'float32'", - ), - ParseErrorCase( - id="unknown-reduce-kind-string", - subject=_BAD_REDUCE_KIND, - raises=VerifyError, - match=r"ReduceKind: unknown value 'plus'", - ), -) - - -_DIM_OPERANDS: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="dim-plus-bool", - subject=lambda: CTX_LEN + True, - raises=TypeError, - match=r"dim arithmetic: bool operand True is not a dimension", - ), - ParseErrorCase( - id="dim-plus-object", - subject=lambda: CTX_LEN + object(), - raises=TypeError, - match=r"unsupported operand type\(s\) for \+: 'DimVar' and 'object'", - ), - ParseErrorCase( - id="simplify-dim-bool-on-the-left", - subject=lambda: simplify_dim(DimAdd, (True, CTX_LEN)), - raises=TypeError, - match=r"simplify_dim: bool operand True is not a ShapeDim", - ), - ParseErrorCase( - id="simplify-dim-bool-on-the-right", - subject=lambda: simplify_dim(DimAdd, (CTX_LEN, False)), - raises=TypeError, - match=r"simplify_dim: bool operand False is not a ShapeDim", - ), -) - - -_CALL_BOUNDARY: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="nested-call-arg-type-mismatch", - subject=_arg_type_mismatch, - raises=VerifyError, - match=r"arg 0 shape/dtype mismatch", - ), - ParseErrorCase( - id="direct-call-arity-mismatch", - subject=_direct_call_of_the_wrong_arity, - raises=VerifyError, - match="nested @func call arity mismatch", - ), - ParseErrorCase( - id="child-call-of-the-wrong-width", - subject=_child_call_of_the_wrong_width, - raises=VerifyError, - match="takes 1 activation", - ), - ParseErrorCase( - id="reach-a-child-entry-by-name", - subject=_reach_a_child_entry_by_name, - raises=VerifyError, - match=r"'leaf.entry': a Module is called through its bare binding", - ), - ParseErrorCase( - id="reach-a-child-helper-by-name", - subject=_reach_a_child_helper_by_name, - raises=VerifyError, - match=r"'leaf.helper': a Module is called through its bare binding", - ), - ParseErrorCase( - id="reach-a-module-member-by-class", - subject=_reach_a_module_member_by_class, - raises=VerifyError, - match=r"'Callee.entry': a Module is called through its bare binding", - ), - ParseErrorCase( - id="call-a-module-with-no-entry", - subject=_call_a_module_with_no_entry, - raises=VerifyError, - match=r"Module 'NoEntry' declares no entry", - ), - ParseErrorCase( - id="call-a-module-whose-entry-is-a-prim-func", - subject=_call_a_module_whose_entry_is_a_prim_func, - raises=VerifyError, - match="rather than an hir Function", - ), - ParseErrorCase( - id="bare-decorator-leaves-the-child-unattached", - subject=_bare_decorator_leaves_the_child_unattached, - raises=VerifyError, - match=r"'leaf': a Module is called only from a function authored in a @module class body", - ), - ParseErrorCase( - id="declaration-left-open-by-a-failed-class-body", - subject=_declaration_left_open_by_a_failed_class_body, - raises=VerifyError, - match=r"'Callee': a Module is called only from a function authored in a @module class body", - ), - ParseErrorCase( - id="module-call-with-no-binding", - subject=_module_call_with_no_binding, - raises=ValueError, - match=r"@module '_Unattached': call\(s\) to Module\(s\) \['Callee'\] that no class-body binding attaches", - ), - ParseErrorCase( - id="module-call-bound-only-inside-a-list", - subject=_module_call_bound_only_inside_a_list, - raises=ValueError, - match=r"@module '_ListAttached': call\(s\) to Module\(s\) \['Callee'\] that no class-body binding attaches", - ), -) - - -_FUNCTION_BINDINGS: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="kernel-binding-underscore", - subject=""" - from tilefoundry import func, module - from tilefoundry.dsl import Tensor - - @module() - class _KernelUnderscore: - @func - def _(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return x - """, - raises=ValueError, - match=r"@module '_KernelUnderscore': a kernel binding may not be named '_'", - ), - ParseErrorCase( - id="variant-binding-underscore", - subject=""" - from tilefoundry import func, module - from tilefoundry.dsl import DimVar, DimVarRangePat, Tensor - - _N = DimVar("N", 1, 8) - - @module() - class _VariantUnderscore: - @func - def dispatch(x: Tensor[(_N,), "f32"]) -> Tensor[(_N,), "f32"]: - pass - - @dispatch.specialize(DimVarRangePat("N", 1, 4)) - def _(x: Tensor[(_N,), "f32"]) -> Tensor[(_N,), "f32"]: - return x - """, - raises=ValueError, - match=( - r"@module '_VariantUnderscore' base 'dispatch': " - r"a variant binding may not be named '_'" - ), - ), - ParseErrorCase( - id="duplicate-kernel-binding", - subject=""" - from tilefoundry import func, module - from tilefoundry.dsl import Tensor - - @module() - class _DuplicateKernel: - @func - def run(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return x - - @func - def run(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return x - """, - raises=ValueError, - match=r"@module '_DuplicateKernel': duplicate kernel binding 'run'", - ), -) - - -_SUBSCRIPTS: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="tile-with-too-many-args", - subject=hir_source( - "for i in tile(1, 2, 3):", - " y = relu(x)", - ), - raises=VerifyError, - match="tile.. takes 2 arguments", - ), - ParseErrorCase( - id="runtime-slice-stride", - subject=hir_source( - "return x[0:8:step, :]", - signature='x: Tensor[(8, 4), "f32"], step: Tensor[(), "i64"]) -> Tensor[(4, 4), "f32"]', - ), - raises=VerifyError, - match=r"tensor subscript axis 0: slice stride must be a compile-time dimension", - ), - ParseErrorCase( - id="runtime-start-with-an-unrelated-stop", - subject=hir_source( - "return x[start:8, :]", - signature=( - 'x: Tensor[(8, 4), "f32"], start: Tensor[(), "i64"]) -> Tensor[(4, 4), "f32"]' - ), - ), - raises=VerifyError, - match=r"tensor subscript axis 0: a run-time start needs the stop endpoint", - ), - ParseErrorCase( - id="non-divisible-tile-window-at-evaluate-time", - subject=_evaluate_a_non_divisible_tile_window, - raises=EvalError, - match="Slice window exceeds axis 0", - ), - ParseErrorCase( - id="window-moved-by-a-runtime-offset", - subject=hir_source( - "o = relu(x[:, 0:2])", - "for n in tile(4, 2):", - " o = relu(x[:, n + k])", - "return o", - signature='x: Tensor[(1, 8), "f32"], k: Tensor[(), "i64"]) -> Tensor[(1, 2), "f32"]', - ), - raises=VerifyError, - match="moves by a compile-time integer", - ), - ParseErrorCase( - id="window-reversed-instead-of-moved", - subject=hir_source( - "o = relu(x[:, 0:2])", - "for n in tile(4, 2):", - " o = relu(x[:, 4 - n])", - "return o", - signature='x: Tensor[(1, 8), "f32"]) -> Tensor[(1, 2), "f32"]', - ), - raises=VerifyError, - match="reverses the window", - ), - ParseErrorCase( - id="window-moved-off-the-end", - subject=hir_source( - "o = relu(x[:, 0:2])", - "for n in tile(4, 2):", - " o = relu(x[:, n + 5])", - "return o", - signature='x: Tensor[(1, 8), "f32"]) -> Tensor[(1, 2), "f32"]', - ), - raises=VerifyError, - match=r"reads \[7, 9\).*axis is 8 long", - ), - ParseErrorCase( - id="window-moved-before-the-front", - subject=hir_source( - "o = relu(x[:, 0:2])", - "for n in tile(4, 2):", - " o = relu(x[:, n - 2])", - "return o", - signature='x: Tensor[(1, 8), "f32"]) -> Tensor[(1, 2), "f32"]', - ), - raises=VerifyError, - match="begin before the axis", - ), - ParseErrorCase( - id="subscript-rank-mismatch", - subject=hir_source( - "o = relu(x)", - "for ok in tile(2048, 512):", - " o = relu(x[ok])", - "return o", - signature='x: Tensor[(1, 2048), "f32"]) -> Tensor[(1, 2048), "f32"]', - ), - raises=VerifyError, - match="rank 1 != tensor rank 2", - ), - ParseErrorCase( - id="runtime-tuple-index", - subject=hir_source( - "out = quant(x)", - "return out[i]", - signature=( - 'x: Tensor[(1, 1536), "bf16"], i: Tensor[(), "i64"])' - ' -> Tensor[(1, 1536), "fp8e4m3"]' - ), - ), - raises=VerifyError, - match="integer constant index", - ), - ParseErrorCase( - id="window-scaled-instead-of-moved", - subject=hir_source( - "o = relu(x[:, 0:2])", - "for n in tile(4, 2):", - " o = relu(x[:, n * 2])", - "return o", - signature='x: Tensor[(1, 8), "f32"]) -> Tensor[(1, 2), "f32"]', - ), - raises=VerifyError, - match="unsupported indexer", - ), -) - - -_GRID_LOOPS: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="single-argument-tile", - subject=hir_source("for i in tile(8):", " y = relu(x)"), - raises=VerifyError, - match=r"use range\(extent\)", - ), - ParseErrorCase( - id="tile-with-a-keyword-step", - subject=hir_source("for i in tile(8, step=2):", " y = relu(x)"), - raises=VerifyError, - match=r"tile\(\) does not accept keyword args", - ), - ParseErrorCase( - id="range-with-a-keyword-stop", - subject=hir_source("for i in range(stop=8):", " y = relu(x)"), - raises=VerifyError, - match=r"range\(\) does not accept keyword args", - ), - ParseErrorCase( - id="range-over-a-non-dim-expr", - subject=hir_source("for i in range(x):", " y = relu(x)"), - raises=VerifyError, - match=r"and extent=x \(Var\) is not one", - ), - ParseErrorCase( - id="range-over-a-float-extent", - subject=hir_source("for i in range(8.5):", " y = relu(x)"), - raises=VerifyError, - match=r"and extent=float is not one", - ), -) - - -_SHARD_SUGAR: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="multi-axis-split-not-divisible", - subject=_multi_axis_split_not_divisible, - raises=ValueError, - match="not divisible", - ), - ParseErrorCase( - id="value-state-not-the-last-outer-item", - subject=_value_state_not_final, - raises=ValueError, - match="last outer item", - ), - ParseErrorCase( - id="value-state-with-a-bare-p", - subject=_value_state_bare_p, - raises=ValueError, - match="reduction argument", - ), - ParseErrorCase( - id="mesh-coordinate-slices-a-placed-tensor", - subject=_mesh_coordinate_slices_a_placed_tensor, - raises=VerifyError, - match="data-dependent mesh ownership", - ), - ParseErrorCase( - id="unresolved-dynamic-split-axis", - subject=_unresolved_dynamic_split_axis, - raises=ValueError, - match=r"split layout dim DimVar\(name='seq_len'.*mesh extent DimVar\(name='other'", - ), - ParseErrorCase( - id="symbolic-mesh-extent", - subject=_symbolic_mesh_extent, - raises=ValueError, - match=r"split layout dim 128 and mesh extent DimVar\(name='W'.*at axis position 0", - ), - ParseErrorCase( - id="bool-layout-extent", - subject=_bool_split_extent_single_axis, - raises=ValueError, - match=r"bool True is not one; bool is an int subclass", - ), - ParseErrorCase( - id="float-layout-extent", - subject=_float_split_extent_single_axis, - raises=ValueError, - match=r"layout dim must be a shape dimension \(int / DimVar / dim-op Expr\), got float", - ), -) - - -_TOPOLOGY: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="duplicate-topology-name", - subject=_duplicate_topology_name, - raises=VerifyError, - match="duplicate topology name", - ), - ParseErrorCase( - id="mesh-on-an-undeclared-topology", - subject=_mesh_on_an_undeclared_topology, - raises=VerifyError, - match="topology.*not declared", - ), - ParseErrorCase( - id="mesh-with-a-bare-topology-name", - subject=_mesh_topology_source('Mesh("cta", layout=(128,))'), - raises=VerifyError, - match="tuple of declared topology names", - ), - ParseErrorCase( - id="mesh-with-a-keyword-topology-name", - subject=_mesh_topology_source('Mesh(topology="cta", layout=(128,))'), - raises=VerifyError, - match="tuple of declared topology names", - ), -) - - -_TUPLE_UNPACK: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="unpack-a-non-tuple-rhs", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def bad_rhs(a: Tensor[(1, 4), "f32"], b: Tensor[(1, 4), "f32"]) -> Tensor[(1, 4), "f32"]: - p, q = add(a, b) - return p - """, - raises=VerifyError, - match="tuple unpack requires RHS of TupleType", - ), - ParseErrorCase( - id="unpack-of-the-wrong-arity", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def bad_targets(x: Tensor[(1, 1536), "bf16"]) -> Tensor[(1, 1536), "fp8e4m3"]: - a, b, c = quant(x) - return a - """, - raises=VerifyError, - match="tuple unpack arity mismatch", - ), - ParseErrorCase( - id="unpack-into-a-nested-target", - subject=""" - from tilefoundry import func - from tilefoundry.dsl.tf import * - from tilefoundry.dsl import Tensor - - @func - def bad_targets(x: Tensor[(1, 1536), "bf16"]) -> Tensor[(1, 1536), "fp8e4m3"]: - (a, b), c = quant(x) - return a - """, - raises=VerifyError, - match="targets must all be plain names", - ), -) - - -_CONSTRAINTS: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="where-with-no-kwargs", - subject=where_source(" y: where() = tf.add(x, x)\n return y"), - raises=VerifyError, - match=r"where\(\.\.\.\) cannot be empty", - ), - ParseErrorCase( - id="where-with-an-empty-layout", - subject=where_source(" y: where(layout=()) = tf.add(x, x)\n return y"), - raises=VerifyError, - match="layout constraint cannot be empty", - ), - ParseErrorCase( - id="where-binding-one-topology-twice", - subject=where_source( - ' y: where(layout=((_, 16 @ cta), {cta @ P("sum")})) = tf.add(x, x)\n return y' - ), - raises=VerifyError, - match="layout constraint cannot bind one topology more than once", - ), - ParseErrorCase( - id="where-with-two-binding-sets", - subject=where_source( - ' y: where(layout=((_, 16), {cta @ P("sum")}, {cta @ B()})) = tf.add(x, x)\n' - " return y" - ), - raises=VerifyError, - match="layout constraint accepts one binding set", - ), - ParseErrorCase( - id="where-with-an-unknown-kwarg", - subject=where_source(' y: where(partial=P("sum")) = tf.add(x, x)\n return y'), - raises=VerifyError, - match=r"where\(\.\.\.\) has unknown field 'partial'", - ), - ParseErrorCase( - id="where-with-a-non-int-extent", - subject=where_source(" y: where(layout=(1.5,)) = tf.add(x, x)\n return y"), - raises=VerifyError, - match="layout dimensions must use `_`, an integer, or a symbolic extent", - ), - ParseErrorCase( - id="where-annotated-twice", - subject=where_source( - ' y: where(storage="gmem") = tf.add(x, x)\n' - ' y: where(storage="gmem")\n' - " return y" - ), - raises=VerifyError, - match="duplicate where annotation for Expr 'y'", - ), - ParseErrorCase( - id="where-on-a-subscript-lvalue", - subject=where_source( - """ value = tf.add(x, x) - value[0]: where(storage="gmem") - return value""" - ), - raises=VerifyError, - match="bound plain Name|annotation lvalue", - ), - ParseErrorCase( - id="where-on-a-whole-tuple-binding", - subject=where_source( - """ pair = tf.topk(x, k=4, axis=-1) - pair: where(storage="gmem") - return x""", - preamble="", - ), - raises=VerifyError, - match="tensor-valued", - ), - ParseErrorCase( - id="where-layout-extent-name-undefined", - subject=where_source(" y: where(layout=(_, N @ cta)) = tf.add(x, x)\n return y"), - raises=VerifyError, - match=r"where layout extent 'N' could not be resolved: undefined name 'N'", - ), - ParseErrorCase( - id="where-layout-extent-name-is-a-string", - subject=where_source( - " y: where(layout=(_, N @ cta)) = tf.add(x, x)\n return y", - preamble=_WHERE_PRELUDE + '\nN = "not-an-int"', - ), - raises=VerifyError, - match="must resolve to an int or DimVar", - ), -) - - -_MMA_SCOPES: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="mma-flat-32-lanes", - subject=_mma_scope(Topology("thread", 32), Layout(shape=(32,), strides=(1,))), - raises=VerifyError, - match=r"shape \(32,\) strides \(1,\).*thread-value decomposition", - ), - ParseErrorCase( - id="mma-wrong-lane-order", - subject=_mma_scope(Topology("thread", 32), Layout(shape=(4, 8), strides=(8, 1))), - raises=VerifyError, - match=r"shape \(4, 8\) strides \(8, 1\).*thread-value decomposition", - ), - ParseErrorCase( - id="mma-cta-not-thread", - subject=_mma_scope(Topology("cta", 32), Layout(shape=(4, 8), strides=(1, 4))), - raises=VerifyError, - match="the scope is a cta scope and the atom needs a thread one", - ), - ParseErrorCase( - id="mma-wrong-lane-count", - subject=_mma_scope(Topology("thread", 64), Layout(shape=(8, 8), strides=(1, 8))), - raises=VerifyError, - match=r"shape \(8, 8\) strides \(1, 8\).*64 lanes and the atom needs 32", - ), - ParseErrorCase( - id="mma-inconsistent-mesh", - subject=_mma_scope(Topology("thread", 64), Layout(shape=(4, 8), strides=(1, 4))), - raises=VerifyError, - match=r"thread\(64\) viewed as shape \(4, 8\).*64 lanes and the atom needs 32", - ), - ParseErrorCase( - id="mma-atom-outside-a-mesh-scope", - subject=_atom_outside_a_mesh_scope, - raises=VerifyError, - match="must be used inside a `with Mesh", - ), -) - - -_SYNC_SCOPES: tuple[ParseErrorCase, ...] = ( - ParseErrorCase( - id="sync-argument-is-not-a-mesh", - subject=_sync_argument_is_not_a_mesh, - raises=VerifyError, - match=r"T.sync expects a Mesh argument \(m or a slice m\[\.\.\.\]\), got Var", - ), - ParseErrorCase( - id="sync-with-no-enclosing-mesh", - subject=_sync_with_no_enclosing_mesh, - raises=VerifyError, - match="no enclosing mesh scope", - ), - ParseErrorCase( - id="sync-on-a-mesh-the-scope-does-not-bind", - subject=_sync_on_a_mesh_the_scope_does_not_bind, - raises=VerifyError, - match="no enclosing scope binds that mesh", - ), - ParseErrorCase( - id="sync-on-a-forged-subbox-exceeding-its-parent", - subject=_sync_on_a_forged_subbox_exceeding_its_parent, - raises=VerifyError, - match=r"T\.sync\(\(thread\(128\)\)\[1, 64\]\): that sub-box is not a slice", - ), - ParseErrorCase( - id="sync-on-a-forged-topology-mismatch", - subject=_sync_on_a_forged_topology_mismatch, - raises=VerifyError, - match=r"T\.sync\(\(warp\(4\)\)\[2, 32\]\): that sub-box is not a slice", - ), - ParseErrorCase( - id="sync-on-a-non-contiguous-slice", - subject=_sync_on_a_non_contiguous_slice, - raises=VerifyError, - match="contiguous", - ), - ParseErrorCase( - id="sync-on-a-cross-warp-unaligned-slice", - subject=_sync_on_a_cross_warp_unaligned_slice, - raises=VerifyError, - match=r"a cross-warp subset must be warp-aligned", - ), - ParseErrorCase( - id="classify-a-partial-cta-slice", - subject=_classify_a_partial_cta_slice, - raises=VerifyError, - match="partial grid", - ), - ParseErrorCase( - id="named-barriers-exhausted", - subject=_exhaust_the_named_barriers, - raises=ValueError, - match="too many distinct named barriers", - ), -) - - -ERROR_CASES: tuple[ParseErrorCase, ...] = ( - *_HIR_BODY_STATEMENTS, - *_OP_CALL_SURFACE, - *_DIM_OPERANDS, - *_CALL_BOUNDARY, - *_FUNCTION_BINDINGS, - *_SUBSCRIPTS, - *_GRID_LOOPS, - *_SHARD_SUGAR, - *_TOPOLOGY, - *_TUPLE_UNPACK, - *_CONSTRAINTS, - *_MMA_SCOPES, - *_SYNC_SCOPES, -) diff --git a/tests/parser/golden/hir_expressions.py b/tests/parser/golden/hir_expressions.py deleted file mode 100644 index d3eab786..00000000 --- a/tests/parser/golden/hir_expressions.py +++ /dev/null @@ -1,199 +0,0 @@ -from __future__ import annotations - -from tilefoundry.module import module -from tilefoundry import func -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 -from tilefoundry.ir.types.shard import ( - B, S, P, ComposedLayout, Layout, Mesh, ShardLayout, Topology, -) -from tilefoundry.ir.types.dim import DimVar, ceildiv - -CTX_LEN = DimVar("CTX_LEN", 1, 4097) - -@module(entry="dim_anchored_twice") -class HirExpressions: - @func - def dim_from_a_static_call( - x: Tensor[(CTX_LEN,), "bf16"] - ) -> Tensor[((128 * ((CTX_LEN - 1) // 128)) + 128,), "bf16"]: - v0 = zeros(type=Tensor[((128 * ((CTX_LEN - 1) // 128)) + 128,), "bf16"]) - return v0 - - @func - def cast_by_dtype_string( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "bf16"]: - v0 = cast(x, dtype="bf16") - return v0 - - @func - def reduce_by_kind_string( - x: Tensor[(8,), "f32"] - ) -> Tensor[(1,), "f32"]: - v0 = reduce(x, axes=(0,), keepdim=True, kind="sum") - return v0 - - @func - def unmaterialized_surface_storage( - x: Tensor[(8,), "f32", "umat"] - ) -> Tensor[(8,), "f32"]: - return x - - @func - def storage_without_a_layout_slot( - x: Tensor[(8,), "f32", "umat"] - ) -> Tensor[(8,), "f32"]: - return x - - @func - def literal_meets_bf16( - x: Tensor[(1, 8), "bf16"] - ) -> Tensor[(1, 8), "bf16"]: - v0 = 1e-06 - v1 = add(x, v0) - return v1 - - @func - def captured_float_meets_bf16( - x: Tensor[(1, 8), "bf16"] - ) -> Tensor[(1, 8), "bf16"]: - v0 = 1e-06 - v1 = add(x, v0) - return v1 - - @func - def compile_time_operands( - x: Tensor[(1, 2048), "bf16"] - ) -> Tensor[(1, 16, 128), "bf16"]: - v0 = 0.08838834764831845 - scaled = mul(x, v0) - v1 = 1e-06 - shifted = add(scaled, v1) - v2 = reshape(shifted, new_shape=(1, 16, 128)) - return v2 - - @func - def unpacked_compile_time_values( - x: Tensor[(1, 32, 128), "f32"] - ) -> Tensor[(1, 64, 64), "f32"]: - v0 = reshape(x, new_shape=(1, 64, 64)) - return v0 - - @func - def offsets_as_a_tuple_literal( - dst: Tensor[(2, 8, 4), "f32"], - upd: Tensor[(1, 3, 4), "f32"], - p: Tensor[(), "i32"] - ) -> Tensor[(2, 8, 4), "f32"]: - v0 = 1 - v1 = 0 - v3 = insert_slice(dst, upd, (1, p, 0)) - return v3 - - @func - def unpacked_multi_output( - x: Tensor[(1, 1536), "bf16"] - ) -> Tensor[(1, 1536), "fp8e4m3"]: - quant_out = quant(x, scheme="per_token_group", group=128, target_dtype="fp8e4m3") - x_fp8 = tuple_get_item(quant_out, index=0) - return x_fp8 - - @func - def index_drops_its_axis( - x: Tensor[(1, 4, 8), "f32"] - ) -> Tensor[(1, 4), "f32"]: - v0 = 0 - v1 = 0 - v2 = 3 - v4 = x[:, :, 3:4] - v5 = reshape(v4, new_shape=(1, 4)) - return v5 - - @func - def slice_keeps_its_axis( - x: Tensor[(1, 4, 8), "f32"] - ) -> Tensor[(1, 4, 1), "f32"]: - v0 = 0 - v1 = 0 - v2 = 3 - v4 = x[:, :, 3:4] - return v4 - - @func - def index_counted_from_the_end( - x: Tensor[(1, 4, 8), "f32"] - ) -> Tensor[(1, 4), "f32"]: - v0 = 0 - v1 = 0 - v2 = 7 - v4 = x[:, :, 7:8] - v5 = reshape(v4, new_shape=(1, 4)) - return v5 - - @func - def slice_strided_and_clamped( - x: Tensor[(1, 4, 8), "f32"] - ) -> Tensor[(1, 4, 3), "f32"]: - v0 = 0 - v1 = 0 - v2 = 1 - v4 = x[:, :, 1:10:3] - return v4 - - @func - def slice_to_symbolic_extents( - x: Tensor[(CTX_LEN, 128), "f32"] - ) -> Tensor[(CTX_LEN, 128), "f32"]: - v0 = 0 - v1 = 0 - v3 = x[:, :] - return v3 - - @func - def full_tile_window( - x: Tensor[(8, 4), "f32"], - seed: Tensor[(4, 4), "f32"] - ) -> Tensor[(4, 4), "f32"]: - out = add(seed, seed) - for row in tile(8, 4): - out_2 = x[row, :] - out_3 = add(out_2, seed) - out = out_3 - return out - - @func - def two_windows_a_fixed_distance_apart( - gu: Tensor[(3, 8), "f32"], - seed: Tensor[(3, 2), "f32"] - ) -> Tensor[(3, 2), "f32"]: - out = add(seed, seed) - for n in tile(4, 2): - out_2 = gu[:, n] - out_3 = gu[:, n + 4] - out_4 = mul(out_2, out_3) - out_5 = add(out, out_4) - out = out_5 - return out - - @func - def a_summed_offset_names_the_same_move( - gu: Tensor[(3, 8), "f32"], - seed: Tensor[(3, 2), "f32"] - ) -> Tensor[(3, 2), "f32"]: - out = add(seed, seed) - for n in tile(4, 2): - out_2 = gu[:, n] - out_3 = gu[:, n + 4] - out_4 = mul(out_2, out_3) - out_5 = add(out, out_4) - out = out_5 - return out - - @func - def dim_anchored_twice( - x: Tensor[(CTX_LEN,), "bf16"], - y: Tensor[(CTX_LEN + 1,), "bf16"] - ) -> Tensor[(CTX_LEN,), "bf16"]: - return x diff --git a/tests/parser/golden/hir_grid.py b/tests/parser/golden/hir_grid.py deleted file mode 100644 index eeec47bf..00000000 --- a/tests/parser/golden/hir_grid.py +++ /dev/null @@ -1,141 +0,0 @@ -from __future__ import annotations - -from tilefoundry.module import module -from tilefoundry import func -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 -from tilefoundry.ir.types.shard import ( - B, S, P, ComposedLayout, Layout, Mesh, ShardLayout, Topology, -) -from tilefoundry.ir.types.dim import DimVar, ceildiv - -seq_len = DimVar("seq_len", 1, 100) - -@module(entry="single_carry") -class HirGrid: - @func - def range_default_step( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - for i in range(8): - y = relu(x) - return () - - @func - def tile_extent_step( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - for i in range(0, 8, 2): - y = relu(x) - return () - - @func - def tile_dimvar_extent( - x: Tensor[(seq_len, 4), "f32"] - ) -> Tensor[(seq_len, 4), "f32"]: - for i in range(0, seq_len, 2): - y = relu(x) - return () - - @func - def range_dim_expr_extent( - x: Tensor[(seq_len, 4), "f32"] - ) -> Tensor[(seq_len, 4), "f32"]: - for i in range(seq_len // 2): - y = relu(x) - return () - - @func - def range_start_stop_step( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - for i in range(2, 8, 3): - y = relu(x) - return () - - @func - def inner_bindings_carry_nothing( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - for i in range(8): - t = relu(x) - z = add(t, x) - return () - - @func - def carry_reads_old_and_new( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - o = relu(x) - m = relu(x) - for i in range(8): - m_new = max(m, x) - correction = sub(m, m_new) - o_2 = add(o, correction) - o = o_2 - m = m_new - return o - - @func - def carry_initialized_from_a_parameter( - acc: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - for i in range(8): - acc_2 = add(acc, acc) - acc = acc_2 - return acc - - @func - def nested_for( - x: Tensor[(8, 4), "f32"] - ) -> Tensor[(8, 4), "f32"]: - o = relu(x) - for r in range(8): - for c in range(4): - o_2 = add(o, x) - o = o_2 - o = o - return o - - @func - def where_on_a_binding( - x: Tensor[(8, 16), "bf16"] - ) -> Tensor[(8, 16), "bf16"]: - y = add(x, x) - y: where(layout=(_, 16 @ cta), mesh=Mesh((Topology("cta", 8),), Layout((8,), (1,))), storage="gmem") - return y - - @func - def where_with_a_partial_value_state( - x: Tensor[(8, 16), "bf16"] - ) -> Tensor[(8, 16), "bf16"]: - y = add(x, x) - y: where(layout=((_, 16), {cta @ P("sum")})) - return y - - @func - def where_on_a_parameter( - x: Tensor[(8, 16), "bf16"] - ) -> Tensor[(8, 16), "bf16"]: - x: where(storage="smem") - return x - - @func - def where_on_a_bound_tuple_element( - x: Tensor[(8, 16), "bf16"] - ) -> Tensor[(8, 4), "i64"]: - values = topk(x, k=4, axis=-1, largest=True, sorted=True) - ids = tuple_get_item(values, index=1) - ids: where(storage="gmem") - return ids - - @func - def single_carry( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - o = relu(x) - for i in range(8): - o_2 = add(o, x) - o = o_2 - return o diff --git a/tests/parser/golden/hir_module.py b/tests/parser/golden/hir_module.py deleted file mode 100644 index 93877cd1..00000000 --- a/tests/parser/golden/hir_module.py +++ /dev/null @@ -1,174 +0,0 @@ -from __future__ import annotations - -from tilefoundry.module import module -from tilefoundry import func -from tilefoundry.target import CudaTarget -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.dsl import ConstTensor, Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 -from tilefoundry.ir.types.shard import ( - B, S, P, ComposedLayout, Layout, Mesh, ShardLayout, Topology, -) -from tilefoundry.ir.types.dim import DimVar, ceildiv - -N = DimVar("N", 1, 64) - -cta = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) -cta_2 = Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) - -@module(entry="calls_a_child", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("cta", 4),)) -class HirModule: - @module(entry="entry") - class leaf: - @func - def helper( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = mul(x, x) - return v0 - - @func - def entry( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = helper(x) - return v0 - - @module(entry="entry") - class first: - @func - def helper( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = mul(x, x) - return v0 - - @func - def entry( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = helper(x) - return v0 - - @module(entry="entry") - class second: - @func - def helper( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = mul(x, x) - return v0 - - @func - def entry( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = helper(x) - return v0 - - @module(entry="run") - class mlp: - @func - def run( - x: Tensor[(4, 8), "f32"], - w: ConstTensor[(8, 8), "f32"] - ) -> Tensor[(4, 8), "f32"]: - v0 = matmul(x, w) - return v0 - - @module(entry="mid") - class deep: - @module(entry="run") - class grand: - @func - def run( - x: Tensor[(8,), "f32"], - w: ConstTensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = mul(x, w) - return v0 - - @func - def mid( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = grand(x) - return v0 - - @module(entry="scale") - class variant_leaf: - @func - def scale( - x: Tensor[(N,), "f32"] - ) -> Tensor[(N,), "f32"]: - v0 = mul(x, x) - return v0 - - @func - def two_bindings_under_a_reshard( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - local = reshard(x, layout=ShardLayout( - layout=Layout((4, 2), None), - attrs=(S(0),), - mesh=cta, - ), storage=rmem) - v0 = first(local) - v1 = second(local) - v2 = add(v0, v1) - return v2 - - @func - def through_the_grandchild( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - local = reshard(x, layout=ShardLayout( - layout=Layout((4, 2), None), - attrs=(S(0),), - mesh=cta_2, - ), storage=gmem) - v0 = deep(local) - return v0 - - @func - def carries_activations_only( - x: Tensor[(4, 8), "f32"] - ) -> Tensor[(4, 8), "f32"]: - v0 = mlp(x) - return v0 - - @func - def converts_its_weight( - x: Tensor[(8,), "f32"], - w: ConstTensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = add(x, w) - return v0 - - @func - def dispatches_to_a_variant( - x: Tensor[(N,), "f32"] - ) -> Tensor[(N,), "f32"]: - pass - - @dispatches_to_a_variant.specialize(DimVarRangePat("N", 1, 64)) - def scaled_variant( - x: Tensor[(N,), "f32"] - ) -> Tensor[(N,), "f32"]: - v0 = variant_leaf(x) - return v0 - - @func - def uses_a_custom_op( - a: Tensor[(8,), "f32"], - b: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = custom_parse_addsq(a, b) - return v0 - - @func - def calls_a_child( - x: Tensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - v0 = leaf(x) - return v0 diff --git a/tests/parser/golden/hir_sharded.py b/tests/parser/golden/hir_sharded.py deleted file mode 100644 index 5eab93b9..00000000 --- a/tests/parser/golden/hir_sharded.py +++ /dev/null @@ -1,76 +0,0 @@ -from __future__ import annotations - -from tilefoundry.module import module -from tilefoundry import func -from tilefoundry.dsl.tf import * # noqa: F401, F403 -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.storage import gmem, host, rmem, smem, tmem # noqa: F401 -from tilefoundry.ir.types.shard import ( - B, S, P, ComposedLayout, Layout, Mesh, ShardLayout, Topology, -) -from tilefoundry.ir.types.dim import DimVar, ceildiv - -seq_len = DimVar("seq_len", 1, 4) - -gpu = Mesh((Topology("gpu", 8192),), Layout((32, 2, 8, 32), (2048, 1024, 32, 1)), names=('cluster', 'cta', 'warp', 'lane')) -thread = Mesh((Topology("thread", 192),), Layout((6, 32), (32, 1)), names=('w', 't')) -thread_2 = Mesh((Topology("thread", 128),), Layout((4, 32), (32, 1)), names=('y', 't')) -cta = Mesh((Topology("cta", 128),), Layout((128,), (1,)), names=('cta',)) -cta_2 = Mesh((Topology("cta", 8),), Layout((8,), (1,)), names=('w',)) -cta_3 = Mesh((Topology("cta", 8),), Layout((8,), (1,)), names=()) - -@module(entry="split_inline_and_default_broadcast", topologies=(Topology("cta", 8),)) -class HirSharded: - @func - def partial_brace_value_state( - a: Tensor[(64, 128), "bf16", ((32 @ gpu.cluster, 64), {gpu.warp @ P("sum")}), "smem"] - ) -> Tensor[(64, 128), "f32"]: - return a - - @func - def multi_axis_split_with_remainder( - a: Tensor[(1, 1536), "f32", (1, 6 @ thread.w, 32 @ thread.t, 8), "smem"] - ) -> Tensor[(1, 1536), "f32"]: - return a - - @func - def explicit_strides( - a: Tensor[(12, 4), "f32", (12 @ thread_2.y, 4), "smem"] - ) -> Tensor[(12, 4), "f32"]: - return a - - @func - def int_at_a_single_axis_mesh( - a: Tensor[(1, 8192), "f32", (1, 128 @ cta.cta, 64), "smem"] - ) -> Tensor[(1, 8192), "f32"]: - return a - - @func - def mesh_axis_as_a_position_coordinate( - ) -> Tensor[(), "i64"]: - v0 = arange(type=Tensor[(8,), "i64"], start=0, step=1) - v1 = reshard(v0, layout=ShardLayout( - layout=Layout((8,), (1,)), - attrs=(S(0),), - mesh=cta_2, - ), storage=rmem) - v2 = local(v1) - v3 = reshape(v2, new_shape=()) - return v3 - - @func - def reshard_with_a_dynamic_and_a_closure_axis( - q: Tensor[(1, seq_len, 32, 128), "bf16"] - ) -> Tensor[(1, seq_len, 32, 128), "bf16"]: - v0 = reshard(q, layout=ShardLayout( - layout=Layout((1, seq_len, 8, 4, 128), None), - attrs=(S(2),), - mesh=Mesh((Topology("cta", 8),), Layout((8,), (1,))), - )) - return v0 - - @func - def split_inline_and_default_broadcast( - a: Tensor[(32, 128), "bf16", (32 @ gpu.cluster, 2 @ gpu.cta, 64), "smem"] - ) -> Tensor[(32, 128), "f32"]: - return a diff --git a/tests/parser/programs.py b/tests/parser/programs.py deleted file mode 100644 index 137b63be..00000000 --- a/tests/parser/programs.py +++ /dev/null @@ -1,631 +0,0 @@ -"""The feature-dense programs the parser tests read. - -Each program is parsed once, at import, and its recorded golden is that program -printed back as DSL source. A parser feature belongs in the program whose -``features`` list names it, so adding one is adding a line to a body here rather -than a new file. What a golden cannot show — node identity, a target's concrete -``Op`` class, a layout the printer renders as the sugar it was written as — is a -named test in ``test_programs.py``. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -from tests.fixtures.logical.hir_composition import Expert -from tests.parser.error_cases import CTX_LEN, Callee -from tilefoundry import func, module, prim_func -from tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, T, Tensor, ceildiv, tf -from tilefoundry.dsl.tf import * # noqa: F401, F403 — bare op names used by the bodies -from tilefoundry.ir.core import Op -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.pattern import Tensor as TensorPattern -from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout, Mesh, P, ShardLayout, Split, Topology -from tilefoundry.ir.types.shard import Mesh as TirMesh -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.target import CpuTarget, CudaTarget -from tilefoundry.visitor_registry import register_typeinfer - -SEQ_LEN = DimVar("seq_len", 1, 100) -N_SCALED = DimVar("N", 1, 64) -SEQ_DYN = DimVar("seq_len", 1, 4) - -_EPS = 1e-6 -_NK, _KD, _VD, _NV = 16, 128, 64, 32 - - -@dataclass(frozen=True) -class _Cfg: - """A model config read at parse time, so its fields arrive as numbers.""" - - head_dim: int = 128 - rms_eps: float = 1e-6 - - -_CFG = _Cfg() -_HALF, _STEP, _ROWS = 4, 2, 3 - -CTA_MESH = Mesh((Topology("cta", 8),), Layout((8,), (1,))) -M_GPU = Mesh( - (Topology("gpu", 8192),), - Layout((32, 2, 8, 32), (2048, 1024, 32, 1)), - names=("cluster", "cta", "warp", "lane"), -) -M_MULTI = Mesh((Topology("thread", 6 * 32),), Layout((6, 32), (32, 1)), names=("w", "t")) -M_STRIDED = Mesh((Topology("thread", 4 * 32),), Layout((4, 32), (32, 1)), names=("y", "t")) -M_CTA = Mesh((Topology("cta", 128),), Layout((128,), (1,)), names=("cta",)) - - -@module(entry="dim_anchored_twice") -class HirExpressions: - """Expression, annotation, and subscript surface, in one Module. - - Dim arithmetic in a signature, the string dtype surface, a value literal's - unmaterialized storage, a declared constant, tuple literals in and out, and - every subscript form: an index that drops its axis, a slice that keeps it, - a negative index, a clamped stride, and a tile window that a compile-time - offset moves. - """ - - @func - def dim_anchored_twice( - x: Tensor[(CTX_LEN,), "bf16"], - y: Tensor[(CTX_LEN + 1,), "bf16"], - ) -> Tensor[(CTX_LEN,), "bf16"]: - return x - - @func - def dim_from_a_static_call(x: Tensor[(CTX_LEN,), "bf16"]): - return tf.zeros(Tensor[(ceildiv(CTX_LEN, 128) * 128,), "bf16"]) - - @func - def cast_by_dtype_string(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "bf16"]: - return cast(x, dtype="bf16") # noqa: F405 - - @func - def reduce_by_kind_string(x: Tensor[(8,), "f32"]) -> Tensor[(1,), "f32"]: - return reduce(x, axes=(0,), keepdim=True, kind="sum") # noqa: F405 - - @func - def unmaterialized_surface_storage( - x: Tensor[(8,), "f32", None, "umat"], - ) -> Tensor[(8,), "f32"]: - return x - - @func - def storage_without_a_layout_slot( - x: Tensor[(8,), "f32", "umat"], - ) -> Tensor[(8,), "f32"]: - return x - - @func - def literal_meets_bf16(x: Tensor[(1, 8), "bf16"]) -> Tensor[(1, 8), "bf16"]: - return add(x, 1e-6) # noqa: F405 - - @func - def captured_float_meets_bf16(x: Tensor[(1, 8), "bf16"]) -> Tensor[(1, 8), "bf16"]: - return add(x, _EPS) # noqa: F405 - - @func - def compile_time_operands(x: Tensor[(1, 2048), "bf16"]) -> Tensor[(1, 16, 128), "bf16"]: - key_dim = _NK * _KD - scaled = mul(x, _CFG.head_dim**-0.5) # noqa: F405 - shifted = add(scaled, _CFG.rms_eps) # noqa: F405 - return reshape(shifted, new_shape=(1, key_dim // _KD, _KD)) # noqa: F405 - - @func - def unpacked_compile_time_values( - x: Tensor[(1, 32, 128), "f32"], - ) -> Tensor[(1, 64, 64), "f32"]: - nv, kd, vd = _NV, _KD, _VD - return reshape(x, new_shape=(1, nv * kd // vd, vd)) # noqa: F405 - - @func - def offsets_as_a_tuple_literal( - dst: Tensor[(2, 8, 4), "f32"], - upd: Tensor[(1, 3, 4), "f32"], - p: Tensor[(), "i32"], - ) -> Tensor[(2, 8, 4), "f32"]: - return insert_slice(dst, upd, (1, p, 0)) # noqa: F405 - - @func - def unpacked_multi_output(x: Tensor[(1, 1536), "bf16"]) -> Tensor[(1, 1536), "fp8e4m3"]: - x_fp8, x_scale = quant(x) # noqa: F405, F841 - return x_fp8 - - @func - def index_drops_its_axis(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4), "f32"]: - return x[:, :, 3] - - @func - def slice_keeps_its_axis(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4, 1), "f32"]: - return x[:, :, 3:4] - - @func - def index_counted_from_the_end(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4), "f32"]: - return x[:, :, -1] - - @func - def slice_strided_and_clamped(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4, 3), "f32"]: - return x[:, :, 1:20:3] - - @func - def slice_to_symbolic_extents( - x: Tensor[(CTX_LEN, _KD), "f32"], - ) -> Tensor[(CTX_LEN, _KD), "f32"]: - return x[0:CTX_LEN, 0:_KD] - - @func - def full_tile_window(x: Tensor[(8, 4), "f32"], seed: Tensor[(4, 4), "f32"]): - out = add(seed, seed) # noqa: F405 - for row in tile(8, 4): # noqa: F405 - out = add(x[row, :], seed) # noqa: F405 - return out - - @func - def two_windows_a_fixed_distance_apart( - gu: Tensor[(_ROWS, 2 * _HALF), "f32"], - seed: Tensor[(_ROWS, _STEP), "f32"], - ): - out = add(seed, seed) # noqa: F405 - for n in tile(_HALF, _STEP): # noqa: F405 - out = add(out, mul(gu[:, n], gu[:, n + _HALF])) # noqa: F405 - return out - - @func - def a_summed_offset_names_the_same_move( - gu: Tensor[(_ROWS, 2 * _HALF), "f32"], - seed: Tensor[(_ROWS, _STEP), "f32"], - ): - out = add(seed, seed) # noqa: F405 - for n in tile(_HALF, _STEP): # noqa: F405 - out = add(out, mul(gu[:, n], gu[:, _HALF + 1 + n - 1])) # noqa: F405 - return out - - -@func -def returns_a_pair(a: Tensor[(4,), "f32"], b: Tensor[(4,), "f32"]): - return (add(a, b), mul(a, b)) # noqa: F405 - - -@func -def doubles_a_constant(w: ConstTensor[(8, 64), "f32"]) -> Tensor[(8, 64), "f32"]: - return add(w, w) # noqa: F405 - - -@module(entry="single_carry") -class HirGrid: - """Grid-region loops and the ``where`` annotations that ride on them. - - Both loop spellings over one domain, every way an extent is written, a - rebinding lifted to a phi carry, a carry initialized straight from a - parameter, nested loops, and layout / mesh / storage intent stated on a - binding, on a parameter, and on a bound tuple element. - """ - - @func - def range_default_step(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - for i in range(8): - y = relu(x) # noqa: F405, F841 - - @func - def tile_extent_step(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - for i in tile(8, 2): # noqa: F405 - y = relu(x) # noqa: F405, F841 - - @func - def tile_dimvar_extent(x: Tensor[(SEQ_LEN, 4), "f32"]) -> Tensor[(SEQ_LEN, 4), "f32"]: - for i in tile(SEQ_LEN, 2): # noqa: F405 - y = relu(x) # noqa: F405, F841 - - @func - def range_dim_expr_extent(x: Tensor[(SEQ_LEN, 4), "f32"]) -> Tensor[(SEQ_LEN, 4), "f32"]: - for i in range(SEQ_LEN // 2): - y = relu(x) # noqa: F405, F841 - - @func - def range_start_stop_step(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - for i in range(2, 8, 3): - y = relu(x) # noqa: F405, F841 - - @func - def single_carry(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - o = relu(x) # noqa: F405 - for i in range(8): - o = add(o, x) # noqa: F405 - return o - - @func - def inner_bindings_carry_nothing(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - for i in range(8): - t = relu(x) # noqa: F405 - z = add(t, x) # noqa: F405, F841 - - @func - def carry_reads_old_and_new(x: Tensor[(8,), "f32"]): - m = relu(x) # noqa: F405 - o = relu(x) # noqa: F405 - for i in range(8): - m_new = maximum(m, x) # noqa: F405 - correction = sub(m, m_new) # noqa: F405 - o = add(o, correction) # noqa: F405 - m = m_new - return o - - @func - def carry_initialized_from_a_parameter(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - acc = x - for i in range(8): - acc = add(acc, x) # noqa: F405 - return acc - - @func - def nested_for(x: Tensor[(8, 4), "f32"]) -> Tensor[(8, 4), "f32"]: - o = relu(x) # noqa: F405 - for r in range(8): - for c in range(4): - o = add(o, x) # noqa: F405 - return o - - @func - def where_on_a_binding(x: Tensor[(8, 16), "bf16"]) -> Tensor[(8, 16), "bf16"]: - y: where( # noqa: F405, F821 - layout=(_, 16 @ cta), mesh=CTA_MESH, storage="gmem" - ) = tf.add(x, x) - return y - - @func - def where_with_a_partial_value_state(x: Tensor[(8, 16), "bf16"]) -> Tensor[(8, 16), "bf16"]: - y: where(layout=((_, 16), {cta @ P("sum")})) = tf.add(x, x) # noqa: F405, F821 - return y - - @func - def where_on_a_parameter(x: Tensor[(8, 16), "bf16"]) -> Tensor[(8, 16), "bf16"]: - x: where(storage="smem") # noqa: F405, F821 - return x - - @func - def where_on_a_bound_tuple_element(x: Tensor[(8, 16), "bf16"]) -> Tensor[(8, 4), "i64"]: - values = tf.topk(x, k=4, axis=-1) - ids = values[1] - ids: where(storage="gmem") # noqa: F405, F821 - return ids - - -@module( - entry="split_inline_and_default_broadcast", - topologies=(Topology("cta", 8),), -) -class HirSharded: - """Mesh placement, shard-layout sugar, and topology declaration in one Module. - - Every annotation sugar form the parser canonicalises — an inline ``Split``, - the ``{...}`` value-state set, a multi-mesh-axis split with a remainder, - explicit strides, and the single-axis ``int @ mesh`` shorthand — plus a mesh - axis read as a position coordinate and a reshard whose split extent is - resolved through the closure. - """ - - @func - def split_inline_and_default_broadcast( - a: Tensor[(32, 128), bf16, (32 @ M_GPU.cluster, 2 @ M_GPU.cta, 64), "smem"], # noqa: F405 - ) -> Tensor[(32, 128), "f32"]: - return a - - @func - def partial_brace_value_state( - a: Tensor[(64, 128), "bf16", ((32 @ M_GPU.cluster, 64), {M_GPU.warp @ P("sum")}), "smem"], - ) -> Tensor[(64, 128), "f32"]: - return a - - @func - def multi_axis_split_with_remainder( - a: Tensor[(1, 1536), "f32", (1, 1536 @ (M_MULTI.w, M_MULTI.t)), "smem"], - ) -> Tensor[(1, 1536), "f32"]: - return a - - @func - def explicit_strides( - a: Tensor[(12, 4), "f32", ((12 @ M_STRIDED.y, 4), (4, 1)), "smem"], - ) -> Tensor[(12, 4), "f32"]: - return a - - @func - def int_at_a_single_axis_mesh( - a: Tensor[(1, 8192), "f32", (1, 8192 @ M_CTA), "smem"], - ) -> Tensor[(1, 8192), "f32"]: - return a - - @func - def mesh_axis_as_a_position_coordinate() -> Tensor[(), "i64"]: - with Mesh(("cta",), layout=(8,), names=("w",)) as cta: - return cta.w - - @func - def reshard_with_a_dynamic_and_a_closure_axis( - q: Tensor[(1, SEQ_DYN, 32, 128), "bf16"], - ) -> Tensor[(1, SEQ_DYN, 32, 128), "bf16"]: - with Mesh(("cta",), layout=Layout((8,), (1,))) as cta: - return reshard(q, layout=(1, SEQ_DYN, 32 @ cta, 128)) # noqa: F405 - - -@register_op(dialect="tf", category="custom", name="custom_parse_addsq") -class CustomParseAddSq(Op): - """Test-only custom op that squares the sum of its inputs.""" - - lhs = ParamDef(kind="input", pattern=TensorPattern) - rhs = ParamDef(kind="input", pattern=TensorPattern) - - -@register_typeinfer(CustomParseAddSq) -def _(call, ctx): - return ctx.type_of(call.args[0]) - - -custom_parse_addsq = CustomParseAddSq - - -@module(entry="scale") -class Scaled: - """A child Module whose entry is shaped by a ``DimVar``.""" - - @func - def scale(x: Tensor[(N_SCALED,), "f32"]) -> Tensor[(N_SCALED,), "f32"]: - return tf.mul(x, x) - - -@module(entry="run") -class Grand: - """The grandchild of ``Deep``: one activation and one declared constant.""" - - @func - def run(x: Tensor[(8,), "f32"], w: ConstTensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return tf.mul(x, w) - - -@module(entry="mid") -class Deep: - """A child Module that itself calls a child.""" - - grand = Grand - - @func - def mid(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return grand(x) # noqa: F821 — the class-body child binding - - -@module( - entry="calls_a_child", - target=CudaTarget("nvidia.h200_sxm"), - topologies=(Topology("cta", 4),), -) -class HirModule: - """A ``@module`` class body and every call it can make. - - A child Module bound by name, two bindings of one child rebuilt for a - resharded argument, a grandchild reached through the middle Module, a - child taking an activation while its constant stays declared, a weight - converter that calls a child of its own, a specialization variant, and a - registered custom op as a call target. - """ - - leaf = Callee - first = Callee - second = Callee - mlp = Expert - deep = Deep - variant_leaf = Scaled - - @func - def calls_a_child(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf(x) # noqa: F821 - - @func - def two_bindings_under_a_reshard(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - with Mesh(("cta",), layout=(4,), names=("tile",)) as cta: - local = tf.reshard(x, (8 @ cta.tile,), "rmem") - return tf.add(first(local), second(local)) # noqa: F821 - - @func - def through_the_grandchild(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - with Mesh(("cta",), layout=(4,), names=("tile",)) as cta: - local = tf.reshard(x, (8 @ cta.tile,), "gmem") - return deep(local) # noqa: F821 - - @func - def carries_activations_only(x: Tensor[(4, 8), "f32"]) -> Tensor[(4, 8), "f32"]: - return mlp(x) # noqa: F821 - - @func - def converts_its_weight( - x: Tensor[(8,), "f32"], w: ConstTensor[(8,), "f32"] - ) -> Tensor[(8,), "f32"]: - return tf.add(x, w) - - @converts_its_weight.converter("w") # noqa: F821 - def _convert_w(w: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return leaf(w) # noqa: F821 - - @func - def dispatches_to_a_variant(x: Tensor[(N_SCALED,), "f32"]) -> Tensor[(N_SCALED,), "f32"]: - pass - - @dispatches_to_a_variant.specialize(DimVarRangePat("N", 1, 64)) # noqa: F821 - def scaled_variant(x: Tensor[(N_SCALED,), "f32"]) -> Tensor[(N_SCALED,), "f32"]: - return variant_leaf(x) # noqa: F821 - - @func - def uses_a_custom_op(a: Tensor[(8,), "f32"], b: Tensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return custom_parse_addsq(a, b) - - -@dataclass(frozen=True) -class ParserProgram: - """One feature-dense program and the features it is here to carry.""" - - name: str - parsed: Module | Function - features: tuple[str, ...] - - -PROGRAMS: tuple[ParserProgram, ...] = ( - ParserProgram( - "hir_expressions", - HirExpressions, - ( - "dim arithmetic", - "dim from a static call", - "str dtype surface", - "str reduce-kind surface", - "storage with and without an empty layout slot", - "value literal dtype", - "compile-time operands", - "compile-time tuple unpack", - "tuple-literal op input", - "multi-output tuple unpack", - "subscript indexing and slicing", - "slice endpoints at symbolic extents", - "tile window", - "window move by a compile-time offset", - ), - ), - ParserProgram( - "hir_grid", - HirGrid, - ( - "range and tile share one loop domain", - "static, DimVar, and dim-expression extents", - "carry lifting", - "carry initialized from a parameter", - "nested for", - "where layout / mesh / storage constraint", - "where on a parameter and on a bound tuple element", - ), - ), - ParserProgram( - "hir_sharded", - HirSharded, - ( - "inline Split and default Broadcast", - "Partial value-state set", - "multi-mesh-axis split with a remainder", - "explicit strides", - "int @ single-axis mesh", - "mesh placement", - "mesh axis as a position coordinate", - "closure-resolved reshard split axis", - ), - ), - ParserProgram( - "hir_module", - HirModule, - ( - "@module class body", - "child Module bound by name", - "two bindings of one child rebuilt at a call site", - "grandchild call through the middle Module", - "activations carried, constants left declared", - "weight converter", - "specialization variant", - "custom op call target", - ), - ), -) - - -_TILE = 12 -_NT = DimVar("Ntile", 1, 64) - - -@prim_func(target=CudaTarget("nvidia.h200_sxm")) -def tir_dynamic_device(a: Tensor[(_NT, _TILE), "f32"]): - with TirMesh((Topology("cta", _NT),), Layout(shape=(_NT,), strides=(1,))) as cta: - view = T.tensor_view( - a, - layout=ShardLayout( - layout=Layout(shape=(_NT, _TILE), strides=(_TILE, 1)), - attrs=(Split(0),), - mesh=cta, - ), - ) - reg = T.alloc_tensor( - TensorType( - shape=(_NT, _TILE), - dtype=DType.f32, - layout=ShardLayout( - layout=Layout(shape=(_NT, _TILE), strides=(_TILE, 1)), - attrs=(Split(0),), - mesh=cta, - ), - storage=StorageKind.RMEM, - ) - ) - T.copy(view, reg) - - -@prim_func(target=CpuTarget()) -def tir_host_entry(a: Tensor[(_NT, _TILE), "f32"]): - launch(tir_dynamic_device, a, grid=(_NT, 1, 1), block=(1, 1, 1)) # noqa: F821 - - -@prim_func(target=CudaTarget("nvidia.h200_sxm")) -def tir_static_device(a: Tensor[(16, 8), "f32"]): - with TirMesh((Topology("thread", 8),), Layout(shape=(8,), strides=(1,))) as t: - view = T.tensor_view( - a, - layout=ShardLayout( - layout=Layout(shape=(16, 8), strides=(8, 1)), attrs=(Split(0),), mesh=t - ), - ) - reg = T.alloc_tensor( - TensorType( - shape=(16, 8), - dtype=DType.f32, - layout=ShardLayout( - layout=Layout(shape=(16, 8), strides=(8, 1)), attrs=(Split(0),), mesh=t - ), - storage=StorageKind.RMEM, - ) - ) - T.copy(view, reg) - - -@prim_func(target=CpuTarget()) -def tir_effect_form_selector(a: Tensor[(128,), "f32"], b: Tensor[(128,), "f32"]): - copy_(a, b) # noqa: F821 — resolved via dispatch.resolve_callable, not the closure - - -@prim_func(target=CudaTarget("nvidia.h200_sxm")) -def tir_param_layout_sugar(a: Tensor[(1, 8192), "f32", (1, 8192 @ M_CTA), "smem"]): - return - - -@prim_func(target=CudaTarget("nvidia.h200_sxm")) -def tir_sync_scopes(a: Tensor[(128,), "f32"]): # noqa: ARG001 - with TirMesh( - (Topology("thread", 128),), Layout(shape=(4, 32), strides=(32, 1)), ("w", "t") - ) as m: - T.sync(m) - T.sync(m[0, :]) - T.sync(m[1:3, :]) - - -@prim_func(target=CudaTarget("nvidia.h200_sxm")) -def tir_static_atom_bindings(a: Tensor[(16, 16), "bf16"]): # noqa: ARG001 - op = T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN - atom = T.cuda.mma.atom(op=op) # noqa: F841 - - -@prim_func(target=CudaTarget("nvidia.h200_sxm")) -def tir_atom_fragment_in_a_warp_scope(a: Tensor[(16, 16), "bf16"]): # noqa: ARG001 - atom = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) - with TirMesh( - (Topology("thread", 32),), Layout(shape=(4, 8), strides=(1, 4)), names=("warp", "lane") - ) as warp: # noqa: F841 - frag = T.alloc_tensor( # noqa: F841 - TensorType(shape=(16, 16), dtype=DType.bf16, layout=atom.A, storage=StorageKind.RMEM) - ) diff --git a/tests/parser/test_programs.py b/tests/parser/test_programs.py deleted file mode 100644 index 4255f207..00000000 --- a/tests/parser/test_programs.py +++ /dev/null @@ -1,637 +0,0 @@ -"""What each parser program must parse to, and what a golden cannot say. - -The golden is the whole assertion for anything the printed source shows. The -tests beside it are the ones a golden cannot carry: node identity, a target's -concrete ``Op`` class, a canonicalisation the printer renders back as the sugar -it was written as, and what a program does once it is evaluated or rebuilt. -""" - -from __future__ import annotations - -import ast -import tempfile -from dataclasses import replace -from pathlib import Path - -import pytest -import torch - -from tests._source import import_dsl -from tests.fixtures.placed.mma_tile import MatmulModule -from tests.parser.conftest import GoldenFiles -from tests.parser.error_cases import literal_reshard_func, mesh_dims_reshard_func -from tests.parser.programs import ( - M_CTA, - M_GPU, - M_MULTI, - M_STRIDED, - PROGRAMS, - SEQ_LEN, - CustomParseAddSq, - HirExpressions, - HirGrid, - HirModule, - HirSharded, - ParserProgram, - Scaled, - doubles_a_constant, - returns_a_pair, -) -from tilefoundry import func, module -from tilefoundry.analysis.preflight import infer_authored_types -from tilefoundry.analysis.walk import postorder -from tilefoundry.dsl import ConstTensor, DimVarRangePat, Tensor -from tilefoundry.dsl._stub_gen import regen_stubs -from tilefoundry.evaluator import evaluate -from tilefoundry.evaluator.dim import resolve_dim -from tilefoundry.inspection import as_script -from tilefoundry.ir.constraints import LayoutConstraint, constraint_metadata -from tilefoundry.ir.core import Call, Constant, Tuple, Var, get_metadata -from tilefoundry.ir.hir.function import Function, elaborate -from tilefoundry.ir.hir.grid_region import GridRegionExpr -from tilefoundry.ir.hir.specialize import origin_of, specialize_function -from tilefoundry.ir.hir.tensor.slice import Slice -from tilefoundry.ir.hir.verify import verify_function -from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type -from tilefoundry.ir.types.dim import DimAdd, DimVar -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Topology, make_mesh -from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, Split -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.parser import hir_parser -from tilefoundry.parser.base import _ModuleCallee -from tilefoundry.parser.sugar import parse_sugar -from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor - - -@pytest.mark.parametrize("program", PROGRAMS, ids=[program.name for program in PROGRAMS]) -def test_a_feature_dense_program_parses_to_its_golden( - program: ParserProgram, golden: GoldenFiles -) -> None: - """Parsing the program yields exactly the recorded IR, printed back as source.""" - golden.check(f"{program.name}.py", as_script(program.parsed)) - - -def test_converter_declared_before_variants_is_retained() -> None: - seq = DimVar("converter_before_variant", 1, 8) - - @module() - class ConverterThenVariants: - @func - def dispatch( - x: Tensor[(seq,), "f32"], w: ConstTensor[(8,), "f32"] - ) -> Tensor[(seq,), "f32"]: - pass - - @dispatch.converter("w") - def convert(w: ConstTensor[(8,), "f32"]) -> Tensor[(8,), "f32"]: - return w - - @dispatch.specialize(DimVarRangePat("converter_before_variant", 1, 4)) - def small(x: Tensor[(seq,), "f32"], w: ConstTensor[(8,), "f32"]) -> Tensor[(seq,), "f32"]: - return x - - @dispatch.specialize(DimVarRangePat("converter_before_variant", 4, 8)) - def large(x: Tensor[(seq,), "f32"], w: ConstTensor[(8,), "f32"]) -> Tensor[(seq,), "f32"]: - return x - - dispatch = ConverterThenVariants.lookup("dispatch") - assert len(dispatch.converters) == 1 - assert len(dispatch.variants) == 2 - - -def test_annotation_sugar_lands_on_the_hand_written_layout() -> None: - """The golden prints each annotation back as the sugar it was written as. - - What it therefore cannot show is what the sugar canonicalises *to*: the - auto-filled C-order strides, and the ``Broadcast`` every mesh axis named in - no ``Split`` falls back to. - """ - expected = { - "split_inline_and_default_broadcast": ShardLayout( - layout=Layout((32, 2, 64), (128, 64, 1)), - attrs=(Split(0), Split(1), Broadcast(), Broadcast()), - mesh=M_GPU, - ), - "partial_brace_value_state": ShardLayout( - layout=Layout((32, 64), (64, 1)), - attrs=(Split(0), Broadcast(), Partial("sum"), Broadcast()), - mesh=M_GPU, - ), - "multi_axis_split_with_remainder": ShardLayout( - layout=Layout((1, 6, 32, 8), (1536, 256, 8, 1)), - attrs=(Split(1), Split(2)), - mesh=M_MULTI, - ), - "explicit_strides": ShardLayout( - layout=Layout((12, 4), (4, 1)), attrs=(Split(0), Broadcast()), mesh=M_STRIDED - ), - "int_at_a_single_axis_mesh": ShardLayout( - layout=Layout((1, 128, 64), (8192, 64, 1)), attrs=(Split(1),), mesh=M_CTA - ), - } - for name, layout in expected.items(): - parsed = HirSharded.lookup(name).params[0].type - assert parsed.storage is StorageKind.SMEM, name - assert parsed.layout == layout, name - - -def test_a_reshard_split_axis_resolved_through_the_closure() -> None: - """``32 @ cta`` keeps the un-factorised shape and defers strides to typeinfer.""" - body = HirSharded.lookup("reshard_with_a_dynamic_and_a_closure_axis").body - assert isinstance(body, Call) - assert body.type.shape == (1, DimVar("seq_len", 1, 4), 32, 128) - assert any(isinstance(attr, Split) for attr in body.target.layout.attrs) - assert body.target.layout.layout.strides is None - - -def test_a_symbolic_extent_on_both_sides_has_local_size_one() -> None: - """Shard-layout sugar on a bare AST node, with no function around it.""" - dyn = DimVar("seq_len", 1, 4) - cta = Mesh((Topology("cta", dyn),), Layout((dyn,), (1,)), names=("cta",)) - node = ast.parse("(1, S @ cta, 32, 128)", mode="eval").body - - actual = parse_sugar( - node, - ShardLayout, - mesh_resolver=lambda n: cta if n == "cta" else None, - closure={"S": dyn}, - ) - - assert actual.layout.shape == (1, dyn, 32, 128) - assert actual.attrs == (Split(1),) - - -def test_the_two_grid_spellings_share_one_loop_domain() -> None: - """The domain fields and the induction var's storage are not printed. - - ``range`` and ``tile`` differ only in their step, which the golden shows as - two loop headers rather than as one domain; and nothing in the source says - the induction var is unmaterialized. - """ - default_step = HirGrid.lookup("range_default_step").body - assert isinstance(default_step, GridRegionExpr) - assert (default_step.start, default_step.extent, default_step.step) == (0, 8, 1) - assert default_step.carried_args == () - assert default_step.init_args == () - assert default_step.yield_values == () - - with_step = HirGrid.lookup("tile_extent_step").body - assert repr(default_step) == repr(replace(with_step, step=1)) - - ranged = HirGrid.lookup("range_start_stop_step").body - assert (ranged.start, ranged.extent, ranged.step) == (2, 8, 3) - assert isinstance(ranged.induction_var, Var) - assert ranged.induction_var.type.storage is StorageKind.UMAT - - assert isinstance(HirGrid.lookup("range_dim_expr_extent").body.extent, Call) - assert HirGrid.lookup("tile_dimvar_extent").body.extent == SEQ_LEN - assert HirGrid.lookup("inner_bindings_carry_nothing").body.carried_args == () - - -def test_the_generated_tile_stub_requires_its_window_step() -> None: - """``tile`` is positional-only at the IR level, so its stub declares no default. - - The stub is generated rather than parsed, so no program can carry this. - """ - with tempfile.TemporaryDirectory() as directory: - stub = ast.parse(regen_stubs(Path(directory))["tf"].read_text()) - - tile_def = next( - node for node in stub.body if isinstance(node, ast.FunctionDef) and node.name == "tile" - ) - assert [arg.arg for arg in tile_def.args.args] == ["extent", "step"] - assert tile_def.args.defaults == [] - - -def test_a_carry_reuses_the_nodes_it_was_built_from() -> None: - """Which node a carry *is* — not what it prints as. - - A rebinding yields the very node its RHS bound, the reader of the old value - still points at the phi, and a carry initialized from a bare name adds no - call of its own. - """ - grid = HirGrid.lookup("carry_reads_old_and_new").body.args[0] - assert isinstance(grid, GridRegionExpr) - carried = {value.name: value for value in grid.carried_args} - yielded = dict(zip((value.name for value in grid.carried_args), grid.yield_values)) - correction = yielded["o"].args[1] - - assert grid.body is yielded["m"] - assert correction.args[0] is carried["m"] - assert correction.args[1] is yielded["m"] - - from_parameter = HirGrid.lookup("carry_initialized_from_a_parameter") - initialized = from_parameter.body - assert initialized.init_args[0] is from_parameter.params[0] - assert [expr for expr in postorder(initialized) if isinstance(expr, Call)] == [initialized.body] - - outer = HirGrid.lookup("nested_for").body - assert [v.name for v in outer.carried_args] == ["o"] - assert [v.name for v in outer.yield_values[0].carried_args] == ["o"] - - -def test_a_where_annotation_attaches_to_the_existing_ssa_node() -> None: - """A constraint is metadata on a node, so nothing about it is a printed line. - - Both readers of the annotated binding reach the same node, and the parser - attaches once rather than rebuilding. - """ - layout = constraint_metadata(HirGrid.lookup("where_on_a_binding").body).constraints[0] - assert isinstance(layout, LayoutConstraint) - assert repr(layout.layout.shape[0]) == "_" - assert layout.layout.shape[1] == 16 - assert layout.bindings == (("cta", Split(1)),) - - partial = constraint_metadata( - HirGrid.lookup("where_with_a_partial_value_state").body - ).constraints[0] - assert partial.bindings == (("cta", Partial("sum")),) - - for name in ("where_on_a_binding", "where_with_a_partial_value_state"): - verify_function(HirGrid.lookup(name)) - - -def test_a_where_layout_extent_is_read_out_of_the_globals(monkeypatch) -> None: - """A named extent resolves through the function's globals, as an int or a DimVar. - - The subject has to be parsed inside the test to count the attachments, so - this one stays on DSL source rather than reading a program. - """ - attached = [] - original = hir_parser._HirBodyVisitor._attach_metadata - - def capture(expr, metadata): - attached.append(expr) - original(expr, metadata) - - monkeypatch.setattr(hir_parser._HirBodyVisitor, "_attach_metadata", staticmethod(capture)) - - preamble = ( - "from tilefoundry.ir.types.shard import Layout, Mesh, Topology\n\n" - 'cta_mesh = Mesh((Topology("cta", 8),), Layout((8,), (1,)))\n' - ) - source = ( - "from __future__ import annotations\n" - "from tilefoundry import func\n" - "from tilefoundry.dsl import Tensor, tf\n\n" - f"{preamble}N = 16\n\n" - "@func\n" - 'def candidate(x: Tensor[(8, 16), "bf16"]) -> Tensor[(8, 16), "bf16"]:\n' - " y: where(layout=(_, N @ cta)) = tf.add(x, x)\n" - " return y\n" - ) - function = import_dsl(source) - - assert len(attached) == 1 - assert function.body is attached[0] - assert constraint_metadata(attached[0]).constraints[0].layout.shape[1] == 16 - - as_dim_var = import_dsl( - source.replace("N = 16", 'N = DimVar("S", 1, 128)').replace( - "from tilefoundry.dsl import Tensor, tf\n", - "from tilefoundry.dsl import Tensor, tf\nfrom tilefoundry.ir.types.dim import DimVar\n", - ) - ) - extent = constraint_metadata(as_dim_var.body).constraints[0].layout.shape[1] - assert isinstance(extent, DimVar) and extent.name == "S" - - -def test_umat_is_an_accepted_surface_storage() -> None: - """An explicit ``umat`` annotation preserves unresolved residency. - - Both spellings live in ``HirExpressions``, whose module declares meshes — - which is the point: the third annotation slot holds a storage name or an - empty layout, and a mesh being in scope does not make either one layout - sugar. The golden shows the storage; that it resolved to ``UMAT`` rather - than defaulting to ``GMEM`` is what is asserted here. - """ - for name in ("unmaterialized_surface_storage", "storage_without_a_layout_slot"): - parsed = HirExpressions.lookup(name) - assert parsed.params[0].type.storage is StorageKind.UMAT, name - assert parsed.params[0].type.layout is None, name - assert parsed.return_type.storage is StorageKind.GMEM, name - - -def test_a_value_literal_takes_the_dtype_of_the_operand_it_meets() -> None: - """The golden prints ``1e-06`` without saying what dtype it carries.""" - for name in ("literal_meets_bf16", "captured_float_meets_bf16"): - fn = HirExpressions.lookup(name) - assert fn.body.args[1].type.dtype == DType.bf16, name - assert fn.body.type.dtype == DType.bf16, name - - -def test_a_dim_expression_stays_resolvable_arithmetic() -> None: - """The golden prints the dim expression; that it still evaluates is separate.""" - padded = HirExpressions.lookup("dim_from_a_static_call").body.target.type.shape[0] - assert resolve_dim(padded, {"CTX_LEN": 128}) == 128 - assert resolve_dim(padded, {"CTX_LEN": 130}) == 256 - - anchored = HirExpressions.lookup("dim_anchored_twice") - assert isinstance(anchored.params[1].type.shape[0], Call) - verify_function(anchored) - - -def test_a_literal_tuple_return_folds_to_a_tuple_typed_body() -> None: - """The caller's golden names the callee; the callee's own body is not in it.""" - assert isinstance(returns_a_pair.body, Tuple) - assert len(returns_a_pair.body.elements) == 2 - assert isinstance(returns_a_pair.return_type, TupleType) - assert all(field.dtype == DType.f32 for field in returns_a_pair.return_type.fields) - - -def test_a_custom_op_call_resolves_to_the_registered_op_class() -> None: - """``as_script`` prints the call by name; the target's class is not visible there.""" - body = HirModule.lookup("uses_a_custom_op").body - assert isinstance(body, Call) - assert isinstance(body.target, CustomParseAddSq) - - -def test_each_child_binding_calls_its_own_attached_entry() -> None: - """Two bindings of one Module print identically and are still different objects.""" - children = {child.name: child for child in HirModule.modules} - assert set(children) == {"leaf", "first", "second", "mlp", "deep", "variant_leaf"} - - call = HirModule.lookup("calls_a_child").body - assert isinstance(call.target, Function) - assert call.target is children["leaf"].entry_function() - assert get_metadata(call, _ModuleCallee) is None - - left, right = HirModule.lookup("two_bindings_under_a_reshard").body.args - assert left.target is not right.target - assert origin_of(left.target) is children["first"].entry_function() - assert origin_of(right.target) is children["second"].entry_function() - - (variant,) = HirModule.lookup("dispatches_to_a_variant").variants - assert variant.body.target is children["variant_leaf"].entry_function() - assert variant.body.target is not Scaled.entry_function() - - ((weight, converter),) = HirModule.lookup("converts_its_weight").converters - assert weight == "w" - assert converter.body.target is children["leaf"].entry_function() - - -def test_a_child_call_carries_activations_and_leaves_the_constants_declared() -> None: - """Which parameters are constants, and whose weights they are, is not printed.""" - (child,) = [c for c in HirModule.modules if c.name == "mlp"] - call = HirModule.lookup("carries_activations_only").body - assert len(call.args) == 1 - assert [(p.name, p.is_const) for p in call.target.params] == [("x", False), ("w", True)] - assert call.target.params[1].type == child.weights["w"] - assert set(child.weights) == {"w"} - assert child.weights["w"].shape == (8, 8) - assert HirModule.weights["w"].shape == (8,) - - grandchild = next(c for c in HirModule.modules if c.name == "deep").modules[0] - inner = HirModule.lookup("through_the_grandchild").body.target.body - assert [p.is_const for p in inner.target.params] == [False, True] - assert origin_of(inner.target) is grandchild.entry_function() - - -def test_a_rebuilt_child_target_is_owned_only_by_that_child() -> None: - """Ownership after a rebuild, which no printed program states.""" - one, two = [c for c in HirModule.modules if c.name in ("first", "second")][:2] - sized = specialize_function(Scaled.entry_function(), {"N": 8}) - resharded = elaborate( - sized, (make_shard_tensor_type((8,), mesh=make_mesh((4,)), attrs=(Split(0),)),) - ) - assert resharded is not sized - assert not one.owns(resharded, derived=True) - assert not two.owns(resharded, derived=True) - - -def test_a_wildcard_chain_reelaborates_the_whole_nested_call() -> None: - """Re-elaboration builds new Functions; the program it started from is unchanged.""" - x_split = make_shard_tensor_type((8, 64), mesh=make_mesh((4,)), attrs=(Split(0),)) - rebuilt = elaborate(doubles_a_constant, (x_split,)) - assert rebuilt is not doubles_a_constant - assert rebuilt.params[0].is_const is True - assert rebuilt.params[0].type == x_split - - -def test_the_printer_falls_back_to_verbose_when_a_mesh_has_no_names() -> None: - """A mesh without ``names=`` cannot use ``@`` sugar, so the layout prints in full.""" - src = as_script(MatmulModule.entry_function()) - assert "@" not in src.split("@func")[1].split("def ")[0] - assert "ShardLayout(" in src - - -def test_both_authoring_spellings_print_to_one_canonical_program() -> None: - """A string and the descriptor it names are one surface over one IR.""" - header = ( - "from tilefoundry import func\n" - "from tilefoundry.dsl import Tensor\n" - "from tilefoundry.dsl.tf import *\n" - ) - string_form = header + ( - "@func\n" - 'def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "bf16"]:\n' - ' return cast(x, dtype="bf16")\n' - ) - descriptor_form = header + ( - "from tilefoundry.ir.types import DType\n" - "@func\n" - 'def f(x: Tensor[(8,), "f32"]) -> Tensor[(8,), "bf16"]:\n' - " return cast(x, dtype=DType.bf16)\n" - ) - printed = as_script(import_dsl(string_form)) - assert printed == as_script(import_dsl(descriptor_form)) - assert 'dtype="bf16"' in printed - - -def test_a_program_still_evaluates_to_what_torch_would_give() -> None: - """Subscript semantics are a runtime contract; the golden only shows the IR.""" - x = torch.arange(32, dtype=torch.float32).reshape(1, 4, 8) - for name, expected in ( - ("index_drops_its_axis", x[:, :, 3]), - ("slice_keeps_its_axis", x[:, :, 3:4]), - ("index_counted_from_the_end", x[:, :, -1]), - ("slice_strided_and_clamped", x[:, :, 1:20:3]), - ): - torch.testing.assert_close( - evaluate(HirExpressions.lookup(name), x, device="cpu"), expected, msg=name - ) - - -def test_symbolic_slice_endpoints_preserve_shape_and_bind_across_a_call() -> None: - """Equivalent full windows retain the authored dimension at a call boundary.""" - prelude = ( - "from tilefoundry import func\n" - "from tilefoundry.dsl.tf import *\n" - "from tilefoundry.dsl import DimVar, Tensor\n" - ) - full_window = import_dsl( - prelude + '\nCTX_LEN = DimVar("CTX_LEN", 1, 4097)\n' - '@func\ndef f(x: Tensor[(CTX_LEN, 128), "f32"]) ' - '-> Tensor[(CTX_LEN, 128), "f32"]:\n' - " return x[:, 0:128]\n" - ) - explicit_window = HirExpressions.lookup("slice_to_symbolic_extents") - assert isinstance(explicit_window.body, Call) - assert isinstance(explicit_window.body.target, Slice) - assert explicit_window.body.target.sizes == full_window.body.target.sizes - assert explicit_window.body.target.strides == full_window.body.target.strides - assert explicit_window.body.type == full_window.body.type - assert explicit_window.body.args[1] == full_window.body.args[1] - authored_extent = explicit_window.params[0].type.shape[0] - assert explicit_window.body.target.sizes[0] is authored_extent - assert explicit_window.body.type.shape[0] is authored_extent - - param = Var(type=explicit_window.return_type, name="window") - consumer = Function.build( - name="consume_window", - params=(param,), - body=param, - return_type=explicit_window.return_type, - ) - call = Call(type=consumer.return_type, target=consumer, args=(explicit_window.body,)) - assert TypeInferVisitor(TypeInferContext()).visit(call) == consumer.return_type - - dimension_start = import_dsl( - prelude + '\nS = DimVar("m2_start_seq", 1, 4097)\n' - '@func\ndef f(x: Tensor[(S + 8, 128), "f32"]) -> Tensor[(8, 128), "f32"]:\n' - " return x[S:S + 8, 0:128]\n" - ) - start = dimension_start.body.args[1].elements[0] - assert isinstance(start, Call) and isinstance(start.target, DimAdd) - assert any(start.args[0] is arg for arg in dimension_start.params[0].type.shape[0].args) - assert isinstance(start.args[1], Constant) and start.args[1].value == 0 - - -def test_a_packed_cache_can_keep_a_symbolic_capacity_axis() -> None: - """A layer index and full symbolic windows can name every packed-cache axis.""" - packed = import_dsl( - "from tilefoundry import func, module\n" - "from tilefoundry.dsl.tf import *\n" - "from tilefoundry.dsl import DimVar, Tensor\n" - '\nCAP = DimVar("m2_capacity", 1, 4097)\n' - '@module(entry="run")\n' - "class PackedCache:\n" - " @func\n" - ' def run(kc: Tensor[(4, CAP, 8, 16), "f32"], ' - 'seed: Tensor[(CAP, 8, 16), "f32"]) -> Tensor[(CAP, 8, 16), "f32"]:\n' - " out = relu(seed)\n" - " for i in range(4):\n" - " out = add(out, kc[i, 0:CAP, 0:8, 0:16])\n" - " return out\n" - ) - entry = packed.entry_function() - infer_authored_types((entry,), packed) - capacity = entry.params[0].type.shape[1] - - assert entry.body.type.shape == (capacity, 8, 16) - - -def _fused_reference(gu, seed): - out = seed * 2 - for lo in range(0, 4, 2): - out = out + gu[:, lo : lo + 2] * gu[:, lo + 4 : lo + 4 + 2] - return out - - -def test_a_moved_tile_window_evaluates_where_it_says_it_reads() -> None: - """Both spellings of the same move land the same data, and re-import unchanged.""" - gu = torch.arange(3 * 8, dtype=torch.float32).reshape(3, 8) - seed = torch.ones((3, 2), dtype=torch.float32) - expected = _fused_reference(gu, seed) - - for name in ("two_windows_a_fixed_distance_apart", "a_summed_offset_names_the_same_move"): - fn = HirExpressions.lookup(name) - torch.testing.assert_close(evaluate(fn, gu, seed, device="cpu"), expected, msg=name) - - script = as_script(HirExpressions.lookup("two_windows_a_fixed_distance_apart")) - assert "gu[:, n + 4]" in script, script - torch.testing.assert_close(evaluate(import_dsl(script), gu, seed, device="cpu"), expected) - - -def test_a_range_scalar_and_a_runtime_endpoint_drive_a_slice_window() -> None: - """A window whose start is only known at run time keeps its static size.""" - prelude = ( - "from tilefoundry import func\n" - "from tilefoundry.dsl.tf import *\n" - "from tilefoundry.dsl import Tensor\n" - ) - runtime_start = import_dsl( - prelude + "\nfrom tilefoundry.ir.types.shard import Layout\n" - "plain_layout = Layout((8, 4), (4, 1))\n" - '\n@func\ndef f(x: Tensor[(8, 4), "f32", plain_layout], ' - 'start: Tensor[(), "i64"]) -> Tensor[(4, 4), "f32"]:\n' - " return x[start:start + 4, :]\n" - ) - assert runtime_start.body.target.sizes == (4, 4) - assert runtime_start.body.args[1].elements[0] is runtime_start.params[1] - assert runtime_start.body.type.layout is None - - shifted_start = import_dsl( - prelude + '\n@func\ndef f(x: Tensor[(16, 4), "f32"], ' - 'start: Tensor[(), "i64"]) -> Tensor[(8, 4), "f32"]:\n' - " return x[start + 1:start + 9, :]\n" - ) - assert shifted_start.body.target.sizes == (8, 4) - assert shifted_start.body.type.shape == (8, 4) - - x = torch.arange(32, dtype=torch.float32).reshape(8, 4) - torch.testing.assert_close( - evaluate(runtime_start, x, torch.tensor(2, dtype=torch.int64), device="cpu"), x[2:6, :] - ) - - sharded = import_dsl( - prelude - + "\nfrom tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology\n" - '\n@func\ndef f(x: Tensor[(8, 4), "f32", ShardLayout(' - "layout=Layout((8, 2, 2), (4, 2, 1)), attrs=(Split(1),), " - 'mesh=Mesh((Topology("gpu", 2),), Layout((2,), (1,)), names=("g",)))], ' - 'start: Tensor[(), "i64"]) -> Tensor[(2, 4), "f32"]:\n' - " return x[start:start + 2, :]\n" - ) - assert isinstance(sharded.body.type.layout, ShardLayout) - assert sharded.body.type.layout.attrs == sharded.body.args[0].type.layout.attrs - assert sharded.body.type.layout.layout.shape == (2, 2, 2) - - -def test_a_compile_time_list_is_indexed_where_it_is_written() -> None: - """A comprehension and a plain list literal both bind a Python list of Exprs. - - Indexing either picks an expression rather than emitting an op, so neither - list survives into the program a golden could show. - """ - prelude = ( - "from tilefoundry import func\n" - "from tilefoundry.dsl.tf import *\n" - "from tilefoundry.dsl import Tensor\n" - '\n@func\ndef f(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4, 1), "f32"]:\n' - ) - x = torch.arange(32, dtype=torch.float32).reshape(1, 4, 8) - - comprehension = import_dsl( - prelude + " taps = [x[:, :, j:j + 1] for j in range(4)]\n" - " return add(taps[0], taps[-1])\n" - ) - torch.testing.assert_close( - evaluate(comprehension, x, device="cpu"), x[:, :, 0:1] + x[:, :, 3:4] - ) - - literal = import_dsl( - prelude + " ends = [x[:, :, 0:1], x[:, :, 7:8]]\n return add(ends[0], ends[1])\n" - ) - torch.testing.assert_close(evaluate(literal, x, device="cpu"), x[:, :, 0:1] + x[:, :, 7:8]) - - -def test_a_closure_int_mesh_dim_resolves_like_the_literal() -> None: - """A closure int in a mesh-shape sugar prints back to the literal form. - - The parser must resolve the ``ast.Name`` rather than reject it, and the - trailing reshard names no mesh axis at all, so it resolves its mesh from the - enclosing scope. - """ - assert as_script(mesh_dims_reshard_func(4, 32)) == as_script(literal_reshard_func()) - - -def test_a_tile_window_survives_the_print_import_trip() -> None: - """The canonical source of a windowed loop evaluates to what it started as.""" - x = torch.arange(32, dtype=torch.float32).reshape(8, 4) - seed = torch.ones((4, 4), dtype=torch.float32) - fn = HirExpressions.lookup("full_tile_window") - expected = evaluate(fn, x, seed, device="cpu") - - torch.testing.assert_close(evaluate(import_dsl(as_script(fn)), x, seed, device="cpu"), expected) diff --git a/tests/parser/test_refused_programs.py b/tests/parser/test_refused_programs.py deleted file mode 100644 index 7a658b6a..00000000 --- a/tests/parser/test_refused_programs.py +++ /dev/null @@ -1,17 +0,0 @@ -"""The one entry point for every subject ``tests/parser`` refuses. - -The table is in ``error_cases.py``; this file only runs it. A refusal that is not -a row here is a refusal nothing pins. -""" - -from __future__ import annotations - -import pytest - -from tests.parser.error_cases import ERROR_CASES, ParseErrorCase, run_parse_error_case - - -@pytest.mark.parametrize("case", ERROR_CASES, ids=[case.id for case in ERROR_CASES]) -def test_a_refused_program_raises_its_diagnostic(case: ParseErrorCase) -> None: - """Every refused subject named in the table is refused, saying why.""" - run_parse_error_case(case) diff --git a/tests/parser/test_roundtrip.py b/tests/parser/test_roundtrip.py new file mode 100644 index 00000000..ed2f0572 --- /dev/null +++ b/tests/parser/test_roundtrip.py @@ -0,0 +1,30 @@ +"""Parser-owned checks for canonical grid-loop bindings.""" + +from tests.fixtures.placed.gqa_decode import GqaOnline +from tilefoundry.inspection import as_script +from tilefoundry.ir.hir.specialize import specialize_concretely + + +def test_gqa_correction_reads_the_old_carry_and_unique_yield() -> None: + function = specialize_concretely(GqaOnline.entry_function(), {"ctx_len": 8}) + printed = as_script(function) + + assert printed.count(" m_new = max(m, score)") == 1 + assert " = sub(m, m_new)" in printed + lines = printed.splitlines() + start = lines.index(" for i in range(8):") + end = next( + index + for index in range(start + 1, len(lines)) + if lines[index].startswith(" ") + and not lines[index].startswith(" ") + ) + yielded = lines[end - 3 : end] + assert [line.split(" = ", 1)[0] for line in yielded] == [ + " l", + " o", + " m", + ] + assert yielded[-1] == " m = m_new" + assert len({line.split(" = ", 1)[1] for line in yielded}) == 3 + assert lines[end] == ' k_n = cast(k_new, dtype="f32")' diff --git a/tests/parser/test_tir_programs.py b/tests/parser/test_tir_programs.py deleted file mode 100644 index d20bd8b5..00000000 --- a/tests/parser/test_tir_programs.py +++ /dev/null @@ -1,173 +0,0 @@ -"""What the TIR half of the parser must build, asserted on the parsed nodes. - -There is no golden here: ``tilefoundry.inspection`` prints HIR and nothing -prints a ``PrimFunction``, so these read the programs in ``programs.py`` -directly. What is left out of the programs is left out on purpose — a -hand-forged node fed to the verifier is not something the parser produced, and -those are rows in ``error_cases.py``. -""" - -from __future__ import annotations - -import tilefoundry.codegen.cuda # noqa: F401 — trigger emitter autodiscovery -from tests.parser.error_cases import scoped_sync, thread_mesh -from tests.parser.programs import ( - M_CTA, - tir_atom_fragment_in_a_warp_scope, - tir_dynamic_device, - tir_effect_form_selector, - tir_host_entry, - tir_param_layout_sugar, - tir_static_atom_bindings, - tir_static_device, - tir_sync_scopes, -) -from tilefoundry.codegen.cuda.context import CodegenContext -from tilefoundry.dsl import T -from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom -from tilefoundry.ir.tir.memory.copy import Copy -from tilefoundry.ir.tir.stmts import Evaluate, LetStmt, MeshScope, Sequential -from tilefoundry.ir.tir.sync import Sync -from tilefoundry.ir.tir.verify import verify_prim_function -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.shard import Layout -from tilefoundry.ir.types.shard.layout import ComposedLayout -from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split -from tilefoundry.ir.types.storage import StorageKind - -_ATOM = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) - - -def _shape_scalars(prim_fn) -> list[str]: - return [p.name for p in prim_fn.params if "_shape_" in p.name] - - -def test_a_dynamic_device_dim_injects_a_hidden_shape_scalar() -> None: - """A device kernel reading a ``DimVar`` axis declares the i32 extent alongside it. - - That is the ABI the HIR-to-TIR lowering appends, so codegen can plumb the - runtime extent. - """ - assert _shape_scalars(tir_dynamic_device) == ["a_shape_0"] - scalar = next(p for p in tir_dynamic_device.params if p.name == "a_shape_0") - assert isinstance(scalar.type, TensorType) - assert scalar.type.shape == () - assert scalar.type.dtype is DType.i32 - - -def test_a_host_entry_and_a_static_kernel_stay_unpolluted() -> None: - """A host entry reads its shapes from its tensor argument at launch time. - - A device kernel over only static dims has nothing to plumb, so neither - grows a hidden scalar. - """ - assert _shape_scalars(tir_host_entry) == [] - assert [p.name for p in tir_host_entry.params] == ["a"] - assert _shape_scalars(tir_static_device) == [] - - -def test_a_prim_func_param_takes_the_same_layout_sugar_as_a_func_param() -> None: - """``8192 @ cta`` canonicalises into ``(128, 64)`` on a device parameter too. - - ``_build_params`` is one walk for both dialects, so this is the TIR twin of - the HIR ``int-at-single-axis-mesh`` case. - """ - assert tir_param_layout_sugar.params[0].type == TensorType( - shape=(1, 8192), - dtype=DType.f32, - storage=StorageKind.SMEM, - layout=ShardLayout( - layout=Layout((1, 128, 64), (8192, 64, 1)), attrs=(Split(1),), mesh=M_CTA - ), - ) - - -def test_a_trailing_underscore_selects_the_effect_form() -> None: - """A bare ``copy_(...)`` strips its suffix and resolves ``Copy`` from the T dialect. - - It is unresolved through the closure, so it goes through - ``dispatch.resolve_callable`` and lands on the same statement ``T.copy`` - would have produced. - """ - assert isinstance(tir_effect_form_selector.body, Sequential) - (stmt,) = tir_effect_form_selector.body.body - assert isinstance(stmt, Evaluate) - assert isinstance(stmt.callable, Copy) - assert stmt.args[0].name == "a" - assert stmt.args[1].name == "b" - - -def test_an_atom_binding_is_a_compile_time_value_not_a_letstmt() -> None: - """``op = ...`` and ``atom = ...`` bind statically, leaving the body empty.""" - assert all(not isinstance(s, LetStmt) for s in tir_static_atom_bindings.body.body) - assert tir_static_atom_bindings.body.body == () - - atom = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) - assert isinstance(atom, MmaAtom) - assert (atom.A, atom.B, atom.C) == (_ATOM.A, _ATOM.B, _ATOM.C) - assert atom.required_scope is _ATOM.required_scope - - -def test_a_fragment_alloc_in_a_valid_scope_keeps_the_atom_layout() -> None: - """A scope that passes the use-point check leaves the fragment layout alone. - - The resolver returns the atom's own object rather than rebinding it to the - caller's mesh, whatever that mesh names its axes. - """ - mesh_scope = next( - s for s in tir_atom_fragment_in_a_warp_scope.body.body if isinstance(s, MeshScope) - ) - let = next(s for s in mesh_scope.body.body if isinstance(s, LetStmt)) - assert let.var.type.layout is _ATOM.A - - -def _syncs(body) -> list[Sync]: - """The Sync ops, in order, from a parsed body.""" - out: list[Sync] = [] - - def walk(s) -> None: - if isinstance(s, Sequential): - for x in s.body: - walk(x) - elif isinstance(s, MeshScope): - walk(s.body) - elif isinstance(s, Evaluate) and isinstance(s.callable, Sync): - out.append(s.callable) - - walk(body) - return out - - -def test_a_mesh_scoped_sync_records_the_participating_sub_box() -> None: - """``T.sync(m)`` lowers to ``Evaluate(Sync(mesh=m))`` carrying a compile-time mesh. - - A sliced sync records its extents and slice origin in a composed layout; the - full sync's layout stays a plain ``Layout``. - """ - mesh_scope = tir_sync_scopes.body.body[0] - assert isinstance(mesh_scope, MeshScope) - first = mesh_scope.body.body[0] - assert isinstance(first, Evaluate) and isinstance(first.callable, Sync) - assert first.callable.mesh == mesh_scope.mesh - - full, warp0, mid = _syncs(tir_sync_scopes.body) - assert not isinstance(full.mesh.layout, ComposedLayout) - assert full.mesh.layout.shape == (4, 32) - assert warp0.mesh.layout.outer.shape == (1, 32) and warp0.mesh.layout.offset == 0 - assert mid.mesh.layout.outer.shape == (2, 32) and mid.mesh.layout.offset == 32 - - -def test_verify_accepts_a_full_and_a_sliced_sync() -> None: - """The scopes the verifier must accept, next to the ones it must refuse.""" - mesh = thread_mesh() - for participating in (mesh, mesh[1:3, :]): - verify_prim_function(scoped_sync(mesh, participating)) - - -def test_codegen_puts_a_multi_warp_subset_behind_a_named_barrier() -> None: - """An aligned multi-warp subset emits a predicated ``bar_sync``, not ``__syncthreads``.""" - ctx = CodegenContext() - ctx.reset_barrier_ids() - ctx.emit_node(Evaluate(callable=Sync(mesh=thread_mesh()[1:3, :]), args=())) - - assert "SyncKind::bar_sync, 32, 64, 0u," in ctx.source()