diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index a7fcfd3d..854ed7e5 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1,310 +1,23 @@ -# TileFoundry Spec — analysis (polyhedral model + per-stage target facts) +# TileFoundry Spec — analysis (authored-HIR metrics + per-stage target facts) This spec owns TileFoundry's fact layer: everything a later stage decides *over*, and nothing that decides anything itself. It has two surfaces: | Surface | Entry | What it states | |---|---|---| -| Polyhedral model | `extract(hir) -> TileGraph` | one HIR `Function` body as isl domains, access relations and auto-inferred dependences — target-independent | | Program check | `check_program(module, function, level=..., budget=..., analyzers=...)` | an inlined Function view after validating one authored program, its declared topology, and what each requested analysis needs of it | | Composed measurement | `analyze(module, function, analysis=...)` | one or more root analyses and their union dependency closure, leaving typed Metadata on the IR | Per-Op semantic derivation — typeinfer, the forward access relation, shard propagation — is owned by [semantic-analysis](./semantic-analysis.md), and the -registries behind it by [visitor-registry](./visitor-registry.md); the -polyhedral model consumes the forward relation +registries behind it by [visitor-registry](./visitor-registry.md); the families +below consume the forward relation ([visitor-registry §4.1](./visitor-registry.md#41-access-relation-service--access_relation)) rather than restating it. -**Layering.** The decisions taken over these facts are owned by -[schedule](./schedule.md#4-kernel-schedule-construction). The dependency is -one-way: the schedule layer reads this layer's facts, and this layer MUST NOT -import or otherwise depend on the schedule layer. The atom catalogue and the -store a tile lives in are the schedule layer's own inputs and are owned there -([schedule §5](./schedule.md#5-scheduling-facts)). +## 1. Authored-HIR metrics -## 1. Polyhedral model - -One extraction models one HIR `Function` body as a set of *statements* at -**element** granularity. A statement is one compute op at one call site; its -iteration domain is the op's own element domain, prefixed by one dimension per -authored loop enclosing it. Nothing here retiles, reorders, or searches: the -model states what the authored program accesses, and what must run before what. - -### 1.1 `TileUnit` - -```python -class TileUnit: - """One statement's identity. - - Attributes: - name: attribute; isl tuple name shared by this statement's domain, reads, writes and deps pieces. - op: attribute; The HIR Call (one op at one call site) that produced this statement. - """ - - name: str - op: object -``` - -- constraints: - - The structure MUST be immutable. - - `name` MUST be a valid isl identifier and MUST be the tuple name of every - `TileGraph` piece belonging to this statement. - - `name` MUST be unique within one extraction: a name two statements would - share MUST be disambiguated with a numeric suffix, and a statement - contributed by a penetrated nested `Function` MUST additionally carry that - call site's own prefix. - - `op` MUST be the `Call` node itself, so a consumer can recover - `op.target`, `op.args` and `op.type`; it MUST NOT be the bare `Op`. - -### 1.2 `TileGraph` - -```python -class TileGraph: - """Provide the polyhedral model of one HIR Function body. - - Attributes: - domain: attribute; Union of every statement's iteration domain, one named tuple per statement. - deps: attribute; Auto-inferred read-after-write must-dependence between statement instances. - reads: attribute; Union of every statement's read access relations, statement tuple to buffer tuple. - writes: attribute; Union of every statement's write access relations. - units: attribute; One TileUnit per statement, in dependence-respecting order. - params: attribute; isl parameter name to the ShapeDim it stands for. - buffer_dtypes: attribute; Buffer tuple name to the DType its elements carry. - parallel_dims: attribute; Statement name to one flag per own domain dimension, set when that dimension carries no dependence. - """ - - domain: "isl.union_set" - deps: "isl.union_map" - reads: "isl.union_map" - writes: "isl.union_map" - units: tuple[TileUnit, ...] - params: dict - buffer_dtypes: dict = field(default_factory=dict) - parallel_dims: dict = field(default_factory=dict) -``` - -- constraints: - - The structure MUST be immutable; a stage that adds a fact MUST return a - replaced copy rather than mutate one. - - `domain`, `reads` and `writes` MUST be unions of per-statement pieces with - one isl tuple name per statement (the domain, and the input side of the - access relations) or per accessed buffer (their output side). The domain - MUST NOT be a single unnamed set: the schedule tree needs one named tuple - per statement. - - `deps` MUST relate statement *instances* and MUST be derived from - `reads` / `writes` by the extraction itself ([§1.3](#13-extract)), never supplied by a - caller. - - `units` MUST be ordered so that every dependence runs forwards: the order - is the body's SSA-DAG postorder, which makes sequencing the statements in - that order legal by construction. - - `params` MUST resolve every isl parameter name appearing in `domain` back - to its `ShapeDim`. One parameter name resolving to two different - `ShapeDim`s across statements MUST raise. - - `buffer_dtypes` MUST record an element `DType` for every buffer an access - relation names, so a byte count over an access relation needs no second - walk of the HIR. - - `parallel_dims` MUST carry exactly one flag per own domain dimension of - every statement, and MUST be measured from `domain` + `deps` ([§1.7](#17-parallel-dimensions)) rather - than reported by a scheduler. - - Schedule trees, ring depths, and decisions MUST remain schedule-owned state - outside `TileGraph`; schedule program views pair those values with the - immutable analysis graph without mutating it - ([schedule §4](./schedule.md#4-kernel-schedule-construction)). - -### 1.3 `extract` - -```python -def extract(hir: Function) -> TileGraph: ... - -class ExtractError(NotImplementedError): - """A construct the polyhedral extraction does not model.""" -``` - -`extract` walks `hir.body` in SSA-DAG postorder (dependencies before -dependents) and classifies every node it meets: - -| Body node | How `extract` models it | -|---|---| -| `Call` of a compute op | one statement; or one statement per output branch when the outputs cannot share one domain (`RoPE`'s grouped-query `q` / `k`, whose head counts differ) | -| `Call` of `TupleGetItem` / `Reshape` / `IndexSelect` / `Slice` | structural view — no statement. It resolves to its source's buffer name, and the coordinate change it expresses is folded into every consumer's access map | -| `Call` of `Zeros` / `FullLike` | buffer declaration — no statement and no access relation: it names a fresh buffer and gives it a starting value | -| `Call` whose target is a `Function` | penetrated, not rejected: the callee's params bind to the caller's already-resolved argument expressions, its body is walked in place, and every statement and buffer it contributes is prefixed with the callee name plus a per-call-site index | -| `GridRegionExpr` | not a statement — it contributes one leading domain dimension to every statement it encloses ([§1.4](#14-authored-loops)) | -| `Tuple` | resolved through the same substitution table; no statement | - -- constraints: - - Each statement's access relations MUST come from the forward - (input-type-driven) relation service - ([visitor-registry §4.1](./visitor-registry.md#41-access-relation-service--access_relation)). - `extract` MUST stamp the statement and buffer tuple names onto each - returned map and restrict it to the paired domain, and MUST reuse each - access map's own formula unchanged — no retiling happens here. - - Before the relation is built, every argument type MUST be narrowed to its - per-shard local shape when it carries a `ShardLayout`: each mesh `Split`'s - target *tensor* axis divided by that mesh axis's extent, tensor rank - preserved. A `Partial` / `Broadcast` / `Dynamic` mesh axis consumes no - tensor axis. Narrowing is centralized in the extraction, so every - registered relation is sharding-aware without knowing sharding exists. - - A `Split`-sharded axis whose extent is not a static integer, or is not - evenly divisible by its mesh extent, MUST raise `ExtractError`. - - A loop-indexed `Slice` read MUST map result coordinate `u` to source - coordinate `start + u * stride`. Its consumer statement domain contains - only full windows (`start + size * stride <= source_extent`). No residual - tail domain is returned. A non-affine start or non-static window step MUST - raise `ExtractError`. - - An op with no registered forward relation MUST raise `ExtractError` naming - the op and the registration remedy. `extract` MUST NOT guess an access - pattern and MUST NOT carry a per-op fallback. - - A statement MUST write at least one value; a relation that produced no - output map MUST raise `ExtractError`. - - A statement with several outputs MUST write them under its own buffer name - suffixed `_`, matching what a downstream `TupleGetItem` read - resolves to. - - An output access map that is **not** injective MUST also be recorded as a - read of the same buffer: two domain points writing one output cell is only - sound as a read-modify-write accumulation, and recording that self-read is - what lets the dependence inference discover a reduction carry. An injective - output map MUST stay a pure write. - - `deps` MUST be inferred by isl dependence analysis over `reads` (sink), - `writes` (must-source) and an initial total execution order, keeping the - resulting read-after-write must-dependence. That initial order MUST place - the authored loop coordinates *before* the per-statement postorder index, so - a value written at one iteration is seen by a read at the next; a - statement-first order would lose every loop carry. - - `extract` MUST raise `ExtractError`, naming the offender, for: a body that - is `None`; a body with no compute op; a self-recursive nested call; a - dispatch prototype (a callee carrying variants or no body); a nested call - whose arity does not match the callee's params; and a `DimVar` that binds to - conflicting shapes at one call site. - -### 1.4 Authored loops - -An authored loop is a `GridRegionExpr` ([hir §1.2](./hir.md#12-gridregionexpr)). -It is modelled as a domain dimension, not as a statement. - -- constraints: - - Each enclosing loop MUST prefix one dimension to the domain of every - statement it encloses, outermost loop first, so the loop axes are the - leading dimensions of that domain. - - The dimension MUST range over the loop's own half-open `[start, extent)` - and MUST carry the raw induction value rather than a normalised trip - counter; a `step` other than `1` MUST appear as a stride constraint on it. - - `start` and `step` MUST be static integers. `extent` MAY be a bare - `DimVar`, which becomes a same-name isl parameter bound to its own - `[lo, hi)` range — the same treatment a dynamic tensor axis gets. - - Which statements a loop encloses MUST be decided by **variance**, not - reachability: a value is inside a loop only when it transitively reads that - loop's induction variable or one of its carried args. A loop-invariant value - MUST lift out, and a value read after the loop MUST be outside it even - though the loop's yield produced it. - - A carried arg and the value yielded into it MUST share one buffer name. The - dependence inference then reports the loop carry as a distance-1 dependence - along that loop's dimension, read at iteration `i` and written at `i - 1`; - nothing states the carry separately. - - An `IndexSelect` whose one-element index is an enclosing loop's induction - variable reshaped to `(1,)` MUST fold into the consumer's access map at the - selected dim, so iteration `i` addresses slice `i`. Any other index MUST - raise `ExtractError`: a data-dependent selection has no affine access map - and MUST NOT be approximated. - -### 1.5 Facts over a time relation - -Four measurements take the time relation as data — one `isl.union_map` from -statement coordinates to a common time space, which the schedule layer owns and -this layer only reads. - -```python -def time_extents(tg: TileGraph, time_map: "isl.union_map") -> tuple[int, ...]: ... - -def statement_time_dims(tg: TileGraph, time_map: "isl.union_map") -> dict[str, tuple[int, ...]]: ... - -def carried_distances(tg: TileGraph, time_map: "isl.union_map", n_dims: int) -> dict[str, tuple[int, ...]]: ... - -def access_footprints(tg: TileGraph, time_map: "isl.union_map") -> tuple[AccessFootprint, ...]: ... -``` - -- constraints: - - `time_extents` MUST return the per-dimension extent of the time relation's - range over `tg.domain`. Every statement MUST share one time space, and every - time dimension MUST start at `0` — tile counting assumes an origin-based - extent — otherwise it MUST raise `ExtractError`. - - `statement_time_dims` MUST report, per statement and per time dimension, the - statement's own domain dimension that time dimension travels with, or `-1` - where it is constant there. A time dimension mixing two domain dimensions (a - skewed band) MUST raise `ExtractError`: no per-axis tile size describes it. - - `carried_distances` MUST report, per buffer, the largest dependence distance - isl reports along each of the first `n_dims` time dimensions. A dependence - MUST be attributed to every buffer its source writes and its sink reads, - which for a read-after-write must-dependence is exactly the memory it - travels through. - - `access_footprints` MUST return one `AccessFootprint` per read and per write - of `tg`, expressed against the time relation's range so that a size per time - dimension sizes it. - -### 1.6 `AxisExtent` / `AccessFootprint` - -```python -class AxisExtent: - """One buffer dimension's reach inside one statement's whole access. - - Attributes: - axes: attribute; Time dimensions that reach this dimension; empty when none does. - extent: attribute; Number of elements of this dimension the access reaches. - """ - - axes: tuple[int, ...] - extent: int - -class AccessFootprint: - """One (statement, buffer) access, sized per buffer dimension. - - Attributes: - statement: attribute; isl tuple name of the accessing statement. - buffer: attribute; isl tuple name of the accessed buffer. - is_read: attribute; True for a read access, False for a write. - dims: attribute; One AxisExtent per buffer dimension. - elem_bytes: attribute; Bytes one element of the buffer occupies. - """ - - statement: str - buffer: str - is_read: bool - dims: tuple[AxisExtent, ...] - elem_bytes: int -``` - -- constraints: - - Both structures MUST be immutable. - - `extent` MUST be measured off the access relation, never derived from a tile - size: a dimension read at a rate reaches fewer elements than the iterations - that reach it. - - `axes` carries no size, only the reuse fact. An empty `axes` MUST mean no - time dimension reaches that dimension — it is re-read in full by every - iteration. - - The element count of one access MUST be the product over `dims`. That count - is the bounding box of the access's range: exact for a box-shaped access, an - upper bound for one that leaves holes inside its own box. - - `elem_bytes` MUST be the buffer's element size in whole bytes, resolved from - `TileGraph.buffer_dtypes`. A buffer with no recorded dtype MUST raise - `ExtractError`. - -### 1.7 Parallel dimensions - -`TileGraph.parallel_dims` is the fact isl names `coincident`, measured here -rather than obtained from a scheduler. - -- constraints: - - Only a statement's **self**-dependence MAY constrain its own dimensions: the - schedule layer sequences statements, so every cross-statement dependence is - already satisfied by that order. - - A dimension MUST be reported parallel when every self-dependence has - distance `0` there, and MUST NOT be otherwise. A statement with no - self-dependence MUST have every dimension reported parallel. - -## 2. Authored-HIR metrics - -The measurement entry is the composed operation ([§3](#3-composed-analysis)). +The measurement entry is the composed operation ([§2](#2-composed-analysis)). Each family below owns its record, field derivation, target facts, and rendered forms. The command line composes one call per requested family and renders those results together ([cli §Analyze](./cli.md#analyze)). @@ -320,7 +33,7 @@ they describe: a record on a `Call` describes that call, while a record on a inherently carries. It states what one analysis found for one invocation, and there MUST be no cross-call cache behind it. -### 2.2 Analysis families +### 1.2 Analysis families The first families are `compute-cost`, `memory`, `roofline`, and `performance`. Each owns its record types and declares its dependencies and output additions. @@ -411,8 +124,8 @@ layer settles is which type a field holds and what its keys name: metadata-free traversal of the derived program answers all of them. Performance readiness requires a positive `ParallelCapacityFacts` value for the selected topology, rates stated for that same level, and one valid - execution placement for every occurrence that will take time. Where the buffers - go is not a readiness question: it is decided with the schedule. + execution placement for every occurrence that will take time. Where the + buffers go is not a readiness question: nothing here decides it. Failing performance readiness MUST NOT make the same unplaced program invalid for `compute-cost`, `memory`, or `roofline`. - Global logical work, per-unit work, and lifetime order MUST remain @@ -423,7 +136,7 @@ layer settles is which type a field holds and what its keys name: - A rendering MUST report what the caller requested. Dependency records nobody requested MUST stay on the IR and MUST NOT be reported except for roofline's bounded evidence defined below. Record ownership MUST come from the - Target-selected descriptor ([§3.2](#32-target-selected-analyzers)). + Target-selected descriptor ([§2.2](#22-target-selected-analyzers)). - Every rendering of one run MUST select records through one shared decision and MUST show only records actually written. - Every reported quantity MUST come from a record, except a total that is the @@ -445,11 +158,11 @@ layer settles is which type a field holds and what its keys name: facts MUST stay on their annotated equations; JSON MAY retain operand names and types in its structured projection. -#### 2.2.1 `compute-cost` +#### 1.2.1 `compute-cost` `compute-cost` measures the logical work of each authored `Call` without reading target hardware facts. What an occurrence moves is the memory family's answer -([§2.2.2](#222-memory)), read off the same registered evaluator. +([§1.2.2](#122-memory)), read off the same registered evaluator. ```python class ComputeCostMetadata(IRMetadata): @@ -508,7 +221,7 @@ Each reported Call's JSON projection is under its `compute-cost` key: - Downstream families MUST read the already-scaled record and MUST NOT apply authored loop trip counts a second time. -#### 2.2.2 `memory` +#### 1.2.2 `memory` `memory` measures whole-Function value lifetimes and footprints, decides where each value's bytes live, and states what every occurrence moves and at which @@ -668,7 +381,7 @@ of this analysis. - A capacity conclusion MUST NOT correct or invent a movement number. What an occurrence moves is counted once from its own boundaries, so a function with no `allocation` still carries traffic -- a different question from whether a - time may be reported for it ([§2.2.4](#224-performance)) -- and a window + time may be reported for it ([§1.2.4](#124-performance)) -- and a window whose start arrives at run time reads that start rather than becoming a full read of its source and a write of its result. @@ -863,7 +576,7 @@ attached only to the Function. Its full JSON projection is under implicit cache capacity MUST instead produce an advisory and MUST NOT fail the call. -#### 2.2.3 `roofline` +#### 1.2.3 `roofline` `roofline` converts recorded work into a lower time bound at the target's published compute and memory rates. @@ -945,7 +658,7 @@ standing an instruction rate in for one. Requesting roofline adds this verdict and the two quantities the bound divides -- the summed `flops` and the bandwidth-level bytes -- as `totals`. Nothing else its dependencies wrote is promoted; asking for `memory` is what states those -([§2.2.2](#222-memory)). +([§1.2.2](#122-memory)). ```text roofline ideal-ns= bound-by= @@ -1021,11 +734,11 @@ as defined in that family's section. reports only what roofline and that dependency wrote. No other family's conclusion is promoted into it. -#### 2.2.4 `performance` +#### 1.2.4 `performance` `performance` places compute-cost-priced occurrences on a CTA-local nominal timeline, holds the buffers they keep live to the levels this model addresses, -and scales the root schedule by a fixed physical parallel capacity. The records +and scales the root timeline by a fixed physical parallel capacity. The records it owns are named for the prediction they carry rather than for the selector, so that the interval stays one nested value with one meaning wherever it appears. @@ -1075,7 +788,7 @@ Occurrence fields are: | Field | How it is computed | Reads the target | |---|---|---| -| `start_ns` | Start of one occurrence in the authored-order local schedule, after its producers end and after the last occurrence sharing any of its participants. | No | +| `start_ns` | Start of one occurrence on the authored-order local timeline, after its producers end and after the last occurrence sharing any of its participants. | No | | `end_ns` | End of that occurrence's first execution. | No | | `trips` | One outside a loop; within a loop, the enclosing loop trip count represented by the interval. | No | | `stride_ns` | Zero outside a loop; within a loop, the makespan of one body execution. | No | @@ -1084,7 +797,7 @@ Function summary fields are: | Field | How it is computed | Reads the target | |---|---|---| -| `timeline` | `[0, local makespan * waves)`, where the local makespan is the end of the CTA-local schedule, or zero with no work. Its duration is the prediction. | Through `waves` | +| `timeline` | `[0, local makespan * waves)`, where the local makespan is the end of the CTA-local timeline, or zero with no work. Its duration is the prediction. | Through `waves` | | `waves` | `ceil(N / P)`, where `N` is the static extent of the root topology selected by `ParallelCapacityFacts.topology` and `P` is `parallel_units`. | `ParallelCapacityFacts` | Occurrence intervals remain CTA-local. They are not copied once per wave, and @@ -1100,7 +813,7 @@ class ParallelCapacityFacts: """Carry the parallel capacity assumed by performance analysis. Attributes: - topology: attribute; Topology level being scheduled. + topology: attribute; Topology level being measured over. parallel_units: attribute; Instances admitted concurrently. """ @@ -1192,7 +905,7 @@ model. bytes is a plan's decision and no plan has been made -- and no occurrence is held back for a write nobody proved happens in place. - Occurrences MUST be laid out in inline occurrence order. Reordering - independent work is a schedule's decision, not an analysis's: what overlaps + independent work is a later decision, not an analysis's: what overlaps is what the program's own placement made independent, and the reported time is the time of the program as written. On a `Function`, the summary's `timeline` MUST start at zero and span the whole local plan, scaled by @@ -1210,10 +923,10 @@ model. - Performance is a modeled plan and MUST NOT be read as a guarantee about lowering, physical occupancy, or runtime performance. -## 3. Composed analysis +## 2. Composed analysis `tilefoundry.analysis.check_program` is the shared, reusable gate before an -analysis or schedule algorithm runs. +analysis runs. ```python def check_program( @@ -1245,19 +958,19 @@ class AnalysisCheckContext: - constraints: - The operation MUST infer types over the full reachable Function graph and validate its caller/callee execution context, and MUST NOT run an analysis - or schedule algorithm or attach derived Metadata to the authored IR. + or attach derived Metadata to the authored IR. - The reachable Function and Mesh geometry and every effective Module - topology extent MUST be concrete before this operation runs. Public Analyze - and Schedule calls with `dims` MUST resolve all three through one binding - pass before calling this gate; a residual dimension expression MUST fail - before any consuming algorithm runs. + topology extent MUST be concrete before this operation runs. A public + Analyze call with `dims` MUST resolve all three through one binding pass + before calling this gate; a residual dimension expression MUST fail before + any consuming algorithm runs. - Every effective Module topology MUST name a level the resolved Target supports. A resolved static extent MUST be positive and within that level's finite hardware limit. A rejection MUST name the level, its extent, and the reason. - A non-`None` `level` MUST name exactly one effective Module topology. - - Analyze and Schedule MUST call this operation before any consuming - algorithm. Analyze MUST pass the whole resolved dependency closure as + - Analyze MUST call this operation before any consuming algorithm, and + MUST pass the whole resolved dependency closure as `analyzers`, so every analysis about to run states its input contract here. - The `analyzers` checkers MUST be bound to the derived Function and run in closure order: every `check_target`, then every `check_call` over one @@ -1282,12 +995,11 @@ class AnalysisCheckContext: expression nodes after inlining. An oversized view MUST fail with both its size and the limit and MUST NOT return a partial Function. - Authored-analysis readiness is not a program-level rejection. Analyze MUST - NOT reject a schedule constraint: `where(...)` is a scheduling input, and a - program carrying one is measured as written. A value whose placement is - deferred contributes its whole-program figures to the per-unit total, - because a deferred layout states no distribution to project through; - values with a resolved layout still project. Schedule MAY consume or - diagnose those inputs under its own algorithm contract. + NOT reject an authored `where(...)` constraint: it is an input to a later + decision, and a program carrying one is measured as written. A value whose + placement is deferred contributes its whole-program figures to the per-unit + total, because a deferred layout states no distribution to project through; + values with a resolved layout still project. `tilefoundry.analysis.api.analyze` is the dependency-composed measurement operation. One call selects one or more root analyses by name; the operation @@ -1413,7 +1125,7 @@ def analyze( renderings of it and of the Metadata on the IR, and MUST NOT be fields of it. -### 3.1 Shared Scope and Access +### 2.1 Shared Scope and Access The normalized HIR is visited once per `analyze()` call. That visit produces a `Scope` tree parallel to Function/GridRegionExpr nesting and `Access` relations @@ -1428,7 +1140,7 @@ call site, source expressions shared by identity remain one shared expression in the clone; sharing never aliases the independently cloned body of another call site. -### 3.2 Target-selected Analyzers +### 2.2 Target-selected Analyzers ```python class AnalyzeContext: diff --git a/docs/spec/architecture.md b/docs/spec/architecture.md index c9808237..a116ee06 100644 --- a/docs/spec/architecture.md +++ b/docs/spec/architecture.md @@ -19,7 +19,6 @@ graph TD hir["hir"] tir["tir"] analysis["analysis"] - schedule["schedule"] passes["passes"] target["target"] runtime["runtime"] @@ -46,15 +45,11 @@ graph TD coreir --> hir coreir --> tir hir --> analysis - hir --> schedule - schedule --> passes hir --> passes tir --> passes passes --> target target --> runtime - analysis -. facts read by .-> schedule target -. projects declared Facts .-> analysis - target -. projects declared Facts .-> schedule types -. carried by Expr type .-> coreir shard -. layout sublayer .-> types @@ -69,14 +64,11 @@ graph TD evaluator -. reference oracle .-> hir cli -. user entry surface .-> parser cli -. reports .-> analysis - cli -. reports .-> schedule ``` A TileFoundry compilation flows left to right along the **pipeline**: -`parser → core-ir → {hir, tir} → passes → target → runtime`. Typed HIR MAY -first pass through the public `schedule` operation before pass sequencing; the -algorithm it selects decides over the facts the `analysis` layer states about the -same HIR, and the direction is one-way. The +`parser → core-ir → {hir, tir} → passes → target → runtime`. The `analysis` +layer states facts about the same typed HIR without joining that flow. The **type system** (types + shard) and the **IR framework** (visitor-mutator + visitor-registry) cut across every pipeline stage: they are co-designed with the IR, not standalone modules. Auxiliary @@ -146,29 +138,15 @@ This stage layers two concerns on top of the same IR: are both ordinary stages in that manager. A pass may use a pass-private intermediate representation without elevating it to a peer IR layer. -3. **Fact layer** — the polyhedral model of one HIR `Function` body and - the authored-HIR metrics ([analysis](./analysis.md)). It is neither a - pass nor an IR layer: it measures, and the scheduling algorithms below - decide over what it measures. The facts a scheduling decision is made - *over* — the atom catalogue and the store a tile lives in — belong to - the scheduling layer that decides, not here - ([schedule](./schedule.md)). +3. **Fact layer** — the authored-HIR metrics measured over one HIR + `Function` ([analysis](./analysis.md)). It is neither a pass nor an IR + layer: it measures, and it decides nothing over what it measures. IR traversal / rewrite utilities (`ExprVisitor` / `ExprCloner` / `StmtVisitor` / `StmtMutator` / mixed stmt-expr rewriters) are shared infrastructure used by both passes and codegen walkers; the framework contract lives in [visitor-mutator](./visitor-mutator.md). -Scheduling is one explicit public operation, not a pass-manager stage and not a -Target-owned service. A caller names the program and one level of the hierarchy -that program declares; the algorithm registered for that exact hardware and level -answers with a Plan it owns entirely. What a Plan states is a decision about a -program, never a rewritten one: no scheduling algorithm materializes its selection -into HIR, and applying a decision is a separate operation a caller asks for. The -invocation contract, the result boundary, and the Plan base are owned by -[schedule](./schedule.md). An algorithm reads the hardware it decides over by -projecting the same Target for the aggregates it declares. - ## 6. Target / codegen Codegen is the back end. **TIR is the lowest IR**, and each target @@ -229,14 +207,13 @@ This table is the authoritative spec-to-box map. Each row lists the | **[evaluator](./evaluator.md)** | HIR reference interpreter: `evaluate` entry, `Value` family (`TensorValue` / `TupleValue`), `register_eval` op registry, node-evaluation + `GridRegionExpr` + layout-domain rules. Logical reference oracle, no codegen / runtime | | **[visitor-registry](./visitor-registry.md)** | Derived-visitor dispatch pattern: `AnalysisRegistry`, per-class handler registration, four instances (`typeinfer` / `verify` / `codegen_` / `cost`) with their Context / Visitor derivations | | **[semantic-analysis](./semantic-analysis.md)** | Static analysis service semantics: type propagation (relation-derived type behavior), access relation analysis, shard propagation (logical shape → layout domain, relation-driven propagation, output storage + mesh/layout compatibility). The registration mechanism itself is owned by visitor-registry | -| **[analysis](./analysis.md)** | Fact layer: the polyhedral model of one HIR Function body (`TileGraph` / `extract`, authored-loop modelling, and the facts measured over a time relation), and the composed authored-HIR measurement — its analysis families, their owned Metadata records, and the narrow Target Facts each family declares | +| **[analysis](./analysis.md)** | Fact layer: the composed authored-HIR measurement — its analysis families, their owned Metadata records, and the narrow Target Facts each family declares | | **[visitor-mutator](./visitor-mutator.md)** | IR traversal / rewrite infrastructure: expr / stmt visitors, mutators, identity-preserving rewrite invariants, mixed stmt-expr traversal | | **[passes](./passes.md)** | Pass framework + implemented passes: `Pass` / `PassManager`, three pass granularities, per-pass subsections (lowering / optimization rules) | -| **[schedule](./schedule.md)** | The public scheduling operation: invocation contract, exact algorithm registration, shared options, result boundary, the extensible Plan base, the typed plan each algorithm family exports, the schedule-tree construction and scaffold emission stages an algorithm composes its solve from, and the facts it projects (`AtomFact`, plus each family's own closed facts) | | **[target](./target.md)** | Target capability descriptors, architecture/device facts, Facts projection, and admitted program topology levels | | **[codegen](./codegen.md)** | Target-selected CodeGenerator services, emit / link products (`LinkableFunction` / `LinkableModule` / `LinkedModule`), dispatch + shape-scalar ABI, program-shape / dynamic-CTA source contract, ShardLayout emission | | **[runtime](./runtime.md)** | `RuntimeModule` / launcher ABI, C++ runtime surface, `runtime.h` umbrella header, runtime op free-function contract | -| **[cli](./cli.md)** | Command-line grammar and behavior for models, spec, tutorial, check, analyze, schedule, and inspect | +| **[cli](./cli.md)** | Command-line grammar and behavior for models, spec, tutorial, check, analyze, and inspect | | **[code-organization](./code-organization.md)** | Implementation guide (not architectural): Python source tree layout | **Cross-spec sync.** Downstream specs link back to the relevant § of diff --git a/docs/spec/cli.md b/docs/spec/cli.md index 57e01600..6d8eca42 100644 --- a/docs/spec/cli.md +++ b/docs/spec/cli.md @@ -334,7 +334,7 @@ explicit analysis; there is no ordinary `--target` option. analysis option. - `analyze` MUST invoke the public operation once with every requested root, so their union dependency closure runs on one inlined Function view - ([analysis §3](./analysis.md#3-composed-analysis)). Each closure member MUST + ([analysis §2](./analysis.md#2-composed-analysis)). Each closure member MUST run once, and requesting another root MUST NOT let one analyzer change Metadata owned by another. - A selection MUST resolve to a Module. A bare Function MUST be rejected @@ -350,7 +350,7 @@ explicit analysis; there is no ordinary `--target` option. It MUST be refused as an argument-combination error when no analysis flag was supplied, naming that a report needs a requested root and printing the `analyze` usage. Both formats MUST carry the same conclusions - ([analysis §2](./analysis.md#2-authored-hir-metrics)). + ([analysis §1](./analysis.md#1-authored-hir-metrics)). - `--operands` MUST add each operand's share of a call's traffic to that call's annotation, and MUST NOT change the JSON report, which carries that split either way. It is the traffic the annotation already states, one layer finer, @@ -368,7 +368,7 @@ explicit analysis; there is no ordinary `--target` option. MUST be repeatable to bind several. One dimension MUST receive one extent; a comma-separated list of extents for one dimension MUST be rejected because several extents together are a `check` request. It MUST be passed through as the - operation's `dims` ([analysis §2.2](./analysis.md#22-analysis-families)); + operation's `dims` ([analysis §1.2](./analysis.md#12-analysis-families)); the CLI MUST NOT specialise the selection itself, because then what it wrote would be about a program the operation never saw. - A `--dim` argument that is not `NAME=EXTENT`, or whose extent is not an @@ -397,7 +397,7 @@ explicit analysis; there is no ordinary `--target` option. - Output MUST report the analyses that were requested. A dependency that ran because a requested root needed it MUST appear in the executed list and, other than the bounded roofline support view defined by - [analysis §2](./analysis.md#2-authored-hir-metrics), MUST NOT have its own + [analysis §1](./analysis.md#1-authored-hir-metrics), MUST NOT have its own measurements reported. - The report's `target` field MUST be the concrete Target value's `identity`, so two products served by one Target class remain distinguishable. @@ -444,7 +444,7 @@ interface that claims to list every Facts projection it supports. already occupied by any other importable module MUST likewise be rejected without replacing that module in `sys.modules`. - Every command MUST replay the registry before doing its own work, so an - added Target is equally available to inspection, analysis, and scheduling. + added Target is equally available to inspection and analysis. A missing or changed source MUST produce a warning naming that entry while valid entries continue to load and the requested command continues. - The default writable registry MUST be diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index a484deb8..cc2c46b7 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -20,13 +20,12 @@ truth for the directory's structure and invariants. | `ir/core/` | [core-ir](./core-ir.md) | Shared node algebra: `Module` / `Expr` / `Var` / `Constant` / `Tuple` / `Op` / `Call` / `Stmt` (base class) / `OpSchema` / `ParamDef` / call-graph and ownership queries / typed metadata attach-detach and diagnostics / `@register_op` / `@register_alias` / `op_registry` / `errors`. | | `ir/types/` | [types](./types.md) | Type-system root: `Type` / `TensorType` / `TupleType` / `UnitType` / `CallableType` / `DType` / `StorageKind` / `resolve_storage` / local projections (`local_type_of`) / tensor-leaf, byte-by-storage, and topology-extent queries / `dim.*` (with their typeinfer). | | `ir/types/shard/` | [shard](./shard.md) | Shard / layout sublayer: `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`). The physical nesting reflects the spec's "sublayer" relationship. | -| `ir/constraints/` | [schedule](./schedule.md) | Authored scheduling constraints: the shared base plus layout, mesh, and storage constraint records. | +| `ir/constraints/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | | `ir/visitor.py` | [visitor-mutator](./visitor-mutator.md) | `ExprFunctor` / `ExprVisitor` / `ExprWalker` / `ExprCollector` / `ExprCloner` / `BindingSubstitutionCloner` / `StmtVisitor` / `StmtMutator` / `StmtExprMutator`, plus `collect_exprs`, value-operand/function-value queries, and the canonical `PrimFunction` walk and rewrite entries. | | `ir/hir/` | [hir](./hir.md) | HIR Op layer; one subdirectory per category (`math/` / `tensor/` / `nn/` / `shape/` / `sharding/`). One real Op per `.py` ([§2](#2-file-naming-and-content-rules) rule 1); surface-alias schemas have no per-name file and live in each category's `aliases.py` ([§2](#2-file-naming-and-content-rules) rule 5). | | `ir/tir/` | [tir](./tir.md) | TIR layer: `stmt.py` re-exports the `Stmt` base from `ir/core/stmt.py`; `stmts.py` hosts the general TIR `Stmt` subclasses (`LetStmt` / `Evaluate` / `Sequential` / `MeshScope` / …), while specialized statement families such as `DispatchCall` may live in their own file; `prim_function.py`; effect Ops and TIR-owned Expr Ops by category (`memory/` / `nn/` / …); `launch.py` owns `Launch` and its authored launch-attribute descriptors; `arith.py` / `reduce.py` for tag-dispatched `Binary` / `Unary` / `Reduce`; `intrinsic.py` for the `@intrinsic` decorator. Target-specific nodes nest under `ir/tir///` (e.g. `ir/tir/cuda/nn/mma.py`) per [§2](#2-file-naming-and-content-rules) Rule 1c. | | `parser/` | [parser](./parser.md) | DSL → IR parsing: `base.py` (shared visitor base + dispatch), `hir_parser.py` (`@func` body), `tir_parser.py` (`@prim_func` body), layout sugar / range-slice / dispatch modules. **Not under `ir/`** — the parser is a producer of IR, not an IR sublayer. | -| `analysis/` | [analysis](./analysis.md) | Fact layer over typed HIR: `poly/` (the polyhedral model — `extract` / `TileGraph` and the facts measured over a time relation), and one module per analysis family. The compact public surface lives in `analysis/__init__.py`; per-target atom catalogues and Facts projections live with their owning Target. | -| `schedule/` | [schedule](./schedule.md) | The public Schedule boundary in `schedule/__init__.py` -- the `schedule()` operation, immutable options, result, and plan base. One directory per algorithm family: `pipeline/` for asynchronous overlap within a cooperating unit, `partition/` for spatial division across a device. Each owns its private program view, projected Facts, closed problem, solve, and typed plan export; concrete Target packages register which families serve which exact levels. | +| `analysis/` | [analysis](./analysis.md) | Fact layer over typed HIR: one module per analysis family. The compact public surface lives in `analysis/__init__.py`; per-target Facts projections live with their owning Target. | | `passes/` | [passes](./passes.md) | Pass framework (`pass_base.py` / `pass_manager.py`) plus concrete transforms (`transforms/.py`, [§2](#2-file-naming-and-content-rules) rule 6). | | `target/` | [target](./target.md) | Compilation Target classes, class registration, service selection, and architecture/device facts: `base.py` owns `Architecture` / `Device` / `Target` / `register_target` / `registered_targets`; `services.py` owns the immutable service descriptors; each backend owns its concrete Target. Authored code constructs Target values; there is no string resolver. | | `target/cpu.py` | [target](./target.md) | The `CpuTarget` backend and its CPU code-generation service selection. | @@ -37,11 +36,10 @@ truth for the directory's structure and invariants. | `analysis/errors.py` | [analysis](./analysis.md) | `AnalysisError`, the one diagnostic the whole analysis layer raises, so catching an analysis failure catches every analysis failure rather than the subset the caller happened to import. | | `analysis/visitor.py` | [analysis](./analysis.md) | The per-call `AnalyzeContext`, carrying the shared root/current lexical `Scope` while a family traverses its work. | | `analysis/scope.py` | [analysis](./analysis.md) | The shared `Scope` tree and `Access` relations built once from normalized HIR; families query these views instead of constructing parallel structure. | -| `analysis/poly/` | [analysis](./analysis.md) | The polyhedral subpackage: `model.py` owns `TileUnit` / `TileGraph`, `extract.py` builds that graph from normalized HIR, `access.py` owns its time/access queries, and `errors.py` owns `ExtractError`; `__init__.py` is the retained public import surface. | -| `analysis/poly/affine.py` | [analysis](./analysis.md) | The shared loop-affine term parser used by poly extraction and authored-loop footprint binding, including constant loop strides and bounded invariant offsets. It does not introduce a second affine graph representation. | -| `analysis/footprint.py` | [analysis](./analysis.md) | Target-independent authored-loop access images, buffer-view folding, and deduplicated versus repeated byte readings. Requires neither a `TileGraph` nor a scheduled time map. | +| `analysis/affine.py` | [analysis](./analysis.md) | The shared loop-affine term parser used by scope binding and authored-loop footprint binding, including constant loop strides and bounded invariant offsets. It does not introduce a second affine graph representation. | +| `analysis/footprint.py` | [analysis](./analysis.md) | Target-independent authored-loop access images, buffer-view folding, and deduplicated versus repeated byte readings. Requires no separate time map. | | `analysis/report.py` | [analysis](./analysis.md) | Structured analysis report data, including record-family registration, field serialization, and target-aware report-only projections. It depends only on analysis/core modules; inspection consumes it to produce text and source annotations. | -| `analysis/check.py` | [analysis](./analysis.md) | The shared authored-program gate for analysis and scheduling: authored-type re-derivation, authored validation, call-context validation, and checker-specific input checks. Established once per public call rather than per family, before analysis or scheduling makes placement decisions. | +| `analysis/check.py` | [analysis](./analysis.md) | The shared authored-program gate for analysis: authored-type re-derivation, authored validation, call-context validation, and checker-specific input checks. Established once per public call rather than per family. | | `analysis/facts.py` | [analysis](./analysis.md) | The narrow Facts aggregates the analysis families declare — the memory hierarchy graph, the throughput rates, and the parallel capacity. It is the record of how much hardware each measurement rests on, and names no backend; a Fact shared across consumer families belongs in `target/facts.py`. | | `analysis/metadata.py` | [analysis](./analysis.md) | The typed records the families leave on the IR, split by what each number depends on rather than by convenience. | | `analysis/compute_cost.py` | [analysis](./analysis.md) | The `compute-cost` family: logical flops per DType and bytes per storage level, from the authored program alone. | @@ -49,15 +47,7 @@ truth for the directory's structure and invariants. | `analysis/roofline.py` | [analysis](./analysis.md) | The `roofline` family: the recorded work divided by the target's published rates, per Call and aggregated per Function. Adds no count of its own. | | `analysis/performance.py` | [analysis](./analysis.md) | The `performance` family: occurrences projected from the shared `Scope` tree into flat timeline records and one function envelope, scaled by parallel capacity. It introduces no second scope tree. | | `visitor_registry/` | [visitor-registry](./visitor-registry.md) | Shared registry instances and derived visitors: access-relation construction, contexts, ISL helpers, relation building, shard propagation, type inference, verification, code generation, and cost evaluation. | -| `schedule/api.py` | [schedule](./schedule.md) | The public `schedule()` operation and `ScheduleResult`: resolve the Target and the requested level from the Module, dispatch once on the exact pair, and verify the returned Plan. Generic -- it names no concrete target. | -| `schedule/plan.py` | [schedule](./schedule.md) | `SchedulePlan`, the extensible semantic base every algorithm's result derives from, and `PlanVerificationError`. It fixes three operations and no shape: there is no shared schema, version, deserializer, or renderer registry. | -| `schedule/errors.py` | [schedule](./schedule.md) | `ScheduleError`, the one diagnostic the schedule layer raises for a request it cannot serve or a solve that failed. Distinct from `PlanVerificationError`, which says a plan was produced and does not hold together. | -| `schedule/kernel_schedule.py` | [schedule](./schedule.md) | ISL schedule-tree construction, band discovery, tiling, and kernel-schedule validation. | -| `schedule/render.py` | [schedule](./schedule.md) | Scaffold emission from a graph, its independently built schedule tree, and selected ring depths. | -| `target//schedule.py` | [target](./target.md) | One backend's scheduling algorithms and immutable Scheduler values. This is where that backend's private problem construction, solve, and Plan export are composed; its Target class selects them. | -| `schedule/partition/` | [schedule](./schedule.md) | The spatial partition family: program extraction, `PartitionFacts`, the closed candidate problem, the CP-SAT solve, and the `PartitionSchedulePlan` export. Every hardware number enters through the Facts, so no module below the family entry holds a Target, and nothing in it rewrites the program it decided about. | | `visitor_registry/op_cost.py` | [analysis](./analysis.md) | Each operation's per-instance flops and bytes, registered into the shared cost-evaluator registry. Owned here rather than by any target package, because the work an operation asks for follows from its own semantics and operand types on every backend. | -| `schedule/facts.py` | [schedule](./schedule.md) | `AtomFact`, the one instruction fact every algorithm family reads the same way. Everything else a family needs from a target is declared by that family, so no aggregate here becomes a vocabulary another family has to satisfy. | | `inspection/analysis_report.py` | [inspection](./inspection.md) | Presentation of analysis-owned report data as text and annotated source. Analysis owns the structured report data and JSON dump; inspection owns how a human reads it. | | `target//facts.py` | [target](./target.md) | One backend's Facts projections selected by its Target's `get_facts`. They restate installed documents in the shape a family declared and measure nothing. | | `target/facts.py` | [target](./target.md) | Facts used across consumer families, such as topology limits, plus validation for values returned by `Target.get_facts`: the requested frozen-dataclass shape and returned type. A Fact used by one family stays with that family; this module holds no projection registry. | @@ -74,10 +64,9 @@ truth for the directory's structure and invariants. | `utils/` | [code-organization](./code-organization.md) | Shared leaf machinery: a module here MUST import nothing from `ir/`, `parser/`, `passes/`, `codegen/`, `runtime/` or `cli/`, and MUST name no layer. It is depended on and depends on nothing, which is what lets a consumer outside the package — a pre-commit hook under an interpreter with nothing installed — load one of these modules by path and get the same implementation the package uses. A helper that needs to know a layer belongs in that layer; this is not a home for anything that did not fit. | **Stage boundary.** The pipeline picture in -[architecture §1](./architecture.md#1-spec-relationship-map) places -`parser/`, `schedule/`, and `codegen/` outside `ir/` (front-end producer, -decision service over typed HIR, and back-end consumer); the physical directory -layout reflects that boundary directly. +[architecture §1](./architecture.md#1-spec-relationship-map) places `parser/` +and `codegen/` outside `ir/` (front-end producer and back-end consumer); the +physical directory layout reflects that boundary directly. **Reading notes:** @@ -94,14 +83,9 @@ layout reflects that boundary directly. [architecture §1](./architecture.md#1-spec-relationship-map) pipeline they are the front-end producer and back-end consumer of IR, not IR sublayers. -- `schedule/` sits outside `ir/` because it defines an operation over typed HIR, - not a new IR layer. `schedule/__init__.py` contains only the public operation - and its shared value structures; the construction stages are imported from - their own modules, and each algorithm family's candidate graph, solver model, - and decoded solution stay private to that family. -- `analysis/` sits outside `ir/` for the same reason: it derives facts about - typed HIR rather than defining an IR layer. It reads the IR and the `Target`, - and never `schedule/` — the dependency between the two runs one way +- `analysis/` sits outside `ir/` because it derives facts about typed HIR + rather than defining an IR layer. It reads the IR and the `Target` and + decides nothing over what it measures ([architecture §5](./architecture.md#5-analysis--optimization)). - `codegen//` consumes only TIR. The subtree mirrors `ir/tir/`: `prim_function` lives in `tir/`, Stmt emitters in @@ -116,7 +100,7 @@ layout reflects that boundary directly. boundary. `ir/constraints/`, `visitor_registry/`, and `dump/` are cross-cutting packages; -their stable responsibilities are owned by [schedule](./schedule.md), +their stable responsibilities are owned by [parser](./parser.md), [visitor-registry](./visitor-registry.md), and [inspection](./inspection.md), respectively. Their internal file layout is not a per-Op contract. diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index adb6f8e5..03e60df9 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -169,9 +169,9 @@ through its owner chain. - A Module reused as an owned child and as an independently analysed root declares a target only in the latter role. - A Module published as an independently analysable root MUST declare both - its target and the topology levels it is aimed at, so that Analyze and - Schedule answer it as it is published rather than after an edit. Naming the - device is enough; the architecture is derived from that device's document + its target and the topology levels it is aimed at, so that Analyze answers + it as it is published rather than after an edit. Naming the device is + enough; the architecture is derived from that device's document ([target §4](./target.md#4-cudatarget)). - Declaring a Module MUST NOT require a target. Whether a Module becomes an owned child is decided by the owner, after the child has been constructed, diff --git a/docs/spec/inspection.md b/docs/spec/inspection.md index e1a7d25c..7297a220 100644 --- a/docs/spec/inspection.md +++ b/docs/spec/inspection.md @@ -307,7 +307,7 @@ key, its type, and where its value comes from. Metadata with no declaration renders as nothing, which is how an annotation that is not a report -- a binding name, a constraint -- stays off the line. -A record an analysis report projects ([analysis §2](./analysis.md#2-authored-hir-metrics)) +A record an analysis report projects ([analysis §1](./analysis.md#1-authored-hir-metrics)) MUST NOT emit a key that projection cannot state; it MAY emit fewer, because a comment is read on a line and JSON is read by a program. A record no report projects is comment-only -- `SourceSpanMetadata` is where an expression was diff --git a/docs/spec/schedule.md b/docs/spec/schedule.md deleted file mode 100644 index 3e83054c..00000000 --- a/docs/spec/schedule.md +++ /dev/null @@ -1,878 +0,0 @@ -# TileFoundry Spec — Schedule - -Scheduling decides how one Function is placed over one level of the parallel -hierarchy its Module declares. One public operation names the program and the -level; one Target-selected Scheduler answers with a Plan it owns entirely. What an -algorithm decides is its own vocabulary: two algorithms placing different -hardware over different levels do not decide the same things, and a shared result -schema would either describe neither or force both to pretend. - -The top-level `schedule` package exports the operation in [§1](#1-the-public-schedule-operation), -its request/result/base-plan types, and the public errors in [§6](#6-public-errors). -Algorithm-specific plan structures in [§2](#2-public-structures) are public from their owning -submodules and are not necessarily re-exported at package top level. The -construction stages an algorithm composes its solve from -([§4](#4-kernel-schedule-construction)) are imported from their own modules; they -read [analysis](./analysis.md) facts, and the dependency is one-way — nothing in -the analysis layer imports the schedule layer. - -## 1. The public Schedule operation - -```python -def schedule( - module: Module, - function: Function, - *, - topology: str, - options: ScheduleOptions | None = None, - dims: Mapping[str, int] | None = None, -) -> ScheduleResult: ... -``` - -- constraints: - - The caller MUST supply the Module, the Function, and one non-empty topology - level name. Nothing about the request MAY be inferred from layouts, - constraints, or the shape of the program. - - The Function MUST be one the Module owns: one it declares, or a - specialization variant of one it declares - ([core-ir §1](./core-ir.md#1-module)). A Function derived by specialising one - of these MUST be refused, so that ownership is settled before anything is - rebuilt. - - `dims` states one extent per dimension reached through the Function graph, - its Mesh geometry, or the effective Module topology expressions. A solver - places work across a level by counting it and holds a tile against a capacity - in bytes, so a range in any of those positions MUST be solved at a chosen - size rather than as authored. - - `dims=None` MUST behave as a call that states no size: the Function is - solved as authored. - - When `dims` is stated it MUST be non-empty; every key MUST name a dimension - reached through the Function graph, its Mesh geometry, or the effective - Module topology expressions; every value MUST be an integer inside that - dimension's declared bounds; every dimension the Function selects a variant - on MUST be given a value; and no dimension MAY remain a range after - substitution. Each of these MUST fail with a Schedule domain error. A stated - `dims` MUST NOT be silently ignored, including when the Function declares no - range at all. - - Variant resolution and substitution MUST happen after the ownership check - and before the shared program check or algorithm runs. Function types, Mesh - geometry, and effective topology extents MUST use the same resolved binding. - Exactly one variant MUST cover the stated size; none and more than one MUST - both fail. - - The Target MUST come from `module.resolve_target()` - ([core-ir §1](./core-ir.md#1-module)); a call MUST NOT override it and MUST - NOT fall back to a default Target when no Module in the owner chain declares - one ([target §6](./target.md#6-target-ownership-and-compile-resolution)). - - The level MUST be resolved from the Module's own effective hierarchy. A name - the hierarchy does not declare MUST fail, and a level whose extent is known - only at launch MUST fail: an algorithm places work across a level by counting - it. - - Target resolution, level resolution, and algorithm resolution MUST all - complete before the algorithm runs, so a request that cannot be served never - leaves a partial solve behind. - - `options=None` MUST mean a fresh default `ScheduleOptions()` for that call. - - The returned Plan MUST be verified ([§2.3](#23-scheduleplan)) before the result reaches the - caller. - - The common Schedule code MUST NOT import a concrete Target implementation. - -### 1.1 Target-selected Scheduler - -```python -class Scheduler: - topology: str - solve: Callable[ - [Module, Function, Target, Topology, object | None], SchedulePlan - ] - - -class Target: - def get_scheduler(self, topology: str) -> Scheduler: ... -``` - -- constraints: - - The public operation MUST call `get_scheduler` on the exact Target resolved - from the Module only after ownership, specialization, Target, and topology - resolution succeed. - - A Target subclass MUST inherit its base Schedulers through normal Python - inheritance. It MAY override one topology, delegate other requests to - `super()`, or refuse inherited behavior that is invalid for its hardware. - - The Target selects only the immutable descriptor. The public operation MUST - retain orchestration, returned-plan type checking, verification, and result - construction. - - There MUST be no public scheduler registration step, exact-concrete-Target - table, `target.schedule()` wrapper, or automatic level selection. - - An algorithm MUST own its whole problem: its private program view, its Facts - query, its constraint problem, its solve, and its Plan type. Those names MUST - remain private to the algorithm's own package. - -## 2. Public structures - -### 2.1 `ScheduleOptions` - -`ScheduleOptions` carries solver runtime controls, independent of which -algorithm runs. - -```python -class ScheduleOptions: - """Configure one schedule call. - - Attributes: - timeout_seconds: attribute; Wall-clock budget for the underlying solver. - workers: attribute; Solver worker count, where zero selects the solver default. - random_seed: attribute; Deterministic solver tie-break seed. - stop_at_first_solution: attribute; Accept the first result satisfying the - constraints instead of searching the budget for the best one. - debug_dump_dir: attribute; Directory for algorithm-private artifacts, or None. - """ - - timeout_seconds: float = 60.0 - workers: int = 0 - random_seed: int = 0 - stop_at_first_solution: bool = False - debug_dump_dir: Path | None = None -``` - -- constraints: - - The structure MUST be immutable. - - `debug_dump_dir` MUST affect artifact emission only and MUST NOT change the - selected result. - - `stop_at_first_solution` MUST change which satisfying result is selected and - MUST NOT change what counts as one: a result accepted under it MUST satisfy - every constraint a result accepted without it satisfies, so a plan obtained - this way is verifiable on the same terms. - - `stop_at_first_solution` MUST NOT lift `timeout_seconds`. An algorithm that has - found no satisfying result yet stays bounded by it, so the option cannot turn a - bounded search into an unbounded one. - -### 2.2 `ScheduleResult` - -`ScheduleResult` is the complete public result of one call. - -```python -class ScheduleResult: - """Carry what was decided, and what it was decided against. - - Attributes: - module: attribute; The Module the call was made on. - function: attribute; The Function that was scheduled. - topology: attribute; The resolved level of that Module's hierarchy. - plan: attribute; The verified Plan the selected algorithm produced. - """ - - module: Module - function: Function - topology: Topology - plan: SchedulePlan -``` - -- constraints: - - The structure MUST be immutable. - - `module` MUST be the same object the caller supplied. Scheduling decides - about a program; it MUST NOT return a rewritten Module in its place. An - algorithm whose decision *is* a rewritten program MUST carry that program in - its own Plan. - - When the call states no `dims`, `function` MUST be the same object the - caller supplied. - - When the call states `dims`, `function` MUST be the concrete Function the - plan was solved for, derived from the Function the caller supplied. It MUST - record that Function as the one it was specialised from, and the plan MUST - verify against it. A caller returned its own symbolic input would hold a - plan it cannot check. - - `topology` MUST be the level at the resolved geometry. With no `dims` it is - the object the Module declares; with `dims` its extent is the value obtained - from that binding. The plan MUST solve and verify against this same level, - while `module` remains the object the caller supplied. - -### 2.3 `SchedulePlan` - -`SchedulePlan` is the extensible semantic base of every algorithm's result. It is -not a union, a shared schema, or a shared JSON envelope. - -```python -class SchedulePlan: - """One solve's decisions, owned by the algorithm that made them.""" - - def verify(self, module: Module, function: Function, topology: Topology) -> None: ... - - def to_json(self) -> str: ... - - def render(self) -> str: ... -``` - -- constraints: - - The base MUST expose exactly these three operations and MUST impose no - concrete field, shared schema, or common rendering on a subtype. - - The base MUST NOT carry a version field, a deserializer, a renderer registry, - or a generic data-export accessor. A Plan is produced by the algorithm that - solved for it, in the process that solved; reading one back from text would - mean trusting a document to describe decisions nobody made in that run. - - A subtype MUST own the whole of its JSON object and the whole of its human - rendering. - - `verify` MUST be a structural check of the exported plan against the request - it answers: that it refers only to things that exist and that its own - references agree. It MUST NOT re-solve, invoke a solver, or state anything - about whether the schedule is good. - - A Plan that does not hold together MUST raise `PlanVerificationError` and MUST - NOT reach the caller. - -Which hardware documents a decision was made against is the same question -whatever was decided, so every Plan states it the same way. - -```python -class TargetSpecRef: - """Stable identity of the installed target facts one plan relies on. - - Attributes: - architecture_id: attribute; Installed architecture document ID, or the architecture's own name. - architecture_digest: attribute; Content digest of that document, empty when none was installed. - device_id: attribute; Installed device document ID, or the device's own name. - device_digest: attribute; Content digest of that document, empty when none was installed. - """ - - architecture_id: str - architecture_digest: str - device_id: str - device_digest: str - - def of(cls, target: object) -> TargetSpecRef: ... -``` - -- constraints: - - The structure MUST be immutable and MUST be shared by every Plan that names - the hardware it relied on, so two plans cannot describe the same target - differently. - - A Target constructed directly rather than installed from documents MUST state - an empty digest rather than a fabricated one. - - `TargetSpecRef` is public from `tilefoundry.schedule.plan`, not re-exported - from `tilefoundry.schedule`. `of` derives IDs and digests from the supplied - target as a classmethod, falling back to architecture/device names and - empty digests. - -### 2.4 `PipelineSchedulePlan` - -```python -class ScheduledStatement: - """One selected instruction and execution interval. - - Attributes: - id: attribute; stable statement identity. - instruction: attribute; selected instruction name. - tile: attribute; selected tile extents. - resources: attribute; selected resource counts. - start: attribute; inclusive start of the execution interval. - end: attribute; exclusive end of the execution interval. - footprint_bytes: attribute; ring-adjusted buffer footprint. - fits_capacity: attribute; whether the footprint fits tile capacity. - """ - - id: str - instruction: str - tile: tuple[int, ...] - resources: tuple[tuple[str, int], ...] - start: int - end: int - footprint_bytes: int - fits_capacity: bool - -class ScheduledBuffer: - """One storage object and its ring allocation. - - Attributes: - id: attribute; stable buffer identity. - storage: attribute; selected storage name. - ring_depth: attribute; dependency-safe positive ring depth. - producer_ids: attribute; producing statement identities. - consumer_ids: attribute; consuming statement identities. - """ - - id: str - storage: str - ring_depth: int - producer_ids: tuple[str, ...] - consumer_ids: tuple[str, ...] - -class KernelHole: - """One statement reference with serialized boundary relations. - - Attributes: - statement_id: attribute; referenced statement identity. - inputs: attribute; input buffer identities. - outputs: attribute; output buffer identities. - relations: attribute; serialized ISL boundary relations. - """ - - statement_id: str - inputs: tuple[str, ...] - outputs: tuple[str, ...] - relations: tuple[str, ...] - -class PipelineSchedulePlan(SchedulePlan): - """Export one closed pipeline schedule.""" - - target: TargetSpecRef - scaffold: str - statements: tuple[ScheduledStatement, ...] - buffers: tuple[ScheduledBuffer, ...] - holes: tuple[KernelHole, ...] -``` - -- constraints: - - `target` MUST record architecture and device IDs with their content digests. - - Each `ScheduledStatement` MUST carry one stable ID, selected instruction, - tile, resource assignment, a half-open interval, the bytes it holds, and - whether the level's tile store holds them. - - `ScheduledStatement.footprint_bytes` MUST count each buffer the statement - touches at the ring depth that buffer was given, so it states what the - statement occupies once the pipeline is deep enough to run. - - `ScheduledStatement.fits_capacity` MUST record `footprint_bytes` against the - tile capacity the level states. A statement that does not fit MUST still - appear in the plan: the plan states what the program costs on the target, - and a solve MUST NOT drop or shrink a statement to make a plan fit. - - Each `ScheduledBuffer` MUST carry one stable ID, storage, positive ring - depth, and typed producer and consumer statement IDs. - - `ScheduledBuffer.ring_depth` MUST be derived from the dependence distance - the buffer carries under the extents of each statement holding it, so that - a buffer whose value survives into a later tile is given enough slots to - keep the earlier tile alive. A buffer that carries no dependence MUST be - given one slot. - - Each `KernelHole` MUST reference one stable statement ID and expose tuple - inputs, tuple outputs, and serialized ISL relations. It MUST NOT expose an - opaque HIR operation reference. - - JSON and text rendering MUST be deterministic views of the same decisions. - - Plan construction MUST finish Target Facts projection before creating the - closed problem. The problem and solve MUST hold no Target object or callback. - -### 2.5 `PartitionSchedulePlan` - -A partition plan states where each value was placed, which operations run over -which placements, and what the solve proved. It names both by identities derived -from the authored program rather than by the indexes its own problem allocated, so -an agent reading the plan can find what it refers to in the program it wrote. - -```python -class PositionInterval: - """The half-open range of parallel positions something occupies. - - Attributes: - start: attribute; First position occupied. - end: attribute; One past the last position occupied. - """ - - start: int - end: int - -class TimeInterval: - """One operation's half-open execution interval. - - Attributes: - start_ns: attribute; Start of the interval, in ns. - end_ns: attribute; One past the end of the interval, in ns. - """ - - start_ns: int - end_ns: int - -class PlacedValue: - """One tensor value, the type it was placed in, and who touches it. - - Attributes: - id: attribute; Stable identity of this placement, derived from the program. - type: attribute; The ordinary IR Type selected for it, carrying its layout and storage. - producer_id: attribute; The operation that produces it, or None when the plan produces none. - consumer_ids: attribute; Every operation that reads it. - positions: attribute; The positions this placement occupies. - """ - - id: str - type: Type - producer_id: str | None - consumer_ids: tuple[str, ...] - positions: PositionInterval - -class PartitionedOperation: - """One operation that runs, where it runs, and when. - - Attributes: - id: attribute; Stable identity of this operation, derived from the program. - operation: attribute; The operation's own kind. - synthesized: attribute; True when the algorithm introduced it rather than the author. - input_ids: attribute; The placements it reads. - output_ids: attribute; The placements it produces. - positions: attribute; The positions it occupies, or None when it occupies none. - interval: attribute; Its execution interval, or None when the model gives it none. - """ - - id: str - operation: str - synthesized: bool - input_ids: tuple[str, ...] - output_ids: tuple[str, ...] - positions: PositionInterval | None - interval: TimeInterval | None - -class PartitionProof: - """What the solve proved about its own objective. - - Attributes: - status: attribute; Whether optimality was proven or only feasibility reached. - objective_ns: attribute; The selected makespan, in ns. - best_bound_ns: attribute; The bound the solve established, in ns. - proven_optimal: attribute; True when the two met. - """ - - status: Literal["OPTIMAL", "FEASIBLE_NOT_PROVEN"] - objective_ns: int - best_bound_ns: int - proven_optimal: bool - -class PartitionSchedulePlan(SchedulePlan): - """The placement one partition solve committed to, and its proof.""" - - topology: str - extent: int - target: TargetSpecRef - values: tuple[PlacedValue, ...] - operations: tuple[PartitionedOperation, ...] - root_results: tuple[str, ...] - proof: PartitionProof -``` - -- constraints: - - Every structure MUST be immutable, and the plan MUST state the level it - decided about, how many positions of it there were, and the identity of the - installed documents it decided against. - - An identity MUST be derived from the authored program and MUST be unique - within the plan. One value MAY hold more than one placement at once -- that is - what a Reshard connects -- so a placement, not a value, is what an identity - names. - - `PlacedValue.type` MUST be the ordinary IR Type that was selected, so the - layout and storage a value was placed in are read off the type system rather - than restated in a parallel vocabulary. - - An operation the algorithm synthesized MUST appear among `operations`, marked - as synthesized, with the placements it moves between. There MUST be no - separate route, report, or debug channel through which a caller would learn - that data moves. - - An operation charged as traffic rather than as occupancy MUST state no - position range. Giving it one would claim it excludes other work from those - positions, which the decision does not. - - `proof` MUST state the objective, the bound the solve established, and whether - the two met. It is a result fact and MUST NOT become a generic report facade. - - An edge MUST be named the same way from both of its ends: a placement's - producer MUST list that placement among its outputs, its consumers MUST list - it among their inputs, and every operation's outputs and inputs MUST be - reflected back by the placements they name. An edge only one end claims is a - claim about a decision nobody made, and a walk that followed it would report a - program flow the operations do not implement. - - `verify` MUST reject: a level or extent other than the one the plan decided; - two placements or two operations sharing an identity; a reference to a - placement or operation the plan does not carry; an edge whose two ends - disagree; a placement in a type that is not addressable global memory; a synthesized move between two placements that - are different logical tensors or that are identical; a position range outside - the level; an interval that ends before it starts; two operations holding the - same position at the same time; a root result the plan does not carry or - cannot reach by following producers; and a bound above the stated objective. - It MUST do all of that without rebuilding candidates and without invoking a - solver. - - JSON rendering MUST be deterministic and MUST state the placements, - operations, intervals, and proof. - - A solver variable or solver-native value MUST NOT appear in the exported plan. - - The plan MUST NOT carry a rewritten program, and producing it MUST NOT rewrite - one: a partition decides where work and its tensors go, and applying that - decision to HIR is a separate operation the caller asks for. - -## 3. Constraint metadata - -Hard schedule constraints are represented by one stage-neutral -`ScheduleConstraintMetadata` record attached to the constrained HIR -expression. The record contains zero or one `LayoutConstraint`, -`MeshConstraint`, and `StorageConstraint` value, represented by the existing -constraint base and source-location fields. - -```python -class LayoutConstraint(ScheduleConstraint): - """Fix a physical layout pattern and shard bindings. - - Attributes: - layout: attribute; required physical layout pattern. - bindings: attribute; named shard-attribute bindings. - """ - - layout: Layout = Layout(shape=()) - bindings: tuple[tuple[str, ShardAttr], ...] = () - -class MeshConstraint(ScheduleConstraint): - """Filter an eventual shard layout by one mesh. - - Attributes: - mesh: attribute; required mesh value. - """ - - mesh: Mesh | None = None - -class StorageConstraint(ScheduleConstraint): - """Filter a value by one storage kind. - - Attributes: - storage: attribute; required storage kind. - """ - - storage: StorageKind | None = None -``` - -- constraints: - - `LayoutConstraint.layout` is constraint-owned and MAY contain the private - wildcard sentinel. Its `bindings` reuse `Split`, `Broadcast`, and `Partial` - from [shard](./shard.md). - - A wildcard MUST NOT be stored as `Layout(None)` and MUST NOT enter a - `TensorType.layout`. - - Constraint metadata MUST NOT participate in expression equality, hashing, - or the printed `repr`. - - These values are hard filters for later scheduling stages. They carry no - preferences, candidate rows, costs, solver state, or CTA capability - decisions, and they do not register a scheduling algorithm for a - `CudaTarget`. - -## 4. Kernel schedule construction - -An algorithm composes its solve from these stages over one immutable -`TileGraph` ([analysis §1.2](./analysis.md#12-tilegraph)). Scheduling state is -held separately; no stage mutates or enriches the analysis graph, and none -re-derives a fact the analysis layer already states. - -```text -extract(root) ──▶ graph ──▶ build_schedule_tree(graph) ──▶ tree - graph + tree + solved ring ──▶ emit_scaffold - │ - ▼ - Skeleton / Swimlane / HoleContract -``` - -Each stage lives in its own module of the `schedule` package and is imported -from there; the compact public package surface ([§1](#1-the-public-schedule-operation)–[§2](#2-public-structures)) carries the public -operation and its results only. - -| Stage | Signature | Error | -|---|---|---| -| tree construction | `build_schedule_tree(tg: TileGraph) -> isl.schedule` | `KernelScheduleError` | -| scaffold emission | `emit_scaffold(graph: TileGraph, tree: isl.schedule, ring: dict[str, int]) -> tuple[Skeleton, Swimlane, list[HoleContract]]` | `EmitScaffoldError` | - -### 4.1 Tree construction - -```python -def build_schedule_tree(tg: TileGraph) -> "isl.schedule": ... - -def schedule_bands(tree: "isl.schedule") -> tuple["isl.schedule_node_band", ...]: ... - -def band_statement(band: "isl.schedule_node_band") -> str: ... - -def tile_band(band: "isl.schedule_node_band", sizes: tuple[int, ...]) -> "isl.schedule": ... - -def tile_bands(tree: "isl.schedule", sizes: dict[str, tuple[int, ...]]) -> "isl.schedule": ... - -class KernelScheduleError(RuntimeError): - """A schedule tree the band operations cannot work on.""" -``` - -- constraints: - - Nothing in this stage solves. The tree MUST be **constructed** from the - statement order the analysis layer already reports: one identity band per - statement, sequenced in `tg.units` order. That order respects every - dependence, so the result is legal by construction. - - An affine scheduling solve MUST NOT be introduced here, and no objective MAY - be smuggled in from isl's own schedule constraints: its - dependence-distance goal is not the one this layer decides for. - - The statements MUST NOT be fused into one band: their ranks differ, so one - padded shared band member would mean a different loop in each of them. - - Each band's `coincident` members MUST be written from - `tg.parallel_dims` ([analysis §1.7](./analysis.md#17-parallel-dimensions)) — - the flags are read, never recomputed. - - `build_schedule_tree` MUST return a new ISL schedule and MUST NOT mutate - `tg`. An empty `tg.units`, or a unit with no matching piece of `tg.domain`, - MUST raise `KernelScheduleError`. - - `schedule_bands` MUST return every band of the tree in top-down order, which - for a constructed tree is `tg.units` order, and MUST raise when the tree - carries no band. - - `band_statement` MUST raise unless the band belongs to exactly one - statement. - - `tile_band` MUST split one band into a tile band over `sizes` plus a point - band holding the remainder. A size count that does not match the band's - member count, or a size below `1`, MUST raise. - - `tile_bands` MUST tile every band by its own statement's sizes and MUST - raise for a statement with no decided size. - -### 4.2 Scaffold emission - -`emit_scaffold` renders the decided tree into what an authoring agent fills: a -holed loop nest, a human-readable swimlane, and one hole contract per statement. - -```python -class Skeleton: - """A holed, C-like loop-nest skeleton. - - Attributes: - text: attribute; The generated loop nest, with one hole call per statement instance. - holes: attribute; Every hole name in text, in first-appearance order. - """ - - text: str - holes: tuple[str, ...] - -class Swimlane: - """A human-readable rendering of the decided schedule. - - Attributes: - text: attribute; One Mermaid gantt section per statement, minimally unrolled. - """ - - text: str - -class BufferAccess: - """One buffer touched by one statement. - - Attributes: - tensor_name: attribute; Buffer tuple name, as the TileGraph names it. - index_map: attribute; Access map from this statement's coordinates to that buffer's elements. - dtype: attribute; Recovered HIR element DType, or None when it could not be resolved. - """ - - tensor_name: str - index_map: "isl.map" - dtype: object | None - -class HoleContract: - """What one hole must compute. - - Attributes: - name: attribute; The hole's own call name in the skeleton. - op_ref: attribute; The HIR Call this hole stands for. - inputs: attribute; Every buffer the statement reads, in source-call argument order. - output: attribute; The single buffer the statement writes. - coords: attribute; The schedule coordinates the hole is parametrised by. - """ - - name: str - op_ref: object - inputs: tuple[BufferAccess, ...] - output: BufferAccess - coords: tuple[str, ...] - -def emit_scaffold( - graph: TileGraph, - tree: "isl.schedule", - ring: dict[str, int], -) -> tuple[Skeleton, Swimlane, list[HoleContract]]: ... - -class EmitScaffoldError(RuntimeError): - """A construct emit_scaffold does not render, or a TileGraph precondition that did not hold.""" -``` - -- constraints: - - Every structure MUST be immutable. - - The skeleton MUST be ISL code generation over `tree`, with each naked - statement call replaced by its hole call. `graph`, `tree`, and `ring` are - separate inputs; the function MUST NOT read scheduling state from or write - scheduling state to the `TileGraph`. - - A hole call MUST name its inputs, its output, and its raw schedule - coordinates, each behind its own marker, so the three groups are - distinguishable without re-deriving them. - - A read-modify-write self-read on the output buffer MUST appear among the - inputs rather than be silently dropped. - - A buffer whose `ring` depth is above `1` MUST be referenced through that - ring, indexed by the innermost coordinate modulo the depth. An empty `ring` - mapping makes every reference the bare buffer name. - - Exactly one `HoleContract` MUST be produced per statement, not per call site - in the generated text, and its `coords` MUST come from the first occurrence. - - A statement whose name has no matching `TileUnit`, or that writes more than - one buffer, MUST raise `EmitScaffoldError`. - - A hole whose statement call cannot be placed in the generated text MUST raise - rather than be dropped. - - `HoleContract` MUST be a pure function contract — inputs, output, coordinates - — and MUST NOT carry indexing or synchronization: the skeleton already - carries those. `op_ref` MUST be the HIR `Call`, so a later stage can fill the - hole and diff it against the [evaluator](./evaluator.md)'s own result for - that op subgraph. - - The swimlane MUST be minimally unrolled — a prologue instance, a handful of - steady-state instances, and an epilogue instance, with the elided count - stated — never the full iteration count: a real kernel's domain runs to - hundreds of millions of points. - -## 5. Scheduling facts - -The polyhedral model is target-independent; the atom catalogue, the rates work is -charged at, and the store a tile lives in are not. Each algorithm family declares -the facts it needs as its own aggregate, so what one family asks for cannot become -a shared vocabulary another family has to satisfy. All of them are obtained by -projecting the Target ([target §11](./target.md#11-target-facts-projection)), so -an algorithm names the facts it needs and never calls into a target through an -object whose shape it must know. - -The level an algorithm is asked about and the level whose store bounds it need not -be the same one, and a projection MUST NOT collapse them. An AMX core both runs -the work and owns the L1d its tile lives in. A CUDA pipeline is asked about -`thread`, because what it decides is how the threads of one CTA overlap their -work, but the store they cooperate in is shared memory, which is a CTA-scoped -resource: reporting that capacity as a per-thread number would claim a limit no -hardware publishes. - -### 5.1 `AtomFact` - -```python -class AtomFact: - """One candidate atom's facts, as the deciding stage consumes them. - - Attributes: - shape: attribute; The atom's own M, N and K extents. - dtype: attribute; The atom's own a, b and c operand DTypes. - duration: attribute; Nominal roofline estimate for one instance, in ns. - compute_duration: attribute; The compute-side half of that estimate alone, in ns. - storage: attribute; Per-role fragment occupancy in bytes. - resource: attribute; Required thread-scope footprint, keyed by scope name. - is_async: attribute; True when the instruction is asynchronous. - atom: attribute; The target's own realized atom descriptor, carried through opaquely. - """ - - shape: tuple[int, int, int] - dtype: tuple[DType, DType, DType] - duration: float - compute_duration: float - storage: dict[str, int] - resource: dict[str, int] - is_async: bool - atom: object -``` - -- constraints: - - The structure MUST be immutable and MUST stay target-independent: `atom` - MUST be kept opaque, so a target package can enumerate its own catalogue - without this type knowing that catalogue's types. - - `shape` / `dtype` MUST mirror the atom's own shape and operand dtypes, so a - consumer can filter and granularise without unpacking `atom`. - - `duration` MUST be a nominal estimate in ns for **one** atom instance, and - `compute_duration` MUST be its compute-side half alone — for a consumer that - models the surrounding traffic itself and would otherwise charge memory - twice. - - `atom` MUST be the realized descriptor a later fill or codegen stage needs, - so that stage never re-resolves it from `shape` / `dtype`. - -### 5.2 `PartitionFacts` - -```python -class PartitionFactsQuery: - """The one topology level a projection is asked to describe. - - Attributes: - topology: attribute; The level being divided. - """ - - topology: str - -class PartitionFacts: - """All concrete hardware information required to close one partition. - - Attributes: - topology: attribute; The level being divided. - spec: attribute; Identity of the installed documents these numbers came from. - parallel_units: attribute; How many positions of that level the plan may occupy. - memory_bandwidth_bytes_per_second: attribute; The rate traffic is charged at. - memory_capacity_bytes: attribute; The capacity resident bytes are charged against. - peak_flops_per_second: attribute; Dense peak rate per compute DType. - """ - - topology: str - spec: TargetSpecRef - parallel_units: int - memory_bandwidth_bytes_per_second: int - memory_capacity_bytes: int - peak_flops_per_second: tuple[tuple[DType, int], ...] - - def peak_flops(self, dtype: DType) -> int: ... -``` - -- constraints: - - The structure MUST be immutable and MUST contain every numerical fact the - closed problem and its solve consume. After it is projected, neither the - problem nor the solve MAY hold a Target, follow one through the program, or - invoke a projection again. - - A capacity MUST be stated once here rather than copied onto each candidate: a - candidate states the demand it makes, and what that demand is compared against - belongs to the hardware. - - A DType the hardware publishes no rate for MUST fail rather than resolve to - zero or to a neighbouring rate. Charging work at a rate no document supports - would put an unsupported number in the plan. - - A level the target does not divide MUST be reported as that, and the algorithm - MUST surface it as its own scheduling diagnostic rather than let a projection - failure escape. - - Compiler policy MUST stay in `ScheduleOptions`, not in these facts: what the - hardware is does not depend on how aggressively the compiler was asked to - schedule it. - -### 5.3 `PipelineFacts` - -```python -class PipelineFactsQuery: - """Facts requested for one pipeline projection. - - Attributes: - topology: attribute; topology level the pipeline is scheduled over. - statements: attribute; stable statement IDs paired with their ops. - """ - - topology: str - statements: tuple[tuple[str, object], ...] - -class PipelineInstructionFacts: - """Instruction choices for one statement. - - Attributes: - statement_id: attribute; stable statement identity. - candidates: attribute; supported atom choices for that statement. - """ - - statement_id: str - candidates: tuple[AtomFact, ...] - -class PipelineFacts: - """All concrete information required to close a pipeline problem. - - Attributes: - topology: attribute; level the pipeline is scheduled over. - tile_capacity_scope: attribute; level that owns the bounded tile store. - tile_capacity_bytes: attribute; capacity of that tile store in bytes. - max_threads_per_warp: attribute; target warp-width limit. - instructions: attribute; instruction choices by statement identity. - """ - - topology: str - tile_capacity_scope: str - tile_capacity_bytes: int - max_threads_per_warp: int - instructions: tuple[PipelineInstructionFacts, ...] -``` - -- constraints: - - Every structure MUST be immutable. - - `PipelineFactsQuery.statements` MUST preserve program statement order and - MUST use the stable IDs later used by the plan. - - `topology` and `tile_capacity_scope` MUST remain distinct: the level whose - threads are scheduled need not own the capacity that bounds their shared - tile. - - `instructions` MUST contain one entry per requested statement. Each atom is - carried through `AtomFact` and remains opaque to the common scheduler. - - Like `PartitionFacts`, this aggregate is projected once with - `Target.get_facts`; the closed problem MUST NOT retain a `Target` callback. - -## 6. Public errors - -```python -class ScheduleError(ValueError): - """A scheduling request that cannot be served, or a solve that failed.""" - -class PlanVerificationError(ValueError): - """A plan that refers to missing state or contradicts itself.""" -``` - -- constraints: - - Both errors are public from `tilefoundry.schedule`. - - Request validation, algorithm resolution, and solve failures MUST surface as - `ScheduleError`. - - `SchedulePlan.verify` failures MUST surface as `PlanVerificationError`, so a - caller can distinguish an unserviceable request from a malformed result. diff --git a/docs/spec/semantic-analysis.md b/docs/spec/semantic-analysis.md index e9fd6519..100519fe 100644 --- a/docs/spec/semantic-analysis.md +++ b/docs/spec/semantic-analysis.md @@ -65,8 +65,6 @@ behavior, shard propagation, dependence and movement: per boundary, an affine access map from the Op's own iteration space to a tensor's index space. Its carrier `AccessRelations` and the registry that produces it are both defined in [visitor-registry §4.1](./visitor-registry.md#41-access-relation-service--access_relation). -The polyhedral model that lifts a whole `Function` body out of these per-op -relations is owned by [analysis §1](./analysis.md#1-polyhedral-model). The rule reads only the access maps' affine structure (which domain dim each axis uses), never the domain bounds, so it is size-agnostic and identical for static and dynamic shapes. diff --git a/docs/spec/target.md b/docs/spec/target.md index dfecb20f..b19e7ef2 100644 --- a/docs/spec/target.md +++ b/docs/spec/target.md @@ -26,7 +26,6 @@ class Target: def available(cls) -> tuple[Target, ...]: ... def get_analyzer(self, selector: str) -> Analyzer: ... - def get_scheduler(self, topology: str) -> Scheduler: ... def get_code_generator(self) -> CodeGenerator: ... def validate_program_topology(self, topology: Topology) -> None: ... def get_facts( @@ -88,13 +87,13 @@ def registered_targets() -> Mapping[str, type[Target]]: ... and any resolved static extent that is not positive or exceeds the finite `TopologyLimitFacts` bound for that level. The shared program check MUST use this method rather than reproduce a backend's topology limits. - - A provider MAY import `Analyzer` and `Scheduler` from `tilefoundry.target` - to construct getter results. That package MUST NOT expose `CodeGenerator` + - A provider MAY import `Analyzer` from `tilefoundry.target` to construct + getter results. That package MUST NOT expose `CodeGenerator` or `LinkableModule` as provider API. - A missing getter capability MUST fail and name the concrete Target class, its registration name, and the requested selector, topology, or Facts type. - Target values MUST NOT own code emission, linking, loading, or the public - Analyze, Schedule, compile, build, or jit orchestration. + Analyze, compile, build, or jit orchestration. ### 1.1 `Architecture` @@ -175,7 +174,6 @@ class CudaTarget(Target): def validate_program_topology(self, topology: Topology) -> None: ... def topology_limit(self, name: str) -> int: ... def get_analyzer(self, selector: str) -> Analyzer: ... - def get_scheduler(self, topology: str) -> Scheduler: ... def get_code_generator(self) -> CodeGenerator: ... def get_facts(self, facts_type: type[FactsT], query=None) -> FactsT: ... def __repr__(self) -> str: ... @@ -209,11 +207,6 @@ class CudaTarget(Target): - `CudaTarget.available()` MUST contain one value per device document whose sole compatible architecture document is available. Its `identity` MUST be that device document's ID. - - CUDA MUST select the pipeline Scheduler at `thread` and the partition - Scheduler at `cta` through `get_scheduler`. A CUDA subclass MUST inherit - those services through ordinary Python inheritance unless it overrides or - refuses one. The algorithms and their Plan types are not part of the public - `schedule` package. - CUDA MUST select its standard Analyzer, Facts, and CodeGenerator services through the corresponding getters. These selections MUST NOT use `Target.name`, an exact-concrete-type table, or a second extension @@ -237,21 +230,12 @@ class CudaTarget(Target): override `__repr__`. - The store the threads of one CTA cooperate in MUST be projected as `architecture.shared_memory_per_cta_bytes`, and MUST be reported as belonging - to the `cta` scope even when the level being scheduled is `thread` - ([schedule §5](./schedule.md#5-scheduling-facts)). + to the `cta` scope even when the level being asked about is `thread`. - The tensor-memory level MUST be projected with `architecture.tensor_memory_per_cta_bytes` only where the architecture states a capacity, and MUST be absent from `explicit_levels` where it states `None`. A level on hardware that has no such store would offer a plan somewhere to hold accumulators that does not exist. - - The partition projection MUST state the device's SM count as the parallel - units, its HBM bandwidth and capacity, and its dense peak rate per DType - ([schedule §5.2](./schedule.md#52-partitionfacts)). Every one of those MUST be - a hardware fact as the installed documents state it. How much of the machine an - algorithm chooses to occupy is a compiler policy and belongs in - `ScheduleOptions` ([schedule §2.1](./schedule.md#21-scheduleoptions)); it MUST - NOT be projected here, because a Facts value that already encodes a policy - cannot be read as what the hardware is. #### Topology levels @@ -271,8 +255,6 @@ thread mesh layouts. `architecture.max_threads_per_cta`. - `Topology.size` MUST be an explicit `ShapeDim`; construction with `None` MUST fail for every topology level. - - An unresolved symbolic topology extent MUST NOT be scheduled, because - scheduling requires a static extent. - Static declared topology extents MUST be positive integers within their target resource limits. - Unsupported topology levels MUST fail at the generic lowering boundary. @@ -484,15 +466,14 @@ class CpuTarget(Target): - A `Target` belongs to a `Module` rather than an authored HIR `Function`. Target inheritance and its declaration rules are defined by [core-ir `target-inheritance`](./core-ir.md#target-inheritance). -- Analyze and Schedule MUST obtain the Target from `Module.resolve_target()` - and from nowhere else. Neither accepts a bare `Function`, and neither - resolves an undeclared Target to a default: both report hardware-dependent - results, so measuring or scheduling against a device the author never - declared is a silent wrong answer. In particular neither reads a Target out - of `Module.metadata`; the `metadata["target"]` the compile pipeline carries - is the codegen boundary's own record - ([passes §6](./passes.md#6-top-level-api)), not a Target source - for Analyze or Schedule. +- Analyze MUST obtain the Target from `Module.resolve_target()` and from + nowhere else. It does not accept a bare `Function`, and it does not resolve + an undeclared Target to a default: it reports hardware-dependent results, so + measuring against a device the author never declared is a silent wrong + answer. In particular it does not read a Target out of `Module.metadata`; the + `metadata["target"]` the compile pipeline carries is the codegen boundary's + own record ([passes §6](./passes.md#6-top-level-api)), not a Target source + for Analyze. - The compile boundary MAY resolve an omitted Module Target to `default_target()` for lowering, because `jit(fn)` on a plain Function is a documented entry point ([runtime §1.3](./runtime.md#13-jit-api)). It MUST @@ -611,7 +592,6 @@ class AmxTarget(Target): def validate_program_topology(self, topology: Topology) -> None: ... def get_analyzer(self, selector: str) -> Analyzer: ... - def get_scheduler(self, topology: str) -> Scheduler: ... def get_facts(self, facts_type: type[FactsT], query=None) -> FactsT: ... ``` @@ -635,25 +615,6 @@ class AmxTarget(Target): MUST NOT be admitted at either level. - Unsupported topology levels MUST raise an actionable error naming the supported levels, from both the limit lookup and topology validation. - - AMX MUST select exactly one Scheduler, for the `core` level, through - `get_scheduler`. A core both runs the work and owns the store its tile lives - in, so the level asked about and the capacity's scope are the same one. The - `amx` level issues one atom at a time, so there is nothing to place across it - and no Scheduler for it. The algorithm and its Plan type are not part of the - public `schedule` package. - - The core atom-candidate projection MUST list an op's candidates by hard - filtering the registered catalogue, and MUST NOT rank them. The filter is - shape divisibility, operand DType, operand layout, and the storage level - the atom's operand roles need — the last is what separates a - register-resident atom from one streaming through cache, so an op too wide - for the register files lists only the streaming atom. - - An op that clears no filter MUST report an empty candidate list, which is a - covered op with no usable atom rather than an error. Only an op kind or a - target the bridge does not model at all MUST raise. - - The core-level algorithm MUST decide resources over the schedule tree - extracted from the Module's entry function and report the objective in ns. It - MUST NOT rewrite the program it decided about, and its Plan MUST carry no - program. ## 10. Installed hardware resources @@ -711,9 +672,9 @@ conditions = "No validated number." recorded as `estimated`. `derived` and `estimated` MUST state how in `conditions`. - Compiler policy and a program's Topology extents MUST NOT appear in a - hardware document. They are inputs to scheduling, not immutable hardware - truth: a fixed-wave parallel capacity is a scheduling policy even when its - current value equals a device count. An explicit memory level's `owner` is + hardware document. They are compiler inputs, not immutable hardware truth: + a fixed-wave parallel capacity is a policy even when its current value + equals a device count. An explicit memory level's `owner` is different: it MUST name one topology from the Target's hardware vocabulary, or the reserved word `target` for an allocation shared by the whole device. The document states that ownership directly and MUST NOT encode it as an diff --git a/docs/spec/types.md b/docs/spec/types.md index c0f1fbae..e2a8cae2 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -135,7 +135,7 @@ def local_type_of( rather than assign one of its layout axes to a guessed level ([shard §5](./shard.md#5-mesh)). - The result MUST remain an ordinary IR Type and MUST NOT introduce a - schedule-specific tensor type. + consumer-specific tensor type. - Unresolved layouts and local extents that are not concrete non-negative integers MUST raise at the projection boundary. @@ -713,4 +713,4 @@ TensorType.umat_tensor(shape, dtype) # ranked: a shape vector equal to `umat_tensor`. - An operand carrying `UMAT` MUST NOT be charged to a memory level by the residency of its own type alone; what charges it is where it is consumed - ([analysis §2.2.1](./analysis.md#221-compute-cost)). + ([analysis §1.2.1](./analysis.md#121-compute-cost)). diff --git a/docs/spec/visitor-registry.md b/docs/spec/visitor-registry.md index 1aa27440..07f0cd3c 100644 --- a/docs/spec/visitor-registry.md +++ b/docs/spec/visitor-registry.md @@ -286,10 +286,10 @@ consumers run. The visitor and its context share the current scope's memo; ### 4.1 Access relation service — `access_relation` One registry over the Op classes says where each Op reads and writes, and every -reader asks it. Typeinfer asks it to derive the result's Type, the polyhedral -model asks it for dependences, and the movement family asks it how much crossed -each boundary. There is no second registry and no fallback: a boundary nobody -can price is a boundary nobody can schedule. +reader asks it. Typeinfer asks it to derive the result's Type, the loop +footprint asks it for the bytes one authored loop touches, and the movement +family asks it how much crossed each boundary. There is no second registry and +no fallback: a boundary nobody can price is a boundary nobody can measure. ```python class AffineAccess: diff --git a/docs/tutorial/index.md b/docs/tutorial/index.md index bff7ae17..5880d306 100644 --- a/docs/tutorial/index.md +++ b/docs/tutorial/index.md @@ -6,9 +6,9 @@ as `ConstTensor` parameters. That is the *reference*. It is finished when it agr with the published implementation on real weights. **Step two — make it fast.** The authored HIR stays the reference. Write a runtime -twin beside it and `check` the two against each other. `analyze` and `schedule` -report what the program costs and what a plan for it looks like; you decide what to -do with that, including changing the authored HIR. +twin beside it and `check` the two against each other. `analyze` reports what the +program costs; you decide what to do with that, including changing the authored +HIR. TileFoundry is **source to source**. The reference is source code, the fast implementation is source code, and either can be pointed at any command. @@ -20,8 +20,6 @@ implementation is source code, and either can be pointed at any command. - `tilefoundry check --help` — the comparison predicates and their bounds, with the arithmetic for choosing a tolerance. - `tilefoundry analyze --help` — flops, traffic, roofline bounds, a predicted time. -- `tilefoundry schedule --help` — placement, resharding and timing for one topology - level. - `tilefoundry models` — the models already described, and their authored source to copy from. - `tilefoundry tutorial orchestrator causal_lm` — the shipped autoregressive decode diff --git a/docs/tutorial/optimize.md b/docs/tutorial/optimize.md index cf3558b9..e999f2c1 100644 --- a/docs/tutorial/optimize.md +++ b/docs/tutorial/optimize.md @@ -14,11 +14,9 @@ The authored HIR stays the reference. Everything here is written beside it. choosing a tolerance. 3. **Ask for evidence when you want it.** `analyze` reports what the authored - program costs: flops, traffic, roofline bounds, a predicted time. `schedule` proposes - a plan for one topology level you name: placement, resharding, timing. Both read - the authored source, both are optional, and neither decides anything. - `tilefoundry analyze --help` and `tilefoundry schedule --help` say what each - reports. + program costs: flops, traffic, roofline bounds, a predicted time. It reads the + authored source, it is optional, and it decides nothing. + `tilefoundry analyze --help` says what it reports. 4. **Change the authored HIR when that is the answer.** Fusion is done by writing the fused form: two equations in one `@func`, two `@func`s merged into one, or @@ -47,11 +45,10 @@ can read: tilefoundry check mine/model.py:MyFused.fused --inputs random --out output \ --fn allclose --atol 1e-6 --rtol 1e-6 tilefoundry analyze mine/model.py:MyFused --roofline -tilefoundry schedule mine/model.py:MyFused --topology cta --first-plan ``` -`analyze` and `schedule` need a Module that reaches a declared target, so name it -from the root down. +`analyze` needs a Module that reaches a declared target, so name it from the root +down. ## What check cannot do diff --git a/src/tilefoundry/analysis/__init__.py b/src/tilefoundry/analysis/__init__.py index f8a746fd..99792735 100644 --- a/src/tilefoundry/analysis/__init__.py +++ b/src/tilefoundry/analysis/__init__.py @@ -33,30 +33,15 @@ from .registry import Analyzer from .api import AnalysisResult, analyze from .check import check_program -from .poly import ( - AccessFootprint, - AxisExtent, - ExtractError, - TileGraph, - TileUnit, - access_footprints, - carried_distances, - extract, - statement_time_dims, - time_extents, -) __all__ = [ - "AccessFootprint", "Analyzer", "AnalysisError", "AnalysisResult", - "AxisExtent", "BufferFootprint", "ComputeCostMetadata", "ExplicitMemoryLevelFacts", - "ExtractError", "ImplicitMemoryLevelFacts", "LevelFootprint", "LoopFootprintMetadata", @@ -72,16 +57,9 @@ "RooflineMetadata", "TrafficMetadata", "ThroughputFacts", - "TileGraph", - "TileUnit", "TimelineMetadata", "TrafficBytes", "ValueLifetime", - "access_footprints", "analyze", - "carried_distances", "check_program", - "extract", - "statement_time_dims", - "time_extents", ] diff --git a/src/tilefoundry/analysis/poly/affine.py b/src/tilefoundry/analysis/affine.py similarity index 100% rename from src/tilefoundry/analysis/poly/affine.py rename to src/tilefoundry/analysis/affine.py diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index aea9d657..f3cfe2c2 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -1,9 +1,8 @@ -"""The shared authored-program gate for analysis and scheduling. +"""The shared authored-program gate for analysis. An analysis reads inferred types and assumes the authored program holds together. Both conditions are established once per public call rather than -per algorithm, so no family can be the one that forgot. The same gate is used -by scheduling before it makes placement decisions. +per family, so no family can be the one that forgot. """ from __future__ import annotations @@ -325,7 +324,7 @@ class PerformanceChecker(ExprVisitor[None]): Two things, and nothing about storage: rates it can put on a clock, and a placement for every occurrence that will take time. Where the buffers go is - the solver's question, answered against the schedule it is choosing. + not this question, and nothing here decides it. """ def check_target_facts(self, ctx: AnalysisCheckContext) -> None: @@ -514,7 +513,7 @@ def _require_concrete_geometry( raise error_type( f"program topology level {topology.name!r} still states " f"symbolic extent {topology.size!r}; bind every dimension " - "before analysis or scheduling" + "before analysis" ) _require_concrete_function_geometry(function, error_type=error_type) @@ -527,7 +526,7 @@ def _require_concrete_function_geometry( if not is_concrete(function): raise error_type( f"{function.name!r} still states symbolic dimensions in its reachable " - "Function or Mesh geometry; bind every dimension before analysis or scheduling" + "Function or Mesh geometry; bind every dimension before analysis" ) diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index baaa673f..aca75eb2 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -9,7 +9,7 @@ from tilefoundry.ir.types import TensorType from tilefoundry.visitor_registry.access_relation import index_set -from .poly.affine import LoopAffineTerm +from .affine import LoopAffineTerm class _Unavailable(Exception): diff --git a/src/tilefoundry/analysis/poly/__init__.py b/src/tilefoundry/analysis/poly/__init__.py deleted file mode 100644 index 3e918f52..00000000 --- a/src/tilefoundry/analysis/poly/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Public polyhedral analysis API.""" - -from .access import ( - AccessFootprint, - AxisExtent, - access_footprints, - carried_distances, - statement_time_dims, - time_extents, -) -from .errors import ExtractError -from .extract import extract -from .model import TileGraph, TileUnit - -__all__ = [ - "AccessFootprint", - "AxisExtent", - "ExtractError", - "TileGraph", - "TileUnit", - "access_footprints", - "carried_distances", - "extract", - "statement_time_dims", - "time_extents", -] diff --git a/src/tilefoundry/analysis/poly/access.py b/src/tilefoundry/analysis/poly/access.py deleted file mode 100644 index e8383203..00000000 --- a/src/tilefoundry/analysis/poly/access.py +++ /dev/null @@ -1,184 +0,0 @@ -"""Polyhedral access and time queries.""" - -from __future__ import annotations - -import math -from dataclasses import dataclass - -import isl - -from .errors import ExtractError -from .extract import ( - TileGraph, - _as_map, - _buffers_by_statement, - _only_out_dim, - _static_extent, - _travels_with, -) - - -@dataclass(frozen=True) -class AxisExtent: - """Describe one buffer dimension within a statement's complete access. - - ``extent`` measures reached elements rather than deriving a tile size. - ``axes`` identifies time dimensions that reach them and carries no size; - when empty, every iteration touches the dimension in full. - """ - - axes: tuple[int, ...] - extent: int - - -@dataclass(frozen=True) -class AccessFootprint: - """One access sized per buffer dimension, so the element count is the product over ``dims``. - - One (statement, buffer) access sized per buffer dimension, so the - element count is the product over ``dims``. - - The count is the bounding box of the access's range, which is exact for a - box-shaped access and an upper bound for one that leaves holes in it (a - diagonal ``b[t0 + t1]`` reaches a band inside its own box). - """ - - statement: str - buffer: str - is_read: bool - dims: tuple[AxisExtent, ...] - elem_bytes: int - - -def time_extents(tg: TileGraph, time_map: "isl.union_map") -> tuple[int, ...]: - """Per-dimension extent of ``time_map``'s range over ``tg.domain``. - - Raises unless every dimension starts at 0, since a tile index counts - from the origin. - """ - sets: list["isl.set"] = [] - time_map.intersect_domain(tg.domain).range().foreach_set(sets.append) - if len(sets) != 1: - raise ExtractError( - f"time_extents: expected one time space, got {len(sets)} -- " - "every statement must share the band's own range space" - ) - box = sets[0] - extents = [] - for i in range(box.dim(isl.dim_type.SET)): - lo, hi = _static_extent(box, i, "time_extents") - if lo != 0: - raise ExtractError( - f"time_extents: time dimension {i} starts at {lo}, not 0 -- " - "tile counting assumes an origin-based extent" - ) - extents.append(hi + 1) - return tuple(extents) - - -def statement_time_dims(tg: TileGraph, time_map: "isl.union_map") -> dict[str, tuple[int, ...]]: - """Per statement, one entry per time dimension. - - Per statement, one entry per time dimension: the statement's own - domain dimension that dimension travels with, or ``-1`` when it is - constant there (``RN[d0] -> [d0, 63, 127]`` gives ``(0, -1, -1)``). - Raises on a skewed time dimension, which no per-axis tile size can - describe. - """ - maps: list["isl.map"] = [] - time_map.foreach_map(maps.append) - out: dict[str, tuple[int, ...]] = {} - for m in maps: - name = m.get_tuple_name(isl.dim_type.IN) - row = [] - for pos in range(m.dim(isl.dim_type.OUT)): - involved = _travels_with(m, pos) - if len(involved) > 1: - raise ExtractError( - f"statement_time_dims: time dimension {pos} of statement " - f"{name!r} mixes domain dimensions {involved} ({m}) -- " - "a skewed band has no per-axis tile size" - ) - row.append(involved[0] if involved else -1) - out[name] = tuple(row) - return out - - -def carried_distances( - tg: TileGraph, time_map: "isl.union_map", n_dims: int -) -> dict[str, tuple[int, ...]]: - """Per buffer, the largest dependence distance isl reports along each time dimension. - - Per buffer, the largest dependence distance isl reports along each - time dimension. A flow dependence ``a -> b`` is attributed to every - buffer ``a`` writes and ``b`` reads, which for a RAW must-dependence is - exactly the memory it travels through. - """ - written = _buffers_by_statement(tg.writes) - read = _buffers_by_statement(tg.reads) - names = {buf for bufs in (*written.values(), *read.values()) for buf in bufs} - distances: dict[str, list[int]] = {buf: [0] * n_dims for buf in names} - deps: list["isl.map"] = [] - tg.deps.foreach_map(deps.append) - for dep in deps: - carriers = written.get(dep.get_tuple_name(isl.dim_type.IN), set()) & read.get( - dep.get_tuple_name(isl.dim_type.OUT), set() - ) - if not carriers: - continue - pieces: list["isl.set"] = [] - dep.apply_domain(time_map).apply_range(time_map).deltas().foreach_set(pieces.append) - for piece in pieces: - for i in range(n_dims): - lo, hi = _static_extent(piece, i, "carried_distances") - reach = max(abs(lo), abs(hi)) - for buf in carriers: - distances[buf][i] = max(distances[buf][i], reach) - return {buf: tuple(dims) for buf, dims in distances.items()} - - -def access_footprints(tg: TileGraph, time_map: "isl.union_map") -> tuple[AccessFootprint, ...]: - """Access footprints. - - Every read and write of ``tg``, expressed against ``time_map``'s - range so a tile size per time dimension sizes it (see - :class:`AccessFootprint`). - """ - out: list[AccessFootprint] = [] - for um, is_read in ((tg.reads, True), (tg.writes, False)): - maps: list["isl.map"] = [] - um.foreach_map(maps.append) - for m in maps: - stmt = m.get_tuple_name(isl.dim_type.IN) - buf = m.get_tuple_name(isl.dim_type.OUT) - dtype = tg.buffer_dtypes.get(buf) - if dtype is None: - raise ExtractError( - f"access_footprints: buffer {buf!r} has no recorded dtype " - "-- extract must resolve every accessed buffer's element type" - ) - timed = _as_map(m.apply_domain(time_map)) - dims = [] - for pos in range(timed.dim(isl.dim_type.OUT)): - lo, hi = _static_extent( - _only_out_dim(timed, pos).range(), 0, f"access_footprints[{buf}]" - ) - dims.append( - AxisExtent(axes=_travels_with(timed, pos), extent=hi - lo + 1) - ) - out.append( - AccessFootprint( - statement=stmt, buffer=buf, is_read=is_read, dims=tuple(dims), - elem_bytes=math.ceil(dtype.bit_width / 8), - ) - ) - return tuple(out) - -__all__ = [ - "AccessFootprint", - "AxisExtent", - "access_footprints", - "carried_distances", - "statement_time_dims", - "time_extents", -] diff --git a/src/tilefoundry/analysis/poly/errors.py b/src/tilefoundry/analysis/poly/errors.py deleted file mode 100644 index 256fc146..00000000 --- a/src/tilefoundry/analysis/poly/errors.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Errors raised while extracting polyhedral analysis data.""" - - -class ExtractError(NotImplementedError): - """A construct `extract` does not (yet) support. - - A construct `extract` does not (yet) support -- always raised with - a specific, actionable message; V1 never silently guesses. - """ - - -__all__ = ["ExtractError"] diff --git a/src/tilefoundry/analysis/poly/extract.py b/src/tilefoundry/analysis/poly/extract.py deleted file mode 100644 index 653babfd..00000000 --- a/src/tilefoundry/analysis/poly/extract.py +++ /dev/null @@ -1,1080 +0,0 @@ -"""Extract an element-granularity polyhedral graph from an HIR function. - -Compute calls become statements; structural views fold into consumer access -maps and nested functions are penetrated with call-site-qualified names. Loop -dimensions preserve carried dependencies, and registered relations supply all -access maps without guessed fallbacks. ISL flow analysis derives dependencies; -scheduling consumes the result but owns its own schedule tree and decisions. -""" - -from __future__ import annotations - -import dataclasses -import itertools -from dataclasses import dataclass - -import isl - -from tilefoundry.ir.core import Call, Expr, Tuple, TypeInferContext, Var, binding_name -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.grid_region import GridRegionExpr -from tilefoundry.ir.hir.tensor.full_like import FullLike -from tilefoundry.ir.hir.tensor.index_select import IndexSelect -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.hir.tensor.zeros import Zeros -from tilefoundry.ir.types import TupleType -from tilefoundry.ir.types.dim import DimVar, is_dim_op_call -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.utils import local_type_of -from tilefoundry.ir.visitor import ExprVisitor, collect_exprs, expr_children -from tilefoundry.visitor_registry.access_relation import ( - AccessRelations, - BoundaryRelation, - access_relation_registry, - local_relations_of, - relation_of, - renaming_relation, -) - -from .affine import loop_affine_term -from .errors import ExtractError -from .model import TileGraph, TileUnit - -_STATEMENT_ABBREV = {"MatMul": "MM", "RMSNorm": "RN"} - - -@dataclass(frozen=True) -class _StatementAccess: - """Extraction-internal. - - Extraction-internal: one statement's domain + tuple-named access - maps, before they are unioned into the returned ``TileGraph``. - ``loops`` is the enclosing loop nest, outermost first -- the first - ``len(loops)`` dimensions of ``domain`` are its axes. - """ - - name: str - domain: "isl.set" - op: Call - reads: tuple["isl.map", ...] - writes: tuple["isl.map", ...] - params: dict - dtypes: dict - loops: tuple[GridRegionExpr, ...] = () - - -def _buffer_namer(): - """Assign stable ISL tuple names to SSA buffers during extraction. - - Authored names are preferred and collisions receive numeric suffixes. - Tuple projections and structural views resolve to their source buffers while - preserving coordinate transforms. Prefixes isolate penetrated call sites; - aliases fuse a loop's carried variable with its yielded value so flow - analysis observes a distance-one dependency. - """ - seen: dict[int, str] = {} - used: set[str] = set() - anonymous = itertools.count() - - class _NameVisitor(ExprVisitor[str]): - def __init__(self, prefix: str) -> None: - super().__init__() - self.prefix = prefix - - def _assign(self, expr, base: str | None = None) -> str: - key = id(expr) - cached = seen.get(key) - if cached is not None: - return cached - if base is None: - base = binding_name(expr) or f"t{next(anonymous)}" - base = f"{self.prefix}{base}" - candidate = base - suffix = 0 - while candidate in used: - suffix += 1 - candidate = f"{base}_{suffix}" - used.add(candidate) - seen[key] = candidate - return candidate - - def visit_Call(self, expr: Call, ctx=None) -> str: - if isinstance(expr.target, TupleGetItem): - base = _NameVisitor("").visit(expr.args[0], ctx) - name = f"{base}_{expr.target.index}" - seen[id(expr)] = name - return name - if isinstance(expr.target, (Reshape, IndexSelect, Slice)): - name = _NameVisitor("").visit(expr.args[0], ctx) - seen[id(expr)] = name - return name - return self._assign(expr) - - def visit_Var(self, expr: Var, ctx=None) -> str: - return self._assign(expr, expr.name) - - def default_visit(self, expr, ctx=None) -> str: - return self._assign(expr) - - def name_for(expr, prefix: str = "") -> str: - return _NameVisitor(prefix).visit(expr) - - def alias(phi: Var, yielded) -> None: - """One buffer for a loop carry: ``phi`` and ``yielded`` share a name. - - Chained carries (a nested loop re-carrying the same value) all land - on whichever name the group already has. - """ - name = seen.get(id(yielded)) or seen.get(id(phi)) or name_for(phi) - seen[id(phi)] = name - seen[id(yielded)] = name - - def pierce(m: "isl.map", expr, loops: tuple[GridRegionExpr, ...] = ()) -> "isl.map": - """Pierce. - - ``m`` (a read of ``expr``) rewritten to address ``expr``'s ultimate - buffer, folding each view hop between the two into the map's range. The - coordinates come from the view's own registered relation, so a reshape's - arithmetic and a window's stride are stated once for every reader rather - than rebuilt here per Op. - """ - ctx = _RankPreserving() - while isinstance(expr, Call) and isinstance(expr.target, (Reshape, IndexSelect, Slice)): - try: - folded = renaming_relation(expr, ctx) - m = m.apply_range(relation_of(folded)) - m = _place_parameters(m, BoundaryRelation(folded), loops) - except (NotImplementedError, TypeError, ValueError, isl.Error) as error: - raise ExtractError( - f"extract: {type(expr.target).__name__} cannot state where it " - f"reads what it renames: {error}" - ) from error - expr = expr.args[0] - return m - - name_for.alias = alias - name_for.pierce = pierce - return name_for - - -def _assign_statement_names(ops: list[object]) -> list[str]: - """One isl tuple name per statement, in call order. - - One isl tuple name per statement, in call order: the bare (possibly - abbreviated) op name when it is unique in this extraction, else - suffixed ``NAME0``, ``NAME1``, ... in first-seen order. - """ - bases = [_STATEMENT_ABBREV.get(type(op).__name__, type(op).__name__) for op in ops] - counts: dict[str, int] = {} - for base in bases: - counts[base] = counts.get(base, 0) + 1 - next_index: dict[str, int] = {} - names = [] - for base in bases: - if counts[base] == 1: - names.append(base) - continue - index = next_index.get(base, 0) - names.append(f"{base}{index}") - next_index[base] = index + 1 - return names - - -def _static_loop_bound(dim, what: str) -> int: - if isinstance(dim, int) and not isinstance(dim, bool): - return dim - raise ExtractError( - f"extract: loop {what} {dim!r} is not a static int -- only a loop " - "extent may be dynamic (a bare DimVar, which becomes an isl parameter)" - ) - - -def _loop_domain(inner: "isl.set", loops: tuple[GridRegionExpr, ...]) -> tuple["isl.set", dict]: - """Prefix ``inner`` with outermost-first enclosing loop dimensions. - - Each dimension spans ``[start, extent)`` by ``step`` and carries the raw - induction value so indexed gathers can use it directly. A ``DimVar`` extent - becomes a same-name ISL parameter bound to its declared range. - """ - if not loops: - return inner, {} - rank = inner.dim(isl.dim_type.SET) - params: dict = {} - bounds: list[str] = [] - for j, loop in enumerate(loops): - start = _static_loop_bound(loop.start, "start") - step = _static_loop_bound(loop.step, "step") - if isinstance(loop.extent, DimVar): - params[loop.extent.name] = loop.extent - hi = loop.extent.name - else: - hi = str(_static_loop_bound(loop.extent, "extent")) - bounds.append(f"{start} <= p{j} < {hi}") - if step != 1: - bounds.append(f"(p{j} - {start}) mod {step} = 0") - for name, dim in params.items(): - bounds.append(f"{dim.lo} <= {name} < {dim.hi}") - dims = [f"p{j}" for j in range(len(loops))] + [f"f{i}" for i in range(rank)] - prefix = f"[{', '.join(params)}] -> " if params else "" - box = isl.set(prefix + f"{{ [{', '.join(dims)}] : {' and '.join(bounds)} }}") - return inner.insert_dims(isl.dim_type.SET, 0, len(loops)).intersect(box), params - - -def _renaming(expr) -> bool: - """Whether a value is another name for a buffer, with coordinates to fold.""" - return isinstance(expr, Call) and isinstance(expr.target, (Reshape, IndexSelect, Slice)) - - -def _whole_views(prefixed: "isl.set", read_maps, args: tuple, namer, loops, within) -> "isl.set": - """The trips whose views are whole, read off the folded relations themselves. - - A window placed by the trip runs past its source on the last few, and the - view's own relation already says so: composing it leaves those coordinates - out. Asking the relation is what keeps one window formula in one place. - """ - held = prefixed - for index, arg in enumerate(args): - if index >= len(read_maps) or read_maps[index] is None or not _renaming(arg): - continue - reached = within(read_maps[index]) - if reached.is_empty(): - continue - held = held.intersect(namer.pierce(reached, arg, loops).domain()) - return held - - -def _lift(m: "isl.map", depth: int) -> "isl.map": - """One access map given ``depth`` extra leading input dimensions -- the enclosing loop axes. - - One access map given ``depth`` extra leading input dimensions -- the - enclosing loop axes, which the op's own relation knows nothing about. - """ - return m if not depth else m.insert_dims(isl.dim_type.IN, 0, depth) - - -def _bind_map(m: "isl.map", stmt_name: str, domain: "isl.set", buffer_name: str) -> "isl.map": - """Name and domain-restrict an element-granularity access map. - - Bounds live in ``domain``. Stamp the statement name before intersection so - the map input tuple identity and parameter space match the set; reversing - that order produces incompatible ISL spaces. - """ - return ( - m.set_tuple_name(isl.dim_type.IN, stmt_name) - .intersect_domain(domain) - .set_tuple_name(isl.dim_type.OUT, buffer_name) - ) - - -def _read_map( - m: "isl.map", - stmt_name: str, - domain: "isl.set", - arg, - namer, - loops=(), -) -> "isl.map": - """Read map. - - A read access for input ``arg``, pierced through any reshape / loop-index - selection view (``namer.pierce``) before binding -- the view-fold's landing - point for every op's input side (an op's own output buffer never needs - piercing, a reshape output is never written to). - """ - return _bind_map(namer.pierce(m, arg, loops), stmt_name, domain, namer(arg)) - - -def _out_dtype(call: Call, out_idx: int): - """The dtype of ``call``'s ``out_idx``-th output. - - The dtype of ``call``'s ``out_idx``-th output: a multi-output op - (``RoPE``) types its result as a ``TupleType``, a single-output one as - the ``TensorType`` itself. - """ - ty = call.type - if isinstance(ty, TupleType) and out_idx < len(ty.fields): - return getattr(ty.fields[out_idx], "dtype", None) - return getattr(ty, "dtype", None) - - -@dataclasses.dataclass -class _RankPreserving(TypeInferContext): - """A context whose local view narrows a Split axis and keeps the rank. - - An isl flow model is indexed by a tensor's logical axes, so a local view that - factored them into layout positions would silently lose dependences. Handlers - ask their context this question, which is what lets one registered relation - answer both the whole program and one participant. - """ - - def local_type_of(self, expr) -> object: - return local_type_of(expr.type) - - -def _placed_value(value, loops: tuple[GridRegionExpr, ...]): - """Where in this loop nest a parameter's value sits, when it sits anywhere. - - A shape parameter is a `DimVar` and stays one: isl carries it symbolically - and the schedule reads it back. An offset is an operand, and this nest may be - the thing that moves it. - """ - if isinstance(value, DimVar): - return None - if not isinstance(value, Expr): - return None - term = loop_affine_term(value, loops, narrow=False) - if term is None or term.low != term.high: - return None - return term.loop_axis, term.stride, term.low - - -def _place_parameters(access: "isl.map", boundary, loops: tuple[GridRegionExpr, ...]) -> "isl.map": - """Say what a relation's parameters are, in this loop nest's terms. - - A window's offset is bound to the operand that states it, and that operand - may be an induction variable: then the parameter is the loop coordinate, and - the dependence moves with the trip. A shape parameter stays symbolic, isl - carrying it through to the schedule. One nobody here can place is projected - out, which unions over every value the Op allows -- more accesses than there - are, which is the safe direction for a dependence. - """ - bound = dict(getattr(boundary.pattern, "parameters", ()) or ()) - for name in _named(access): - value = bound.get(name) - position = _named(access).index(name) if name in _named(access) else None - if position is None: - continue - if isinstance(value, DimVar): - access = _under_its_own_name(access, position, value.name) - continue - number = static_dim_value(value) - loop_position = None - stride = 0 - if number is None: - placed = _placed_value(value, loops) - if placed is None: - access = access.project_out(isl.dim_type.PARAM, position, 1) - continue - loop_position, stride, number = placed - local = isl.local_space.from_space(access.get_space()) - equality = isl.constraint.alloc_equality(local).set_coefficient_si( - isl.dim_type.PARAM, position, 1 - ) - if loop_position is not None: - equality = equality.set_coefficient_si(isl.dim_type.IN, loop_position, -stride) - access = access.add_constraint(equality.set_constant_si(-number)) - access = access.project_out(isl.dim_type.PARAM, position, 1) - return access - - -def _under_its_own_name(access: "isl.map", position: int, name: str) -> "isl.map": - """One symbolic parameter renamed to the dimension it stands for. - - A relation names its own parameters, and two relations naming the same - dimension differently do not meet: the schedule would carry two symbols for - one extent. Where the name is already taken it is the same dimension, so the - two are equated and one of them leaves. - """ - held = _named(access) - if held[position] == name: - return access - if name not in held: - return access.set_dim_name(isl.dim_type.PARAM, position, name) - local = isl.local_space.from_space(access.get_space()) - equality = ( - isl.constraint.alloc_equality(local) - .set_coefficient_si(isl.dim_type.PARAM, position, 1) - .set_coefficient_si(isl.dim_type.PARAM, held.index(name), -1) - ) - return access.add_constraint(equality).project_out(isl.dim_type.PARAM, position, 1) - - -def _named(access: "isl.map") -> list: - """The parameters one relation names, in the order isl holds them.""" - return [ - access.get_dim_name(isl.dim_type.PARAM, index) - for index in range(access.dim(isl.dim_type.PARAM)) - ] - - -def _walked(call: Call, relations: AccessRelations, walked: "isl.set") -> tuple["isl.set", dict]: - """One statement's own space, and the values its isl parameters name. - - An access map's domain is the Op's iteration space, so the space a statement - walks is the one its output relation was written over -- a product keeps the - axis it sums, a normalisation is asked once per row -- rather than the shape - of what it produced. - """ - named = { - walked.get_dim_name(isl.dim_type.PARAM, index) - for index in range(walked.dim(isl.dim_type.PARAM)) - } - bound = { - name: value - for boundary in (*relations.outputs, *relations.inputs) - for name, value in (getattr(boundary.pattern, "parameters", ()) or ()) - } - missing = named - set(bound) - if missing: - raise ExtractError( - f"extract: {type(call.target).__name__} walks a space naming " - f"{sorted(missing)}, and nothing here says what those are" - ) - return walked.coalesce(), {name: bound[name] for name in named} - - -def _registered_access( - call: Call, - stmt_name: str, - relations: AccessRelations, - ctx, - namer, - prefix: str, - loops: tuple[GridRegionExpr, ...] = (), -) -> list[_StatementAccess]: - """Statement extraction from the Op's own registered boundary relations. - - Every boundary is asked by the same coordinates, so one space serves them - all -- and a boundary may answer on only part of it. Each output is its own - statement over the part its own relation answers on, which is how one - rotation over two head counts lifts into two, and what a hole writing one - buffer needs. A relation's parameters are placed in this loop nest's terms - first, because a window that moves with the trip carries the dependence that - makes the loop a loop. - """ - depth = len(loops) - - def placed(boundary) -> "isl.map": - return _place_parameters(_lift(relation_of(boundary.pattern), depth), boundary, loops) - - written = tuple(placed(boundary) for boundary in relations.outputs) - if not written: - raise ExtractError( - f"extract: {type(call.target).__name__} states no output boundary; " - "a compute-op statement must write at least one value" - ) - walks = tuple(raw_map.domain().project_out(isl.dim_type.SET, 0, depth) for raw_map in written) - whole = walks[0] - for own in walks[1:]: - whole = whole.union(own) - read_maps = tuple( - placed(relations.inputs[index]) if index < len(relations.inputs) else None - for index in range(len(call.args)) - ) - return [ - _one_statement( - call, - stmt_name if len(written) == 1 else f"{stmt_name}_{index}", - index, - len(written), - walks[index], - _settled_within(whole, walks[index]), - read_maps, - relations, - namer, - prefix, - loops, - ) - for index in range(len(written)) - ] - - -def _settled_within(whole: "isl.set", piece: "isl.set") -> tuple[int, ...]: - """The coordinates one output piece holds still while its Call varies them. - - A Call that rotates two values walks a coordinate saying which it is - rotating, and one output answers at one value of it. That coordinate is not - part of the statement -- the statement is the piece -- so it comes off, and - the rank a reader sees is the rank the work has. An extent of one that the - whole Call also holds still is a real axis and stays. - """ - return tuple( - position - for position in range(piece.dim(isl.dim_type.SET)) - if _held_still(piece, position) and not _held_still(whole, position) - ) - - -def _held_still(walked: "isl.set", position: int) -> bool: - """Whether one coordinate takes a single value everywhere in a space.""" - low, high = walked.dim_min_val(position), walked.dim_max_val(position) - return low.is_int() and high.is_int() and low.get_num_si() == high.get_num_si() - - -def _without(walked: "isl.set", settled: tuple[int, ...]) -> "isl.set": - """One space with the coordinates an output piece holds still taken out.""" - for position in reversed(settled): - walked = walked.project_out(isl.dim_type.SET, position, 1) - return walked - - -def _one_statement( - call: Call, - stmt_name: str, - out_idx: int, - outputs: int, - walks: "isl.set", - settled: tuple[int, ...], - read_maps: tuple, - relations: AccessRelations, - namer, - prefix: str, - loops: tuple[GridRegionExpr, ...], -) -> _StatementAccess: - """One statement: the part of a Call's space one output answers on. - - A boundary that answers nowhere in this piece is not this statement's: the - other value a fused rotation writes is read by the other statement, and - copying it here would invent a dependence on bytes this work never touches. - """ - depth = len(loops) - here = walks.insert_dims(isl.dim_type.SET, 0, depth) - - def within(access: "isl.map") -> "isl.map": - held = access.intersect_domain(here) - for position in reversed(settled): - held = held.project_out(isl.dim_type.IN, depth + position, 1) - return held - - own, shape_params = _walked(call, relations, _without(walks, settled)) - prefixed, loop_params = _loop_domain(own, loops) - prefixed = _whole_views(prefixed, read_maps, call.args, namer, loops, within) - domain = prefixed.set_tuple_name(stmt_name) - - reads: list["isl.map"] = [] - writes: list["isl.map"] = [] - dtypes: dict = {} - for index, arg in enumerate(call.args): - if read_maps[index] is None: - continue - reached = within(read_maps[index]) - if reached.is_empty(): - continue - read = _read_map(reached, stmt_name, domain, arg, namer, loops) - reads.append(read) - dtypes[read.get_tuple_name(isl.dim_type.OUT)] = getattr(arg.type, "dtype", None) - - out_buf = namer(call, prefix) if outputs == 1 else f"{namer(call, prefix)}_{out_idx}" - bound = _bind_map( - within( - _place_parameters( - _lift(relation_of(relations.outputs[out_idx].pattern), depth), - relations.outputs[out_idx], - loops, - ) - ), - stmt_name, - domain, - out_buf, - ) - writes.append(bound) - dtypes[out_buf] = _out_dtype(call, out_idx) - if not bound.is_injective(): - reads.append(bound) - - return _StatementAccess( - name=stmt_name, - domain=domain, - op=call, - reads=tuple(reads), - writes=tuple(writes), - params={**shape_params, **loop_params}, - dtypes=dtypes, - loops=loops, - ) - - -def _extract_statement( - call: Call, - stmt_name: str, - namer, - prefix: str, - loops: tuple[GridRegionExpr, ...] = (), -) -> list[_StatementAccess]: - if access_relation_registry.lookup(type(call.target)) is None: - raise ExtractError( - f"extract: op {type(call.target).__name__!r} has no registered " - "access relation -- register one via tilefoundry.visitor_registry." - "access_relation.register_access_relation(...); extract has no " - "per-op fallback." - ) - ctx = _RankPreserving() - try: - relations = local_relations_of(call, ctx) - except (NotImplementedError, TypeError, ValueError, isl.Error) as error: - raise ExtractError( - f"extract: {type(call.target).__name__} cannot state its boundary " - f"relations here: {error}" - ) from error - return _registered_access(call, stmt_name, relations, ctx, namer, prefix, loops) - - -def _initial_schedule(accesses: list[_StatementAccess]) -> "isl.union_map": - """Build a total order that seeds flow analysis, not the final schedule. - - Coordinates are ``[*loop_dims, stage, *own_dims, 0-pad]``. Loop dimensions - precede the collect_exprs stage so statements interleave per iteration and a - read at ``i + 1`` observes a write at ``i``; placing stage first would lose - loop-carried dependencies. - """ - slots: list[GridRegionExpr] = [] - for acc in accesses: - for loop in acc.loops: - if not any(loop is seen for seen in slots): - slots.append(loop) - own_rank = max((a.domain.dim(isl.dim_type.SET) - len(a.loops) for a in accesses), default=0) - sched = isl.union_map("{}") - for stage, acc in enumerate(accesses): - rank = acc.domain.dim(isl.dim_type.SET) - depth = len(acc.loops) - dims = [f"d{i}" for i in range(rank)] - head = ["0"] * len(slots) - for j, loop in enumerate(acc.loops): - head[next(s for s, x in enumerate(slots) if x is loop)] = dims[j] - tail = dims[depth:] + ["0"] * (own_rank - (rank - depth)) - src = f"[{', '.join(dims)}]" if dims else "[]" - dst = f"[{', '.join([*head, str(stage), *tail])}]" - m = isl.map(f"{{ {src} -> {dst} }}").set_tuple_name(isl.dim_type.IN, acc.name) - sched = sched.union(m.intersect_domain(acc.domain)) - return sched - - -def _parallel_dims(domain: "isl.union_set", deps: "isl.union_map") -> dict[str, tuple[bool, ...]]: - """Per statement, per own domain dimension, whether that dimension is free of dependence. - - Per statement, per own domain dimension, whether that dimension is - free of dependence -- the fact ``coincident`` names in isl. - - Only a statement's *self*-dependence can constrain its own loop - dimensions: the schedule layer sequences statements, so every - cross-statement dependence is already satisfied by that order. A - dimension is parallel when every self-dependence has distance 0 there - (a matmul's k, which accumulates, is the one that is not). - """ - sets: list["isl.set"] = [] - domain.foreach_set(sets.append) - out: dict[str, tuple[bool, ...]] = {} - for s in sets: - own = s.to_union_set() - rank = s.dim(isl.dim_type.SET) - self_deps = deps.intersect_domain(own).intersect_range(own) - if self_deps.is_empty(): - out[s.get_tuple_name()] = (True,) * rank - continue - pieces: list["isl.set"] = [] - self_deps.deltas().foreach_set(pieces.append) - out[s.get_tuple_name()] = tuple( - all(p.dim_min_val(d).is_zero() and p.dim_max_val(d).is_zero() for p in pieces) - for d in range(rank) - ) - return out - - -def _resolve(expr, table: dict[int, object]): - """``expr`` if it is not (transitively) bound in ``table``, else the expression it resolves to. - - ``expr`` if it is not (transitively) bound in ``table``, else the - expression it resolves to -- a penetrated callee's own param Var - bound to the caller's argument, or a penetrated wrapper Call aliased - to whatever its body ultimately resolves to. - """ - return table.get(id(expr), expr) - - -def _bind_dim_vars(params: tuple[Var, ...], args: tuple, callee_name: str) -> None: - """Mirrors ``evaluator.interpreter._bind_dim_vars`` at the type level. - - Mirrors ``evaluator.interpreter._bind_dim_vars`` at the type level: - each ``DimVar`` in a callee param's declared shape binds to the - caller argument's ``ShapeDim`` at that axis. A conflicting bind for - the same name is an actionable error naming the callee -- defense in - depth, since ``hir.function.elaborate`` already requires a call's - argument shapes to equal the callee's own declared ones exactly. - """ - binding: dict[str, object] = {} - for p, a in zip(params, args): - p_shape = getattr(p.type, "shape", None) - a_shape = getattr(a.type, "shape", None) - if p_shape is None or a_shape is None: - continue - for axis, dim in enumerate(p_shape): - if not isinstance(dim, DimVar) or axis >= len(a_shape): - continue - bound = a_shape[axis] - prev = binding.get(dim.name) - if prev is not None and prev != bound: - raise ExtractError( - f"extract: call to {callee_name!r}: DimVar " - f"{dim.name!r} binds to conflicting shapes {prev!r} vs {bound!r}" - ) - binding[dim.name] = bound - - -@dataclass(frozen=True) -class _Gathered: - """Represent Gathered. - - One compute-op ``Call``, args already resolved against every - enclosing penetrated call's argument substitution, ready for - ``_extract_statement``. ``prefix`` is its owning scope's call-site - tag (empty at the top level); ``stmt_name`` already carries it. - ``loops`` is the loop nest the call sits inside, outermost first. - """ - - call: Call - stmt_name: str - prefix: str - loops: tuple[GridRegionExpr, ...] = () - - -def _maybe_replace_args(e: Call, resolved_args: tuple) -> Call: - if all(r is a for r, a in zip(resolved_args, e.args)): - return e - return dataclasses.replace(e, args=resolved_args) - - -def _loop_axes(root): - """Every ``GridRegionExpr`` reachable from ``root``. - - Every ``GridRegionExpr`` reachable from ``root``, as ``(axis per grid, - (axis, initial value) per induction/carry Var, nesting depth per axis)``. - A carry phi's initial value is what makes it variant in an *enclosing* - loop as well as its own. Depth is the number of enclosing grids, taken at - first sight. - """ - axis_of: dict[int, GridRegionExpr] = {} - seed: dict[int, tuple] = {} - depth: dict[int, int] = {} - - class _LoopAxisVisitor(ExprVisitor[None]): - def __init__(self) -> None: - super().__init__() - self.level = 0 - - def _visit_at(self, expr, level: int, ctx=None) -> None: - previous = self.level - self.level = level - try: - self.visit(expr, ctx) - finally: - self.level = previous - - def visit_Call(self, expr: Call, ctx=None) -> None: - for arg in expr.args: - self._visit_at(arg, self.level, ctx) - - def visit_Tuple(self, expr: Tuple, ctx=None) -> None: - for element in expr.elements: - self._visit_at(element, self.level, ctx) - - def visit_GridRegionExpr(self, expr: GridRegionExpr, ctx=None) -> None: - axis_of[id(expr)] = expr - depth[id(expr)] = self.level - seed[id(expr.induction_var)] = (expr, None) - for phi, init in zip(expr.carried_args, expr.init_args): - seed[id(phi)] = (expr, init) - for init in expr.init_args: - self._visit_at(init, self.level, ctx) - self._visit_at(expr.body, self.level + 1, ctx) - for value in expr.yield_values: - self._visit_at(value, self.level + 1, ctx) - - def default_visit(self, expr, ctx=None) -> None: - return None - - _LoopAxisVisitor()._visit_at(root, 0) - return axis_of, seed, depth - - -def _loop_scopes(root) -> dict[int, tuple[GridRegionExpr, ...]]: - """Per expression of one function body, the loop axes it varies with, outermost first. - - Per expression of one function body, the loop axes it varies with, - outermost first -- its iteration domain's leading dimensions. - - Variance, not reachability: a value is inside a loop only when it - (transitively) reads that loop's induction variable or one of its - carried args. Everything else lifts out, and a grid node itself absorbs - its own axis, so a value read *after* the loop is outside it even though - the loop's yield is what produced it. - """ - axis_of, seed, depth = _loop_axes(root) - if not axis_of: - return {} - by_id = {id(axis): axis for axis in axis_of.values()} - variance: dict[int, frozenset] = {} - - for e in collect_exprs(root): - if isinstance(e, GridRegionExpr): - own = set() - for child in (*e.init_args, e.body, *e.yield_values): - own |= variance.get(id(child), frozenset()) - own.discard(id(axis_of[id(e)])) - elif isinstance(e, Var) and id(e) in seed: - axis, init = seed[id(e)] - own = {id(axis)} - if init is not None: - own |= set(variance.get(id(init), frozenset())) - else: - own = set() - for child in expr_children(e): - own |= variance.get(id(child), frozenset()) - variance[id(e)] = frozenset(own) - - return { - key: tuple(by_id[a] for a in sorted(axes, key=lambda a: depth[a])) - for key, axes in variance.items() - if axes - } - - -def _walk_calls( - body, - prefix: str, - active: tuple[int, ...], - site_counter: dict[str, int], - table: dict[int, object], - loops: tuple[GridRegionExpr, ...] = (), - carries: list | None = None, -) -> list["_Gathered"]: - """Walk a body in collect_exprs while penetrating nested function calls. - - Bind resolved arguments, qualify the callee scope, splice its statements, - and alias wrapper results to real producers. Reject recursion, prototypes, - and arity mismatch. Grid regions reuse precomputed loop scopes; their nodes - wire final yields and record carried-value aliases rather than statements. - """ - scope = _loop_scopes(body) - order = collect_exprs(body) - grid_yields: dict[int, tuple] = {} - - own_targets: list[object] = [] - pending: list[object] = [] - for e in order: - if isinstance(e, Tuple): - resolved_elems = tuple(_resolve(x, table) for x in e.elements) - same = all(r is x for r, x in zip(resolved_elems, e.elements)) - table[id(e)] = e if same else dataclasses.replace(e, elements=resolved_elems) - continue - if isinstance(e, GridRegionExpr): - if e.carried_args: - grid_yields[id(e)] = e.yield_values - table[id(e)] = _resolve(e.yield_values[0], table) - if carries is not None: - for phi, value in zip(e.carried_args, e.yield_values): - carries.append((phi, _resolve(value, table))) - else: - table[id(e)] = _resolve(e.body, table) - continue - if not isinstance(e, Call): - continue - - target = e.target - if isinstance(target, TupleGetItem) and id(e.args[0]) in grid_yields: - table[id(e)] = _resolve(grid_yields[id(e.args[0])][target.index], table) - continue - resolved_args = tuple(_resolve(a, table) for a in e.args) - own_loops = loops + scope.get(id(e), ()) - - if isinstance(target, Function): - callee = target - if id(callee) in active: - raise ExtractError( - f"extract: self-recursive call to {callee.name!r} " - "-- extract cannot unroll a function that (transitively) " - "calls itself" - ) - if callee.variants or callee.body is None: - raise ExtractError( - f"extract: {callee.name!r} is a dispatch " - "prototype (has variants / no body) -- extract has no " - "runtime shape to pick a variant statically" - ) - if len(resolved_args) != len(callee.params): - raise ExtractError( - f"extract: call to {callee.name!r} expects " - f"{len(callee.params)} arg(s), got {len(resolved_args)}" - ) - for p, a in zip(callee.params, resolved_args): - table[id(p)] = a - _bind_dim_vars(callee.params, resolved_args, callee.name) - idx = site_counter.get(callee.name, 0) - site_counter[callee.name] = idx + 1 - nested = _walk_calls( - callee.body, - f"{prefix}{callee.name}{idx}_", - active + (id(callee),), - site_counter, - table, - own_loops, - carries, - ) - pending.extend(nested) - table[id(e)] = _resolve(callee.body, table) - continue - - if isinstance(target, (TupleGetItem, Reshape, IndexSelect, Slice)): - table[id(e)] = _maybe_replace_args(e, resolved_args) - continue - - if is_dim_op_call(e): - table[id(e)] = _maybe_replace_args(e, resolved_args) - continue - - if isinstance(target, (Zeros, FullLike)): - table[id(e)] = _maybe_replace_args(e, resolved_args) - continue - - resolved = _maybe_replace_args(e, resolved_args) - table[id(e)] = resolved - own_targets.append(target) - pending.append((resolved, own_loops)) - - own_names = iter(_assign_statement_names(own_targets)) - gathered: list[_Gathered] = [] - for item in pending: - if isinstance(item, _Gathered): - gathered.append(item) - else: - call, own_loops = item - gathered.append( - _Gathered( - call=call, - stmt_name=f"{prefix}{next(own_names)}", - prefix=prefix, - loops=own_loops, - ) - ) - return gathered - - -def extract(hir: Function) -> TileGraph: - """Lift ``hir``'s body into a :class:`TileGraph`. - - Lift ``hir``'s body into a :class:`TileGraph`: one statement per - compute op at element granularity, penetrating every nested ``@func`` - call and every authored loop, with ``deps`` auto-inferred from - ``reads``/``writes`` (see module docstring for the full algorithm). - """ - if hir.body is None: - raise ExtractError( - f"extract: hir Function {hir.name!r} has no body " - "(a dispatch prototype cannot be extracted)" - ) - - carries: list[tuple[Var, object]] = [] - gathered = _walk_calls(hir.body, "", (id(hir),), {}, {}, (), carries) - if not gathered: - raise ExtractError(f"extract: hir Function {hir.name!r} body has no compute ops to extract") - - namer = _buffer_namer() - for phi, yielded in carries: - namer.alias(phi, yielded) - accesses: list[_StatementAccess] = [] - for g in gathered: - accesses.extend(_extract_statement(g.call, g.stmt_name, namer, g.prefix, g.loops)) - - domain = isl.union_set("{}") - reads = isl.union_map("{}") - writes = isl.union_map("{}") - params: dict = {} - buffer_dtypes: dict = {} - for acc in accesses: - domain = domain.union(acc.domain) - for m in acc.reads: - reads = reads.union(m) - for m in acc.writes: - writes = writes.union(m) - for buf, dtype in acc.dtypes.items(): - if dtype is not None: - buffer_dtypes.setdefault(buf, dtype) - for name, dim in acc.params.items(): - prev = params.get(name) - if prev is not None and prev != dim: - raise ExtractError( - f"extract: isl parameter {name!r} resolves to " - f"conflicting ShapeDims across statements: {prev!r} vs " - f"{dim!r}" - ) - params[name] = dim - - schedule_map = _initial_schedule(accesses) - info = isl.union_access_info(reads).set_must_source(writes).set_schedule_map(schedule_map) - deps = info.compute_flow().get_must_dependence() - - units = tuple(TileUnit(name=acc.name, op=acc.op) for acc in accesses) - return TileGraph( - domain=domain, - deps=deps, - reads=reads, - writes=writes, - units=units, - params=params, - buffer_dtypes=buffer_dtypes, - parallel_dims=_parallel_dims(domain, deps), - ) - - -def _as_map(value) -> "isl.map": - """The single ``isl.map`` in ``value``. - - The single ``isl.map`` in ``value``, which isl-python returns as a - ``union_map`` when a ``map`` is composed with one. - """ - if not hasattr(value, "foreach_map"): - return value - maps: list["isl.map"] = [] - value.foreach_map(maps.append) - if len(maps) != 1: - raise ExtractError(f"expected a single map, got {len(maps)}: {value}") - return maps[0] - - -def _only_out_dim(m: "isl.map", pos: int) -> "isl.map": - n_out = m.dim(isl.dim_type.OUT) - return m.project_out(isl.dim_type.OUT, pos + 1, n_out - pos - 1).project_out( - isl.dim_type.OUT, 0, pos - ) - - -def _travels_with(m: "isl.map", pos: int) -> tuple[int, ...]: - """The input dimensions output dimension ``pos`` of ``m`` moves with. - - The map's own domain bounds mention every input dimension, so they are - dropped first -- only the constraints that tie ``pos`` to an input can - answer this. - """ - coupled = _only_out_dim(m.drop_constraints_not_involving_dims(isl.dim_type.OUT, pos, 1), pos) - return tuple( - i for i in range(m.dim(isl.dim_type.IN)) if coupled.involves_dims(isl.dim_type.IN, i, 1) - ) - - -def _static_extent(s: "isl.set", pos: int, what: str) -> tuple[int, int]: - lo, hi = s.dim_min_val(pos), s.dim_max_val(pos) - if not (lo.is_int() and hi.is_int()): - raise ExtractError( - f"{what}: dimension {pos} of {s} is not statically bounded " - "-- a parametric extent has no integer tile count" - ) - return int(lo.num_si()), int(hi.num_si()) - - -def _buffers_by_statement(um: "isl.union_map") -> dict[str, set[str]]: - maps: list["isl.map"] = [] - um.foreach_map(maps.append) - out: dict[str, set[str]] = {} - for m in maps: - stmt = m.get_tuple_name(isl.dim_type.IN) - out.setdefault(stmt, set()).add(m.get_tuple_name(isl.dim_type.OUT)) - return out - - -__all__ = [ - "ExtractError", - "TileGraph", - "TileUnit", - "extract", -] diff --git a/src/tilefoundry/analysis/poly/model.py b/src/tilefoundry/analysis/poly/model.py deleted file mode 100644 index 68befdf6..00000000 --- a/src/tilefoundry/analysis/poly/model.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Polyhedral graph data structures.""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -import isl - - -@dataclass(frozen=True) -class TileUnit: - """One statement's identity. - - ``name`` is the isl tuple name shared by this statement's pieces of - ``TileGraph.domain``/``reads``/``writes``/``deps`` (e.g. ``"MM"``). - ``op`` is the HIR ``Call`` (op@site) that produced this statement -- - the call node itself, not just its bare ``Op``, so a consumer can - still recover ``op.target`` / ``op.args`` / ``op.type``. - """ - - name: str - op: object - - -@dataclass(frozen=True) -class TileGraph: - """Represent one HIR function body as a polyhedral analysis result. - - Domain and access unions use one tuple name per statement or buffer. - ``deps`` contains inferred RAW must-dependencies and ``params`` resolves - dynamic ISL parameters. Buffer dtypes support byte counts without another - HIR walk; ``parallel_dims`` reports dependence-free statement dimensions. - Scheduling owns all schedule trees and resource decisions. - """ - - domain: "isl.union_set" - deps: "isl.union_map" - reads: "isl.union_map" - writes: "isl.union_map" - units: tuple[TileUnit, ...] - params: dict - buffer_dtypes: dict = field(default_factory=dict) - parallel_dims: dict = field(default_factory=dict) - - -__all__ = ["TileGraph", "TileUnit"] diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index e549c661..8222fee3 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -30,10 +30,10 @@ ) from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext +from .affine import loop_affine_term from .errors import AnalysisError from .footprint import _widest_allowed from .metadata import BufferFootprint, LoopFootprintMetadata -from .poly.affine import loop_affine_term @dataclass(frozen=True) diff --git a/src/tilefoundry/cli/analyze.py b/src/tilefoundry/cli/analyze.py index 5fe81886..a6ea3e9c 100644 --- a/src/tilefoundry/cli/analyze.py +++ b/src/tilefoundry/cli/analyze.py @@ -101,10 +101,10 @@ def guidance() -> str: is an observation, not a bound. Each family's record, how every field is computed, and what it prints: - tilefoundry spec analysis 2.2.1 compute-cost - tilefoundry spec analysis 2.2.2 memory - tilefoundry spec analysis 2.2.3 roofline - tilefoundry spec analysis 2.2.4 performance + tilefoundry spec analysis 1.2.1 compute-cost + tilefoundry spec analysis 1.2.2 memory + tilefoundry spec analysis 1.2.3 roofline + tilefoundry spec analysis 1.2.4 performance """ ) diff --git a/src/tilefoundry/schedule/__init__.py b/src/tilefoundry/schedule/__init__.py deleted file mode 100644 index c7563028..00000000 --- a/src/tilefoundry/schedule/__init__.py +++ /dev/null @@ -1,46 +0,0 @@ -"""The public Schedule boundary. - -One call names a program and a level of its parallel hierarchy; one registered -algorithm answers with a plan it owns entirely. The names re-exported here are -that boundary and nothing else: how an algorithm reaches its answer, and what it -looks at on the way, are its own. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path - -from .api import ScheduleResult, schedule -from .errors import ScheduleError -from .plan import PlanVerificationError, SchedulePlan - - -@dataclass(frozen=True) -class ScheduleOptions: - """Solver runtime and debug controls, independent of which algorithm runs. - - `stop_at_first_solution` asks for a plan rather than the best plan. The search - minimises a makespan, so on a model it cannot prove optimal for it keeps - improving until the time limit -- which makes `timeout_seconds` the runtime of - every solve rather than a limit that rarely fires. A caller that needs a plan to - exist and to verify, and not to be optimal, says so here and gets the first - feasible one. The time limit still applies: a model with no solution found yet - is still bounded, so this cannot turn a slow search into an unbounded one. - """ - - timeout_seconds: float = 60.0 - workers: int = 0 - random_seed: int = 0 - stop_at_first_solution: bool = False - debug_dump_dir: Path | None = None - - -__all__ = [ - "PlanVerificationError", - "ScheduleError", - "ScheduleOptions", - "SchedulePlan", - "ScheduleResult", - "schedule", -] diff --git a/src/tilefoundry/schedule/api.py b/src/tilefoundry/schedule/api.py deleted file mode 100644 index 77400214..00000000 --- a/src/tilefoundry/schedule/api.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Expose the public Schedule operation. - -A call names a module, one of its functions, and a declared topology level. The -module owns hardware and topology resolution. Resolution and scheduler dispatch -finish before solving, and the returned plan records every algorithm decision. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass - -from tilefoundry.analysis.check import _resolve_program_geometry, check_program -from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.specialize import SpecializationError -from tilefoundry.ir.types.shard import Topology -from tilefoundry.target import Target, UnsupportedCapabilityError -from tilefoundry.target.services import Scheduler - -from .errors import ScheduleError -from .plan import SchedulePlan - - -@dataclass(frozen=True) -class ScheduleResult: - """What one public Schedule call decided, and what it decided it against. - - Scheduling is a decision about a program, not a rewrite of one, so the - Module is the one that was passed in. - - The Function is too, unless an extent was chosen for it -- then it is the - concrete function derived from that input, because that is the program the - plan was solved for and the one the plan can be verified against. A caller - handed back its own symbolic input would have a plan it could not check. - """ - - module: Module - function: Function - topology: Topology - plan: SchedulePlan - - -def _topology(module: Module, name: str) -> Topology: - """Return one named level after the local scheduling preconditions. - - The launch-extent rule is defined by docs/spec/target.md § Topology levels. - """ - if not isinstance(name, str) or not name: - raise ScheduleError( - f"schedule: topology must be a non-empty level name, got {name!r}" - ) - try: - level = module.resolve_topology(name) - except ValueError as error: - raise ScheduleError(f"schedule: {error}") from None - return level - - -def _algorithm(target: Target, topology: str) -> Scheduler: - """The scheduler selected by the resolved Target for *topology*.""" - try: - return target.get_scheduler(topology) - except UnsupportedCapabilityError as error: - raise ScheduleError(f"schedule: {error}") from None - - -def _options(options: object | None) -> object: - """The common options every registered algorithm receives.""" - from . import ScheduleOptions # noqa: PLC0415 - - if options is None: - return ScheduleOptions() - if not isinstance(options, ScheduleOptions): - raise ScheduleError( - "schedule: options must be ScheduleOptions, got " - f"{type(options).__name__}" - ) - return options - - -def schedule( - module: Module, - function: Function, - *, - topology: str, - options: object | None = None, - dims: "Mapping[str, int] | None" = None, -) -> ScheduleResult: - """Solve *function* at the *topology* level of *module*'s hierarchy. - - *dims* selects concrete extents and the matching specialization before the - solver counts work or capacity. The input must be a prototype or variant - owned by *module*; substitution does not widen that ownership boundary. - """ - if not isinstance(module, Module): - raise TypeError( - f"schedule: expected a Module, got {type(module).__name__}. A " - "Function declares neither hardware nor a topology hierarchy; " - "select the Module that owns it." - ) - if not isinstance(function, Function): - raise TypeError( - f"schedule: expected an hir.Function, got {type(function).__name__}" - ) - if not module.owns(function): - raise ScheduleError( - f"schedule: {function.name!r} is not a function of module " - f"{module.name!r}" - ) - result_module = module - try: - module, function = _resolve_program_geometry(module, function, dims) - except SpecializationError as error: - raise ScheduleError(f"schedule: {error}") from None - - target = module.resolve_target() - level = _topology(module, topology) - try: - check_program(module, function, level=topology) - except AnalysisError as error: - raise ScheduleError(f"schedule: {error}") from None - algorithm = _algorithm(target, topology) - resolved_options = _options(options) - - plan = algorithm.solve(module, function, target, level, resolved_options) - if not isinstance(plan, SchedulePlan): - raise ScheduleError( - f"schedule: the {topology!r} algorithm for " - f"{type(target).__name__} returned a {type(plan).__name__}, not a " - "SchedulePlan" - ) - plan.verify(result_module, function, level) - return ScheduleResult( - module=result_module, function=function, topology=level, plan=plan - ) - - -__all__ = ["ScheduleResult", "schedule"] diff --git a/src/tilefoundry/schedule/errors.py b/src/tilefoundry/schedule/errors.py deleted file mode 100644 index 759f7c2a..00000000 --- a/src/tilefoundry/schedule/errors.py +++ /dev/null @@ -1,19 +0,0 @@ -"""The diagnostic the schedule layer raises. - -There is one such class for the whole layer, so a caller that catches a -scheduling failure catches every scheduling failure rather than the subset that -happens to come from the entry it imported. - -Plan verification failures are separate, because they say something different: a -schedule error means the request could not be served, while a verification error -means a plan was produced and does not hold together. -""" - -from __future__ import annotations - - -class ScheduleError(ValueError): - """A request the schedule layer cannot serve, or a solve that failed.""" - - -__all__ = ["ScheduleError"] diff --git a/src/tilefoundry/schedule/facts.py b/src/tilefoundry/schedule/facts.py deleted file mode 100644 index 96ea2954..00000000 --- a/src/tilefoundry/schedule/facts.py +++ /dev/null @@ -1,36 +0,0 @@ -"""What the schedule layer asks a target about one instruction. - -An atom is target-specific, but *asking* for one must not be: an algorithm names -the facts it needs and the target package registers the conversion that supplies -them, so nothing reaches into a target through an object whose shape it has to -know. Each algorithm family owns the rest of its own facts vocabulary. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -from tilefoundry.ir.types import DType - - -@dataclass(frozen=True) -class AtomFact: - """One atom's facts, for later CP-SAT atom selection. - - Shape and dtype describe MNK and A/B/C. Durations are nominal nanoseconds; - ``compute_duration`` excludes traffic. Storage is per-thread bytes and - resource is thread-scope occupancy. ``atom`` remains target-owned and - opaque so later stages can use the realized descriptor directly. - """ - - shape: tuple[int, int, int] - dtype: tuple[DType, DType, DType] - duration: float - compute_duration: float - storage: dict[str, int] - resource: dict[str, int] - is_async: bool - atom: object - - -__all__ = ["AtomFact"] diff --git a/src/tilefoundry/schedule/kernel_schedule.py b/src/tilefoundry/schedule/kernel_schedule.py deleted file mode 100644 index c0ea86c5..00000000 --- a/src/tilefoundry/schedule/kernel_schedule.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Build and transform an isl schedule tree for a ``TileGraph``. - -The dependence-respecting unit order becomes a sequence of identity bands; -nothing here solves an objective. Band helpers enumerate, identify, and tile -those bands. Parallel dimensions set each member's ``coincident`` flag. -""" -from __future__ import annotations - -import isl - -from tilefoundry.analysis.poly import TileGraph - - -class KernelScheduleError(RuntimeError): - """A schedule tree the band operations cannot work on. - - A schedule tree the band operations cannot work on -- always raised - with a message naming what was found instead. - """ - - -def _domain_sets(domain: "isl.union_set") -> dict[str, "isl.set"]: - sets: list["isl.set"] = [] - domain.foreach_set(sets.append) - return {s.get_tuple_name(): s for s in sets} - - -def _statement_schedule(s: "isl.set") -> "isl.schedule": - """One statement's own domain under one identity band. - - One statement's own domain under one identity band, so the band's - members are that statement's own dimensions, in order. - """ - sched = isl.schedule.from_domain(s.to_union_set()) - if not s.dim(isl.dim_type.SET): - return sched - identity = isl.multi_union_pw_aff.from_union_map(s.to_union_set().identity()) - return sched.insert_partial_schedule(identity) - - -def _mark_coincident(tree: "isl.schedule", parallel: dict[str, tuple[bool, ...]]): - """``parallel``'s per-dimension flags written onto every band member.""" - - def mark(node): - if not isinstance(node, isl.schedule_node_band): - return node - flags = parallel.get(band_statement(node), ()) - for member, is_parallel in enumerate(flags[: node.n_member()]): - if is_parallel: - node = node.member_set_coincident(member, 1) - return node - - return tree.get_root().map_descendant_bottom_up(mark).get_schedule() - - -def build_schedule_tree(tg: TileGraph) -> "isl.schedule": - """Build a private sequence of identity bands from immutable analysis. - - The statements are not fused into one band: their ranks differ, so a - padded shared band member would mean a different loop in each of them. - """ - if not tg.units: - raise KernelScheduleError("build_schedule_tree: tg.units is empty -- nothing to schedule") - by_name = _domain_sets(tg.domain) - missing = [unit.name for unit in tg.units if unit.name not in by_name] - if missing: - raise KernelScheduleError( - f"build_schedule_tree: statements {missing} have no domain piece -- " - "tg.units and tg.domain must come from one extract() run" - ) - tree = _statement_schedule(by_name[tg.units[0].name]) - for unit in tg.units[1:]: - tree = tree.sequence(_statement_schedule(by_name[unit.name])) - return _mark_coincident(tree, tg.parallel_dims) - - -def schedule_bands(tree: "isl.schedule") -> tuple["isl.schedule_node_band", ...]: - """Every band in ``tree``, in top-down order. - - Every band in ``tree``, in top-down order -- which for a - :func:`build_schedule_tree` tree is ``tg.units`` order. - """ - found: list["isl.schedule_node_band"] = [] - - def visit(node) -> bool: - if isinstance(node, isl.schedule_node_band): - found.append(node) - return True - - tree.get_root().foreach_descendant_top_down(visit) - if not found: - raise KernelScheduleError(f"schedule_bands: schedule tree carries no band node: {tree}") - return tuple(found) - - -def band_statement(band: "isl.schedule_node_band") -> str: - """The one statement ``band`` schedules.""" - sets: list["isl.set"] = [] - band.get_domain().foreach_set(sets.append) - unique = sorted({s.get_tuple_name() for s in sets}) - if len(unique) != 1: - raise KernelScheduleError( - f"band_statement: band covers statements {unique} -- every band a " - "sequenced schedule tree carries belongs to exactly one statement" - ) - return unique[0] - - -def tile_band(band: "isl.schedule_node_band", sizes: tuple[int, ...]) -> "isl.schedule": - """Tile band. - - ``band`` split into a tile band over ``sizes`` plus a point band - holding the remainder, returned as the whole schedule. - """ - if band.n_member() != len(sizes): - raise KernelScheduleError( - f"tile_band: band has {band.n_member()} member(s) but got " - f"{len(sizes)} tile size(s)" - ) - space = band.get_partial_schedule().get_space() - multi = isl.multi_val.zero(space) - for i, size in enumerate(sizes): - if size < 1: - raise KernelScheduleError(f"tile_band: tile size {size} at member {i} must be >= 1") - multi = multi.set_at(i, isl.val(size)) - return band.tile(multi).get_schedule() - - -def tile_bands(tree: "isl.schedule", sizes: dict[str, tuple[int, ...]]) -> "isl.schedule": - """Every band in ``tree`` tiled by its own statement's ``sizes``. - - Tiling replaces one band with two, shifting every band below it in - top-down order, so the walk runs bottom-up over the positions instead - of over live nodes (an ``isl.schedule_node`` does not survive the tree - it was taken from being rebuilt). - """ - for position in reversed(range(len(schedule_bands(tree)))): - band = schedule_bands(tree)[position] - name = band_statement(band) - if name not in sizes: - raise KernelScheduleError(f"tile_bands: no tile size decided for statement {name!r}") - tree = tile_band(band, sizes[name]) - return tree - - -__all__ = [ - "KernelScheduleError", - "band_statement", - "build_schedule_tree", - "schedule_bands", - "tile_band", - "tile_bands", -] diff --git a/src/tilefoundry/schedule/partition/__init__.py b/src/tilefoundry/schedule/partition/__init__.py deleted file mode 100644 index eac31e73..00000000 --- a/src/tilefoundry/schedule/partition/__init__.py +++ /dev/null @@ -1,45 +0,0 @@ -"""The spatial partition scheduling algorithm family. - -The names here are the stages one registered algorithm composes its solve from, -in the order it composes them: extract what the program states, ask the hardware -once, close the problem, solve it, export the plan. They are this family's own -vocabulary and no other algorithm reads them. -""" - -from __future__ import annotations - -from .facts import PartitionFacts, PartitionFactsError, PartitionFactsQuery -from .plan import ( - PartitionedOperation, - PartitionProof, - PartitionSchedulePlan, - PlacedValue, - PositionInterval, - TimeInterval, - export_partition_plan, -) -from .problem import PartitionProblem, PartitionProblemError, build_partition_problem -from .program import PartitionProgram, PartitionProgramError, build_partition_program -from .solve import PartitionSolution, PartitionSolveError, solve_partition_problem - -__all__ = [ - "PartitionFacts", - "PartitionFactsError", - "PartitionFactsQuery", - "PartitionProblem", - "PartitionProblemError", - "PartitionProgram", - "PartitionProgramError", - "PartitionProof", - "PartitionSchedulePlan", - "PartitionedOperation", - "PlacedValue", - "PositionInterval", - "TimeInterval", - "PartitionSolution", - "PartitionSolveError", - "build_partition_problem", - "build_partition_program", - "export_partition_plan", - "solve_partition_problem", -] diff --git a/src/tilefoundry/schedule/partition/facts.py b/src/tilefoundry/schedule/partition/facts.py deleted file mode 100644 index 6e70ee34..00000000 --- a/src/tilefoundry/schedule/partition/facts.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Typed target facts consumed by the private partition scheduler. - -These are every number the partition problem is closed with. Once they are -projected, neither the problem nor the solve holds a Target: what the hardware -contributes to a decision is exactly this record, and it is readable in one -place instead of inferred from the call sites that used to reach for it. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -from tilefoundry.ir.types import DType -from tilefoundry.schedule.plan import TargetSpecRef - - -class PartitionFactsError(ValueError): - """A projected fact the partition algorithm needs is absent or unusable.""" - - -@dataclass(frozen=True) -class PartitionFactsQuery: - """The one topology level a projection is asked to describe.""" - - topology: str - - -@dataclass(frozen=True) -class PartitionFacts: - """All concrete hardware information required to close one partition.""" - - topology: str - spec: TargetSpecRef - parallel_units: int - memory_bandwidth_bytes_per_second: int - memory_capacity_bytes: int - peak_flops_per_second: tuple[tuple[DType, int], ...] - - def peak_flops(self, dtype: DType) -> int: - """The dense peak rate stated for *dtype*. - - A dtype the hardware documents no rate for is an error rather than a - zero or a substituted neighbour: costing work at a rate nobody published - would put a number in the plan that no document supports. - """ - for candidate, value in self.peak_flops_per_second: - if candidate == dtype: - return value - stated = ", ".join(sorted(item[0].name for item in self.peak_flops_per_second)) - raise PartitionFactsError( - f"{self.spec.device_id} states no dense peak rate for {dtype.name}; " - f"it states {stated or 'none'}" - ) - - -__all__ = ["PartitionFacts", "PartitionFactsError", "PartitionFactsQuery"] diff --git a/src/tilefoundry/schedule/partition/plan.py b/src/tilefoundry/schedule/partition/plan.py deleted file mode 100644 index 0e5f65f7..00000000 --- a/src/tilefoundry/schedule/partition/plan.py +++ /dev/null @@ -1,603 +0,0 @@ -"""Represent a solved spatial partition without rewriting its program. - -Values and operations use authored identities and retain their IR types. -Verification resolves references, checks placement edges and level bounds, and -rejects overlapping operations without rebuilding candidates or re-solving. -""" - -from __future__ import annotations - -import json -from dataclasses import asdict, dataclass -from typing import Literal - -from tilefoundry.ir.core.metadata import binding_name -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.types import TensorType, Type -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import ShardLayout, Topology -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.schedule.plan import ( - PlanVerificationError, - SchedulePlan, - TargetSpecRef, -) - -from .problem import PartitionProblem -from .solve import PartitionSolution - - -@dataclass(frozen=True) -class PositionInterval: - """The half-open range of parallel positions something occupies.""" - - start: int - end: int - - -@dataclass(frozen=True) -class TimeInterval: - """One operation's half-open execution interval.""" - - start_ns: int - end_ns: int - - -@dataclass(frozen=True) -class PlacedValue: - """One tensor value, the type it was placed in, and who touches it.""" - - id: str - type: Type - producer_id: str | None - consumer_ids: tuple[str, ...] - positions: PositionInterval - - -@dataclass(frozen=True) -class PartitionedOperation: - """One operation that runs, where it runs, and when. - - `synthesized` marks a Reshard the algorithm introduced to connect two - otherwise unconnected placements. It is one of these records like any other - operation, so there is no second channel an agent would have to read to learn - that data moves. - - `positions` is absent for an operation that occupies no parallel position of - its own. Moving a value between placements is charged as traffic rather than - as occupancy, so a Reshard has no position range to state. - """ - - id: str - operation: str - synthesized: bool - input_ids: tuple[str, ...] - output_ids: tuple[str, ...] - positions: PositionInterval | None - interval: TimeInterval | None - - -@dataclass(frozen=True) -class PartitionProof: - """What the solve proved about its own objective. - - This is a result fact, not a summary of the run: the objective value, the - bound the solver could establish, and whether the two met. - """ - - status: Literal["OPTIMAL", "FEASIBLE_NOT_PROVEN"] - objective_ns: int - best_bound_ns: int - proven_optimal: bool - - -def _logical_tensor(type: Type) -> tuple[object, ...] | None: - """What stays the same across every placement of one tensor.""" - if not isinstance(type, TensorType): - return None - return (type.shape, type.dtype, type.storage) - - -def _layout_json(type: Type) -> object: - """The placement part of one type, as plain data.""" - if not isinstance(type, TensorType) or not isinstance(type.layout, ShardLayout): - return None - layout = type.layout - return { - "topology": layout.mesh.topologies[0].name, - "mesh_shape": [str(dim) for dim in layout.mesh.layout.shape], - "attrs": [attr.__class__.__name__ for attr in layout.attrs], - "shape": [str(dim) for dim in layout.layout.shape], - "strides": ( - None - if layout.layout.strides is None - else [str(stride) for stride in layout.layout.strides] - ), - } - - -def _type_json(type: Type) -> object: - """One selected type as plain data, in the plan's own vocabulary.""" - if not isinstance(type, TensorType): - return {"kind": type.__class__.__name__} - return { - "kind": "tensor", - "shape": [str(dim) for dim in type.shape], - "dtype": type.dtype.name, - "storage": type.storage.name.lower(), - "layout": _layout_json(type), - } - - -@dataclass(frozen=True) -class PartitionSchedulePlan(SchedulePlan): - """The placement one partition solve committed to, and its proof.""" - - topology: str - extent: int - target: TargetSpecRef - values: tuple[PlacedValue, ...] - operations: tuple[PartitionedOperation, ...] - root_results: tuple[str, ...] - proof: PartitionProof - - def verify(self, module: Module, function: Function, topology: Topology) -> None: - """Check the decision holds together, without invoking a solver.""" - self._check_request(topology) - values = self._checked_index() - self._check_references(values) - self._check_edges(values) - self._check_positions(values) - self._check_exclusion() - self._check_roots(values) - if self.proof.best_bound_ns > self.proof.objective_ns: - raise PlanVerificationError( - "partition plan states a bound above its own objective" - ) - - def _check_request(self, topology: Topology) -> None: - if topology.name != self.topology: - raise PlanVerificationError( - f"partition plan decided topology {self.topology!r}, not " - f"{topology.name!r}" - ) - extent = static_dim_value(topology.size) - if extent != self.extent: - raise PlanVerificationError( - f"partition plan decided over {self.extent} positions of " - f"{self.topology!r}, but the level declares {extent}" - ) - - def _checked_index(self) -> dict[str, PlacedValue]: - values: dict[str, PlacedValue] = {} - for value in self.values: - if value.id in values: - raise PlanVerificationError( - f"partition plan places value {value.id!r} twice" - ) - values[value.id] = value - seen: set[str] = set() - for operation in self.operations: - if operation.id in seen: - raise PlanVerificationError( - f"partition plan runs operation {operation.id!r} twice" - ) - seen.add(operation.id) - return values - - def _check_references(self, values: dict[str, PlacedValue]) -> None: - """Every edge must be named the same way from both of its ends. - - Naming an operation that exists is not enough: an edge that one end - claims and the other does not is a claim about a decision nobody made, - and a reachability walk that followed it would report a program flow that - the operations do not implement. - """ - operations = {operation.id: operation for operation in self.operations} - for value in self.values: - producer = value.producer_id - if producer is not None: - if producer not in operations: - raise PlanVerificationError( - f"partition plan value {value.id!r} names producer " - f"{producer!r}, which the plan does not run" - ) - if value.id not in operations[producer].output_ids: - raise PlanVerificationError( - f"partition plan value {value.id!r} names producer " - f"{producer!r}, which does not produce it" - ) - for consumer_id in value.consumer_ids: - if consumer_id not in operations: - raise PlanVerificationError( - f"partition plan value {value.id!r} names consumer " - f"{consumer_id!r}, which the plan does not run" - ) - if value.id not in operations[consumer_id].input_ids: - raise PlanVerificationError( - f"partition plan value {value.id!r} names consumer " - f"{consumer_id!r}, which does not read it" - ) - for operation in self.operations: - for value_id in (*operation.input_ids, *operation.output_ids): - if value_id not in values: - raise PlanVerificationError( - f"partition plan operation {operation.id!r} refers to " - f"unplaced value {value_id!r}" - ) - for value_id in operation.output_ids: - if values[value_id].producer_id != operation.id: - raise PlanVerificationError( - f"partition plan operation {operation.id!r} produces " - f"{value_id!r}, which names producer " - f"{values[value_id].producer_id!r}" - ) - for value_id in operation.input_ids: - if operation.id not in values[value_id].consumer_ids: - raise PlanVerificationError( - f"partition plan operation {operation.id!r} reads " - f"{value_id!r}, which does not name it as a consumer" - ) - - def _check_edges(self, values: dict[str, PlacedValue]) -> None: - for value in self.values: - if not isinstance(value.type, TensorType): - raise PlanVerificationError( - f"partition plan places value {value.id!r} in " - f"{type(value.type).__name__}, which is not a tensor type" - ) - if value.type.storage is not StorageKind.GMEM: - raise PlanVerificationError( - f"partition plan places value {value.id!r} in " - f"{value.type.storage.name}, and a partitioned value is " - "addressable global memory" - ) - for operation in self.operations: - if not operation.synthesized: - continue - inputs = tuple(values[value_id] for value_id in operation.input_ids) - outputs = tuple(values[value_id] for value_id in operation.output_ids) - if len(inputs) != 1 or len(outputs) != 1: - raise PlanVerificationError( - f"partition plan synthesized {operation.id!r} with " - f"{len(inputs)} inputs and {len(outputs)} outputs; moving one " - "value takes one of each" - ) - source, target = inputs[0], outputs[0] - if _logical_tensor(source.type) != _logical_tensor(target.type): - raise PlanVerificationError( - f"partition plan synthesized {operation.id!r} between " - f"{source.id!r} and {target.id!r}, which are different logical " - "tensors" - ) - if source.type == target.type: - raise PlanVerificationError( - f"partition plan synthesized {operation.id!r} between two " - "identical placements, which moves nothing" - ) - - def _check_positions(self, values: dict[str, PlacedValue]) -> None: - for value in self.values: - self._check_interval(value.id, value.positions) - for operation in self.operations: - if operation.positions is not None: - self._check_interval(operation.id, operation.positions) - if operation.interval is not None and ( - operation.interval.end_ns < operation.interval.start_ns - ): - raise PlanVerificationError( - f"partition plan operation {operation.id!r} ends before it starts" - ) - - def _check_interval(self, owner: str, positions: PositionInterval) -> None: - if positions.end <= positions.start: - raise PlanVerificationError( - f"partition plan gives {owner!r} the empty position range " - f"[{positions.start}, {positions.end})" - ) - if positions.start < 0 or positions.end > self.extent: - raise PlanVerificationError( - f"partition plan places {owner!r} on positions [{positions.start}, " - f"{positions.end}), outside the {self.extent} positions of " - f"{self.topology!r}" - ) - - def _check_exclusion(self) -> None: - """No two operations may hold the same position at the same time. - - Only operations that occupy positions and run over an interval take part: - an operation charged as traffic states no occupancy to conflict over. - """ - placed = tuple( - operation - for operation in self.operations - if operation.interval is not None and operation.positions is not None - ) - for index, left in enumerate(placed): - for right in placed[index + 1 :]: - if not _overlap( - left.interval.start_ns, - left.interval.end_ns, - right.interval.start_ns, - right.interval.end_ns, - ): - continue - if _overlap( - left.positions.start, - left.positions.end, - right.positions.start, - right.positions.end, - ): - raise PlanVerificationError( - f"partition plan runs {left.id!r} and {right.id!r} at the " - "same time on overlapping positions" - ) - - def _check_roots(self, values: dict[str, PlacedValue]) -> None: - """Every root result must be reachable by following producer edges. - - A placement the plan does not produce is where the walk stops: it is - either the program's own input or a value a region carries, and neither is - an operation this plan decided about. - """ - available = {value.id for value in self.values if value.producer_id is None} - producer_of = { - value.id: value.producer_id - for value in self.values - if value.producer_id is not None - } - inputs_of = {operation.id: operation.input_ids for operation in self.operations} - - def reachable(value_id: str, active: frozenset[str]) -> bool: - if value_id in available: - return True - if value_id in active: - return False - producer = producer_of.get(value_id) - if producer is None: - return False - return all( - reachable(source, active | {value_id}) for source in inputs_of[producer] - ) - - for value_id in self.root_results: - if value_id not in values: - raise PlanVerificationError( - f"partition plan leaves root result {value_id!r} unplaced" - ) - if not reachable(value_id, frozenset()): - raise PlanVerificationError( - f"partition plan cannot reach root result {value_id!r} from the " - "program's own inputs" - ) - - def to_json(self) -> str: - """Render the whole decision as sorted-key JSON.""" - payload = { - "topology": self.topology, - "extent": self.extent, - "target": asdict(self.target), - "proof": asdict(self.proof), - "root_results": list(self.root_results), - "values": [ - { - "id": value.id, - "type": _type_json(value.type), - "producer_id": value.producer_id, - "consumer_ids": list(value.consumer_ids), - "positions": asdict(value.positions), - } - for value in self.values - ], - "operations": [ - { - "id": operation.id, - "operation": operation.operation, - "synthesized": operation.synthesized, - "input_ids": list(operation.input_ids), - "output_ids": list(operation.output_ids), - "positions": ( - asdict(operation.positions) - if operation.positions is not None - else None - ), - "interval": ( - asdict(operation.interval) - if operation.interval is not None - else None - ), - } - for operation in self.operations - ], - } - return json.dumps(payload, sort_keys=True) - - -def _overlap(left_start: int, left_end: int, right_start: int, right_end: int) -> bool: - return left_start < right_end and right_start < left_end - - -def _base_value_name(problem: PartitionProblem, value_id: int) -> str: - """A readable name for one value, derived from the authored program.""" - info = problem.values[value_id] - name = binding_name(info.source) or getattr(info.source, "name", None) - if not name: - target = getattr(info.source, "target", None) - name = type(target).__name__.lower() if target is not None else "value" - for index in info.leaf_path: - name = f"{name}.{index}" - if info.role != "normal": - name = f"{name}.{info.role}" - return name - - -def _placement_tag(type: Type) -> str: - """A short readable name for how one placement divides its value.""" - if not isinstance(type, TensorType) or not isinstance(type.layout, ShardLayout): - return "whole" - kinds = {"Split": "split", "Broadcast": "bcast", "Partial": "partial"} - attrs = "".join( - kinds.get(attr.__class__.__name__, attr.__class__.__name__.lower()) - for attr in type.layout.attrs - ) - extent = type.layout.mesh.layout.shape[0] - return f"{attrs or 'whole'}{extent}" - - -def _placement_ids( - problem: PartitionProblem, selected_buckets: tuple[int, ...] -) -> dict[int, str]: - """One stable readable identity per selected placement. - - A value may be resident in more than one placement at once: that is exactly - what a Reshard connects. The value's own name is used alone while it has a - single placement, and is qualified by how each placement divides it when it - has several, so a plan naming two placements of one tensor stays readable. - """ - by_value: dict[int, list[int]] = {} - for bucket_id in sorted(selected_buckets): - by_value.setdefault(problem.buckets[bucket_id].value_id, []).append(bucket_id) - used: dict[str, int] = {} - ids: dict[int, str] = {} - for value_id in sorted(by_value): - buckets = by_value[value_id] - name = _base_value_name(problem, value_id) - for bucket_id in buckets: - base = name - if len(buckets) > 1: - type = problem.types[problem.buckets[bucket_id].type_id] - base = f"{name}@{_placement_tag(type)}" - count = used.get(base, 0) - used[base] = count + 1 - ids[bucket_id] = base if count == 0 else f"{base}#{count}" - return ids - - -def _operation_ids( - problem: PartitionProblem, - placement_ids: dict[int, str], - selected: tuple[int, ...], -) -> dict[int, str]: - """One stable readable identity per selected operation.""" - ids: dict[int, str] = {} - used: dict[str, int] = {} - for candidate_id in selected: - candidate = problem.candidates[candidate_id] - kind = type(candidate.op).__name__.lower() - produced = tuple( - placement_ids[bucket_id] - for bucket_id in candidate.output_bucket_ids - if bucket_id in placement_ids - ) - base = f"{kind}:{produced[0]}" if produced else kind - count = used.get(base, 0) - used[base] = count + 1 - ids[candidate_id] = base if count == 0 else f"{base}#{count}" - return ids - - -def _placed_positions(problem: PartitionProblem, bucket_id: int | None) -> int: - """How many positions one selected placement occupies.""" - if bucket_id is None: - return 1 - type = problem.types[problem.buckets[bucket_id].type_id] - if not isinstance(type, TensorType) or not isinstance(type.layout, ShardLayout): - return 1 - count = type.layout.mesh.layout.shape[0] - return count if isinstance(count, int) and count > 0 else 1 - - -def export_partition_plan( - problem: PartitionProblem, solution: PartitionSolution -) -> PartitionSchedulePlan: - """State the solved selection in the plan's own stable vocabulary.""" - placement_ids = _placement_ids(problem, solution.selected_bucket_ids) - operation_ids = _operation_ids( - problem, placement_ids, solution.selected_candidate_ids - ) - intervals = dict(solution.candidate_intervals_ns) - offsets = dict(solution.bucket_offsets) - - producers: dict[str, str] = {} - consumers: dict[str, list[str]] = {} - operations: list[PartitionedOperation] = [] - for candidate_id in solution.selected_candidate_ids: - candidate = problem.candidates[candidate_id] - operation_id = operation_ids[candidate_id] - input_ids = tuple( - placement_ids[bucket_id] for bucket_id in candidate.input_bucket_ids - ) - output_ids = tuple( - placement_ids[bucket_id] for bucket_id in candidate.output_bucket_ids - ) - for placement in output_ids: - producers[placement] = operation_id - for placement in input_ids: - consumers.setdefault(placement, []).append(operation_id) - anchor = candidate.output_bucket_ids[0] if candidate.output_bucket_ids else None - interval = intervals.get(candidate_id) - positions = None - if candidate.topology_count > 0 and anchor is not None: - start = offsets.get(anchor, 0) - positions = PositionInterval(start, start + candidate.topology_count) - operations.append( - PartitionedOperation( - id=operation_id, - operation=type(candidate.op).__name__, - synthesized=candidate.site_id is None, - input_ids=input_ids, - output_ids=output_ids, - positions=positions, - interval=( - TimeInterval(interval.start_ns, interval.end_ns) - if interval is not None - else None - ), - ) - ) - - values = tuple( - PlacedValue( - id=placement_ids[bucket_id], - type=problem.types[problem.buckets[bucket_id].type_id], - producer_id=producers.get(placement_ids[bucket_id]), - consumer_ids=tuple(consumers.get(placement_ids[bucket_id], ())), - positions=PositionInterval( - offsets.get(bucket_id, 0), - offsets.get(bucket_id, 0) + _placed_positions(problem, bucket_id), - ), - ) - for bucket_id in sorted(solution.selected_bucket_ids) - ) - root_placements = tuple( - placement_ids[bucket_id] - for bucket_id in sorted(solution.selected_bucket_ids) - if problem.buckets[bucket_id].value_id in set(problem.root_value_ids) - ) - - return PartitionSchedulePlan( - topology=problem.topology.name, - extent=problem.extent, - target=problem.facts.spec, - values=values, - operations=tuple(operations), - root_results=root_placements, - proof=PartitionProof( - status=solution.status, - objective_ns=solution.makespan_ns, - best_bound_ns=solution.best_bound_ns, - proven_optimal=solution.status == "OPTIMAL", - ), - ) - - -__all__ = [ - "PartitionProof", - "PartitionSchedulePlan", - "PartitionedOperation", - "PlacedValue", - "PositionInterval", - "TimeInterval", - "export_partition_plan", -] diff --git a/src/tilefoundry/schedule/partition/problem.py b/src/tilefoundry/schedule/partition/problem.py deleted file mode 100644 index 18b028cc..00000000 --- a/src/tilefoundry/schedule/partition/problem.py +++ /dev/null @@ -1,818 +0,0 @@ -"""The closed, target-free constraint input for partition scheduling. - -Given what the program states and what the hardware was asked once, this -enumerates every legal placement of every value, every operation that can produce -it, and the Reshard operations needed where no direct edge connects two otherwise -legal choices. Each candidate carries its own already-computed duration and -traffic, so the solve that follows reads numbers instead of rates: after this -stage nothing consults a Target again. -""" - -from __future__ import annotations - -import itertools -from dataclasses import dataclass, replace -from types import MappingProxyType -from typing import Literal, Mapping - -from tilefoundry.ir.constraints import ( - LayoutConstraint, - MeshConstraint, - ScheduleConstraintMetadata, - StorageConstraint, - is_layout_wildcard, -) -from tilefoundry.ir.core import Call, Expr, Op, Tuple, Var, VerifyError, source_metadata -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.sharding.reshard import Reshard -from tilefoundry.ir.hir.tensor.reshape import Reshape -from tilefoundry.ir.hir.tensor.transpose import Transpose -from tilefoundry.ir.types import TensorType, Type, make_shard_tensor_type -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import ( - Broadcast, - Layout, - Mesh, - Partial, - ShardLayout, - Split, - Topology, - try_c_order_strides, -) -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.ir.visitor import ExprCloner -from tilefoundry.visitor_registry.contexts import ( - Cost, - CostContext, - FunctionScope, - TypeInferContext, -) -from tilefoundry.visitor_registry.visitors import CostEvaluator, TypeInferVisitor - -from ..errors import ScheduleError -from .facts import PartitionFacts -from .program import ( - OperationSite, - PartitionProgram, - RegionInfo, - ValueInfo, - ceil_div, - expr_location, - tensor_leaves, -) - -PlacementRelation = Literal["SAME_INTERVAL", "CONTAINED"] - - -class PartitionProblemError(ScheduleError): - """The program and projected facts cannot form a finite partition problem. - - A scheduling failure, and reachable as one: a caller asking this - layer to schedule something catches what the layer raises, and a - capability that cannot be scheduled is recorded against that. Sitting - outside `ScheduleError` made a limit of this algorithm unstateable - except as a bare `ValueError`, which is also what a caller passing - nonsense gets -- so the two could not be told apart. - """ - - -@dataclass(frozen=True) -class CandidateBucket: - """One value held in one concrete type, and who can produce it that way.""" - - value_id: int - type_id: int - candidate_ids: tuple[int, ...] - fixed_offset: int | None - is_source: bool - - -@dataclass(frozen=True) -class BucketRequirement: - """The buckets an authored placement constraint admits for one value.""" - - value_id: int - bucket_ids: tuple[int, ...] - source: Expr - metadata: ScheduleConstraintMetadata - - -@dataclass(frozen=True) -class CandidateDependency: - """One candidate's demand on one input bucket, and how they must overlap.""" - - parent_candidate_id: int - input_index: int - child_bucket_id: int - placement_relation: PlacementRelation | None - - -@dataclass(frozen=True) -class OpCandidate: - """One way to run one operation, priced against the projected facts. - - Authored computation and synthesized Reshard are the same concept here: both - produce buckets from buckets at a stated cost. A Reshard is recognised by - having no authored site rather than by a separate record type. - """ - - op: Op - input_bucket_ids: tuple[int, ...] - output_bucket_ids: tuple[int, ...] - output_alias_input_indices: tuple[int | None, ...] - active_mesh: Mesh - topology_count: int - local_cost: Cost - duration_ns: int - total_hbm_bytes: int - hbm_demand_bytes_per_ns: int - moved_bytes: int - site_id: int | None = None - source_call: Call | None = None - source_types: tuple[Type, ...] = () - output_types: tuple[Type, ...] = () - - -@dataclass(frozen=True) -class PartitionProblem: - """A complete finite problem with no Target object or callback.""" - - module: Module - root: Function - topology: Topology - extent: int - facts: PartitionFacts - types: tuple[Type, ...] - values: Mapping[int, ValueInfo] - buckets: Mapping[int, CandidateBucket] - candidates: Mapping[int, OpCandidate] - authored_candidates: Mapping[int, tuple[int, ...]] - dependencies: tuple[CandidateDependency, ...] - requirements: tuple[BucketRequirement, ...] - root_value_ids: tuple[int, ...] - regions: Mapping[int, RegionInfo] - candidate_enclosing_regions: Mapping[int, int | None] = MappingProxyType({}) - value_availability_regions: Mapping[int, int | None] = MappingProxyType({}) - site_order: tuple[int, ...] = () - function_instances: tuple[tuple[tuple[int, ...], Function], ...] = () - diagnostics: tuple[str, ...] = () - - -def _mesh(count: int, topology: str) -> Mesh: - return Mesh((Topology(topology, count),), Layout(shape=(count,), strides=(1,))) - - -def _type_mesh(type: TensorType, fallback: Mesh) -> Mesh: - if isinstance(type.layout, ShardLayout): - mesh = type.layout.mesh - if len(mesh.layout.shape) != 1: - raise PartitionProblemError( - "partition requires rank-one candidate meshes" - ) - return mesh - return fallback - - -def _same_logical_tensor(a: TensorType, b: TensorType) -> bool: - return a.shape == b.shape and a.dtype == b.dtype and a.storage == b.storage - - -def _layout_matches(actual: object, constraint: LayoutConstraint) -> bool: - if not isinstance(actual, ShardLayout): - return False - expected = constraint.layout - shape = tuple(actual.layout.shape) - if len(shape) != len(expected.shape): - return False - for got, want in zip(shape, expected.shape): - if not is_layout_wildcard(want) and got != want: - return False - for topology, attr in constraint.bindings: - if actual.mesh.topologies[0].name != topology: - return False - try: - index = actual.mesh.names.index(topology) - except ValueError: - index = 0 - if index >= len(actual.attrs) or actual.attrs[index] != attr: - return False - return True - - -def _bucket_matches(type: TensorType, constraints: tuple[object, ...]) -> bool: - for constraint in constraints: - if isinstance(constraint, LayoutConstraint) and not _layout_matches( - type.layout, constraint - ): - return False - if isinstance(constraint, MeshConstraint): - if ( - not isinstance(type.layout, ShardLayout) - or type.layout.mesh != constraint.mesh - ): - return False - if isinstance(constraint, StorageConstraint) and type.storage != constraint.storage: - return False - return True - - -def _placement_relation(type: TensorType, mesh: Mesh) -> PlacementRelation | None: - if not isinstance(type.layout, ShardLayout): - return "SAME_INTERVAL" if mesh.layout.shape == (1,) else "CONTAINED" - if type.layout.mesh == mesh: - return "SAME_INTERVAL" - attrs = type.layout.attrs - if attrs and all(isinstance(attr, Broadcast) for attr in attrs): - return "CONTAINED" - return None - - -class _Closer: - def __init__( - self, - program: PartitionProgram, - facts: PartitionFacts, - topology: Topology, - extent: int, - ) -> None: - self.program = program - self.facts = facts - self.topology = topology - self.extent = extent - self.types: list[Type] = [] - self.type_ids: dict[Type, int] = {} - self.values: dict[int, ValueInfo] = dict(program.values) - self.value_types: dict[int, tuple[Type, ...]] = {} - self.buckets: dict[int, CandidateBucket] = {} - self._bucket_candidates: dict[int, list[int]] = {} - self._bucket_by_value_type: dict[tuple[int, int], int] = {} - self.candidates: dict[int, OpCandidate] = {} - self.authored_candidates: dict[int, tuple[int, ...]] = {} - self.dependencies: list[CandidateDependency] = [] - self.requirements: list[BucketRequirement] = [] - self.candidate_enclosing_regions: dict[int, int | None] = {} - self._next_candidate = 0 - self._next_bucket = 0 - self.counts = self._resource_counts() - - def _resource_counts(self) -> tuple[int, ...]: - """The parallel-position counts a candidate may divide work over. - - An extent the program mentions contributes each of its divisors up to the - topology's own extent: dividing a dimension over more positions than the - level has is not a placement this level can hold. - """ - result = {1, self.extent} - for extent in self.program.observed_extents: - for count in range(1, min(extent, self.extent) + 1): - if extent % count == 0: - result.add(count) - return tuple(sorted(result)) - - def _intern(self, type: Type) -> int: - type_id = self.type_ids.get(type) - if type_id is None: - type_id = len(self.types) - self.type_ids[type] = type_id - self.types.append(type) - return type_id - - def _legal_types(self, base: TensorType) -> tuple[Type, ...]: - result: list[TensorType] = [base] - if isinstance(base.layout, ShardLayout): - return tuple(result) - level = self.topology.name - for count in self.counts: - if count == 1: - continue - meshes = [_mesh(count, level)] - meshes.extend( - mesh - for mesh in self.program.required_meshes - if mesh.layout.shape == (count,) and mesh not in meshes - ) - for mesh in meshes: - replicated = TensorType( - shape=base.shape, - dtype=base.dtype, - storage=base.storage, - layout=ShardLayout( - layout=Layout( - shape=base.shape, strides=try_c_order_strides(base.shape) - ), - attrs=(Broadcast(),), - mesh=mesh, - ), - ) - result.append(replicated) - result.append( - replace( - replicated, - layout=replace( - replicated.layout, - attrs=(Partial("sum"),), # type: ignore[arg-type] - ), - ) - ) - for axis, dim in enumerate(base.shape): - if isinstance(dim, int) and not isinstance(dim, bool) and dim % count == 0: - result.append( - make_shard_tensor_type( - base.shape, base.dtype, base.storage, mesh, (Split(axis),) - ) - ) - dedup: list[TensorType] = [] - for type in result: - if type not in dedup: - dedup.append(type) - self._intern(type) - return tuple(dedup) - - def _init_buckets(self) -> None: - for value_id, base in self.program.value_base_types.items(): - self.value_types[value_id] = self._legal_types(base) - for value_id, types in self.value_types.items(): - source_bucket_ids: list[int] = [] - for type in types: - type_id = self._intern(type) - bucket_id = self._next_bucket - self._next_bucket += 1 - self._bucket_by_value_type[(value_id, type_id)] = bucket_id - self._bucket_candidates[bucket_id] = [] - self.buckets[bucket_id] = CandidateBucket( - value_id=value_id, - type_id=type_id, - candidate_ids=(), - fixed_offset=None, - is_source=( - self.values[value_id].role == "normal" - and self.values[value_id].producer_site_id is None - ), - ) - if self.buckets[bucket_id].is_source: - source_bucket_ids.append(bucket_id) - if source_bucket_ids: - self.values[value_id] = replace( - self.values[value_id], source_bucket_ids=tuple(source_bucket_ids) - ) - - def _retag(self, expr: Expr, types: "itertools.chain | object") -> Expr: - class _RetagMutator(ExprCloner): - def __init__(self, type_iter) -> None: - super().__init__() - self.type_iter = type_iter - - def visit_Tuple(self, value: Tuple, ctx=None) -> Expr: - new_elements = tuple(self.visit(element, ctx) for element in value.elements) - if new_elements == value.elements: - return value - return replace(value, elements=new_elements) - - def default_visit(self, value: Expr, ctx=None) -> Expr: - if isinstance(value.type, TensorType): - try: - type = next(self.type_iter) # type: ignore[call-overload] - except StopIteration: - return value - return replace(value, type=type) - return value - - return _RetagMutator(types).visit(expr) - - def _candidate_call( - self, site: OperationSite, input_types: tuple[Type, ...] - ) -> tuple[Call, tuple[Type, ...]]: - type_iter = iter(input_types) - args = tuple(self._retag(arg, type_iter) for arg in site.call.args) - return replace(site.call, args=args), input_types - - def _active_mesh_for_outputs( - self, outputs: tuple[TensorType, ...] - ) -> tuple[Mesh, int]: - fallback = _mesh(1, self.topology.name) - meshes = tuple(_type_mesh(output, fallback) for output in outputs) - mesh = meshes[0] if meshes else fallback - if any(other != mesh for other in meshes[1:]): - raise PartitionProblemError( - "multi-output candidate leaves require one shared Mesh" - ) - if len(mesh.layout.shape) != 1: - raise PartitionProblemError("candidate Mesh must be rank one") - count = mesh.layout.shape[0] - if not isinstance(count, int) or not 1 <= count <= self.extent: - raise PartitionProblemError( - f"candidate Mesh extent {count!r} is outside topology " - f"{self.topology.name!r} extent {self.extent}" - ) - return mesh, count - - def _price( - self, call: Call, cost: Cost, count: int - ) -> tuple[int, int, int, int]: - """What one candidate costs, in the units the solve reasons in. - - A Reshard is charged as pure traffic, so its duration follows the - bandwidth alone. Everything else is charged the worse of its compute and - its traffic, with compute scaled by how much of the device the candidate - occupies. - """ - facts = self.facts - if isinstance(call.target, Reshard): - moved = cost.bytes - duration = ceil_div( - moved * 1_000_000_000, facts.memory_bandwidth_bytes_per_second - ) - demand = ceil_div(moved, duration) if duration else 0 - return duration, moved, demand, moved - if count == 0: - raise PartitionProblemError( - "only Reshard candidates may have topology_count=0" - ) - compute = 0 - for dtype, flops in cost.flops.items(): - compute += ceil_div( - flops * count * 1_000_000_000 * facts.parallel_units, - facts.peak_flops(dtype) * count, - ) - total_bytes = cost.bytes * count - memory = ( - ceil_div(total_bytes * 1_000_000_000, facts.memory_bandwidth_bytes_per_second) - if total_bytes - else 0 - ) - duration = max(compute, memory, 1) if cost.flops or cost.bytes else 0 - demand = ceil_div(total_bytes, duration) if duration else 0 - return duration, total_bytes, demand, 0 - - def _add_candidate( - self, - site_id: int | None, - call: Call, - input_bucket_ids: tuple[int, ...], - output_bucket_ids: tuple[int, ...], - source_types: tuple[Type, ...], - output_types: tuple[Type, ...], - cost: Cost, - *, - reshard: bool = False, - ) -> int: - tensor_outputs = tuple( - type for type in output_types if isinstance(type, TensorType) - ) - mesh, count = self._active_mesh_for_outputs(tensor_outputs) - if isinstance(call.target, Reshard): - count = 0 - duration, total_bytes, demand, moved = self._price(call, cost, count) - aliases = tuple( - 0 if isinstance(call.target, (Reshape, Transpose)) and index == 0 else None - for index, _ in enumerate(output_bucket_ids) - ) - candidate_id = self._next_candidate - self._next_candidate += 1 - self.candidates[candidate_id] = OpCandidate( - op=call.target, - input_bucket_ids=input_bucket_ids, - output_bucket_ids=output_bucket_ids, - output_alias_input_indices=aliases, - active_mesh=mesh, - topology_count=count, - local_cost=cost, - duration_ns=duration, - total_hbm_bytes=total_bytes, - hbm_demand_bytes_per_ns=demand, - moved_bytes=moved, - site_id=site_id, - source_call=None if reshard else call, - source_types=source_types, - output_types=output_types, - ) - if site_id is None: - value_id = self.buckets[output_bucket_ids[0]].value_id - self.candidate_enclosing_regions[candidate_id] = ( - self.program.value_availability_regions.get(value_id) - ) - else: - self.candidate_enclosing_regions[candidate_id] = ( - self.program.site_enclosing_regions.get(site_id) - ) - for bucket_id in output_bucket_ids: - self._bucket_candidates[bucket_id].append(candidate_id) - for input_index, bucket_id in enumerate(input_bucket_ids): - value_type = self.types[self.buckets[bucket_id].type_id] - relation = ( - _placement_relation(value_type, mesh) - if isinstance(value_type, TensorType) - else None - ) - self.dependencies.append( - CandidateDependency(candidate_id, input_index, bucket_id, relation) - ) - return candidate_id - - def _generate_site(self, site: OperationSite) -> None: - choices: list[int] = [] - for refs in site.input_value_ids: - choices.extend(refs) - bucket_options = [] - for value_id in choices: - bucket_options.append( - tuple( - bucket_id - for (candidate_value, _), bucket_id in self._bucket_by_value_type.items() - if candidate_value == value_id - ) - ) - combinations = itertools.product(*bucket_options) if bucket_options else ((),) - found: list[int] = [] - seen: set[tuple] = set() - for input_bucket_ids in combinations: - input_types = tuple( - self.types[self.buckets[bucket_id].type_id] - for bucket_id in input_bucket_ids - ) - candidate_call, selected_types = self._candidate_call(site, input_types) - scope = FunctionScope(self.program.module, self.program.root) - ctx = TypeInferContext(scope=scope) - try: - output_type = TypeInferVisitor().visit(candidate_call, ctx) - except (TypeError, ValueError, NotImplementedError, VerifyError, IndexError): - continue - output_leaves = tensor_leaves(output_type) - if len(output_leaves) != len(site.output_value_ids): - continue - output_types = tuple(type for _, type in output_leaves) - try: - output_bucket_ids = tuple( - self._bucket_by_value_type[(value_id, self._intern(type))] - for value_id, type in zip(site.output_value_ids, output_types) - ) - cost_ctx = CostContext( - scope=scope, - selected_types={ - id(arg): type - for arg, type in zip(candidate_call.args, selected_types) - }, - selected_output_type=output_type, - level=self.topology.name, - topologies=self.program.module.effective_topologies(), - ) - cost = CostEvaluator().visit_Call(candidate_call, cost_ctx) - mesh, count = self._active_mesh_for_outputs(output_types) - except KeyError: - continue - except (TypeError, ValueError) as exc: - raise PartitionProblemError( - f"cost evaluation for {type(site.call.target).__name__} at " - f"{expr_location(site.call)} failed: {exc}" - ) from exc - key = ( - input_bucket_ids, - output_bucket_ids, - mesh, - count, - tuple(sorted(cost.flops.items(), key=lambda item: item[0].name)), - cost.bytes, - ) - if key in seen: - continue - seen.add(key) - found.append( - self._add_candidate( - site.site_id, - candidate_call, - tuple(input_bucket_ids), - output_bucket_ids, - selected_types, - output_types, - cost, - ) - ) - if not found: - raise PartitionProblemError( - f"operation {type(site.call.target).__name__} at " - f"{expr_location(site.call)} has no legal candidates" - ) - self.authored_candidates[site.site_id] = tuple(found) - - def _synthesized_reshards(self) -> None: - """Connect legal buckets no authored operation can reach directly. - - A bucket that some authored candidate already produces needs nothing: a - Reshard is introduced only where a legal placement would otherwise have - no producer at all, and only from a source holding the same logical - tensor. - """ - for value_id, types in self.value_types.items(): - source_buckets = tuple( - bucket_id - for (candidate_value, _), bucket_id in self._bucket_by_value_type.items() - if candidate_value == value_id and self._bucket_candidates[bucket_id] - ) - for target in types: - if not isinstance(target, TensorType): - continue - target_bucket = self._bucket_by_value_type[(value_id, self._intern(target))] - if self._bucket_candidates[target_bucket]: - continue - for source_bucket in source_buckets: - source = self.types[self.buckets[source_bucket].type_id] - if not isinstance(source, TensorType) or not _same_logical_tensor( - source, target - ): - continue - if source == target: - continue - op = Reshard(layout=target.layout, storage=StorageKind.GMEM) - source_expr = self.values[value_id].source - metadata = source_metadata(source_expr) - source_var = Var(type=source, name="reshard_source", metadata=metadata) - call = Call( - type=target, target=op, args=(source_var,), metadata=metadata - ) - cost_ctx = CostContext( - scope=FunctionScope(self.program.module, self.program.root), - selected_types={id(source_var): source}, - selected_output_type=target, - level=self.topology.name, - topologies=self.program.module.effective_topologies(), - ) - try: - cost = CostEvaluator().visit_Call(call, cost_ctx) - except (TypeError, ValueError) as exc: - raise PartitionProblemError( - f"cost evaluation for synthesized Reshard at " - f"{expr_location(source_expr)} failed: {exc}" - ) from exc - self._add_candidate( - None, - call, - (source_bucket,), - (target_bucket,), - (source,), - (target,), - cost, - reshard=True, - ) - break - - def _refs_from_annotation(self, source: Expr) -> tuple[int, ...]: - for refs, candidate_source, _ in self.program.requirement_annotations: - if candidate_source is source: - return refs - return () - - def _finish_buckets(self) -> None: - for bucket_id, bucket in tuple(self.buckets.items()): - fixed_offset = None - for _, source, metadata in self.program.requirement_annotations: - if bucket.value_id not in self._refs_from_annotation(source): - continue - for constraint in metadata.constraints: - if isinstance(constraint, MeshConstraint) and constraint.mesh is not None: - layout = constraint.mesh.layout - if hasattr(layout, "offset"): - fixed_offset = layout.offset - self.buckets[bucket_id] = replace( - bucket, - candidate_ids=tuple(self._bucket_candidates[bucket_id]), - fixed_offset=fixed_offset, - ) - for refs, source, metadata in self.program.requirement_annotations: - value_id = refs[0] - matching = tuple( - bucket_id - for (candidate_value, type_id), bucket_id in self._bucket_by_value_type.items() - if candidate_value == value_id - and isinstance(self.types[type_id], TensorType) - and _bucket_matches(self.types[type_id], metadata.constraints) - ) - if not matching: - raise PartitionProblemError( - f"no candidate bucket satisfies where constraint at " - f"{expr_location(source)}" - ) - self.requirements.append( - BucketRequirement(value_id, matching, source, metadata) - ) - - def _root_connected(self, value_id: int, seen: set[int]) -> bool: - if value_id in seen: - return True - seen.add(value_id) - if self.values[value_id].role != "normal": - return True - bucket_ids = tuple( - bucket_id - for (candidate_value, _), bucket_id in self._bucket_by_value_type.items() - if candidate_value == value_id - ) - if not bucket_ids: - return False - for bucket_id in bucket_ids: - bucket = self.buckets[bucket_id] - if bucket.is_source: - return True - for candidate_id in self._bucket_candidates[bucket_id]: - candidate = self.candidates[candidate_id] - if all( - self._root_connected(self.buckets[child].value_id, seen) - for child in candidate.input_bucket_ids - ): - return True - return False - - def build(self) -> PartitionProblem: - self._init_buckets() - for site in self.program.sites: - self._generate_site(site) - self._synthesized_reshards() - self._finish_buckets() - for value_id in self.program.root_value_ids: - self.values[value_id] = replace( - self.values[value_id], is_final_output=True - ) - if not self._root_connected(value_id, set()): - raise PartitionProblemError( - f"no legal root-connected candidate path for value {value_id}" - ) - reshards = sum( - type(candidate.op) is Reshard and candidate.site_id is None - for candidate in self.candidates.values() - ) - return PartitionProblem( - module=self.program.module, - root=self.program.root, - topology=self.topology, - extent=self.extent, - facts=self.facts, - types=tuple(self.types), - values=MappingProxyType(dict(self.values)), - buckets=MappingProxyType(dict(self.buckets)), - candidates=MappingProxyType(dict(self.candidates)), - authored_candidates=MappingProxyType(dict(self.authored_candidates)), - dependencies=tuple(self.dependencies), - requirements=tuple(self.requirements), - root_value_ids=self.program.root_value_ids, - regions=MappingProxyType(dict(self.program.regions)), - candidate_enclosing_regions=MappingProxyType( - dict(self.candidate_enclosing_regions) - ), - value_availability_regions=MappingProxyType( - dict(self.program.value_availability_regions) - ), - site_order=self.program.site_order, - function_instances=self.program.function_instances, - diagnostics=( - f"ops={len(self.program.sites)}", - f"candidates={len(self.candidates)}", - f"buckets={len(self.buckets)}", - f"reshards={reshards}", - ), - ) - - -def build_partition_problem( - program: PartitionProgram, facts: PartitionFacts, topology: Topology -) -> PartitionProblem: - """Close the problem from an immutable program and already-projected facts.""" - if facts.topology != topology.name: - raise PartitionProblemError( - f"partition facts describe {facts.topology!r}, not topology " - f"{topology.name!r}" - ) - if facts.parallel_units < 1: - raise PartitionProblemError( - "partition facts require at least one parallel unit" - ) - if facts.memory_bandwidth_bytes_per_second < 1: - raise PartitionProblemError( - "partition facts require a positive memory bandwidth" - ) - if facts.memory_capacity_bytes < 1: - raise PartitionProblemError( - "partition facts require a positive memory capacity" - ) - extent = static_dim_value(topology.size) - if extent is None or extent < 1: - raise PartitionProblemError( - f"topology {topology.name!r} requires a static positive extent, got " - f"{topology.size!r}" - ) - if extent > facts.parallel_units: - raise PartitionProblemError( - f"topology {topology.name!r} extent {extent} exceeds the " - f"{facts.parallel_units} parallel units {facts.spec.device_id} states" - ) - return _Closer(program, facts, topology, extent).build() - - -__all__ = [ - "BucketRequirement", - "CandidateBucket", - "CandidateDependency", - "OpCandidate", - "PartitionProblem", - "PartitionProblemError", - "build_partition_problem", -] diff --git a/src/tilefoundry/schedule/partition/program.py b/src/tilefoundry/schedule/partition/program.py deleted file mode 100644 index d0636ffe..00000000 --- a/src/tilefoundry/schedule/partition/program.py +++ /dev/null @@ -1,644 +0,0 @@ -"""Target-independent extraction of one immutable partition program view. - -This walks the authored HIR once and records what is there: every tensor value a -scheduled operation produces or consumes, every operation site, every grid region -and what it carries, and every authored placement constraint. Nothing here -enumerates a choice or reads a machine. The extents observed on the way are kept -raw, because which of them are usable divisors depends on a topology this stage -has not been told about. -""" - -from __future__ import annotations - -from dataclasses import dataclass, replace -from types import MappingProxyType -from typing import Literal, Mapping - -from tilefoundry.ir.constraints import ( - LayoutConstraint, - MeshConstraint, - ScheduleConstraintMetadata, - constraint_metadata, -) -from tilefoundry.ir.core import Call, Constant, Expr, Tuple, Var, diagnostic_location -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.grid_region import GridRegionExpr -from tilefoundry.ir.hir.tensor.reshape import is_induction_var_singleton_reshape -from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem -from tilefoundry.ir.tir.launch import Launch -from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types import TensorType, TupleType, Type -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.shard import Mesh, ShardLayout, Split -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.ir.visitor import ExprVisitor, ExprWalker - -from .facts import PartitionFactsQuery - -ValueRole = Literal["normal", "carry", "yield", "result"] - - -class PartitionProgramError(ValueError): - """The authored program cannot be extracted into a partition program.""" - - -class _ProcessExprVisitor(ExprVisitor[tuple[int, ...]]): - """Process one Expr context while the owner keeps partition side effects.""" - - def __init__(self, owner, function, function_path, env) -> None: - super().__init__() - self.owner = owner - self.function = function - self.function_path = function_path - self.env = env - - def _key(self, expr: Expr) -> tuple[tuple[int, ...], int]: - return self.function_path, id(expr) - - def visit(self, expr: Expr, ctx=None) -> tuple[int, ...]: - cached = self.owner._expr_values.get(self._key(expr)) - if cached is not None: - return cached - return super().visit(expr, ctx) - - def visit_Var(self, expr: Var, ctx=None) -> tuple[int, ...]: - key = self._key(expr) - refs = self.env.get(id(expr)) - if refs is None: - refs = self.owner._param_values.get(key, ()) - self.owner._expr_values[key] = refs - return refs - - def visit_Constant(self, expr: Constant, ctx=None) -> tuple[int, ...]: - self.owner._expr_values[self._key(expr)] = () - return () - - def visit_Tuple(self, expr: Tuple, ctx=None) -> tuple[int, ...]: - refs = tuple(ref for element in expr.elements for ref in self.visit(element, ctx)) - key = self._key(expr) - self.owner._expr_values[key] = refs - self.owner._record_requirement(refs, expr) - return refs - - def visit_GridRegionExpr(self, expr: GridRegionExpr, ctx=None) -> tuple[int, ...]: - refs = self.owner._process_region( - expr, self.function, self.function_path, self.env - ) - key = self._key(expr) - self.owner._expr_values[key] = refs - self.owner._record_requirement(refs, expr) - return refs - - def visit_Call(self, expr: Call, ctx=None) -> tuple[int, ...]: - key = self._key(expr) - arg_refs = tuple(self.visit(arg, ctx) for arg in expr.args) - target = expr.target - if isinstance(target, Function): - call_path = self.function_path + (len(self.owner.function_instances),) - helper_env = dict(zip((id(param) for param in target.params), arg_refs)) - self.owner._process_function(target, call_path, helper_env, parent_call=expr) - refs = self.owner._process_expr(target.body, target, call_path, helper_env) - self.owner._expr_values[key] = refs - self.owner._record_requirement(refs, expr) - return refs - if isinstance(target, (PrimFunction, Launch, SymbolRef)): - raise PartitionProgramError( - f"kernel boundary {type(target).__name__} at " - f"{expr_location(expr)} is not a partitioned operation" - ) - if is_induction_var_singleton_reshape(expr): - refs = arg_refs[0] if arg_refs else () - self.owner._expr_values[key] = refs - self.owner._record_requirement(refs, expr) - return refs - if isinstance(target, TupleGetItem): - source_refs = arg_refs[0] if arg_refs else () - fields = tensor_leaves(expr.args[0].type) - index = target.index - if index < 0 or index >= len(fields): - raise PartitionProgramError( - f"TupleGetItem index {index} is out of range at " - f"{expr_location(expr)}" - ) - start = sum( - len(tensor_leaves(field)) - for field in expr.args[0].type.fields[:index] # type: ignore[union-attr] - ) - refs = source_refs[ - start : start - + len(tensor_leaves(expr.args[0].type.fields[index])) # type: ignore[union-attr] - ] - 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 output_types - ) - self.owner.sites.append( - OperationSite(site_id, expr, self.function_path, arg_refs, output_refs) - ) - self.owner.site_order.append(site_id) - self.owner.site_enclosing_regions[site_id] = ( - self.owner._active_regions[-1] if self.owner._active_regions else None - ) - if self.owner._active_regions: - region_id = self.owner._active_regions[-1] - info = self.owner.regions[region_id] - self.owner.regions[region_id] = replace( - info, operation_site_ids=(*info.operation_site_ids, site_id) - ) - self.owner._expr_values[key] = output_refs - self.owner._record_requirement(output_refs, expr) - return output_refs - - def default_visit(self, expr: Expr, ctx=None) -> tuple[int, ...]: - self.owner._expr_values[self._key(expr)] = () - return () - - -def expr_location(expr: Expr) -> str: - """The clearest location this expression can name itself by.""" - return diagnostic_location(expr) or getattr(expr, "name", None) or "" - - -def ceil_div(a: int, b: int) -> int: - """Integer division that rounds away from zero.""" - return (a + b - 1) // b - - -def tensor_leaves( - type: Type, path: tuple[int, ...] = () -) -> tuple[tuple[tuple[int, ...], TensorType], ...]: - """Every tensor leaf of a possibly nested type, with its index path.""" - if isinstance(type, TensorType): - return ((path, type),) - if isinstance(type, TupleType): - return tuple( - leaf - for index, field in enumerate(type.fields) - for leaf in tensor_leaves(field, path + (index,)) - ) - return () - - -@dataclass(frozen=True) -class ValueInfo: - """One tensor value the partition decides a placement for.""" - - source: Expr - leaf_path: tuple[int, ...] - is_const: bool - source_bucket_ids: tuple[int, ...] = () - producer_site_id: int | None = None - function_path: tuple[int, ...] = () - is_final_output: bool = False - role: ValueRole = "normal" - - -@dataclass(frozen=True) -class OperationSite: - """One authored call, and the value IDs on each side of it.""" - - site_id: int - call: Call - function_path: tuple[int, ...] - input_value_ids: tuple[tuple[int, ...], ...] - output_value_ids: tuple[int, ...] - - -@dataclass(frozen=True) -class RegionInfo: - """One grid region, its trip count, and the sites inside it.""" - - source: GridRegionExpr - parent_region_id: int | None - trip_count: int - operation_site_ids: tuple[int, ...] - init_use_ids: tuple[int, ...] - backedge_use_ids: tuple[int, ...] - carry_infos: tuple["RegionCarryInfo", ...] = () - result_value_ids: tuple[int, ...] = () - - -@dataclass(frozen=True) -class RegionCarryInfo: - """One value carried around a region, at each of its four positions.""" - - init_value_id: int - carried_value_id: int - yield_value_id: int - result_value_id: int - - -@dataclass(frozen=True) -class PartitionProgram: - """What the authored program states, before any choice is enumerated.""" - - module: Module - root: Function - values: Mapping[int, ValueInfo] - value_base_types: Mapping[int, TensorType] - sites: tuple[OperationSite, ...] - site_order: tuple[int, ...] - site_enclosing_regions: Mapping[int, int | None] - value_availability_regions: Mapping[int, int | None] - regions: Mapping[int, RegionInfo] - requirement_annotations: tuple[ - tuple[tuple[int, ...], Expr, ScheduleConstraintMetadata], ... - ] - function_instances: tuple[tuple[tuple[int, ...], Function], ...] - root_value_ids: tuple[int, ...] - observed_extents: tuple[int, ...] - required_meshes: tuple[Mesh, ...] - - def facts_query(self, topology: str) -> PartitionFactsQuery: - """Return the explicit query used to project target facts once.""" - return PartitionFactsQuery(topology=topology) - - -class _Extractor: - def __init__(self, module: Module, root: Function) -> None: - self.module = module - self.root = root - self.values: dict[int, ValueInfo] = {} - self.value_base_types: dict[int, TensorType] = {} - self.regions: dict[int, RegionInfo] = {} - self.value_availability_regions: dict[int, int | None] = {} - self.site_enclosing_regions: dict[int, int | None] = {} - self.sites: list[OperationSite] = [] - self.site_order: list[int] = [] - self.function_instances: list[tuple[tuple[int, ...], Function]] = [] - self.requirement_annotations: list[ - tuple[tuple[int, ...], Expr, ScheduleConstraintMetadata] - ] = [] - self._expr_values: dict[tuple[tuple[int, ...], int], tuple[int, ...]] = {} - self._param_values: dict[tuple[tuple[int, ...], int], tuple[int, ...]] = {} - self._active_functions: set[int] = set() - self._active_regions: list[int] = [] - self._next_value = 0 - self._next_site = 0 - self._next_region = 0 - self._next_use = 0 - self.required_meshes: list[Mesh] = [] - self.observed_extents = self._observed_extents() - - def _observed_extents(self) -> tuple[int, ...]: - """Every static extent the program mentions a division by. - - Split axes, mesh shapes, and authored placement constraints all name a - number of parallel positions. They are collected raw: how many of them a - candidate may actually use is a property of the topology, which this - stage does not see. - """ - extents: set[int] = set() - seen_functions: set[int] = set() - - def record(dim: object) -> None: - if isinstance(dim, int) and not isinstance(dim, bool) and dim > 0: - extents.add(dim) - - def visit_type(type: Type) -> None: - for _, tensor in tensor_leaves(type): - if isinstance(tensor.layout, ShardLayout): - for attr in tensor.layout.attrs: - if isinstance(attr, Split) and attr.axis < len( - tensor.layout.layout.shape - ): - record(tensor.layout.layout.shape[attr.axis]) - for dim in tensor.layout.mesh.layout.shape: - record(dim) - - owner = self - - class _ExtentVisitor(ExprWalker[None]): - def visit(self, expr, ctx=None): - if expr is not None and not isinstance(expr, Function): - self._record_expr(expr) - return super().visit(expr, ctx) - - def _record_expr(self, expr: Expr) -> None: - visit_type(expr.type) - metadata = constraint_metadata(expr) - if metadata is None: - return - for constraint in metadata.constraints: - if isinstance(constraint, LayoutConstraint): - for _, attr in constraint.bindings: - if isinstance(attr, Split) and attr.axis < len( - constraint.layout.shape - ): - record(constraint.layout.shape[attr.axis]) - elif isinstance(constraint, MeshConstraint) and constraint.mesh is not None: - owner.required_meshes.append(constraint.mesh) - for dim in constraint.mesh.layout.shape: - record(dim) - - def visit_function(self, function: Function, ctx=None) -> None: - if id(function) in seen_functions: - return - seen_functions.add(id(function)) - for param in function.params: - self.visit(param, ctx) - if function.body is not None: - self.visit(function.body, ctx) - - def visit_Call(self, expr: Call, ctx=None) -> None: - if isinstance(expr.target, Function): - self.visit_function(expr.target, ctx) - self.visit_operands(expr, ctx) - - def visit_Tuple(self, expr: Tuple, ctx=None) -> None: - self.visit_operands(expr, ctx) - - def visit_GridRegionExpr(self, expr: GridRegionExpr, ctx=None) -> None: - self.visit_operands(expr, ctx) - - def visit_Function(self, expr: Function, ctx=None) -> None: - self.visit_function(expr, ctx) - - def visit_Var(self, expr: Var, ctx=None) -> None: - return None - - def visit_Constant(self, expr: Constant, ctx=None) -> None: - return None - - def visit_SymbolRef(self, expr: SymbolRef, ctx=None) -> None: - return None - - def visit_ShapeOf(self, expr: Expr, ctx=None) -> None: - return None - - _ExtentVisitor().visit_function(self.root) - return tuple(sorted(extents)) - - def _new_value( - self, - source: Expr, - type: TensorType, - leaf_path: tuple[int, ...], - function_path: tuple[int, ...], - *, - is_const: bool = False, - producer_site_id: int | None = None, - role: ValueRole = "normal", - ) -> int: - if type.storage is not StorageKind.GMEM: - loc = diagnostic_location(source) or expr_location(self.root) - raise PartitionProgramError( - f"tensor storage {type.storage!r} at {loc} is unsupported; " - "partitioned tensor values must reside in GMEM" - ) - value_id = self._next_value - self._next_value += 1 - self.values[value_id] = ValueInfo( - source=source, - leaf_path=leaf_path, - is_const=is_const, - producer_site_id=producer_site_id, - function_path=function_path, - role=role, - ) - self.value_availability_regions[value_id] = ( - self._active_regions[-1] if self._active_regions else None - ) - self.value_base_types[value_id] = type - return value_id - - def _source_value(self, param: Var, function_path: tuple[int, ...]) -> tuple[int, ...]: - refs = tuple( - self._new_value(param, type, path, function_path, is_const=param.is_const) - for path, type in tensor_leaves(param.type) - ) - self._param_values[(function_path, id(param))] = refs - self._record_requirement(refs, param) - return refs - - def _record_requirement(self, refs: tuple[int, ...], source: Expr) -> None: - metadata = constraint_metadata(source) - if metadata is not None: - for ref in refs: - self.requirement_annotations.append(((ref,), source, metadata)) - - def _process_expr( - self, - expr: Expr | None, - function: Function, - function_path: tuple[int, ...], - env: Mapping[int, tuple[int, ...]], - ) -> tuple[int, ...]: - if expr is None: - return () - key = (function_path, id(expr)) - cached = self._expr_values.get(key) - if cached is not None: - return cached - return _ProcessExprVisitor(self, function, function_path, env).visit(expr) - - def _process_region( - self, - region: GridRegionExpr, - function: Function, - function_path: tuple[int, ...], - env: Mapping[int, tuple[int, ...]], - ) -> tuple[int, ...]: - start = static_dim_value(region.start) - stop = static_dim_value(region.extent) - step = static_dim_value(region.step) - context = ( - f"GridRegion at {diagnostic_location(region) or expr_location(function)}" - ) - if start is None or stop is None or step is None: - raise PartitionProgramError( - f"{context} requires static start, stop, and step" - ) - if start < 0 or step <= 0: - raise PartitionProgramError( - f"{context} has invalid start/step ({start}, {step})" - ) - trip_count = ceil_div(stop - start, step) if stop > start else 0 - if trip_count <= 0: - raise PartitionProgramError(f"{context} has non-positive trip count") - region_id = self._next_region - self._next_region += 1 - info = RegionInfo( - source=region, - parent_region_id=self._active_regions[-1] if self._active_regions else None, - trip_count=trip_count, - operation_site_ids=(), - init_use_ids=tuple(self._new_use() for _ in region.init_args), - backedge_use_ids=tuple(self._new_use() for _ in region.yield_values), - ) - self.regions[region_id] = info - init_ref_groups = tuple( - self._process_expr(value, function, function_path, env) - for value in region.init_args - ) - init_refs = tuple(ref for refs in init_ref_groups for ref in refs) - self._active_regions.append(region_id) - try: - carried_ref_groups = tuple( - tuple( - self._new_value( - phi, type, path, function_path + (region_id,), role="carry" - ) - for path, type in tensor_leaves(phi.type) - ) - for phi in region.carried_args - ) - phi_env = dict(env) - for phi, refs in zip(region.carried_args, carried_ref_groups): - phi_env[id(phi)] = refs - body_refs = self._process_expr( - region.body, function, function_path + (region_id,), phi_env - ) - yield_ref_groups = tuple( - self._process_expr( - value, function, function_path + (region_id,), phi_env - ) - for value in region.yield_values - ) - yield_refs = tuple(ref for refs in yield_ref_groups for ref in refs) - finally: - self._active_regions.pop() - parent_region_id = info.parent_region_id - for ref in yield_refs: - if self.values[ref].role == "normal": - self.values[ref] = replace(self.values[ref], role="yield") - if region.carried_args: - result_refs = tuple( - self._new_value( - self.values[yield_ref].source, - self.value_base_types[yield_ref], - self.values[yield_ref].leaf_path, - function_path, - role="result", - ) - for yield_ref in yield_refs - ) - else: - result_refs = body_refs - for ref in result_refs: - self.value_availability_regions[ref] = parent_region_id - carry_infos = tuple( - RegionCarryInfo( - init_value_id=init_ref, - carried_value_id=carried_ref, - yield_value_id=yield_ref, - result_value_id=result_ref, - ) - for init_ref, carried_ref, yield_ref, result_ref in zip( - init_refs, - (ref for refs in carried_ref_groups for ref in refs), - yield_refs, - result_refs, - ) - ) - self.regions[region_id] = replace( - self.regions[region_id], - carry_infos=carry_infos, - result_value_ids=result_refs, - ) - return result_refs - - def _new_use(self) -> int: - value = self._next_use - self._next_use += 1 - return value - - def _process_function( - self, - function: Function, - function_path: tuple[int, ...], - env: Mapping[int, tuple[int, ...]], - *, - parent_call: Call | None = None, - ) -> None: - if function.body is None: - raise PartitionProgramError( - f"helper function {function.name!r} has no body at " - f"{expr_location(parent_call or function)}" - ) - if id(function) in self._active_functions: - raise PartitionProgramError( - f"recursive helper call to {function.name!r} at " - f"{expr_location(parent_call or function)}" - ) - self._active_functions.add(id(function)) - self.function_instances.append((function_path, function)) - try: - function_env = dict(env) - if function is self.root: - for param in function.params: - function_env[id(param)] = self._source_value(param, function_path) - else: - for param, refs in zip(function.params, env.values()): - self._param_values[(function_path, id(param))] = refs - self._process_expr(function.body, function, function_path, function_env) - finally: - self._active_functions.remove(id(function)) - - def build(self) -> PartitionProgram: - self._process_function(self.root, (), {}, parent_call=None) - root_refs = self._expr_values.get(((), id(self.root.body)), ()) - if not root_refs: - raise PartitionProgramError("root function has no tensor result value") - return PartitionProgram( - module=self.module, - root=self.root, - values=MappingProxyType(dict(self.values)), - value_base_types=MappingProxyType(dict(self.value_base_types)), - sites=tuple(self.sites), - site_order=tuple(self.site_order), - site_enclosing_regions=MappingProxyType(dict(self.site_enclosing_regions)), - value_availability_regions=MappingProxyType( - dict(self.value_availability_regions) - ), - regions=MappingProxyType(dict(self.regions)), - requirement_annotations=tuple(self.requirement_annotations), - function_instances=tuple(self.function_instances), - root_value_ids=tuple(root_refs), - observed_extents=self.observed_extents, - required_meshes=tuple(self.required_meshes), - ) - - -def build_partition_program(module: Module, function: Function) -> PartitionProgram: - """Extract one deterministic view of what the authored program states.""" - if not isinstance(function, Function): - raise TypeError( - f"partition program: root must be a HIR Function, got " - f"{type(function).__name__}" - ) - if not module.owns(function, derived=True): - raise PartitionProgramError( - f"{function.name!r} is not a function of module {module.name!r}" - ) - return _Extractor(module, function).build() - - -__all__ = [ - "OperationSite", - "PartitionProgram", - "PartitionProgramError", - "RegionCarryInfo", - "RegionInfo", - "ValueInfo", - "build_partition_program", - "ceil_div", - "expr_location", - "tensor_leaves", -] diff --git a/src/tilefoundry/schedule/partition/solve.py b/src/tilefoundry/schedule/partition/solve.py deleted file mode 100644 index 4acbc905..00000000 --- a/src/tilefoundry/schedule/partition/solve.py +++ /dev/null @@ -1,969 +0,0 @@ -"""The CP-SAT solve over one closed partition problem. - -Every number this model needs is already in the problem: durations, traffic -demands, capacities, and how many parallel positions there are. Nothing here -resolves a Target, projects a fact, or asks the hardware a question, so what the -solve minimises is fully determined by its input. -""" - -from __future__ import annotations - -import json -import math -from dataclasses import dataclass -from typing import Literal - -from ortools.sat.python import cp_model - -from tilefoundry.ir.hir.sharding.reshard import Reshard -from tilefoundry.ir.hir.tensor.reshape import Reshape -from tilefoundry.ir.hir.tensor.transpose import Transpose -from tilefoundry.ir.types import TensorType, Type, tensor_bytes -from tilefoundry.ir.types.shard import ComposedLayout, ShardLayout -from tilefoundry.schedule import ScheduleOptions - -from .problem import OpCandidate, PartitionProblem -from .program import RegionInfo - -_INT64_MAX = (1 << 63) - 1 - - -class PartitionSolveError(RuntimeError): - """The closed problem has no schedule, or the solver could not decide.""" - - -@dataclass(frozen=True) -class ExecutionInterval: - """One selected candidate's half-open execution interval.""" - - start_ns: int - end_ns: int - - -@dataclass(frozen=True) -class PartitionSolution: - """What the solve selected, and how sure it is of the objective.""" - - status: Literal["OPTIMAL", "FEASIBLE_NOT_PROVEN"] - selected_candidate_ids: tuple[int, ...] - selected_bucket_ids: tuple[int, ...] - candidate_intervals_ns: tuple[tuple[int, ExecutionInterval], ...] - bucket_offsets: tuple[tuple[int, int], ...] - makespan_ns: int - best_bound_ns: int - gap: float - - -@dataclass -class _CpModelState: - model: cp_model.CpModel - pick_candidates: dict[int, cp_model.IntVar] - pick_buckets: dict[int, cp_model.IntVar] - terminal_buckets: dict[int, cp_model.IntVar] - starts: dict[int, cp_model.IntVar] - ends: dict[int, cp_model.IntVar] - ready: dict[int, cp_model.IntVar] - offsets: dict[int, cp_model.IntVar] - makespan: cp_model.IntVar - horizon_ns: int - - -def _is_reshard(candidate: OpCandidate) -> bool: - return isinstance(candidate.op, Reshard) - - -def _is_view(candidate: OpCandidate) -> bool: - return isinstance(candidate.op, (Reshape, Transpose)) and candidate.duration_ns == 0 - - -def _checked_add(total: int, value: int, context: str) -> int: - result = total + value - if result < 0 or result > _INT64_MAX: - raise PartitionSolveError(f"{context} exceeds the solver integer domain") - return result - - -def _checked_mul(left: int, right: int, context: str) -> int: - if left < 0 or right < 0 or (left and right > _INT64_MAX // left): - raise PartitionSolveError(f"{context} exceeds the solver integer domain") - return left * right - - -def _region_chain( - problem: PartitionProblem, region_id: int | None -) -> tuple[RegionInfo, ...]: - chain: list[RegionInfo] = [] - while region_id is not None: - region = problem.regions[region_id] - chain.append(region) - region_id = region.parent_region_id - return tuple(reversed(chain)) - - -def _horizon(problem: PartitionProblem) -> int: - horizon = 0 - for candidate_id, candidate in problem.candidates.items(): - if candidate.duration_ns <= 0: - continue - duration = candidate.duration_ns - for region in _region_chain( - problem, problem.candidate_enclosing_regions.get(candidate_id) - ): - duration = _checked_mul(duration, region.trip_count, "horizon") - horizon = _checked_add(horizon, duration, "horizon") - return horizon - - -def _tensor_mesh_count(type: Type) -> int: - if not isinstance(type, TensorType) or not isinstance(type.layout, ShardLayout): - return 1 - shape = type.layout.mesh.layout.shape - count = shape[0] - if not isinstance(count, int) or count <= 0: - raise PartitionSolveError( - f"bucket Mesh count must be a positive integer, got {count!r}" - ) - return count - - -def _mesh_offset(type: Type) -> int | None: - if not isinstance(type, TensorType) or not isinstance(type.layout, ShardLayout): - return None - layout = type.layout.mesh.layout - return layout.offset if isinstance(layout, ComposedLayout) else None - - -def _buckets_for_value(problem: PartitionProblem, value_id: int) -> tuple[int, ...]: - return tuple( - bucket_id - for bucket_id, bucket in problem.buckets.items() - if bucket.value_id == value_id - ) - - -def _buckets_by_type(problem: PartitionProblem, value_id: int) -> dict[int, int]: - return { - bucket.type_id: bucket_id - for bucket_id, bucket in problem.buckets.items() - if bucket.value_id == value_id - } - - -def _source_value_ids(problem: PartitionProblem) -> tuple[int, ...]: - return tuple( - value_id - for value_id, value in problem.values.items() - if value.role == "normal" - and value.producer_site_id is None - and value.function_path == () - ) - - -def _result_region_ids(problem: PartitionProblem) -> dict[int, int]: - result_regions: dict[int, int] = {} - for region_id, region in problem.regions.items(): - for value_id in region.result_value_ids: - result_regions[value_id] = region_id - return result_regions - - -def _descendant_regions(problem: PartitionProblem, region_id: int) -> set[int]: - descendants = {region_id} - changed = True - while changed: - changed = False - for candidate_id, region in problem.regions.items(): - if region.parent_region_id in descendants and candidate_id not in descendants: - descendants.add(candidate_id) - changed = True - return descendants - - -def _allocation_groups( - problem: PartitionProblem, -) -> tuple[tuple[int, tuple[int, ...]], ...]: - """Conservative physical groups covering all possible bucket selections. - - Carry facts are unconditional in-place aliases. View aliases are selected - candidate facts, so they are intentionally kept as singleton groups in the - pre-solve capacity model. This can overestimate resident bytes, but cannot - merge an unselected view path and undercount them. - """ - bucket_ids = tuple(sorted(problem.buckets)) - parent = {bucket_id: bucket_id for bucket_id in bucket_ids} - - def find(bucket_id: int) -> int: - while parent[bucket_id] != bucket_id: - parent[bucket_id] = parent[parent[bucket_id]] - bucket_id = parent[bucket_id] - return bucket_id - - def union(left: int, right: int) -> None: - left_root = find(left) - right_root = find(right) - if left_root != right_root: - parent[right_root] = left_root - - for region in problem.regions.values(): - for carry in region.carry_infos: - carry_values = ( - carry.init_value_id, - carry.carried_value_id, - carry.yield_value_id, - carry.result_value_id, - ) - for left_value, right_value in zip(carry_values, carry_values[1:]): - left_by_type = _buckets_by_type(problem, left_value) - right_by_type = _buckets_by_type(problem, right_value) - for type_id in left_by_type.keys() & right_by_type.keys(): - union(left_by_type[type_id], right_by_type[type_id]) - - groups: dict[int, list[int]] = {} - for bucket_id in bucket_ids: - groups.setdefault(find(bucket_id), []).append(bucket_id) - return tuple( - (group_id, tuple(sorted(group_bucket_ids))) - for group_id, group_bucket_ids in sorted(groups.items()) - ) - - -def _add_exactly_one( - model: cp_model.CpModel, literals: list[cp_model.IntVar], label: str -) -> None: - if not literals: - raise PartitionSolveError(f"no selectable {label}") - model.AddExactlyOne(literals) - - -def _build_model(problem: PartitionProblem) -> _CpModelState: - horizon = _horizon(problem) - extent = problem.extent - bandwidth_per_ns = math.ceil( - problem.facts.memory_bandwidth_bytes_per_second / 1_000_000_000 - ) - model = cp_model.CpModel() - pick_candidates = { - candidate_id: model.NewBoolVar(f"pick_candidate_{candidate_id}") - for candidate_id in sorted(problem.candidates) - } - pick_buckets = { - bucket_id: model.NewBoolVar(f"pick_bucket_{bucket_id}") - for bucket_id in sorted(problem.buckets) - } - - for site_id in problem.site_order: - _add_exactly_one( - model, - [ - pick_candidates[candidate_id] - for candidate_id in problem.authored_candidates[site_id] - ], - f"authored candidates for site {site_id}", - ) - for value_id in _source_value_ids(problem): - _add_exactly_one( - model, - [ - pick_buckets[bucket_id] - for bucket_id in _buckets_for_value(problem, value_id) - if problem.buckets[bucket_id].is_source - ], - f"source buckets for value {value_id}", - ) - for requirement in problem.requirements: - _add_exactly_one( - model, - [pick_buckets[bucket_id] for bucket_id in requirement.bucket_ids], - f"requirement buckets for value {requirement.value_id}", - ) - terminal_buckets: dict[int, cp_model.IntVar] = {} - for value_id, value in problem.values.items(): - if value.is_final_output: - value_bucket_ids = _buckets_for_value(problem, value_id) - reshard_output_buckets = tuple( - bucket_id - for bucket_id in value_bucket_ids - if any( - _is_reshard(problem.candidates[candidate_id]) - for candidate_id in problem.buckets[bucket_id].candidate_ids - ) - ) - if not reshard_output_buckets: - _add_exactly_one( - model, - [pick_buckets[bucket_id] for bucket_id in value_bucket_ids], - f"function result buckets for value {value_id}", - ) - continue - terminal_literals = [] - for bucket_id in value_bucket_ids: - terminal = model.NewBoolVar(f"terminal_root_bucket_{bucket_id}") - terminal_buckets[bucket_id] = terminal - model.AddImplication(terminal, pick_buckets[bucket_id]) - terminal_literals.append(terminal) - _add_exactly_one( - model, terminal_literals, f"function result buckets for value {value_id}" - ) - - for bucket_id, bucket in problem.buckets.items(): - if bucket.is_source or problem.values[bucket.value_id].role != "normal": - continue - producers = [ - pick_candidates[candidate_id] for candidate_id in bucket.candidate_ids - ] - model.Add(sum(producers) == pick_buckets[bucket_id]) - for candidate_id, candidate in problem.candidates.items(): - present = pick_candidates[candidate_id] - for bucket_id in (*candidate.input_bucket_ids, *candidate.output_bucket_ids): - model.AddImplication(present, pick_buckets[bucket_id]) - - demand_literals_by_bucket: dict[int, list[cp_model.IntVar]] = {} - for output_bucket, terminal in terminal_buckets.items(): - demand_literals_by_bucket.setdefault(output_bucket, []).append(terminal) - for other_candidate_id, other_candidate in problem.candidates.items(): - for input_bucket in other_candidate.input_bucket_ids: - demand_literals_by_bucket.setdefault(input_bucket, []).append( - pick_candidates[other_candidate_id] - ) - for requirement in problem.requirements: - for bucket_id in requirement.bucket_ids: - demand_literals_by_bucket.setdefault(bucket_id, []).append( - pick_buckets[bucket_id] - ) - - for candidate_id, candidate in problem.candidates.items(): - if not _is_reshard(candidate) or candidate.site_id is not None: - continue - output_bucket = candidate.output_bucket_ids[0] - demand_literals = tuple( - dict.fromkeys(demand_literals_by_bucket.get(output_bucket, ())) - ) - if demand_literals: - demand = model.NewBoolVar(f"reshard_demand_{output_bucket}") - for literal in demand_literals: - model.AddImplication(literal, demand) - model.AddBoolOr([demand.Not(), *demand_literals]) - model.AddImplication(pick_candidates[candidate_id], demand) - else: - model.Add(pick_candidates[candidate_id] == 0) - - starts: dict[int, cp_model.IntVar] = {} - ends: dict[int, cp_model.IntVar] = {} - ready = { - bucket_id: model.NewIntVar(0, horizon, f"ready_{bucket_id}") - for bucket_id in sorted(problem.buckets) - } - merged_geometry_sites = { - site_id - for site_id in problem.site_order - if all( - problem.candidates[candidate_id].duration_ns > 0 - and not _is_reshard(problem.candidates[candidate_id]) - for candidate_id in problem.authored_candidates[site_id] - ) - } - merged_geometry_candidates = { - candidate_id - for site_id in merged_geometry_sites - for candidate_id in problem.authored_candidates[site_id] - } - positive_intervals: dict[int, cp_model.IntervalVar] = {} - for candidate_id, candidate in problem.candidates.items(): - if candidate.duration_ns <= 0: - continue - duration = candidate.duration_ns - for region in _region_chain( - problem, problem.candidate_enclosing_regions.get(candidate_id) - ): - duration = _checked_mul(duration, region.trip_count, "candidate duration") - start = model.NewIntVar(0, horizon, f"start_{candidate_id}") - end = model.NewIntVar(0, horizon, f"end_{candidate_id}") - starts[candidate_id] = start - ends[candidate_id] = end - if candidate_id not in merged_geometry_candidates: - positive_intervals[candidate_id] = model.NewOptionalIntervalVar( - start, - duration, - end, - pick_candidates[candidate_id], - f"execution_{candidate_id}", - ) - else: - model.Add(end == start + duration).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - model.Add(start == horizon).OnlyEnforceIf(pick_candidates[candidate_id].Not()) - model.Add(end == 0).OnlyEnforceIf(pick_candidates[candidate_id].Not()) - for input_bucket in candidate.input_bucket_ids: - model.Add(start >= ready[input_bucket]).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - - result_regions = _result_region_ids(problem) - for bucket_id, bucket in problem.buckets.items(): - if bucket.is_source: - model.Add(ready[bucket_id] == 0) - for candidate_id, candidate in problem.candidates.items(): - present = pick_candidates[candidate_id] - if candidate.duration_ns == 0: - for output_bucket in candidate.output_bucket_ids: - if candidate.input_bucket_ids: - model.Add( - ready[output_bucket] == ready[candidate.input_bucket_ids[0]] - ).OnlyEnforceIf(present) - continue - for output_bucket in candidate.output_bucket_ids: - model.Add(ready[output_bucket] == ends[candidate_id]).OnlyEnforceIf(present) - - offsets: dict[int, cp_model.IntVar] = {} - for bucket_id, bucket in problem.buckets.items(): - bucket_type = problem.types[bucket.type_id] - count = _tensor_mesh_count(bucket_type) - if count > extent: - raise PartitionSolveError( - f"bucket {bucket_id} count {count} exceeds topology extent {extent}" - ) - offset = model.NewIntVar(0, extent - count, f"offset_{bucket_id}") - offsets[bucket_id] = offset - fixed_offset = bucket.fixed_offset - if fixed_offset is None: - fixed_offset = _mesh_offset(bucket_type) - if fixed_offset is not None: - if not 0 <= fixed_offset <= extent - count: - raise PartitionSolveError( - f"bucket {bucket_id} fixed offset {fixed_offset} is outside " - f"topology extent {extent}" - ) - model.Add(offset == fixed_offset).OnlyEnforceIf(pick_buckets[bucket_id]) - - region_starts: dict[int, cp_model.IntVar] = { - region_id: model.NewIntVar(0, horizon, f"region_start_{region_id}") - for region_id in problem.regions - } - region_ends: dict[int, cp_model.IntVar] = { - region_id: model.NewIntVar(0, horizon, f"region_end_{region_id}") - for region_id in problem.regions - } - region_members = { - region_id: _descendant_regions(problem, region_id) - for region_id in problem.regions - } - for region_id, region in problem.regions.items(): - members = region_members[region_id] - member_candidates = [ - candidate_id - for candidate_id, candidate_region in problem.candidate_enclosing_regions.items() - if candidate_region in members and candidate_id in starts - ] - child_regions = [ - child_id - for child_id, child in problem.regions.items() - if child.parent_region_id == region_id - ] - starts_for_min = [starts[candidate_id] for candidate_id in member_candidates] - starts_for_min.extend( - region_starts[child_id] - for child_id in child_regions - if child_id in region_starts - ) - ends_for_max = [ends[candidate_id] for candidate_id in member_candidates] - ends_for_max.extend( - region_ends[child_id] for child_id in child_regions if child_id in region_ends - ) - if not starts_for_min or not ends_for_max: - raise PartitionSolveError( - f"GridRegion {region_id} has no positive-duration work" - ) - model.AddMinEquality(region_starts[region_id], starts_for_min) - model.AddMaxEquality(region_ends[region_id], ends_for_max) - for carry in region.carry_infos: - carry_values = ( - carry.init_value_id, - carry.carried_value_id, - carry.yield_value_id, - carry.result_value_id, - ) - for left_value, right_value in zip(carry_values, carry_values[1:]): - left_by_type = _buckets_by_type(problem, left_value) - right_by_type = _buckets_by_type(problem, right_value) - for type_id in left_by_type.keys() & right_by_type.keys(): - left_bucket = left_by_type[type_id] - right_bucket = right_by_type[type_id] - model.Add(pick_buckets[left_bucket] == pick_buckets[right_bucket]) - model.Add( - offsets[left_bucket] == offsets[right_bucket] - ).OnlyEnforceIf( - [pick_buckets[left_bucket], pick_buckets[right_bucket]] - ) - for carry in region.carry_infos: - for bucket_id in _buckets_for_value(problem, carry.init_value_id): - model.Add(ready[bucket_id] <= region_starts[region_id]).OnlyEnforceIf( - pick_buckets[bucket_id] - ) - - for candidate_id, candidate in problem.candidates.items(): - if candidate_id not in starts: - continue - candidate_region = problem.candidate_enclosing_regions.get(candidate_id) - for bucket_id in candidate.input_bucket_ids: - result_region = result_regions.get(problem.buckets[bucket_id].value_id) - if result_region is None or candidate_region in _descendant_regions( - problem, result_region - ): - continue - model.Add( - starts[candidate_id] >= region_ends[result_region] - ).OnlyEnforceIf(pick_candidates[candidate_id]) - - makespan = model.NewIntVar(0, horizon, "makespan") - makespan_terms = list(ends.values()) - makespan_terms.extend(region_ends.values()) - if makespan_terms: - model.AddMaxEquality(makespan, makespan_terms) - else: - model.Add(makespan == 0) - - topology_intervals: list[cp_model.IntervalVar] = [] - time_intervals: list[cp_model.IntervalVar] = [] - for site_id in sorted(merged_geometry_sites): - site_start = model.NewIntVar(0, horizon, f"site_start_{site_id}") - site_end = model.NewIntVar(0, horizon, f"site_end_{site_id}") - site_duration = model.NewIntVar(0, horizon, f"site_duration_{site_id}") - site_offset = model.NewIntVar(0, extent, f"site_offset_{site_id}") - site_count = model.NewIntVar(1, extent, f"site_count_{site_id}") - site_offset_end = model.NewIntVar(0, extent, f"site_offset_end_{site_id}") - model.Add(site_offset_end == site_offset + site_count) - for candidate_id in problem.authored_candidates[site_id]: - candidate = problem.candidates[candidate_id] - present = pick_candidates[candidate_id] - model.Add(site_start == starts[candidate_id]).OnlyEnforceIf(present) - model.Add(site_end == ends[candidate_id]).OnlyEnforceIf(present) - model.Add( - site_duration == ends[candidate_id] - starts[candidate_id] - ).OnlyEnforceIf(present) - model.Add( - site_offset == offsets[candidate.output_bucket_ids[0]] - ).OnlyEnforceIf(present) - model.Add(site_count == candidate.topology_count).OnlyEnforceIf(present) - output_offsets = [ - offsets[bucket_id] for bucket_id in candidate.output_bucket_ids - ] - for output_offset in output_offsets[1:]: - model.Add(output_offset == output_offsets[0]).OnlyEnforceIf(present) - for dependency in ( - item - for item in problem.dependencies - if item.parent_candidate_id == candidate_id - ): - input_offset = offsets[dependency.child_bucket_id] - output_offset = output_offsets[0] - if dependency.placement_relation == "SAME_INTERVAL": - model.Add(input_offset == output_offset).OnlyEnforceIf(present) - elif dependency.placement_relation == "CONTAINED": - input_count = _tensor_mesh_count( - problem.types[problem.buckets[dependency.child_bucket_id].type_id] - ) - model.Add(input_offset <= output_offset).OnlyEnforceIf(present) - model.Add( - output_offset + candidate.topology_count - <= input_offset + input_count - ).OnlyEnforceIf(present) - time_intervals.append( - model.NewIntervalVar( - site_start, site_duration, site_end, f"site_execution_{site_id}" - ) - ) - topology_intervals.append( - model.NewIntervalVar( - site_offset, site_count, site_offset_end, f"site_topology_{site_id}" - ) - ) - for candidate_id, candidate in problem.candidates.items(): - if ( - candidate_id not in starts - or _is_reshard(candidate) - or candidate_id in merged_geometry_candidates - ): - continue - output_offsets = [offsets[bucket_id] for bucket_id in candidate.output_bucket_ids] - for output_offset in output_offsets[1:]: - model.Add(output_offset == output_offsets[0]).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - if _is_view(candidate) and candidate.input_bucket_ids: - for output_offset in output_offsets: - model.Add( - output_offset == offsets[candidate.input_bucket_ids[0]] - ).OnlyEnforceIf(pick_candidates[candidate_id]) - if candidate.input_bucket_ids: - for dependency in ( - item - for item in problem.dependencies - if item.parent_candidate_id == candidate_id - ): - input_offset = offsets[dependency.child_bucket_id] - output_offset = output_offsets[0] - if dependency.placement_relation == "SAME_INTERVAL": - model.Add(input_offset == output_offset).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - elif dependency.placement_relation == "CONTAINED": - input_count = _tensor_mesh_count( - problem.types[problem.buckets[dependency.child_bucket_id].type_id] - ) - model.Add(input_offset <= output_offset).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - model.Add( - output_offset + candidate.topology_count - <= input_offset + input_count - ).OnlyEnforceIf(pick_candidates[candidate_id]) - topology_intervals.append( - model.NewOptionalIntervalVar( - output_offsets[0], - candidate.topology_count, - output_offsets[0] + candidate.topology_count, - pick_candidates[candidate_id], - f"topology_{candidate_id}", - ) - ) - time_intervals.append(positive_intervals[candidate_id]) - if time_intervals: - model.AddNoOverlap2D(time_intervals, topology_intervals) - - for region_id, region in problem.regions.items(): - parent_region = region.parent_region_id - direct_candidates = [ - candidate_id - for candidate_id, candidate_region in problem.candidate_enclosing_regions.items() - if candidate_region == parent_region and candidate_id in starts - ] - for candidate_id in direct_candidates: - before = model.NewBoolVar(f"candidate_{candidate_id}_before_region_{region_id}") - present = pick_candidates[candidate_id] - model.Add(starts[candidate_id] >= region_ends[region_id]).OnlyEnforceIf( - [present, before.Not()] - ) - model.Add(ends[candidate_id] <= region_starts[region_id]).OnlyEnforceIf( - [present, before] - ) - sibling_regions = [ - other_id - for other_id, other in problem.regions.items() - if other.parent_region_id == parent_region and other_id != region_id - ] - for other_id in sibling_regions: - if other_id < region_id: - continue - before = model.NewBoolVar(f"region_{region_id}_before_{other_id}") - model.Add(region_ends[region_id] <= region_starts[other_id]).OnlyEnforceIf( - before - ) - model.Add(region_ends[other_id] <= region_starts[region_id]).OnlyEnforceIf( - before.Not() - ) - - bandwidth_intervals: list[cp_model.IntervalVar] = [] - bandwidth_demands: list[int] = [] - - def add_bandwidth_group(candidate_ids: list[int], demand: int, label: str) -> None: - literals = [pick_candidates[candidate_id] for candidate_id in candidate_ids] - if len(literals) == 1: - active = literals[0] - else: - active = model.NewBoolVar(f"bandwidth_active_{label}") - for literal in literals: - model.AddImplication(literal, active) - model.AddBoolOr([active.Not(), *literals]) - start = model.NewIntVar(0, horizon, f"bandwidth_start_{label}") - end = model.NewIntVar(0, horizon, f"bandwidth_end_{label}") - duration = model.NewIntVar(0, horizon, f"bandwidth_duration_{label}") - for candidate_id in candidate_ids: - model.Add(start == starts[candidate_id]).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - model.Add(end == ends[candidate_id]).OnlyEnforceIf( - pick_candidates[candidate_id] - ) - model.Add( - duration == ends[candidate_id] - starts[candidate_id] - ).OnlyEnforceIf(pick_candidates[candidate_id]) - bandwidth_intervals.append( - model.NewOptionalIntervalVar( - start, duration, end, active, f"bandwidth_{label}" - ) - ) - bandwidth_demands.append(demand) - - for site_id in problem.site_order: - groups: dict[int, list[int]] = {} - for candidate_id in problem.authored_candidates[site_id]: - candidate = problem.candidates[candidate_id] - if candidate_id not in starts or candidate.hbm_demand_bytes_per_ns <= 0: - continue - demand = ( - bandwidth_per_ns - if _is_reshard(candidate) - else candidate.hbm_demand_bytes_per_ns - ) - groups.setdefault(demand, []).append(candidate_id) - for demand, candidate_ids in sorted(groups.items()): - add_bandwidth_group( - candidate_ids, demand, f"site_{site_id}_demand_{demand}" - ) - - for candidate_id, candidate in problem.candidates.items(): - if candidate.site_id is not None or candidate_id not in starts: - continue - if candidate.hbm_demand_bytes_per_ns <= 0: - continue - add_bandwidth_group( - [candidate_id], bandwidth_per_ns, f"candidate_{candidate_id}" - ) - if bandwidth_intervals: - model.AddCumulative(bandwidth_intervals, bandwidth_demands, bandwidth_per_ns) - - _add_capacity_resource( - problem, model, pick_buckets, starts, ends, makespan, horizon - ) - return _CpModelState( - model=model, - pick_candidates=pick_candidates, - pick_buckets=pick_buckets, - terminal_buckets=terminal_buckets, - starts=starts, - ends=ends, - ready=ready, - offsets=offsets, - makespan=makespan, - horizon_ns=horizon, - ) - - -def _add_capacity_resource( - problem: PartitionProblem, - model: cp_model.CpModel, - pick_buckets: dict[int, cp_model.IntVar], - starts: dict[int, cp_model.IntVar], - ends: dict[int, cp_model.IntVar], - makespan: cp_model.IntVar, - horizon: int, -) -> None: - """Charge each allocation group its widest resident type for its lifetime.""" - intervals: list[cp_model.IntervalVar] = [] - demands: list[int] = [] - for group_id, bucket_ids in _allocation_groups(problem): - selected = [pick_buckets[bucket_id] for bucket_id in bucket_ids] - active = model.NewBoolVar(f"allocation_active_{group_id}") - for literal in selected: - model.AddImplication(literal, active) - model.AddBoolOr([active.Not(), *selected]) - start_terms: list[cp_model.IntVar] = [] - end_terms: list[cp_model.IntVar] = [] - for bucket_id in bucket_ids: - value_id = problem.buckets[bucket_id].value_id - if problem.values[value_id].producer_site_id is None: - start_terms.append( - _constant_or_selected( - model, - 0, - pick_buckets[bucket_id], - horizon, - f"source_start_{bucket_id}", - default=horizon, - ) - ) - if problem.values[value_id].is_const: - start_terms.append( - _constant_or_selected( - model, - 0, - pick_buckets[bucket_id], - horizon, - f"constant_start_{bucket_id}", - default=horizon, - ) - ) - for candidate_id, candidate in problem.candidates.items(): - if bucket_id in candidate.output_bucket_ids and candidate_id in starts: - start_terms.append(starts[candidate_id]) - if bucket_id in candidate.input_bucket_ids and candidate_id in ends: - end_terms.append(ends[candidate_id]) - if bucket_id in candidate.output_bucket_ids and candidate_id in ends: - end_terms.append(ends[candidate_id]) - if problem.values[value_id].is_final_output: - final_end = model.NewIntVar(0, horizon, f"final_output_end_{bucket_id}") - model.Add(final_end == makespan).OnlyEnforceIf(pick_buckets[bucket_id]) - model.Add(final_end == 0).OnlyEnforceIf(pick_buckets[bucket_id].Not()) - end_terms.append(final_end) - if problem.values[value_id].is_const: - constant_end = model.NewIntVar(0, horizon, f"constant_end_{bucket_id}") - model.Add(constant_end == makespan).OnlyEnforceIf( - pick_buckets[bucket_id] - ) - model.Add(constant_end == 0).OnlyEnforceIf( - pick_buckets[bucket_id].Not() - ) - end_terms.append(constant_end) - if not start_terms: - start_terms.append(model.NewConstant(0)) - if not end_terms: - end_terms.append(model.NewConstant(0)) - minimum_start = model.NewIntVar(0, horizon, f"allocation_min_start_{group_id}") - maximum_end = model.NewIntVar(0, horizon, f"allocation_max_end_{group_id}") - allocation_size = model.NewIntVar(0, horizon, f"allocation_size_{group_id}") - model.AddMinEquality(minimum_start, start_terms) - model.AddMaxEquality(maximum_end, end_terms) - model.Add(allocation_size == maximum_end - minimum_start).OnlyEnforceIf(active) - model.Add(allocation_size == 0).OnlyEnforceIf(active.Not()) - intervals.append( - model.NewOptionalIntervalVar( - minimum_start, - allocation_size, - maximum_end, - active, - f"allocation_{group_id}", - ) - ) - type_values = [ - problem.types[problem.buckets[bucket_id].type_id] for bucket_id in bucket_ids - ] - byte_counts = [ - tensor_bytes(type) for type in type_values if isinstance(type, TensorType) - ] - demands.append(max(byte_counts, default=0)) - if intervals: - model.AddCumulative(intervals, demands, problem.facts.memory_capacity_bytes) - - -def _constant_or_selected( - model: cp_model.CpModel, - constant: int, - present: cp_model.IntVar, - horizon: int, - name: str, - *, - default: int = 0, -) -> cp_model.IntVar: - value = model.NewIntVar(0, horizon, name) - model.Add(value == constant).OnlyEnforceIf(present) - model.Add(value == default).OnlyEnforceIf(present.Not()) - return value - - -def _decode( - problem: PartitionProblem, - state: _CpModelState, - solver: cp_model.CpSolver, - status: int, -) -> PartitionSolution: - selected_candidates = tuple( - candidate_id - for candidate_id in sorted(problem.candidates) - if solver.Value(state.pick_candidates[candidate_id]) - ) - selected_buckets = tuple( - bucket_id - for bucket_id in sorted(problem.buckets) - if solver.Value(state.pick_buckets[bucket_id]) - ) - intervals = tuple( - ( - candidate_id, - ExecutionInterval( - solver.Value(state.starts[candidate_id]), - solver.Value(state.ends[candidate_id]), - ), - ) - for candidate_id in selected_candidates - if candidate_id in state.starts - ) - offsets = tuple( - (bucket_id, solver.Value(state.offsets[bucket_id])) - for bucket_id in selected_buckets - ) - makespan = solver.Value(state.makespan) - if status == cp_model.OPTIMAL: - return PartitionSolution( - "OPTIMAL", - selected_candidates, - selected_buckets, - intervals, - offsets, - makespan, - makespan, - 0.0, - ) - best_bound = math.floor(solver.BestObjectiveBound()) - best_bound = max(0, min(best_bound, makespan)) - gap = (makespan - best_bound) / max(makespan, 1) - return PartitionSolution( - "FEASIBLE_NOT_PROVEN", - selected_candidates, - selected_buckets, - intervals, - offsets, - makespan, - best_bound, - gap, - ) - - -def _write_failure( - options: ScheduleOptions, problem: PartitionProblem, error: Exception -) -> None: - if options.debug_dump_dir is None: - return - options.debug_dump_dir.mkdir(parents=True, exist_ok=True) - payload = { - "error": str(error), - "root": problem.root.name, - "status": type(error).__name__, - "target": problem.facts.spec.device_id, - } - (options.debug_dump_dir / "solve_failure.json").write_text( - json.dumps(payload, indent=2, sort_keys=True) + "\n" - ) - - -def solve_partition_problem( - problem: PartitionProblem, options: ScheduleOptions -) -> PartitionSolution: - """Build and solve one makespan model over an already-closed problem.""" - try: - state = _build_model(problem) - state.model.Minimize(state.makespan) - solver = cp_model.CpSolver() - solver.parameters.max_time_in_seconds = options.timeout_seconds - solver.parameters.num_search_workers = options.workers - solver.parameters.random_seed = options.random_seed - solver.parameters.stop_after_first_solution = options.stop_at_first_solution - status = solver.Solve(state.model) - if status == cp_model.INFEASIBLE: - raise PartitionSolveError( - f"no feasible partition for root {problem.root.name!r} at " - f"topology {problem.topology.name!r} on {problem.facts.spec.device_id}" - ) - if status == cp_model.MODEL_INVALID: - raise PartitionSolveError("the solver reported an invalid partition model") - if status == cp_model.UNKNOWN: - raise PartitionSolveError( - f"the solver returned no incumbent for root {problem.root.name!r} " - "within its time limit" - ) - if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE): - raise PartitionSolveError( - f"unexpected solver status {solver.StatusName(status)}" - ) - return _decode(problem, state, solver, status) - except (ValueError, RuntimeError) as error: - _write_failure(options, problem, error) - raise - - -__all__ = [ - "ExecutionInterval", - "PartitionSolution", - "PartitionSolveError", - "solve_partition_problem", -] diff --git a/src/tilefoundry/schedule/pipeline/__init__.py b/src/tilefoundry/schedule/pipeline/__init__.py deleted file mode 100644 index 1df2f2a6..00000000 --- a/src/tilefoundry/schedule/pipeline/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -"""The private typed pipeline used by the CUDA CTA schedule algorithm.""" - -from .facts import PipelineFacts, PipelineFactsQuery -from .plan import PipelineSchedulePlan, export_pipeline_plan -from .problem import PipelineProblem, PipelineProblemError, build_pipeline_problem -from .program import PipelineProgram, build_pipeline_program -from .solve import PipelineSolution, solve_pipeline_problem - -__all__ = [ - "PipelineFacts", - "PipelineFactsQuery", - "PipelineProblem", - "PipelineProblemError", - "PipelineProgram", - "PipelineSolution", - "PipelineSchedulePlan", - "build_pipeline_problem", - "build_pipeline_program", - "solve_pipeline_problem", - "export_pipeline_plan", -] diff --git a/src/tilefoundry/schedule/pipeline/facts.py b/src/tilefoundry/schedule/pipeline/facts.py deleted file mode 100644 index 5a4457a8..00000000 --- a/src/tilefoundry/schedule/pipeline/facts.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Typed target facts consumed by the private pipeline scheduler. - -The level a pipeline is asked about and the level whose resources bound it are -two different things, and this record keeps them apart. A CUDA pipeline is asked -about `thread`, because what it decides is how the warps of one CTA overlap; the -tile they cooperate on is shared memory, which is a CTA-scoped resource. Stating -one name for both would claim a per-thread capacity that no hardware publishes. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -from tilefoundry.schedule.facts import AtomFact - - -@dataclass(frozen=True) -class PipelineFactsQuery: - """The target-independent program facts needed for one projection.""" - - topology: str - statements: tuple[tuple[str, object], ...] - - -@dataclass(frozen=True) -class PipelineInstructionFacts: - """The supported instruction choices for one stable statement ID.""" - - statement_id: str - candidates: tuple[AtomFact, ...] - - -@dataclass(frozen=True) -class PipelineFacts: - """All concrete information required to close one pipeline problem. - - `topology` is the level that was asked about. `tile_capacity_scope` names the - level the capacity below it belongs to, which may be a coarser one: the - cooperating threads of one such unit share that store between them. - """ - - topology: str - tile_capacity_scope: str - tile_capacity_bytes: int - max_threads_per_warp: int - instructions: tuple[PipelineInstructionFacts, ...] - - -__all__ = ["PipelineFacts", "PipelineFactsQuery", "PipelineInstructionFacts"] diff --git a/src/tilefoundry/schedule/pipeline/plan.py b/src/tilefoundry/schedule/pipeline/plan.py deleted file mode 100644 index fc298c8f..00000000 --- a/src/tilefoundry/schedule/pipeline/plan.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Typed export of a solved pipeline schedule.""" - -from __future__ import annotations - -import dataclasses -import json -from dataclasses import dataclass - -from tilefoundry.schedule.plan import ( - PlanVerificationError, - SchedulePlan, - TargetSpecRef, -) - -from .program import PipelineProgram -from .solve import PipelineSolution - - -@dataclass(frozen=True) -class ScheduledStatement: - """One selected instruction and its half-open execution interval. - - `footprint_bytes` already counts the rings its buffers were given, and - `fits_capacity` says whether the level's tile store holds that. A statement - that does not fit is still scheduled and still reported, because the plan's - job is to say what this program costs on this machine, not to hide the - parts of it that are expensive. - """ - - id: str - instruction: str - tile: tuple[int, ...] - resources: tuple[tuple[str, int], ...] - start: int - end: int - footprint_bytes: int - fits_capacity: bool - - -@dataclass(frozen=True) -class ScheduledBuffer: - """One named storage object and its dependency-safe ring allocation.""" - - id: str - storage: str - ring_depth: int - producer_ids: tuple[str, ...] - consumer_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class KernelHole: - """One stable statement reference with serialized input/output relations.""" - - statement_id: str - inputs: tuple[str, ...] - outputs: tuple[str, ...] - relations: tuple[str, ...] - - -@dataclass(frozen=True) -class PipelineSchedulePlan(SchedulePlan): - """The complete, deterministic exported result of a CTA pipeline solve.""" - - target: TargetSpecRef - scaffold: str - statements: tuple[ScheduledStatement, ...] - buffers: tuple[ScheduledBuffer, ...] - holes: tuple[KernelHole, ...] - - def verify(self, module, function, topology) -> None: - statement_ids = {item.id for item in self.statements} - if len(statement_ids) != len(self.statements): - raise PlanVerificationError("pipeline plan has duplicate statement IDs") - for item in self.statements: - if item.start < 0 or item.end <= item.start: - raise PlanVerificationError(f"statement {item.id!r} has invalid interval") - for buffer in self.buffers: - if buffer.ring_depth < 1: - raise PlanVerificationError(f"buffer {buffer.id!r} has invalid ring depth") - unknown = set(buffer.producer_ids + buffer.consumer_ids) - statement_ids - if unknown: - raise PlanVerificationError( - f"buffer {buffer.id!r} references unknown statements {sorted(unknown)!r}" - ) - for hole in self.holes: - if hole.statement_id not in statement_ids: - raise PlanVerificationError( - f"hole references unknown statement {hole.statement_id!r}" - ) - if not all(isinstance(item, str) for item in hole.inputs + hole.outputs + hole.relations): - raise PlanVerificationError(f"hole {hole.statement_id!r} has malformed relations") - - def to_json(self) -> str: - return json.dumps(dataclasses.asdict(self), sort_keys=True) - - def render(self) -> str: - lines = ["pipeline schedule"] - lines.extend( - f"{item.id}: {item.instruction} [{item.start}, {item.end})" - for item in self.statements - ) - return "\n".join(lines) - - -def export_pipeline_plan( - program: PipelineProgram, solution: PipelineSolution, target: object -) -> PipelineSchedulePlan: - """Export stable values only; no target or solver value escapes the plan.""" - from tilefoundry.schedule.render import emit_scaffold # noqa: PLC0415 - - by_id = {item.id: item for item in solution.statements} - if tuple(by_id) != tuple(unit.name for unit in program.units): - raise PlanVerificationError("solution statements do not match the pipeline program") - ring = {item.id: item.ring_depth for item in solution.buffers} - skeleton, _swimlane, contracts = emit_scaffold(program.graph, program.tree, ring) - statements = tuple( - ScheduledStatement( - id=unit.name, - instruction=by_id[unit.name].instruction.atom.op.name, - tile=tuple( - int(domain.dim_max_val(axis).num_si()) - int(domain.dim_min_val(axis).num_si()) + 1 - for domain in _domains(program, unit.name) - for axis in range(domain.dim(1)) - ), - resources=by_id[unit.name].resources, - start=by_id[unit.name].start, - end=by_id[unit.name].end, - footprint_bytes=by_id[unit.name].footprint_bytes, - fits_capacity=by_id[unit.name].fits_capacity, - ) - for unit in program.units - ) - buffers = tuple( - ScheduledBuffer( - id=item.id, - storage="smem", - ring_depth=item.ring_depth, - producer_ids=item.producer_ids, - consumer_ids=item.consumer_ids, - ) - for item in solution.buffers - ) - holes = tuple( - KernelHole( - statement_id=contract.name.removeprefix("HOLE_"), - inputs=tuple(item.tensor_name for item in contract.inputs), - outputs=(contract.output.tensor_name,), - relations=tuple(str(item.index_map) for item in contract.inputs) - + (str(contract.output.index_map),), - ) - for contract in contracts - ) - return PipelineSchedulePlan( - target=TargetSpecRef.of(target), - scaffold=skeleton.text, - statements=statements, - buffers=buffers, - holes=holes, - ) - - -def _domains(program: PipelineProgram, name: str) -> tuple[object, ...]: - """Return the one ISL domain piece named by a stable statement ID.""" - domains: list[object] = [] - program.graph.domain.foreach_set( - lambda domain: domains.append(domain) if domain.get_tuple_name() == name else None - ) - if len(domains) != 1: - raise PlanVerificationError(f"statement {name!r} has {len(domains)} domain pieces") - return tuple(domains) - - -__all__ = [ - "KernelHole", - "PipelineSchedulePlan", - "ScheduledBuffer", - "ScheduledStatement", - "TargetSpecRef", - "export_pipeline_plan", -] diff --git a/src/tilefoundry/schedule/pipeline/problem.py b/src/tilefoundry/schedule/pipeline/problem.py deleted file mode 100644 index 9209c3be..00000000 --- a/src/tilefoundry/schedule/pipeline/problem.py +++ /dev/null @@ -1,255 +0,0 @@ -"""The closed, target-free constraint input for pipeline scheduling. - -Everything here is measured off the program's own schedule tree. Two of those -measurements are what make an intra-CTA pipeline a pipeline rather than a -sequence: the dependence distance a buffer carries, which is how many tiles of -it have to be alive at once, and the bytes each statement holds, which is what -the level's tile store has to fit. Neither is a decision, so both are closed -into the problem and neither is left for the solver to guess. -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass - -import isl - -from tilefoundry.analysis.poly import ( - AccessFootprint, - access_footprints, - carried_distances, -) -from tilefoundry.ir.types.shard import Topology -from tilefoundry.schedule.kernel_schedule import band_statement, schedule_bands - -from ..errors import ScheduleError -from .facts import PipelineFacts -from .program import PipelineProgram - - -class PipelineProblemError(ScheduleError): - """The program and projected facts cannot form a finite schedule problem. - - A scheduling failure, and reachable as one: a caller asking this - layer to schedule something catches what the layer raises, and a - capability that cannot be scheduled is recorded against that. Sitting - outside `ScheduleError` made a limit of this algorithm unstateable - except as a bare `ValueError`, which is also what a caller passing - nonsense gets -- so the two could not be told apart. - """ - - -@dataclass(frozen=True) -class PipelineStatementProblem: - """One statement's explicit legal instruction and resource choices. - - `footprint_bytes` is what one instance of this statement holds in each - buffer it touches, before any ring multiplies it. A dimension is counted at - the widest extent any of this statement's accesses reaches there, so two - accesses to one buffer are counted once. - """ - - id: str - extents: tuple[int, ...] - candidates: tuple[object, ...] - resources: tuple[tuple[str, int], ...] - footprint_bytes: tuple[tuple[str, int], ...] - - -@dataclass(frozen=True) -class PipelineBufferProblem: - """One storage object, its users, and measured carried distances. - - `carried_distances` is per holding statement, because a distance is only - meaningful against the extents of the statement whose band reported it: the - same buffer carried two iterations spans a different number of tiles under - a statement tiled 2 wide than under one tiled 64 wide. - """ - - id: str - producer_ids: tuple[str, ...] - consumer_ids: tuple[str, ...] - carried_distances: tuple[tuple[str, tuple[int, ...]], ...] - - -@dataclass(frozen=True) -class PipelineProblem: - """A complete finite problem with no Target object or callback.""" - - topology: str - capacity_bytes: int - statements: tuple[PipelineStatementProblem, ...] - buffers: tuple[PipelineBufferProblem, ...] - - -def build_pipeline_problem( - program: PipelineProgram, facts: PipelineFacts, topology: Topology -) -> PipelineProblem: - """Close the problem from immutable analysis and already-projected facts.""" - if facts.topology != topology.name: - raise PipelineProblemError( - f"pipeline facts describe {facts.topology!r}, not topology " - f"{topology.name!r}" - ) - if facts.tile_capacity_bytes <= 0: - raise PipelineProblemError("pipeline facts require a positive tile capacity") - candidates = {item.statement_id: item.candidates for item in facts.instructions} - expected = tuple(unit.name for unit in program.units) - if tuple(candidates) != expected: - raise PipelineProblemError( - f"pipeline facts statements {tuple(candidates)!r} do not match program {expected!r}" - ) - statements: list[PipelineStatementProblem] = [] - time_maps = program.tree.get_map() - accesses = access_footprints(program.graph, time_maps) - bands = _bands_by_statement(program) - held = _held_by_statement(accesses) - distances: dict[str, dict[str, tuple[int, ...]]] = {} - domains: dict[str, object] = {} - program.graph.domain.foreach_set( - lambda domain: domains.__setitem__(domain.get_tuple_name(), domain) - ) - for unit in program.units: - domain = domains.get(unit.name) - if domain is None: - raise PipelineProblemError(f"statement {unit.name!r} has no domain") - extents = tuple( - int(domain.dim_max_val(axis).num_si()) - - int(domain.dim_min_val(axis).num_si()) - + 1 - for axis in range(domain.dim(isl.dim_type.SET)) - ) - if not extents: - raise PipelineProblemError(f"statement {unit.name!r} has no finite extent") - choices = candidates[unit.name] - if not choices: - raise PipelineProblemError( - f"statement {unit.name!r} has no supported instruction candidates" - ) - resources = tuple( - sorted( - (name, value) - for fact in choices - for name, value in fact.resource.items() - if value > 0 - ) - ) - distances[unit.name] = _carried_for(program, unit.name, bands, extents) - statements.append( - PipelineStatementProblem( - unit.name, - extents, - choices, - resources, - footprint_bytes=tuple( - (buffer, _occupancy_bytes(occupancy)) - for buffer, occupancy in sorted(held.get(unit.name, {}).items()) - ), - ) - ) - writers: dict[str, list[str]] = {} - readers: dict[str, list[str]] = {} - for footprint in accesses: - if footprint.is_read: - readers.setdefault(footprint.buffer, []).append(footprint.statement) - maps: list[object] = [] - program.graph.writes.foreach_map(maps.append) - for mapping in maps: - writers.setdefault( - mapping.get_tuple_name(isl.dim_type.OUT), [] - ).append(mapping.get_tuple_name(isl.dim_type.IN)) - buffers = tuple( - PipelineBufferProblem( - id=name, - producer_ids=tuple(sorted(set(writers.get(name, ()))),), - consumer_ids=tuple(sorted(set(readers.get(name, ()))),), - carried_distances=tuple( - (statement, per_buffer[name]) - for statement, per_buffer in sorted(distances.items()) - if name in per_buffer - ), - ) - for name in sorted(set(writers) | set(readers)) - ) - return PipelineProblem(topology.name, facts.tile_capacity_bytes, tuple(statements), buffers) - - -def _bands_by_statement(program: PipelineProgram) -> dict[str, object]: - """One band per statement, keyed by the statement it schedules.""" - bands = {band_statement(band): band for band in schedule_bands(program.tree)} - missing = sorted(unit.name for unit in program.units if unit.name not in bands) - if missing: - raise PipelineProblemError( - f"statements {missing!r} have no band in the schedule tree" - ) - return bands - - -def _carried_for( - program: PipelineProgram, - name: str, - bands: dict[str, object], - extents: tuple[int, ...], -) -> dict[str, tuple[int, ...]]: - """Per buffer, the distances statement *name*'s own band carries. - - Only the buffers that carry something are kept: a buffer with no carried - dependence needs one slot, and saying so with a tuple of zeros would put - every buffer in the program into every statement's record. - """ - time_map = bands[name].get_partial_schedule_union_map() - return { - buffer: carried - for buffer, carried in carried_distances( - program.graph, time_map, len(extents) - ).items() - if any(carried) - } - - -def _held_by_statement( - footprints: tuple[AccessFootprint, ...], -) -> dict[str, dict[str, tuple[tuple[int, ...], int]]]: - """Held by statement. - - Group *footprints* by statement then buffer, widening each buffer - dimension to every extent any access in that group needs there. - """ - dims: dict[tuple[str, str], list[set[int]]] = {} - elem_bytes: dict[tuple[str, str], int] = {} - for footprint in footprints: - group = (footprint.statement, footprint.buffer) - if group not in dims: - dims[group] = [set() for _ in footprint.dims] - elem_bytes[group] = footprint.elem_bytes - if len(dims[group]) != len(footprint.dims): - raise PipelineProblemError( - f"buffer {footprint.buffer!r} is accessed with " - f"{len(footprint.dims)} and {len(dims[group])} dimension(s)" - ) - for position, extent in enumerate(footprint.dims): - dims[group][position].add(extent.extent) - held: dict[str, dict[str, tuple[tuple[int, ...], int]]] = {} - for (statement, buffer), per_dim in dims.items(): - widest = tuple(max(options) for options in per_dim) - held.setdefault(statement, {})[buffer] = ( - widest, - elem_bytes[(statement, buffer)], - ) - return held - - -def _occupancy_bytes(occupancy: tuple[tuple[int, ...], int]) -> int: - """Bytes one instance of a statement holds in one buffer.""" - widest, elem_bytes = occupancy - return math.prod(widest) * elem_bytes - - -__all__ = [ - "PipelineBufferProblem", - "PipelineProblem", - "PipelineProblemError", - "PipelineStatementProblem", - "build_pipeline_problem", -] diff --git a/src/tilefoundry/schedule/pipeline/program.py b/src/tilefoundry/schedule/pipeline/program.py deleted file mode 100644 index 9955103e..00000000 --- a/src/tilefoundry/schedule/pipeline/program.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Target-independent construction of one immutable pipeline program view.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import isl - -from tilefoundry.analysis.poly import TileGraph, TileUnit, extract -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.schedule.kernel_schedule import build_schedule_tree - -from .facts import PipelineFactsQuery - - -@dataclass(frozen=True) -class PipelineProgram: - """Analysis facts plus a private ISL tree, without mutating TileGraph.""" - - graph: TileGraph - tree: "isl.schedule" - units: tuple[TileUnit, ...] - - def facts_query(self, topology: str) -> PipelineFactsQuery: - """Return the explicit query used to project target facts once.""" - return PipelineFactsQuery( - topology=topology, - statements=tuple((unit.name, unit.op) for unit in self.units), - ) - - -def build_pipeline_program(module: Module, function: Function) -> PipelineProgram: - """Extract and tree one function without attaching schedule state to analysis.""" - if not module.owns(function, derived=True): - raise ValueError( - f"pipeline program: {function.name!r} is not owned by module {module.name!r}" - ) - graph = extract(function) - return PipelineProgram(graph=graph, tree=build_schedule_tree(graph), units=graph.units) - - -__all__ = ["PipelineProgram", "build_pipeline_program"] diff --git a/src/tilefoundry/schedule/pipeline/solve.py b/src/tilefoundry/schedule/pipeline/solve.py deleted file mode 100644 index de7ddded..00000000 --- a/src/tilefoundry/schedule/pipeline/solve.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Deterministic closed-problem scheduling decisions.""" - -from __future__ import annotations - -from dataclasses import dataclass - -from .problem import PipelineBufferProblem, PipelineProblem - - -class PipelineSolveError(ValueError): - """A closed pipeline problem has no legal finite solution.""" - - -@dataclass(frozen=True) -class PipelineStatementSolution: - """One chosen instruction and its half-open interval. - - `footprint_bytes` counts the rings this statement's buffers were given, so - it is what the statement really occupies once the pipeline is deep enough - to run. `fits_capacity` records that against the level's tile store rather - than enforcing it: a statement too wide for the store still has a - schedule, only a worse one, and a solver that silently dropped it would - report a plan for a program nobody asked for. - """ - - id: str - instruction: object - start: int - end: int - resources: tuple[tuple[str, int], ...] - footprint_bytes: int - fits_capacity: bool - - -@dataclass(frozen=True) -class PipelineBufferSolution: - """One solved ring depth.""" - - id: str - ring_depth: int - producer_ids: tuple[str, ...] - consumer_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class PipelineSolution: - """All typed decisions exported by the pipeline plan.""" - - statements: tuple[PipelineStatementSolution, ...] - buffers: tuple[PipelineBufferSolution, ...] - - -def solve_pipeline_problem(problem: PipelineProblem) -> PipelineSolution: - """Select the unique lowest-cost legal choice and sequence resources. - - Choice is a deterministic optimization over the complete candidate set, not - a catalogue-order fallback. The problem has no target reference at this - point; all legality and capacity data were frozen in PipelineFacts. - """ - extents = {statement.id: statement.extents for statement in problem.statements} - ring = { - buffer.id: _ring_depth(buffer, extents) for buffer in problem.buffers - } - clock = 0 - statements: list[PipelineStatementSolution] = [] - for statement in problem.statements: - ranked = sorted( - statement.candidates, - key=lambda fact: (fact.duration, getattr(getattr(fact, "atom", None), "op", object()).name), - ) - if not ranked: - raise PipelineSolveError(f"statement {statement.id!r} has no legal candidate") - chosen = ranked[0] - duration = max(1, round(chosen.duration * 1000)) - footprint = sum( - held * ring.get(buffer, 1) for buffer, held in statement.footprint_bytes - ) - statements.append( - PipelineStatementSolution( - statement.id, - chosen, - clock, - clock + duration, - statement.resources, - footprint_bytes=footprint, - fits_capacity=footprint <= problem.capacity_bytes, - ) - ) - clock += duration - buffers = tuple( - PipelineBufferSolution( - item.id, ring[item.id], item.producer_ids, item.consumer_ids - ) - for item in problem.buffers - ) - return PipelineSolution(tuple(statements), buffers) - - -def _ring_depth( - buffer: PipelineBufferProblem, extents: dict[str, tuple[int, ...]] -) -> int: - """One buffer's ring depth, measured rather than searched for. - - A dependence carried `distance` iterations along a dimension tiled `tile` - wide spans `ceil(distance / tile)` tiles, and the ring holds one slot more - than that so the older tile is still alive while the newer one fills. A - buffer that carries nothing needs the single slot every buffer needs. - """ - depths = [1] - for statement, carried in buffer.carried_distances: - tile = extents.get(statement) - if tile is None: - raise PipelineSolveError( - f"buffer {buffer.id!r} carries a distance for unknown statement " - f"{statement!r}" - ) - if len(tile) != len(carried): - raise PipelineSolveError( - f"buffer {buffer.id!r} carries {len(carried)} distance(s) for " - f"statement {statement!r}, which spans {len(tile)} dimension(s)" - ) - depths.extend( - -(-distance // width) + 1 - for distance, width in zip(carried, tile) - if distance - ) - return max(depths) - - -__all__ = [ - "PipelineBufferSolution", - "PipelineSolution", - "PipelineSolveError", - "PipelineStatementSolution", - "solve_pipeline_problem", -] diff --git a/src/tilefoundry/schedule/plan.py b/src/tilefoundry/schedule/plan.py deleted file mode 100644 index 042a72c8..00000000 --- a/src/tilefoundry/schedule/plan.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Define the common behavior of algorithm-owned schedule plans. - -A plan is the typed result of one solve. Each subtype owns its decisions, JSON, -and human rendering because algorithms at different levels have no shared -decision schema. Plans are produced and consumed in the solving process. -""" - -from __future__ import annotations - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class TargetSpecRef: - """Stable identity of the installed target facts one plan relies on. - - Which hardware documents a decision was made against is the same question - whatever was decided, so every plan states it the same way. A target that - was constructed directly rather than installed from documents has a name but - no digest, and says so with an empty one rather than a fabricated value. - """ - - architecture_id: str - architecture_digest: str - device_id: str - device_digest: str - - @classmethod - def of(cls, target: object) -> "TargetSpecRef": - """The identity *target* publishes for its installed documents.""" - architecture = target.architecture # type: ignore[attr-defined] - device = target.device # type: ignore[attr-defined] - return cls( - architecture_id=getattr(target, "architecture_id", None) or architecture.name, - architecture_digest=getattr(target, "architecture_digest", None) or "", - device_id=getattr(target, "device_id", None) or device.name, - device_digest=getattr(target, "device_digest", None) or "", - ) - - -class SchedulePlan: - """One solve's decisions, owned by the algorithm that made them. - - A subtype carries the typed decisions of its own algorithm, and owns the - whole of its JSON object and its human rendering. Nothing here constrains - what those decisions are. - """ - - def verify(self, module: "Module", function: "Function", topology: "Topology") -> None: - """Check that this plan is internally consistent against its inputs. - - This is a structural check, not a second solve: it confirms that the - plan refers only to things that exist and that its own references agree - with each other. It states nothing about whether the schedule is good. - """ - raise NotImplementedError - - def to_json(self) -> str: - """The whole plan as JSON, owned entirely by the subtype.""" - raise NotImplementedError - - def render(self) -> str: - """The whole plan as human-readable text, owned entirely by the subtype.""" - raise NotImplementedError - - -class PlanVerificationError(ValueError): - """A plan referred to something that does not exist, or contradicted itself.""" - - -__all__ = ["PlanVerificationError", "SchedulePlan", "TargetSpecRef"] diff --git a/src/tilefoundry/schedule/render.py b/src/tilefoundry/schedule/render.py deleted file mode 100644 index dd665b7b..00000000 --- a/src/tilefoundry/schedule/render.py +++ /dev/null @@ -1,420 +0,0 @@ -"""Render a schedule as a holed C-like skeleton and Mermaid swimlane. - -isl owns loop construction and indentation; user nodes become statement holes. -``BufferAccess`` carries a statement-local polyhedral footprint rather than a -TIR memory expression. Statement names and coordinates come from call -expressions because AST annotations do not survive traversal. -""" -from __future__ import annotations - -import itertools -import math -from dataclasses import dataclass - -import isl - -from tilefoundry.analysis.poly import TileGraph, TileUnit -from tilefoundry.ir.core import Var, binding_name - - -class EmitScaffoldError(RuntimeError): - """Represent EmitScaffoldError. - - A construct `emit_scaffold` does not (yet) support, or a ``tg`` - consistency precondition that did not hold -- always raised with - a specific, actionable message; V1 never silently guesses. - """ - - -@dataclass(frozen=True) -class _RenderProgram: - """Private renderer state kept separate from analysis output.""" - - graph: TileGraph - tree: "isl.schedule" - ring: dict[str, int] - - def __getattr__(self, name: str): - return getattr(self.graph, name) - - - - - - - -@dataclass(frozen=True) -class Skeleton: - """A holed, C-like loop-nest skeleton. - - A holed, C-like loop-nest skeleton: isl ``ast_build`` codegen (PoC - 11) with every naked statement call (``MM(c0, c1, c2);``) rendered as - a ``HOLE_(...)`` call instead (see ``_build_skeleton``). - ``holes`` names every hole in ``text``, in first-appearance order. - """ - - text: str - holes: tuple[str, ...] - - -@dataclass(frozen=True) -class Swimlane: - """A human-readable Mermaid ``gantt`` rendering of the schedule. - - A human-readable Mermaid ``gantt`` rendering of the schedule: one - section (swimlane) per statement, minimally unrolled (prologue + a - handful of steady-state iterations + epilogue, see - ``_illustrative_instances``) rather than the full iteration count. - """ - - text: str - - -@dataclass(frozen=True) -class BufferAccess: - """V1 fallback view of one buffer touched by one statement. - - V1 fallback view of one buffer touched by one statement (see the - module docstring's "BufferAccess reuse note" for why the existing TIR - ``TensorView`` is not reused here). ``index_map`` is the ``isl.map`` - (from ``TileGraph.reads``/``writes``) taking this statement's - coordinates to elements of buffer ``tensor_name``; ``dtype`` is - best-effort HIR dtype recovery (``None`` if it could not be resolved, - e.g. an unbound intermediate -- see ``_dtype_table``). - """ - - tensor_name: str - index_map: "isl.map" - dtype: object | None - - -@dataclass(frozen=True) -class HoleContract: - """One per statement: what a hole must compute. - - One per statement: what a hole must compute, as a pure function - ``(inputs, coords) -> output`` -- no side effects, no indexing/sync (the - skeleton already carries those). ``op_ref`` is ``TileUnit.op`` (the HIR - ``Call``) so M3 can fill the hole and diff it against the HIR - Evaluator's own op subgraph result; M2 itself never runs it. - """ - - name: str - op_ref: object - inputs: tuple[BufferAccess, ...] - output: BufferAccess - coords: tuple[str, ...] - - - - - - - -def _dtype_table(units: tuple[TileUnit, ...]) -> dict[str, object]: - """Recover HIR dtypes for named statement inputs and outputs. - - This reapplies the common ``Var.name`` / ``binding_name`` rule because the - graph does not expose its identity-name table. Collisions are last-write - wins and unnamed intermediates retain ``dtype=None``. - """ - table: dict[str, object] = {} - for unit in units: - call = unit.op - out_name = binding_name(call) - if out_name is not None: - table[out_name] = getattr(call.type, "dtype", None) - for arg in call.args: - name = arg.name if isinstance(arg, Var) else binding_name(arg) - if name is not None: - table[name] = getattr(arg.type, "dtype", None) - return table - - -def _by_buf(union_map: "isl.union_map", stmt_name: str) -> dict[str, "isl.map"]: - """By buf. - - Decompose ``union_map`` (``tg.reads`` or ``tg.writes``) into the - per-buffer maps whose ``IN`` tuple is ``stmt_name``, keyed by ``OUT`` - tuple (buffer) name. - """ - maps: list["isl.map"] = [] - union_map.foreach_map(maps.append) - return { - m.get_tuple_name(isl.dim_type.OUT): m - for m in maps - if m.get_tuple_name(isl.dim_type.IN) == stmt_name - } - - -def _ordered_inputs( - unit: TileUnit, read_by_buf: dict[str, "isl.map"], dtype_table: dict[str, object] -) -> tuple[BufferAccess, ...]: - """Order buffer reads to match the operation arguments. - - Argument reads retain source-call order. Other reads, including an output - buffer's read-modify-write access, follow in buffer-name order because isl - union-map iteration is not stable. - """ - ordered: list[str] = [] - used: set[str] = set() - for arg in unit.op.args: - name = arg.name if isinstance(arg, Var) else binding_name(arg) - if name in read_by_buf and name not in used: - ordered.append(name) - used.add(name) - for name in sorted(read_by_buf): - if name not in used: - ordered.append(name) - used.add(name) - return tuple( - BufferAccess(tensor_name=name, index_map=read_by_buf[name], dtype=dtype_table.get(name)) - for name in ordered - ) - - -def _output_view( - unit: TileUnit, write_by_buf: dict[str, "isl.map"], dtype_table: dict[str, object] -) -> BufferAccess: - if len(write_by_buf) != 1: - raise EmitScaffoldError( - f"emit_scaffold: statement {unit.name!r} writes {len(write_by_buf)} " - "buffers -- HoleContract.output is singular in V1 (a real " - "multi-output statement is future work, see poly.py's " - "n_outputs support in _registered_access)" - ) - ((name, m),) = write_by_buf.items() - return BufferAccess(tensor_name=name, index_map=m, dtype=dtype_table.get(name)) - - - - - - - -def _call_coords(expr: "isl.ast_expr") -> tuple[str, tuple[str, ...]]: - """Decode an isl ``ast_expr_op_call``. - - Argument zero is the statement id; remaining arguments are coordinates. - Each coordinate renders through isl rather than parsing text, preserving - nontrivial affine expressions. - """ - name = expr.op_arg(0).id().name() - coords = tuple(expr.op_arg(i).to_C_str() for i in range(1, expr.op_n_arg())) - return name, coords - - -def _ring_ref(buf_name: str, coords: tuple[str, ...], ring: dict) -> str: - """Render one buffer reference for a hole call. - - A graph whose ring depths were never decided leaves ``tg.ring`` empty, and - every reference is then the bare buffer name; a decided depth above one - indexes the buffer by its innermost coordinate mod that depth. - """ - n = ring.get(buf_name) - if not n or n <= 1: - return buf_name - - - return f"{buf_name}[({coords[-1]}) % {n}]" - - -def _render_hole_call( - hole_name: str, in_refs: tuple[str, ...], out_ref: str, coords: tuple[str, ...] -) -> str: - """``HOLE_(/*in*/ a, b, /*out*/ c, /*coords*/ c0, c1, c2);``. - - ``HOLE_(/*in*/ a, b, /*out*/ c, /*coords*/ c0, c1, c2);`` -- - the hole's *inputs* (all reads, including any RMW self-read on the - output buffer -- included honestly rather than silently dropped, see - ``_ordered_inputs``), its *output*, and the raw schedule coordinates - it is parametrised by, each behind its own comment marker. - """ - sections = [ - "/*in*/ " + ", ".join(in_refs), - "/*out*/ " + out_ref, - "/*coords*/ " + ", ".join(coords), - ] - return f"{hole_name}({', '.join(sections)});" - - -def _print_to_str(node: "isl.ast_node", options: "isl.ast_print_options") -> str: - """``node`` printed as C through ``options``.""" - printer = isl.printer.to_str().set_output_format(isl.format.C) - return node.print(printer, options).get_str() - - -def _build_skeleton( - tg: TileGraph, dtype_table: dict[str, object] -) -> tuple[Skeleton, dict[str, HoleContract]]: - """Isl ``ast_build`` codegen (PoC 11) over ``tg.tree``. - - Isl ``ast_build`` codegen (PoC 11) over ``tg.tree``, with an - ``at_each_domain`` hook (validated in ``m2_hook_probe.py``) that - records each statement's hole-call replacement text in visit order, - then splices those replacements into the final ``to_C_str()`` text. - Also builds each statement's ``HoleContract`` along the way (first - occurrence only -- one contract per *statement*, not per call site). - """ - if tg.tree is None: - raise EmitScaffoldError( - "emit_scaffold: tg.tree is None -- call build_schedule_tree(tg) " - "before emit_scaffold(tg)" - ) - units_by_name = {u.name: u for u in tg.units} - contracts: dict[str, HoleContract] = {} - - def contract_for(stmt_name: str, coords: tuple[str, ...]) -> HoleContract: - contract = contracts.get(stmt_name) - if contract is not None: - return contract - unit = units_by_name.get(stmt_name) - if unit is None: - raise EmitScaffoldError( - f"emit_scaffold: schedule tree statement {stmt_name!r} has no " - "matching TileUnit in tg.units -- tg.tree and tg.units must come " - "from the same extract()/build_schedule_tree() pipeline run" - ) - contract = HoleContract( - name=f"HOLE_{stmt_name}", - op_ref=unit.op, - inputs=_ordered_inputs(unit, _by_buf(tg.reads, stmt_name), dtype_table), - output=_output_view(unit, _by_buf(tg.writes, stmt_name), dtype_table), - coords=coords, - ) - contracts[stmt_name] = contract - return contract - - def line(printer, text: str): - return printer.start_line().print_str(text).end_line() - - def print_user(printer, options, node): - """Isl asks for one statement's text and owns the loop nest and the indentation around it. - - Isl asks for one statement's text and owns the loop nest and the - indentation around it, so the hole is emitted here rather than spliced - into finished output. - - The hole and its sync sit in their own brace block: isl prints a - single-statement loop body without braces, and two bare statements - there would put the sync outside the loop. - """ - stmt_name, coords = _call_coords(node.get_expr()) - contract = contract_for(stmt_name, coords) - in_refs = tuple(_ring_ref(v.tensor_name, coords, tg.ring) for v in contract.inputs) - out_ref = _ring_ref(contract.output.tensor_name, coords, tg.ring) - printer = line(printer, "{") - printer = printer.indent(2) - printer = line(printer, _render_hole_call(contract.name, in_refs, out_ref, coords)) - printer = line(printer, "// barrier") - printer = printer.indent(-2) - return line(printer, "}") - - ast = isl.ast_build.from_context(isl.set("{ : }")).node_from(tg.tree) - options = isl.ast_print_options.alloc().set_print_user(print_user) - text = _print_to_str(ast, options) - - holes = tuple(contract.name for contract in contracts.values()) - return Skeleton(text=text, holes=holes), contracts - - - - - - - -def _statement_extents(tg: TileGraph, stmt_name: str) -> tuple[int, ...]: - """Per-axis ``[lo, hi]`` extent of ``stmt_name``'s own domain piece. - - Per-axis ``[lo, hi]`` extent of ``stmt_name``'s own domain piece - (same ``dim_min_val``/``dim_max_val`` technique, applied to the one - ``isl.set`` in ``tg.domain`` whose tuple name matches). - """ - sets: list["isl.set"] = [] - tg.domain.foreach_set(sets.append) - for s in sets: - if s.get_tuple_name() == stmt_name: - rank = s.dim(isl.dim_type.SET) - return tuple( - int(s.dim_max_val(i).num_si()) - int(s.dim_min_val(i).num_si()) + 1 - for i in range(rank) - ) - raise EmitScaffoldError(f"emit_scaffold: no domain set found for statement {stmt_name!r}") - - -def _illustrative_instances( - extents: tuple[int, ...], -) -> tuple[list[tuple[int, ...]], int]: - """Minimal loop unrolling for the swimlane. - - Show the prologue, up to ``depth + 1`` steady-state instances, and the - epilogue. The Cartesian product stays lazy; the last coordinate is derived - directly from extents. Returns shown coordinates and the omitted count. - """ - depth = len(extents) - total = math.prod(extents) - head_n = min(1 + (depth + 1), total) - head = list(itertools.islice(itertools.product(*(range(e) for e in extents)), head_n)) - last = tuple(e - 1 for e in extents) - shown = head + ([last] if total > head_n and last not in head else []) - return shown, total - len(shown) - - -def _swimlane_lines(tg: TileGraph, contracts: dict[str, HoleContract]) -> list[str]: - lines = [ - "```mermaid", - "gantt", - " title tilefoundry scaffold -- statement swimlanes", - " dateFormat X", - " axisFormat %s", - ] - tick = 0 - for stmt_name in contracts: - extents = _statement_extents(tg, stmt_name) - shown, collapsed = _illustrative_instances(extents) - lines.append(f" section {stmt_name}") - n_shown = len(shown) - for idx, coord in enumerate(shown): - is_last = idx == n_shown - 1 - if collapsed and is_last: - lines.append(f" ... x{collapsed} elided :{tick}, 1d") - tick += 1 - if not collapsed: - role = "" - elif idx == 0: - role = " (prologue)" - elif is_last: - role = " (epilogue)" - else: - role = " (steady)" - label = f"{stmt_name}({', '.join(str(c) for c in coord)}){role}" - lines.append(f" {label} :{tick}, 1d") - tick += 1 - lines.append("```") - return lines - - -def _build_swimlane(tg: TileGraph, contracts: dict[str, HoleContract]) -> Swimlane: - return Swimlane(text="\n".join(_swimlane_lines(tg, contracts))) - - - - - - - -def emit_scaffold( - graph: TileGraph, tree: "isl.schedule", ring: dict[str, int] -) -> tuple[Skeleton, Swimlane, list[HoleContract]]: - """Render private scheduling state over immutable analysis output.""" - program = _RenderProgram(graph, tree, ring) - dtype_table = _dtype_table(program.units) - skeleton, contracts = _build_skeleton(program, dtype_table) - swimlane = _build_swimlane(program, contracts) - return skeleton, swimlane, list(contracts.values()) - - -__all__ = [ - "EmitScaffoldError", - "emit_scaffold", -] diff --git a/src/tilefoundry/target/__init__.py b/src/tilefoundry/target/__init__.py index f5909ff1..8b839dab 100644 --- a/src/tilefoundry/target/__init__.py +++ b/src/tilefoundry/target/__init__.py @@ -19,7 +19,7 @@ TopologyLimitFacts, facts_result, ) -from tilefoundry.target.services import Analyzer, Scheduler +from tilefoundry.target.services import Analyzer _ANALYSIS_FACTS = { "MemoryHierarchyFacts", @@ -74,7 +74,6 @@ def validate_cuda_topology_levels(target: Target, names) -> None: "MemoryHierarchyFacts", "ParallelCapacityFacts", "PerformanceServiceFacts", - "Scheduler", "Target", "TargetFactsError", "TopologyLimitFacts", diff --git a/src/tilefoundry/target/amx/atoms.py b/src/tilefoundry/target/amx/atoms.py deleted file mode 100644 index 092373f0..00000000 --- a/src/tilefoundry/target/amx/atoms.py +++ /dev/null @@ -1,253 +0,0 @@ -"""``candidate_atoms(op, target) -> list[AtomFact]``. - -``candidate_atoms(op, target) -> list[AtomFact]`` -- bridge one HIR -compute op to the atom catalogue of the AMX target, which spans two -execution units: the AMX coprocessor and the core's own NEON SIMD pipes. -It only *lists* candidates (a hard filter over shape, dtype, layout and -operand storage); it never picks one, that choice is the schedule layer's. -""" -from __future__ import annotations - -import math -from dataclasses import dataclass - -from tilefoundry.ir.core import Call -from tilefoundry.ir.hir.nn.matmul import MatMul, matmul_axes -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.schedule.facts import AtomFact -from tilefoundry.target import Target -from tilefoundry.target.amx.spec import APPLE_AMX_ID -from tilefoundry.target.amx.target import AmxTarget -from tilefoundry.target.base import target_instance - - -@dataclass(frozen=True) -class StorageLevel: - """Where an atom's operands sit while it executes. - - Where an atom's operands sit while it executes: per operand role, the - bytes that role has to fit into. A level backed by a larger store only - streams its operands through, so it budgets no role and holds anything. - """ - - name: str - budget: tuple[tuple[str, int], ...] = () - - def holds(self, operand_bytes: dict[str, int]) -> bool: - """Whether every budgeted role fits -- vacuously so when none is.""" - return all(operand_bytes[role] <= limit for role, limit in self.budget) - - -_ISA = AmxTarget.hardware.resolve(APPLE_AMX_ID).value - - - -AMX_REGISTERS = StorageLevel( - name="amx_xyz_registers", - budget=( - ("a_bytes", _ISA.staging_bytes), - ("b_bytes", _ISA.staging_bytes), - ("c_bytes", _ISA.accumulator_bytes), - ), -) - - - - -CORE_CACHE = StorageLevel(name="core_cache") - - -@dataclass(frozen=True) -class AmxOpSpec: - """A named, fully-specified matrix instruction: which execution unit issues it. - - A named, fully-specified matrix instruction: which execution unit issues - it, and which storage level has to hold the operands it is handed. - """ - - name: str - unit: str - level: StorageLevel - shape_mnk: tuple[int, int, int] - dtype_a: DType - dtype_b: DType - dtype_c: DType - - -@dataclass(frozen=True) -class AmxAtom: - """Realized atom -- op plus the bytes each of its own operands occupies.""" - - op: AmxOpSpec - a_bytes: int - b_bytes: int - c_bytes: int - - - - -AMX_FMA32_16x16x1_F32 = AmxOpSpec( - name="AMX_FMA32_16x16x1_F32", - unit="amx", - level=AMX_REGISTERS, - shape_mnk=(16, 16, 1), - dtype_a=DType.f32, - dtype_b=DType.f32, - dtype_c=DType.f32, -) - - - - -NEON_FMLA_4x4x1_F32 = AmxOpSpec( - name="NEON_FMLA_4x4x1_F32", - unit="neon", - level=CORE_CACHE, - shape_mnk=(4, 4, 1), - dtype_a=DType.f32, - dtype_b=DType.f32, - dtype_c=DType.f32, -) - - - -_AMX_OP_CATALOG: tuple[AmxOpSpec, ...] = (AMX_FMA32_16x16x1_F32, NEON_FMLA_4x4x1_F32) - - -def _dense_bytes(shape: tuple[int, ...], dtype: DType) -> int: - """Bytes for one dense ``shape``/``dtype`` operand.""" - return math.ceil(math.prod(shape) * dtype.bit_width / 8) - - -def _operand_bytes(shape_mnk: tuple[int, int, int], op: AmxOpSpec) -> dict[str, int]: - """Bytes each of a ``shape_mnk`` matmul's operand roles has to hold at ``op``'s dtypes. - - Bytes each of a ``shape_mnk`` matmul's operand roles has to hold at - ``op``'s dtypes. - - A and B are *staged* one reduction step at a time, so their roles hold one - column and one row rather than the whole K extent. C accumulates, so its - role holds the entire M by N block for as long as the reduction runs -- - which is what makes a wide matmul unable to sit in a register file that a - narrow one fits exactly. - """ - m, n, _ = shape_mnk - return { - "a_bytes": _dense_bytes((m,), op.dtype_a), - "b_bytes": _dense_bytes((n,), op.dtype_b), - "c_bytes": _dense_bytes((m, n), op.dtype_c), - } - - -def make_atom(op: AmxOpSpec) -> AmxAtom: - """Realize ``op``: the bytes its own three operand roles need.""" - return AmxAtom(op=op, **_operand_bytes(op.shape_mnk, op)) - - -def _roofline_duration_ns(atom: AmxAtom, target: AmxTarget) -> tuple[float, float]: - """Estimate nominal compute and traffic time for one AMX atom. - - Compute uses MNK flops and measured unit throughput; memory uses operand - bytes and unified bandwidth. Return the maximum and compute-only time in - nanoseconds so callers that account for traffic do not charge it twice. - This ranks candidates rather than predicting measured performance. Values - remain sub-nanosecond instead of applying a one-nanosecond floor. - """ - m, n, k = atom.op.shape_mnk - flops = 2 * m * n * k - moved_bytes = atom.a_bytes + atom.b_bytes + atom.c_bytes - device = target.device - compute_ns = flops * 1_000_000_000 / device.throughput_for(atom.op.unit, atom.op.dtype_a) - memory_ns = ( - moved_bytes * 1_000_000_000 / device.unified_memory_bandwidth_bytes_per_second - ) - return max(compute_ns, memory_ns), compute_ns - - -def _operands_layout_ok(lhs: TensorType, rhs: TensorType) -> bool: - """Layout hard filter: the X/Y operand packing is derived for dense, unsharded operands. - - Layout hard filter: the X/Y operand packing is derived for dense, - unsharded operands. A ShardLayout-carrying operand may need a repack step to - feed this atom, which is an agent-filled hole rather than a candidate here. - """ - return lhs.layout is None and rhs.layout is None - - -def _static_positive(*dims: object) -> bool: - return all(isinstance(d, int) and not isinstance(d, bool) and d > 0 for d in dims) - - -def candidate_atoms(op: Call, target: Target | None = None) -> list[AtomFact]: - """List AMX atoms eligible for an HIR ``MatMul`` Call. - - This is a hard filter: static M/N/K divisibility, input dtypes, layouts, and - three operands fitting the atom storage level. The storage test separates a - register-resident unit from cache streaming. ``[]`` is valid; unsupported - operation kinds and Targets raise. - """ - target = AmxTarget() if target is None else target - target = target_instance(target) - if not isinstance(op, Call) or not isinstance(op.target, MatMul): - got = type(op).__name__ - if isinstance(op, Call): - got += f" (target={type(op.target).__name__})" - raise NotImplementedError( - f"candidate_atoms: only a MatMul Call is supported, got {got}" - ) - if not isinstance(target, AmxTarget): - raise NotImplementedError( - "candidate_atoms: only AmxTarget is supported, got " - f"{type(target).__name__}" - ) - - lhs_type, rhs_type = op.args[0].type, op.args[1].type - a_m, a_k, b_n, _b_k = matmul_axes(op.target) - m, k = lhs_type.shape[a_m], lhs_type.shape[a_k] - n = rhs_type.shape[b_n] - if not _static_positive(m, n, k) or not _operands_layout_ok(lhs_type, rhs_type): - return [] - - facts: list[AtomFact] = [] - for amx_op in _AMX_OP_CATALOG: - atom_m, atom_n, atom_k = amx_op.shape_mnk - if m % atom_m != 0 or n % atom_n != 0 or k % atom_k != 0: - continue - if lhs_type.dtype != amx_op.dtype_a or rhs_type.dtype != amx_op.dtype_b: - continue - if not amx_op.level.holds(_operand_bytes((m, n, k), amx_op)): - continue - atom = make_atom(amx_op) - duration, compute_duration = _roofline_duration_ns(atom, target) - facts.append( - AtomFact( - shape=amx_op.shape_mnk, - dtype=(amx_op.dtype_a, amx_op.dtype_b, amx_op.dtype_c), - duration=duration, - compute_duration=compute_duration, - storage={ - "a_bytes": atom.a_bytes, - "b_bytes": atom.b_bytes, - "c_bytes": atom.c_bytes, - "operand_bytes": atom.a_bytes + atom.b_bytes + atom.c_bytes, - }, - - resource={amx_op.unit: 1}, - is_async=False, - atom=atom, - ) - ) - return facts - - -__all__ = [ - "AMX_FMA32_16x16x1_F32", - "AMX_REGISTERS", - "CORE_CACHE", - "NEON_FMLA_4x4x1_F32", - "AmxAtom", - "AmxOpSpec", - "StorageLevel", - "candidate_atoms", - "make_atom", -] diff --git a/src/tilefoundry/target/amx/facts.py b/src/tilefoundry/target/amx/facts.py index ea058085..ab5cf975 100644 --- a/src/tilefoundry/target/amx/facts.py +++ b/src/tilefoundry/target/amx/facts.py @@ -7,8 +7,6 @@ from __future__ import annotations -from types import SimpleNamespace - from tilefoundry.analysis.facts import ( ExplicitMemoryLevelFacts, ImplicitMemoryLevelFacts, @@ -19,17 +17,12 @@ ThroughputFacts, ) from tilefoundry.ir.types import DType -from tilefoundry.schedule.facts import AtomFact -from .atoms import candidate_atoms from .target import AmxTarget _ROOFLINE_UNIT = "amx" -_PIPELINE_TOPOLOGY = "core" - - def memory_hierarchy(target: AmxTarget, query: object = None) -> MemoryHierarchyFacts: """The AMX memory levels and how they are related. @@ -122,54 +115,8 @@ def parallel_capacity( ) -def pipeline_facts(target: AmxTarget, query: object) -> object: - """Project the finite AMX instruction catalogue before solving. - - An AMX core both runs the work and owns the L1d the tile lives in, so here - the level asked about and the level the capacity belongs to are the same one. - """ - from tilefoundry.schedule.pipeline.facts import ( # noqa: PLC0415 - PipelineFacts, - PipelineFactsQuery, - PipelineInstructionFacts, - ) - - if not isinstance(query, PipelineFactsQuery): - raise TypeError("AMX pipeline facts need a PipelineFactsQuery") - if query.topology != _PIPELINE_TOPOLOGY: - raise ValueError(f"AMX states no pipeline facts for {query.topology!r}") - instructions: list[PipelineInstructionFacts] = [] - for statement_id, op in query.statements: - try: - candidates = tuple(candidate_atoms(op, target)) - except NotImplementedError: - candidates = () - if not candidates: - candidates = ( - AtomFact( - shape=(1, 1, 1), - dtype=(DType.f32, DType.f32, DType.f32), - duration=1.0, - compute_duration=1.0, - storage={}, - resource={"core": 1}, - is_async=False, - atom=SimpleNamespace(op=SimpleNamespace(name="amx.scalar")), - ), - ) - instructions.append(PipelineInstructionFacts(statement_id, candidates)) - return PipelineFacts( - topology=query.topology, - tile_capacity_scope=_PIPELINE_TOPOLOGY, - tile_capacity_bytes=target.device.l1d_bytes_per_performance_core, - max_threads_per_warp=1, - instructions=tuple(instructions), - ) - - __all__ = [ "memory_hierarchy", "parallel_capacity", - "pipeline_facts", "throughput", ] diff --git a/src/tilefoundry/target/amx/schedule.py b/src/tilefoundry/target/amx/schedule.py deleted file mode 100644 index beb68734..00000000 --- a/src/tilefoundry/target/amx/schedule.py +++ /dev/null @@ -1,55 +0,0 @@ -"""The AMX scheduler registered at the core level.""" - -from __future__ import annotations - -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.schedule import ScheduleError, ScheduleOptions -from tilefoundry.target.services import Scheduler - -from .target import AmxTarget - -TOPOLOGY = "core" - - -def schedule_core( - module: Module, - function: Function, - target: AmxTarget, - topology: object, - options: object | None = None, -) -> object: - """Build, close, solve, and export the AMX core pipeline schedule.""" - from tilefoundry.schedule.pipeline import ( # noqa: PLC0415 - PipelineFacts, - build_pipeline_problem, - build_pipeline_program, - export_pipeline_plan, - solve_pipeline_problem, - ) - - if function is not module.entry_function(): - raise ScheduleError( - f"{TOPOLOGY} schedule requires the module entry function, got " - f"{function.name!r}" - ) - if options is not None and not isinstance(options, ScheduleOptions): - raise ScheduleError( - f"{TOPOLOGY} schedule options must be ScheduleOptions, got " - f"{type(options).__name__}" - ) - program = build_pipeline_program(module, function) - facts = target.get_facts(PipelineFacts, program.facts_query(TOPOLOGY)) - problem = build_pipeline_problem(program, facts, topology) - solution = solve_pipeline_problem(problem) - return export_pipeline_plan(program, solution, target) - - -def amx_scheduler(topology: str) -> Scheduler | None: - """Construct the AMX scheduler requested for one topology level.""" - if topology == TOPOLOGY: - return Scheduler(TOPOLOGY, schedule_core) - return None - - -__all__ = ["TOPOLOGY", "amx_scheduler", "schedule_core"] diff --git a/src/tilefoundry/target/amx/target.py b/src/tilefoundry/target/amx/target.py index 3b9dc629..a4464d13 100644 --- a/src/tilefoundry/target/amx/target.py +++ b/src/tilefoundry/target/amx/target.py @@ -27,7 +27,6 @@ ) from tilefoundry.target.facts import TopologyLimitFacts, facts_result from tilefoundry.target.hardware.envelope import HardwareDocument -from tilefoundry.target.services import Scheduler from tilefoundry.utils.python_source import PythonExpr @@ -145,24 +144,8 @@ def get_facts(self, facts_type: type, query: object | None = None): return facts_result(self, facts_type, throughput(self, query)) if facts_type is ParallelCapacityFacts: return facts_result(self, facts_type, parallel_capacity(self, query)) - - from tilefoundry.schedule.pipeline import PipelineFacts # noqa: PLC0415 - - if facts_type is PipelineFacts: - from tilefoundry.target.amx.facts import pipeline_facts # noqa: PLC0415 - - return facts_result(self, facts_type, pipeline_facts(self, query)) return super().get_facts(facts_type, query) - def get_scheduler(self, topology: str) -> Scheduler: - """Select the AMX core scheduler.""" - from tilefoundry.target.amx.schedule import amx_scheduler # noqa: PLC0415 - - scheduler = amx_scheduler(topology) - if scheduler is not None: - return scheduler - return super().get_scheduler(topology) - def _python_import_module(self) -> str: if type(self) is AmxTarget: return "tilefoundry.target.amx" diff --git a/src/tilefoundry/target/base.py b/src/tilefoundry/target/base.py index 80b4dbdc..0ca01031 100644 --- a/src/tilefoundry/target/base.py +++ b/src/tilefoundry/target/base.py @@ -21,7 +21,7 @@ UnknownSchemaError, parse_document, ) -from tilefoundry.target.services import Analyzer, CodeGenerator, Scheduler +from tilefoundry.target.services import Analyzer, CodeGenerator from tilefoundry.utils.python_source import PythonExpr, dataclass_to_python FactsT = TypeVar("FactsT") @@ -203,12 +203,6 @@ def get_analyzer(self, selector: str) -> Analyzer: f"{_target_summary(self)}: no analyzer for {selector!r}" ) - def get_scheduler(self, topology: str) -> Scheduler: - """Return the scheduler selected by this concrete Target.""" - raise UnsupportedCapabilityError( - f"{_target_summary(self)}: no scheduler for {topology!r}" - ) - def get_code_generator(self) -> CodeGenerator: """Return the code-generation service selected by this Target.""" raise UnsupportedCapabilityError( diff --git a/src/tilefoundry/target/cuda/atoms.py b/src/tilefoundry/target/cuda/atoms.py deleted file mode 100644 index d73da84c..00000000 --- a/src/tilefoundry/target/cuda/atoms.py +++ /dev/null @@ -1,170 +0,0 @@ -"""``candidate_atoms(op, target) -> list[AtomFact]``. - -``candidate_atoms(op, target) -> list[AtomFact]`` -- bridge one HIR -compute op to the TIR MMA atom catalogue it could run on. It only *lists* -candidates (a hard filter over shape/dtype/layout); it never picks one, -that ranking is the schedule layer's CP-SAT job. See each helper's -docstring below for exactly where its numbers come from. -""" -from __future__ import annotations - -from tilefoundry.ir.core import Call -from tilefoundry.ir.hir.nn.matmul import MatMul, matmul_axes -from tilefoundry.ir.tir.cuda.nn.mma import SM80_16x8x16_F32BF16BF16F32_TN, make_atom -from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaOpSpec -from tilefoundry.ir.types import DType, TensorType, tensor_bytes -from tilefoundry.ir.types.shard import ShardLayout -from tilefoundry.ir.types.shard.shard_layout import shard_layout_local_shape -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.schedule.facts import AtomFact -from tilefoundry.target import Target, default_target -from tilefoundry.target.base import target_instance -from tilefoundry.target.cuda.target import CudaTarget - -_MMA_OP_CATALOG: tuple[MmaOpSpec, ...] = (SM80_16x8x16_F32BF16BF16F32_TN,) - - -def _is_async_op(op: MmaOpSpec) -> bool: - """wgmma-family instructions issue asynchronously; the sole registered op today is synchronous. - - wgmma-family instructions (SM90+) issue asynchronously; the sole - registered op today (SM80 ``mma.sync``) is synchronous. No wgmma - ``MmaOpSpec`` is registered yet in ``ir.tir.cuda.nn.mma._ATOM_TABLE``, so - this substring check is forward-compatible naming, not a real dispatch - -- it always returns ``False`` today. - """ - return "wgmma" in op.name.lower() - - -def _dense_bytes(shape: tuple[int, ...], dtype: DType) -> int: - """Bytes for one dense (unsharded) ``shape``/``dtype`` tile. - - Bytes for one dense (unsharded) ``shape``/``dtype`` tile -- reuses - ``target.cuda.cost.tensor_bytes`` (the exact helper the ``MatMul`` cost - evaluator uses for its own HBM-traffic accounting), so op-level and - atom-level byte counts share one formula. ``storage``/``layout`` are - irrelevant to a byte count, so a placeholder ``GMEM`` is used. - """ - type = TensorType(shape=shape, dtype=dtype, layout=None, storage=StorageKind.GMEM) - return tensor_bytes(type) - - -def _fragment_reg_bytes(fragment: ShardLayout, dtype: DType) -> int: - """Per-thread register bytes for one A/B/C fragment. - - Per-thread register bytes for one A/B/C fragment: divide the - fragment's global layout shape down to its per-thread local shape - (``shard_layout_local_shape`` -- quotients out the Split-bound lane - axes; see ``mma.py``'s fragment derivation comments for the 8/4/4 - elements-per-thread this yields for A/B/C), then reuse ``tensor_bytes`` - for the numel*bitwidth->bytes conversion (the same formula as - ``_dense_bytes``, at per-thread instead of whole-tile granularity). - """ - local_shape = shard_layout_local_shape(fragment) - type = TensorType(shape=local_shape, dtype=dtype, layout=None, storage=StorageKind.RMEM) - return tensor_bytes(type) - - -def _roofline_duration_ns(op: MmaOpSpec, target: CudaTarget) -> tuple[float, float]: - """Estimate nominal compute and traffic time for one CUDA MMA atom. - - Compute uses the atom's MNK flops, device peak throughput, and SM count; - memory uses dense A+B+C bytes and HBM bandwidth. Return the maximum with a - one-nanosecond floor and the compute-only time so callers that account for - traffic do not charge it twice. This ranks candidates rather than - predicting measured performance. - """ - m, n, k = op.shape_mnk - flops = 2 * m * n * k - moved_bytes = ( - _dense_bytes((m, k), op.dtype_a) - + _dense_bytes((k, n), op.dtype_b) - + _dense_bytes((m, n), op.dtype_c) - ) - device = target.device - compute_ns = flops * 1_000_000_000 * device.sm_count / device.peak_for(op.dtype_a) - memory_ns = moved_bytes * 1_000_000_000 / device.hbm_bandwidth_bytes_per_second - return max(compute_ns, memory_ns, 1.0), compute_ns - - -def _operands_layout_ok(lhs: TensorType, rhs: TensorType) -> bool: - """V1 layout hard filter. - - V1 layout hard filter: the registered fragment layouts - (``ir.tir.cuda.nn.mma``) are derived for dense, unsharded row-major - operands -- the sole ``operand_layout="TN"`` convention today. A - ShardLayout-carrying operand may need a pack/repack step to feed this - atom; V1 defers that to an agent-filled hole (out of scope for this - bridge) and simply excludes such operands from candidacy here, rather - than erroring. - """ - return lhs.layout is None and rhs.layout is None - - -def _static_positive(*dims: object) -> bool: - return all(isinstance(d, int) and not isinstance(d, bool) and d > 0 for d in dims) - - -def candidate_atoms(op: Call, target: Target | None = None) -> list[AtomFact]: - """List CUDA MMA atoms eligible for an HIR ``MatMul`` Call. - - This is a hard filter, not ranking: static M/N/K must divide ``shape_mnk``; - input dtypes and operand layouts must agree. ``dtype_c`` is the atom's - internal accumulator, so narrowing output is an epilogue concern. ``[]`` is - a valid outcome; unsupported operation kinds and Targets raise. - """ - target = default_target() if target is None else target - target = target_instance(target) - if not isinstance(op, Call) or not isinstance(op.target, MatMul): - got = type(op).__name__ - if isinstance(op, Call): - got += f" (target={type(op.target).__name__})" - raise NotImplementedError( - f"candidate_atoms: only a MatMul Call is supported, got {got}" - ) - if not isinstance(target, CudaTarget): - raise NotImplementedError( - "candidate_atoms: only CudaTarget is supported, got " - f"{type(target).__name__}" - ) - - lhs_type, rhs_type = op.args[0].type, op.args[1].type - a_m, a_k, b_n, _b_k = matmul_axes(op.target) - m, k = lhs_type.shape[a_m], lhs_type.shape[a_k] - n = rhs_type.shape[b_n] - if not _static_positive(m, n, k) or not _operands_layout_ok(lhs_type, rhs_type): - return [] - - facts: list[AtomFact] = [] - for mma_op in _MMA_OP_CATALOG: - atom_m, atom_n, atom_k = mma_op.shape_mnk - if m % atom_m != 0 or n % atom_n != 0 or k % atom_k != 0: - continue - if lhs_type.dtype != mma_op.dtype_a or rhs_type.dtype != mma_op.dtype_b: - continue - atom = make_atom(mma_op) - a_bytes = _fragment_reg_bytes(atom.A, mma_op.dtype_a) - b_bytes = _fragment_reg_bytes(atom.B, mma_op.dtype_b) - c_bytes = _fragment_reg_bytes(atom.C, mma_op.dtype_c) - duration, compute_duration = _roofline_duration_ns(mma_op, target) - facts.append( - AtomFact( - shape=mma_op.shape_mnk, - dtype=(mma_op.dtype_a, mma_op.dtype_b, mma_op.dtype_c), - duration=duration, - compute_duration=compute_duration, - storage={ - "a_reg_bytes": a_bytes, - "b_reg_bytes": b_bytes, - "c_reg_bytes": c_bytes, - "reg_bytes": a_bytes + b_bytes + c_bytes, - }, - resource={"lane": atom.required_scope.topologies[0].size}, - is_async=_is_async_op(mma_op), - atom=atom, - ) - ) - return facts - - -__all__ = ["candidate_atoms"] diff --git a/src/tilefoundry/target/cuda/facts.py b/src/tilefoundry/target/cuda/facts.py index 67d65b74..98909955 100644 --- a/src/tilefoundry/target/cuda/facts.py +++ b/src/tilefoundry/target/cuda/facts.py @@ -2,14 +2,12 @@ Each conversion answers exactly one aggregate's question, so what a consumer can read is visible here rather than spread through the consumers themselves. Nothing -in this module decides anything: it restates the installed hardware documents, and -the target's own atom catalogue, in the shape the asking algorithm declared. +in this module decides anything: it restates the installed hardware documents in +the shape the asking analysis declared. """ from __future__ import annotations -from types import SimpleNamespace - from tilefoundry.analysis.facts import ( ExplicitMemoryLevelFacts, ImplicitMemoryLevelFacts, @@ -20,22 +18,9 @@ PerformanceServiceFacts, ThroughputFacts, ) -from tilefoundry.ir.types import DType -from tilefoundry.schedule.facts import AtomFact -from tilefoundry.schedule.plan import TargetSpecRef -from .atoms import candidate_atoms from .target import CudaTarget -_PIPELINE_TOPOLOGY = "thread" - - - -_TILE_CAPACITY_SCOPE = "cta" - - -_PARTITION_TOPOLOGY = "cta" - def memory_hierarchy(target: CudaTarget, query: object = None) -> MemoryHierarchyFacts: """The CUDA memory levels and how they are related. @@ -150,105 +135,15 @@ def parallel_capacity( """How many CTAs the plan assumes run at once. This is a compiler policy, not CUDA's grid limit and not the hardware - resident-CTA maximum: one active CTA per SM. A tighter policy is a - scheduling input, and changing it changes the plan rather than the program. + resident-CTA maximum: one active CTA per SM. A tighter policy changes what + analysis concludes rather than the program. """ return ParallelCapacityFacts(topology="cta", parallel_units=target.device.sm_count) -def pipeline_facts(target: CudaTarget, query: object) -> object: - """Project every instruction and capacity fact before pipeline solving. - - The capacity is per-CTA shared memory and is reported as CTA-scoped, because - that is whose store it is; the level being decided about is finer. - """ - from tilefoundry.schedule.pipeline.facts import ( # noqa: PLC0415 - PipelineFacts, - PipelineFactsQuery, - PipelineInstructionFacts, - ) - - if not isinstance(query, PipelineFactsQuery): - raise TypeError( - "CudaTarget pipeline facts need a PipelineFactsQuery, got " - f"{type(query).__name__}" - ) - if query.topology != _PIPELINE_TOPOLOGY: - raise ValueError( - f"CudaTarget states no pipeline facts for {query.topology!r}; it " - f"pipelines {_PIPELINE_TOPOLOGY!r}" - ) - instructions: list[PipelineInstructionFacts] = [] - for statement_id, op in query.statements: - if not isinstance(statement_id, str) or not statement_id: - raise ValueError(f"pipeline statement id must be a non-empty string, got {statement_id!r}") - try: - candidates = tuple(candidate_atoms(op, target)) - except NotImplementedError: - candidates = () - if not candidates: - candidates = ( - AtomFact( - shape=(1, 1, 1), - dtype=(DType.f32, DType.f32, DType.f32), - duration=1.0, - compute_duration=1.0, - storage={}, - resource={"lane": 1}, - is_async=False, - atom=SimpleNamespace(op=SimpleNamespace(name="cuda.scalar")), - ), - ) - instructions.append(PipelineInstructionFacts(statement_id, candidates)) - return PipelineFacts( - topology=query.topology, - tile_capacity_scope=_TILE_CAPACITY_SCOPE, - tile_capacity_bytes=target.architecture.shared_memory_per_cta_bytes, - max_threads_per_warp=target.architecture.max_threads_per_warp, - instructions=tuple(instructions), - ) - - -def partition_facts(target: CudaTarget, query: object) -> object: - """Project every rate, capacity, and position count before partitioning. - - All of it is what the installed documents state: how many SMs the device has, - how fast and how large its memory is, and what it peaks at per dtype. How much - of that machine to occupy is the caller's to decide, so no policy is encoded - here. After this call the partition holds numbers, not a target. - """ - from tilefoundry.schedule.partition.facts import ( # noqa: PLC0415 - PartitionFacts, - PartitionFactsQuery, - ) - - if not isinstance(query, PartitionFactsQuery): - raise TypeError( - "CudaTarget partition facts need a PartitionFactsQuery, got " - f"{type(query).__name__}" - ) - if query.topology != _PARTITION_TOPOLOGY: - raise ValueError( - f"CudaTarget states no partition facts for {query.topology!r}; it " - f"partitions {_PARTITION_TOPOLOGY!r}" - ) - device = target.device - return PartitionFacts( - topology=query.topology, - spec=TargetSpecRef.of(target), - parallel_units=device.sm_count, - memory_bandwidth_bytes_per_second=device.hbm_bandwidth_bytes_per_second, - memory_capacity_bytes=device.hbm_capacity_bytes, - peak_flops_per_second=tuple( - sorted(device.dense_flops_per_second.items(), key=lambda item: item[0].name) - ), - ) - - __all__ = [ "memory_hierarchy", "parallel_capacity", - "partition_facts", - "pipeline_facts", + "performance_service", "throughput", ] diff --git a/src/tilefoundry/target/cuda/schedule.py b/src/tilefoundry/target/cuda/schedule.py deleted file mode 100644 index 672c62b3..00000000 --- a/src/tilefoundry/target/cuda/schedule.py +++ /dev/null @@ -1,125 +0,0 @@ -"""What CUDA scheduling is, at each level CUDA schedules. - -Two levels, two decisions. At `thread` the question is how the warps of one CTA -overlap their work asynchronously; the tile they cooperate on is the CTA's, which -is why the capacity fact projected there is per-CTA shared memory. At `cta` the -question is how work and its tensors spread across the device. - -Each entry does the same four things in the same order -- build the program view, -ask the hardware once, close and solve the problem, export the plan -- and the -steps below that boundary stay private to the algorithm that owns them. -""" - -from __future__ import annotations - -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.specialize import origin_of -from tilefoundry.schedule import ScheduleError, ScheduleOptions -from tilefoundry.target.services import Scheduler - -from .target import CudaTarget - -PIPELINE_TOPOLOGY = "thread" -PARTITION_TOPOLOGY = "cta" - - -def _options(topology: str, options: object | None) -> ScheduleOptions: - """The solver controls to run under, defaulted rather than inferred.""" - if options is None: - return ScheduleOptions() - if not isinstance(options, ScheduleOptions): - raise ScheduleError( - f"{topology} schedule options must be ScheduleOptions, got " - f"{type(options).__name__}" - ) - return options - - -def _entry_function(topology: str, module: Module, function: Function) -> None: - """Refuse a leaf: these algorithms decide the whole launch. - - The entry at a chosen size is a different object and the same program, so - the comparison follows what a specialised function records about where it - came from. Comparing objects alone would refuse it; comparing names would - accept anything called the same. - """ - entry = module.entry_function() - permitted = (entry, *entry.variants) - candidate: object | None = function - while candidate is not None: - if any(candidate is allowed for allowed in permitted): - return - candidate = origin_of(candidate) - raise ScheduleError( - f"{topology} schedule requires the module entry function, got " - f"{function.name!r}" - ) - - -def schedule_thread( - module: Module, - function: Function, - target: CudaTarget, - topology: object, - options: object | None = None, -) -> object: - """Build, close, solve, and export the intra-CTA pipeline schedule.""" - from tilefoundry.schedule.pipeline import ( # noqa: PLC0415 - PipelineFacts, - build_pipeline_problem, - build_pipeline_program, - export_pipeline_plan, - solve_pipeline_problem, - ) - - _entry_function(PIPELINE_TOPOLOGY, module, function) - _options(PIPELINE_TOPOLOGY, options) - program = build_pipeline_program(module, function) - facts = target.get_facts(PipelineFacts, program.facts_query(PIPELINE_TOPOLOGY)) - problem = build_pipeline_problem(program, facts, topology) - solution = solve_pipeline_problem(problem) - return export_pipeline_plan(program, solution, target) - - -def schedule_cta( - module: Module, - function: Function, - target: CudaTarget, - topology: object, - options: object | None = None, -) -> object: - """Build, close, solve, and export the device-wide partition schedule.""" - from tilefoundry.schedule.partition import ( # noqa: PLC0415 - PartitionFacts, - build_partition_problem, - build_partition_program, - export_partition_plan, - solve_partition_problem, - ) - - _entry_function(PARTITION_TOPOLOGY, module, function) - resolved = _options(PARTITION_TOPOLOGY, options) - program = build_partition_program(module, function) - facts = target.get_facts(PartitionFacts, program.facts_query(PARTITION_TOPOLOGY)) - problem = build_partition_problem(program, facts, topology) - solution = solve_partition_problem(problem, resolved) - return export_partition_plan(problem, solution) - - -__all__ = [ - "PARTITION_TOPOLOGY", - "PIPELINE_TOPOLOGY", - "schedule_cta", - "cuda_scheduler", - "schedule_thread", -] - - -def cuda_scheduler(topology: str) -> Scheduler | None: - """Construct the CUDA scheduler requested for one topology level.""" - if topology == PIPELINE_TOPOLOGY: - return Scheduler(PIPELINE_TOPOLOGY, schedule_thread) - if topology == PARTITION_TOPOLOGY: - return Scheduler(PARTITION_TOPOLOGY, schedule_cta) - return None diff --git a/src/tilefoundry/target/cuda/target.py b/src/tilefoundry/target/cuda/target.py index 6b4365f5..42c82237 100644 --- a/src/tilefoundry/target/cuda/target.py +++ b/src/tilefoundry/target/cuda/target.py @@ -27,7 +27,7 @@ ) from tilefoundry.target.facts import TopologyLimitFacts, facts_result from tilefoundry.target.hardware.envelope import HardwareDocument -from tilefoundry.target.services import CodeGenerator, Scheduler +from tilefoundry.target.services import CodeGenerator from tilefoundry.utils.python_source import PythonExpr @@ -145,31 +145,8 @@ def get_facts(self, facts_type: type, query: object | None = None): return facts_result(self, facts_type, parallel_capacity(self, query)) if facts_type is PerformanceServiceFacts: return facts_result(self, facts_type, performance_service(self, query)) - - from tilefoundry.schedule.pipeline import PipelineFacts # noqa: PLC0415 - - if facts_type is PipelineFacts: - from tilefoundry.target.cuda.facts import pipeline_facts # noqa: PLC0415 - - return facts_result(self, facts_type, pipeline_facts(self, query)) - - from tilefoundry.schedule.partition import PartitionFacts # noqa: PLC0415 - - if facts_type is PartitionFacts: - from tilefoundry.target.cuda.facts import partition_facts # noqa: PLC0415 - - return facts_result(self, facts_type, partition_facts(self, query)) return super().get_facts(facts_type, query) - def get_scheduler(self, topology: str) -> Scheduler: - """Select a CUDA scheduler; subclasses inherit these solvers.""" - from tilefoundry.target.cuda.schedule import cuda_scheduler # noqa: PLC0415 - - scheduler = cuda_scheduler(topology) - if scheduler is not None: - return scheduler - return super().get_scheduler(topology) - def get_code_generator(self) -> CodeGenerator: from tilefoundry.codegen.cuda.module import ( # noqa: PLC0415 CUDA_CODE_GENERATOR, diff --git a/src/tilefoundry/target/services.py b/src/tilefoundry/target/services.py index a09d5e48..700f0edd 100644 --- a/src/tilefoundry/target/services.py +++ b/src/tilefoundry/target/services.py @@ -61,24 +61,6 @@ def get_checker(self) -> AnalysisChecker | None: return None -ScheduleCallable = Callable[ - ["Module", "Function", "Target", "Topology", object | None], - "SchedulePlan", -] - - -@dataclass(frozen=True) -class Scheduler: - """One scheduler service: the level it solves for, and the solve.""" - - topology: str - solve: ScheduleCallable - - def __post_init__(self) -> None: - if not self.topology: - raise ValueError("a schedule algorithm needs a non-empty topology name") - - @dataclass(frozen=True) class CodeGenerator: """One immutable Target-owned code-generation service.""" @@ -93,6 +75,4 @@ class CodeGenerator: "AnalysisChecker", "Analyzer", "CodeGenerator", - "ScheduleCallable", - "Scheduler", ] diff --git a/tests/analysis/test_analysis_invariants.py b/tests/analysis/test_analysis_invariants.py index 93cf72f1..91f27aac 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -77,8 +77,8 @@ def test_every_callable_op_states_its_coordinates_exactly_once() -> None: """One canonical relation per Op, enumerated rather than listed by hand. - Where an Op reads and writes is stated once, so type inference, the - polyhedral model, the loop footprint and the movement half read one answer. + Where an Op reads and writes is stated once, so type inference, the loop + footprint and the movement half read one answer. The set is taken from the Op registry itself, so an Op added to the surface joins this without anybody adding it here. The other dialect is not a Call target of these analyses and an Op a test registers is not part of the diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index d6b609db..5f94ac59 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -1,9 +1,8 @@ -"""Analysing and scheduling a function authored for a range of sizes. +"""Analysing a function authored for a range of sizes. -An analysis counts elements and holds them against a machine; a solver lays -work across a level by counting it. Neither has an answer for a dimension that -is still a range, so the size is stated at the call and the program that gets -measured is the one at that size. +An analysis counts elements and holds them against a machine. It has no answer +for a dimension that is still a range, so the size is stated at the call and the +program that gets measured is the one at that size. What the call accepts stays narrow: a function this Module owns. Choosing the size happens after that, so nothing here widens which programs a Module will @@ -48,7 +47,6 @@ Topology, ) from tilefoundry.ir.visitor import collect_exprs -from tilefoundry.schedule import ScheduleError, ScheduleOptions, schedule from tilefoundry.target import CudaTarget, PerformanceServiceFacts, ThroughputFacts CONTEXT = 32 @@ -57,7 +55,6 @@ INVENTORY = [pytest.param(case, id=case.id) for case in placed_cases()] -SOLVER = ScheduleOptions(timeout_seconds=60, workers=4, random_seed=0, stop_at_first_solution=True) def _aimed(): @@ -369,21 +366,12 @@ def test_a_size_states_nothing_about_a_function_from_elsewhere() -> None: def test_the_entry_at_a_chosen_size_is_still_the_entry() -> None: - """The device-wide solver admits only the entry, and it decides that by name. + """Choosing a size does not rename the entry. - The device-wide solver admits only the entry, and it decides that by - name: a function specialised from the entry is a different object and the - same program. + A function specialised from the entry is a different object and the same + program, so anything that identifies the entry by name still finds it. """ module = _aimed() variant = variant_for(module.entry_function(), DIMS) assert variant.name == module.entry_function().name - with pytest.raises(ScheduleError, match="requires the module entry"): - schedule( - module, - module.lookup("_ctx_partials"), - topology="cta", - options=SOLVER, - dims=DIMS, - ) diff --git a/tests/analysis/test_poly_dynamic_shape.py b/tests/analysis/test_poly_dynamic_shape.py deleted file mode 100644 index 3a7c5efa..00000000 --- a/tests/analysis/test_poly_dynamic_shape.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Pin ``extract`` behavior for dynamic ``DimVar`` shapes. - -A ``DimVar`` flows through as a bounded isl parameter and resolves to its -``ShapeDim`` in ``TileGraph.params``. Emitted loops name that parameter rather -than inventing a fixed trip count. ``test_analysis_invariants.py`` pins the -static counterpart. -""" - -from __future__ import annotations - -import re - -import isl - -from tests.fixtures.shapes.matmul_programs import DYNAMIC_M as SEQ -from tests.fixtures.shapes.matmul_programs import dynamic_bf16_gemm as dyn_matmul -from tilefoundry.analysis import TileGraph, extract -from tilefoundry.schedule.kernel_schedule import build_schedule_tree -from tilefoundry.schedule.render import emit_scaffold - - -def test_dynamic_matmul_extract_params_and_domain(): - """Test dynamic matmul extract params and domain. - - A DimVar M axis extracts a parametrised ``[seq]->{...}`` domain - (``0 <= i < seq``, straight from ``to_domain`` -- no tiling), resolves - ``TileGraph.params['seq']`` back to the exact ``DimVar``, and the M - axis is still bounded (``dim_max_val`` a finite 126, not ``infty``, - since ``seq``'s own half-open range ``[1, 128)`` tops out at 127 -- - an unbounded ``DimVar`` is not constructible in the first place). - ``build_schedule_tree()`` stays parametrised too. - """ - tg = extract(dyn_matmul) - assert isinstance(tg, TileGraph) - - assert tg.params == {"seq": SEQ} - assert tg.domain.space().dim(isl.dim_type.PARAM) == 1 - assert "[seq]" in str(tg.domain) - assert "x[" in str(tg.reads) and "w[" in str(tg.reads) - assert "h[" in str(tg.writes) - - sets: list = [] - tg.domain.foreach_set(sets.append) - assert len(sets) == 1 - (mm_set,) = sets - assert mm_set.get_tuple_name() == "MM" - - assert int(mm_set.dim_min_val(0).num_si()) == 0 - assert int(mm_set.dim_max_val(0).num_si()) == 126 - - assert int(mm_set.dim_max_val(1).num_si()) == 1 - assert int(mm_set.dim_max_val(2).num_si()) == 3 - - tree = build_schedule_tree(tg) - assert "[seq]" in str(tree) - - -def test_dynamic_matmul_end_to_end_emits_symbolic_loop(): - """Extract -> build_schedule_tree -> emit_scaffold. - - Extract -> build_schedule_tree -> emit_scaffold: the M loop's upper bound - names the isl parameter directly, never a fixed integer trip count. - """ - tg = extract(dyn_matmul) - tree = build_schedule_tree(tg) - skeleton, _swimlane, contracts = emit_scaffold(tg, tree, {}) - - print("\n=== dynamic matmul skeleton ===") - print(skeleton.text) - - assert "HOLE_MM" in skeleton.text - assert len(contracts) == 1 - - m = re.search(r"for \(int c0 = 0; c0 <=? ([^;]+); c0 \+= 1\)", skeleton.text) - assert m is not None, skeleton.text - bound = m.group(1) - assert "seq" in bound diff --git a/tests/analysis/test_poly_grid_region.py b/tests/analysis/test_poly_grid_region.py deleted file mode 100644 index 8bcb5fa5..00000000 --- a/tests/analysis/test_poly_grid_region.py +++ /dev/null @@ -1,274 +0,0 @@ -"""Pin ``extract`` behavior for authored ``GridRegionExpr`` loops. - -Loop axes prefix enclosed statements only; nested loops order outermost first. -A carried buffer creates a distance-one dependence that scheduling must order. -``DimVar`` extents become isl parameters. Data-dependent selections fail closed -instead of claiming a known slice. Small loops keep expected delta sets -hand-transcribable; the corpus Analyze witness covers real tiled kernels. -""" - -from __future__ import annotations - -import isl -import pytest - -from tests.fixtures.shapes.window_programs import ( - WINDOW_SEQ as SEQ, -) -from tests.fixtures.shapes.window_programs import ( - dynamic_tile_window_add, - moved_tile_window_add, - tile_window_add, - unspecialized_tile_window_add, -) -from tilefoundry import func -from tilefoundry.analysis import extract -from tilefoundry.analysis.poly import ExtractError -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import * # noqa: F401,F403 -- op names resolved dynamically -from tilefoundry.schedule.kernel_schedule import build_schedule_tree - - -@func -def carry_loop(x: Tensor[(8, 4), "f32"], y: Tensor[(8, 4), "f32"]) -> Tensor[(8, 4), "f32"]: - o = mul(x, y) - for i in range(6): - o = add(o, x) - return o - - -@func -def nested_carry(x: Tensor[(8, 4), "f32"], y: Tensor[(8, 4), "f32"]) -> Tensor[(8, 4), "f32"]: - o = mul(x, y) - for r in range(4): - for c in range(2): - o = add(o, x) - return o - - -@func -def dyn_carry(x: Tensor[(SEQ, 4), "f32"], y: Tensor[(SEQ, 4), "f32"]) -> Tensor[(SEQ, 4), "f32"]: - o = mul(x, y) - for i in range(SEQ): - o = add(o, x) - return o - - -@func -def data_index_select( - x: Tensor[(8, 4), "f32"], - y: Tensor[(4,), "f32"], - idx: Tensor[(), "i32"], -) -> Tensor[(4,), "f32"]: - o = mul(y, y) - for i in range(8): - selected = index_select(x, reshape(idx, new_shape=(1,)), dim=0) - o = add(o, reshape(selected, new_shape=(4,))) - return o - - -def _domains(tg) -> dict[str, "isl.set"]: - sets: list["isl.set"] = [] - tg.domain.foreach_set(sets.append) - return {s.get_tuple_name(): s for s in sets} - - -def _extent(s: "isl.set", pos: int) -> tuple[int, int]: - return int(s.dim_min_val(pos).num_si()), int(s.dim_max_val(pos).num_si()) - - -def _maps(um: "isl.union_map") -> list["isl.map"]: - out: list["isl.map"] = [] - um.foreach_map(out.append) - return out - - -def _writer_of(tg, buffer: str) -> str: - """The statement that writes ``buffer``.""" - names = { - m.get_tuple_name(isl.dim_type.IN) - for m in _maps(tg.writes) - if m.get_tuple_name(isl.dim_type.OUT) == buffer - } - assert len(names) == 1, f"{buffer}: expected one writer, got {sorted(names)}" - return next(iter(names)) - - -def _self_deltas(tg, statement: str) -> "isl.set": - """Self deltas. - - ``statement``'s own dependence distances, tuple name dropped so the - expected set reads as plain coordinates. - """ - own = _domains(tg)[statement].to_union_set() - pieces: list["isl.set"] = [] - tg.deps.intersect_domain(own).intersect_range(own).deltas().foreach_set(pieces.append) - assert len(pieces) == 1, f"{statement}: expected one delta piece, got {pieces}" - return pieces[0].reset_tuple_id() - - -def _violations(tg) -> "isl.union_set": - """The dependence deltas ``build_schedule_tree`` does not order strictly.""" - sched = build_schedule_tree(tg).get_map() - timed = tg.deps.apply_domain(sched).apply_range(sched) - assert not timed.is_empty(), "every dependence must survive into time space" - return _lex_nonpositive(timed.deltas()) - - -def _lex_nonpositive(deltas: "isl.union_set") -> "isl.union_set": - out = isl.union_set("{}") - pieces: list = [] - deltas.foreach_set(pieces.append) - for piece in pieces: - rank = piece.dim(isl.dim_type.SET) - dims = ", ".join(f"d{i}" for i in range(rank)) - positive = isl.set( - "{ " - + "; ".join( - f"[{dims}] : " - + " and ".join([*(f"d{i} = 0" for i in range(index)), f"d{index} > 0"]) - for index in range(rank) - ) - + " }" - ) - out = out.union(piece.subtract(positive)) - return out - - -def test_carried_arg_is_a_distance_one_dependence_on_the_loop_axis(): - """The carry is one buffer: the ``add`` reads and writes ``o``. - - The ``add`` reads and writes ``o``, so isl reports ``[1, 0, 0]``: iteration - ``i`` needs what ``i - 1`` wrote. The loop-invariant ``mul`` keeps its 2-d - domain while ``add`` gains a leading loop axis. The schedule must order that - carry strictly. - """ - tg = extract(carry_loop) - doms = _domains(tg) - assert sorted(doms) == ["Binary0", "Binary1"] - - assert doms["Binary0"].dim(isl.dim_type.SET) == 2 - assert _extent(doms["Binary0"], 0) == (0, 7) - assert doms["Binary1"].dim(isl.dim_type.SET) == 3 - assert _extent(doms["Binary1"], 0) == (0, 5) - assert _extent(doms["Binary1"], 1) == (0, 7) - assert _extent(doms["Binary1"], 2) == (0, 3) - - assert _writer_of(tg, "o") == "Binary1" - assert "-> o[" in str(tg.reads), tg.reads - assert _self_deltas(tg, "Binary1").is_equal(isl.set("{ [1, 0, 0] }")) - - assert tg.parallel_dims["Binary1"] == (False, True, True) - assert tg.parallel_dims["Binary0"] == (True, True) - - assert _violations(tg).is_empty() - - -def test_nested_loops_contribute_one_dimension_each(): - """Two axes, innermost last: the carry advances by one inner step. - - Two axes, innermost last: the carry advances by one inner step, and - wraps to the next outer step from the last inner one. - """ - tg = extract(nested_carry) - dom = _domains(tg)["Binary1"] - assert dom.dim(isl.dim_type.SET) == 4 - assert _extent(dom, 0) == (0, 3) - assert _extent(dom, 1) == (0, 1) - assert _self_deltas(tg, "Binary1").is_equal(isl.set("{ [0, 1, 0, 0]; [1, -1, 0, 0] }")) - assert tg.parallel_dims["Binary1"] == (False, False, True, True) - - -def test_dynamic_extent_becomes_an_isl_parameter(): - """A loop whose trip count is only known at the call. - - A loop whose trip count is only known at the call: the axis is bounded by - the parameter itself, and the carry distance is still one step of it. - """ - tg = extract(dyn_carry) - assert tg.params == {"seq": SEQ} - dom = _domains(tg)["Binary1"] - assert "0 <= i0 < seq" in str(dom), dom - assert _self_deltas(tg, "Binary1").is_equal(isl.set("[seq] -> { [1, 0, 0] : 4 <= seq <= 63 }")) - - -def test_windowed_loop_analyzes_only_full_tiles_and_offsets_its_read(): - tg = extract(tile_window_add) - domain = _domains(tg)["Binary1"] - - assert domain.is_equal( - isl.set("{ Binary1[i, r, c] : 0 <= i <= 4 and i mod 4 = 0 " - "and 0 <= r < 4 and 0 <= c < 4 }") - ) - source_reads = tg.reads.intersect_range( - isl.union_set("{ x[r, c] : 0 <= r < 10 and 0 <= c < 4 }") - ) - assert source_reads.is_equal( - isl.union_map( - "{ Binary1[i, r, c] -> x[i + r, c] : " - "0 <= i <= 4 and i mod 4 = 0 and 0 <= r < 4 and 0 <= c < 4 }" - ) - ) - - -def test_a_moved_window_carries_its_offset_into_the_access_map(): - """A window moved by a compile-time offset reads the same loop dimension. - - A window moved by a compile-time offset reads the same loop dimension shifted - by that offset, so the offset belongs in the access map rather than in a - separate statement -- the move is an address, not a computation. - """ - tg = extract(moved_tile_window_add) - domain = _domains(tg)["Binary1"] - - assert domain.is_equal( - isl.set("{ Binary1[i, r, c] : 0 <= i <= 3 and i mod 3 = 0 " - "and 0 <= r < 3 and 0 <= c < 4 }") - ) - source_reads = tg.reads.intersect_range( - isl.union_set("{ x[r, c] : 0 <= r < 12 and 0 <= c < 4 }") - ) - assert source_reads.is_equal( - isl.union_map( - "{ Binary1[i, r, c] -> x[i + r + 6, c] : " - "0 <= i <= 3 and i mod 3 = 0 and 0 <= r < 3 and 0 <= c < 4 }" - ) - ) - - -def test_symbolic_extent_keeps_only_parameterized_full_windows(): - domain = _domains(extract(dynamic_tile_window_add))["Binary1"] - - assert domain.is_equal( - isl.set( - "[seq] -> { Binary1[i, r, c] : 4 <= seq < 64 and 0 <= i " - "and i + 4 <= seq and i mod 4 = 0 and 0 <= r < 4 and 0 <= c < 4 }" - ) - ) - - -def test_unspecialized_window_step_fails_closed(): - with pytest.raises(ExtractError, match="loop step.*not a static int"): - extract(unspecialized_tile_window_add) - - -def test_data_dependent_index_select_reads_every_row_it_could_name(): - """A gather whose index is a value reads every row that value could name. - - No relation holds the deciding element, so the coordinate it lands on is not - one extraction can state. What it states instead is every row the axis - could legally name: more dependences than the program has, which is the safe - direction, and the same answer every other reader of that relation gets. - """ - tg = extract(data_index_select) - - gathered = tg.reads.intersect( - isl.union_map("{ Binary1[i, d] -> x[r, d] }") - ) - assert not gathered.is_empty(), "the gather's read was dropped rather than widened" - assert gathered.is_equal( - isl.union_map( - "{ Binary1[i, d] -> x[r, d] : 0 <= i <= 7 and 0 <= d <= 3 " - "and 0 <= r <= 7 }" - ) - ), "every row of the table, for every coordinate of the result" diff --git a/tests/analysis/test_poly_nested_call.py b/tests/analysis/test_poly_nested_call.py deleted file mode 100644 index c3cfff5a..00000000 --- a/tests/analysis/test_poly_nested_call.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Pin the nested ``@func`` call shapes that ``extract`` refuses. - -Normal calls bind arguments, recurse, and prefix contributed statements and -buffers by callee and call-site index. Failures must name the stopped callee. -Self-recursion and arity mismatches require hand-built HIR because the authoring -surface rejects forward references and bad arity. A dispatch prototype is -authorable normally, but extraction cannot resolve its ``pass`` body statically. -""" - -from __future__ import annotations - -import pytest - -from tilefoundry import func -from tilefoundry.analysis import ExtractError, extract -from tilefoundry.dsl import Tensor -from tilefoundry.ir.core import Call, Var -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.types import DType, make_tensor_type - - -def test_self_recursive_call_raises_naming_the_callee(): - """A hand-built self-recursive function reports its callee name.""" - ty = make_tensor_type((2, 2), DType.f32) - x = Var(type=ty, name="x") - stub = Function.build(name="loopy", params=(x,), body=x, return_type=ty) - object.__setattr__(stub, "body", Call(type=ty, target=stub, args=(x,))) - - with pytest.raises(ExtractError, match="loopy"): - extract(stub) - - -@func -def _dispatch_prototype_helper(x: Tensor[(4, 4), "f32"]) -> Tensor[(4, 4), "f32"]: - pass - - -@func -def _call_dispatch_prototype(x: Tensor[(4, 4), "f32"]) -> Tensor[(4, 4), "f32"]: - return _dispatch_prototype_helper(x) - - -def test_dispatch_prototype_call_raises_naming_the_callee(): - """A callee with no body cannot be penetrated statically. - - A dispatch prototype is unresolved - without a concrete runtime shape) cannot be penetrated statically. - """ - with pytest.raises(ExtractError, match="_dispatch_prototype_helper"): - extract(_call_dispatch_prototype) - - -def test_arity_mismatch_call_raises_naming_the_callee(): - """A call passing fewer arguments than declared reports its callee.""" - ty = make_tensor_type((2, 2), DType.f32) - x = Var(type=ty, name="x") - y = Var(type=ty, name="y") - callee = Function.build(name="needs_two", params=(x, y), body=x, return_type=ty) - only_arg = Var(type=ty, name="only_one") - bad_call = Call(type=ty, target=callee, args=(only_arg,)) - caller = Function.build(name="caller", params=(only_arg,), body=bad_call, return_type=ty) - - with pytest.raises(ExtractError, match="needs_two"): - extract(caller) diff --git a/tests/analysis/test_poly_reshape.py b/tests/analysis/test_poly_reshape.py deleted file mode 100644 index ef417db2..00000000 --- a/tests/analysis/test_poly_reshape.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Pin ``extract``'s zero-op view fold for ``Reshape``. - -Reshape contributes no statement or buffer: consumers resolve through to the -source and fold the coordinates through the Op's own registered access -relation, which is the only place the renaming is stated. Distinct toy -dimensions mirror decoder head split/merge shapes while keeping exact -per-element maps readable. -""" - -from __future__ import annotations - -import isl -import pytest - -from tilefoundry import func -from tilefoundry.analysis import ExtractError, TileGraph, extract -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import * # noqa: F401,F403 -- reshape/sigmoid resolved dynamically - -B, S, H, D = 1, 4, 3, 8 -HD = H * D - - -@func -def split_then_sigmoid(x: Tensor[(B, S, HD), "f32"]) -> Tensor[(B, S, H, D), "f32"]: - y = reshape(x, new_shape=(B, S, H, D)) - z = sigmoid(y) - return z - - -@func -def merge_then_sigmoid(v: Tensor[(B, S, H, D), "f32"]) -> Tensor[(B, S, HD), "f32"]: - w = reshape(v, new_shape=(B, S, HD)) - u = sigmoid(w) - return u - - -@func -def bare_reshape_return(x4: Tensor[(B, S, HD), "f32"]) -> Tensor[(B, S, H, D), "f32"]: - y5 = reshape(x4, new_shape=(B, S, H, D)) - return y5 - - -def test_a_reshape_is_a_view_its_consumer_reads_through(): - """Both directions, because they are not the same construction. - - Splitting a merged head axis composes the consumer read through to the source - with ``D*h + d``. Merging reconstructs ``(h, d)`` with div/mod. In both - directions reshape contributes no statement and its SSA name is not a buffer. - """ - split = extract(split_then_sigmoid) - assert isinstance(split, TileGraph) - assert [type(u.op.target).__name__ for u in split.units] == ["Sigmoid"] - bounds = f"0<=b<{B} and 0<=s<{S} and 0<=h<{H} and 0<=d<{D}" - assert split.domain.is_equal(isl.union_set(f"{{ Sigmoid[b,s,h,d] : {bounds} }}")) - assert split.reads.is_equal( - isl.union_map(f"{{ Sigmoid[b,s,h,d] -> x[b,s,{D}*h+d] : {bounds} }}") - ) - assert split.writes.is_equal(isl.union_map(f"{{ Sigmoid[b,s,h,d] -> z[b,s,h,d] : {bounds} }}")) - assert "y[" not in str(split.reads) and "y[" not in str(split.writes) - assert split.deps.is_empty() - - merged = extract(merge_then_sigmoid) - assert [type(u.op.target).__name__ for u in merged.units] == ["Sigmoid"] - assert merged.reads.is_equal( - isl.union_map( - f"{{ Sigmoid[b,s,e] -> v[b,s,floor(e/{D}),e mod {D}] : " - f"0<=b<{B} and 0<=s<{S} and 0<=e<{HD} }}" - ) - ) - assert "w[" not in str(merged.reads) - - -def test_boundary_reshape_with_no_consumer_fails_closed(): - """Test boundary reshape with no consumer fails closed. - - A body that is *nothing but* a reshape (``return reshape(x, ...)``) has no - compute op left once the reshape folds away -- ``extract`` fails closed with - its empty-body error rather than fabricating a copy statement. - """ - with pytest.raises(ExtractError, match="no compute ops to extract"): - extract(bare_reshape_return) diff --git a/tests/analysis/test_poly_rmsnorm.py b/tests/analysis/test_poly_rmsnorm.py deleted file mode 100644 index dd9b439d..00000000 --- a/tests/analysis/test_poly_rmsnorm.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Pin ``extract`` behavior for the registered ``RMSNorm`` relation. - -Its domain contains batch axes only; the reduced last axis is existential in -read/write maps, including the weight read. ``local_type_of`` resolves -sharding before the relation sees the type. ``test_analysis_invariants.py`` -pins the corresponding ``SoftMax`` shape. -""" - -from __future__ import annotations - -import isl -import pytest - -from tilefoundry import func -from tilefoundry.analysis import ExtractError, TileGraph, extract -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import * # noqa: F401,F403 -- rms_norm resolved dynamically -from tilefoundry.ir.types import local_type_of, make_shard_tensor_type, make_tensor_type -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Split, Topology -from tilefoundry.ir.visitor import collect_exprs - -_MESH = Mesh((Topology("gpu", 2),), Layout((2,), (1,)), names=("a",)) - - -@func -def rmsnorm_only(x: Tensor[(2, 64), "f32"], weight: Tensor[(64,), "f32"]) -> Tensor[(2, 64), "f32"]: - y = rms_norm(x, weight) - return y - - -@func -def local_type_boundary(x: Tensor[(7,), "f32"]) -> Tensor[(7,), "f32"]: - view = reshape(x, new_shape=(7,)) - return sigmoid(view) - - -def test_extract_rmsnorm_single_statement(): - """``y = rms_norm(x, weight)`` extracts to one statement. - - ``y = rms_norm(x, weight)`` extracts to one statement: domain = - the batch axis only, reads/writes both range over the whole - (existentially-quantified) row -- exactly like ``SoftMax``'s shape. - """ - tg = extract(rmsnorm_only) - assert isinstance(tg, TileGraph) - assert len(tg.units) == 1 - assert tg.units[0].name == "RN" - assert type(tg.units[0].op.target).__name__ == "RMSNorm" - - print("\n=== rmsnorm: domain ===") - print(tg.domain) - print("=== rmsnorm: reads ===") - print(tg.reads) - print("=== rmsnorm: writes ===") - print(tg.writes) - - assert tg.domain.is_equal(isl.union_set("{ RN[i] : 0 <= i < 2 }")) - expected_reads = ( - isl.union_map("{}") - .union(isl.map("{ RN[i] -> x[i, j] : 0 <= i < 2 and 0 <= j < 64 }")) - .union(isl.map("{ RN[i] -> weight[j] : 0 <= i < 2 and 0 <= j < 64 }")) - ) - expected_writes = isl.union_map("{ RN[i] -> y[i, j] : 0 <= i < 2 and 0 <= j < 64 }") - assert tg.reads.is_equal(expected_reads) - assert tg.writes.is_equal(expected_writes) - - assert tg.deps.is_empty() - - -def test_local_type_divides_the_split_axis_and_keeps_tensor_rank(): - """A ``Split`` axis contributes its per-shard extent. - - Canonical layouts may outrank tensors: split ``(8, 16)`` becomes layout - ``(2, 4, 16)``. Localization must use ``split_target_axes`` to divide the - tensor axis while preserving rank, or reader and writer enter different isl - spaces and lose dependencies. The test also covers a trailing split and the - identity behavior for an unsharded type. - """ - x = make_shard_tensor_type((8, 16), mesh=_MESH, attrs=(Split(0),)) - assert len(x.layout.layout.shape) == 3 - - local = local_type_of(x) - - assert local.shape == (4, 16) - assert len(local.shape) == len(x.shape) - - trailing = make_shard_tensor_type((8, 16), mesh=_MESH, attrs=(Split(1),)) - assert local_type_of(trailing).shape == (8, 8) - - plain = make_tensor_type((8, 16)) - assert local_type_of(plain) is plain - - -def _set_function_type(function, type_): - function.params[0].type = type_ - function.return_type = type_ - for expr in collect_exprs(function.body): - if hasattr(expr, "type"): - expr.type = type_ - - -def test_extract_rejects_a_dynamic_split_extent_with_context(): - dynamic = make_shard_tensor_type( - (DimVar("local_type_dynamic", 1, 65),), mesh=_MESH, attrs=(Split(0),) - ) - _set_function_type(local_type_boundary, dynamic) - - with pytest.raises( - ExtractError, - match=r"tensor axis 0.*extent .*not a static int", - ): - extract(local_type_boundary) - - -def test_extract_rejects_a_non_divisible_split_extent_with_context(): - layout = ShardLayout(layout=Layout((2, 7), (7, 1)), attrs=(Split(0),), mesh=_MESH) - non_divisible = make_tensor_type((7,), layout=layout) - _set_function_type(local_type_boundary, non_divisible) - - with pytest.raises( - ExtractError, - match=r"tensor axis 0.*extent 7.*mesh extent 2", - ): - extract(local_type_boundary) diff --git a/tests/fixtures/shapes/matmul_programs.py b/tests/fixtures/shapes/matmul_programs.py index 5ebf23d3..8ad558f0 100644 --- a/tests/fixtures/shapes/matmul_programs.py +++ b/tests/fixtures/shapes/matmul_programs.py @@ -1,4 +1,4 @@ -"""Matmul programs spanning analysis, CUDA scheduling, and AMX target boundaries.""" +"""Matmul programs spanning analysis, CUDA, and AMX target boundaries.""" from __future__ import annotations @@ -48,13 +48,6 @@ def cuda_odd_m_bf16_gemm( return tf.matmul(x, w) -@func -def scheduling_gemm( - x: Tensor[(64, 128), "f32"], w: Tensor[(128, 64), "f32"] -) -> Tensor[(64, 64), "f32"]: - return tf.matmul(x, w) - - @func def dynamic_bf16_gemm( x: Tensor[(DYNAMIC_M, 4), "bf16"], w: Tensor[(4, 2), "bf16"] diff --git a/tests/installed/models/smoke_kimi_linear_48b_a3b.py b/tests/installed/models/smoke_kimi_linear_48b_a3b.py index 56acf836..da657674 100644 --- a/tests/installed/models/smoke_kimi_linear_48b_a3b.py +++ b/tests/installed/models/smoke_kimi_linear_48b_a3b.py @@ -1,8 +1,8 @@ """Kimi-Linear-48B-A3B, as the installation ships it, asked through the commands. Three Modules of kernels are reached from one root, so this model states three -schedule cases where most state one, and its selectors are dotted paths through -the tree. It ships no ``hf_alias.py``: it is not loaded from a raw published +cases where most state one, and its selectors are dotted paths through the +tree. It ships no ``hf_alias.py``: it is not loaded from a raw published checkpoint. """ diff --git a/tests/installed/models/smoke_qwen3_1_7b.py b/tests/installed/models/smoke_qwen3_1_7b.py index 41073280..b5273938 100644 --- a/tests/installed/models/smoke_qwen3_1_7b.py +++ b/tests/installed/models/smoke_qwen3_1_7b.py @@ -23,7 +23,7 @@ ANALYSED = [ pytest.param(case, selected, id=selected.id) for case in CASES for selected in case.analyze ] -JSON_CASES = [pytest.param(case, case.schedule[0], id=case.id) for case in CASES] +JSON_CASES = [pytest.param(case, case.analyze[0], id=case.id) for case in CASES] SIZED = [pytest.param(case, sized, id=sized.id) for case in CASES for sized in case.sized] diff --git a/tests/installed/smoke_target/vendor_npu/__init__.py b/tests/installed/smoke_target/vendor_npu/__init__.py index aea387c5..ae3e534f 100644 --- a/tests/installed/smoke_target/vendor_npu/__init__.py +++ b/tests/installed/smoke_target/vendor_npu/__init__.py @@ -2,20 +2,15 @@ from __future__ import annotations -import json -import os from dataclasses import dataclass -from pathlib import Path from typing import ClassVar from tilefoundry import DType from tilefoundry.analysis import ExplicitMemoryLevelFacts -from tilefoundry.schedule import SchedulePlan from tilefoundry.target import ( MemoryHierarchyFacts, ParallelCapacityFacts, PerformanceServiceFacts, - Scheduler, Target, ThroughputFacts, TopologyLimitFacts, @@ -24,30 +19,6 @@ ) -@dataclass(frozen=True) -class VendorNpuPlan(SchedulePlan): - topology: str - extent: int - - def verify(self, module, function, topology) -> None: - if topology.name != self.topology or topology.size != self.extent: - raise ValueError("vendor NPU plan does not match its topology") - - def to_json(self) -> str: - return json.dumps({"topology": self.topology, "extent": self.extent}) - - def render(self) -> str: - return f"vendor NPU schedule: {self.extent} {self.topology}" - - -def _schedule_vendor_npu(module, function, target, topology, options) -> SchedulePlan: - marker = os.environ.get("TF_VENDOR_NPU_SCHEDULER_CALLS") - if marker is not None: - with Path(marker).open("a", encoding="utf-8") as calls: - calls.write(f"{topology.name}\n") - return VendorNpuPlan(topology.name, topology.size) - - @register_target @dataclass(frozen=True) class VendorNpuTarget(Target): @@ -94,8 +65,3 @@ def get_facts(self, facts_type: type, query: object | None = None): else: return super().get_facts(facts_type, query) return facts_result(self, facts_type, value) - - def get_scheduler(self, topology: str) -> Scheduler: - if topology == "core": - return Scheduler("core", _schedule_vendor_npu) - return super().get_scheduler(topology) diff --git a/tests/integration/models/deepseek_v4_flash/test_case.py b/tests/integration/models/deepseek_v4_flash/test_case.py index d7452d92..f33aa214 100644 --- a/tests/integration/models/deepseek_v4_flash/test_case.py +++ b/tests/integration/models/deepseek_v4_flash/test_case.py @@ -17,15 +17,13 @@ def test_the_case_selects_every_function_the_description_defines(): - """Analyze has no reason to leave a function out, and schedule admits only the entry function. + """Analyze has no reason to leave a function out. - Analyze has no reason to leave a function out, and schedule admits only - the entry function -- so the other one is untested rather than blocked. + Every function this Module defines is selected, so the model's own inventory + and the case's selection agree and nothing escapes the report. """ module = CASE.build() assert CASE.untested("analyze", module) == () - assert CASE.selected("schedule") == (module.entry_function().name,) - assert CASE.untested("schedule", module) == ("mla_kv_update",) def test_the_context_lengths_the_case_names_are_ones_the_model_has(): diff --git a/tests/ir/types/test_mesh.py b/tests/ir/types/test_mesh.py index 6a031b1e..f0bc698c 100644 --- a/tests/ir/types/test_mesh.py +++ b/tests/ir/types/test_mesh.py @@ -2,11 +2,9 @@ import pytest -from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard import ( Layout, Mesh, - Split, Topology, make_mesh, product, @@ -16,8 +14,6 @@ mesh_scope_matches_required_scope, states_consistent_positions, ) -from tilefoundry.ir.types.shard.shard_layout import ShardLayout -from tilefoundry.schedule.partition.problem import _placement_relation def test_mesh_position_consistency_is_an_explicit_predicate() -> None: @@ -67,17 +63,9 @@ def test_mesh_slice_keeps_the_parent_topologies() -> None: assert sliced.layout.shape == (1, 32) -def test_mesh_value_equality_is_usable_by_partition() -> None: +def test_mesh_value_equality_is_by_value() -> None: left = make_mesh((8,), topology="thread") right = make_mesh((8,), topology="thread") - layout = Layout(shape=(8,), strides=(1,)) - consumer = TensorType( - shape=(8,), - dtype=DType.f32, - layout=ShardLayout(layout, (Split(0),), right), - storage="rmem", - ) assert left == right assert hash(left) == hash(right) - assert _placement_relation(consumer, left) == "SAME_INTERVAL" diff --git a/tests/models/corpus.py b/tests/models/corpus.py index c58b2beb..822a64a7 100644 --- a/tests/models/corpus.py +++ b/tests/models/corpus.py @@ -1,7 +1,7 @@ """The shared vocabulary every model-driven test is written against. One model is described once, as a `ModelCase`, and every kind of test reads that -one description: the reference run, the analyses, the schedules, the CLI, and the +one description: the reference run, the analyses, the CLI, and the end-to-end witnesses. Nothing here copies a model into a smaller graph for one subsystem's convenience, because a result measured on a copy says nothing about the program a user would actually hand us. @@ -206,7 +206,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class FunctionCase: - """One function of one model, selected to be analysed or scheduled. + """One function of one model, selected to be analysed. `selector` names it root-relative to the case's Module: a bare function name for a Module that owns the kernel itself, or a dotted path through the child @@ -282,7 +282,6 @@ class ModelCase: prototype: Module reference: ReferenceCase | None = None analyze: tuple[FunctionCase, ...] = () - schedule: tuple[FunctionCase, ...] = () sized: tuple[SizedCase, ...] = () model: str = "" scope: str = "" @@ -332,12 +331,12 @@ def inventory(self, module: Module | None = None) -> tuple[str, ...]: built = self.build() if module is None else module return tuple(selector for selector, _ in function_selectors(built)) - def selected(self, kind: Literal["analyze", "schedule"]) -> tuple[str, ...]: - cases = self.analyze if kind == "analyze" else self.schedule + def selected(self, kind: Literal["analyze"]) -> tuple[str, ...]: + cases = self.analyze return tuple(dict.fromkeys(case.selector for case in cases)) def untested( - self, kind: Literal["analyze", "schedule"], module: Module | None = None + self, kind: Literal["analyze"], module: Module | None = None ) -> tuple[str, ...]: """The model's own functions that no case of *kind* selected.""" chosen = set(self.selected(kind)) diff --git a/tests/models/deepseek_v4_flash/case.py b/tests/models/deepseek_v4_flash/case.py index 8c444761..add2f41f 100644 --- a/tests/models/deepseek_v4_flash/case.py +++ b/tests/models/deepseek_v4_flash/case.py @@ -1,5 +1,5 @@ -"""DeepSeek-V4-Flash as one corpus case: what runs, what is analysed, what is -scheduled, and at what context length. +"""DeepSeek-V4-Flash as one corpus case: what runs, what is analysed, and at +what context length. The boundary this model states is the sliding-window MLA attention submodule of its first sliding layer, named as that -- ``DeepseekV4Flash.layer0.attention``, @@ -8,8 +8,8 @@ embedding, 43 decoder layers, MoE, final norm, head -- is a real end-to-end path with a real checkpoint pipeline behind it (`test_causal_lm_e2e.py`), and it is a tree of Modules walked by orchestration methods rather than a Module of Functions: -a `ModelCase` names one Module and analysis and scheduling select Functions of -that one Module, so naming the root would put every kernel that does arithmetic +a `ModelCase` names one Module and analysis selects Functions of that one +Module, so naming the root would put every kernel that does arithmetic out of reach and report the model as three norms and an add. `ctx_len` is bounded by the window rather than by the position embedding, so the @@ -72,14 +72,6 @@ dims=ANALYZED_AT, ), ), - schedule=( - FunctionCase( - id="deepseek_v4_flash/schedule/mla_attend", - selector="mla_attend", - topology="cta", - dims=ANALYZED_AT, - ), - ), sized=( SizedCase( id="deepseek_v4_flash/sized/mla_attend", diff --git a/tests/models/deepseek_v4_flash/model.py b/tests/models/deepseek_v4_flash/model.py index 25153104..38aa85fc 100644 --- a/tests/models/deepseek_v4_flash/model.py +++ b/tests/models/deepseek_v4_flash/model.py @@ -1,8 +1,8 @@ """DeepSeek-V4-Flash as IR Modules: the sliding-window MLA attention submodule, the two MoE blocks, the decoder layer that composes them, and the causal-LM root. -One authored source at two configurations: the corpus analyses and schedules the -attention submodule at the real checkpoint's dimensions, and +One authored source at two configurations: the corpus analyses the attention +submodule at the real checkpoint's dimensions, and `test_causal_lm_e2e.py` runs the whole tree at a shape small enough to be affordable. The class bodies sit inside functions rather than at file scope, so each is evaluated once per configuration, and ``build_deepseek_v4_flash`` is how a diff --git a/tests/models/fixtures.py b/tests/models/fixtures.py index c73809ae..b12d4c80 100644 --- a/tests/models/fixtures.py +++ b/tests/models/fixtures.py @@ -36,7 +36,7 @@ def _parallel_units(target: Target, level: str) -> int: def h200_sxm(*, threads_per_cta: int = 512) -> TargetFixture: - """One H200 SXM, at both levels CUDA schedules. + """One H200 SXM, at both of its topology levels. The CTA extent is the device's own SM count, so the fixture divides work over exactly as much machine as the documents say exists. The thread extent is a diff --git a/tests/models/gemma2_2b/case.py b/tests/models/gemma2_2b/case.py index 3018aeb5..b5e2c65b 100644 --- a/tests/models/gemma2_2b/case.py +++ b/tests/models/gemma2_2b/case.py @@ -8,29 +8,22 @@ Analyse selects every function the model defines -- what is not selected here is untested, and the report derives that from the model's own function inventory. -Schedule admits only the module entry function, because the device-wide partition -algorithm decides the launch and a leaf is not something it has an answer to. -``sized`` is a third question: whether the model can be analysed at a context +``sized`` is a second question: whether the model can be analysed at a context length of the caller's choosing. -Every case below is declared as passing, and two of the analyse cases, the -schedule case and the sized case do not pass yet -- not for a reason in this -package. ``Gelu`` has an evaluation handler and no *cost* evaluator, so anything -that costs ``mlp`` or ``decoder_layer`` stops at -``no cost evaluator registered for Gelu``: all four analysis families, the -partition, and the sized question alike. It is one registration next to +Every case below is declared as passing, and two of the analyse cases and the +sized case do not pass yet -- not for a reason in this package. ``Gelu`` has an +evaluation handler and no *cost* evaluator, so anything that costs ``mlp`` or +``decoder_layer`` stops at ``no cost evaluator registered for Gelu``: all four +analysis families and the sized question alike. It is one registration next to ``Sigmoid`` / ``Softplus`` / ``Tanh`` / ``ReLU`` in ``src/tilefoundry/visitor_registry/op_cost.py`` (``_elementwise``, as those are), and it is deliberately not made here -- this package's boundary is ``tests/models/gemma2_2b/``. -They are not recorded as ``BLOCKED`` because a gate could not carry the claim -honestly: the analyse and sized gates would absorb it, but ``schedule`` fails -with ``PartitionProblemError`` -- a ``ValueError``, not a ``ScheduleError`` -- -which ``CapabilityGate.expected_failure`` cannot express, and the schedule case -cannot be dropped either (the harness requires each model to select its module -entry). Half a matrix of blocks around one missing registration would describe a -limit of this model, which this is not. +They are not recorded as ``BLOCKED`` because a block would describe a limit of +this model, which this is not: it is one missing cost registration, shared by +every case that touches ``Gelu``. """ from __future__ import annotations @@ -84,14 +77,6 @@ dims=ANALYZED_AT, ), ), - schedule=( - FunctionCase( - id="gemma2_2b/schedule/decoder_layer", - selector="decoder_layer", - topology="cta", - dims=ANALYZED_AT, - ), - ), sized=( SizedCase( id="gemma2_2b/sized/decoder_layer", diff --git a/tests/models/kimi_linear_48b_a3b/case.py b/tests/models/kimi_linear_48b_a3b/case.py index 3b761136..09c78568 100644 --- a/tests/models/kimi_linear_48b_a3b/case.py +++ b/tests/models/kimi_linear_48b_a3b/case.py @@ -14,10 +14,8 @@ fail: two for MLA, and three for the MoE covering the router bias, the order the routed scaling is applied in, and the shared expert's contribution. -Schedule admits one function per execution Module, so it selects each child's -entry and not its leaves; analyze selects everything the tree defines. What is not -selected is untested, and the report derives that from the model's own function -inventory. +Analyze selects everything the tree defines. What is not selected is untested, +and the report derives that from the model's own function inventory. """ from __future__ import annotations @@ -90,24 +88,6 @@ selector="moe.shared_expert", ), ), - schedule=( - FunctionCase( - id="kimi_linear_48b_a3b/schedule/kda_attention", - selector="kda.kda_attention", - topology="cta", - ), - FunctionCase( - id="kimi_linear_48b_a3b/schedule/mla_attention", - selector="mla.mla_attention", - topology="cta", - dims=ANALYZED_AT, - ), - FunctionCase( - id="kimi_linear_48b_a3b/schedule/moe", - selector="moe.moe", - topology="cta", - ), - ), #: Only MLA leaves a dimension open. KDA's state is fixed-size and the MoE's #: expert count is a constant of the published model, so neither has a size to #: be asked at -- which is what those shapes mean, not a capability they lack. diff --git a/tests/models/kimi_linear_48b_a3b/reference.py b/tests/models/kimi_linear_48b_a3b/reference.py index 1cb12f49..3b7af900 100644 --- a/tests/models/kimi_linear_48b_a3b/reference.py +++ b/tests/models/kimi_linear_48b_a3b/reference.py @@ -374,7 +374,7 @@ class KdaReferenceUnavailable(RuntimeError): #: Why the KDA reference is blocked, as measured on 2026-07-28. #: #: It is the *reference* that is blocked, not the model: `model.py` -#: describes `kda_attention` completely, and it analyses and schedules. What is +#: describes `kda_attention` completely, and it analyses. What is #: missing is anything to check its values against. #: #: `transformers` 5.14.1 has no `kimi_linear` implementation: `KimiLinearForCausalLM` diff --git a/tests/models/minicpm3_4b/case.py b/tests/models/minicpm3_4b/case.py index 1d68fb0a..f584f044 100644 --- a/tests/models/minicpm3_4b/case.py +++ b/tests/models/minicpm3_4b/case.py @@ -9,17 +9,8 @@ makes equal parts, so the model states them with ``tf.slice`` -- and no analysis has a cost evaluator for ``Slice``. Measured per (function, family): the two Slice-free functions analyse under all four families; the two that carry a Slice -fail under all four with ``no cost evaluator registered for Slice``, and the -partition path fails on the same op. - -One caveat about the schedule gate, which nothing in this package can fix: the -partition path raises ``PartitionProblemError``, which is a ``ValueError`` and not -a ``ScheduleError``, while ``test_schedule_coverage.py`` holds a blocked schedule -case to ``expect=ScheduleError``. So the gate below states the right reason and -would still be recorded as a plain failure rather than as the expected block. A -``Slice`` cost evaluator retires both problems at once; until one exists, adding -this case to ``registry.CORPUS`` needs that or a harness that sees the error the -partitioner actually raises. +fail under all four with ``no cost evaluator registered for Slice``. A ``Slice`` +cost evaluator retires the problem and lifts every gate below at once. """ from __future__ import annotations @@ -76,14 +67,6 @@ dims=ANALYZED_AT, ), ), - schedule=( - FunctionCase( - id="minicpm3_4b/schedule/decoder_layer", - selector="decoder_layer", - topology="cta", - dims=ANALYZED_AT, - ), - ), sized=( SizedCase( id="minicpm3_4b/sized/decoder_layer", diff --git a/tests/models/qwen2_5_1_5b/case.py b/tests/models/qwen2_5_1_5b/case.py index 922f8569..986cfde1 100644 --- a/tests/models/qwen2_5_1_5b/case.py +++ b/tests/models/qwen2_5_1_5b/case.py @@ -58,14 +58,6 @@ dims=ANALYZED_AT, ), ), - schedule=( - FunctionCase( - id="qwen2_5_1_5b/schedule/decoder_layer", - selector="decoder_layer", - topology="cta", - dims=ANALYZED_AT, - ), - ), sized=( SizedCase( id="qwen2_5_1_5b/sized/decoder_layer", diff --git a/tests/models/qwen3_1_7b/case.py b/tests/models/qwen3_1_7b/case.py index 1d252518..5e5a64ad 100644 --- a/tests/models/qwen3_1_7b/case.py +++ b/tests/models/qwen3_1_7b/case.py @@ -62,14 +62,6 @@ dims=ANALYZED_AT, ), ), - schedule=( - FunctionCase( - id="qwen3_1_7b/schedule/decoder_layer", - selector="decoder_layer", - topology="cta", - dims=ANALYZED_AT, - ), - ), sized=( SizedCase( id="qwen3_1_7b/sized/decoder_layer", diff --git a/tests/models/qwen3_5_35b_a3b/case.py b/tests/models/qwen3_5_35b_a3b/case.py index 3d7fafc0..62cf1569 100644 --- a/tests/models/qwen3_5_35b_a3b/case.py +++ b/tests/models/qwen3_5_35b_a3b/case.py @@ -1,5 +1,5 @@ -"""Qwen3.5-35B-A3B as corpus cases: what runs, what is analysed, what is -scheduled, and at what context length. +"""Qwen3.5-35B-A3B as corpus cases: what runs, what is analysed, and at what +context length. Three Modules, three cases. The published stack is a hybrid -- three linear attention layers to every full attention one, each ending in the same 256-expert @@ -96,13 +96,6 @@ id="qwen3_5_35b_a3b/analyze/l2_normalise", selector="l2_normalise" ), ), - schedule=( - FunctionCase( - id="qwen3_5_35b_a3b/schedule/linear_attention", - selector="linear_attention", - topology="cta", - ), - ), #: The Gated DeltaNet's state is fixed-size, so no extent in this Module is #: left open and there is no context length to ask it about. That is what a #: recurrent state means rather than a missing capability, so `sized` is empty @@ -145,14 +138,6 @@ id="qwen3_5_35b_a3b/analyze/partial_rope_kv", selector="partial_rope_kv" ), ), - schedule=( - FunctionCase( - id="qwen3_5_35b_a3b/schedule/full_attention", - selector="full_attention", - topology="cta", - dims=ANALYZED_AT, - ), - ), sized=( SizedCase( id="qwen3_5_35b_a3b/sized/full_attention", @@ -182,8 +167,8 @@ #: through a command: one read the selected experts back as a *set*, and every #: predicate `check` offers compares positionally; the other required the routed #: branch and the shared branch to sum to the whole block, which is two commands' - #: outputs added together. What this case contributes here is the block's analysis - #: and schedule coverage and its function inventory. + #: outputs added together. What this case contributes here is the block's + #: analysis coverage and its function inventory. analyze=( FunctionCase(id="qwen3_5_35b_a3b/analyze/experts", selector="experts"), FunctionCase(id="qwen3_5_35b_a3b/analyze/post_norm", selector="post_norm"), @@ -197,11 +182,6 @@ id="qwen3_5_35b_a3b/analyze/shared_expert", selector="shared_expert" ), ), - schedule=( - FunctionCase( - id="qwen3_5_35b_a3b/schedule/experts", selector="experts", topology="cta" - ), - ), #: One token through a router; nothing here is authored over a context. sized=(), ) diff --git a/tests/models/registry.py b/tests/models/registry.py index cdd56757..2510e0b7 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -19,15 +19,12 @@ case starts passing, the case fails until the package is corrected, so the matrix cannot quietly drift into describing a system nobody has. -Analyze selects every function a model defines. Schedule cannot: the device-wide -partition algorithm decides the launch, so it admits only the module entry -function, and selecting a leaf for it would be selecting something the algorithm -has no answer to rather than something it answers badly. - -`sized` is a third question, asked separately because a model can answer the -others without answering it: whether it can be analysed at a context length of -the caller's choosing. A model authored as one fixed shape analyses and schedules -perfectly well and has no context length to state, and the two facts must not be +Analyze selects every function a model defines. + +`sized` is a second question, asked separately because a model can answer the +other without answering it: whether it can be analysed at a context length of +the caller's choosing. A model authored as one fixed shape analyses perfectly +well and has no context length to state, and the two facts must not be collapsed -- a working analysis recorded as broken, or a missing capability recorded as nothing at all. It stays its own row once a model answers both, so there is somewhere to record the next model that answers only one. diff --git a/tests/schedule/test_atom_facts.py b/tests/schedule/test_atom_facts.py deleted file mode 100644 index 157aac99..00000000 --- a/tests/schedule/test_atom_facts.py +++ /dev/null @@ -1,108 +0,0 @@ -"""``candidate_atoms(op, target) -> list[AtomFact]`` -- the CUDA target's own candidate enumeration. - -``candidate_atoms(op, target) -> list[AtomFact]`` -- the CUDA target's -own candidate enumeration: HIR ``MatMul`` op + target -> the MMA atom -candidates it could run on (a hard filter over shape/dtype/layout; no -CP-SAT ranking, which is the solver's own subject). - -Builds a bf16 gemm HIR function -- bf16 being the sole dtype the one registered -SM80 atom accepts -- and checks the listed ``AtomFact`` against that atom's -real, known numbers, not just non-empty/non-zero placeholders. -""" - -from __future__ import annotations - -import pytest - -from tests.fixtures.shapes.matmul_programs import ( - cuda_bf16_gemm as bf16_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - cuda_f32_gemm as f32_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - cuda_odd_m_bf16_gemm as odd_shape_bf16_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - gemm_rms_norm, -) -from tilefoundry.dsl.tf import * # noqa: F401,F403 -- matmul/rms_norm resolved dynamically -from tilefoundry.ir.tir.cuda.nn.mma import SM80_16x8x16_F32BF16BF16F32_TN -from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom -from tilefoundry.ir.types import DType -from tilefoundry.schedule.facts import AtomFact -from tilefoundry.target import CpuTarget -from tilefoundry.target.cuda.atoms import candidate_atoms - - -def test_bf16_gemm_lists_the_sm80_atom_with_real_numbers(): - """Test bf16 gemm lists the sm80 atom with real numbers. - - The sole registered atom (SM80 16x8x16, bf16 x bf16 -> f32) is a - candidate for a bf16 gemm whose M/N/K (64, 64, 128) all divide its - (16, 8, 16) shape; every ``AtomFact`` field is checked against the - atom's own known real numbers. - """ - facts = candidate_atoms(bf16_gemm.entry_function().body, bf16_gemm.resolve_target()) - - print("\n=== candidate AtomFacts (bf16 gemm, M=64 N=64 K=128) ===") - for fact in facts: - print(fact) - - assert len(facts) == 1 - fact = facts[0] - assert isinstance(fact, AtomFact) - assert fact.shape == (16, 8, 16) - assert fact.dtype == (DType.bf16, DType.bf16, DType.f32) - assert fact.duration > 0 - assert isinstance(fact.atom, MmaAtom) - assert fact.atom.op is SM80_16x8x16_F32BF16BF16F32_TN - - assert fact.storage == { - "a_reg_bytes": 16, - "b_reg_bytes": 8, - "c_reg_bytes": 16, - "reg_bytes": 40, - } - assert fact.resource == {"lane": 32} - assert fact.is_async is False - - -def test_a_gemm_the_atom_cannot_run_lists_no_candidate(): - """Both halves of the hard filter, each an empty list rather than an error. - - The SM80 atom is bf16 x bf16 -> f32, so an all-f32 gemm's operand dtypes do - not match it at all; and a gemm whose M is 15 does not divide the atom's M of - 16 even though its dtypes do match. An empty list is the answer a caller can - act on -- an error here would make "this atom does not apply" indistinguishable - from "this op cannot be asked about". - """ - assert candidate_atoms(f32_gemm.entry_function().body, f32_gemm.resolve_target()) == [] - assert ( - candidate_atoms( - odd_shape_bf16_gemm.entry_function().body, - odd_shape_bf16_gemm.resolve_target(), - ) - == [] - ) - - -def test_non_matmul_op_raises(): - """V1 supports a MatMul Call only. - - V1 supports a MatMul Call only; any other op (here RMSNorm, from - ``gemm_rms_norm``'s body) raises a clear ``NotImplementedError`` rather - than silently returning ``[]``. - """ - with pytest.raises(NotImplementedError): - candidate_atoms(gemm_rms_norm.body) - - -def test_non_cuda_target_raises(): - """V1 supports ``CudaTarget`` only -- no per-atom device facts exist for a CPU target. - - V1 supports ``CudaTarget`` only -- no per-atom device facts - (sm_count / hbm_bandwidth / peak_for) exist for a CPU target. - """ - with pytest.raises(NotImplementedError): - candidate_atoms(bf16_gemm.entry_function().body, target=CpuTarget()) diff --git a/tests/schedule/test_partition.py b/tests/schedule/test_partition.py deleted file mode 100644 index a9e29a0e..00000000 --- a/tests/schedule/test_partition.py +++ /dev/null @@ -1,492 +0,0 @@ -"""Cover closed spatial partition scheduling through the public boundary. - -Returned plans must verify and invalid plans identify the conflicting edge, -position, or operation. A small program keeps each mutation isolated; corpus -tests cover real model scheduling. -""" - -from __future__ import annotations - -import dataclasses -import json -from dataclasses import replace - -import pytest -from ortools.sat.python import cp_model - -from tests.fixtures.logical.authored_constraint import AuthoredConstraint -from tests.fixtures.logical.gqa_static import static_online_attend -from tests.fixtures.shapes.matmul_programs import bf16_gemm_rms_norm -from tilefoundry.inspection.python_printer import as_script -from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.shard import ShardLayout, Topology -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.schedule import PlanVerificationError, ScheduleError, ScheduleOptions, schedule -from tilefoundry.schedule.partition import ( - PartitionedOperation, - PartitionFacts, - PartitionFactsError, - PartitionProblemError, - PlacedValue, - PositionInterval, - build_partition_problem, - build_partition_program, -) -from tilefoundry.schedule.partition import solve as solve_module -from tilefoundry.schedule.pipeline.problem import PipelineProblemError - -_SOLVER = ScheduleOptions(workers=1, stop_at_first_solution=True) - - -def _module(extent: int = 4): - return replace(bf16_gemm_rms_norm, topologies=(Topology("cta", extent),)) - - -def _closed(extent: int = 4): - module = _module(extent) - function = module.entry_function() - program = build_partition_program(module, function) - facts = module.resolve_target().get_facts(PartitionFacts, program.facts_query("cta")) - return module, function, program, facts - - -@pytest.fixture(scope="module") -def solved(): - """One solved plan, shared by every test that only mutates a copy of it. - - Verification is a structural check over an immutable plan, so the tests below - build broken variants with `replace` rather than re-solving. Solving once is - not an optimisation of the assertions -- it is the same plan under every - mutation, which is what makes the failures comparable. - """ - module, function, _, _ = _closed() - return module, function, schedule(module, function, topology="cta").plan - - -def test_partition_schedules_through_the_public_operation_without_rewriting() -> None: - """The plan is a decision about the program. - - The plan is a decision about the program, so the program comes back as the - same objects and prints identically -- nothing about scheduling rewrites it. - """ - module, function, _, _ = _closed() - before = as_script(module) - - result = schedule(module, function, topology="cta") - - assert result.module is module - assert result.function is function - assert result.topology == Topology("cta", 4) - assert result.plan.topology == "cta" - assert result.plan.proof.objective_ns > 0 - assert result.plan.proof.best_bound_ns <= result.plan.proof.objective_ns - assert result.plan.root_results - assert as_script(result.module) == before - - -def test_partition_accepts_an_authored_where_constraint_through_schedule() -> None: - module = AuthoredConstraint - function = module.entry_function() - - plan = schedule( - module, - function, - topology="cta", - options=_SOLVER, - ).plan - - assert plan.root_results - - -def test_partition_program_states_the_program_without_asking_the_hardware() -> None: - _, _, program, _ = _closed() - - assert program.sites - assert program.root_value_ids - assert all(base.storage.name.lower() == "gmem" for base in program.value_base_types.values()) - assert not any( - field.name in {"target", "facts", "device"} for field in dataclasses.fields(program) - ) - - -def test_partition_problem_closes_every_hardware_number_before_solving() -> None: - _, _, program, facts = _closed() - - problem = build_partition_problem(program, facts, Topology("cta", 4)) - - assert problem.facts is facts - assert facts.parallel_units > 0 - assert facts.memory_bandwidth_bytes_per_second > 0 - assert facts.memory_capacity_bytes > 0 - assert facts.peak_flops_per_second - assert all( - not hasattr(candidate, "capacity_bytes") for candidate in problem.candidates.values() - ) - assert all(candidate.duration_ns >= 0 for candidate in problem.candidates.values()) - - -def test_partition_keeps_synthesized_gmem_views_zero_copy_and_only_where_needed() -> None: - """The public partition boundary makes every synthesized move a zero-copy view. - - Partition accepts only GMEM tensor values. Its synthesized Reshards therefore - change placement without changing storage, so ``moved_bytes`` is always zero. - They exist only where no authored candidate already produces the placement. - """ - _, _, program, facts = _closed() - - problem = build_partition_problem(program, facts, Topology("cta", 4)) - - authored = [ - candidate for candidate in problem.candidates.values() if candidate.site_id is not None - ] - synthesized = [ - candidate for candidate in problem.candidates.values() if candidate.site_id is None - ] - assert authored - assert all(type(candidate).__name__ == "OpCandidate" for candidate in authored) - assert all(type(candidate).__name__ == "OpCandidate" for candidate in synthesized) - for candidate in synthesized: - tensors = (*candidate.source_types, *candidate.output_types) - assert tensors - assert all( - isinstance(type_, TensorType) and type_.storage is StorageKind.GMEM - for type_ in tensors - ) - assert candidate.moved_bytes == 0 - assert candidate.topology_count == 0 - - for bucket in problem.buckets.values(): - producers = [problem.candidates[cid] for cid in bucket.candidate_ids] - assert not ( - [item for item in producers if item.site_id is not None] - and [item for item in producers if item.site_id is None] - ) -def test_partition_refuses_a_level_the_facts_and_the_program_do_not_share() -> None: - """Three ways of asking about the wrong level, each answered before a solve. - - Facts projected for another level describe another machine's parallelism; a - level the target does not partition has no facts to project at all; and a - level the program never declared has no extent to place work across. All - three used to be servable by whatever numbers happened to be at hand. - """ - module, _, program, facts = _closed() - - with pytest.raises(PartitionProblemError, match="describe 'core'"): - build_partition_problem(program, replace(facts, topology="core"), Topology("cta", 4)) - - with pytest.raises(ValueError, match="no partition facts for 'thread'"): - module.resolve_target().get_facts(PartitionFacts, program.facts_query("thread")) - - thread_only = replace(bf16_gemm_rms_norm, topologies=(Topology("thread", 128),)) - with pytest.raises(ScheduleError, match="cta"): - schedule(thread_only, thread_only.entry_function(), topology="cta") - - -def test_partition_refuses_hardware_it_cannot_charge_the_work_against() -> None: - """An extent wider than the machine states and a missing rate are both refusals, not defaults. - - An extent wider than the machine states and a missing rate are both - refusals, not defaults: a problem that guessed either would return a plan - priced against a machine nobody has. - """ - _, _, program, facts = _closed() - - with pytest.raises(PartitionProblemError, match="exceeds the 2 parallel units"): - build_partition_problem(program, replace(facts, parallel_units=2), Topology("cta", 4)) - - with pytest.raises(PartitionFactsError, match="no dense peak rate"): - build_partition_problem( - program, replace(facts, peak_flops_per_second=()), Topology("cta", 4) - ) - - -def test_partition_plan_names_values_and_operations_from_the_authored_program( - solved, -) -> None: - _, _, plan = solved - - names = {value.id for value in plan.values} - assert {"x", "w", "weight"} <= names - for value in plan.values: - assert isinstance(value.type, TensorType) - producers = {value.producer_id for value in plan.values if value.producer_id} - assert producers <= {operation.id for operation in plan.operations} - assert plan.root_results - assert set(plan.root_results) <= names - - -def test_partition_plan_states_a_reshard_as_an_operation_with_both_placements() -> None: - """A moved value is one of the plan's own operations. - - A moved value is one of the plan's own operations, and both placements of - it are named values: same shape and dtype, different type, and more than one - placement sharing a base name. A plan that reported the move on the side - would leave a reader unable to say where a value is. - """ - module = static_online_attend - plan = schedule( - module, - module.entry_function(), - topology="cta", - options=_SOLVER, - ).plan - - reshards = tuple(operation for operation in plan.operations if operation.operation == "Reshard") - assert reshards - assert not hasattr(plan, "routes") - assert not hasattr(plan, "report") - values = {value.id: value for value in plan.values} - for reshard in reshards: - assert reshard.positions is None - synthesized = tuple(item for item in reshards if item.synthesized) - assert synthesized - for reshard in synthesized: - assert len(reshard.input_ids) == 1 - assert len(reshard.output_ids) == 1 - source = values[reshard.input_ids[0]] - target = values[reshard.output_ids[0]] - assert source.type != target.type - assert (source.type.shape, source.type.dtype) == ( - target.type.shape, - target.type.dtype, - ) - - qualified = tuple(value for value in plan.values if "@" in value.id) - assert qualified - for base in {value.id.split("@", 1)[0] for value in qualified}: - placements = tuple(value for value in qualified if value.id.split("@", 1)[0] == base) - assert len(placements) > 1 - - -def test_verification_rejects_an_edge_the_two_ends_do_not_agree_on(solved) -> None: - """Every producer/consumer edge is stated twice, and both statements have to say the same thing. - - Every producer/consumer edge is stated twice, and both statements have to - say the same thing. - - A named producer that does not list the placement among its outputs, a named - consumer that does not read it, and either side disowning an edge the other - still names, are four separate corruptions with four separate messages -- and - all four are invisible to a check that only walked one direction. - """ - module, function, plan = solved - level = Topology("cta", 4) - - produced = next(value for value in plan.values if value.producer_id) - other_operation = next( - operation for operation in plan.operations if produced.id not in operation.output_ids - ) - with pytest.raises(PlanVerificationError, match="does not produce it"): - _with_value(plan, produced, producer_id=other_operation.id).verify(module, function, level) - - read = next(value for value in plan.values if value.consumer_ids) - not_a_reader = next( - operation for operation in plan.operations if read.id not in operation.input_ids - ) - with pytest.raises(PlanVerificationError, match="does not read it"): - _with_value(plan, read, consumer_ids=(*read.consumer_ids, not_a_reader.id)).verify( - module, function, level - ) - - with pytest.raises(PlanVerificationError, match="which names producer None"): - _with_value(plan, produced, producer_id=None).verify(module, function, level) - - with pytest.raises(PlanVerificationError, match="does not name it as a consumer"): - _with_value(plan, read, consumer_ids=()).verify(module, function, level) - - -def _with_value(plan, target, **changes): - """*plan* with one placement replaced -- the only difference from the plan the solver returned. - - *plan* with one placement replaced -- the only difference from the plan the - solver returned, so a rejection can only be about that one field. - """ - return replace( - plan, - values=tuple( - replace(value, **changes) if value is target else value for value in plan.values - ), - ) - - -def test_verification_rejects_a_value_nothing_runs_and_a_root_nothing_reaches( - solved, -) -> None: - """A plan has to place every value it uses and reach every result it claims. - - An operation over a value the plan does not place, a producer the plan does - not run, a root that is not a placement at all, and a root reachable only - through an edge one end does not confirm: each is a plan that describes work - nobody could carry out, and each says which. - """ - module, function, plan = solved - level = Topology("cta", 4) - - with pytest.raises(PlanVerificationError, match="unplaced value"): - replace(plan, values=()).verify(module, function, level) - - produced = next(value for value in plan.values if value.producer_id) - with pytest.raises(PlanVerificationError, match="which the plan does not run"): - _with_value(plan, produced, producer_id="nobody").verify(module, function, level) - - with pytest.raises(PlanVerificationError, match="unplaced"): - replace(plan, root_results=("nowhere",)).verify(module, function, level) - - real = next(operation for operation in plan.operations if operation.output_ids) - orphan = PlacedValue( - id="orphan", - type=plan.values[0].type, - producer_id=real.id, - consumer_ids=(), - positions=plan.values[0].positions, - ) - with pytest.raises(PlanVerificationError, match="does not produce it"): - replace(plan, values=(*plan.values, orphan), root_results=("orphan",)).verify( - module, function, level - ) - - -def test_verification_rejects_a_placement_the_level_does_not_contain(solved) -> None: - """A plan is solved for one extent of one level, so it is only meaningful against that extent. - - A plan is solved for one extent of one level, so it is only meaningful - against that extent: a placement outside the positions the level declares, and - a level of a different width than the one solved for, are both refused. - """ - module, function, plan = solved - - first = plan.values[0] - with pytest.raises(PlanVerificationError, match="outside the 4 positions"): - _with_value(plan, first, positions=PositionInterval(3, 9)).verify( - module, function, Topology("cta", 4) - ) - - with pytest.raises(PlanVerificationError, match="the level declares 8"): - plan.verify(module, function, Topology("cta", 8)) - - -def test_verification_rejects_two_operations_sharing_a_position(solved) -> None: - """Test verification rejects two operations sharing a position. - - Two operations placed on one position at one time is the one thing a - spatial partition exists to decide, so an overlap is not a preference. - """ - module, function, plan = solved - placed = next( - operation - for operation in plan.operations - if operation.positions is not None and operation.interval is not None - ) - - twin = replace(placed, id=placed.id + ".twin", input_ids=(), output_ids=()) - broken = replace(plan, operations=(*plan.operations, twin)) - with pytest.raises(PlanVerificationError, match="overlapping positions"): - broken.verify(module, function, Topology("cta", 4)) - - -def test_verification_rejects_a_reshard_that_moves_nothing(solved) -> None: - """Test verification rejects a reshard that moves nothing. - - A move whose source and destination are the same placement is a cost the - plan charges for work it does not do. - """ - module, function, plan = solved - held = next(value for value in plan.values if value.producer_id is None) - identity = PartitionedOperation( - id="reshard:identity", - operation="Reshard", - synthesized=True, - input_ids=(held.id,), - output_ids=(held.id,), - positions=None, - interval=None, - ) - - broken = replace( - _with_value( - plan, - held, - producer_id=identity.id, - consumer_ids=(*held.consumer_ids, identity.id), - ), - operations=(*plan.operations, identity), - ) - with pytest.raises(PlanVerificationError, match="moves nothing"): - broken.verify(module, function, Topology("cta", 4)) - - -def test_verification_rejects_a_bound_above_its_own_objective(solved) -> None: - """The proof is part of the plan. - - The proof is part of the plan: a lower bound above the objective it bounds - is an arithmetic impossibility, and reading one as optimality would report a - plan as proven that is not. - """ - module, function, plan = solved - - broken = replace(plan, proof=replace(plan.proof, best_bound_ns=plan.proof.objective_ns + 1)) - with pytest.raises(PlanVerificationError, match="bound above its own objective"): - broken.verify(module, function, Topology("cta", 4)) - - -def test_verification_needs_no_solver(solved) -> None: - """Verification is a structural check, so it must not reach the solver.""" - module, function, plan = solved - calls: list[object] = [] - - original = cp_model.CpSolver.Solve - - def refuse(self, model, *args, **kwargs): # pragma: no cover - must not run - calls.append(model) - raise AssertionError("verification solved a model") - - cp_model.CpSolver.Solve = refuse - try: - plan.verify(module, function, Topology("cta", 4)) - finally: - cp_model.CpSolver.Solve = original - assert calls == [] - - -def test_partition_plan_json_preserves_the_decision(solved) -> None: - """The library plan still exposes its deterministic machine representation.""" - _, _, plan = solved - - assert plan.to_json() == plan.to_json() - data = json.loads(plan.to_json()) - - assert data["topology"] == plan.topology == "cta" - assert data["extent"] == plan.extent - assert data["proof"]["status"] == plan.proof.status - assert data["root_results"] == list(plan.root_results) - assert {item["id"] for item in data["values"]} == {value.id for value in plan.values} - assert {item["id"] for item in data["operations"]} == { - operation.id for operation in plan.operations - } - by_id = {item["id"]: item for item in data["values"]} - for value in plan.values: - stated = by_id[value.id]["type"] - assert stated["dtype"] == value.type.dtype.name - assert stated["storage"] == value.type.storage.name.lower() - assert stated["shape"] == [str(dim) for dim in value.type.shape] - if isinstance(value.type.layout, ShardLayout): - assert stated["layout"]["topology"] == "cta" - else: - assert stated["layout"] is None - - -def test_a_problem_that_cannot_be_formed_is_a_schedule_failure() -> None: - """The algorithms' own failures are reachable as `ScheduleError`. - - A caller asks this layer to schedule something and catches what the layer - raises; a capability the layer cannot serve is recorded against the same - type. While these sat outside it, a limit of an algorithm could only be - stated as a bare `ValueError` -- which is also what a caller passing nonsense - gets, so a recorded limit and a caller's mistake were indistinguishable. - """ - for error in (PartitionProblemError, PipelineProblemError): - assert issubclass(error, ScheduleError), error.__name__ - - assert issubclass(error, ValueError), error.__name__ - - assert issubclass(solve_module.PartitionSolveError, RuntimeError) diff --git a/tests/schedule/test_pipeline.py b/tests/schedule/test_pipeline.py deleted file mode 100644 index 5199f5d0..00000000 --- a/tests/schedule/test_pipeline.py +++ /dev/null @@ -1,179 +0,0 @@ -"""Closed CUDA pipeline scheduling through the public boundary.""" - -from __future__ import annotations - -from dataclasses import replace - -import pytest - -from tests.fixtures.shapes.matmul_programs import bf16_gemm_rms_norm -from tilefoundry import func -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import sigmoid -from tilefoundry.ir.types.shard import Topology -from tilefoundry.schedule import schedule -from tilefoundry.schedule.pipeline import ( - PipelineFacts, - PipelineProblemError, - build_pipeline_problem, - build_pipeline_program, -) -from tilefoundry.schedule.pipeline.problem import ( - PipelineBufferProblem, - PipelineProblem, - PipelineStatementProblem, -) -from tilefoundry.schedule.pipeline.solve import ( - PipelineSolveError, - solve_pipeline_problem, -) -from tilefoundry.target import CudaTarget - - -def _module(): - return replace(bf16_gemm_rms_norm, topologies=(Topology("cta", 1), Topology("thread", 128))) - - -@func(target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 8),)) -def local_sigmoid(x: Tensor[(8,), "f32", None, "rmem"]): - return sigmoid(x) - - -def test_pipeline_accepts_a_local_value_before_layout_is_resolved() -> None: - result = schedule(local_sigmoid, local_sigmoid.entry_function(), topology="thread") - - assert result.topology == Topology("thread", 8) - assert tuple(statement.id for statement in result.plan.statements) == ("Sigmoid",) - - -def test_pipeline_closes_target_facts_before_solving_and_exports_stable_values() -> None: - module = _module() - function = module.entry_function() - program = build_pipeline_program(module, function) - facts = module.resolve_target().get_facts(PipelineFacts, program.facts_query("thread")) - problem = build_pipeline_problem(program, facts, Topology("thread", 128)) - - assert tuple(item.id for item in problem.statements) == ("MM", "RN") - assert all(item.candidates for item in problem.statements) - assert not hasattr(problem, "target") - - result = schedule(module, function, topology="thread") - plan = result.plan - assert tuple(item.id for item in plan.statements) == ("MM", "RN") - assert all(item.end > item.start >= 0 for item in plan.statements) - assert plan.to_json() == plan.to_json() - assert {hole.statement_id for hole in plan.holes} == {"MM", "RN"} - assert all(isinstance(relation, str) for hole in plan.holes for relation in hole.relations) - - -def _problem(): - module = _module() - program = build_pipeline_program(module, module.entry_function()) - facts = module.resolve_target().get_facts(PipelineFacts, program.facts_query("thread")) - return build_pipeline_problem(program, facts, Topology("thread", 128)) - - -def test_a_buffer_that_carries_a_dependence_gets_more_than_one_slot() -> None: - """The property that makes this a pipeline: `h` is the accumulator the matmul carries along k. - - The property that makes this a pipeline: `h` is the accumulator the - matmul carries along k, so it has to hold two tiles at once. - """ - problem = _problem() - carried = {buffer.id: buffer.carried_distances for buffer in problem.buffers} - assert carried["h"] == (("MM", (0, 0, 1)),) - assert carried["x"] == () - - ring = {buffer.id: buffer.ring_depth for buffer in solve_pipeline_problem(problem).buffers} - assert ring["h"] == 2 - assert ring["x"] == 1 - assert set(ring.values()) != {1} - - -def test_ring_depth_counts_tiles_not_iterations() -> None: - """A distance shorter than the tile it runs inside spans one tile.""" - - def depth(distance: int, tile: int) -> int: - problem = PipelineProblem( - topology="thread", - capacity_bytes=1024, - statements=( - PipelineStatementProblem( - id="S", - extents=(tile,), - candidates=(_candidate(),), - resources=(), - footprint_bytes=(), - ), - ), - buffers=( - PipelineBufferProblem( - id="b", - producer_ids=("S",), - consumer_ids=("S",), - carried_distances=(("S", (distance,)),), - ), - ), - ) - return solve_pipeline_problem(problem).buffers[0].ring_depth - - assert depth(0, 8) == 1 - assert depth(1, 8) == 2 - assert depth(8, 8) == 2 - assert depth(9, 8) == 3 - assert depth(64, 8) == 9 - - -def test_a_statement_records_what_it_holds_against_the_tile_store() -> None: - """Capacity is recorded, not enforced -- but it has to be recorded.""" - problem = _problem() - assert problem.capacity_bytes > 0 - held = dict(next(item for item in problem.statements if item.id == "MM").footprint_bytes) - assert held == {"h": 8192, "w": 16384, "x": 16384} - - solution = solve_pipeline_problem(problem) - ring = {buffer.id: buffer.ring_depth for buffer in solution.buffers} - matmul_solution = next(item for item in solution.statements if item.id == "MM") - assert matmul_solution.footprint_bytes == sum(held[name] * ring[name] for name in held) - assert matmul_solution.fits_capacity is ( - matmul_solution.footprint_bytes <= problem.capacity_bytes - ) - - -def test_a_statement_too_wide_for_the_store_is_reported_not_dropped() -> None: - problem = replace(_problem(), capacity_bytes=1) - solution = solve_pipeline_problem(problem) - - assert len(solution.statements) == len(problem.statements) - assert not any(item.fits_capacity for item in solution.statements) - - -def test_a_distance_measured_against_an_unknown_statement_is_refused() -> None: - problem = _problem() - broken = replace( - problem, - buffers=( - replace(problem.buffers[0], carried_distances=(("nobody", (1,)),)), - *problem.buffers[1:], - ), - ) - with pytest.raises(PipelineSolveError, match="unknown statement"): - solve_pipeline_problem(broken) - - -def _candidate(): - """The one thing the solver reads off a candidate: how long it takes.""" - module = _module() - program = build_pipeline_program(module, module.entry_function()) - facts = module.resolve_target().get_facts(PipelineFacts, program.facts_query("thread")) - return facts.instructions[0].candidates[0] - - -def test_pipeline_rejects_missing_statement_facts_before_solving() -> None: - module = _module() - program = build_pipeline_program(module, module.entry_function()) - facts = module.resolve_target().get_facts(PipelineFacts, program.facts_query("thread")) - incomplete = replace(facts, instructions=facts.instructions[:-1]) - - with pytest.raises(PipelineProblemError, match="do not match"): - build_pipeline_problem(program, incomplete, Topology("thread", 128)) diff --git a/tests/schedule/test_schedule_api.py b/tests/schedule/test_schedule_api.py deleted file mode 100644 index c4cfde2b..00000000 --- a/tests/schedule/test_schedule_api.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Schedule obtains one solver from the exact Target value.""" - -from __future__ import annotations - -import json -from dataclasses import dataclass -from typing import ClassVar - -import pytest - -from tilefoundry import func, module -from tilefoundry.dsl import Tensor -from tilefoundry.dsl.tf import * # noqa: F401,F403 -from tilefoundry.ir.types.shard import Topology -from tilefoundry.schedule.api import ScheduleResult, schedule -from tilefoundry.schedule.errors import ScheduleError -from tilefoundry.schedule.plan import PlanVerificationError, SchedulePlan -from tilefoundry.target import CudaTarget, Target, TopologyLimitFacts, register_target -from tilefoundry.target.services import Scheduler - - -@dataclass(frozen=True) -class _Plan(SchedulePlan): - width: int - valid: bool = True - - def verify(self, module, function, topology) -> None: - if not self.valid: - raise PlanVerificationError("invalid plan") - - def to_json(self) -> str: - return json.dumps({"width": self.width}) - - def render(self) -> str: - return str(self.width) - - -_CALLS: list[str] = [] - - -def _solve(module, function, target, topology, options): - _CALLS.append(topology.name) - return _Plan(topology.size) - - -@dataclass(frozen=True) -class _TopologyTarget(Target): - topology_levels: ClassVar[tuple[str, ...]] = ("tile",) - - def get_facts(self, facts_type: type, query: object | None = None): - if facts_type is TopologyLimitFacts and query == "tile": - return TopologyLimitFacts("tile", 4) - return super().get_facts(facts_type, query) - - -@dataclass(frozen=True) -class _SchedulingTarget(_TopologyTarget): - name: ClassVar[str] = "test.scheduler" - - def get_scheduler(self, topology: str) -> Scheduler: - if topology == "tile": - return Scheduler("tile", _solve) - return super().get_scheduler(topology) - - -@dataclass(frozen=True) -class _UnsupportedTarget(_TopologyTarget): - name: ClassVar[str] = "test.unsupported-scheduler" - - -@func -def scale(x: Tensor[(64, 64), "f32"]) -> Tensor[(64, 64), "f32"]: - return relu(x) # noqa: F405 - - -@module(entry="scale", target=_SchedulingTarget(), topologies=(Topology("tile", 4),)) -class Widget: - scale = scale - - -@module(entry="scale", target=_UnsupportedTarget(), topologies=(Topology("tile", 4),)) -class Unsupported: - scale = scale - - -@dataclass(frozen=True) -class _BrokenSchedulerTarget(_TopologyTarget): - name: ClassVar[str] = "test.broken-scheduler" - - def get_scheduler(self, topology: str) -> Scheduler: - raise ValueError("provider scheduler failure") - - -@module(entry="scale", target=_BrokenSchedulerTarget(), topologies=(Topology("tile", 4),)) -class BrokenScheduler: - scale = scale - - -@module( - entry="scale", - target=CudaTarget("nvidia.h200_sxm"), - topologies=(Topology("cta", 1), Topology("thread", 1025)), -) -class OverLimitCuda: - scale = scale - - -def test_schedule_uses_the_target_selected_solver() -> None: - _CALLS.clear() - result = schedule(Widget, scale, topology="tile") - - assert isinstance(result, ScheduleResult) - assert result.plan == _Plan(4) - assert _CALLS == ["tile"] - - -def test_an_unsupported_scheduler_fails_before_a_solver_runs() -> None: - _CALLS.clear() - with pytest.raises(ScheduleError, match="Target .*no scheduler for 'tile'"): - schedule(Unsupported, scale, topology="tile") - assert _CALLS == [] - - -def test_schedule_preserves_plan_verification() -> None: - @register_target - @dataclass(frozen=True) - class _InvalidTarget(_SchedulingTarget): - name: ClassVar[str] = "test.invalid-scheduler" - - def get_scheduler(self, topology: str) -> Scheduler: - return Scheduler(topology, lambda *_args: _Plan(4, valid=False)) - - @module(entry="scale", target=_InvalidTarget(), topologies=(Topology("tile", 4),)) - class Invalid: - scale = scale - - with pytest.raises(PlanVerificationError, match="invalid plan"): - schedule(Invalid, scale, topology="tile") - - -def test_schedule_keeps_provider_failures_and_rejects_over_limit_topology() -> None: - with pytest.raises(ValueError, match="provider scheduler failure"): - schedule(BrokenScheduler, scale, topology="tile") - with pytest.raises(ValueError, match="1 <= extent <= 1024"): - schedule(OverLimitCuda, scale, topology="thread") diff --git a/tests/target/test_amx_target.py b/tests/target/test_amx_target.py index efec840b..11da660f 100644 --- a/tests/target/test_amx_target.py +++ b/tests/target/test_amx_target.py @@ -1,46 +1,18 @@ -"""Pin the AMX target's facts, atom candidates, and evidence origins. +"""Pin the AMX target's facts and evidence origins. Assertions use exact measured device values, including performance-core cache -sizes that unqualified sysctls under-report. They prevent silent schedules priced -from unmeasured rates or facts belonging to another hardware level. +sizes that unqualified sysctls under-report. They prevent facts being priced from +unmeasured rates or from another hardware level. """ from __future__ import annotations import pytest -from tests.fixtures.shapes.matmul_programs import ( - amx_bf16_gemm as bf16_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - amx_coarse_m_f32_gemm as coarse_m_f32_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - amx_f32_gemm as f32_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - amx_odd_m_f32_gemm as odd_m_f32_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - amx_odd_n_f32_gemm as odd_n_f32_gemm, -) -from tests.fixtures.shapes.matmul_programs import ( - amx_register_sized_f32_gemm as register_sized_f32_gemm, -) from tilefoundry.ir.types import DType from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.shard import Topology -from tilefoundry.schedule.facts import AtomFact -from tilefoundry.schedule.pipeline.facts import PipelineFacts, PipelineFactsQuery from tilefoundry.target import AmxTarget, TopologyLimitFacts, UnsupportedCapabilityError -from tilefoundry.target.amx.atoms import ( - AMX_REGISTERS, - CORE_CACHE, - AMX_FMA32_16x16x1_F32, - AmxAtom, - NEON_FMLA_4x4x1_F32, - candidate_atoms, -) def test_amx_target_reports_and_validates_its_own_topology_levels(): @@ -78,87 +50,6 @@ def test_amx_target_reports_and_validates_its_own_topology_levels(): target.validate_program_topology(Topology("core", DimVar("cores", 1, 8))) -def test_both_catalogue_atoms_are_priced_at_their_own_measured_rates(): - """Price both core units using their own measured rates. - - A 16x16 gemm fitting X/Y/Z registers admits both AMX 16x16 and NEON 4x4 - outer-product atoms. Each field and duration uses that atom's own geometry, - throughput, and operand-byte cost. Both are arithmetically memory-bound at - this granularity. - """ - facts = candidate_atoms( - register_sized_f32_gemm.entry_function().body, - register_sized_f32_gemm.resolve_target(), - ) - assert [fact.atom.op.name for fact in facts] == [ - "AMX_FMA32_16x16x1_F32", - "NEON_FMLA_4x4x1_F32", - ] - - amx, neon = facts - assert isinstance(amx, AtomFact) - assert isinstance(amx.atom, AmxAtom) - assert amx.atom.op is AMX_FMA32_16x16x1_F32 - assert amx.atom.op.level is AMX_REGISTERS - assert amx.shape == (16, 16, 1) - assert amx.dtype == (DType.f32, DType.f32, DType.f32) - assert amx.storage == { - "a_bytes": 64, - "b_bytes": 64, - "c_bytes": 1024, - "operand_bytes": 1152, - } - assert amx.resource == {"amx": 1} - assert amx.is_async is False - assert amx.compute_duration == pytest.approx(512e9 / 504_900_000_000) - assert amx.duration == pytest.approx(1152e9 / 200_000_000_000) - assert amx.duration > amx.compute_duration - - assert neon.atom.op is NEON_FMLA_4x4x1_F32 - assert neon.atom.op.level is CORE_CACHE - assert neon.shape == (4, 4, 1) - assert neon.storage == {"a_bytes": 16, "b_bytes": 16, "c_bytes": 64, "operand_bytes": 96} - assert neon.resource == {"neon": 1} - assert neon.is_async is False - assert neon.compute_duration == pytest.approx(32e9 / 107_700_000_000) - assert neon.duration == pytest.approx(96e9 / 200_000_000_000) - - -def test_the_hard_filter_is_per_atom_and_covers_storage_shape_and_dtype(): - """Reject atoms independently for storage, shape, and dtype mismatches. - - X/Y/Z geometry rejects a whole 64x128 operand from AMX while cache-streaming - NEON survives. Shape is evaluated per atom: M=8 fits storage but not AMX's - 16-row granularity, while NEON's 4 divides it. Indivisible extents and bf16 - match neither atom and return an empty list rather than an unsupported query. - """ - architecture = AmxTarget().architecture - assert AMX_REGISTERS.budget == ( - ("a_bytes", architecture.staging_bytes), - ("b_bytes", architecture.staging_bytes), - ("c_bytes", architecture.accumulator_bytes), - ) - assert CORE_CACHE.budget == () - assert CORE_CACHE.holds({}) and CORE_CACHE.holds({"a_bytes": 10**12}) - - whole_tensor = candidate_atoms(f32_gemm.entry_function().body, f32_gemm.resolve_target()) - assert [fact.atom.op.name for fact in whole_tensor] == ["NEON_FMLA_4x4x1_F32"] - assert not AMX_REGISTERS.holds({"a_bytes": 64 * 128 * 4, "b_bytes": 512, "c_bytes": 512}) - assert CORE_CACHE.holds({"a_bytes": 64 * 128 * 4, "b_bytes": 10**9, "c_bytes": 10**9}) - - coarse = candidate_atoms( - coarse_m_f32_gemm.entry_function().body, coarse_m_f32_gemm.resolve_target() - ) - assert [fact.atom.op.name for fact in coarse] == ["NEON_FMLA_4x4x1_F32"] - assert AMX_REGISTERS.holds({"a_bytes": 8 * 8 * 4, "b_bytes": 8 * 16 * 4, "c_bytes": 8 * 16 * 4}) - - for indivisible in (odd_m_f32_gemm, odd_n_f32_gemm): - assert ( - candidate_atoms(indivisible.entry_function().body, indivisible.resolve_target()) == [] - ) - assert candidate_atoms(bf16_gemm.entry_function().body, bf16_gemm.resolve_target()) == [] - - def test_amx_values_stand_on_the_installed_documents_and_say_how_they_were_got(): """Test amx values stand on the installed documents and say how they were got. @@ -198,20 +89,6 @@ def test_amx_values_stand_on_the_installed_documents_and_say_how_they_were_got() assert fact.conditions -def test_a_core_tile_is_bounded_by_the_l1d_not_by_the_register_files(): - """A core-level tile's resident working set is bounded by the performance core's L1d. - - A core-level tile's resident working set is bounded by the performance - core's L1d; the register files bound one atom instance instead, which the - storage filter enforces rather than a per-tile capacity. - """ - target = AmxTarget() - facts = target.get_facts(PipelineFacts, PipelineFactsQuery(topology="core", statements=())) - assert facts.tile_capacity_bytes == target.device.l1d_bytes_per_performance_core - assert facts.tile_capacity_scope == "core" - assert target.architecture.accumulator_bytes < (target.device.l1d_bytes_per_performance_core) - - def test_unmeasured_units_and_dtypes_have_no_throughput_entry(): """FMA16 exists in the instruction set, but no f16 rate was measured on either unit. diff --git a/tests/target/test_target.py b/tests/target/test_target.py index fb3d9b9f..e89e5bb4 100644 --- a/tests/target/test_target.py +++ b/tests/target/test_target.py @@ -1,7 +1,6 @@ -"""What a Target validates about a program, which scheduler it reaches. +"""What a Target validates about a program, and how codegen groups by one. -What a Target validates about a program, which scheduler it reaches, and how -codegen groups by one. +What a Target validates about a program, and how codegen groups by one. The composed hardware facts themselves -- which documents a default target resolves, and which limits belong to the architecture rather than the device -- @@ -12,14 +11,13 @@ from __future__ import annotations import typing -from dataclasses import dataclass, replace +from dataclasses import replace import pytest from tests.fixtures.placed.moe_mega_kernel import MoEMegaKernel from tests.fixtures.placed.rmsnorm import RmsnormModule from tests.fixtures.placed.square_cuda import Model as SquareCudaModel -from tests.fixtures.shapes.matmul_programs import scheduling_gemm from tests.installed.smoke_target.vendor_npu import VendorNpuTarget from tilefoundry import CompilerOptions, DType, build, jit, lower, module from tilefoundry.analysis import AnalysisError, analyze @@ -29,10 +27,7 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Sequential from tilefoundry.ir.types.shard import Topology -from tilefoundry.schedule import ScheduleError, schedule -from tilefoundry.schedule.plan import SchedulePlan from tilefoundry.target import ( - AmxTarget, CpuTarget, CudaTarget, MemoryHierarchyFacts, @@ -49,7 +44,7 @@ validate_cuda_topology_levels, ) from tilefoundry.target.cuda.spec import SM90_ID -from tilefoundry.target.services import CodeGenerator, Scheduler +from tilefoundry.target.services import CodeGenerator class ExternalCudaTarget(CudaTarget): @@ -76,30 +71,8 @@ def get_facts(self, facts_type: type, query: object | None = None): return facts -@dataclass(frozen=True) -class _CustomSchedulePlan(SchedulePlan): - topology: str - - def verify(self, module, function, topology) -> None: - assert topology.name == self.topology - - def to_json(self) -> str: - return self.topology - - def render(self) -> str: - return self.topology - - -_CUSTOM_SOLVES: list[str] = [] - - -def _solve_custom_topology(module, function, target, topology, options): - _CUSTOM_SOLVES.append(topology.name) - return _CustomSchedulePlan(topology.name) - - -class CustomSchedulerCudaTarget(CudaTarget): - name = "tests.target.custom_scheduler_cuda" +class ExtraTopologyCudaTarget(CudaTarget): + name = "tests.target.extra_topology_cuda" topology_levels = (*CudaTarget.topology_levels, "custom", "unknown") def get_facts(self, facts_type: type, query: object | None = None): @@ -107,22 +80,6 @@ def get_facts(self, facts_type: type, query: object | None = None): return TopologyLimitFacts(query, 1) return super().get_facts(facts_type, query) - def get_scheduler(self, topology: str) -> Scheduler: - if topology == "custom": - return Scheduler("custom", _solve_custom_topology) - return super().get_scheduler(topology) - - -class RefusingCudaTarget(CustomSchedulerCudaTarget): - name = "tests.target.refusing_cuda" - - def get_scheduler(self, topology: str) -> Scheduler: - if topology == "thread": - raise UnsupportedCapabilityError( - f"{type(self).__name__} ({type(self).name}): no scheduler for {topology!r}" - ) - return super().get_scheduler(topology) - def _provider_target(module_name: str, provider_name: str, registered_name: str): return type( @@ -155,7 +112,6 @@ def test_target_registration_and_service_annotations_resolve() -> None: assert typing.get_type_hints(PrimFunction) for getter in ( Target.get_analyzer, - Target.get_scheduler, Target.get_code_generator, ): assert typing.get_type_hints(getter) @@ -241,7 +197,7 @@ def get_facts(self, facts_type: type, query: object | None = None): def test_cuda_mesh_topology_validation_uses_the_emission_target() -> None: - custom = CustomSchedulerCudaTarget("nvidia.h200_sxm") + custom = ExtraTopologyCudaTarget("nvidia.h200_sxm") validate_cuda_topology_levels(custom, ("custom",)) with pytest.raises(ValueError, match=r"supports \{cta, thread, custom, unknown\}"): @@ -323,78 +279,6 @@ def test_lower_rejects_a_topology_level_unsupported_by_the_target() -> None: lower(unsupported) -@pytest.mark.parametrize( - ("target", "topology", "extent"), - ( - (ExternalCudaTarget("nvidia.h200_sxm"), "thread", 128), - (AmxTarget(), "core", 1), - ), -) -def test_public_schedule_uses_inherited_target_schedulers( - target: Target, topology: str, extent: int -) -> None: - scheduled = Module( - "inherited_scheduler", - (scheduling_gemm,), - scheduling_gemm.name, - target=target, - topologies=(Topology(topology, extent),), - ) - - result = schedule(scheduled, scheduling_gemm, topology=topology) - - assert result.module is scheduled - assert result.function is scheduling_gemm - assert result.topology == Topology(topology, extent) - assert isinstance(result.plan, SchedulePlan) - - -def test_public_schedule_overrides_refuses_and_rejects_unknown_topologies( - monkeypatch: pytest.MonkeyPatch, -) -> None: - def scheduled(target: Target, topology: str) -> Module: - return Module( - "custom_scheduler", - (scheduling_gemm,), - scheduling_gemm.name, - target=target, - topologies=(Topology(topology, 1),), - ) - - _CUSTOM_SOLVES.clear() - overridden = scheduled(CustomSchedulerCudaTarget("nvidia.h200_sxm"), "custom") - result = schedule(overridden, scheduling_gemm, topology="custom") - assert result.plan == _CustomSchedulePlan("custom") - assert _CUSTOM_SOLVES == ["custom"] - - solver_calls: list[str] = [] - - def unexpected_thread_solver(*args): - solver_calls.append("thread") - raise AssertionError("refused topology reached a solver") - - monkeypatch.setattr( - "tilefoundry.target.cuda.schedule.schedule_thread", - unexpected_thread_solver, - ) - refused = scheduled(RefusingCudaTarget("nvidia.h200_sxm"), "thread") - with pytest.raises(ScheduleError) as refusal: - schedule(refused, scheduling_gemm, topology="thread") - assert str(refusal.value) == ( - "schedule: RefusingCudaTarget (tests.target.refusing_cuda): no scheduler for 'thread'" - ) - assert solver_calls == [] - - unknown = scheduled(CustomSchedulerCudaTarget("nvidia.h200_sxm"), "unknown") - with pytest.raises(ScheduleError) as unknown_error: - schedule(unknown, scheduling_gemm, topology="unknown") - assert str(unknown_error.value) == ( - "schedule: CustomSchedulerCudaTarget " - "(tests.target.custom_scheduler_cuda): no scheduler for 'unknown'" - ) - assert solver_calls == [] - - def test_program_topologies_use_target_resource_facts() -> None: """A declared extent is validated against the level's own resource fact. diff --git a/tests/target/test_target_facts.py b/tests/target/test_target_facts.py index 28394a76..bd6a603a 100644 --- a/tests/target/test_target_facts.py +++ b/tests/target/test_target_facts.py @@ -2,8 +2,6 @@ from __future__ import annotations -import subprocess -import sys from dataclasses import dataclass from typing import ClassVar @@ -109,22 +107,6 @@ def get_facts(self, facts_type: type, query: object | None = None): _DirectTarget().validate_program_topology(Topology("unit", 5)) -def test_cuda_throughput_projection_leaves_scheduler_families_unloaded() -> None: - source = """ -import sys -from tilefoundry.analysis.facts import ThroughputFacts -from tilefoundry.target import CudaTarget - -CudaTarget(\"nvidia.h200_sxm\").get_facts(ThroughputFacts) -assert \"tilefoundry.schedule.partition\" not in sys.modules -assert \"tilefoundry.schedule.pipeline\" not in sys.modules -""" - completed = subprocess.run( - (sys.executable, "-c", source), text=True, capture_output=True, check=False - ) - assert completed.returncode == 0, completed.stderr - - def test_projection_results_are_still_immutable_aggregates_of_the_requested_type() -> None: @dataclass(frozen=True) class _Facts: