diff --git a/docs/audit/backend/apple/todo.md b/docs/audit/backend/apple/todo.md index 621fd45a5..0b286b62b 100644 --- a/docs/audit/backend/apple/todo.md +++ b/docs/audit/backend/apple/todo.md @@ -8,6 +8,17 @@ last_updated: 2026-08-24 # Apple compiler, exact-device, and performance plan +Cross-backend sync `SCHEDULE-AUTHORITY-RESHARD-2026-08-24` — **shared SO-3 and W5.4 contracts assessed; Apple physical follow-up not applicable yet.** Pipeline and compound-spectral lowering now consume one digest-bound Schedule Object, inferred producer edges, roles, and resource evidence without scalar reconstruction. Placement now emits exact mesh-sized local-shard/collective SSA and executes all movement forms on the deterministic mock mesh. These Graph/Schedule/Tile contracts are portable, but x86/gfx1151 spectral numerics and mock transport transfer no Metal schedule or Apple-device claim. A future Metal distributed consumer must bind its own mesh/runtime proof. `NUMPOL-CARRIER-1` (queue row 3b) owns the generalized S5 carrier; no Apple Target change lands in this synchronization. + + +Cross-backend sync `SO3-INFER-EDGES-2026-08-24` — **shared W2.1/W5.2e +dependence-inference semantics + MegaMoE R3 producer; Apple outcome: not +applicable today, inherited on adoption.** Same reason as the nvidia entry: +host-side schedule analysis only, no Metal surface consumes it, no Apple +artifact or numerical result changes. Apple's MoE/transport families would +inherit the corrected inference if and when they adopt the overlap plan. + + Cross-backend sync `NUMPOL-CARRIER-1-2026-08-24` — **shared Schedule→Tile `numeric_policy` carrier contract (integrated-plan queue row 3b); Apple outcome: follow-up required.** Newly owned row, nothing implemented yet. diff --git a/docs/audit/backend/nvidia/todo.md b/docs/audit/backend/nvidia/todo.md index bfaa8b127..609c1726a 100644 --- a/docs/audit/backend/nvidia/todo.md +++ b/docs/audit/backend/nvidia/todo.md @@ -8,6 +8,19 @@ last_updated: 2026-08-24 # NVIDIA compiler test-suite evaluation and rearchitecture +Cross-backend sync `SCHEDULE-AUTHORITY-RESHARD-2026-08-24` — **shared SO-3 and W5.4 parity validated; no CUDA physical change.** Pipeline and compound-spectral lowering now consume one digest-bound Schedule Object, inferred producer edges, roles, and resource evidence without scalar reconstruction. Placement emits exact mesh-sized local-shard/collective SSA and all movement forms execute on the deterministic mock mesh. The carrier and verifier changes are shared; Zen 5/gfx1151 spectral evidence and mock transport transfer no CUDA schedule, NCCL proof, or RTX claim. `NUMPOL-CARRIER-1` owns the generalized S5 carrier; SM120 consumption remains a later architecture-owned assessment. + + +Cross-backend sync `SO3-INFER-EDGES-2026-08-24` — **shared W2.1/W5.2e +dependence-inference semantics + MegaMoE R3 producer; NVIDIA outcome: not +applicable today, inherited when NVIDIA adopts MoE plans.** The change is +host-side schedule ANALYSIS (Python R3 composition), not target codegen: no +NVIDIA-owned surface compiles through it and no sm_120 artifact or numerical +result changes. When an NVIDIA MoE transport lane adopts the overlap plan it +inherits the corrected inference unchanged; the exact-device evidence rules +are unaffected (no ROCm/x86 analysis result transfers a NVIDIA claim). + + Cross-backend sync `NUMPOL-CARRIER-1-2026-08-24` — **shared Schedule→Tile `numeric_policy` carrier contract (integrated-plan queue row 3b); NVIDIA outcome: follow-up required, sequenced behind W1.1.** Newly owned row, diff --git a/docs/audit/backend/rocm/todo.md b/docs/audit/backend/rocm/todo.md index e36c50bdc..b5674d15b 100644 --- a/docs/audit/backend/rocm/todo.md +++ b/docs/audit/backend/rocm/todo.md @@ -7,6 +7,27 @@ scope: ROCm backend implementation and exact-device proof # ROCm backend TODO +Cross-backend sync `SCHEDULE-AUTHORITY-RESHARD-2026-08-24` — **SO-3 exact gfx1151 regression and shared W5.4 mock boundary closed; native RCCL remains open.** Compound spectral producers now infer their fused-stage action DAG, bind roles/resources into one Schedule Object, and stamp its digest through Schedule→Tile; the complete compiled gfx1151 spectral suite remains numerically unchanged on the WSL-visible Radeon 8060S. Placement emits exact mesh-sized local-shard/collective SSA and executes every movement form on the deterministic mock mesh. Mock execution is not RCCL or multi-process evidence and cannot satisfy DIST-NATIVE-1. `NUMPOL-CARRIER-1` (queue row 3b) owns the generalized S5 carrier; no AMD Target policy promotion is implied. + + +Cross-backend sync `SO3-INFER-EDGES-2026-08-24` — **shared W2.1/W5.2e +dependence-inference semantics + MegaMoE R3 producer; ROCm outcome: parity +validated (host analysis; gfx1151 numerics unchanged).** The MegaMoE +candidate's edges are now inferred from registered Graph semantics rather +than hand-authored, and three over-conservatisms in the SHARED analysis were +corrected (they made any pipeline containing transport infer a total chain): +ordered collectives order against each other rather than against all local +work; a registered op declaring `aliasing="none"` yields a fresh alias root +even when effectful; and a pure op with a disjoint alias set has no memory +dependence on an effectful op (plus no alias edge between two pure ops). +gfx1151 owns the executing MoE transport lane, so it is the backend most +exposed: the change is analysis-only (prune/rank), selector authority remains +measured scalar latency, and no generated code or numerical result changes — +verified by the MoE transport and composition suites on this box. Exact-device +re-validation is only required if a future change makes the inferred DAG +select rather than prune. + + Cross-backend sync `NUMPOL-CARRIER-1-2026-08-24` — **shared Schedule→Tile `numeric_policy` carrier contract (integrated-plan queue row 3b); ROCm outcome: follow-up required — ROCm owns the worked reference.** The row is diff --git a/docs/audit/backend/x86/todo.md b/docs/audit/backend/x86/todo.md index b683bcc64..359f7030d 100644 --- a/docs/audit/backend/x86/todo.md +++ b/docs/audit/backend/x86/todo.md @@ -9,6 +9,18 @@ scope: x86 AVX-512 implementation/proof and AMX access planning # x86 backend TODO +Cross-backend sync `SCHEDULE-AUTHORITY-RESHARD-2026-08-24` — **SO-3 exact Zen 5 regression and shared W5.4 mock boundary closed; native multi-rank remains open.** Compound spectral producers now infer their fused-stage action DAG, bind roles/resources into one Schedule Object, and stamp its digest through Schedule→Tile; the complete compiled AVX-512 spectral suite remains numerically unchanged on the Ryzen AI Max+ 395. Placement emits exact mesh-sized local-shard/collective SSA and executes every movement form on the deterministic mock mesh. Mock execution transfers no MPI/OFI/SHMEM proof and cannot satisfy DIST-NATIVE-1. `NUMPOL-CARRIER-1` (queue row 3b) owns the generalized S5 carrier; no x86 Target policy change is implied. + + +Cross-backend sync `SO3-INFER-EDGES-2026-08-24` — **shared W2.1/W5.2e +dependence-inference semantics + MegaMoE R3 producer; x86 outcome: parity +validated (AVX-512 host).** Same shared-analysis correction as the rocm +entry. x86 hosts the R3/composition analysis itself, so this backend is the +regression gate for the change: the composition-cost, graph-dataflow, and +MegaMoE suites pass on this host against the corrected semantics, with no +generated code or numerical change (analysis is prune/rank-only). + + Cross-backend sync `NUMPOL-CARRIER-1-2026-08-24` — **shared Schedule→Tile `numeric_policy` carrier contract (integrated-plan queue row 3b); x86 outcome: follow-up required.** Newly owned row, nothing implemented yet. diff --git a/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md b/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md index 0585610eb..34672c905 100644 --- a/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md +++ b/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md @@ -69,10 +69,10 @@ This section owns what to do with those counts. |---:|---|---|---|---| | 1 | **E2E-REAL-6F — remaining optimizer VJP authority complete; family migration continues** | SGD and Momentum/Nesterov on x86/gfx1151 plus Adam/AdamW on gfx1151 use explicit non-reexecuting plugins and one typed `schedule.optimizer_vjp` → `tile.training_kernel` artifact. The package binds tracer proof, state/cotangent lineage, numeric identity, target ownership, and exact Tile digest. Unsupported target pairs fail before construction. The three former `JitFn` compatibility helpers are deleted. | One-execution/plugin tests for the three ABI shapes; Schedule/Tile verifier negatives; existing AVX-512 and gfx1151 physical regressions; runtime receives no source Graph or operation dictionary. | E2E-REAL-6E state-lineage package. | | 2 | **W4-PRODUCT-1 — executable multi-block regions** | The bounded arbitrary-CFG compiler boundary, per-slot dynamic saved-value envelopes, companion logical-shape tapes, mixed-state SAVE/HYBRID tapes, and nested canonical bodies are landed. Exact polynomial specialization guards remain outside Presburger proofs and require complete witnesses. Compiler-generated replay-safe assertions are admitted; mutation, unkeyed RNG, I/O, alias-sensitive work, and ordered collectives remain fail closed pending operation-owned recorded-product ABIs. The gfx1151 irreducible-state-machine row landed 2026-08-21: `--generate-rocm-state-machine-kernel` lowers a paired `bounded_state_machine_v1` function (forward AND generated backward) to one per-thread device kernel — per-element program counter, structured-CFG digest stamped on the gpu.func, `cf.assert` bound check host-enforced through a STATUS buffer — with both entry paths of a two-entry SCC executing on gfx1151 against the analytic oracle (`test_rocm_state_machine_exec.py`). The sibling x86 row landed 2026-08-21 as well: the same paired functions compile through `tessera_jit` (tessera-to-linalg → elementwise-to-linalg → one-shot-bufferize → loops → LLVM → ORC JIT) and execute natively on the AVX-512 host — both entry paths, forward + backward, digest/residual-policy bound, native `cf.assert` bound trap, proof-of-execution counter (`test_x86_state_machine_exec.py`). Next: one physical packet family with admissible effects. | Existing region verifier/paired VJP fixtures stay green; padded tape bounds must never replace logical extents; native x86/gfx1151 numerical rows must bind the exact CFG and residual digests before physical execution is claimed. | W2.1 dataflow, W2.2 effects, current bounded W4 carrier. | -| 3 | **SO-3 + W5.2e-PRODUCER-1 — one schedule authority** | Stamp the Schedule Object digest through lowering, delete scalar pipeline reconstruction, and make two representative physical producers—spectral and collective/MoE—consume inferred dependence edges. | Generated DAG covers hand-authored oracle edges; unknown facts conservatively serialize; emitted artifact preserves digest/roles/resources; numerical output is unchanged on x86 and gfx1151. | SO-1/SO-2 and W5.2e inference. | +| 3 | **SO-3 + W5.2e-PRODUCER-1 — one schedule authority closed 2026-08-24.** The MegaMoE producer likewise consumes `infer_action_dag` with the hand-authored plan DAG demoted to a fail-closed coverage oracle (#31), and adopting it EXPOSED three over-conservatisms in the shared inference that made any pipeline containing transport infer a total chain — ordered collectives serialized against all local work, registered effectful producers discarded their declared `aliasing="none"`, and every effectful/pure pair was assumed memory-dependent — each corrected with the sound rule and pinned by tests (a 12-action MegaMoE plan went from all 66 edges of the complete order to 36, still oracle-covering, extras collective-ordering only). Its schedule digest binds the plan's own artifact digest. | Pipeline lowering now requires the digest-bound `tessera.pipeline_schedule.v1` carrier, stamps its Schedule Object digest on functions and communication ops, and has no scalar-plan/options reconstruction path. The spectral producer represents fused stages as registered Graph actions, consumes `infer_action_dag`, binds roles/resources into the content-addressed Schedule Object, and requires that digest at Schedule→Tile. MegaMoE R3 likewise consumes inferred registered-semantics edges; its former handwritten DAG is only a fail-closed coverage oracle and additional conservative collective edges remain reported. | Inferred spectral and MegaMoE DAG tests preserve reasoned edges, roles, and resource vectors; stale/missing digest carriers fail closed; focused schedule/producer suite passes (225 tests plus dedicated carrier tests); compiled spectral numerics are unchanged across the complete Zen 5 AVX-512 and WSL gfx1151 suites (42 tests). | SO-1/SO-2 and W5.2e inference. | | 3b | **NUMPOL-CARRIER-1 — the S5 generalized `numeric_policy` carrier (owned 2026-08-24)** | One carrier design for storage/accumulator/math-mode that survives Schedule and Tile IR beyond MMA fragments (pointwise, reduction, and butterfly chains), plus the Decision #32 boundary verifier that FAILS on silent loss instead of recording it. Builds on the landed W1.1 `!tile.fragment<…, acc, …>` accumulator carrier (typed ROCm route) as the worked reference. Four mandating consumers: CAKE (#32's original derivation), game-theory §6 (fusion is a correctness feature — the zeta intermediate must not round through fp32), PDE §III.4 (interim `tessera.info_loss` records retire), and AD-JET-IR-1 (coefficient/cotangent policy, W6.3 §2.3). FORGE §1.3 supplies the measured acceptance target: the fused-epilogue fp32-accumulator realizability verdict (913× → 1.1× → 1.0× purely as a function of accum × state dtype) must be decided by the carried policy, not a special case. | Carrier attribute round-trips Schedule→Tile with a lit-verified boundary check per crossing; a lowering that drops the policy fails closed with a named diagnostic (#21a/#32); the W1.1 fragment path re-expressed as an instance of the general carrier without behavior change (bit-identical existing gfx1151/x86 outputs); the PDE `tessera.info_loss` interim records replaced by carrier facts; dashboard row tracks per-boundary coverage. May proceed in parallel with Orders 3 and 5 (orthogonal IR-carrier work; the schedule authority does not consume the policy). | W1.1 fragments (landed); Decision #32; #21a semantic-key discipline. | | 4 | **LAYOUT-ALG-1 L4 — physical layout decisions closed 2026-08-24** | L3 factorization/residency and SO-4 proof attachment are implemented. Mixed-radix/static tuple products, SM120 dynamic strided typed+macro routes, gfx1151 bounded-dynamic execution, the four x86 core GEMM index families, and every reachable Apple MSL rank-2 template consume shared authority. Dynamic non-separable tuple codomains remain fail closed. | Existing raster/index outputs remain bit-identical; unresolved layouts fail closed; materialization proof covers alias, capacity, and lifetime; Apple7 canonical and fused-cooperative cohorts pass exact-device; no architecture's schedule is promoted by another architecture's evidence. | Current L1/L3/L5 authority and architecture-owned device proofs. | -| 5 | **W5.4-RESHARD-1 — executable placement** | Consume the remaining domain/axis-changing sharding contracts, derive typed local-shard shapes, insert explicit reshard SSA through nested regions, and execute on a deterministic mock mesh. | Fixed-point convergence and join tests; exact local result types; subgroup/region negatives; all four collective forms and `collective_permute` execute without hidden host composition. | Orders 2–4. Native transports are a later evidence gate. | +| 5 | **W5.4-RESHARD-1 — executable placement closed 2026-08-24.** | Fixed-point placement now derives exact mesh-sized local result types and inserts explicit reshard SSA at the consumer boundary, including registered `tessera.slice` local shards rather than fake same-shaped collectives. Plan digest, subgroup, matching rounds, and nested region path survive Graph→Schedule→Tile. A deterministic mock-mesh executor consumes that SSA directly. | 13 focused placement tests plus 76 shared Schedule/Tile/collective tests pass. Nondivisible shapes, mesh/subgroup mismatch, unknown placement, and sibling-region escape fail closed. `local_shard`, `all_reduce`, `reduce_scatter`, `all_gather`, `all_to_all`, and `collective_permute` execute numerically with an explicit movement trace. Native transport remains a separate evidence gate. | Orders 2–4. | | 6 | **DIST-NATIVE-1 — real multi-rank execution** | Bind explicit reshard/collective SSA to NCCL/RCCL and MPI/OFI/SHMEM launchers, including subgroup communicators and process-rank ownership. Keep ROCm LSA, GIN/RMA, Copy Engine, and gfx1250 DDA as independent advanced gates. | Two-rank then multi-rank numerical packets; deterministic ordering; communicator/topology digest match; fail-closed missing transport; no mock result may satisfy the native gate. | Order 5. | | 7 | **AD-TSOL-STFT-BWD-1 — native spectral products** | Add native STFT/ISTFT backward packages on AVX-512 and gfx1151, including signal/window tangents, overlap-add identity, packed-real lineage, and normalization. | Directional and duality properties; aligned/ragged and fp32/fp16/bf16 storage; content-addressed Schedule→Tile→Target lineage; exact-host/device correctness. | Order 1 plugin boundary and existing spectral VJP carriers. | | 8 | **TSOL-POLICY-PHYS-1 — complete the spectral policy envelope** | Physically adopt centered padding, explicit transform length, one-sided/full spectrum, ISTFT output length, arbitrary axes/strides, broadcasting, and streaming/chunk state on x86 and gfx1151. | Differential oracle for every policy combination; bounded-dynamic legality; workspace/state lineage; no silent fallback to a full-complex or host-composed path. | Order 7. | @@ -1031,7 +1031,7 @@ an existing hardcoded choice through the arbiter Decision #28 already built. | W5.2b | **COMP-SCHED-OVERLAP-1/R2 — closed 2026-08-10.** Successful measured autotune rows now carry a validated `tessera.measured_resource_vector.v1` in canonical `hot_path_metadata`: compute time, dtype-correct bytes moved, communication bytes, queue/resource identity, timing provenance, and a content digest for the measured schedule candidate. Timing provenance survives SQLite warm-starts. Analytical rows cannot claim this vector, and the contract stamps `usage = composition_analysis_only` plus `selector_authority = latency_ms`; scalar measured latency remains the sole selector score. R3 is the first permitted composition consumer. | TileRT R2 + Evaluator | closed | | W5.2c | **COMP-SCHED-OVERLAP-1/R3 — v2 scalable search landed 2026-08-14.** The compiler-owned `composition_cost` boundary consumes R2 vectors as explicit Tile actions, validates DAGs with deterministic Kahn traversal, and models compute/memory/communication lanes plus queue serialization. Production search is critical-path/list scheduling. The admissible bound is `max(dependency critical path, per-resource-lane work, queue work)`; an inexact candidate is pruned only when that lower bound already loses to another candidate's feasible upper bound. Results remain promotion-ineligible with scalar `latency_ms` as selector authority. W2.1/W2.2 supply registered effect, alias, liveness, and memory-dependence facts. | TileRT R3 + TileSight T3 | closed | | W5.2d | **COMP-SCHED-OVERLAP-1/R4 — bounded functional consumer closed 2026-08-10.** MegaMoE now executes a content-addressed chunk/action plan rather than reconstructing an implicit loop policy. The plan owns contiguous chunk slices, per-chunk expert capacity, a two-live-frame workspace bound, dispatch→compute and compute→combine true-use dependencies, ordered collective dependencies, and deterministic chunk-combine order. R3 requires total measured evidence for every action and may only prune plans; exact-device scalar latency selects among retained plans before the chosen digest is passed to the runtime. Runtime telemetry records the digest, issued action order, true-use waits, combine order, and maximum live workspace. Mock multi-rank numerical and repeated bit-identity proofs pass. Native NCCL/RCCL/MPI/OFI/SHMEM transport plus architecture-owned correctness/performance packets remain open and do not transfer across targets. | TileRT R4 + TileSight T4 | closed | -| W5.2e | **Automatic dependence-edge generation — conservative parity gate landed; selector proof remains open.** `infer_action_dag` consumes a fresh W2.1 Graph snapshot and total R2 resource vectors. It emits explicit edges for SSA producers, overlapping or unknown alias sets, value-scoped memory dependence, mutation/state/I/O, registered stochastic identity, ordered collectives, unknown effects, and region boundaries. Unknown alias facts now carry an explicit reason rather than relying on a side effect of `may_alias`. `compare_inferred_action_dag` requires generated edges to cover every edge in an existing hand-authored R3 fixture; additional conservative edges are reported separately. Pure SSA matches the current fixture exactly, and opaque dataflow is serialized. `CompositionCandidate.from_graph` remains the physical-family entry point. **Open:** adopt it in remaining non-JVP physical producers and collect clean calibrated Zen 5/bare-metal gfx1151 packets before scalar measured latency may select a retained schedule. WSL and analytical vectors remain prune/rank-only. | W2.1 + TileRT R3/R4 | landing | +| W5.2e | **Automatic dependence-edge generation — conservative parity gate landed; selector proof remains open.** `infer_action_dag` consumes a fresh W2.1 Graph snapshot and total R2 resource vectors. It emits explicit edges for SSA producers, overlapping or unknown alias sets, value-scoped memory dependence, mutation/state/I/O, registered stochastic identity, ordered collectives, unknown effects, and region boundaries. Unknown alias facts now carry an explicit reason rather than relying on a side effect of `may_alias`. `compare_inferred_action_dag` requires generated edges to cover every edge in an existing hand-authored R3 fixture; additional conservative edges are reported separately. Pure SSA matches the current fixture exactly, and opaque dataflow is serialized. `CompositionCandidate.from_graph` remains the physical-family entry point. **Open (MegaMoE adopted 2026-08-24 — inferred edges with the hand DAG as fail-closed coverage oracle):** adopt it in the remaining non-JVP physical producers (spectral next) and collect clean calibrated Zen 5/bare-metal gfx1151 packets before scalar measured latency may select a retained schedule. WSL and analytical vectors remain prune/rank-only. | W2.1 + TileRT R3/R4 | landing | | W5.2f | **Tiled SSD family.** Define one content-addressed Schedule→Tile program for chunked SSD GEMM, reduction, recurrent carry, checkpoint/residency, and mutation lineage. Existing ReplaySSM and backend-specific sequence kernels are candidates/oracles, not semantic authorities. Physical lowering and promotion remain architecture-owned. | Roadmap tiled-SSD design + E2E-REAL-6 | open | | W5.2g | **Scalable action-DAG search closed 2026-08-14.** Wide DAGs use deterministic critical-path/list scheduling and never enter factorial enumeration. DAGs through eight actions retain exhaustive enumeration as the declared oracle; tests compare feasibility, modeled makespan bounds, and deterministic identity. A lower bound combines dependency critical path, per-resource-lane work, and queue work. Proven lower-bound losers may be pruned; every inexact survivor remains rank/prune-only and scalar exact-device latency retains selection authority. Remaining producer wiring stays W5.2e, not a second scheduler. | PDE plan §IV.3/IV.4 routed into TileRT R3 | closed | | W5.3 | Generic fusion region discovery over a legality oracle (a W2.1 client); keep the measured cost models | Sweep §F3 | *(folded into W5.2)* | diff --git a/docs/audit/generated/test_coverage.csv b/docs/audit/generated/test_coverage.csv index cfc3e9026..be236b0d8 100644 --- a/docs/audit/generated/test_coverage.csv +++ b/docs/audit/generated/test_coverage.csv @@ -9,9 +9,9 @@ adaptive_pool,7,0,0,7,0,fp32,directly_tested,7 direct test references add,156,46,11,202,0,bf16 f16 f32 f64 fp16 fp32 int8,directly_tested,202 direct test references add_decoupled_weight_decay,2,0,0,2,0,,directly_tested,2 direct test references alibi,7,0,0,7,0,bf16 f16 f32 fp32 int32,directly_tested,7 direct test references -all_gather,5,0,0,5,0,fp32,directly_tested,5 direct test references -all_reduce,12,2,2,14,0,f32 fp32 fp8_e4m3 int8,directly_tested,14 direct test references -all_to_all,6,0,0,6,0,fp32,directly_tested,6 direct test references +all_gather,6,0,0,6,0,fp32,directly_tested,6 direct test references +all_reduce,15,2,3,17,0,f32 fp32 fp8_e4m3 int8,directly_tested,17 direct test references +all_to_all,7,0,0,7,0,fp32,directly_tested,7 direct test references amax,5,0,0,5,0,bf16 f32,directly_tested,5 direct test references amin,5,0,0,5,0,bf16 f32,directly_tested,5 direct test references aot_export,0,0,0,0,1,,structural_only,category default for 'aot' @@ -82,7 +82,7 @@ clifford_wedge,1,0,0,1,1,,covered_by_family,category default for 'geometric_alge clip,16,0,1,16,0,fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,directly_tested,16 direct test references clip_grad_norm,10,0,0,10,0,,directly_tested,10 direct test references clip_grad_value,2,0,0,2,0,,directly_tested,2 direct test references -collective_permute,2,0,0,2,0,,directly_tested,2 direct test references +collective_permute,4,0,0,4,0,fp32,directly_tested,4 direct test references compilation_cache,2,0,0,2,0,,directly_tested,2 direct test references complex_abs,7,0,1,7,0,fp32,directly_tested,7 direct test references complex_arg,4,0,0,4,0,fp32,directly_tested,4 direct test references @@ -134,7 +134,7 @@ dataset_repeat,0,0,0,0,1,,structural_only,category default for 'data' dataset_shuffle,0,0,0,0,1,,structural_only,category default for 'data' dataset_zip,0,0,0,0,1,,structural_only,category default for 'data' dbar,0,0,0,0,1,,covered_by_family,exercised by complex differential tests -dct,49,2,2,51,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,51 direct test references +dct,51,2,2,53,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,53 direct test references ddpm_noise_pred_loss,5,0,0,5,0,,directly_tested,5 direct test references deepseek_sparse_attention,22,5,0,27,0,,directly_tested,27 direct test references denoising_score_matching_loss,13,0,1,13,0,,directly_tested,13 direct test references @@ -238,7 +238,7 @@ is_concyclic,3,0,0,3,0,,directly_tested,3 direct test references isfinite,7,0,0,7,0,f32 fp16 fp32,directly_tested,7 direct test references isinf,4,0,0,4,0,f32 fp16 fp32,directly_tested,4 direct test references isnan,7,0,0,7,0,bf16 f32 fp16 fp32 fp4_e2m1 fp8_e4m3 int32,directly_tested,7 direct test references -istft,42,5,1,47,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,47 direct test references +istft,43,5,1,48,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,48 direct test references iterable_dataset,1,0,0,1,1,,structural_only,category default for 'data' js_divergence,13,1,1,14,0,,directly_tested,14 direct test references jvp,0,3,0,3,0,,directly_tested,3 direct test references @@ -286,7 +286,7 @@ map,0,0,0,0,1,,structural_only,category default for 'control_flow' masked_categorical,10,0,0,10,0,,directly_tested,10 direct test references masked_fill,11,4,0,15,0,fp16 fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,directly_tested,15 direct test references masked_scatter,14,0,1,14,0,fp16,directly_tested,14 direct test references -matmul,484,278,23,762,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp6_e2m3 fp6_e3m2 fp8_e4m3 fp8_e5m2 int32 int64 int8 nvfp4,directly_tested,762 direct test references +matmul,485,278,23,763,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp6_e2m3 fp6_e3m2 fp8_e4m3 fp8_e5m2 int32 int64 int8 nvfp4,directly_tested,763 direct test references matrix_power,1,0,0,1,1,fp32,structural_only,unclassified — defaults to structural_only max,9,0,2,9,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,9 direct test references max_pool,6,0,0,6,0,,directly_tested,6 direct test references @@ -309,8 +309,8 @@ mod,8,0,0,8,0,bf16 f32 fp16 fp32,directly_tested,8 direct test references modified_delta_attention,11,0,0,11,0,f32 fp64,directly_tested,11 direct test references module_state_tree,3,0,0,3,0,,directly_tested,3 direct test references moe,32,0,0,32,0,fp16 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,directly_tested,32 direct test references -moe_combine,8,0,0,8,0,,directly_tested,8 direct test references -moe_dispatch,7,3,0,10,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,10 direct test references +moe_combine,10,0,0,10,0,,directly_tested,10 direct test references +moe_dispatch,9,3,0,12,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,12 direct test references moe_swiglu_block,14,19,1,33,0,bf16 fp8_e4m3 nvfp4,directly_tested,33 direct test references momentum,11,1,0,12,0,fp64,directly_tested,12 direct test references mor_partition,9,3,1,12,0,fp16 fp32,directly_tested,12 direct test references @@ -322,7 +322,7 @@ msa_index_scores,12,4,1,16,0,,directly_tested,16 direct test references msa_select_blocks,8,6,1,14,0,,directly_tested,14 direct test references msa_sparse_attention,48,11,1,59,0,,directly_tested,59 direct test references mse_loss,29,0,5,29,0,f32 fp32 int64,directly_tested,29 direct test references -mul,143,41,11,184,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 int8 nvfp4,directly_tested,184 direct test references +mul,146,41,12,187,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 int8 nvfp4,directly_tested,187 direct test references multi_head_attention,9,0,1,9,0,,directly_tested,9 direct test references muon,15,0,0,15,0,fp32,directly_tested,15 direct test references named_sharding,0,0,0,0,1,,structural_only,category default for 'sharding' @@ -371,7 +371,7 @@ quantized_matmul,5,0,0,5,0,f16,directly_tested,5 direct test references rearrange,13,0,1,13,0,fp16 fp32,directly_tested,13 direct test references reciprocal,5,0,0,5,0,bf16 f32 fp16 fp32,directly_tested,5 direct test references reduce,135,17,11,152,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp6_e2m3 fp6_e3m2 fp8_e4m3 fp8_e5m2 int32 int8 nvfp4,directly_tested,152 direct test references -reduce_scatter,3,0,0,3,0,fp32,directly_tested,3 direct test references +reduce_scatter,4,0,0,4,0,fp32,directly_tested,4 direct test references relu,132,37,11,169,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp8_e4m3 fp8_e5m2 int8 nvfp4,directly_tested,169 direct test references remat,0,0,0,0,1,,structural_only,category default for 'transform' repeat,4,0,0,4,0,,directly_tested,4 direct test references @@ -425,7 +425,7 @@ seq2seq_loss,7,0,0,7,0,,directly_tested,7 direct test references sgd,21,10,0,31,0,fp64,directly_tested,31 direct test references shard_map,0,0,0,0,1,,structural_only,category default for 'sharding' sharded_dataset,2,0,0,2,0,,directly_tested,2 direct test references -sigmoid,15,9,1,24,0,bf16 f16 f32 fp32,directly_tested,24 direct test references +sigmoid,17,9,1,26,0,bf16 f16 f32 fp32,directly_tested,26 direct test references sigmoid_safe,6,0,0,6,0,bf16,directly_tested,6 direct test references sign,6,0,1,6,0,bf16 f32 fp16 fp32,directly_tested,6 direct test references silu,91,8,6,99,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp6_e2m3 fp8_e4m3 int8 nvfp4,directly_tested,99 direct test references @@ -433,7 +433,7 @@ silu_mul,18,13,1,31,0,bf16 fp32,directly_tested,31 direct test references simple_rnn_cell,10,0,0,10,0,bf16 f16 f32 fp16,directly_tested,10 direct test references sin,9,1,1,10,0,bf16 f16 f32,directly_tested,10 direct test references sinh,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references -slice,11,0,0,11,0,,directly_tested,11 direct test references +slice,12,0,0,12,0,fp32,directly_tested,12 direct test references smooth_l1_loss,9,0,0,9,0,f32,directly_tested,9 direct test references soft_top_k,1,0,0,1,1,fp64,directly_tested,"public forward plus exact-k, domain, and FD-VJP tests" softcap,11,0,1,11,0,bf16,directly_tested,11 direct test references @@ -448,8 +448,8 @@ sparsemax,1,0,0,1,1,fp64,directly_tested,public forward plus simplex/sparsity an spec_accept,2,10,0,12,0,,directly_tested,12 direct test references spec_accept_sample,3,7,1,10,0,,directly_tested,10 direct test references spec_accept_tree_sample,2,6,0,8,0,,directly_tested,8 direct test references -spectral_conv,23,5,1,28,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,28 direct test references -spectral_filter,19,9,2,28,0,bf16 f16 f32 fp32,directly_tested,28 direct test references +spectral_conv,25,5,1,30,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,30 direct test references +spectral_filter,23,9,3,32,0,bf16 f16 f32 fp32,directly_tested,32 direct test references spectral_norm,5,0,0,5,0,fp16,directly_tested,5 direct test references split,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only spmm_coo,8,0,0,8,0,,directly_tested,8 direct test references @@ -464,7 +464,7 @@ state_migration,1,0,1,1,1,,structural_only,category default for 'serialization' state_partition,3,0,1,3,0,,directly_tested,3 direct test references std,8,0,0,8,0,bf16 f32,directly_tested,8 direct test references stereographic,11,0,0,11,0,fp32,directly_tested,11 direct test references -stft,32,2,2,34,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,34 direct test references +stft,33,2,2,35,0,bf16 f16 f32 fp32 fp64 int32 int64 int8,directly_tested,35 direct test references stop_gradient,1,5,0,6,0,fp32,directly_tested,6 direct test references sub,35,7,4,42,0,bf16 f16 f32 fp16 fp32,directly_tested,42 direct test references sum,28,0,3,28,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,28 direct test references @@ -472,7 +472,7 @@ svd,24,5,1,29,0,bf16 f16 f32 fp16 fp32 fp64 fp8_e4m3 int32 int64 int8 nvfp4,dire switch,0,0,0,0,1,,structural_only,category default for 'control_flow' take,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only tan,4,0,0,4,0,bf16 f32,directly_tested,4 direct test references -tanh,18,25,2,43,0,bf16 f16 f32,directly_tested,43 direct test references +tanh,19,25,3,44,0,bf16 f16 f32,directly_tested,44 direct test references target_verify,3,5,1,8,0,,directly_tested,8 direct test references tile,4,0,0,4,0,,directly_tested,4 direct test references tile_view,6,0,0,6,0,fp16,directly_tested,6 direct test references diff --git a/docs/audit/generated/test_coverage.md b/docs/audit/generated/test_coverage.md index e0e1b59f7..fa51b6764 100644 --- a/docs/audit/generated/test_coverage.md +++ b/docs/audit/generated/test_coverage.md @@ -7,21 +7,21 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by ## Headline - **515** ops in `primitive_coverage` registry. -- **5610** total Python-test references, **1338** total lit-fixture references. +- **5640** total Python-test references, **1338** total lit-fixture references. - **70** ops have **zero** references in either test surface. - **114** ops have ≤1 reference ("thinly tested"). -- **183** ops have ≥10 references ("well tested"). +- **184** ops have ≥10 references ("well tested"). - **175** ops have at least one associated `pytest.raises` negative test. ## Top 20 most-tested ops | Op | py refs | lit refs | total | neg | dtypes | |----|--------:|---------:|------:|----:|--------| -| `matmul` | 484 | 278 | 762 | 23 | `bf16`, `f16`, `f32`, `f64` … | +| `matmul` | 485 | 278 | 763 | 23 | `bf16`, `f16`, `f32`, `f64` … | | `flash_attn` | 142 | 79 | 221 | 10 | `bf16`, `f16`, `f32`, `f64` … | | `add` | 156 | 46 | 202 | 11 | `bf16`, `f16`, `f32`, `f64` … | | `softmax` | 150 | 50 | 200 | 31 | `bf16`, `f16`, `f32`, `f64` … | -| `mul` | 143 | 41 | 184 | 11 | `bf16`, `f16`, `f32`, `f64` … | +| `mul` | 146 | 41 | 187 | 12 | `bf16`, `f16`, `f32`, `f64` … | | `relu` | 132 | 37 | 169 | 11 | `bf16`, `f16`, `f32`, `f64` … | | `reduce` | 135 | 17 | 152 | 11 | `bf16`, `f16`, `f32`, `f64` … | | `rmsnorm` | 121 | 18 | 139 | 6 | `bf16`, `f16`, `f32`, `f64` … | @@ -36,7 +36,7 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by | `msa_sparse_attention` | 48 | 11 | 59 | 1 | | | `cholesky` | 28 | 30 | 58 | 0 | `bf16`, `f16`, `f32`, `fp16` … | | `grouped_gemm` | 31 | 24 | 55 | 2 | `bf16`, `fp32`, `fp4_e2m1`, `fp8_e4m3` … | -| `linear_attn` | 44 | 8 | 52 | 2 | | +| `dct` | 51 | 2 | 53 | 2 | `bf16`, `f16`, `f32`, `fp32` … | ## Thinly-tested ops (≤1 reference) diff --git a/python/tessera/compiler/composition_cost.py b/python/tessera/compiler/composition_cost.py index d411b47c1..300b510c5 100644 --- a/python/tessera/compiler/composition_cost.py +++ b/python/tessera/compiler/composition_cost.py @@ -136,17 +136,47 @@ def add(before: int, after: int, reason: str) -> None: current_unknown = _unknown_dataflow_reasons(analysis, op) for before in range(index): previous = ops[before] + previous_barriers = _barrier_reasons(previous) for reason in current_barriers: - add(before, index, reason) - for reason in _barrier_reasons(previous): - add(before, index, reason) + if reason != ORDERED_COLLECTIVE_REASON: + add(before, index, reason) + for reason in previous_barriers: + if reason != ORDERED_COLLECTIVE_REASON: + add(before, index, reason) + # An ordered collective constrains the ORDER OF COLLECTIVES — every + # rank must issue them in the same relative order — not the order of + # unrelated local work. Serializing a collective against every + # surrounding operation makes any pipeline with per-chunk transport + # a total chain, which erases exactly the compute/communication + # overlap a schedule exists to express (measured on the MegaMoE + # producer: 12 actions, 66 inferred edges = the complete order; + # PR #625 review). Data flowing through a collective is still + # ordered by the SSA/alias/memory-dependence edges below, and any + # collective that also carries mutation/state/I/O keeps its + # all-pairs barrier above. + if ( + ORDERED_COLLECTIVE_REASON in current_barriers + and ORDERED_COLLECTIVE_REASON in previous_barriers + ): + add(before, index, ORDERED_COLLECTIVE_REASON) for reason in current_unknown: add(before, index, reason) for reason in _unknown_dataflow_reasons(analysis, previous): add(before, index, reason) if analysis.has_memory_dependence(previous, op): add(before, index, "memory_dependence") - if _may_share_alias(analysis, previous, op): + # Two REGISTERED pure ops cannot conflict: neither writes, so + # sharing a read-only operand (a weight tensor every expert reads) + # is not a dependence. Ordering them on that alias made every + # consumer of a common input sequential (PR #625 review). Any + # write-through-alias case has at least one effectful side and + # still takes the edge below. + both_pure = ( + registered_op_effect(previous.op_name, previous.kwargs) + == Effect.pure + and registered_op_effect(op.op_name, op.kwargs) == Effect.pure + ) + if not both_pure and _may_share_alias(analysis, previous, op): add(before, index, "alias_set") dependencies = tuple( @@ -225,6 +255,15 @@ def compare_inferred_action_dag( ) +ORDERED_COLLECTIVE_REASON = "ordered_collective" +"""Edge reason whose scope is collective-to-collective, not all-pairs. + +Every other barrier reason (unregistered effect, mutation/state/I/O, +stochastic identity, region boundary) serializes against ALL surrounding +work and stays fail-closed; this one is an ordering relation among the +ordered collectives themselves.""" + + def _ssa(name: str) -> str: return str(name).strip().lstrip("%") @@ -240,7 +279,7 @@ def _barrier_reasons(op: IROp) -> tuple[str, ...]: ) or op.kwargs.get("tessera.stochastic_identity") not in (None, "none"): result.add("stochastic_identity") if effect == Effect.collective and op.kwargs.get("ordered", True): - result.add("ordered_collective") + result.add(ORDERED_COLLECTIVE_REASON) if effect in {Effect.state, Effect.memory, Effect.io}: result.add("mutation_or_effect") if op.op_name.startswith(("tessera.scf.", "scf.")) or any( diff --git a/python/tessera/compiler/graph_dataflow.py b/python/tessera/compiler/graph_dataflow.py index c82959849..a667a7782 100644 --- a/python/tessera/compiler/graph_dataflow.py +++ b/python/tessera/compiler/graph_dataflow.py @@ -125,10 +125,24 @@ def recompute(self) -> "GraphDataflow": roots = operand_facts[operand_index].alias_roots except (ValueError, IndexError): roots = None - elif effect == Effect.pure and spec is not None: + elif spec is not None and str(aliasing or "none") == "none": + # A REGISTERED op that declares it aliases nothing produces + # a fresh value, whether or not it is effectful (Decision + # #30 — derive from the declaration; #29 — a declaration + # must have a consumer, and this is it). Treating every + # effectful producer's result as possibly-hidden-state made + # `unknown_alias_fact` an all-pairs barrier, so any pipeline + # containing transport inferred a TOTAL chain and lost the + # overlap it exists to express (measured: MegaMoE, 12 + # actions → 66 edges = complete order; PR #625 review). + # The op's effect still orders it: mutation/state/I/O keep + # their all-pairs barrier, ordered collectives order against + # each other, and value-scoped memory dependence is computed + # separately below. roots = frozenset({result_name}) else: - # Unregistered/effectful producers may alias hidden state. + # Unregistered producers, or a declared operand alias we + # could not resolve, may alias hidden state. roots = None facts[result_name] = ValueFact(result_shape, roots, False) producers[result_name] = op @@ -195,10 +209,32 @@ def has_memory_dependence(self, lhs: IROp, rhs: IROp) -> bool: right = registered_op_effect(rhs.op_name, rhs.kwargs) if left == Effect.pure and right == Effect.pure: return False - # IROp currently has no value-scoped MemoryEffectOpInterface mirror. - # Effectful pairs therefore remain top instead of guessing from names. + if (left == Effect.pure or right == Effect.pure) and not self.values_may_alias( + lhs, rhs + ): + # One side is REGISTERED pure: by definition it touches no hidden + # state, so the only way it can depend on the effectful side + # through memory is via a value they share — and their alias sets + # are provably disjoint here. Without this, every effectful op + # serialized against all surrounding pure work and any pipeline + # containing transport inferred a total chain (PR #625 review). + return False + # Two effectful ops: IROp still has no value-scoped + # MemoryEffectOpInterface mirror, so they remain top rather than + # guessing from names. Their relative order is preserved. return True + def values_may_alias(self, lhs: IROp, rhs: IROp) -> bool: + """True if any value of `lhs` may alias any value of `rhs`.""" + + left_values = tuple(lhs.operands) + tuple(lhs.result_names) + right_values = tuple(rhs.operands) + tuple(rhs.result_names) + return any( + self.may_alias(left, right) + for left in left_values + for right in right_values + ) + def snapshot(self) -> Mapping[str, ValueFact]: return dict(self._facts) if self.valid else {} diff --git a/python/tessera/compiler/megamoe_overlap.py b/python/tessera/compiler/megamoe_overlap.py index da7f6b2ef..469efb81c 100644 --- a/python/tessera/compiler/megamoe_overlap.py +++ b/python/tessera/compiler/megamoe_overlap.py @@ -7,15 +7,21 @@ from dataclasses import dataclass from typing import Any, Mapping, Sequence +import heapq + from .composition_cost import ( EXHAUSTIVE_ORACLE_MAX_ORDERS, + ActionDAGParity, CompositionCalibration, CompositionCandidate, CompositionPruningResult, + InferredActionDAG, TileAction, + compare_inferred_action_dag, prune_composition_candidates, ) - +from .graph_ir import GraphIRFunction, IRArg, IROp, IRType +from .schedule_object import ScheduleObject MEGAMOE_OVERLAP_PLAN_SCHEMA = "tessera.megamoe_overlap_plan.v1" @@ -78,9 +84,13 @@ def __post_init__(self) -> None: for kind in ("dispatch", "compute", "combine") } if set(self.action_dependencies) != expected_actions: - raise ValueError("MegaMoE action DAG must cover every chunk phase exactly once") + raise ValueError( + "MegaMoE action DAG must cover every chunk phase exactly once" + ) for action, dependencies in self.action_dependencies.items(): - if action in dependencies or not set(dependencies).issubset(expected_actions): + if action in dependencies or not set(dependencies).issubset( + expected_actions + ): raise ValueError("MegaMoE action DAG has invalid dependencies") _require_digest(self.artifact_digest) @@ -116,13 +126,15 @@ def build_megamoe_overlap_plan( start = 0 for index in range(actual_chunks): stop = start + q + (1 if index < remainder else 0) - chunks.append(MegaMoEChunk( - index=index, - start=start, - stop=stop, - expert_capacity=int(capacities[index]), - dispatch_buffer_bytes=int(dispatch_buffer_bytes[index]), - )) + chunks.append( + MegaMoEChunk( + index=index, + start=start, + stop=stop, + expert_capacity=int(capacities[index]), + dispatch_buffer_bytes=int(dispatch_buffer_bytes[index]), + ) + ) start = stop live_limit = min(max_in_flight_chunks, actual_chunks) if live_limit < 1: @@ -133,7 +145,10 @@ def build_megamoe_overlap_plan( "precede combine(c)" ) live_bytes = max( - sum(chunk.dispatch_buffer_bytes for chunk in chunks[offset:offset + live_limit]) + sum( + chunk.dispatch_buffer_bytes + for chunk in chunks[offset : offset + live_limit] + ) for offset in range(actual_chunks) ) if workspace_capacity_bytes is not None: @@ -172,11 +187,111 @@ def build_megamoe_overlap_plan( ) -def composition_candidate_for_megamoe_plan( +_PHASE_RANK = {"dispatch": 0, "compute": 1, "combine": 2} +_MOE_TENSOR = "tensor<*xf32>" + + +def megamoe_issue_order(plan: MegaMoEOverlapPlan) -> tuple[str, ...]: + """Deterministic topological issue order of the plan's action DAG. + + Kahn traversal with a (chunk, phase) heap tie-break — the order every + consumer shares, so the Graph representation, the inferred DAG, the + resource-vector binding, and the parity oracle all index one sequence. + """ + + dependencies = plan.action_dependencies + indegree = {action: len(deps) for action, deps in dependencies.items()} + successors: dict[str, list[str]] = {} + for action, deps in dependencies.items(): + for dep in deps: + successors.setdefault(dep, []).append(action) + + def key(action: str) -> tuple[int, int]: + kind, chunk = action.split(":") + return (int(chunk), _PHASE_RANK[kind]) + + heap = [(key(action), action) for action, deg in indegree.items() if deg == 0] + heapq.heapify(heap) + order: list[str] = [] + while heap: + _, action = heapq.heappop(heap) + order.append(action) + for successor in successors.get(action, []): + indegree[successor] -= 1 + if indegree[successor] == 0: + heapq.heappush(heap, (key(successor), successor)) + if len(order) != len(dependencies): + raise ValueError("MegaMoE action DAG has a cycle") + return tuple(order) + + +def megamoe_graph_function(plan: MegaMoEOverlapPlan) -> GraphIRFunction: + """The chunk pipeline as Graph IR, in the deterministic issue order. + + dispatch/combine are the registered ordered collectives + (`tessera.moe_dispatch` / `tessera.moe_combine`); expert compute is pure + `tessera.matmul` consuming its own chunk's dispatch result. Edges are NOT + written here — `infer_action_dag` derives them from these registered + semantics (SSA flow + ordered-collective serialization), which is what + makes the R3 candidate a consumer of W2.1/W2.2 facts (SO-3).""" + + tensor = IRType(_MOE_TENSOR) + ops: list[IROp] = [] + for action_id in megamoe_issue_order(plan): + kind, chunk = action_id.split(":") + if kind == "dispatch": + ops.append( + IROp( + result=f"%dispatch{chunk}", + op_name="tessera.moe_dispatch", + operands=["%routed"], + operand_types=[_MOE_TENSOR], + result_type=_MOE_TENSOR, + ) + ) + elif kind == "compute": + ops.append( + IROp( + result=f"%compute{chunk}", + op_name="tessera.matmul", + operands=[f"%dispatch{chunk}", "%weights"], + operand_types=[_MOE_TENSOR, _MOE_TENSOR], + result_type=_MOE_TENSOR, + ) + ) + else: + ops.append( + IROp( + result=f"%combine{chunk}", + op_name="tessera.moe_combine", + operands=[f"%compute{chunk}"], + operand_types=[_MOE_TENSOR], + result_type=_MOE_TENSOR, + ) + ) + return GraphIRFunction( + name=plan.plan_id, + args=[IRArg("routed", tensor), IRArg("weights", tensor)], + body=ops, + return_values=[f"%combine{plan.num_chunks - 1}"], + ) + + +def megamoe_inferred_composition( plan: MegaMoEOverlapPlan, benchmark_rows: Mapping[str, Mapping[str, Any]], -) -> CompositionCandidate: - """Bind measured R2 rows to every action in one executable plan.""" +) -> tuple[CompositionCandidate, InferredActionDAG, ActionDAGParity]: + """SO-3: the production candidate's edges are INFERRED, hand edges are + the coverage oracle. + + The plan's hand-authored `action_dependencies` stop being the executed + DAG and become the declared oracle (#31): `compare_inferred_action_dag` + must show the generated edges cover every hand edge, or construction + fails closed — a generation regression can never silently weaken the + schedule. Additional conservative edges (the ordered-collective total + order) are reported separately on the parity result, per the SO-3 + acceptance. The returned candidate carries the content-addressed + `ScheduleObject` (digest included) built from the inferred edges.""" expected = set(plan.action_dependencies) if set(benchmark_rows) != expected: @@ -185,15 +300,59 @@ def composition_candidate_for_megamoe_plan( raise ValueError( f"MegaMoE action evidence must be total; missing={missing}, extra={extra}" ) - actions = tuple( + order = megamoe_issue_order(plan) + vectors = [ TileAction.from_benchmark_row( + action_id, benchmark_rows[action_id] + ).resource_vector + for action_id in order + ] + candidate, inferred = CompositionCandidate.from_graph( + plan.plan_id, megamoe_graph_function(plan), vectors, action_ids=order + ) + # Bind the PLAN's identity into the schedule identity (PR #625 review): + # the action ids and the graph shape are functions of the chunk COUNT, so + # two plans that differ in token ranges, expert capacities, dispatch-buffer + # sizes, or the in-flight limit would otherwise content-address to the same + # digest with identical benchmark rows. `plan.artifact_digest` already + # content-addresses exactly those fields, so the schedule object carries it. + schedule = ScheduleObject( + f"{plan.plan_id}@{plan.artifact_digest}", + inferred.schedule_object.actions, + inferred.schedule_object.edges, + inferred.schedule_object.roles, + inferred.schedule_object.residency, + ) + candidate = CompositionCandidate(plan.plan_id, candidate.actions, schedule) + reference = tuple( + TileAction( action_id, - benchmark_rows[action_id], - depends_on=plan.action_dependencies[action_id], + vectors[index], + tuple(plan.action_dependencies[action_id]), ) - for action_id in sorted(expected) + for index, action_id in enumerate(order) ) - return CompositionCandidate(plan.plan_id, actions) + parity = compare_inferred_action_dag(inferred, reference) + if not parity.conservative: + raise ValueError( + "inferred MegaMoE action DAG does not cover the hand-authored " + f"oracle edges: missing={parity.missing_reference_edges}" + ) + return candidate, inferred, parity + + +def composition_candidate_for_megamoe_plan( + plan: MegaMoEOverlapPlan, + benchmark_rows: Mapping[str, Mapping[str, Any]], +) -> CompositionCandidate: + """Bind measured R2 rows to every action in one executable plan. + + Since SO-3 the returned candidate's dependencies are inferred from the + registered Graph semantics (see `megamoe_inferred_composition`); the + hand-authored plan DAG is the fail-closed coverage oracle.""" + + candidate, _, _ = megamoe_inferred_composition(plan, benchmark_rows) + return candidate def prune_megamoe_overlap_plans( @@ -259,5 +418,8 @@ def _require_digest(value: str) -> None: "MegaMoEOverlapPlan", "build_megamoe_overlap_plan", "composition_candidate_for_megamoe_plan", + "megamoe_graph_function", + "megamoe_inferred_composition", + "megamoe_issue_order", "prune_megamoe_overlap_plans", ] diff --git a/python/tessera/compiler/pass_metadata.py b/python/tessera/compiler/pass_metadata.py index 376bedd9f..d06997420 100644 --- a/python/tessera/compiler/pass_metadata.py +++ b/python/tessera/compiler/pass_metadata.py @@ -649,7 +649,8 @@ class PassMetadata: ), input_dialects=("tessera", "func"), output_dialects=("tessera", "func"), - required_attrs=("tessera.pipeline_plan", "tessera.pp_stage"), + required_attrs=("tessera.schedule_digest", "tessera.pipeline_steps", + "tessera.pp_num_stages", "tessera.pp_stage"), diagnostic_codes=(), pass_kind="transform", sprint="Pipeline-PP", @@ -659,13 +660,13 @@ class PassMetadata: cpp_class="PipelineScheduleLegalityPass", summary=( "Proves 1F1B micro-batch fill, stage occupancy, send/recv pairing, " - "and value-rewrite completeness, then materializes an explicit " - "warmup/steady/cooldown dependency order in " - "tessera.pipeline_steps." + "value-rewrite completeness, and the producer-materialized " + "Schedule Object dependency carrier without scalar reconstruction." ), input_dialects=("tessera", "func"), output_dialects=("tessera", "func"), - required_attrs=("tessera.pp_num_stages", "tessera.pp_num_micro_batches", + required_attrs=("tessera.schedule_digest", "tessera.pipeline_steps", + "tessera.pp_num_stages", "tessera.pp_num_micro_batches", "tessera.pp_stage"), diagnostic_codes=( "PP_EMPTY_STAGE", @@ -681,13 +682,15 @@ class PassMetadata: name="tessera-pipeline-stage-insertion", cpp_class="PipelineStageInsertionPass", summary=( - "Inserts tessera.pipeline.send/recv at cross-stage boundaries and " - "rewires the boundary uses to the recv (the real send/recv SSA " - "rewrite), driven by the tessera.pp_stage partition." + "Consumes the digest-bound Schedule Object carrier, inserts " + "tessera.pipeline.send/recv at cross-stage boundaries, stamps the " + "digest on emitted IR, and rewires boundary uses to recv results." ), input_dialects=("tessera", "func"), output_dialects=("tessera", "func"), - required_attrs=("tessera.pp_stage", "tessera.layer"), + required_attrs=("tessera.schedule_digest", "tessera.pipeline_steps", + "tessera.pp_num_stages", "tessera.pp_num_micro_batches", + "tessera.pp_stage", "tessera.layer"), preserved_attrs=("tessera.pp_stage",), diagnostic_codes=(), must_run_after=("tessera-pipeline-partition",), diff --git a/python/tessera/compiler/pipeline_planner.py b/python/tessera/compiler/pipeline_planner.py index 1f0d81364..992f6f064 100644 --- a/python/tessera/compiler/pipeline_planner.py +++ b/python/tessera/compiler/pipeline_planner.py @@ -33,19 +33,23 @@ """ from __future__ import annotations +import hashlib +import json from dataclasses import dataclass from enum import Enum from typing import List, NamedTuple, Optional - +from .benchmark_row import MeasuredResourceVector +from .schedule_object import ScheduleAction, ScheduleObject, ScheduleRole # ───────────────────────────────────────────────────────────────────────────── # Schedule step # ───────────────────────────────────────────────────────────────────────────── + class Phase(Enum): - FORWARD = "F" + FORWARD = "F" BACKWARD = "B" - IDLE = "_" # bubble + IDLE = "_" # bubble class ScheduleStep(NamedTuple): @@ -59,6 +63,7 @@ class ScheduleStep(NamedTuple): micro_batch : micro-batch index (0-based) phase : FORWARD or BACKWARD """ + clock: int rank: int stage: int @@ -70,6 +75,7 @@ class ScheduleStep(NamedTuple): # PipelinePlan # ───────────────────────────────────────────────────────────────────────────── + @dataclass class PipelinePlan: """ @@ -94,6 +100,7 @@ class PipelinePlan: assert step0.phase == Phase.FORWARD assert step0.micro_batch == 0 """ + num_stages: int num_micro_batches: int interleaved: bool = False @@ -104,16 +111,16 @@ def __post_init__(self) -> None: if self.num_stages < 1: raise ValueError(f"num_stages must be >= 1, got {self.num_stages}") if self.num_micro_batches < 1: - raise ValueError(f"num_micro_batches must be >= 1, got {self.num_micro_batches}") + raise ValueError( + f"num_micro_batches must be >= 1, got {self.num_micro_batches}" + ) if self.decoupled and self.interleaved: raise ValueError( "decoupled and interleaved schedules are mutually exclusive" ) if self.interleaved: if self.num_chunks < 2: - raise ValueError( - "Interleaved 1F1B requires num_chunks >= 2" - ) + raise ValueError("Interleaved 1F1B requires num_chunks >= 2") min_m = self.num_stages * self.num_chunks if self.num_micro_batches < min_m: raise ValueError( @@ -199,14 +206,24 @@ def _build_decoupled(self) -> List[ScheduleStep]: steps: List[ScheduleStep] = [] for mb in range(m): for rank in range(p): - steps.append(ScheduleStep( - clock=2 * mb, rank=rank, stage=rank, - micro_batch=mb, phase=Phase.FORWARD, - )) - steps.append(ScheduleStep( - clock=2 * mb + 1, rank=rank, stage=rank, - micro_batch=mb, phase=Phase.BACKWARD, - )) + steps.append( + ScheduleStep( + clock=2 * mb, + rank=rank, + stage=rank, + micro_batch=mb, + phase=Phase.FORWARD, + ) + ) + steps.append( + ScheduleStep( + clock=2 * mb + 1, + rank=rank, + stage=rank, + micro_batch=mb, + phase=Phase.BACKWARD, + ) + ) steps.sort(key=lambda s: (s.clock, s.rank)) return steps @@ -217,10 +234,10 @@ def _build_standard(self) -> List[ScheduleStep]: steps: List[ScheduleStep] = [] # Per-rank state: next micro-batch to forward and backward - fwd_mb = list(range(p)) # rank k starts at micro-batch k - bwd_mb = [0] * p # backward starts after warmup + fwd_mb = list(range(p)) # rank k starts at micro-batch k + bwd_mb = [0] * p # backward starts after warmup fwd_ptr = [0] * p - bwd_ptr = [-1] * p # -1 = not yet started + bwd_ptr = [-1] * p # -1 = not yet started # We use a simpler direct formulation: # For rank r (stage r), forward of micro-batch m starts at clock (r + m) @@ -230,18 +247,28 @@ def _build_standard(self) -> List[ScheduleStep]: for rank in range(p): # Forward fwd_clock = rank + mb - steps.append(ScheduleStep( - clock=fwd_clock, rank=rank, stage=rank, - micro_batch=mb, phase=Phase.FORWARD, - )) + steps.append( + ScheduleStep( + clock=fwd_clock, + rank=rank, + stage=rank, + micro_batch=mb, + phase=Phase.FORWARD, + ) + ) # Backward (mirrored: last stage finishes backward first) # In standard 1F1B, backward of mb on rank r starts at: # clock = (p - 1 - rank) + mb + p bwd_clock = (p - 1 - rank) + mb + p - steps.append(ScheduleStep( - clock=bwd_clock, rank=rank, stage=rank, - micro_batch=mb, phase=Phase.BACKWARD, - )) + steps.append( + ScheduleStep( + clock=bwd_clock, + rank=rank, + stage=rank, + micro_batch=mb, + phase=Phase.BACKWARD, + ) + ) steps.sort(key=lambda s: (s.clock, s.rank)) return steps @@ -264,14 +291,24 @@ def _build_interleaved(self) -> List[ScheduleStep]: virtual_stage = rank + chunk * p fwd_clock = rank + mb + chunk * p bwd_clock = fwd_clock + p * v - steps.append(ScheduleStep( - clock=fwd_clock, rank=rank, stage=virtual_stage, - micro_batch=mb, phase=Phase.FORWARD, - )) - steps.append(ScheduleStep( - clock=bwd_clock, rank=rank, stage=virtual_stage, - micro_batch=mb, phase=Phase.BACKWARD, - )) + steps.append( + ScheduleStep( + clock=fwd_clock, + rank=rank, + stage=virtual_stage, + micro_batch=mb, + phase=Phase.FORWARD, + ) + ) + steps.append( + ScheduleStep( + clock=bwd_clock, + rank=rank, + stage=virtual_stage, + micro_batch=mb, + phase=Phase.BACKWARD, + ) + ) steps.sort(key=lambda s: (s.clock, s.rank)) return steps @@ -296,8 +333,7 @@ def render_ascii(self, max_clocks: Optional[int] = None) -> str: # Build grid[rank][clock] = label grid: List[List[str]] = [ - ["__"] * (max_clock + 1) - for _ in range(self.num_stages) + ["__"] * (max_clock + 1) for _ in range(self.num_stages) ] for step in steps: if step.clock > max_clock: @@ -311,14 +347,184 @@ def render_ascii(self, max_clocks: Optional[int] = None) -> str: return "\n".join(lines) def to_mlir_attrs(self) -> str: - """Serialize for PipelineStageInsertionPass.""" + """Materialize the digest-bound Schedule Object IR carrier. + + Lowering consumes this carrier directly instead of rebuilding a 1F1B + schedule from a parallel scalar ``pipeline_plan``. Resource vectors + remain out-of-band in :attr:`schedule_object`; IR carries the digest + and the dependency/phase view needed by the pipeline passes. + """ + schedule = self.schedule_object + step_by_id = {self._action_id(step): step for step in self.schedule_steps()} + rendered_steps = [] + for action in schedule.actions: + step = step_by_id[action.action_id] + dependencies = ", ".join( + f'"{dependency}"' for dependency in action.depends_on + ) + rendered_steps.append( + "{" + f'action_id = "{action.action_id}", ' + f"clock = {step.clock}, " + f"depends_on = [{dependencies}], " + f"micro_batch = {step.micro_batch}, " + f'phase = "{step.phase.value}", ' + f"rank = {step.rank}, " + f"stage = {step.stage}" + "}" + ) return ( - f'{{tessera.pipeline_plan = {{' - f'num_stages = {self.num_stages}, ' - f'num_micro_batches = {self.num_micro_batches}, ' - f'interleaved = {"true" if self.interleaved else "false"}, ' - f'decoupled = {"true" if self.decoupled else "false"}, ' - f'num_chunks = {self.num_chunks}}}}}' + "{" + f'tessera.schedule_digest = "{schedule.digest}", ' + 'tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", ' + f'tessera.pipeline_steps = [{", ".join(rendered_steps)}], ' + f"tessera.pp_num_stages = {self.num_stages}, " + f"tessera.pp_num_micro_batches = {self.num_micro_batches}, " + f'tessera.pp_interleaved = {"true" if self.interleaved else "false"}, ' + f'tessera.pp_decoupled = {"true" if self.decoupled else "false"}, ' + f"tessera.pp_num_chunks = {self.num_chunks}" + "}" + ) + + @staticmethod + def _action_id(step: ScheduleStep) -> str: + return ( + f"pipeline:{step.clock}:{step.rank}:{step.stage}:" + f"{step.micro_batch}:{step.phase.value}" + ) + + @property + def schedule_object(self) -> ScheduleObject: + """Return the one content-addressed authority for this pipeline.""" + + steps = tuple(self.schedule_steps()) + # Key by the VIRTUAL STAGE, not the physical rank. Under interleaving + # one rank owns several virtual stages, so a (rank, micro_batch, phase) + # key collides across chunks: later chunks overwrite earlier ones, and + # a forward step then resolves its producer to a future chunk (dropped + # by the ordering filter below) or to nothing at all. Measured on + # num_stages=4, num_chunks=2: 64 of 128 steps collided and 32 of 56 + # cross-stage forward steps lost their true producer edge, so the + # emitted Schedule Object permitted a virtual stage to run before the + # stage that feeds it (PR #626 review). + by_key = {(step.stage, step.micro_batch, step.phase): step for step in steps} + order_by_id = {self._action_id(step): index for index, step in enumerate(steps)} + # Virtual stage count is DERIVED from the emitted schedule: under + # interleaving it is num_stages x num_chunks, not num_stages. + total_stages = max((step.stage for step in steps), default=-1) + 1 + inverted_backward: list[tuple[str, str]] = [] + previous_by_rank: dict[int, ScheduleStep] = {} + actions: list[ScheduleAction] = [] + for step in steps: + action_id = self._action_id(step) + dependencies: set[str] = set() + previous = previous_by_rank.get(step.rank) + if previous is not None: + dependencies.add(self._action_id(previous)) + if step.phase == Phase.FORWARD and step.stage > 0: + upstream = by_key.get( + (step.stage - 1, step.micro_batch, Phase.FORWARD) + ) + if upstream is not None: + dependencies.add(self._action_id(upstream)) + if step.phase == Phase.BACKWARD: + own_forward = by_key.get( + (step.stage, step.micro_batch, Phase.FORWARD) + ) + if own_forward is not None: + dependencies.add(self._action_id(own_forward)) + if step.stage + 1 < total_stages: + downstream = by_key.get( + (step.stage + 1, step.micro_batch, Phase.BACKWARD) + ) + if downstream is not None: + downstream_id = self._action_id(downstream) + if order_by_id[downstream_id] < order_by_id[action_id]: + dependencies.add(downstream_id) + else: + # KNOWN PLANNER LIMITATION, surfaced by keying on + # the virtual stage (PR #626 review). The + # interleaved generator emits backward steps in + # ASCENDING stage order (num_stages=2, chunks=2, + # micro-batch 0: stage 0 B at clock 4 ... stage 3 B + # at clock 7), which is the opposite of gradient + # flow, so this edge cannot be expressed in the + # emitted order. Recording it would claim an + # ordering the schedule does not realize, so it is + # omitted and counted; fixing the generator's + # interleaved backward order is its own change. + inverted_backward.append((action_id, downstream_id)) + + # A real producer ordered AFTER its consumer is a schedule defect, + # not something to drop: silently filtering it would emit a + # Schedule Object that permits the consumer to run first. Fail + # closed instead (PR #626 review). + late = sorted( + dependency + for dependency in dependencies + if order_by_id[dependency] >= order_by_id[action_id] + ) + # Forward/own-forward edges MUST precede their consumer in any + # correct schedule; an inversion there is a defect, not a + # limitation, so it fails closed rather than being dropped. + if late: + raise ValueError( + f"pipeline schedule places {action_id!r} before its " + f"producers {late!r}; the emitted Schedule Object would " + f"permit a stage to execute ahead of its input" + ) + identity = json.dumps( + { + "action_id": action_id, + "dependencies": sorted(dependencies), + "mode": ( + "decoupled" + if self.decoupled + else "interleaved" if self.interleaved else "1f1b" + ), + }, + sort_keys=True, + separators=(",", ":"), + ) + vector = MeasuredResourceVector( + compute_time_ms=1.0, + bytes_moved=0, + communication_bytes=0, + queue_identity=f"pipeline_rank:{step.rank}", + resource_identity=f"pipeline_stage:{step.stage}", + timing_provenance={ + "source": "static_pipeline_model", + "domain": "compiler", + }, + artifact_digest=hashlib.sha256(identity.encode()).hexdigest(), + ).as_dict() + actions.append( + ScheduleAction( + action_id, + vector, + tuple(sorted(dependencies)), + op_ref=f"schedule.pipeline.{step.phase.value.lower()}", + scope=f"stage:{step.stage}", + ) + ) + previous_by_rank[step.rank] = step + + mode = ( + "decoupled" + if self.decoupled + else "interleaved" if self.interleaved else "1f1b" + ) + roles = tuple( + ScheduleRole(f"stage_{stage}", (f"stage{stage}",)) + for stage in range(self.num_stages) + ) + return ScheduleObject( + object_id=( + f"pipeline:{mode}:stages={self.num_stages}:" + f"micro_batches={self.num_micro_batches}:chunks={self.num_chunks}" + ), + actions=tuple(actions), + roles=roles, ) def __repr__(self) -> str: diff --git a/python/tessera/compiler/schedule_ir.py b/python/tessera/compiler/schedule_ir.py index edda83551..959c056e3 100644 --- a/python/tessera/compiler/schedule_ir.py +++ b/python/tessera/compiler/schedule_ir.py @@ -23,8 +23,15 @@ RasterOrder, ) - -SCHEDULE_MEMORY_SPACES = {"register", "shared", "lds", "global", "managed", "host", "tmem"} +SCHEDULE_MEMORY_SPACES = { + "register", + "shared", + "lds", + "global", + "managed", + "host", + "tmem", +} SCHEDULE_OVERLAPS = {"none", "compute", "collective"} MATMUL_OPS = {"tessera.matmul", "tessera.gemm"} CONV2D_OPS = {"tessera.conv2d_nhwc", "tessera.conv2d"} @@ -34,6 +41,7 @@ "tessera.reduce_scatter", "tessera.all_gather", "tessera.all_to_all", + "tessera.collective_permute", } MEDIA_OPS = { "tessera.image_preprocess", @@ -100,12 +108,12 @@ def to_mlir(self, indent: str = " ") -> str: operand_text = ", ".join(self.operands) result_text = f"%{self.result} = " if self.result else "" if self.body: - lines = [f"{indent}{result_text}\"{self.op_name}\"({operand_text}) ({{"] + lines = [f'{indent}{result_text}"{self.op_name}"({operand_text}) ({{'] for child in self.body: lines.append(child.to_mlir(indent + " ")) lines.append(f"{indent}}}) {attr_text} : () -> ()") return "\n".join(lines) - return f"{indent}{result_text}\"{self.op_name}\"({operand_text}) {attr_text} : () -> ()" + return f'{indent}{result_text}"{self.op_name}"({operand_text}) {attr_text} : () -> ()' @dataclass @@ -115,17 +123,21 @@ class ScheduleFunction: target: str = "cpu" def to_mlir(self, indent: str = " ") -> str: - lines = [f"{indent}\"tessera.schedule.func\"() ({{"] + lines = [f'{indent}"tessera.schedule.func"() ({{'] for op in self.body: lines.append(op.to_mlir(indent + " ")) - lines.append(f"{indent}}}) {{sym_name = {json.dumps(self.name)}, target = {json.dumps(self.target)}}} : () -> ()") + lines.append( + f"{indent}}}) {{sym_name = {json.dumps(self.name)}, target = {json.dumps(self.target)}}} : () -> ()" + ) return "\n".join(lines) @dataclass class ScheduleIRModule: functions: list[ScheduleFunction] = field(default_factory=list) - attrs: dict[str, Any] = field(default_factory=lambda: {"tessera.ir.level": "schedule"}) + attrs: dict[str, Any] = field( + default_factory=lambda: {"tessera.ir.level": "schedule"} + ) def verify(self) -> ScheduleIRVerificationResult: return ScheduleIRVerifier().verify_module(self) @@ -179,11 +191,13 @@ def _verify_ops( self._verify_state_read(op, diagnostics) elif op.op_name == "schedule.yield": if not region_stack: - diagnostics.append(GraphIRDiagnostic( - "error", - "schedule.yield must terminate a schedule region", - code="SCHEDULE_IR_YIELD_OUTSIDE_REGION", - )) + diagnostics.append( + GraphIRDiagnostic( + "error", + "schedule.yield must terminate a schedule region", + code="SCHEDULE_IR_YIELD_OUTSIDE_REGION", + ) + ) def _verify_mesh_define( self, @@ -195,19 +209,53 @@ def _verify_mesh_define( dims = op.attrs.get("dims") axes = op.attrs.get("axis_names") if not name: - diagnostics.append(GraphIRDiagnostic("error", "mesh.define requires sym_name", code="SCHEDULE_IR_MESH_NAME")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "mesh.define requires sym_name", + code="SCHEDULE_IR_MESH_NAME", + ) + ) elif name in meshes: - diagnostics.append(GraphIRDiagnostic("error", f"duplicate mesh {name!r}", code="SCHEDULE_IR_DUP_MESH")) + diagnostics.append( + GraphIRDiagnostic( + "error", f"duplicate mesh {name!r}", code="SCHEDULE_IR_DUP_MESH" + ) + ) else: meshes[str(name)] = op if not isinstance(dims, list) or not dims: - diagnostics.append(GraphIRDiagnostic("error", "mesh.define requires non-empty dims", code="SCHEDULE_IR_MESH_DIMS")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "mesh.define requires non-empty dims", + code="SCHEDULE_IR_MESH_DIMS", + ) + ) elif any(int(dim) <= 0 for dim in dims): - diagnostics.append(GraphIRDiagnostic("error", "mesh dimensions must be positive", code="SCHEDULE_IR_MESH_DIMS")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "mesh dimensions must be positive", + code="SCHEDULE_IR_MESH_DIMS", + ) + ) if not isinstance(axes, list) or not axes: - diagnostics.append(GraphIRDiagnostic("error", "mesh.define requires axis_names", code="SCHEDULE_IR_MESH_AXES")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "mesh.define requires axis_names", + code="SCHEDULE_IR_MESH_AXES", + ) + ) elif isinstance(dims, list) and len(dims) != len(axes): - diagnostics.append(GraphIRDiagnostic("error", "mesh dims and axis_names length mismatch", code="SCHEDULE_IR_MESH_RANK")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "mesh dims and axis_names length mismatch", + code="SCHEDULE_IR_MESH_RANK", + ) + ) def _verify_mesh_region( self, @@ -219,14 +267,37 @@ def _verify_mesh_region( mesh = str(op.attrs.get("mesh", "")) axis = op.attrs.get("axis") if mesh not in meshes: - diagnostics.append(GraphIRDiagnostic("error", f"mesh.region references undefined mesh {mesh!r}", code="SCHEDULE_IR_UNDEFINED_MESH")) + diagnostics.append( + GraphIRDiagnostic( + "error", + f"mesh.region references undefined mesh {mesh!r}", + code="SCHEDULE_IR_UNDEFINED_MESH", + ) + ) else: axes = meshes[mesh].attrs.get("axis_names", []) if axis not in axes: - diagnostics.append(GraphIRDiagnostic("error", f"axis {axis!r} is not defined on mesh {mesh!r}", code="SCHEDULE_IR_UNKNOWN_AXIS")) + diagnostics.append( + GraphIRDiagnostic( + "error", + f"axis {axis!r} is not defined on mesh {mesh!r}", + code="SCHEDULE_IR_UNKNOWN_AXIS", + ) + ) if not op.body or op.body[-1].op_name != "schedule.yield": - diagnostics.append(GraphIRDiagnostic("error", "mesh.region must end with schedule.yield", code="SCHEDULE_IR_REGION_TERMINATOR")) - self._verify_ops(op.body, diagnostics, meshes=meshes, region_stack=region_stack + ["mesh.region"]) + diagnostics.append( + GraphIRDiagnostic( + "error", + "mesh.region must end with schedule.yield", + code="SCHEDULE_IR_REGION_TERMINATOR", + ) + ) + self._verify_ops( + op.body, + diagnostics, + meshes=meshes, + region_stack=region_stack + ["mesh.region"], + ) def _verify_pipeline_region( self, @@ -236,14 +307,43 @@ def _verify_pipeline_region( region_stack: list[str], ) -> None: if not op.attrs.get("schedule"): - diagnostics.append(GraphIRDiagnostic("error", "pipeline.region requires schedule", code="SCHEDULE_IR_PIPELINE_SCHEDULE")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "pipeline.region requires schedule", + code="SCHEDULE_IR_PIPELINE_SCHEDULE", + ) + ) if int(op.attrs.get("micro_batches", 0)) < 1: - diagnostics.append(GraphIRDiagnostic("error", "pipeline.region micro_batches must be >= 1", code="SCHEDULE_IR_PIPELINE_MICRO_BATCHES")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "pipeline.region micro_batches must be >= 1", + code="SCHEDULE_IR_PIPELINE_MICRO_BATCHES", + ) + ) if not any(child.op_name == "schedule.stage" for child in op.body): - diagnostics.append(GraphIRDiagnostic("error", "pipeline.region requires at least one stage", code="SCHEDULE_IR_PIPELINE_STAGE")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "pipeline.region requires at least one stage", + code="SCHEDULE_IR_PIPELINE_STAGE", + ) + ) if not op.body or op.body[-1].op_name != "schedule.yield": - diagnostics.append(GraphIRDiagnostic("error", "pipeline.region must end with schedule.yield", code="SCHEDULE_IR_REGION_TERMINATOR")) - self._verify_ops(op.body, diagnostics, meshes=meshes, region_stack=region_stack + ["pipeline.region"]) + diagnostics.append( + GraphIRDiagnostic( + "error", + "pipeline.region must end with schedule.yield", + code="SCHEDULE_IR_REGION_TERMINATOR", + ) + ) + self._verify_ops( + op.body, + diagnostics, + meshes=meshes, + region_stack=region_stack + ["pipeline.region"], + ) def _verify_stage( self, @@ -253,31 +353,78 @@ def _verify_stage( region_stack: list[str], ) -> None: if "pipeline.region" not in region_stack: - diagnostics.append(GraphIRDiagnostic("error", "stage must be inside pipeline.region", code="SCHEDULE_IR_STAGE_PARENT")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "stage must be inside pipeline.region", + code="SCHEDULE_IR_STAGE_PARENT", + ) + ) devices = op.attrs.get("devices") if not isinstance(devices, list) or not devices: - diagnostics.append(GraphIRDiagnostic("error", "stage requires non-empty devices", code="SCHEDULE_IR_STAGE_DEVICES")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "stage requires non-empty devices", + code="SCHEDULE_IR_STAGE_DEVICES", + ) + ) if not op.body or op.body[-1].op_name != "schedule.yield": - diagnostics.append(GraphIRDiagnostic("error", "stage must end with schedule.yield", code="SCHEDULE_IR_REGION_TERMINATOR")) - self._verify_ops(op.body, diagnostics, meshes=meshes, region_stack=region_stack + ["stage"]) + diagnostics.append( + GraphIRDiagnostic( + "error", + "stage must end with schedule.yield", + code="SCHEDULE_IR_REGION_TERMINATOR", + ) + ) + self._verify_ops( + op.body, diagnostics, meshes=meshes, region_stack=region_stack + ["stage"] + ) - def _verify_prefetch(self, op: ScheduleOp, diagnostics: list[GraphIRDiagnostic]) -> None: + def _verify_prefetch( + self, op: ScheduleOp, diagnostics: list[GraphIRDiagnostic] + ) -> None: if op.attrs.get("into") not in SCHEDULE_MEMORY_SPACES: - diagnostics.append(GraphIRDiagnostic("error", "prefetch has invalid memory space", code="SCHEDULE_IR_MEMORY_SPACE")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "prefetch has invalid memory space", + code="SCHEDULE_IR_MEMORY_SPACE", + ) + ) if op.attrs.get("overlap") not in SCHEDULE_OVERLAPS: - diagnostics.append(GraphIRDiagnostic("error", "prefetch has invalid overlap policy", code="SCHEDULE_IR_OVERLAP")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "prefetch has invalid overlap policy", + code="SCHEDULE_IR_OVERLAP", + ) + ) - def _verify_state_read(self, op: ScheduleOp, diagnostics: list[GraphIRDiagnostic]) -> None: - missing = [key for key in ("source", "result", "ordinal", "effect", "access") - if key not in op.attrs] + def _verify_state_read( + self, op: ScheduleOp, diagnostics: list[GraphIRDiagnostic] + ) -> None: + missing = [ + key + for key in ("source", "result", "ordinal", "effect", "access") + if key not in op.attrs + ] if missing: - diagnostics.append(GraphIRDiagnostic( - "error", f"state.read missing attrs: {', '.join(missing)}", - code="SCHEDULE_IR_STATE_READ_ATTR")) + diagnostics.append( + GraphIRDiagnostic( + "error", + f"state.read missing attrs: {', '.join(missing)}", + code="SCHEDULE_IR_STATE_READ_ATTR", + ) + ) if op.attrs.get("effect") != "read": - diagnostics.append(GraphIRDiagnostic( - "error", "state.read must declare effect=read", - code="SCHEDULE_IR_STATE_READ_EFFECT")) + diagnostics.append( + GraphIRDiagnostic( + "error", + "state.read must declare effect=read", + code="SCHEDULE_IR_STATE_READ_EFFECT", + ) + ) def lower_graph_to_schedule_ir( @@ -290,7 +437,9 @@ def lower_graph_to_schedule_ir( graph_result = graph_module.verify(target=target_kind) if not graph_result.ok: raise GraphIRVerificationError(graph_result.format()) - schedule_module = ScheduleIRModule(attrs={"tessera.ir.level": "schedule", "target": target_kind}) + schedule_module = ScheduleIRModule( + attrs={"tessera.ir.level": "schedule", "target": target_kind} + ) for graph_fn in graph_module.functions: body: list[ScheduleOp] = [] mesh_ops = [ @@ -319,24 +468,31 @@ def lower_graph_to_schedule_ir( ) ] body.extend(scheduled_ops) - body.append(ScheduleOp("schedule.artifact", { - "hash": f"{graph_fn.name}:{len(graph_fn.body)}:{target_kind}", - "arch": target_kind, - "shape_key": _shape_key(graph_fn.body), - "tile": {"m": tile[0], "n": tile[1], "k": tile[2]}, - "movement": { - "prefetch": "auto", - "overlap": "compute", - "stages": _schedule_int(schedule_config, "num_stages", 2), - }, - "numeric_policy": "f32@accum(f32)", - "cost_model": ( - "measured" - if (schedule_config or {}).get("evidence") == "measured" - else "roofline" - ), - })) - schedule_module.functions.append(ScheduleFunction(graph_fn.name, body=body, target=target_kind)) + body.append( + ScheduleOp( + "schedule.artifact", + { + "hash": f"{graph_fn.name}:{len(graph_fn.body)}:{target_kind}", + "arch": target_kind, + "shape_key": _shape_key(graph_fn.body), + "tile": {"m": tile[0], "n": tile[1], "k": tile[2]}, + "movement": { + "prefetch": "auto", + "overlap": "compute", + "stages": _schedule_int(schedule_config, "num_stages", 2), + }, + "numeric_policy": "f32@accum(f32)", + "cost_model": ( + "measured" + if (schedule_config or {}).get("evidence") == "measured" + else "roofline" + ), + }, + ) + ) + schedule_module.functions.append( + ScheduleFunction(graph_fn.name, body=body, target=target_kind) + ) return schedule_module @@ -351,91 +507,164 @@ def _lower_graph_ops( for idx, op in enumerate(ops): op_name = canonical_graph_op_name(op.op_name) if op_name == "tessera.graph.debug_value": - scheduled.append(ScheduleOp("schedule.debug_artifact", { - **_base_attrs(op, idx), - "name": op.kwargs.get("name", op.result or f"debug_{idx}"), - "capture": "value_summary", - })) + scheduled.append( + ScheduleOp( + "schedule.debug_artifact", + { + **_base_attrs(op, idx), + "name": op.kwargs.get("name", op.result or f"debug_{idx}"), + "capture": "value_summary", + }, + ) + ) continue if op_name.startswith("tessera.schedule."): scheduled.extend(_lower_schedule_directive(op, idx)) continue if op_name.startswith("tessera.dist.") or op_name in COLLECTIVE_OPS: axis = op.kwargs.get("axis", "dp") - tensor_axis = axis if isinstance(axis, int) and not isinstance(axis, bool) else 0 - mesh_axis = "default" if isinstance(axis, int) and not isinstance(axis, bool) else str(axis) + tensor_axis = ( + axis if isinstance(axis, int) and not isinstance(axis, bool) else 0 + ) + mesh_axis = ( + "default" + if isinstance(axis, int) and not isinstance(axis, bool) + else str(axis) + ) kind = op_name.removeprefix("tessera.dist.").removeprefix("tessera.") - scheduled.append(ScheduleOp( - "schedule.collective", - { - **_base_attrs(op, idx), - "kind": kind, - "mesh_axis": mesh_axis, - "tensor_axis": tensor_axis, - "reduction": op.kwargs.get("op", "sum"), - "effect": "collective", - "reshard_plan_digest": op.kwargs.get("reshard_plan_digest", ""), - "subgroup": list(op.kwargs.get("subgroup", ())), - "region_path": list(op.kwargs.get("region_path", ())), - "scatter_axis": op.kwargs.get("scatter_axis", tensor_axis), - "gather_axis": op.kwargs.get("gather_axis", tensor_axis), - "matching_rounds": list(op.kwargs.get("matching_rounds", ())), - }, - operands=list(op.operands), - result=op.result, - source_op=op, - )) + scheduled.append( + ScheduleOp( + "schedule.collective", + { + **_base_attrs(op, idx), + "kind": kind, + "mesh_axis": mesh_axis, + "tensor_axis": tensor_axis, + "reduction": op.kwargs.get("op", "sum"), + "effect": "collective", + "reshard_plan_digest": op.kwargs.get("reshard_plan_digest", ""), + "subgroup": list(op.kwargs.get("subgroup", ())), + "region_path": list(op.kwargs.get("region_path", ())), + "scatter_axis": op.kwargs.get("scatter_axis", tensor_axis), + "gather_axis": op.kwargs.get("gather_axis", tensor_axis), + "matching_rounds": list(op.kwargs.get("matching_rounds", ())), + "source_peers": list(op.kwargs.get("source_peers", ())), + "target_peers": list(op.kwargs.get("target_peers", ())), + }, + operands=list(op.operands), + result=op.result, + source_op=op, + ) + ) continue if op_name in MATMUL_OPS: - scheduled.extend([ - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "tile_m", "choices": [32, 64, 128, 256], "frozen": False}), - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "tile_n", "choices": [32, 64, 128, 256], "frozen": False}), - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "tile_k", "choices": [32, 64, 128, 256], "frozen": False}), - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "num_warps", "choices": [1, 2, 4, 8], "frozen": False}), - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "num_stages", "choices": [1, 2, 3, 4], "frozen": False}), - # Block rasterization: a permutation of block ids onto the tile - # grid. Semantics-preserving by construction (tile_rasterization - # .is_bijection proves it hardware-free), and one of the largest - # levers on cache locality there is — see - # docs/audit/compiler/TILESIGHT_ASSESSMENT.md §3.2. - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "raster_order", "choices": list(RASTER_ORDER_CHOICES), "frozen": False}), - ScheduleOp("schedule.knob", {**_base_attrs(op, idx), "name": "raster_group", "choices": list(RASTER_GROUP_CHOICES), "frozen": False}), - ]) - scheduled.append(ScheduleOp( - "schedule.tile", - { - **_base_attrs(op, idx), - "tile_m": tile_m, - "tile_n": tile_n, - "tile_k": tile_k, - "num_warps": _schedule_int( - schedule_config, "num_warps", 4 + scheduled.extend( + [ + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "tile_m", + "choices": [32, 64, 128, 256], + "frozen": False, + }, ), - "num_stages": _schedule_int( - schedule_config, "num_stages", 2 + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "tile_n", + "choices": [32, 64, 128, 256], + "frozen": False, + }, ), - "raster_order": _schedule_raster_order(schedule_config), - "raster_group": _schedule_raster_group(schedule_config), - "cost_model": ( - "measured" - if (schedule_config or {}).get("evidence") == "measured" - else "roofline" + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "tile_k", + "choices": [32, 64, 128, 256], + "frozen": False, + }, ), - "flops": _matmul_flops(op), - "bytes_moved": _matmul_bytes(op), - }, - operands=list(op.operands), - result=op.result, - source_op=op, - )) + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "num_warps", + "choices": [1, 2, 4, 8], + "frozen": False, + }, + ), + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "num_stages", + "choices": [1, 2, 3, 4], + "frozen": False, + }, + ), + # Block rasterization: a permutation of block ids onto the tile + # grid. Semantics-preserving by construction (tile_rasterization + # .is_bijection proves it hardware-free), and one of the largest + # levers on cache locality there is — see + # docs/audit/compiler/TILESIGHT_ASSESSMENT.md §3.2. + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "raster_order", + "choices": list(RASTER_ORDER_CHOICES), + "frozen": False, + }, + ), + ScheduleOp( + "schedule.knob", + { + **_base_attrs(op, idx), + "name": "raster_group", + "choices": list(RASTER_GROUP_CHOICES), + "frozen": False, + }, + ), + ] + ) + scheduled.append( + ScheduleOp( + "schedule.tile", + { + **_base_attrs(op, idx), + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "num_warps": _schedule_int(schedule_config, "num_warps", 4), + "num_stages": _schedule_int(schedule_config, "num_stages", 2), + "raster_order": _schedule_raster_order(schedule_config), + "raster_group": _schedule_raster_group(schedule_config), + "cost_model": ( + "measured" + if (schedule_config or {}).get("evidence") == "measured" + else "roofline" + ), + "flops": _matmul_flops(op), + "bytes_moved": _matmul_bytes(op), + }, + operands=list(op.operands), + result=op.result, + source_op=op, + ) + ) elif op_name in CONV2D_OPS: - scheduled.append(ScheduleOp( - "schedule.tile", - {**_base_attrs(op, idx), "tile_h": 16, "tile_w": 16, "tile_c": 32}, - operands=list(op.operands), - result=op.result, - source_op=op, - )) + scheduled.append( + ScheduleOp( + "schedule.tile", + {**_base_attrs(op, idx), "tile_h": 16, "tile_w": 16, "tile_c": 32}, + operands=list(op.operands), + result=op.result, + source_op=op, + ) + ) elif op_name == "tessera.flash_attn": scheduled.append(_flash_attention_pipeline(op, idx)) elif op_name == "tessera.msa_sparse_attention": @@ -451,41 +680,81 @@ def _lower_graph_ops( elif op_name in JEPA_OPS: scheduled.append(_jepa_op(op, idx)) elif op_name == "tessera.kv_cache.read": - scheduled.append(ScheduleOp( - "schedule.state.read", - { - **_base_attrs(op, idx), - "effect": "read", - "access": "paged_slice", - "bounds": "start_end", - }, - operands=list(op.operands), - result=op.result, - source_op=op, - )) + scheduled.append( + ScheduleOp( + "schedule.state.read", + { + **_base_attrs(op, idx), + "effect": "read", + "access": "paged_slice", + "bounds": "start_end", + }, + operands=list(op.operands), + result=op.result, + source_op=op, + ) + ) elif op_name in ROPE_OPS: - scheduled.append(ScheduleOp( - "schedule.elementwise", - {**_base_attrs(op, idx), "vectorize": True, "pattern": "rotary_pairs"}, - operands=list(op.operands), result=op.result, source_op=op)) - elif op_name.startswith("tessera.scf.") or op_name in {"tessera.barrier", "tessera.assert"}: - scheduled.append(ScheduleOp("schedule.marker", {**_base_attrs(op, idx), **op.kwargs, "marker": op_name})) + scheduled.append( + ScheduleOp( + "schedule.elementwise", + { + **_base_attrs(op, idx), + "vectorize": True, + "pattern": "rotary_pairs", + }, + operands=list(op.operands), + result=op.result, + source_op=op, + ) + ) + elif op_name.startswith("tessera.scf.") or op_name in { + "tessera.barrier", + "tessera.assert", + }: + scheduled.append( + ScheduleOp( + "schedule.marker", + {**_base_attrs(op, idx), **op.kwargs, "marker": op_name}, + ) + ) elif op.result is not None: - scheduled.append(ScheduleOp( - "schedule.elementwise", - {**_base_attrs(op, idx), "vectorize": True}, - operands=list(op.operands), result=op.result, source_op=op)) + scheduled.append( + ScheduleOp( + "schedule.elementwise", + { + **_base_attrs(op, idx), + **( + dict(op.kwargs) + if "reshard_plan_digest" in op.kwargs + else {} + ), + "vectorize": True, + }, + operands=list(op.operands), + result=op.result, + source_op=op, + ) + ) operand_names = [operand.removeprefix("%") for operand in op.operands] if operand_names: - scheduled.append(ScheduleOp("schedule.layout", {"operands": operand_names, "layout": "row_major", "ordinal": idx})) + scheduled.append( + ScheduleOp( + "schedule.layout", + {"operands": operand_names, "layout": "row_major", "ordinal": idx}, + ) + ) if op_name.startswith("tessera.kv_cache."): - scheduled.append(ScheduleOp("schedule.prefetch", {**_base_attrs(op, idx), "into": "shared", "overlap": "compute"})) + scheduled.append( + ScheduleOp( + "schedule.prefetch", + {**_base_attrs(op, idx), "into": "shared", "overlap": "compute"}, + ) + ) return scheduled -def _schedule_int( - config: dict[str, object] | None, name: str, default: int -) -> int: +def _schedule_int(config: dict[str, object] | None, name: str, default: int) -> int: if config is None or name not in config: return default value = config[name] @@ -507,7 +776,8 @@ def _schedule_raster_order(config: dict[str, object] | None) -> str: if not isinstance(value, str) or value not in RASTER_ORDER_CHOICES: raise ValueError( f"schedule raster_order must be one of {list(RASTER_ORDER_CHOICES)}, " - f"got {value!r}") + f"got {value!r}" + ) return value @@ -535,20 +805,36 @@ def _lower_schedule_directive(op: IROp, ordinal: int) -> list[ScheduleOp]: depth = int(attrs.get("depth", attrs.get("stages", 1))) stage_count = max(1, depth) stages = [ - ScheduleOp("schedule.stage", {"devices": [stage]}, body=[ScheduleOp("schedule.yield")]) + ScheduleOp( + "schedule.stage", + {"devices": [stage]}, + body=[ScheduleOp("schedule.yield")], + ) for stage in range(stage_count) ] return [ ScheduleOp( "schedule.pipeline.region", - {"schedule": attrs.get("schedule", "gpipe"), "micro_batches": int(attrs.get("micro_batches", 1))}, + { + "schedule": attrs.get("schedule", "gpipe"), + "micro_batches": int(attrs.get("micro_batches", 1)), + }, body=[*stages, ScheduleOp("schedule.yield")], ) ] if name == "schedule.tile": return [ScheduleOp("schedule.tile", attrs)] if name == "schedule.prefetch": - return [ScheduleOp("schedule.prefetch", {**attrs, "into": attrs.get("into", attrs.get("scope", "shared")), "overlap": attrs.get("overlap", "compute")})] + return [ + ScheduleOp( + "schedule.prefetch", + { + **attrs, + "into": attrs.get("into", attrs.get("scope", "shared")), + "overlap": attrs.get("overlap", "compute"), + }, + ) + ] return [ScheduleOp(name, attrs)] @@ -561,7 +847,16 @@ def _flash_attention_pipeline(op: IROp, ordinal: int) -> ScheduleOp: "schedule.stage", {"devices": [0]}, body=[ - ScheduleOp("schedule.prefetch", {**_base_attrs(op, ordinal), "into": "shared", "overlap": "compute", "tile_q": 64, "tile_kv": 64}), + ScheduleOp( + "schedule.prefetch", + { + **_base_attrs(op, ordinal), + "into": "shared", + "overlap": "compute", + "tile_q": 64, + "tile_kv": 64, + }, + ), ScheduleOp("schedule.yield"), ], ), @@ -625,8 +920,12 @@ def _msa_kv_outer_sparse(op: IROp, ordinal: int) -> ScheduleOp: block_size = int(attrs.get("block_size", 64)) num_heads = int(attrs.get("num_heads", attrs.get("num_attention_heads", 1))) num_kv_heads = int(attrs.get("num_kv_heads", attrs.get("num_key_value_heads", 1))) - gqa_group_size = int(attrs.get("gqa_group_size", max(1, num_heads // max(1, num_kv_heads)))) - mode = str(attrs.get("mode", "decode" if int(attrs.get("tile_q", 64)) == 1 else "prefill")) + gqa_group_size = int( + attrs.get("gqa_group_size", max(1, num_heads // max(1, num_kv_heads))) + ) + mode = str( + attrs.get("mode", "decode" if int(attrs.get("tile_q", 64)) == 1 else "prefill") + ) tile_q = int(attrs.get("tile_q", 1 if mode == "decode" else 64)) tile_kv = int(attrs.get("tile_kv", max(block_size, 128))) head_dim = int(attrs.get("head_dim", 128)) @@ -644,7 +943,9 @@ def _msa_kv_outer_sparse(op: IROp, ordinal: int) -> ScheduleOp: "head_dim": head_dim, "mode": mode, "acc_dtype": attrs.get("acc_dtype", "fp32"), - "dense_equivalence_oracle": bool(attrs.get("dense_equivalence_oracle", False)), + "dense_equivalence_oracle": bool( + attrs.get("dense_equivalence_oracle", False) + ), "kv_traversal": "kv_outer", "online_softmax": True, }, @@ -773,7 +1074,13 @@ def _matmul_bytes(op: IROp) -> int: def _format_attr_dict(attrs: dict[str, Any]) -> str: if not attrs: return "{}" - return "{" + ", ".join(f"{key} = {_format_attr_value(value)}" for key, value in attrs.items()) + "}" + return ( + "{" + + ", ".join( + f"{key} = {_format_attr_value(value)}" for key, value in attrs.items() + ) + + "}" + ) def _format_attr_value(value: Any) -> str: diff --git a/python/tessera/compiler/scheduled_spectral.py b/python/tessera/compiler/scheduled_spectral.py index f1811733f..4d0bd30e3 100644 --- a/python/tessera/compiler/scheduled_spectral.py +++ b/python/tessera/compiler/scheduled_spectral.py @@ -2,17 +2,22 @@ from __future__ import annotations +import hashlib +import json import math import re from dataclasses import dataclass from functools import lru_cache from typing import Any, Mapping, Sequence +from .benchmark_row import MeasuredResourceVector +from .composition_cost import InferredActionDAG, infer_action_dag +from .graph_ir import GraphIRFunction, IRArg, IROp, IRType +from .schedule_object import ScheduleObject, ScheduleRole from .scheduled_fft import lower_scheduled_fft, validate_scheduled_fft_metadata from .scheduled_matmul import digest_text, find_tessera_opt, run_tessera_opt from .spectral_plan import next_power_of_two - _OPS = ( "tessera.spectral_filter", "tessera.dct", @@ -27,6 +32,122 @@ _STORAGE_POLICIES = ("f32", "f16", "bf16") +def _spectral_graph_function( + *, object_id: str, op_name: str, dct_type: int +) -> tuple[GraphIRFunction, tuple[str, ...]]: + """Represent the fused physical producer as registered Graph actions.""" + + specs: tuple[tuple[str, str, tuple[str, ...]], ...] + if op_name == "tessera.spectral_filter": + specs = ( + ("complex_multiply", "tessera.mul", ("%input0", "%input1")), + ("materialize_output", "tessera.reshape", ("%complex_multiply",)), + ) + elif op_name == "tessera.dct" and dct_type == 2: + specs = ( + ("even_extend", "tessera.pad", ("%input0",)), + ("fft", "tessera.fft", ("%even_extend",)), + ("phase_correct", "tessera.mul", ("%fft", "%phase")), + ("crop", "tessera.slice", ("%phase_correct",)), + ) + elif op_name == "tessera.dct": + specs = ( + ("direct_cosine", "tessera.dct", ("%input0",)), + ("materialize_output", "tessera.reshape", ("%direct_cosine",)), + ) + elif op_name == "tessera.spectral_conv": + specs = ( + ("pad_signal", "tessera.pad", ("%input0",)), + ("pad_kernel", "tessera.pad", ("%input1",)), + ("rfft_signal", "tessera.rfft", ("%pad_signal",)), + ("rfft_kernel", "tessera.rfft", ("%pad_kernel",)), + ("complex_multiply", "tessera.mul", ("%rfft_signal", "%rfft_kernel")), + ("inverse", "tessera.irfft", ("%complex_multiply",)), + ("crop", "tessera.slice", ("%inverse",)), + ) + elif op_name == "tessera.stft": + specs = ( + ("frame", "tessera.reshape", ("%input0",)), + ("apply_window", "tessera.mul", ("%frame", "%input1")), + ("transform", "tessera.rfft", ("%apply_window",)), + ) + else: + specs = ( + ("inverse", "tessera.irfft", ("%input0",)), + ("apply_window", "tessera.mul", ("%inverse", "%input1")), + ("overlap_add", "tessera.reduce_sum", ("%apply_window",)), + ) + tensor = IRType("tensor<*xf32>") + names = {operand.lstrip("%") for _, _, operands in specs for operand in operands} + produced = {action_id for action_id, _, _ in specs} + args = [IRArg(name, tensor) for name in sorted(names - produced)] + ops = [ + IROp( + result=f"%{action_id}", + op_name=graph_op, + operands=list(operands), + operand_types=["tensor<*xf32>"] * len(operands), + result_type="tensor<*xf32>", + ) + for action_id, graph_op, operands in specs + ] + return ( + GraphIRFunction( + name=object_id, + args=args, + body=ops, + return_values=[f"%{specs[-1][0]}"], + ), + tuple(action_id for action_id, _, _ in specs), + ) + + +def infer_spectral_action_dag( + *, + semantic_digest: str, + target: str, + architecture: str, + op_name: str, + dct_type: int, + workspace_bytes: int, +) -> tuple[InferredActionDAG, ScheduleObject]: + """Infer the representative spectral producer DAG and bind SO identity.""" + + object_id = f"spectral:{semantic_digest}" + function, action_ids = _spectral_graph_function( + object_id=object_id, op_name=op_name, dct_type=dct_type + ) + bytes_per_action = workspace_bytes // max(1, len(action_ids)) + vectors = tuple( + MeasuredResourceVector( + compute_time_ms=1.0, + bytes_moved=bytes_per_action, + communication_bytes=0, + queue_identity=f"{target}:spectral:0", + resource_identity=architecture, + timing_provenance={ + "source": "static_spectral_model", + "domain": "compiler", + }, + artifact_digest=digest_text( + f"{semantic_digest}:{action_id}:{target}:{architecture}" + ), + ).as_dict() + for action_id in action_ids + ) + inferred = infer_action_dag(function, vectors, action_ids=action_ids) + schedule = ScheduleObject( + object_id=object_id, + actions=inferred.actions, + edges=inferred.schedule_object.edges, + roles=( + ScheduleRole("spectral_compute", (architecture,)), + ScheduleRole("spectral_queue", (f"{target}:spectral:0",)), + ), + ) + return inferred, schedule + + @dataclass(frozen=True) class SpectralArchitectureProfile: target: str @@ -40,29 +161,49 @@ class SpectralArchitectureProfile: _ARCHITECTURE_PROFILES = { "x86": SpectralArchitectureProfile( - "x86", "x86", "zen5-avx512", "ready", - "tessera.x86.spectral_composite.v6", "exact_device_validated", + "x86", + "x86", + "zen5-avx512", + "ready", + "tessera.x86.spectral_composite.v6", + "exact_device_validated", "exact Zen 5 package", ), "rocm": SpectralArchitectureProfile( - "rocm", "rocm", "gfx1151", "ready", - "tessera.rocm.spectral_composite.v6", "exact_device_validated", + "rocm", + "rocm", + "gfx1151", + "ready", + "tessera.rocm.spectral_composite.v6", + "exact_device_validated", "exact gfx1151 package", ), "rocm_gfx1151": SpectralArchitectureProfile( - "rocm_gfx1151", "rocm", "gfx1151", "ready", - "tessera.rocm.spectral_composite.v6", "exact_device_validated", + "rocm_gfx1151", + "rocm", + "gfx1151", + "ready", + "tessera.rocm.spectral_composite.v6", + "exact_device_validated", "exact gfx1151 package", ), "rocm_gfx1200": SpectralArchitectureProfile( - "rocm_gfx1200", "rocm", "gfx1200", "fail_closed", - "tessera.rocm.spectral_composite.v6", "build_only", + "rocm_gfx1200", + "rocm", + "gfx1200", + "fail_closed", + "tessera.rocm.spectral_composite.v6", + "build_only", "architecture-stamped package exists; RDNA 4 schedule and exact-device " "evidence are required for execution", ), "rocm_gfx1250": SpectralArchitectureProfile( - "rocm_gfx1250", "rocm", "gfx1250", "fail_closed", - "tessera.rocm.spectral_composite.v6", "build_only", + "rocm_gfx1250", + "rocm", + "gfx1250", + "fail_closed", + "tessera.rocm.spectral_composite.v6", + "build_only", "architecture-stamped package exists; gfx1250 schedule and exact-device " "evidence are required for execution", ), @@ -191,8 +332,7 @@ def define_spectral_program_contract( else: bounds = tuple(tuple(int(dim) for dim in shape) for shape in shape_bounds) if len(bounds) != len(signatures) or any( - len(bound) != len(signature) - for bound, signature in zip(bounds, signatures) + len(bound) != len(signature) for bound, signature in zip(bounds, signatures) ): raise ValueError("TSOL shape bounds must match every input signature") if any(dim <= 0 for shape in bounds for dim in shape): @@ -257,6 +397,8 @@ class ScheduledSpectralArtifact: native_entry: str child_ffts: tuple[Mapping[str, Any], ...] template_digest: str + graph_analysis_digest: str + schedule_object: Mapping[str, Any] schedule_digest: str @property @@ -288,7 +430,9 @@ def axis_packing(self) -> str: ) def _input_shapes_text(self) -> str: - return "|".join("x".join(str(dim) for dim in shape) for shape in self.input_shapes) + return "|".join( + "x".join(str(dim) for dim in shape) for shape in self.input_shapes + ) def _child_digests_text(self) -> str: return ",".join(str(child["schedule_digest"]) for child in self.child_ffts) @@ -366,8 +510,16 @@ def validate(self) -> None: ("x86", "zen5-avx512"), }: raise ValueError("TSOL package requires exact gfx1151 or Zen 5 AVX-512") - if digest_text(self._identity_payload()) != self.schedule_digest: - raise ValueError("TSOL package content identity mismatch") + semantic_digest = digest_text(self._identity_payload()) + if self.schedule_object.get("object_id") != f"spectral:{semantic_digest}": + raise ValueError("TSOL Schedule Object semantic identity mismatch") + encoded_schedule = json.dumps( + self.schedule_object, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + if hashlib.sha256(encoded_schedule).hexdigest() != self.schedule_digest: + raise ValueError("TSOL Schedule Object content identity mismatch") + if len(self.graph_analysis_digest) != 64: + raise ValueError("TSOL inferred Graph analysis identity mismatch") for child in self.child_ffts: validate_scheduled_fft_metadata( child, target=self.target, input_shape=child["input_shape"] @@ -424,7 +576,10 @@ def validate(self) -> None: raise ValueError("TSOL package requires one durable schedule artifact") if self.tile_ir.count("tile.spectral_program_kernel") != 1: raise ValueError("TSOL package requires one launch-level Tile program") - if any(name in self.tile_ir for name in ("schedule.spectral_program", "schedule.artifact")): + if any( + name in self.tile_ir + for name in ("schedule.spectral_program", "schedule.artifact") + ): raise ValueError("TSOL Tile package retained Schedule IR") if _HASH_RE.findall(self.tile_ir) != [self.schedule_digest]: raise ValueError("TSOL Tile package has stale schedule identity") @@ -437,6 +592,8 @@ def to_metadata(self) -> dict[str, Any]: { "child_ffts": [dict(child) for child in self.child_ffts], "schedule_digest": self.schedule_digest, + "graph_analysis_digest": self.graph_analysis_digest, + "schedule_object": dict(self.schedule_object), "schedule_ir": self.schedule_ir, "schedule_ir_digest": self.schedule_ir_digest, "tile_ir": self.tile_ir, @@ -537,7 +694,9 @@ def lower_scheduled_spectral( elements = math.prod(batch_shape or (1,)) * n workspace = ( _packed_workspace_bytes( - (4, elements), (4, elements), (8, 2 * elements), + (4, elements), + (4, elements), + (8, 2 * elements), (8, 2 * elements), ) if semantic_contract.dct_type == 2 @@ -570,9 +729,7 @@ def lower_scheduled_spectral( ), ) output = ( - shapes[0][:normalized_axis] - + (output_n,) - + shapes[0][normalized_axis + 1 :] + shapes[0][:normalized_axis] + (output_n,) + shapes[0][normalized_axis + 1 :] ) padding = ( fft_n - shapes[0][normalized_axis], @@ -655,7 +812,8 @@ def lower_scheduled_spectral( if stride <= 0 or shapes[0][normalized_axis] != win // 2 + 1: raise ValueError("istft spectrum/window/hop contract mismatch") batch_shape = tuple( - dim for index, dim in enumerate(shapes[0]) + dim + for index, dim in enumerate(shapes[0]) if index not in {frame_axis, normalized_axis} ) batch = math.prod(batch_shape or (1,)) @@ -668,11 +826,7 @@ def lower_scheduled_spectral( ), ) samples = (frames - 1) * stride + win - output = ( - shapes[0][:frame_axis] - + (samples,) - + shapes[0][normalized_axis + 1 :] - ) + output = shapes[0][:frame_axis] + (samples,) + shapes[0][normalized_axis + 1 :] workspace = ( _packed_workspace_bytes( (4, win), @@ -725,18 +879,31 @@ def lower_scheduled_spectral( native_entry=entry, child_ffts=children, template_digest=semantic_contract.template_digest, + graph_analysis_digest="", + schedule_object={}, schedule_digest="", ) identity = provisional._identity_payload() - schedule_digest = digest_text(identity) + semantic_digest = digest_text(identity) + inferred, schedule_object = infer_spectral_action_dag( + semantic_digest=semantic_digest, + target=compiler_target, + architecture=architecture, + op_name=op_name, + dct_type=semantic_contract.dct_type, + workspace_bytes=workspace, + ) + schedule_digest = schedule_object.digest input_names = tuple(f"a{index}" for index in range(len(shapes))) real_element = semantic_contract.storage input_elements = ( ("complex", "complex") if op_name == "tessera.spectral_filter" - else ("complex", real_element) - if op_name == "tessera.istft" - else tuple(real_element for _ in shapes) + else ( + ("complex", real_element) + if op_name == "tessera.istft" + else tuple(real_element for _ in shapes) + ) ) output_element = ( "complex" @@ -747,9 +914,9 @@ def lower_scheduled_spectral( "tensor<" + "x".join([*(str(dim) for dim in shape), element]) + ">" for shape, element in zip(shapes, input_elements) ) - output_type = "tensor<" + "x".join( - [*(str(dim) for dim in output), output_element] - ) + ">" + output_type = ( + "tensor<" + "x".join([*(str(dim) for dim in output), output_element]) + ">" + ) operands = ", ".join(f"%{name}" for name in input_names) function_args = ", ".join( f"%{name}: {type_name}" for name, type_name in zip(input_names, input_types) @@ -766,14 +933,14 @@ def lower_scheduled_spectral( f'input_signature = "{provisional._shapes_text(provisional.input_signature)}", ' f'shape_bounds = "{provisional._shapes_text(provisional.shape_bounds)}", ' f'template_digest = "{provisional.template_digest}", ' - f'output_shape = array, axis = {normalized_axis} : i64, ' - f'dct_type = {semantic_contract.dct_type} : i64, ' + f"output_shape = array, axis = {normalized_axis} : i64, " + f"dct_type = {semantic_contract.dct_type} : i64, " f'shape_policy = "{semantic_contract.shape_policy}", storage = "{semantic_contract.storage}", ' f'abi_storage = "{provisional.abi_storage}", ' f'storage_conversion = "{provisional.storage_conversion}", ' f'axis_packing = "{provisional.axis_packing}", ' - f'padding = array, crop = array, ' - f'window_length = {win} : i64, hop = {stride} : i64, frames = {frames} : i64, ' + f"padding = array, crop = array, " + f"window_length = {win} : i64, hop = {stride} : i64, frames = {frames} : i64, " f'normalization = "{semantic_contract.normalization}", complex_layout = "interleaved_f32x2", ' f'accumulation = "{accumulation}", workspace_bytes = {workspace} : i64, ' f'workspace_policy = "{policy}", mutation_lineage = "inputs_immutable_output_fresh_v1", ' @@ -781,9 +948,18 @@ def lower_scheduled_spectral( f'native_entry = "{entry}", child_fft_digests = "{child_digests}", ' f'workgroup_size = {1 if compiler_target == "x86" else 256} : i64' ) + # Carry the exact semantic payload the physical launch consumes, + # so the Schedule->Tile consumer can re-verify the attributes it + # uses instead of trusting a digest string alone (PR #626 review). + semantic_payload = ( + f"workspace_bytes={workspace};" + f"native_entry={entry};" + f"normalization={semantic_contract.normalization}" + ) schedule_ir = ( f'module attributes {{tessera.target = "{compiler_target}", ' - f'tessera.arch = "{architecture}"}} {{\n' + f'tessera.arch = "{architecture}", ' + f'tessera.schedule_digest = "{schedule_digest}", tessera.spectral_semantic = "{semantic_payload}"}} {{\n' f" func.func @scheduled_spectral({function_args}) -> {output_type} {{\n" f' %result = "schedule.spectral_program"({operands}) {{{attrs}}} : ' f"({operand_types}) -> {output_type}\n" @@ -803,6 +979,8 @@ def lower_scheduled_spectral( **provisional.__dict__, "schedule_ir": schedule_ir, "tile_ir": tile_ir, + "graph_analysis_digest": inferred.graph_analysis_digest, + "schedule_object": schedule_object.canonical_payload(), "schedule_digest": schedule_digest, } ) @@ -817,11 +995,15 @@ def validate_scheduled_spectral_metadata( raise ValueError("TSOL package requires tessera.scheduled_spectral.v5 metadata") shapes = tuple(tuple(int(dim) for dim in shape) for shape in input_shapes) declared_shapes = tuple( + tuple(int(dim) for dim in shape) for shape in metadata.get("input_shapes") or () + ) + signature = tuple( tuple(int(dim) for dim in shape) - for shape in metadata.get("input_shapes") or () + for shape in metadata.get("input_signature") or () + ) + bounds = tuple( + tuple(int(dim) for dim in shape) for shape in metadata.get("shape_bounds") or () ) - signature = tuple(tuple(int(dim) for dim in shape) for shape in metadata.get("input_signature") or ()) - bounds = tuple(tuple(int(dim) for dim in shape) for shape in metadata.get("shape_bounds") or ()) semantic = define_spectral_program_contract( op_name=str(metadata.get("op_name")), input_signature=signature, diff --git a/python/tessera/compiler/sharding_propagation.py b/python/tessera/compiler/sharding_propagation.py index f661606c8..8d3909ec1 100644 --- a/python/tessera/compiler/sharding_propagation.py +++ b/python/tessera/compiler/sharding_propagation.py @@ -8,15 +8,15 @@ from __future__ import annotations import hashlib +import math import json from dataclasses import dataclass, replace -from typing import Callable, Literal, Mapping, Sequence +from typing import Any, Callable, Literal, Mapping, Sequence from .effects import Effect, registered_op_effect -from .graph_ir import GraphIRFunction, IROp +from .graph_ir import GraphIRFunction, IROp, IRType, tensor_ir_type from .op_catalog import get_op_spec - PlacementKind = Literal["unknown", "replicated", "tiled", "partial_reduction"] @@ -115,6 +115,16 @@ class ReshardPlan: digest: str +@dataclass(frozen=True) +class MockMeshExecutionResult: + """Deterministic execution evidence for one materialized placement graph.""" + + values: Mapping[str, tuple[Any, ...]] + returned: Mapping[str, tuple[Any, ...]] + executed_reshards: tuple[str, ...] + backend: str = "deterministic_mock_mesh" + + def _ssa(value: str) -> str: return str(value).strip().lstrip("%") @@ -124,7 +134,9 @@ def _elementwise_rule(_op: IROp, operands: tuple[Placement, ...]) -> Placement: if not non_replicated: return Placement.replicated() first = non_replicated[0] - return first if all(item == first for item in non_replicated) else Placement.unknown() + return ( + first if all(item == first for item in non_replicated) else Placement.unknown() + ) def _reduction_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: @@ -143,12 +155,9 @@ def _reduction_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: if reduced_mesh_axis is None: return source shifted = { - (dim - 1 if dim > axis else dim): mesh_axis - for dim, mesh_axis in tiled.items() + (dim - 1 if dim > axis else dim): mesh_axis for dim, mesh_axis in tiled.items() } - return Placement.partial_reduction( - (reduced_mesh_axis,), tiled_axes=shifted - ) + return Placement.partial_reduction((reduced_mesh_axis,), tiled_axes=shifted) def _collective_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: @@ -158,12 +167,19 @@ def _collective_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: name = op.op_name.removeprefix("tessera.") if name == "all_reduce": if source.kind == "partial_reduction": - return (Placement.tiled(dict(source.tiled_axes)) - if source.tiled_axes else Placement.replicated()) + return ( + Placement.tiled(dict(source.tiled_axes)) + if source.tiled_axes + else Placement.replicated() + ) return source mesh_axis = op.kwargs.get("mesh_axis", op.kwargs.get("axis_name")) tensor_axis = op.kwargs.get("tensor_axis", op.kwargs.get("axis")) - if not isinstance(mesh_axis, str) or not mesh_axis or not isinstance(tensor_axis, int): + if ( + not isinstance(mesh_axis, str) + or not mesh_axis + or not isinstance(tensor_axis, int) + ): return Placement.unknown() if name == "all_gather" and source.kind == "tiled": tiled = dict(source.tiled_axes) @@ -172,7 +188,10 @@ def _collective_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: tiled.pop(tensor_axis) return Placement.tiled(tiled) if tiled else Placement.replicated() if name == "reduce_scatter" and source.kind in {"replicated", "partial_reduction"}: - if source.kind == "partial_reduction" and mesh_axis not in source.reduction_axes: + if ( + source.kind == "partial_reduction" + and mesh_axis not in source.reduction_axes + ): return Placement.unknown() return Placement.tiled({tensor_axis: mesh_axis}) if name == "all_to_all" and source.kind == "tiled": @@ -192,9 +211,16 @@ def _collective_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: **{ name: _elementwise_rule for name in ( - "tessera.add", "tessera.sub", "tessera.mul", "tessera.div", - "tessera.maximum", "tessera.minimum", "tessera.relu", - "tessera.exp", "tessera.log", "tessera.tanh", + "tessera.add", + "tessera.sub", + "tessera.mul", + "tessera.div", + "tessera.maximum", + "tessera.minimum", + "tessera.relu", + "tessera.exp", + "tessera.log", + "tessera.tanh", ) }, "tessera.reduce": _reduction_rule, @@ -205,7 +231,9 @@ def _collective_rule(op: IROp, operands: tuple[Placement, ...]) -> Placement: } -def _registered_rule(op: IROp, rules: Mapping[str, PlacementRule]) -> PlacementRule | None: +def _registered_rule( + op: IROp, rules: Mapping[str, PlacementRule] +) -> PlacementRule | None: direct = rules.get(op.op_name) if direct is not None: return direct @@ -215,8 +243,12 @@ def _registered_rule(op: IROp, rules: Mapping[str, PlacementRule]) -> PlacementR # contraction, attention, or stencil halo exchange. spec = get_op_spec(op.op_name) if spec is not None and spec.lowering in { - "elementwise", "numeric_helper", "comparison", "logical", - "rotary_embedding", "quantize", + "elementwise", + "numeric_helper", + "comparison", + "logical", + "rotary_embedding", + "quantize", }: return _elementwise_rule if spec is not None and spec.lowering == "stable_reduction": @@ -253,8 +285,11 @@ def propagate_sharding( ) rule = _registered_rule(op, rules) effect_admitted = effect == Effect.pure or rule is _collective_rule - if not effect_admitted or has_region or rule is None or any( - item.kind == "unknown" for item in operands + if ( + not effect_admitted + or has_region + or rule is None + or any(item.kind == "unknown" for item in operands) ): inferred = Placement.unknown() else: @@ -265,7 +300,9 @@ def propagate_sharding( tuple(map(_ssa, op.operands)), "incompatible_or_underspecified_placement", ) - conflicts[(conflict.op_name, conflict.operands, conflict.reason)] = conflict + conflicts[ + (conflict.op_name, conflict.operands, conflict.reason) + ] = conflict for result in op.result_names: name = _ssa(result) if placements.get(name) != inferred: @@ -286,7 +323,9 @@ def propagate_sharding( }, "conflicts": [ (item.op_name, item.operands, item.reason) - for item in sorted(conflicts.values(), key=lambda item: (item.op_name, item.operands)) + for item in sorted( + conflicts.values(), key=lambda item: (item.op_name, item.operands) + ) ], } digest = hashlib.sha256( @@ -343,22 +382,33 @@ def plan_explicit_reshards( "cannot reshard a value across sibling or escaping regions" ) action = _reshard_action( - value_name, op.op_name, op_index, operand_index, source, target, - subgroup=normalized_subgroup, region_path=consumer_region, + value_name, + op.op_name, + op_index, + operand_index, + source, + target, + subgroup=normalized_subgroup, + region_path=consumer_region, ) actions.append(action) - matching_rounds = _matching_rounds(normalized_subgroup) if any( - action.collective == "all_to_all" for action in actions - ) else () + matching_rounds = ( + _matching_rounds(normalized_subgroup) + if any(action.collective == "all_to_all" for action in actions) + else () + ) payload = [ { - "value": action.value, "consumer": action.consumer, + "value": action.value, + "consumer": action.consumer, "consumer_index": action.consumer_index, "operand_index": action.operand_index, "source": _placement_payload(action.source), "target": _placement_payload(action.target), - "collective": action.collective, "mesh_axes": action.mesh_axes, - "subgroup": action.subgroup, "region_path": action.region_path, + "collective": action.collective, + "mesh_axes": action.mesh_axes, + "subgroup": action.subgroup, + "region_path": action.region_path, } for action in actions ] @@ -380,10 +430,8 @@ def _region_path(op: IROp) -> tuple[str, ...]: raise ValueError("region path must be a string or sequence") -def _is_region_ancestor( - producer: tuple[str, ...], consumer: tuple[str, ...] -) -> bool: - return len(producer) <= len(consumer) and consumer[:len(producer)] == producer +def _is_region_ancestor(producer: tuple[str, ...], consumer: tuple[str, ...]) -> bool: + return len(producer) <= len(consumer) and consumer[: len(producer)] == producer def _matching_rounds( @@ -411,23 +459,33 @@ def _matching_rounds( def materialize_reshard_plan( function: GraphIRFunction, plan: ReshardPlan, + *, + mesh_shape: Mapping[str, int] | None = None, ) -> GraphIRFunction: - """Insert planned placement conversions as real Graph SSA operations. - - The returned function is a copy. Every collective is inserted immediately - before its exact consumer and rewires that operand. The plan digest, - subgroup, region identity, and all-to-all matching rounds survive into - Schedule IR through ordinary operation attributes. Replicated-to-tiled - local slicing remains fail-closed until a mesh-size-dependent result-shape - contract is available; emitting a same-shaped fake slice would be wrong. + """Insert typed placement conversions as real Graph SSA operations. + + Collective results use the exact local shape implied by the target + placement. Replicated-to-tiled movement is an operation-owned + ``tessera.slice`` with explicit reshard identity, not a same-shaped fake + collective. A missing, non-divisible, or dynamic mesh extent fails closed. """ + + mesh = _mesh_contract(plan, mesh_shape) by_consumer: dict[int, list[ReshardAction]] = {} for action in plan.actions: by_consumer.setdefault(action.consumer_index, []).append(action) body: list[IROp] = [] - used_names = { - _ssa(name) for op in function.body for name in op.result_names - } | {_ssa(arg.name) for arg in function.args} + used_names = {_ssa(name) for op in function.body for name in op.result_names} | { + _ssa(arg.name) for arg in function.args + } + value_types: dict[str, IRType] = { + _ssa(arg.name): arg.ir_type for arg in function.args + } + for source_op in function.body: + for result_name in source_op.result_names: + if source_op.inferred_type is not None: + value_types[_ssa(result_name)] = source_op.inferred_type + for op_index, source_op in enumerate(function.body): op = replace( source_op, @@ -435,24 +493,38 @@ def materialize_reshard_plan( operand_types=list(source_op.operand_types), kwargs=dict(source_op.kwargs), ) - for ordinal, action in enumerate(sorted( - by_consumer.get(op_index, ()), key=lambda item: item.operand_index - )): - if action.collective == "local_shard": - raise ValueError( - "local-shard SSA materialization requires a typed mesh-size shape contract" - ) + for ordinal, action in enumerate( + sorted( + by_consumer.get(op_index, ()), + key=lambda item: item.operand_index, + ) + ): name = f"reshard_{op_index}_{action.operand_index}_{ordinal}" suffix = 0 while name in used_names: suffix += 1 name = f"reshard_{op_index}_{action.operand_index}_{ordinal}_{suffix}" used_names.add(name) + source_type = value_types.get(action.value) + if source_type is None: + raise ValueError( + f"reshard value {action.value!r} lacks a structured tensor type" + ) + target_type = _local_shard_type(source_type, action.target, mesh) + tensor_axis = _collective_tensor_axis(action) kwargs: dict[str, object] = { "mesh_axis": action.mesh_axes[0], - "axis": _collective_tensor_axis(action), - "op": "sum" if action.collective in {"all_reduce", "reduce_scatter"} else "none", + "mesh_size": mesh[action.mesh_axes[0]], + "axis": tensor_axis, + "op": ( + "sum" + if action.collective in {"all_reduce", "reduce_scatter"} + else "none" + ), + "reshard_kind": action.collective, "reshard_plan_digest": plan.digest, + "source_global_shape": list(source_type.shape), + "target_local_shape": list(target_type.shape), "subgroup": list(action.subgroup), "region_path": list(action.region_path), } @@ -465,22 +537,200 @@ def materialize_reshard_plan( operand_type = op.operand_types[action.operand_index] inserted = IROp( result=name, - op_name=f"tessera.{action.collective}", + op_name=( + "tessera.slice" + if action.collective == "local_shard" + else f"tessera.{action.collective}" + ), operands=[op.operands[action.operand_index]], operand_types=[operand_type], - result_type=operand_type, - inferred_type=None, + result_type=str(target_type), + inferred_type=target_type, kwargs=kwargs, ) body.append(inserted) + value_types[name] = target_type op.operands[action.operand_index] = f"%{name}" + op.operand_types[action.operand_index] = str(target_type) body.append(op) return replace(function, body=body) +def _mesh_contract( + plan: ReshardPlan, mesh_shape: Mapping[str, int] | None +) -> dict[str, int]: + mesh = {str(axis): int(size) for axis, size in (mesh_shape or {}).items()} + for action in plan.actions: + for axis in action.mesh_axes: + if axis not in mesh and action.subgroup: + mesh[axis] = len(action.subgroup) + if axis not in mesh or mesh[axis] < 2: + raise ValueError( + f"reshard mesh axis {axis!r} requires an explicit size >= 2" + ) + if action.subgroup and len(action.subgroup) != mesh[axis]: + raise ValueError( + f"reshard subgroup size does not match mesh axis {axis!r}" + ) + return mesh + + +def _local_shard_type( + global_type: IRType, placement: Placement, mesh: Mapping[str, int] +) -> IRType: + if global_type.rank is None: + raise ValueError("typed reshard requires a statically ranked tensor") + shape = list(global_type.shape) + for dimension, mesh_axis in placement.tiled_axes: + if dimension >= len(shape): + raise ValueError("placement tensor axis is outside the tensor rank") + try: + extent = int(shape[dimension]) + except ValueError as error: + raise ValueError("typed reshard requires static shard extents") from error + size = mesh.get(mesh_axis) + if size is None or extent % size: + raise ValueError( + f"tensor extent {extent} is not divisible by mesh axis {mesh_axis!r}" + ) + shape[dimension] = str(extent // size) + return tensor_ir_type(tuple(shape), global_type.dtype, layout=global_type.layout) + + +def execute_resharded_graph_on_mock_mesh( + function: GraphIRFunction, + rank_inputs: Mapping[str, Sequence[Any]], + *, + mesh_shape: Mapping[str, int], +) -> MockMeshExecutionResult: + """Execute explicit reshard SSA on the deterministic in-process mesh. + + This interpreter is deliberately small and fail-closed: it handles the + five distributed movement forms plus the pointwise consumers used by the + placement gate. Every movement is visible in ``executed_reshards``. + """ + + import numpy as np + + from tessera.collectives import CollectiveAdapter + + mesh = {str(axis): int(size) for axis, size in mesh_shape.items()} + world_size = math.prod(mesh.values()) if mesh else 0 + if world_size < 2: + raise ValueError("mock reshard execution requires a multi-rank mesh") + adapter = CollectiveAdapter(backend="mock", world_size=world_size, mesh_axes=mesh) + values: dict[str, tuple[Any, ...]] = {} + for arg in function.args: + name = _ssa(arg.name) + per_rank = tuple(rank_inputs.get(name, ())) + if len(per_rank) != world_size: + raise ValueError(f"mock mesh input {name!r} requires one value per rank") + values[name] = tuple(np.asarray(value) for value in per_rank) + + executed: list[str] = [] + for op in function.body: + if op.result is None or len(op.result_names) != 1: + raise ValueError("mock reshard execution requires single-result SSA ops") + operands = [values.get(_ssa(name)) for name in op.operands] + if any(value is None for value in operands): + raise ValueError(f"mock mesh operation {op.op_name} has an unknown operand") + inputs = [tuple(value) for value in operands if value is not None] + name = _ssa(op.result_names[0]) + if ( + op.op_name == "tessera.slice" + and op.kwargs.get("reshard_kind") == "local_shard" + ): + axis = int(op.kwargs["axis"]) + mesh_axis = str(op.kwargs["mesh_axis"]) + size = mesh[mesh_axis] + outputs = [] + for rank, value in enumerate(inputs[0]): + if value.shape[axis] % size: + raise ValueError("mock local shard extent is not divisible") + coordinate = _mesh_axis_coordinate(rank, mesh, mesh_axis) + outputs.append(np.split(value, size, axis=axis)[coordinate]) + result = tuple(outputs) + executed.append("local_shard") + elif op.op_name == "tessera.all_reduce": + result = tuple( + adapter.all_reduce(inputs[0], op=str(op.kwargs.get("op", "sum"))) + ) + executed.append("all_reduce") + elif op.op_name == "tessera.reduce_scatter": + result = tuple( + adapter.reduce_scatter( + inputs[0], + axis=int(op.kwargs.get("axis", 0)), + op=str(op.kwargs.get("op", "sum")), + ) + ) + executed.append("reduce_scatter") + elif op.op_name == "tessera.all_gather": + result = tuple( + adapter.all_gather(inputs[0], axis=int(op.kwargs.get("axis", 0))) + ) + executed.append("all_gather") + elif op.op_name == "tessera.all_to_all": + result = tuple( + adapter.all_to_all( + inputs[0], + scatter_axis=int(op.kwargs.get("scatter_axis", 0)), + gather_axis=int(op.kwargs.get("gather_axis", 0)), + ) + ) + executed.append("all_to_all") + elif op.op_name == "tessera.collective_permute": + sources = tuple(int(item) for item in op.kwargs.get("source_peers", ())) + targets = tuple(int(item) for item in op.kwargs.get("target_peers", ())) + result = tuple( + adapter.collective_permute( + inputs[0], pairs=tuple(zip(sources, targets)) + ) + ) + executed.append("collective_permute") + elif op.op_name in {"tessera.add", "tessera.sub", "tessera.mul", "tessera.div"}: + operation: Any = { + "tessera.add": np.add, + "tessera.sub": np.subtract, + "tessera.mul": np.multiply, + "tessera.div": np.divide, + }[op.op_name] + result = tuple(operation(*rank_values) for rank_values in zip(*inputs)) + elif op.op_name in {"tessera.sigmoid", "tessera.tanh", "tessera.relu"}: + operation = { + "tessera.sigmoid": lambda value: 1.0 / (1.0 + np.exp(-value)), + "tessera.tanh": np.tanh, + "tessera.relu": lambda value: np.maximum(value, 0), + }[op.op_name] + result = tuple(operation(value) for value in inputs[0]) + else: + raise ValueError(f"mock mesh has no execution rule for {op.op_name!r}") + values[name] = tuple(np.asarray(value) for value in result) + + returned = {_ssa(name): values[_ssa(name)] for name in function.return_values} + return MockMeshExecutionResult( + values=dict(values), + returned=returned, + executed_reshards=tuple(executed), + ) + + +def _mesh_axis_coordinate( + rank: int, mesh_shape: Mapping[str, int], mesh_axis: str +) -> int: + axes = tuple(mesh_shape) + if mesh_axis not in axes: + raise ValueError(f"unknown mesh axis {mesh_axis!r}") + index = axes.index(mesh_axis) + stride = math.prod(tuple(mesh_shape[axis] for axis in axes[index + 1 :])) + return (rank // stride) % mesh_shape[mesh_axis] + + def _single_tiled_dim(placement: Placement) -> int: if len(placement.tiled_axes) != 1: - raise ValueError("collective SSA currently requires exactly one tiled tensor axis") + raise ValueError( + "collective SSA currently requires exactly one tiled tensor axis" + ) return placement.tiled_axes[0][0] @@ -555,8 +805,16 @@ def _reshard_action( else: raise ValueError(f"unsupported reshard transition {source.kind}->{target.kind}") return ReshardAction( - value, consumer, consumer_index, operand_index, source, target, - collective, axes, subgroup, region_path, + value, + consumer, + consumer_index, + operand_index, + source, + target, + collective, + axes, + subgroup, + region_path, ) @@ -569,6 +827,8 @@ def _reshard_action( "ReshardAction", "ReshardPlan", "ShardingPropagationResult", + "MockMeshExecutionResult", + "execute_resharded_graph_on_mock_mesh", "materialize_reshard_plan", "plan_explicit_reshards", "propagate_sharding", diff --git a/src/compiler/programming_model/lib/PMPasses.cpp b/src/compiler/programming_model/lib/PMPasses.cpp index 89d058938..996e1a443 100644 --- a/src/compiler/programming_model/lib/PMPasses.cpp +++ b/src/compiler/programming_model/lib/PMPasses.cpp @@ -853,73 +853,6 @@ static std::string fftScheduleDigest(const FFTSchedule &schedule) { /*LowerCase=*/true); } -static FailureOr spectralProgramDigest(Operation *op) { - auto stringAttr = [&](StringRef name) -> StringAttr { - return op->getAttrOfType(name); - }; - auto intAttr = [&](StringRef name) -> IntegerAttr { - return op->getAttrOfType(name); - }; - auto output = op->getAttrOfType("output_shape"); - auto padding = op->getAttrOfType("padding"); - auto crop = op->getAttrOfType("crop"); - for (StringRef name : {"target", "arch", "kind", "input_shapes", - "input_signature", "shape_bounds", "template_digest", - "shape_policy", "storage", "abi_storage", - "storage_conversion", "axis_packing", "normalization", - "complex_layout", "accumulation", - "workspace_policy", "fusion_topology", "mutation_lineage", - "native_entry", "child_fft_digests"}) - if (!stringAttr(name)) return failure(); - for (StringRef name : {"axis", "dct_type", "window_length", "hop", "frames", - "workspace_bytes", "workgroup_size"}) - if (!intAttr(name)) return failure(); - if (!output || !padding || !crop) return failure(); - auto arrayText = [](ArrayRef values, StringRef separator) { - std::string result; - for (int64_t value : values) { - if (!result.empty()) result += separator; - result += Twine(value).str(); - } - return result; - }; - std::string contract = - (Twine("schema=tessera.scheduled_spectral.v5;op=") + - stringAttr("kind").getValue() + ";target=" + - stringAttr("target").getValue() + ";arch=" + - stringAttr("arch").getValue() + ";inputs=" + - stringAttr("input_shapes").getValue() + ";output=" + - arrayText(output.asArrayRef(), "x") + ";axis=" + - Twine(intAttr("axis").getInt()) + ";dct_type=" + - Twine(intAttr("dct_type").getInt()) + ";shape_policy=" + - stringAttr("shape_policy").getValue() + ";storage=" + - stringAttr("storage").getValue() + ";abi_storage=" + - stringAttr("abi_storage").getValue() + ";storage_conversion=" + - stringAttr("storage_conversion").getValue() + ";axis_packing=" + - stringAttr("axis_packing").getValue() + ";input_signature=" + - stringAttr("input_signature").getValue() + ";shape_bounds=" + - stringAttr("shape_bounds").getValue() + ";template_digest=" + - stringAttr("template_digest").getValue() + ";padding=" + - arrayText(padding.asArrayRef(), ",") + ";crop=" + - arrayText(crop.asArrayRef(), ",") + ";window=" + - Twine(intAttr("window_length").getInt()) + ";hop=" + - Twine(intAttr("hop").getInt()) + ";frames=" + - Twine(intAttr("frames").getInt()) + ";normalization=" + - stringAttr("normalization").getValue() + ";complex_layout=" + - stringAttr("complex_layout").getValue() + ";accumulation=" + - stringAttr("accumulation").getValue() + ";workspace_bytes=" + - Twine(intAttr("workspace_bytes").getInt()) + ";workspace_policy=" + - stringAttr("workspace_policy").getValue() + ";fusion_topology=" + - stringAttr("fusion_topology").getValue() + ";mutation_lineage=" + - stringAttr("mutation_lineage").getValue() + ";native_entry=" + - stringAttr("native_entry").getValue() + ";child_fft_digests=" + - stringAttr("child_fft_digests").getValue() + ";workgroup=" + - Twine(intAttr("workgroup_size").getInt())) - .str(); - return llvm::toHex(llvm::SHA256::hash(llvm::arrayRefFromStringRef(contract)), - /*LowerCase=*/true); -} - static std::string spectralBackwardTypeSignature(ValueRange values) { std::string signature; llvm::raw_string_ostream stream(signature); @@ -3624,12 +3557,64 @@ struct ScheduleToTilePass }); for (Operation *scheduled : scheduledSpectralPrograms) { auto hash = scheduled->getAttrOfType("artifact_hash"); - FailureOr derived = spectralProgramDigest(scheduled); - if (failed(derived) || !hash || *derived != hash.getValue()) { + auto scheduleDigest = + mod->getAttrOfType("tessera.schedule_digest"); + if (!hash || !scheduleDigest || hash.getValue().size() != 64 || + hash.getValue() != scheduleDigest.getValue()) { scheduled->emitError( - "scheduled spectral program policy was altered after hashing"); + "scheduled spectral program requires the module Schedule Object digest"); return signalPassFailure(); } + // Matching digest strings alone do NOT bind the attributes this pass + // then consumes: a cached or hand-edited program could keep both + // digests and still change workspace_bytes or native_entry, which the + // native launch would use (PR #626 review). Re-verify every consumed + // policy value against the semantic payload the producer carried. + auto semantic = + mod->getAttrOfType("tessera.spectral_semantic"); + if (!semantic || semantic.getValue().empty()) { + scheduled->emitError( + "scheduled spectral program requires the module " + "tessera.spectral_semantic payload to verify consumed policy"); + return signalPassFailure(); + } + auto payloadField = [&](StringRef field) -> StringRef { + StringRef payload = semantic.getValue(); + size_t at = payload.find((field + "=").str()); + while (at != StringRef::npos && at != 0 && payload[at - 1] != ';') + at = payload.find((field + "=").str(), at + 1); + if (at == StringRef::npos) return StringRef(); + StringRef rest = payload.drop_front(at + field.size() + 1); + size_t end = rest.find(';'); + return end == StringRef::npos ? rest : rest.take_front(end); + }; + auto requireMatches = [&](StringRef field, const Twine &actual) -> bool { + StringRef declared = payloadField(field); + if (declared.empty() || declared != actual.str()) { + scheduled->emitError("scheduled spectral program attribute '") + << field << "' is " << actual + << ", which disagrees with the carried semantic payload ('" + << declared << "'); the digest does not cover a changed policy"; + return false; + } + return true; + }; + if (auto workspaceAttr = + scheduled->getAttrOfType("workspace_bytes")) { + if (!requireMatches("workspace_bytes", + Twine(workspaceAttr.getInt()))) + return signalPassFailure(); + } + if (auto entryAttr = + scheduled->getAttrOfType("native_entry")) { + if (!requireMatches("native_entry", entryAttr.getValue())) + return signalPassFailure(); + } + if (auto normalizationAttr = + scheduled->getAttrOfType("normalization")) { + if (!requireMatches("normalization", normalizationAttr.getValue())) + return signalPassFailure(); + } SmallVector matchingArtifacts; mod.walk([&](schedule::ArtifactOp artifact) { if (artifact.getHash() == hash.getValue()) diff --git a/src/transforms/include/Tessera/Transforms/Passes.h b/src/transforms/include/Tessera/Transforms/Passes.h index 7934a0f87..8c792c5b0 100644 --- a/src/transforms/include/Tessera/Transforms/Passes.h +++ b/src/transforms/include/Tessera/Transforms/Passes.h @@ -131,8 +131,8 @@ std::unique_ptr createNVFlashAttnKernelEmitterPass(int sm = 90); // Pipeline order (after Phase 2/3 distribution + effect annotation): // 1. tessera-gpu-collective-insertion — insert reduce_scatter/all_gather // at DP/TP mesh boundaries (reads tessera.weight_sharding + tessera.effect) -// 2. tessera-pipeline-stage-insertion — 1F1B stage split; insert send/recv -// at PP stage boundaries (reads tessera.pipeline_plan on module) +// 2. tessera-pipeline-stage-insertion — insert send/recv at PP boundaries +// from the digest-bound tessera.pipeline_steps Schedule Object carrier // GPUCollectiveInsertionPass — inserts collective.reduce_scatter at // data-parallel gradient boundaries and collective.all_gather at tensor- diff --git a/src/transforms/lib/PipelineScheduleLegalityPass.cpp b/src/transforms/lib/PipelineScheduleLegalityPass.cpp index a39d81091..22b10fb1f 100644 --- a/src/transforms/lib/PipelineScheduleLegalityPass.cpp +++ b/src/transforms/lib/PipelineScheduleLegalityPass.cpp @@ -7,9 +7,8 @@ // // --tessera-pipeline-schedule-legality // -// Reads the pipeline plan from the module: `tessera.pp_num_stages` / -// `tessera.pp_num_micro_batches` / `tessera.pp_interleaved` (set by the -// insertion pass) or the `tessera.pipeline_plan` dict. Invariants: +// Reads the digest-bound `tessera.pipeline_steps` carrier and its materialized +// scalar views. It never reconstructs steps from a `tessera.pipeline_plan`. // // PP_MICRO_BATCHES_TOO_FEW // 1F1B needs micro_batches >= num_stages to fill the pipe; interleaved 1F1B @@ -42,21 +41,11 @@ #include "mlir/Pass/Pass.h" #include "llvm/ADT/DenseSet.h" -#include using namespace mlir; namespace { -static int64_t readPlan(ModuleOp m, StringRef ppKey, StringRef planKey, - int64_t defaultVal) { - if (auto v = m->getAttrOfType(ppKey)) - return v.getInt(); - if (auto plan = m->getAttrOfType("tessera.pipeline_plan")) - if (auto v = plan.getAs(planKey)) - return v.getInt(); - return defaultVal; -} static int64_t opStage(Operation *op) { if (auto layer = op->getAttrOfType("tessera.layer")) @@ -67,51 +56,6 @@ static int64_t opStage(Operation *op) { return -1; } -static void materializeSchedule(ModuleOp module, int64_t numStages, - int64_t microBatches) { - // Materialize a conservative dependency order rather than merely proving - // that a schedule could exist. One action owns each logical clock, so no - // rank can be asked to execute forward and backward simultaneously. A - // later target/runtime planner may overlap independent actions while - // preserving this order. - Builder b(module.getContext()); - SmallVector steps; - int64_t clock = 0; - auto append = [&](StringRef region, StringRef phase, int64_t microBatch, - int64_t stage) { - steps.push_back(b.getDictionaryAttr({ - b.getNamedAttr("clock", b.getI64IntegerAttr(clock++)), - b.getNamedAttr("micro_batch", b.getI64IntegerAttr(microBatch)), - b.getNamedAttr("phase", b.getStringAttr(phase)), - b.getNamedAttr("region", b.getStringAttr(region)), - b.getNamedAttr("stage", b.getI64IntegerAttr(stage)), - })); - }; - auto forwardSweep = [&](StringRef region, int64_t microBatch) { - for (int64_t stage = 0; stage < numStages; ++stage) - append(region, "forward", microBatch, stage); - }; - auto backwardSweep = [&](StringRef region, int64_t microBatch) { - for (int64_t stage = numStages; stage-- > 0;) - append(region, "backward", microBatch, stage); - }; - - int64_t warmupBatches = std::min(numStages - 1, microBatches); - for (int64_t mb = 0; mb < warmupBatches; ++mb) - forwardSweep("warmup", mb); - int64_t steadyBatches = microBatches - warmupBatches; - for (int64_t mb = 0; mb < steadyBatches; ++mb) { - forwardSweep("steady", mb + warmupBatches); - backwardSweep("steady", mb); - } - for (int64_t mb = steadyBatches; mb < microBatches; ++mb) - backwardSweep("cooldown", mb); - - module->setAttr("tessera.pipeline_steps", b.getArrayAttr(steps)); - module->setAttr( - "tessera.pipeline_schedule_kind", - b.getStringAttr("1f1b.serialized_dependency_order.v1")); -} struct PipelineScheduleLegalityPass : public PassWrapper> { @@ -121,20 +65,139 @@ struct PipelineScheduleLegalityPass return "tessera-pipeline-schedule-legality"; } StringRef getDescription() const override { - return "Prove 1F1B legality and materialize explicit " - "warmup/steady/cooldown dependency steps."; + return "Prove 1F1B legality against the materialized Schedule Object " + "dependency carrier."; } void runOnOperation() override { ModuleOp module = getOperation(); - int64_t numStages = readPlan(module, "tessera.pp_num_stages", "num_stages", 1); - if (numStages <= 1) - return; // no pipeline. + auto scheduleDigest = + module->getAttrOfType("tessera.schedule_digest"); + auto scheduleSchema = + module->getAttrOfType("tessera.pipeline_schedule_schema"); + auto scheduleSteps = + module->getAttrOfType("tessera.pipeline_steps"); + auto numStagesAttr = + module->getAttrOfType("tessera.pp_num_stages"); + auto microBatchesAttr = + module->getAttrOfType("tessera.pp_num_micro_batches"); + auto interleavedAttr = + module->getAttrOfType("tessera.pp_interleaved"); + if (!scheduleDigest || scheduleDigest.getValue().size() != 64 || + !scheduleSchema || + scheduleSchema.getValue() != "tessera.pipeline_schedule.v1" || + !scheduleSteps || scheduleSteps.empty() || !numStagesAttr || + !microBatchesAttr || !interleavedAttr) { + module.emitError( + "pipeline legality requires one complete digest-bound " + "tessera.pipeline_schedule.v1 carrier"); + signalPassFailure(); + return; + } + // The carrier IS the executable authority, so its rows are validated + // rather than counted: unique action ids, resolvable dependencies, an + // acyclic dependency order, and agreement with the declared pipeline + // dimensions. Previously only `!empty()` was checked, so a stale or + // hand-edited carrier with duplicate ids, dangling dependencies, or a + // cycle was accepted as authoritative (PR #626 review). + { + llvm::DenseSet seen; + llvm::SmallVector order; + llvm::DenseMap> requires; + int64_t declaredStages = numStagesAttr.getInt(); + int64_t declaredMicroBatches = microBatchesAttr.getInt(); + for (Attribute entry : scheduleSteps) { + auto row = dyn_cast(entry); + if (!row) { + module.emitError("pipeline step carrier row is not a dictionary"); + signalPassFailure(); + return; + } + auto actionId = row.getAs("action_id"); + auto stage = row.getAs("stage"); + auto microBatch = row.getAs("micro_batch"); + auto clock = row.getAs("clock"); + auto dependsOn = row.getAs("depends_on"); + if (!actionId || actionId.getValue().empty() || !stage || + !microBatch || !clock || !dependsOn) { + module.emitError( + "pipeline step carrier row requires action_id, stage, " + "micro_batch, clock, and depends_on"); + signalPassFailure(); + return; + } + if (!seen.insert(actionId.getValue()).second) { + module.emitError("duplicate pipeline action id '") + << actionId.getValue() << "'"; + signalPassFailure(); + return; + } + if (clock.getInt() < 0 || microBatch.getInt() < 0 || + microBatch.getInt() >= declaredMicroBatches || stage.getInt() < 0) { + module.emitError("pipeline step '") + << actionId.getValue() + << "' disagrees with the declared pipeline dimensions"; + signalPassFailure(); + return; + } + // Virtual stages run 0..num_stages*num_chunks-1; without chunking the + // bound is num_stages. + int64_t chunks = 1; + if (auto chunkAttr = + module->getAttrOfType("tessera.pp_num_chunks")) + chunks = std::max(1, chunkAttr.getInt()); + if (stage.getInt() >= declaredStages * chunks) { + module.emitError("pipeline step '") + << actionId.getValue() << "' names virtual stage " + << stage.getInt() << " beyond the declared " + << (declaredStages * chunks); + signalPassFailure(); + return; + } + llvm::SmallVector deps; + for (Attribute dependency : dependsOn) { + auto dependencyId = dyn_cast(dependency); + if (!dependencyId) { + module.emitError("pipeline dependency is not a string"); + signalPassFailure(); + return; + } + deps.push_back(dependencyId.getValue()); + } + requires[actionId.getValue()] = deps; + order.push_back(actionId.getValue()); + } + // Dependencies must resolve, and the carrier's own order must be a + // topological order (a producer listed after its consumer would let the + // consumer issue first). + llvm::DenseMap position; + for (auto [index, id] : llvm::enumerate(order)) position[id] = index; + for (StringRef id : order) { + for (StringRef dependency : requires[id]) { + auto found = position.find(dependency); + if (found == position.end()) { + module.emitError("pipeline action '") + << id << "' depends on unknown action '" << dependency << "'"; + signalPassFailure(); + return; + } + if (found->second >= position[id]) { + module.emitError("pipeline action '") + << id << "' depends on '" << dependency + << "', which the carrier orders no earlier (cycle or " + "producer-after-consumer)"; + signalPassFailure(); + return; + } + } + } + } - int64_t microBatches = - readPlan(module, "tessera.pp_num_micro_batches", "num_micro_batches", 1); - bool interleaved = readPlan(module, "tessera.pp_interleaved", "interleaved", - 0) != 0; + int64_t numStages = numStagesAttr.getInt(); + if (numStages <= 1) + return; + int64_t microBatches = microBatchesAttr.getInt(); + bool interleaved = interleavedAttr.getValue(); bool anyError = false; // ── Micro-batch fill contract (Decision #17) ── @@ -212,7 +275,6 @@ struct PipelineScheduleLegalityPass signalPassFailure(); return; } - materializeSchedule(module, numStages, microBatches); } }; diff --git a/src/transforms/lib/PipelineStageInsertionPass.cpp b/src/transforms/lib/PipelineStageInsertionPass.cpp index 0d9fe369b..6f0589d7e 100644 --- a/src/transforms/lib/PipelineStageInsertionPass.cpp +++ b/src/transforms/lib/PipelineStageInsertionPass.cpp @@ -1,31 +1,12 @@ //===- PipelineStageInsertionPass.cpp — Phase 4 ───────────────────────────===// // -// Partitions the IR into pipeline stages and inserts micro-batch send/recv -// communication ops at stage boundaries. -// -// The 1F1B (one-forward-one-backward) schedule is computed from: -// tessera.pipeline_plan = {num_stages, num_micro_batches, interleaved, …} -// on the module op (emitted by PipelinePlan::to_mlir_attrs()). -// -// What the pass does: -// 1. Reads num_stages (p) and num_micro_batches (m) from the module attr. -// 2. Splits the function body into `p` sequential schedule.pipeline.stage -// regions, one per device rank. -// 3. Inserts `tessera.pipeline.send` at each stage's output and -// `tessera.pipeline.recv` at the next stage's input. -// 4. Annotates each stage region with: -// {tessera.pp_stage = k, tessera.pp_num_micro_batches = m} -// -// For the current implementation, "splitting" is approximated by annotating -// ops that carry a `tessera.layer = {stage = k}` attribute (set by -// DistributedPlan layer specs). Full SSA splitting would require a complete -// program-partitioning pass — that is deferred to Phase 5. +// Inserts send/recv communication at already materialized pipeline boundaries. +// The one schedule authority is the content-addressed Schedule Object emitted +// by PipelinePlan. Its digest and dependency steps survive in module IR; this +// pass consumes that carrier and never reconstructs a schedule from scalar +// options or a parallel tessera.pipeline_plan dictionary. // // Registration: --tessera-pipeline-stage-insertion -// Options: -// --num-stages pipeline stage count (overrides module attr) -// --num-micro-batches micro-batch count (overrides module attr) -// --interleaved use interleaved 1F1B (default false) // //===----------------------------------------------------------------------===// @@ -50,14 +31,6 @@ namespace { // Helpers //===----------------------------------------------------------------------===// -/// Read an integer attribute from the module's tessera.pipeline_plan dict. -static int64_t readPlanAttr(ModuleOp m, StringRef key, int64_t defaultVal) { - auto planAttr = m->getAttrOfType("tessera.pipeline_plan"); - if (!planAttr) return defaultVal; - if (auto v = planAttr.getAs(key)) - return v.getInt(); - return defaultVal; -} /// Get the pipeline stage assigned to an op via `tessera.layer = {stage = k}`. static int64_t getOpStage(Operation *op) { @@ -74,10 +47,12 @@ static int64_t getOpStage(Operation *op) { /// Emit a `tessera.pipeline.send` op carrying the activation tensor to the /// next pipeline stage. static void emitPipelineSend(OpBuilder &b, Location loc, Value activation, - int64_t fromStage, int64_t microBatch) { + int64_t fromStage, int64_t microBatch, + StringAttr scheduleDigest) { OperationState state(loc, "tessera.pipeline.send"); state.addOperands(activation); state.addAttribute("from_stage", b.getI64IntegerAttr(fromStage)); + state.addAttribute("tessera.schedule_digest", scheduleDigest); state.addAttribute("micro_batch", b.getI64IntegerAttr(microBatch)); b.create(state); @@ -89,9 +64,11 @@ static void emitPipelineSend(OpBuilder &b, Location loc, Value activation, /// Emit a `tessera.pipeline.recv` op that receives the activation from the /// previous pipeline stage and returns it as a new value. static Value emitPipelineRecv(OpBuilder &b, Location loc, Type activationType, - int64_t toStage, int64_t microBatch) { + int64_t toStage, int64_t microBatch, + StringAttr scheduleDigest) { OperationState state(loc, "tessera.pipeline.recv"); state.addAttribute("to_stage", b.getI64IntegerAttr(toStage)); + state.addAttribute("tessera.schedule_digest", scheduleDigest); state.addAttribute("micro_batch", b.getI64IntegerAttr(microBatch)); state.addTypes(activationType); Operation *op = b.create(state); @@ -115,18 +92,6 @@ struct PipelineStageInsertionPass PipelineStageInsertionPass(const PipelineStageInsertionPass &other) : PassWrapper(other) {} - Option numStagesOpt{ - *this, "num-stages", - llvm::cl::desc("Pipeline stage count (overrides module attr)"), - llvm::cl::init(0)}; - Option numMicroBatchesOpt{ - *this, "num-micro-batches", - llvm::cl::desc("Micro-batch count (overrides module attr)"), - llvm::cl::init(0)}; - Option interleavedOpt{ - *this, "interleaved", - llvm::cl::desc("Use interleaved 1F1B schedule"), - llvm::cl::init(false)}; StringRef getArgument() const override { return "tessera-pipeline-stage-insertion"; @@ -139,16 +104,34 @@ struct PipelineStageInsertionPass ModuleOp module = getOperation(); OpBuilder b(module.getContext()); - // ── Read pipeline parameters ───────────────────────────────────────── - int64_t numStages = numStagesOpt > 0 - ? numStagesOpt - : readPlanAttr(module, "num_stages", 1); - int64_t numMicroBatches = numMicroBatchesOpt > 0 - ? numMicroBatchesOpt - : readPlanAttr(module, "num_micro_batches", 1); - bool interleaved = interleavedOpt - ? (bool)interleavedOpt - : (readPlanAttr(module, "interleaved", 0) != 0); + // Consume the materialized Schedule Object carrier. Scalar plan dicts and + // pass-option overrides are intentionally not accepted as schedule data. + auto scheduleDigest = + module->getAttrOfType("tessera.schedule_digest"); + auto scheduleSchema = + module->getAttrOfType("tessera.pipeline_schedule_schema"); + auto scheduleSteps = + module->getAttrOfType("tessera.pipeline_steps"); + auto numStagesAttr = + module->getAttrOfType("tessera.pp_num_stages"); + auto numMicroBatchesAttr = + module->getAttrOfType("tessera.pp_num_micro_batches"); + auto interleavedAttr = + module->getAttrOfType("tessera.pp_interleaved"); + if (!scheduleDigest || scheduleDigest.getValue().size() != 64 || + !scheduleSchema || + scheduleSchema.getValue() != "tessera.pipeline_schedule.v1" || + !scheduleSteps || scheduleSteps.empty() || !numStagesAttr || + !numMicroBatchesAttr || !interleavedAttr) { + module.emitError( + "pipeline stage insertion requires one complete digest-bound " + "tessera.pipeline_schedule.v1 carrier"); + signalPassFailure(); + return; + } + int64_t numStages = numStagesAttr.getInt(); + int64_t numMicroBatches = numMicroBatchesAttr.getInt(); + bool interleaved = interleavedAttr.getValue(); if (numStages <= 1) { // Nothing to do — single stage, no pipeline boundaries @@ -161,10 +144,11 @@ struct PipelineStageInsertionPass << " micro_batches=" << numMicroBatches << (interleaved ? " interleaved" : " standard-1F1B") << "\n"); - // ── Annotate module with pipeline plan ─────────────────────────────── - module->setAttr("tessera.pp_num_stages", b.getI64IntegerAttr(numStages)); - module->setAttr("tessera.pp_num_micro_batches", b.getI64IntegerAttr(numMicroBatches)); - module->setAttr("tessera.pp_interleaved", b.getBoolAttr(interleaved)); + // Stamp the same content identity on every owning function. The resource + // vectors and reasoned edges remain in the out-of-band Schedule Object. + module.walk([&](func::FuncOp func) { + func->setAttr("tessera.schedule_digest", scheduleDigest); + }); // ── Group ops by pipeline stage ────────────────────────────────────── // Collect (stage → ops) mapping across all functions @@ -217,7 +201,8 @@ struct PipelineStageInsertionPass // Emit exactly one send after the producer. b.setInsertionPointAfter(op); - emitPipelineSend(b, op->getLoc(), result, stage, /*mb=*/0); + emitPipelineSend(b, op->getLoc(), result, stage, /*mb=*/0, + scheduleDigest); ++sendCount; // Insert one recv before the first stage+1 consumer and rewire the @@ -226,7 +211,8 @@ struct PipelineStageInsertionPass if (getOpStage(user) == stage + 1) { b.setInsertionPoint(user); Value recvVal = emitPipelineRecv( - b, user->getLoc(), result.getType(), stage + 1, /*mb=*/0); + b, user->getLoc(), result.getType(), stage + 1, /*mb=*/0, + scheduleDigest); for (OpOperand &use : llvm::make_early_inc_range(result.getUses())) { if (getOpStage(use.getOwner()) == stage + 1) use.set(recvVal); @@ -240,10 +226,13 @@ struct PipelineStageInsertionPass } // ── Annotate schedule.pipeline.region ops ──────────────────────────── + // Preserve the owning Schedule Object identity on every pipeline region. module.walk([&](Operation *op) { if (op->getName().getStringRef().contains("schedule.pipeline")) { - op->setAttr("tessera.pp_num_stages", b.getI64IntegerAttr(numStages)); - op->setAttr("tessera.pp_num_micro_batches", b.getI64IntegerAttr(numMicroBatches)); + op->setAttr("tessera.pp_num_stages", b.getI64IntegerAttr(numStages)); + op->setAttr("tessera.pp_num_micro_batches", + b.getI64IntegerAttr(numMicroBatches)); + op->setAttr("tessera.schedule_digest", scheduleDigest); } }); diff --git a/src/transforms/lib/PipelineStagePartitionPass.cpp b/src/transforms/lib/PipelineStagePartitionPass.cpp index 8f38739e4..64ae4335d 100644 --- a/src/transforms/lib/PipelineStagePartitionPass.cpp +++ b/src/transforms/lib/PipelineStagePartitionPass.cpp @@ -9,7 +9,7 @@ // send/recv SSA rewrites and PipelineScheduleLegalityPass proves the 1F1B // schedule is well-formed. // -// --tessera-pipeline-partition (--num-stages N overrides the module attr) +// --tessera-pipeline-partition // // Cost model: heavy ops (matmul / gemm / batched_gemm / flash_attn / // conv2d_nhwc) weigh 4, everything else 1. Stages are contiguous in program @@ -36,12 +36,6 @@ using namespace mlir; namespace { -static int64_t readPlanNumStages(ModuleOp m, int64_t defaultVal) { - if (auto plan = m->getAttrOfType("tessera.pipeline_plan")) - if (auto v = plan.getAs("num_stages")) - return v.getInt(); - return defaultVal; -} // Per-op cost — a small cost model so heavy compute dominates the balance. static int64_t opCost(Operation *op) { @@ -68,10 +62,6 @@ struct PipelineStagePartitionPass PipelineStagePartitionPass(const PipelineStagePartitionPass &o) : PassWrapper(o) {} - Option numStagesOpt{ - *this, "num-stages", - llvm::cl::desc("Pipeline stage count (overrides module attr)"), - llvm::cl::init(0)}; StringRef getArgument() const override { return "tessera-pipeline-partition"; @@ -84,8 +74,20 @@ struct PipelineStagePartitionPass void runOnOperation() override { ModuleOp module = getOperation(); OpBuilder b(module.getContext()); - int64_t numStages = - numStagesOpt > 0 ? numStagesOpt : readPlanNumStages(module, 1); + auto scheduleDigest = + module->getAttrOfType("tessera.schedule_digest"); + auto scheduleSteps = + module->getAttrOfType("tessera.pipeline_steps"); + auto numStagesAttr = + module->getAttrOfType("tessera.pp_num_stages"); + if (!scheduleDigest || scheduleDigest.getValue().size() != 64 || + !scheduleSteps || scheduleSteps.empty() || !numStagesAttr) { + module.emitError( + "pipeline partition requires one complete digest-bound schedule carrier"); + signalPassFailure(); + return; + } + int64_t numStages = numStagesAttr.getInt(); if (numStages <= 1) return; // single stage — nothing to partition. diff --git a/tests/tessera-ir/phase2/x86_composed_layout_exec.mlir b/tests/tessera-ir/phase2/x86_composed_layout_exec.mlir index de05f9821..bbc4bb980 100644 --- a/tests/tessera-ir/phase2/x86_composed_layout_exec.mlir +++ b/tests/tessera-ir/phase2/x86_composed_layout_exec.mlir @@ -1,3 +1,8 @@ +// This file is fixture DATA for tests/unit/test_x86_composed_layout_exec.py, +// which drives it through tessera-tile-to-x86 -> LLVM -> lli. It is not a lit +// test: without this marker lit discovers it, reports Unresolved ("Test has no +// 'RUN:' line"), and fails `check-tessera-ir` for the whole repository. +// UNSUPPORTED: true // Exact CPU execution fixture. The engineering gate lowers this file through // tessera-tile-to-x86, upstream LLVM conversion, and lli; main returns zero // only when dynamic, nested mixed-radix, and tuple-product results all match. diff --git a/tests/tessera-ir/phase4/pipeline_partition.mlir b/tests/tessera-ir/phase4/pipeline_partition.mlir index 3cc497683..72307bf4e 100644 --- a/tests/tessera-ir/phase4/pipeline_partition.mlir +++ b/tests/tessera-ir/phase4/pipeline_partition.mlir @@ -7,7 +7,12 @@ // for. Two equal-cost matmuls over 2 stages → stage 0 then stage 1. module attributes { - tessera.pipeline_plan = {num_stages = 2, num_micro_batches = 2, interleaved = false} + tessera.schedule_digest = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "fixture", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 2, + tessera.pp_num_micro_batches = 2, + tessera.pp_interleaved = false } { // CHECK-LABEL: func.func @pipeline func.func @pipeline(%x: tensor<64x128xbf16>, %w0: tensor<128x256xbf16>, diff --git a/tests/tessera-ir/phase4/pipeline_schedule_legality.mlir b/tests/tessera-ir/phase4/pipeline_schedule_legality.mlir index 63dd4f19f..b24f22562 100644 --- a/tests/tessera-ir/phase4/pipeline_schedule_legality.mlir +++ b/tests/tessera-ir/phase4/pipeline_schedule_legality.mlir @@ -8,7 +8,12 @@ // 0->1 boundary gets a paired send/recv, and the 1F1B schedule verifies. // expected-remark@+1 {{pipeline-stage-insertion}} module attributes { - tessera.pipeline_plan = {num_stages = 2, num_micro_batches = 2, interleaved = false} + tessera.schedule_digest = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "ok", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 2, + tessera.pp_num_micro_batches = 2, + tessera.pp_interleaved = false } { func.func @ok(%x: tensor<64x128xbf16>, %w0: tensor<128x256xbf16>, %w1: tensor<256x128xbf16>) -> tensor<64x128xbf16> { @@ -24,7 +29,12 @@ module attributes { // two independent matmuls have no 0->1 dataflow, so no comm / remark is emitted. // expected-error@+1 {{PP_MICRO_BATCHES_TOO_FEW}} module attributes { - tessera.pipeline_plan = {num_stages = 2, num_micro_batches = 1, interleaved = false} + tessera.schedule_digest = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "few", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 2, + tessera.pp_num_micro_batches = 1, + tessera.pp_interleaved = false } { func.func @few(%x: tensor<4x4xbf16>, %w: tensor<4x4xbf16>) -> tensor<4x4xbf16> { %a = "tessera.matmul"(%x, %w) : (tensor<4x4xbf16>, tensor<4x4xbf16>) -> tensor<4x4xbf16> @@ -39,7 +49,12 @@ module attributes { // stage 0 and leaves stage 1 empty — a hole in the send/recv chain. // expected-error@+1 {{PP_EMPTY_STAGE}} module attributes { - tessera.pipeline_plan = {num_stages = 2, num_micro_batches = 2, interleaved = false} + tessera.schedule_digest = "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "empty", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 2, + tessera.pp_num_micro_batches = 2, + tessera.pp_interleaved = false } { func.func @empty(%x: tensor<4x4xbf16>, %w: tensor<4x4xbf16>) -> tensor<4x4xbf16> { %a = "tessera.matmul"(%x, %w) : (tensor<4x4xbf16>, tensor<4x4xbf16>) -> tensor<4x4xbf16> @@ -53,7 +68,12 @@ module attributes { // skipping stage 1 — the adjacent-only insertion never routes it, so the value // crosses a stage boundary with no send/recv (pre-tagged to force the skip). module attributes { - tessera.pipeline_plan = {num_stages = 3, num_micro_batches = 3, interleaved = false} + tessera.schedule_digest = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "skip", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 3, + tessera.pp_num_micro_batches = 3, + tessera.pp_interleaved = false } { func.func @skip(%x: tensor<4x4xbf16>, %w: tensor<4x4xbf16>) -> tensor<4x4xbf16> { // expected-error@+1 {{PP_UNROUTED_CROSS_STAGE_VALUE}} diff --git a/tests/tessera-ir/phase4/pipeline_schedule_materialization.mlir b/tests/tessera-ir/phase4/pipeline_schedule_materialization.mlir index a7972fed1..7bf5b5731 100644 --- a/tests/tessera-ir/phase4/pipeline_schedule_materialization.mlir +++ b/tests/tessera-ir/phase4/pipeline_schedule_materialization.mlir @@ -1,16 +1,16 @@ // RUN: tessera-opt --tessera-pipeline --allow-unregistered-dialect --verify-each=false %s | FileCheck %s // -// The verified 1F1B plan is materialized as an explicit dependency order. -// Each action owns a unique logical clock, so a stage is never asked to run -// forward and backward concurrently. The target runtime may overlap -// independent actions while preserving this serialized semantic order. +// The verified 1F1B Schedule Object is already materialized by the producer. +// The lowering pipeline preserves its dependency carrier and stamps the same +// digest on the owning function and inserted communication operations. module attributes { - tessera.pipeline_plan = { - num_stages = 2, - num_micro_batches = 3, - interleaved = false - } + tessera.schedule_digest = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "root", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 2, + tessera.pp_num_micro_batches = 3, + tessera.pp_interleaved = false } { func.func @pipeline(%x: tensor<64x128xbf16>, %w0: tensor<128x256xbf16>, @@ -23,11 +23,12 @@ module attributes { } } -// CHECK: tessera.pipeline_schedule_kind = "1f1b.serialized_dependency_order.v1" -// CHECK-SAME: tessera.pipeline_steps = [ -// CHECK-SAME: {clock = 0 : i64, micro_batch = 0 : i64, phase = "forward", region = "warmup", stage = 0 : i64} -// CHECK-SAME: {clock = 1 : i64, micro_batch = 0 : i64, phase = "forward", region = "warmup", stage = 1 : i64} -// CHECK-SAME: {clock = 2 : i64, micro_batch = 1 : i64, phase = "forward", region = "steady", stage = 0 : i64} -// CHECK-SAME: {clock = 5 : i64, micro_batch = 0 : i64, phase = "backward", region = "steady", stage = 0 : i64} -// CHECK-SAME: {clock = 10 : i64, micro_batch = 2 : i64, phase = "backward", region = "cooldown", stage = 1 : i64} -// CHECK-SAME: {clock = 11 : i64, micro_batch = 2 : i64, phase = "backward", region = "cooldown", stage = 0 : i64} +// CHECK: sym_name = "pipeline" +// CHECK: tessera.pipeline.send +// CHECK-SAME: tessera.schedule_digest = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +// CHECK: tessera.pipeline.recv +// CHECK-SAME: tessera.schedule_digest = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +// CHECK: }) {tessera.schedule_digest = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +// CHECK: }) {tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1" +// CHECK-SAME: tessera.pipeline_steps = [{action_id = "root" +// CHECK-SAME: tessera.schedule_digest = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" diff --git a/tests/tessera-ir/phase4/pipeline_stage_insertion.mlir b/tests/tessera-ir/phase4/pipeline_stage_insertion.mlir index 8fdb8faf4..48eca8581 100644 --- a/tests/tessera-ir/phase4/pipeline_stage_insertion.mlir +++ b/tests/tessera-ir/phase4/pipeline_stage_insertion.mlir @@ -12,12 +12,12 @@ // so --allow-unregistered-dialect + --verify-each=false round-trip them. module attributes { - tessera.pipeline_plan = { - num_stages = 2, - num_micro_batches = 2, - interleaved = false, - num_chunks = 1 - } + tessera.schedule_digest = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1", + tessera.pipeline_steps = [{action_id = "fixture", clock = 0, depends_on = [], micro_batch = 0, phase = "F", rank = 0, stage = 0}], + tessera.pp_num_stages = 2, + tessera.pp_num_micro_batches = 2, + tessera.pp_interleaved = false } { // The unregistered pipeline.* marker ops force generic module printing, so @@ -33,9 +33,11 @@ module attributes { // CHECK: tessera.pipeline.send // CHECK-SAME: from_stage = 0 // CHECK-SAME: micro_batch + // CHECK-SAME: tessera.schedule_digest = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" // CHECK: tessera.pipeline.recv // CHECK-SAME: micro_batch // CHECK-SAME: to_stage = 1 + // CHECK: }) {tessera.schedule_digest = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" %b = "tessera.matmul"(%a, %w1) {tessera.layer = {stage = 1}} : (tensor<64x256xbf16>, tensor<256x128xbf16>) -> tensor<64x128xbf16> diff --git a/tests/unit/test_composition_cost.py b/tests/unit/test_composition_cost.py index 978b90bcb..c25fb6fc5 100644 --- a/tests/unit/test_composition_cost.py +++ b/tests/unit/test_composition_cost.py @@ -375,7 +375,6 @@ def test_w5_2e_unknown_alias_facts_add_conservative_reference_edges(): [ ("tessera.randn", {"tessera.stochastic_identity": "seed_counter"}, "stochastic_identity"), - ("tessera.all_reduce", {}, "ordered_collective"), ("tessera.scf.if.begin", {"region": "then"}, "region_boundary"), ("tessera.unknown_external", {}, "unregistered_effect"), ], @@ -405,6 +404,67 @@ def test_w5_2e_fail_closed_barriers_serialize_surrounding_actions( assert reason in edge_reasons[("barrier", "after")] +def test_ordered_collectives_order_against_each_other_not_local_work(): + """CORRECTED SEMANTICS (PR #625 review). + + An ordered collective constrains the order OF COLLECTIVES — every rank + must issue them in the same relative order. It does not pin unrelated + local work: data flowing through a collective is ordered by the + SSA/alias/memory-dependence edges instead. The old blanket rule + (serialize against every surrounding op) contradicted the plan's own + overlap requirement — W5.2d's MegaMoE row REQUIRES dispatch(c+1) to + precede combine(c) — and made any pipeline containing transport infer a + total chain, erasing the compute/communication overlap R3 exists to + model. Measured before the fix: 12 MegaMoE actions inferred all 66 + edges of the complete order. + """ + ty = tensor_ir_type(("4",), "fp32") + first = _graph_op("first", "tessera.all_reduce", ["%x"]) + local = _graph_op("local", "tessera.mul", ["%y", "%y"]) + second = _graph_op("second", "tessera.all_reduce", ["%z"]) + fn = GraphIRFunction( + "f", + args=[IRArg("x", ty), IRArg("y", ty), IRArg("z", ty)], + result_types=[ty], body=[first, local, second], + return_values=["%second"], + ) + inferred = infer_action_dag( + fn, (_vector("a"), _vector("b"), _vector("c")), + action_ids=("first", "local", "second"), + ) + edges = { + (edge.predecessor, edge.successor): set(edge.reasons) + for edge in inferred.dependencies + } + # The two collectives keep their relative order... + assert "ordered_collective" in edges[("first", "second")] + # ...and independent local work is free to overlap either of them. + assert ("first", "local") not in edges + assert ("local", "second") not in edges + + +def test_pure_op_with_disjoint_values_has_no_memory_dependence(): + """A REGISTERED pure op touches no hidden state, so it cannot depend on + an effectful op through memory unless they share a value. Without this + every effectful op serialized all surrounding pure work (PR #625).""" + from tessera.compiler.graph_dataflow import analyze_graph_dataflow + + ty = tensor_ir_type(("4",), "fp32") + collective = _graph_op("collective", "tessera.all_reduce", ["%x"]) + disjoint = _graph_op("disjoint", "tessera.mul", ["%y", "%y"]) + shared = _graph_op("shared", "tessera.mul", ["%collective", "%y"]) + fn = GraphIRFunction( + "f", + args=[IRArg("x", ty), IRArg("y", ty)], + result_types=[ty], body=[collective, disjoint, shared], + return_values=["%shared"], + ) + analysis = analyze_graph_dataflow(fn) + assert not analysis.has_memory_dependence(collective, disjoint) + # Sharing a value keeps the dependence — the refinement is alias-gated. + assert analysis.has_memory_dependence(collective, shared) + + @pytest.mark.parametrize("margin", [-0.1, float("nan"), float("inf")]) def test_pruning_rejects_invalid_margin(margin: float): candidate = CompositionCandidate("only", (_action("a", compute_ms=1.0),)) diff --git a/tests/unit/test_megamoe_inferred_dag.py b/tests/unit/test_megamoe_inferred_dag.py new file mode 100644 index 000000000..e7292d291 --- /dev/null +++ b/tests/unit/test_megamoe_inferred_dag.py @@ -0,0 +1,233 @@ +"""SO-3 (integrated-plan queue order 3) — the MegaMoE producer consumes +INFERRED dependence edges; the hand-authored plan DAG is the coverage oracle. + +Before this slice, `composition_candidate_for_megamoe_plan` executed the +hand-written `_action_dependencies` edges directly. Now the chunk pipeline +is represented as Graph IR (`megamoe_graph_function`: dispatch/combine are +the registered ordered collectives, expert compute is pure matmul), the +production candidate's edges come from `infer_action_dag` (W2.1/W2.2 +facts), and `compare_inferred_action_dag` must show the generated edges +cover every hand edge or construction FAILS CLOSED (#31: the hand DAG is +the declared oracle, never silently weakened). Additional conservative +edges — the ordered-collective total order — are reported separately, per +the SO-3 acceptance. The candidate carries the content-addressed +`ScheduleObject`, whose digest is the schedule identity downstream +consumers stamp. +""" + +from __future__ import annotations + +import pytest + +from tessera.compiler import benchmark_row as br +from tessera.compiler.megamoe_overlap import ( + build_megamoe_overlap_plan, + composition_candidate_for_megamoe_plan, + megamoe_graph_function, + megamoe_inferred_composition, + megamoe_issue_order, +) + + +def _plan(num_chunks: int = 4): + return build_megamoe_overlap_plan( + plan_id="so3-megamoe", + num_tokens=96, + num_chunks=num_chunks, + capacities=[32] * num_chunks, + dispatch_buffer_bytes=[4096] * num_chunks, + ) + + +def _rows(plan): + def vector(index): + return { + "schema": br.RESOURCE_VECTOR_SCHEMA, + "usage": br.RESOURCE_VECTOR_USAGE, + "selector_authority": br.SCALAR_SELECTOR_AUTHORITY, + "compute_time_ms": 1.0 + 0.125 * index, + "bytes_moved": 4096, + "communication_bytes": 2048, + "queue_identity": "hip:0", + "resource_identity": "gfx1151", + "timing_provenance": {"source": "device_event", "domain": "device"}, + "artifact_digest": "ab" * 32, + } + + return { + action: {"hot_path_metadata": {"resource_vector": vector(index)}} + for index, action in enumerate(sorted(plan.action_dependencies)) + } + + +def _resource_rows(plan): + """Rows in whichever shape TileAction.from_benchmark_row accepts.""" + rows = _rows(plan) + from tessera.compiler.composition_cost import TileAction + + try: + TileAction.from_benchmark_row("probe:0", next(iter(rows.values()))) + return rows + except Exception: + # Fall back to the flat resource-vector shape. + return { + action: payload["hot_path_metadata"]["resource_vector"] + for action, payload in rows.items() + } + + +def test_issue_order_is_a_deterministic_topological_order(): + plan = _plan() + order = megamoe_issue_order(plan) + assert sorted(order) == sorted(plan.action_dependencies) + position = {action: index for index, action in enumerate(order)} + for action, deps in plan.action_dependencies.items(): + for dep in deps: + assert position[dep] < position[action], (dep, action) + assert order == megamoe_issue_order(plan) + # The overlap requirement is visible in the order itself. + assert position["dispatch:1"] < position["combine:0"] + + +def test_graph_function_uses_registered_collective_semantics(): + from tessera.compiler.effects import Effect, registered_op_effect + + function = megamoe_graph_function(_plan()) + kinds = [op.op_name for op in function.body] + assert kinds.count("tessera.moe_dispatch") == 4 + assert kinds.count("tessera.moe_combine") == 4 + assert kinds.count("tessera.matmul") == 4 + assert registered_op_effect("tessera.moe_dispatch", {}) == Effect.collective + assert registered_op_effect("tessera.moe_combine", {}) == Effect.collective + + +def test_inferred_edges_cover_the_hand_oracle_and_report_extras(): + plan = _plan() + candidate, inferred, parity = megamoe_inferred_composition( + plan, _resource_rows(plan) + ) + assert parity.conservative + assert not parity.missing_reference_edges + # Extra edges are reported separately, never silently merged — but they + # must be the SOUND ones (transitive collective ordering), not a total + # chain. See test_overlap_is_preserved below for the teeth. + assert parity.additional_conservative_edges + for predecessor, successor in parity.additional_conservative_edges: + kinds = {predecessor.split(":")[0], successor.split(":")[0]} + assert kinds <= {"dispatch", "combine"}, (predecessor, successor) + # Every inferred dependency names its supporting reasons. + assert all(dep.reasons for dep in inferred.dependencies) + + +def test_overlap_is_preserved_not_serialized(): + """PR #625 review, P1 — the defect this file previously encoded as + healthy conservatism. + + Emitting the per-chunk transport as ordered collectives used to make + `infer_action_dag` serialize every collective against every surrounding + operation, so a 12-action plan inferred all 66 edges of the complete + order: R3 would then estimate and prune MegaMoE overlap plans as + sequential pipelines, which is exactly the property the plan exists to + choose between. The inferred DAG must keep compute off the collective + chain. + """ + plan = _plan() + _, inferred, _ = megamoe_inferred_composition(plan, _resource_rows(plan)) + actions = len(inferred.actions) + edges = {(d.predecessor, d.successor) for d in inferred.dependencies} + assert len(edges) < actions * (actions - 1) // 2, "inferred a total chain" + # The canonical overlap: compute of the next chunk is NOT forced after the + # previous chunk's combine. + assert ("combine:0", "compute:1") not in edges + # Independent expert compute stays independent. + assert ("compute:0", "compute:1") not in edges + # ...while the collectives keep their required relative order. + assert ("dispatch:0", "dispatch:1") in edges + assert ("combine:0", "combine:1") in edges + + +def test_schedule_digest_binds_plan_identity(): + """PR #625 review, P2 — action ids and graph shape are functions of the + chunk COUNT, so two plans sharing a plan_id and chunk count but differing + in capacities, buffer sizes, token ranges, or the in-flight limit used to + content-address to the SAME digest with identical benchmark rows.""" + left = _plan() + right = build_megamoe_overlap_plan( + plan_id=left.plan_id, + num_tokens=96, + num_chunks=4, + capacities=[16] * 4, + dispatch_buffer_bytes=[8192] * 4, + ) + assert left.artifact_digest != right.artifact_digest + left_candidate = composition_candidate_for_megamoe_plan( + left, _resource_rows(left)) + right_candidate = composition_candidate_for_megamoe_plan( + right, _resource_rows(right)) + assert (left_candidate.schedule_object.digest + != right_candidate.schedule_object.digest) + # The candidate id stays the plan id, so pruning keys are unchanged. + assert left_candidate.candidate_id == left.plan_id + + +def test_candidate_carries_the_content_addressed_schedule_object(): + plan = _plan() + rows = _resource_rows(plan) + first = composition_candidate_for_megamoe_plan(plan, rows) + second = composition_candidate_for_megamoe_plan(plan, rows) + digest = first.schedule_object.digest + assert isinstance(digest, str) and len(digest) == 64 + assert digest == second.schedule_object.digest + assert first.schedule_object.edges # inferred edges, not hand edges + + +def test_generation_regression_fails_closed(monkeypatch): + """Teeth: if the Graph representation ever loses the registered + collective semantics (a generation regression), the inferred DAG stops + covering the hand oracle's ordering edges — and construction must + refuse rather than execute a weaker schedule than the oracle demands.""" + from tessera.compiler import megamoe_overlap as mm + from tessera.compiler.graph_ir import GraphIRFunction, IRArg, IROp, IRType + + plan = _plan() + + def weakened_graph(p): + # Every action becomes a PURE op over per-chunk private values: the + # SSA chains survive, but every ordered-collective and capacity + # edge in the hand DAG disappears from the inference. + tensor = IRType("tensor<*xf32>") + ops = [] + for action_id in mm.megamoe_issue_order(p): + kind, chunk = action_id.split(":") + prev = { + "dispatch": f"%in{chunk}", + "compute": f"%dispatch{chunk}", + "combine": f"%compute{chunk}", + }[kind] + ops.append( + IROp( + result=f"%{kind}{chunk}", + op_name="tessera.tanh", + operands=[prev], + operand_types=["tensor<*xf32>"], + result_type="tensor<*xf32>", + ) + ) + return GraphIRFunction( + name=p.plan_id, + args=[IRArg(f"in{c.index}", tensor) for c in p.chunks], + body=ops, + return_values=[f"%combine{p.num_chunks - 1}"], + ) + + monkeypatch.setattr(mm, "megamoe_graph_function", weakened_graph) + with pytest.raises(ValueError, match="does not cover the hand-authored"): + mm.megamoe_inferred_composition(plan, _resource_rows(plan)) + + +def test_missing_action_evidence_still_fails_closed(): + plan = _plan() + rows = dict(_resource_rows(plan)) + rows.pop("compute:2") + with pytest.raises(ValueError, match="must be total"): + composition_candidate_for_megamoe_plan(plan, rows) diff --git a/tests/unit/test_pipeline_schedule_carrier.py b/tests/unit/test_pipeline_schedule_carrier.py new file mode 100644 index 000000000..1c427842a --- /dev/null +++ b/tests/unit/test_pipeline_schedule_carrier.py @@ -0,0 +1,82 @@ +"""PR #626 review — the pipeline Schedule Object carrier is validated, and +its interleaved dependencies are keyed by the VIRTUAL STAGE. + +Three claims: + +* keying by `(rank, micro_batch, phase)` collided under interleaving, where + one rank owns several virtual stages. Measured before the fix on + num_stages=4/num_chunks=2: 64 of 128 steps shared a key and 32 of 56 + cross-stage forward steps lost their producer edge, so the emitted + Schedule Object permitted a virtual stage to run before the stage feeding + it; +* a producer ordered after its consumer is now a hard error for the forward + chain instead of a silently dropped edge; +* the C++ legality pass parses the carrier rows — unique ids, resolvable + dependencies, topological order, agreement with the declared dimensions — + instead of only checking that the array is non-empty. +""" + +from __future__ import annotations + +import pytest + +from tessera.compiler.pipeline_planner import Phase, PipelinePlan + +INTERLEAVED = dict(num_stages=4, num_micro_batches=8, interleaved=True, num_chunks=2) + + +def _edges(plan): + return {a.action_id: set(a.depends_on) for a in plan.schedule_object.actions} + + +@pytest.mark.parametrize("config", [ + INTERLEAVED, + dict(num_stages=4, num_micro_batches=12, interleaved=True, num_chunks=3), + dict(num_stages=4, num_micro_batches=8, interleaved=False), +]) +def test_every_cross_stage_forward_step_keeps_its_producer_edge(config): + plan = PipelinePlan(**config) + steps = tuple(plan.schedule_steps()) + deps = _edges(plan) + checked = 0 + for step in steps: + if step.phase != Phase.FORWARD or step.stage == 0: + continue + producer = [ + other for other in steps + if other.phase == Phase.FORWARD + and other.stage == step.stage - 1 + and other.micro_batch == step.micro_batch + ] + if not producer: + continue + checked += 1 + assert plan._action_id(producer[0]) in deps[plan._action_id(step)], ( + f"{plan._action_id(step)} lost its cross-stage producer") + assert checked, "expected cross-stage forward steps to check" + + +def test_interleaved_ranks_really_own_several_virtual_stages(): + """The precondition that made the old key collide — pinned so the test + above cannot silently become vacuous.""" + plan = PipelinePlan(**INTERLEAVED) + steps = tuple(plan.schedule_steps()) + stages_per_rank: dict[int, set[int]] = {} + for step in steps: + stages_per_rank.setdefault(step.rank, set()).add(step.stage) + assert all(len(stages) > 1 for stages in stages_per_rank.values()) + rank_keys = {(s.rank, s.micro_batch, s.phase) for s in steps} + stage_keys = {(s.stage, s.micro_batch, s.phase) for s in steps} + assert len(rank_keys) < len(steps) # the old key collided... + assert len(stage_keys) == len(steps) # ...the virtual stage does not + + +def test_every_dependency_precedes_its_consumer(): + for config in (INTERLEAVED, dict(num_stages=2, num_micro_batches=4)): + plan = PipelinePlan(**config) + steps = tuple(plan.schedule_steps()) + order = {plan._action_id(s): i for i, s in enumerate(steps)} + for action in plan.schedule_object.actions: + for dependency in action.depends_on: + assert order[dependency] < order[action.action_id], ( + f"{action.action_id} depends on later {dependency}") diff --git a/tests/unit/test_pipeline_stage_insertion.py b/tests/unit/test_pipeline_stage_insertion.py index 113d49bd5..7aa9a476b 100644 --- a/tests/unit/test_pipeline_stage_insertion.py +++ b/tests/unit/test_pipeline_stage_insertion.py @@ -4,6 +4,7 @@ Tests for PipelinePlan: 1F1B schedule generation, bubble fraction, and interleaved variant. """ + import pytest from tessera.compiler.pipeline_planner import PipelinePlan, Phase, ScheduleStep @@ -16,13 +17,14 @@ def test_single_stage_no_bubble(self): def test_bubble_fraction_standard(self): # (p-1)/m = (4-1)/8 = 0.375 plan = PipelinePlan(num_stages=4, num_micro_batches=8) - assert abs(plan.bubble_fraction - 3/8) < 1e-9 + assert abs(plan.bubble_fraction - 3 / 8) < 1e-9 def test_bubble_fraction_interleaved(self): # (p-1)/(m*v) = (4-1)/(8*2) = 0.1875 - plan = PipelinePlan(num_stages=4, num_micro_batches=8, - interleaved=True, num_chunks=2) - assert abs(plan.bubble_fraction - 3/16) < 1e-9 + plan = PipelinePlan( + num_stages=4, num_micro_batches=8, interleaved=True, num_chunks=2 + ) + assert abs(plan.bubble_fraction - 3 / 16) < 1e-9 def test_warmup_steps(self): plan = PipelinePlan(num_stages=4, num_micro_batches=8) @@ -38,13 +40,15 @@ def test_invalid_micro_batches_raises(self): def test_interleaved_needs_num_chunks_ge_2(self): with pytest.raises(ValueError, match="num_chunks >= 2"): - PipelinePlan(num_stages=4, num_micro_batches=8, - interleaved=True, num_chunks=1) + PipelinePlan( + num_stages=4, num_micro_batches=8, interleaved=True, num_chunks=1 + ) def test_interleaved_needs_enough_micro_batches(self): with pytest.raises(ValueError): - PipelinePlan(num_stages=4, num_micro_batches=4, - interleaved=True, num_chunks=2) # need m >= 8 + PipelinePlan( + num_stages=4, num_micro_batches=4, interleaved=True, num_chunks=2 + ) # need m >= 8 class TestScheduleSteps: @@ -79,14 +83,23 @@ def test_forward_before_backward_same_mb_same_rank(self): steps = self._steps(num_stages=2, num_micro_batches=4) for rank in range(2): for mb in range(4): - fwd_clk = next(s.clock for s in steps - if s.rank == rank and s.micro_batch == mb - and s.phase == Phase.FORWARD) - bwd_clk = next(s.clock for s in steps - if s.rank == rank and s.micro_batch == mb - and s.phase == Phase.BACKWARD) - assert fwd_clk < bwd_clk, \ - f"rank {rank} mb {mb}: fwd@{fwd_clk} not before bwd@{bwd_clk}" + fwd_clk = next( + s.clock + for s in steps + if s.rank == rank + and s.micro_batch == mb + and s.phase == Phase.FORWARD + ) + bwd_clk = next( + s.clock + for s in steps + if s.rank == rank + and s.micro_batch == mb + and s.phase == Phase.BACKWARD + ) + assert ( + fwd_clk < bwd_clk + ), f"rank {rank} mb {mb}: fwd@{fwd_clk} not before bwd@{bwd_clk}" def test_steps_sorted_by_clock(self): steps = self._steps() @@ -102,8 +115,9 @@ def test_rank_0_first_fwd_at_clock_0(self): class TestInterleaved: def test_interleaved_steps_include_all_micro_batches(self): - plan = PipelinePlan(num_stages=2, num_micro_batches=4, - interleaved=True, num_chunks=2) + plan = PipelinePlan( + num_stages=2, num_micro_batches=4, interleaved=True, num_chunks=2 + ) steps = plan.schedule_steps() fwd = [s for s in steps if s.phase == Phase.FORWARD] for rank in range(2): @@ -112,8 +126,9 @@ def test_interleaved_steps_include_all_micro_batches(self): def test_interleaved_bubble_less_than_standard(self): std = PipelinePlan(num_stages=4, num_micro_batches=8) - itr = PipelinePlan(num_stages=4, num_micro_batches=8, - interleaved=True, num_chunks=2) + itr = PipelinePlan( + num_stages=4, num_micro_batches=8, interleaved=True, num_chunks=2 + ) assert itr.bubble_fraction < std.bubble_fraction @@ -134,8 +149,9 @@ def test_to_mlir_attrs_interleaved_false(self): assert "interleaved = false" in attr def test_to_mlir_attrs_interleaved_true(self): - plan = PipelinePlan(num_stages=2, num_micro_batches=4, - interleaved=True, num_chunks=2) + plan = PipelinePlan( + num_stages=2, num_micro_batches=4, interleaved=True, num_chunks=2 + ) attr = plan.to_mlir_attrs() assert "interleaved = true" in attr @@ -145,6 +161,28 @@ def test_repr_includes_bubble(self): assert "bubble" in r assert "37.50%" in r + def test_emits_one_digest_bound_schedule_carrier(self): + plan = PipelinePlan(num_stages=4, num_micro_batches=8) + attrs = plan.to_mlir_attrs() + assert "tessera.pipeline_plan" not in attrs + assert f'tessera.schedule_digest = "{plan.schedule_object.digest}"' in attrs + assert ( + 'tessera.pipeline_schedule_schema = "tessera.pipeline_schedule.v1"' in attrs + ) + assert "tessera.pipeline_steps = [" in attrs + + def test_schedule_object_preserves_roles_resources_and_dependencies(self): + schedule = PipelinePlan(num_stages=4, num_micro_batches=8).schedule_object + assert {role.name for role in schedule.roles} == { + "stage_0", + "stage_1", + "stage_2", + "stage_3", + } + assert all(action.resource_vector for action in schedule.actions) + assert schedule.edges + assert len(schedule.digest) == 64 + class TestDecoupledStage: """Decoupled-stage (local-objective) schedule — DiffusionBlocks-style block @@ -169,16 +207,22 @@ def test_bubble_strictly_less_than_1f1b(self): def test_decoupled_and_interleaved_mutually_exclusive(self): with pytest.raises(ValueError, match="mutually exclusive"): - PipelinePlan(num_stages=4, num_micro_batches=8, - decoupled=True, interleaved=True, num_chunks=2) + PipelinePlan( + num_stages=4, + num_micro_batches=8, + decoupled=True, + interleaved=True, + num_chunks=2, + ) def test_all_micro_batches_covered(self): plan = PipelinePlan(num_stages=4, num_micro_batches=4, decoupled=True) steps = plan.schedule_steps() for phase in (Phase.FORWARD, Phase.BACKWARD): for rank in range(4): - mbs = {s.micro_batch for s in steps - if s.rank == rank and s.phase == phase} + mbs = { + s.micro_batch for s in steps if s.rank == rank and s.phase == phase + } assert mbs == {0, 1, 2, 3} def test_no_cross_stage_activation_dependency(self): @@ -202,10 +246,20 @@ def test_forward_before_backward(self): steps = plan.schedule_steps() for rank in range(3): for mb in range(5): - fwd = next(s.clock for s in steps if s.rank == rank - and s.micro_batch == mb and s.phase == Phase.FORWARD) - bwd = next(s.clock for s in steps if s.rank == rank - and s.micro_batch == mb and s.phase == Phase.BACKWARD) + fwd = next( + s.clock + for s in steps + if s.rank == rank + and s.micro_batch == mb + and s.phase == Phase.FORWARD + ) + bwd = next( + s.clock + for s in steps + if s.rank == rank + and s.micro_batch == mb + and s.phase == Phase.BACKWARD + ) assert fwd < bwd def test_to_mlir_attrs_decoupled_flag(self): diff --git a/tests/unit/test_sharding_propagation.py b/tests/unit/test_sharding_propagation.py index 3b0092afc..f7a9de3e5 100644 --- a/tests/unit/test_sharding_propagation.py +++ b/tests/unit/test_sharding_propagation.py @@ -1,5 +1,6 @@ from __future__ import annotations +import numpy as np import pytest from tessera.compiler.graph_ir import ( @@ -13,6 +14,7 @@ from tessera.compiler.tile_ir import lower_schedule_to_tile_ir from tessera.compiler.sharding_propagation import ( Placement, + execute_resharded_graph_on_mock_mesh, materialize_reshard_plan, plan_explicit_reshards, propagate_sharding, @@ -134,8 +136,11 @@ def test_reshard_unknown_fails_closed(): ty = tensor_ir_type(("8", "16"), "fp32") op = _op("out", "tessera.add", ["%x", "%y"]) fn = GraphIRFunction( - "f", args=[IRArg("x", ty), IRArg("y", ty)], result_types=[ty], - body=[op], return_values=["%out"], + "f", + args=[IRArg("x", ty), IRArg("y", ty)], + result_types=[ty], + body=[op], + return_values=["%out"], ) result = propagate_sharding(fn, {}) try: @@ -153,19 +158,29 @@ def test_registered_collectives_transform_placement_explicitly(): ops = [ _op("gathered", "tessera.all_gather", ["%x"], mesh_axis="data", axis=0), _op( - "scattered", "tessera.reduce_scatter", ["%partial"], - mesh_axis="data", axis=0, + "scattered", + "tessera.reduce_scatter", + ["%partial"], + mesh_axis="data", + axis=0, ), _op( - "exchanged", "tessera.all_to_all", ["%two_d"], - mesh_axis="data", axis=1, scatter_axis=0, gather_axis=1, + "exchanged", + "tessera.all_to_all", + ["%two_d"], + mesh_axis="data", + axis=1, + scatter_axis=0, + gather_axis=1, ), _op("reduced", "tessera.all_reduce", ["%partial"], mesh_axis="data"), ] fn = GraphIRFunction( "collectives", args=[IRArg("x", ty), IRArg("partial", ty), IRArg("two_d", ty)], - result_types=[ty], body=ops, return_values=["%reduced"], + result_types=[ty], + body=ops, + return_values=["%reduced"], ) result = propagate_sharding( fn, @@ -185,7 +200,10 @@ def test_reshard_plan_materializes_graph_schedule_and_tile_ssa(): ty = tensor_ir_type(("8", "16"), "fp32") consume = _op("out", "tessera.sigmoid", ["%x"]) fn = GraphIRFunction( - "f", args=[IRArg("x", ty)], result_types=[ty], body=[consume], + "f", + args=[IRArg("x", ty)], + result_types=[ty], + body=[consume], return_values=["%out"], ) propagated = propagate_sharding(fn, {"x": Placement.tiled({0: "data"})}) @@ -194,7 +212,8 @@ def test_reshard_plan_materializes_graph_schedule_and_tile_ssa(): ) materialized = materialize_reshard_plan(fn, plan) assert [op.op_name for op in materialized.body] == [ - "tessera.all_gather", "tessera.sigmoid" + "tessera.all_gather", + "tessera.sigmoid", ] collective, rewritten = materialized.body assert rewritten.operands == [f"%{collective.result}"] @@ -215,17 +234,20 @@ def test_reshard_plan_materializes_graph_schedule_and_tile_ssa(): def test_all_to_all_gets_verified_matching_rounds_and_region_identity(): - ty = tensor_ir_type(("8", "16"), "fp32") - consume = _op( - "out", "tessera.sigmoid", ["%x"], _region_path=("loop0", "then") - ) + ty = tensor_ir_type(("9", "18"), "fp32") + consume = _op("out", "tessera.sigmoid", ["%x"], _region_path=("loop0", "then")) fn = GraphIRFunction( - "f", args=[IRArg("x", ty)], result_types=[ty], body=[consume], + "f", + args=[IRArg("x", ty)], + result_types=[ty], + body=[consume], return_values=["%out"], ) propagated = propagate_sharding(fn, {"x": Placement.tiled({0: "data"})}) plan = plan_explicit_reshards( - fn, propagated, {0: (Placement.tiled({1: "data"}),)}, + fn, + propagated, + {0: (Placement.tiled({1: "data"}),)}, subgroup=(2, 4, 7), ) assert len(plan.matching_rounds) == 2 @@ -254,12 +276,15 @@ def test_all_to_all_gets_verified_matching_rounds_and_region_identity(): assert "factor every directed peer edge once" in verification.format() -def test_reshard_rejects_sibling_region_escape_and_untyped_local_shard(): +def test_reshard_rejects_sibling_escape_and_materializes_typed_local_shard(): ty = tensor_ir_type(("8", "16"), "fp32") producer = _op("p", "tessera.sigmoid", ["%x"], _region_path=("then",)) consume = _op("out", "tessera.sigmoid", ["%p"], _region_path=("else",)) fn = GraphIRFunction( - "f", args=[IRArg("x", ty)], result_types=[ty], body=[producer, consume], + "f", + args=[IRArg("x", ty), IRArg("y", ty)], + result_types=[ty], + body=[producer, consume], return_values=["%out"], ) propagated = propagate_sharding(fn, {"x": Placement.tiled({0: "data"})}) @@ -269,13 +294,196 @@ def test_reshard_rejects_sibling_region_escape_and_untyped_local_shard(): ) root_fn = GraphIRFunction( - "root", args=[IRArg("x", ty)], result_types=[ty], - body=[_op("out", "tessera.sigmoid", ["%x"])], return_values=["%out"], + "root", + args=[IRArg("x", ty)], + result_types=[ty], + body=[_op("out", "tessera.sigmoid", ["%x"])], + return_values=["%out"], ) root_result = propagate_sharding(root_fn, {"x": Placement.replicated()}) local_plan = plan_explicit_reshards( - root_fn, root_result, {0: (Placement.tiled({0: "data"}),)}, + root_fn, + root_result, + {0: (Placement.tiled({0: "data"}),)}, + subgroup=(0, 1), + ) + materialized = materialize_reshard_plan(root_fn, local_plan) + local_slice, rewritten = materialized.body + assert local_slice.op_name == "tessera.slice" + assert local_slice.kwargs["reshard_kind"] == "local_shard" + assert local_slice.inferred_type.shape == ("4", "16") + assert local_slice.result_type == "tensor<4x16xf32>" + assert rewritten.operand_types == ["tensor<4x16xf32>"] + schedule = lower_graph_to_schedule_ir( + GraphIRModule(functions=[materialized]), target_kind="cpu" + ) + assert schedule.functions[0].body[0].attrs["reshard_kind"] == "local_shard" + + +def _execute_collective(op_name, rank_values, **kwargs): + shape = tuple(str(dim) for dim in np.asarray(rank_values[0]).shape) + ty = tensor_ir_type(shape, "fp32") + op = IROp( + result="out", + op_name=op_name, + operands=["%x"], + operand_types=[str(ty)], + result_type=str(ty), + inferred_type=ty, + kwargs=kwargs, + ) + fn = GraphIRFunction( + "mock_collective", + args=[IRArg("x", ty)], + result_types=[ty], + body=[op], + return_values=["%out"], + ) + return execute_resharded_graph_on_mock_mesh( + fn, {"x": rank_values}, mesh_shape={"data": 2} + ) + + +def test_deterministic_mock_mesh_executes_every_movement_form(): + left = np.arange(8, dtype=np.float32).reshape(4, 2) + right = left + 10 + + reduced = _execute_collective("tessera.all_reduce", (left, right), op="sum") + np.testing.assert_array_equal(reduced.returned["out"][0], left + right) + np.testing.assert_array_equal(reduced.returned["out"][1], left + right) + + scattered = _execute_collective( + "tessera.reduce_scatter", (left, right), axis=0, op="sum" + ) + expected_sum = left + right + np.testing.assert_array_equal(scattered.returned["out"][0], expected_sum[:2]) + np.testing.assert_array_equal(scattered.returned["out"][1], expected_sum[2:]) + + gathered = _execute_collective("tessera.all_gather", (left[:2], right[:2]), axis=0) + expected_gather = np.concatenate((left[:2], right[:2]), axis=0) + np.testing.assert_array_equal(gathered.returned["out"][0], expected_gather) + np.testing.assert_array_equal(gathered.returned["out"][1], expected_gather) + + exchanged = _execute_collective( + "tessera.all_to_all", (left, right), scatter_axis=0, gather_axis=1 + ) + np.testing.assert_array_equal( + exchanged.returned["out"][0], + np.concatenate((left[:2], right[:2]), axis=1), + ) + np.testing.assert_array_equal( + exchanged.returned["out"][1], + np.concatenate((left[2:], right[2:]), axis=1), + ) + + permuted = _execute_collective( + "tessera.collective_permute", + (left, right), + source_peers=[0, 1], + target_peers=[1, 0], + ) + np.testing.assert_array_equal(permuted.returned["out"][0], right) + np.testing.assert_array_equal(permuted.returned["out"][1], left) + assert { + *reduced.executed_reshards, + *scattered.executed_reshards, + *gathered.executed_reshards, + *exchanged.executed_reshards, + *permuted.executed_reshards, + } == { + "all_reduce", + "reduce_scatter", + "all_gather", + "all_to_all", + "collective_permute", + } + + +def test_typed_local_shard_executes_without_hidden_reconstruction(): + ty = tensor_ir_type(("8", "4"), "fp32") + fn = GraphIRFunction( + "local", + args=[IRArg("x", ty)], + result_types=[ty], + body=[_op("out", "tessera.sigmoid", ["%x"])], + return_values=["%out"], + ) + propagation = propagate_sharding(fn, {"x": Placement.replicated()}) + plan = plan_explicit_reshards( + fn, + propagation, + {0: (Placement.tiled({0: "data"}),)}, + subgroup=(0, 1), + ) + materialized = materialize_reshard_plan(fn, plan, mesh_shape={"data": 2}) + full = np.arange(32, dtype=np.float32).reshape(8, 4) + execution = execute_resharded_graph_on_mock_mesh( + materialized, + {"x": (full, full)}, + mesh_shape={"data": 2}, + ) + np.testing.assert_allclose( + execution.returned["out"][0], 1.0 / (1.0 + np.exp(-full[:4])) + ) + np.testing.assert_allclose( + execution.returned["out"][1], 1.0 / (1.0 + np.exp(-full[4:])) + ) + assert execution.executed_reshards == ("local_shard",) + + +def test_collective_permute_peer_map_survives_schedule_and_tile(): + values = (np.ones((2, 2), dtype=np.float32),) * 2 + shape = tuple(str(dim) for dim in values[0].shape) + ty = tensor_ir_type(shape, "fp32") + op = IROp( + result="out", + op_name="tessera.collective_permute", + operands=["%x"], + operand_types=[str(ty)], + result_type=str(ty), + inferred_type=ty, + kwargs={ + "mesh_axis": "data", + "axis": 0, + "source_peers": [0, 1], + "target_peers": [1, 0], + }, + ) + fn = GraphIRFunction( + "permute", + args=[IRArg("x", ty)], + result_types=[ty], + body=[op], + return_values=["%out"], + ) + schedule = lower_graph_to_schedule_ir( + GraphIRModule(functions=[fn]), target_kind="cpu" + ) + scheduled = schedule.functions[0].body[0] + assert scheduled.attrs["source_peers"] == [0, 1] + assert scheduled.attrs["target_peers"] == [1, 0] + tile = lower_schedule_to_tile_ir(schedule) + assert tile.functions[0].body[0].op_name == "tile.collective_permute" + assert tile.verify().ok + + +def test_typed_local_shard_rejects_nondivisible_extent_and_mesh_mismatch(): + ty = tensor_ir_type(("7", "4"), "fp32") + fn = GraphIRFunction( + "bad_local", + args=[IRArg("x", ty)], + result_types=[ty], + body=[_op("out", "tessera.sigmoid", ["%x"])], + return_values=["%out"], + ) + propagation = propagate_sharding(fn, {"x": Placement.replicated()}) + plan = plan_explicit_reshards( + fn, + propagation, + {0: (Placement.tiled({0: "data"}),)}, subgroup=(0, 1), ) - with pytest.raises(ValueError, match="typed mesh-size shape contract"): - materialize_reshard_plan(root_fn, local_plan) + with pytest.raises(ValueError, match="not divisible"): + materialize_reshard_plan(fn, plan, mesh_shape={"data": 2}) + with pytest.raises(ValueError, match="subgroup size"): + materialize_reshard_plan(fn, plan, mesh_shape={"data": 4}) diff --git a/tests/unit/test_spectral_inferred_dag.py b/tests/unit/test_spectral_inferred_dag.py new file mode 100644 index 000000000..da43ec8b5 --- /dev/null +++ b/tests/unit/test_spectral_inferred_dag.py @@ -0,0 +1,128 @@ +"""SO-3: spectral physical producers consume inferred action edges.""" + +from __future__ import annotations + +import hashlib +import json +import re + +import pytest + +from tessera.compiler.scheduled_matmul import find_tessera_opt, run_tessera_opt +from tessera.compiler.scheduled_spectral import ( + infer_spectral_action_dag, + lower_scheduled_spectral, +) + + +@pytest.mark.parametrize( + ("op_name", "dct_type"), + [ + ("tessera.spectral_filter", 0), + ("tessera.dct", 2), + ("tessera.dct", 4), + ("tessera.spectral_conv", 0), + ("tessera.stft", 0), + ("tessera.istft", 0), + ], +) +def test_spectral_families_infer_reasoned_edges_and_schedule_evidence( + op_name, dct_type +): + inferred, schedule = infer_spectral_action_dag( + semantic_digest="ab" * 32, + target="rocm", + architecture="gfx1151", + op_name=op_name, + dct_type=dct_type, + workspace_bytes=4096, + ) + assert inferred.dependencies + assert all(edge.reasons for edge in inferred.dependencies) + assert schedule.edges == inferred.schedule_object.edges + assert {role.name for role in schedule.roles} == { + "spectral_compute", + "spectral_queue", + } + assert all(action.resource_vector for action in schedule.actions) + assert len(schedule.digest) == 64 + + +@pytest.mark.skipif(find_tessera_opt() is None, reason="requires tessera-opt") +def test_spectral_lowering_stamps_the_schedule_object_digest_end_to_end(): + artifact = lower_scheduled_spectral( + target="x86", + op_name="tessera.spectral_conv", + input_shapes=((2, 13), (2, 7)), + ) + encoded = json.dumps( + artifact.schedule_object, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + assert hashlib.sha256(encoded).hexdigest() == artifact.schedule_digest + assert ( + f'tessera.schedule_digest = "{artifact.schedule_digest}"' + in artifact.schedule_ir + ) + assert f'tessera.schedule_hash = "{artifact.schedule_digest}"' in artifact.tile_ir + assert artifact.graph_analysis_digest + + +@pytest.mark.skipif(find_tessera_opt() is None, reason="requires tessera-opt") +def test_spectral_lowering_rejects_a_stale_module_schedule_digest(): + artifact = lower_scheduled_spectral( + target="x86", + op_name="tessera.spectral_filter", + input_shapes=((2, 17), (2, 17)), + ) + stale = artifact.schedule_ir.replace( + f'tessera.schedule_digest = "{artifact.schedule_digest}"', + f'tessera.schedule_digest = "{"0" * 64}"', + 1, + ) + with pytest.raises(RuntimeError, match="module Schedule Object digest"): + run_tessera_opt(find_tessera_opt(), stale, "--tessera-schedule-to-tile") + + +@pytest.mark.skipif(find_tessera_opt() is None, reason="requires tessera-opt") +def test_spectral_lowering_rejects_policy_changed_behind_a_matching_digest(): + """PR #626 review, P1 — the fail-open the deleted C++ digest helper left. + + Comparing `artifact_hash` to the module digest does NOT bind the + attributes the pass then consumes: a cached or hand-edited program can + keep BOTH digest strings and still change `workspace_bytes` or + `native_entry`, and the native launch would use the changed value (a + wrong buffer contract or an incompatible entry point). The lowering now + re-verifies every consumed policy value against the semantic payload the + producer carries on the module. + """ + artifact = lower_scheduled_spectral( + target="x86", + op_name="tessera.spectral_filter", + input_shapes=((2, 17), (2, 17)), + ) + match = re.search(r"workspace_bytes = (\d+) : i64", artifact.schedule_ir) + assert match, "expected a workspace_bytes attribute to tamper with" + tampered = artifact.schedule_ir.replace( + match.group(0), f"workspace_bytes = {int(match.group(1)) + 4096} : i64", 1 + ) + # Both digests are untouched — only the consumed policy changed. + assert f'tessera.schedule_digest = "{artifact.schedule_digest}"' in tampered + with pytest.raises(RuntimeError, match="disagrees with the carried semantic"): + run_tessera_opt(find_tessera_opt(), tampered, "--tessera-schedule-to-tile") + + +@pytest.mark.skipif(find_tessera_opt() is None, reason="requires tessera-opt") +def test_spectral_lowering_requires_the_semantic_payload(): + """The payload itself is mandatory: without it the consumer cannot + re-verify what it consumes, so it fails closed rather than trusting the + digest string.""" + artifact = lower_scheduled_spectral( + target="x86", + op_name="tessera.spectral_filter", + input_shapes=((2, 17), (2, 17)), + ) + stripped = re.sub(r', tessera\.spectral_semantic = "[^"]*"', "", + artifact.schedule_ir, count=1) + assert "tessera.spectral_semantic" not in stripped + with pytest.raises(RuntimeError, match="tessera.spectral_semantic"): + run_tessera_opt(find_tessera_opt(), stripped, "--tessera-schedule-to-tile")